From 88340f057ab35a41f4f0d8a365745d2ce049522a Mon Sep 17 00:00:00 2001 From: HyWang Date: Thu, 10 Sep 2026 13:44:06 +0000 Subject: [PATCH 1/2] Add P6 T04 shared/residual merge reference and conformance tests --- docs/.nav.yml | 1 + docs/operators/shared-residual-merge.md | 151 +++++++ rl_engine/kernels/ops/pytorch/moe/__init__.py | 24 + .../ops/pytorch/moe/shared_residual_merge.py | 398 +++++++++++++++++ rl_engine/kernels/registry.py | 23 + tests/fixtures/p6_t04_merge.json | 33 ++ tests/test_shared_residual_merge.py | 420 ++++++++++++++++++ 7 files changed, 1050 insertions(+) create mode 100644 docs/operators/shared-residual-merge.md create mode 100644 rl_engine/kernels/ops/pytorch/moe/__init__.py create mode 100644 rl_engine/kernels/ops/pytorch/moe/shared_residual_merge.py create mode 100644 tests/fixtures/p6_t04_merge.json create mode 100644 tests/test_shared_residual_merge.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 73281687..e50c7524 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -15,6 +15,7 @@ nav: - Operators: - operators/README.md - operators/activation.md + - operators/shared-residual-merge.md - operators/attention.md - operators/fused-logp.md - operators/linear-logp.md diff --git a/docs/operators/shared-residual-merge.md b/docs/operators/shared-residual-merge.md new file mode 100644 index 00000000..afd818c8 --- /dev/null +++ b/docs/operators/shared-residual-merge.md @@ -0,0 +1,151 @@ +# P6 / T04 Shared and Residual Merge Reference + +T04 merges three token-local rows using exactly `(routed + shared) + residual`. +Both additions materialize FP32; only the final output is converted to BF16. +The routed row is already weighted by P5 and summed by T03. This operator never +multiplies route weights, runs a collective, computes gradients, or performs mHC +post-mixing. It is an eager **forward reference**, not a fused production kernel. + +## Interface and scope + +```python +from rl_engine.kernels.ops.pytorch.moe import shared_residual_merge_fwd + +result = shared_residual_merge_fwd( + routed, shared, residual, context=context, debug=True, +) +y = result.token_moe_row_bf16 +receipt = result.receipt +``` + +| Value | Required shape / type | +| --- | --- | +| `routed` | Contiguous `[T,H]` FP32, already combined in canonical slot order | +| `shared`, `residual` | Same shape and device; BF16 or FP32 | +| `context` | `MergeContext`, carrying validated producer receipts and expected identities | +| `y` | `[T,H]` BF16, no autograd graph | + +CPU and CUDA-device PyTorch eager execution are implemented; ROCm uses PyTorch's +CUDA device API. Hardware qualification requires running the named GPU tests. +Ascend, tensor subclasses, non-contiguous tensors, empty batches, broadcasting, +FP16 inputs, compiled execution, and CUDA Graph capture are rejected. A zero-count +EP rank does not call the local merge; its participation belongs to the P4/T08 gate. + +The semantic registry exposes the explicit backend +`rlkernel.moe.shared_residual_merge.reference.v1` for `shared_residual_merge`. +It is marked `reference_only`, `forward_only`, and not production-certified. It +is not inserted into generic fallback priorities or a model's FFN path. A caller +requesting strict GPU launch observability cannot use this reference as proof. + +## Evidence and exactly-once + +`MergeIdentity` binds case/run/pass, checkpoint and weight fingerprints, opaque +RoutePlan/ExchangePlan/CombinePlan fingerprints, fixture checksum, and ordered +global token IDs. `MergeSource` attaches a role, a stable expected source ID, +the same identity, and a checksum over the tensor's dtype/shape/raw bytes. +The caller obtains these from validated producer artifacts. The reference does +not create a RoutePlan, ExchangePlan, CombinePlan, inverse map, or input adapter. + +`MergeContext` requires one source for each of routed/shared/residual. It also +requires an upstream `route_weight` event exactly once, P5 ownership and weighted +input, and no prior shared/residual application or P6 local downcast. That local +downcast history does **not** count BF16 storage conversions inside P5. The fixed +local order is shared, residual, final BF16 cast. Unknown operations and mHC post +input roles fail closed. + +All rows in a call have the same application history and must belong to the +executing token owner. Shared-expert replicas are recorded separately; their +number never scales the shared contribution. `rank` and replica coordinates are +caller-supplied P4 evidence, not measurements of a process group. The reference +cannot prove rank completeness, global uniqueness of ownership across independent +calls, or validate a false upstream claim. T08/live integration must establish +those facts. Repeating a pure reference call for comparison is allowed; the +output receipt's executed events prevent reusing its result as an unmerged input +when receipts are propagated correctly. + +Identity and discrete checks precede tensor checksum and numerical checks. The +function then checks finite values at every boundary, including BF16 overflow. +Receipt counts derive from recorded local events and observed ATen dispatches; +they are not a certificate that upstream metadata was truthful. + +The three boundary keys are: + +- `P6.merge.after_shared` +- `P6.merge.after_residual` +- `P6.merge.final_bf16` + +Every boundary records identity, rank, shape/dtype/layout, event index, backend, +implementation hash and tensor checksum. Slot is not applicable at this token-row +boundary; T03 owns the slot-order hash. `merge_order_hash` fingerprints only T04's +versioned three operations and is not a replacement for P6's full combine hash. +Debug off returns only hashes and the output. Debug on also returns intermediate +tensors; the execution and hashes are identical in both modes. Treat returned +tensors as immutable when associating them with receipts. + +ATen adds and casts are observed through `TorchDispatchMode`. Actual GPU kernel +symbols, tile/warp/stage/vector width are not exposed by this layer and are +explicitly unobserved, never populated with guessed configuration values. The +reference synchronizes and copies tensor bytes to CPU for evidence, so it must +not be used as a production latency baseline. + +## Validation and handoff + +Use the existing pytest entry point from the repository root: + +```bash +python3 -m pytest -q tests/test_shared_residual_merge.py + +# Retain append-only local reference artifacts with the existing ArtifactStore. +RL_KERNEL_T04_ARTIFACT_DIR=/tmp/p6-t04-evidence \ + python3 -m pytest -q tests/test_shared_residual_merge.py \ + --junitxml=/tmp/p6-t04-tests.xml + +# Run the named GPU slice on a GPU host; skips do not constitute a pass. +python3 -m pytest -q tests/test_shared_residual_merge.py -k gpu_golden +``` + +`tests/fixtures/p6_t04_merge.json` contains only synthetic inputs and manually +specified intermediate FP32 values and BF16 bit patterns. Its SHA256 is pinned +in the test. Goldens cover cancellation/order, early-rounding loss, signed zero, +FP32 and BF16 ties, largest finite BF16, and smallest normal BF16. Negative tests +exercise source duplication/missing data, identity drift, ownership, mHC misuse, +invalid types/layouts, non-finite inputs and intermediate/final overflow. +Batch/chunk/permutation/padding checks inspect raw bytes, preserving signed zero. + +EP=1/2/4/8 cases simulate token ownership and shared replication on CPU. They +are local semantic tests, **not** a multi-GPU WS2 certificate. CUDA/ROCm golden +tests compare the same final and intermediate bit patterns when hardware exists. +No T05 fused-forward, T06 backward or live producer claim is made by this suite. + +The artifact test uses the existing `ArtifactStore` and its completion seal, +checksum validation and immutable experiment identity. It writes T04 receipts +and debug bundles for the three synthetic golden cases. Retained files can be +read with `json` and `ArtifactStore.load_tensor_bundle`; each boundary checksum +is validated after reloading. A seal means complete local artifacts, not complete +P6 certification. Reusing an evidence root after a fixture, context or build +change fails resume; use a new experiment root for a new revision. + +The broad `check_operator.py` dtype/tolerance harness is not modified: this +checked reference requires producer receipts and strict raw-byte tests, while +T01 has not yet frozen the P6 runner. T07 should consume these receipts in its +existing comparator/store. T05 should compare all three boundaries, not just the +final BF16 tensor. T03 is the designated reviewer for the T04 handoff. + +## Contract delta for T01 review + +| Field | Proposal | +| --- | --- | +| Current contract / ABI | `p6-task-contract.v1` / `foundation-moe-return.v1` | +| Proposed local receipt version | `p6.t04.merge-receipt.v1-proposal` | +| Change | Add `MergeIdentity`, `MergeSource`, `MergeContext`, `MergeResult` and keyword `context`/`debug` at the reference boundary | +| Compatibility | Does not change arithmetic, weight owner, return layout, gradient ownership or Foundation ABI | +| Affected consumers | T01, T03 reviewer, T05, T07, T08 | +| Input dtype profile | Routed FP32; shared/residual BF16 or FP32; T01 must approve this admitted subset | +| Fixture IDs | `T04-LOCAL-*` are provisional local IDs; T01 assigns canonical catalog IDs/checksums | +| Blocking scope | Formal P6 ABI certification and live integration; local reference can be tested independently | + +No T01 start kit or public CombinePlan implementation was present in the base +checkout. These small receipt dataclasses are an explicit review proposal rather +than a claim to be that missing schema. Map them to the canonical start kit once +available; do not maintain a competing backend-specific protocol. Review, +approved fixture IDs, GPU evidence and live integration remain separate gates. diff --git a/rl_engine/kernels/ops/pytorch/moe/__init__.py b/rl_engine/kernels/ops/pytorch/moe/__init__.py new file mode 100644 index 00000000..4d17c436 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/moe/__init__.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .shared_residual_merge import ( + MergeContext, + MergeContractError, + MergeIdentity, + MergeResult, + MergeSource, + SharedResidualMergeReferenceOp, + shared_residual_merge_fwd, + tensor_sha256, +) + +__all__ = [ + "MergeContext", + "MergeContractError", + "MergeIdentity", + "MergeResult", + "MergeSource", + "SharedResidualMergeReferenceOp", + "shared_residual_merge_fwd", + "tensor_sha256", +] diff --git a/rl_engine/kernels/ops/pytorch/moe/shared_residual_merge.py b/rl_engine/kernels/ops/pytorch/moe/shared_residual_merge.py new file mode 100644 index 00000000..87008d8c --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/moe/shared_residual_merge.py @@ -0,0 +1,398 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""P6/T04 eager forward reference; no routing, collective, or mHC arithmetic. + +The receipt/context interface is a T01 review proposal, not a second CombinePlan. +Upstream validation binds its opaque plan fingerprint and source receipts. This +module validates the local merge boundary; it cannot authenticate dishonest +producer receipts or certify unobserved GPU launch parameters. +""" + +from __future__ import annotations + +import hashlib +import json +from collections import Counter +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import torch +from torch import Tensor +from torch.utils._python_dispatch import TorchDispatchMode + +CONTRACT_VERSION = "p6-task-contract.v1" +FOUNDATION_ABI = "foundation-moe-return.v1" +RECEIPT_VERSION = "p6.t04.merge-receipt.v1-proposal" +BACKEND_ID = "rlkernel.moe.shared_residual_merge.reference.v1" +ORDER = ("shared", "residual", "cast_bf16") +BOUNDARIES = ( + "P6.merge.after_shared", + "P6.merge.after_residual", + "P6.merge.final_bf16", +) + + +class MergeContractError(ValueError): + def __init__(self, status: str, field: str): + self.status = status + self.field = field + super().__init__(f"{status}: {field}") + + +def _require(condition: bool, status: str, field: str) -> None: + if not condition: + raise MergeContractError(status, field) + + +def _hash_json(value: Any) -> str: + return hashlib.sha256( + json.dumps(value, sort_keys=True, separators=(",", ":"), allow_nan=False).encode() + ).hexdigest() + + +def tensor_sha256(value: Tensor) -> str: + """Hash dtype, shape and raw logical bytes, including the sign bit of zero.""" + snapshot = value.detach().cpu().contiguous() + digest = hashlib.sha256() + digest.update(_hash_json([str(snapshot.dtype), list(snapshot.shape)]).encode()) + digest.update(snapshot.reshape(-1).view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +@dataclass(frozen=True) +class MergeIdentity: + case_id: str + run_id: str + pass_id: str + checkpoint_fingerprint: str + weight_fingerprint: str + route_plan_fingerprint: str + exchange_plan_fingerprint: str + combine_plan_fingerprint: str + fixture_checksum: str + global_token_ids: tuple[int, ...] + + +@dataclass(frozen=True) +class MergeSource: + """Producer evidence supplied with a tensor, never synthesized by the merge.""" + + role: str + source_id: str + identity: MergeIdentity + tensor_checksum: str + + +@dataclass(frozen=True) +class MergeContext: + identity: MergeIdentity + # Exact source ids expected by the caller, in routed/shared/residual order. + expected_source_ids: tuple[str, str, str] + sources: tuple[MergeSource, ...] + # Validated upstream application events for this token batch. All rows must + # share this history; heterogeneous histories require separate calls. + upstream_events: tuple[str, ...] + rank: int + token_owner_ranks: tuple[int, ...] + shared_replica_ranks: tuple[int, ...] + contract_version: str = CONTRACT_VERSION + foundation_abi: str = FOUNDATION_ABI + receipt_version: str = RECEIPT_VERSION + merge_order: tuple[str, ...] = ORDER + input_row_weighted: bool = True + route_weight_owner: str = "P5" + + +@dataclass +class MergeResult: + token_moe_row_bf16: Tensor + receipt: dict[str, Any] + debug_boundaries: dict[str, Tensor] + + +def _validate_context(context: MergeContext) -> dict[str, MergeSource]: + _require(isinstance(context, MergeContext), "INVALID_COMBINE_PLAN", "context") + for field, expected in ( + ("contract_version", CONTRACT_VERSION), + ("foundation_abi", FOUNDATION_ABI), + ("receipt_version", RECEIPT_VERSION), + ): + _require(getattr(context, field) == expected, "SCHEMA_MISMATCH", field) + identity = context.identity + _require(isinstance(identity, MergeIdentity), "IDENTITY_DRIFT", "identity") + for field, value in asdict(identity).items(): + if field != "global_token_ids": + _require(isinstance(value, str) and bool(value.strip()), "IDENTITY_DRIFT", field) + tokens = identity.global_token_ids + _require( + isinstance(tokens, tuple) + and bool(tokens) + and all(type(t) is int and t >= 0 for t in tokens) + and len(set(tokens)) == len(tokens), + "AMBIGUOUS_GLOBAL_TOKEN_MAPPING", + "global_token_ids", + ) + expected_ids = context.expected_source_ids + _require( + isinstance(expected_ids, tuple) + and len(expected_ids) == 3 + and all(isinstance(s, str) and s.strip() for s in expected_ids) + and len(set(expected_ids)) == 3, + "IDENTITY_DRIFT", + "expected_source_ids", + ) + _require( + isinstance(context.sources, tuple) + and all(isinstance(s, MergeSource) for s in context.sources), + "IDENTITY_DRIFT", + "sources", + ) + # Compare identities before looking at payload values or arithmetic policy. + for source in context.sources: + _require( + isinstance(source.role, str) + and isinstance(source.source_id, str) + and isinstance(source.tensor_checksum, str), + "IDENTITY_DRIFT", + "source fields", + ) + _require(source.identity == identity, "IDENTITY_DRIFT", f"{source.role}.identity") + roles = Counter(s.role for s in context.sources) + _require( + not (set(roles) - {"routed", "shared", "residual"}), + "INVALID_COMBINE_PLAN", + "sources.role (mHC post is not a merge input)", + ) + for role, status in ( + ("shared", "SHARED_APPLIED_NOT_ONCE"), + ("residual", "RESIDUAL_APPLIED_NOT_ONCE"), + ("routed", "INVALID_COMBINE_PLAN"), + ): + _require(roles[role] == 1, status, f"sources.{role}") + sources = {s.role: s for s in context.sources} + for role, source_id in zip(("routed", "shared", "residual"), expected_ids, strict=True): + _require(sources[role].source_id == source_id, "IDENTITY_DRIFT", f"{role}.source_id") + + events = context.upstream_events + _require( + context.input_row_weighted is True and context.route_weight_owner == "P5", + "INVALID_COMBINE_PLAN", + "weighted row owner", + ) + _require( + isinstance(events, tuple) and all(isinstance(e, str) for e in events), + "INVALID_COMBINE_PLAN", + "upstream_events", + ) + counts = Counter(events) + _require( + counts["route_weight"] == 1, + "INVALID_COMBINE_PLAN" if counts["route_weight"] == 0 else "ROUTE_WEIGHT_APPLIED_TWICE", + "route_weight", + ) + for name, status in ( + ("shared", "SHARED_APPLIED_NOT_ONCE"), + ("residual", "RESIDUAL_APPLIED_NOT_ONCE"), + ("cast_bf16", "EARLY_OR_MULTIPLE_DOWNCAST"), + ): + _require(counts[name] == 0, status, f"upstream_events.{name}") + _require(set(events) == {"route_weight"}, "INVALID_COMBINE_PLAN", "upstream_events") + _require(context.merge_order == ORDER, "ADDITION_ORDER_MISMATCH", "merge_order") + _require(type(context.rank) is int and context.rank >= 0, "MISSING_PROVENANCE", "rank") + _require( + isinstance(context.token_owner_ranks, tuple) + and len(context.token_owner_ranks) == len(tokens) + and all(type(r) is int and r == context.rank for r in context.token_owner_ranks), + "SHARED_APPLIED_NOT_ONCE", + "token_owner_ranks", + ) + replicas = context.shared_replica_ranks + _require( + isinstance(replicas, tuple) + and bool(replicas) + and all(type(r) is int and r >= 0 for r in replicas) + and len(set(replicas)) == len(replicas), + "INVALID_COMBINE_PLAN", + "shared_replica_ranks", + ) + return sources + + +def _finite(value: Tensor, boundary: str) -> None: + _require(bool(torch.isfinite(value).all().item()), "NON_FINITE", boundary) + + +class _ArithmeticTrace(TorchDispatchMode): + """Observe actual ATen adds/casts without replacing the underlying operations.""" + + def __init__(self): + super().__init__() + self.operations: list[dict[str, Any]] = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + result = func(*args, **(kwargs or {})) + if func in (torch.ops.aten.add.Tensor, torch.ops.aten._to_copy.default): + self.operations.append( + { + "operator": str(func), + "output_dtype": str(result.dtype), + "shape": list(result.shape), + } + ) + return result + + +@torch.no_grad() +def shared_residual_merge_fwd( + routed: Tensor, + shared: Tensor, + residual: Tensor, + *, + context: MergeContext, + debug: bool = False, +) -> MergeResult: + """Merge an owner-local batch; return BF16 output and hash-only evidence. + + CPU and CUDA eager reference only. Hashing synchronizes/copies to CPU; this + is deliberately not a production kernel, CUDA Graph path, or backward op. + ``context`` is required: three bare tensors cannot prove exactly-once. + """ + _require(not torch.compiler.is_compiling(), "UNSUPPORTED_CAPABILITY", "compiled reference") + sources = _validate_context(context) + inputs = {"routed": routed, "shared": shared, "residual": residual} + for role, value in inputs.items(): + _require(type(value) is Tensor, "UNSUPPORTED_CAPABILITY", f"{role}.tensor") + _require( + value.layout == torch.strided and value.ndim == 2 and value.is_contiguous(), + "UNSUPPORTED_CAPABILITY", + f"{role}.layout", + ) + _require(value.device.type in {"cpu", "cuda"}, "UNSUPPORTED_CAPABILITY", f"{role}.device") + _require( + value.shape[0] == len(context.identity.global_token_ids) and value.shape[1] > 0, + "INVALID_COMBINE_PLAN", + f"{role}.shape", + ) + _require( + value.shape == routed.shape and value.device == routed.device, + "INVALID_COMBINE_PLAN", + f"{role}.shape/device", + ) + if role == "routed": + _require(value.dtype == torch.float32, "EARLY_OR_MULTIPLE_DOWNCAST", "routed.dtype") + else: + _require( + value.dtype in {torch.float32, torch.bfloat16}, + "UNSUPPORTED_CAPABILITY", + f"{role}.dtype", + ) + if routed.device.type == "cuda": + with torch.cuda.device(routed.device): + _require( + not torch.cuda.is_current_stream_capturing(), + "UNSUPPORTED_CAPABILITY", + "CUDA Graph reference", + ) + for role, value in inputs.items(): + _require( + tensor_sha256(value) == sources[role].tensor_checksum, + "IDENTITY_DRIFT", + f"{role}.tensor_checksum", + ) + for role, value in inputs.items(): + _finite(value, role) + + events: list[str] = [] + # Disable an enclosing autocast context; each eager add materializes FP32. + trace = _ArithmeticTrace() + with torch.autocast(device_type=routed.device.type, enabled=False), trace: + after_shared = torch.add(routed, shared.float()) + events.append("shared") + _finite(after_shared, BOUNDARIES[0]) + after_residual = torch.add(after_shared, residual.float()) + events.append("residual") + _finite(after_residual, BOUNDARIES[1]) + output = after_residual.to(torch.bfloat16) + events.append("cast_bf16") + _finite(output, BOUNDARIES[2]) + + actual_adds = [op for op in trace.operations if op["operator"] == "aten.add.Tensor"] + actual_downcasts = [ + op + for op in trace.operations + if op["operator"] == "aten._to_copy.default" and op["output_dtype"] == "torch.bfloat16" + ] + _require( + len(actual_adds) == 2 and all(op["output_dtype"] == "torch.float32" for op in actual_adds), + "ADDITION_ORDER_MISMATCH", + "observed ATen adds", + ) + _require(len(actual_downcasts) == 1, "EARLY_OR_MULTIPLE_DOWNCAST", "observed ATen casts") + + boundary_tensors = dict(zip(BOUNDARIES, (after_shared, after_residual, output), strict=True)) + implementation_hash = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + provenance = { + "backend": BACKEND_ID, + "implementation_sha256": implementation_hash, + "torch_version": str(torch.__version__), + "torch_git_version": torch.version.git_version, + "device": str(routed.device), + "device_name": ( + torch.cuda.get_device_name(routed.device) if routed.device.type == "cuda" else "cpu" + ), + "build_runtime": torch.version.hip or torch.version.cuda, + "observed_aten_operations": trace.operations, + "accumulator_dtype": "torch.float32", + "output_dtype": "torch.bfloat16", + "launch_parameters": None, + "launch_provenance_status": "NOT_OBSERVED_AT_REFERENCE_LEVEL", + } + boundaries = [ + { + "key": key, + "event_index": index, + "dtype": str(value.dtype), + "shape": list(value.shape), + "layout": "contiguous", + "checksum": tensor_sha256(value), + "rank": context.rank, + "identity": asdict(context.identity), + "backend": BACKEND_ID, + "implementation_sha256": implementation_hash, + } + for index, (key, value) in enumerate(boundary_tensors.items()) + ] + applied = Counter(context.upstream_events + tuple(events)) + receipt = { + "schema_version": RECEIPT_VERSION, + "contract_version": CONTRACT_VERSION, + "foundation_abi": FOUNDATION_ABI, + "identity": asdict(context.identity), + "sources": [asdict(sources[role]) for role in inputs], + "upstream_events": list(context.upstream_events), + "input_row_weighted": context.input_row_weighted, + "route_weight_owner": context.route_weight_owner, + "executed_events": events, + "route_weight_applied_count": applied["route_weight"], + "shared_applied_count": applied["shared"], + "residual_applied_count": applied["residual"], + "local_downcast_count": len(actual_downcasts), + "merge_order_hash": _hash_json({"version": RECEIPT_VERSION, "order": events}), + "rank": context.rank, + "token_owner_ranks": list(context.token_owner_ranks), + "shared_replica_ranks": list(context.shared_replica_ranks), + "identity_discrete_gate": "PASS", + "scope": "local_forward_reference", + "certification": "INCOMPLETE: T01 ABI approval and GPU/integration evidence required", + "actual_provenance": provenance, + "boundaries": boundaries, + } + return MergeResult(output, receipt, boundary_tensors if debug else {}) + + +class SharedResidualMergeReferenceOp: + """Explicit reference backend factory for the existing semantic registry.""" + + __call__ = staticmethod(shared_residual_merge_fwd) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 7070728a..fb1b3716 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -184,6 +184,29 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: return ( + OperatorBackendDescriptor( + semantic_op="shared_residual_merge", + backend_id="rlkernel.moe.shared_residual_merge.reference.v1", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu", "cuda", "rocm"}), + supported_dtypes=frozenset({"float32"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "algorithm": "fp32_routed_plus_shared_then_residual_final_bf16", + "reference_only": True, + "forward_only": True, + "requires_merge_context": True, + "strict_observable": True, + "gpu_launch_observable": False, + "certification": "pending_T01_and_GPU_evidence", + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=( + "rl_engine.kernels.ops.pytorch.moe.SharedResidualMergeReferenceOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="p6.t04.merge-receipt.v1-proposal", + ), OperatorBackendDescriptor( semantic_op="selected_logprob", backend_id="rlkernel.reference_logp", diff --git a/tests/fixtures/p6_t04_merge.json b/tests/fixtures/p6_t04_merge.json new file mode 100644 index 00000000..0f9f4eae --- /dev/null +++ b/tests/fixtures/p6_t04_merge.json @@ -0,0 +1,33 @@ +{ + "schema_version": "p6.t04.synthetic-fixtures.v1-proposal", + "approval": "T01 case IDs and checksum approval pending", + "cases": [ + { + "id": "T04-LOCAL-MERGE", + "routed": [[1.0, 2.0, 3.0]], + "shared": [[2.0, 3.0, 4.0]], + "residual": [[3.0, 4.0, 5.0]], + "after_shared": [[3.0, 5.0, 7.0]], + "after_residual": [[6.0, 9.0, 12.0]], + "bf16_words": [[16576, 16656, 16704]] + }, + { + "id": "T04-LOCAL-ORDER-CAST-ZERO", + "routed": [[16777216.0, 256.0, 1.0, -0.0]], + "shared": [[1.0, 1.0, 0.000000059604644775390625, -0.0]], + "residual": [[-16777216.0, -256.0, 0.0, -0.0]], + "after_shared": [[16777216.0, 257.0, 1.0, -0.0]], + "after_residual": [[0.0, 1.0, 1.0, -0.0]], + "bf16_words": [[0, 16256, 16256, 32768]] + }, + { + "id": "T04-LOCAL-BF16-TIES", + "routed": [[1.00390625, 1.01171875, 3.3895313892515355e38, 1.1754943508222875e-38]], + "shared": [[0.0, 0.0, 0.0, 0.0]], + "residual": [[0.0, 0.0, 0.0, 0.0]], + "after_shared": [[1.00390625, 1.01171875, 3.3895313892515355e38, 1.1754943508222875e-38]], + "after_residual": [[1.00390625, 1.01171875, 3.3895313892515355e38, 1.1754943508222875e-38]], + "bf16_words": [[16256, 16258, 32639, 128]] + } + ] +} diff --git a/tests/test_shared_residual_merge.py b/tests/test_shared_residual_merge.py new file mode 100644 index 00000000..6d7189fe --- /dev/null +++ b/tests/test_shared_residual_merge.py @@ -0,0 +1,420 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""T04 reference acceptance; fixed goldens, injected faults, and owner-local EP mocks.""" + +from __future__ import annotations + +import hashlib +import json +import os +import shutil +from dataclasses import asdict, replace +from pathlib import Path + +import pytest +import torch + +from rl_engine.alignment.cross_config.artifacts import ArtifactError, ArtifactStore +from rl_engine.kernels.ops.pytorch.moe import ( + MergeContext, + MergeContractError, + MergeIdentity, + MergeSource, + shared_residual_merge_fwd, + tensor_sha256, +) +from rl_engine.kernels.ops.pytorch.moe.shared_residual_merge import BOUNDARIES, ORDER + +FIXTURE = Path(__file__).parent / "fixtures" / "p6_t04_merge.json" +FIXTURE_CHECKSUM = hashlib.sha256(FIXTURE.read_bytes()).hexdigest() +CASES = json.loads(FIXTURE.read_text())["cases"] + + +def test_fixture_checksum(): + assert FIXTURE_CHECKSUM == "79813b6999c56489b4e08575fd71ed62d6d65cdd7ad61a90d7727a83bceb20d4" + + +def make_inputs(case=CASES[0], *, device="cpu", dtype=torch.float32): + return tuple( + torch.tensor(case[role], dtype=torch.float32 if role == "routed" else dtype, device=device) + for role in ("routed", "shared", "residual") + ) + + +def make_context(inputs, *, case_id="T04-LOCAL-MERGE", rank=0, ep=1, tokens=None): + identity = MergeIdentity( + case_id=case_id, + run_id="synthetic-t04", + pass_id="forward-0", + checkpoint_fingerprint="synthetic-no-checkpoint", + weight_fingerprint="synthetic-no-weights", + route_plan_fingerprint="synthetic-route", + exchange_plan_fingerprint="synthetic-exchange", + combine_plan_fingerprint="synthetic-combine", + fixture_checksum=FIXTURE_CHECKSUM, + global_token_ids=tokens or tuple(range(inputs[0].shape[0])), + ) + sources = tuple( + MergeSource(role, f"synthetic-{role}", identity, tensor_sha256(value)) + for role, value in zip(("routed", "shared", "residual"), inputs, strict=True) + ) + return MergeContext( + identity=identity, + expected_source_ids=tuple(s.source_id for s in sources), + sources=sources, + upstream_events=("route_weight",), + rank=rank, + token_owner_ranks=(rank,) * len(identity.global_token_ids), + shared_replica_ranks=tuple(range(ep)), + ) + + +def assert_bytes(a, b): + assert a.shape == b.shape and a.dtype == b.dtype + assert torch.equal( + a.detach().cpu().contiguous().view(torch.uint8), + b.detach().cpu().contiguous().view(torch.uint8), + ) + + +@pytest.mark.parametrize("case", CASES, ids=lambda c: c["id"]) +def test_frozen_golden_and_debug(case): + inputs = make_inputs(case) + context = make_context(inputs, case_id=case["id"]) + result = shared_residual_merge_fwd(*inputs, context=context, debug=True) + golden = torch.tensor(case["bf16_words"], dtype=torch.uint16).view(torch.bfloat16) + assert_bytes(result.token_moe_row_bf16, golden) + for key, name in zip(BOUNDARIES[:2], ("after_shared", "after_residual"), strict=True): + assert_bytes(result.debug_boundaries[key], torch.tensor(case[name], dtype=torch.float32)) + quiet = shared_residual_merge_fwd(*inputs, context=context) + assert quiet.debug_boundaries == {} + assert quiet.receipt == result.receipt + assert_bytes(quiet.token_moe_row_bf16, golden) + for count in ( + "route_weight_applied_count", + "shared_applied_count", + "residual_applied_count", + "local_downcast_count", + ): + assert result.receipt[count] == 1 + assert result.receipt["executed_events"] == list(ORDER) + + +def test_wrong_arithmetic_is_detectable(): + r, s, x = make_inputs(CASES[1]) + good = shared_residual_merge_fwd(r, s, x, context=make_context((r, s, x))) + wrong_order = (r + (s + x)).bfloat16() + early_cast = ((r + s).bfloat16().float() + x).bfloat16() + assert wrong_order[0, 0].item() == 1 and good.token_moe_row_bf16[0, 0].item() == 0 + assert early_cast[0, 1].item() == 0 and good.token_moe_row_bf16[0, 1].item() == 1 + + +@pytest.mark.parametrize( + "event,status", + [ + ("shared", "SHARED_APPLIED_NOT_ONCE"), + ("residual", "RESIDUAL_APPLIED_NOT_ONCE"), + ("cast_bf16", "EARLY_OR_MULTIPLE_DOWNCAST"), + ("route_weight", "ROUTE_WEIGHT_APPLIED_TWICE"), + ("mhc_post", "INVALID_COMBINE_PLAN"), + ], +) +def test_reject_prior_application(event, status): + inputs = make_inputs() + context = make_context(inputs) + context = replace(context, upstream_events=context.upstream_events + (event,)) + with pytest.raises(MergeContractError, match=status): + shared_residual_merge_fwd(*inputs, context=context) + + +@pytest.mark.parametrize( + "role,status", + [("shared", "SHARED_APPLIED_NOT_ONCE"), ("residual", "RESIDUAL_APPLIED_NOT_ONCE")], +) +@pytest.mark.parametrize("duplicate", [False, True]) +def test_missing_duplicate_source(role, status, duplicate): + inputs = make_inputs() + context = make_context(inputs) + source = next(s for s in context.sources if s.role == role) + sources = ( + context.sources + (source,) + if duplicate + else tuple(s for s in context.sources if s.role != role) + ) + with pytest.raises(MergeContractError, match=status): + shared_residual_merge_fwd(*inputs, context=replace(context, sources=sources)) + + +@pytest.mark.parametrize( + "field", + [ + "case_id", + "run_id", + "pass_id", + "checkpoint_fingerprint", + "weight_fingerprint", + "route_plan_fingerprint", + "exchange_plan_fingerprint", + "combine_plan_fingerprint", + "fixture_checksum", + "global_token_ids", + ], +) +def test_identity_precedes_numerics(field): + inputs = make_inputs() + context = make_context(inputs) + source = context.sources[1] + identity = replace( + source.identity, **{field: (42,) if field == "global_token_ids" else "wrong"} + ) + context = replace( + context, + sources=(context.sources[0], replace(source, identity=identity), context.sources[2]), + ) + inputs[0].fill_(float("nan")) + with pytest.raises(MergeContractError, match="IDENTITY_DRIFT: shared.identity"): + shared_residual_merge_fwd(*inputs, context=context) + + +@pytest.mark.parametrize( + "mutation,status", + [ + ({"merge_order": ("residual", "shared", "cast_bf16")}, "ADDITION_ORDER_MISMATCH"), + ({"upstream_events": ()}, "INVALID_COMBINE_PLAN"), + ({"token_owner_ranks": (1,)}, "SHARED_APPLIED_NOT_ONCE"), + ({"rank": -1}, "MISSING_PROVENANCE"), + ({"shared_replica_ranks": (0, 0)}, "INVALID_COMBINE_PLAN"), + ({"foundation_abi": "wrong"}, "SCHEMA_MISMATCH"), + ({"receipt_version": "wrong"}, "SCHEMA_MISMATCH"), + ({"input_row_weighted": False}, "INVALID_COMBINE_PLAN"), + ({"route_weight_owner": "P6"}, "INVALID_COMBINE_PLAN"), + ( + {"expected_source_ids": ("wrong", "synthetic-shared", "synthetic-residual")}, + "IDENTITY_DRIFT", + ), + ], +) +def test_invalid_context(mutation, status): + inputs = make_inputs() + with pytest.raises(MergeContractError, match=status): + shared_residual_merge_fwd(*inputs, context=replace(make_context(inputs), **mutation)) + + +def test_mhc_post_input_rejected(): + inputs = make_inputs() + context = make_context(inputs) + context = replace( + context, sources=(*context.sources[:2], replace(context.sources[2], role="mhc_post")) + ) + with pytest.raises(MergeContractError, match="INVALID_COMBINE_PLAN"): + shared_residual_merge_fwd(*inputs, context=context) + + +def test_checksum_detects_stale_tensor(): + inputs = make_inputs() + context = make_context(inputs) + inputs[1][0, 0] += 1 + with pytest.raises(MergeContractError, match="IDENTITY_DRIFT: shared.tensor_checksum"): + shared_residual_merge_fwd(*inputs, context=context) + + +@pytest.mark.parametrize("role", range(3)) +@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")]) +def test_nonfinite_inputs(role, value): + inputs = make_inputs() + inputs[role][0, 0] = value + with pytest.raises(MergeContractError, match="NON_FINITE"): + shared_residual_merge_fwd(*inputs, context=make_context(inputs)) + + +@pytest.mark.parametrize("stage", ["after_shared", "after_residual", "final_bf16"]) +def test_intermediate_and_cast_overflow(stage): + maxval = torch.finfo(torch.float32).max + values = { + "after_shared": (maxval, maxval, -maxval), + "after_residual": (maxval, 0.0, maxval), + "final_bf16": (maxval, 0.0, 0.0), + }[stage] + inputs = tuple(torch.tensor([[v]], dtype=torch.float32) for v in values) + with pytest.raises(MergeContractError, match=f"NON_FINITE: P6.merge.{stage}"): + shared_residual_merge_fwd(*inputs, context=make_context(inputs)) + + +@pytest.mark.parametrize("kind", ["routed_bf16", "shared_fp16", "broadcast", "strided"]) +def test_unsupported_payloads(kind): + inputs = list(make_inputs()) + if kind == "routed_bf16": + inputs[0] = inputs[0].bfloat16() + elif kind == "shared_fp16": + inputs[1] = inputs[1].half() + elif kind == "broadcast": + inputs[1] = inputs[1][:, :1].contiguous() + else: + inputs[1] = torch.ones(1, 6)[:, ::2] + with pytest.raises(MergeContractError): + shared_residual_merge_fwd(*inputs, context=make_context(inputs)) + + +@pytest.mark.parametrize("ep", [1, 2, 4, 8]) +def test_ep_replication_owner_local_mock(ep): + # No collectives: replicas exist on every mock rank; exactly one owner merges + # each token. This is WS2 semantic evidence, not multi-GPU certification. + inputs = tuple(v.repeat(16, 1) for v in make_inputs()) + canonical = shared_residual_merge_fwd(*inputs, context=make_context(inputs)) + gathered = torch.empty_like(canonical.token_moe_row_bf16) + for rank in range(ep): + tokens = tuple(range(rank, 16, ep)) + local = tuple(v[list(tokens)].contiguous() for v in inputs) + context = make_context(local, rank=rank, ep=ep, tokens=tokens) + result = shared_residual_merge_fwd(*local, context=context) + gathered[list(tokens)] = result.token_moe_row_bf16 + assert result.receipt["shared_applied_count"] == 1 + assert result.receipt["merge_order_hash"] == canonical.receipt["merge_order_hash"] + if ep > 1: + with pytest.raises(MergeContractError, match="SHARED_APPLIED_NOT_ONCE"): + shared_residual_merge_fwd(*local, context=replace(context, rank=(rank + 1) % ep)) + assert_bytes(gathered, canonical.token_moe_row_bf16) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_input_immutability_and_autocast(dtype): + inputs = make_inputs(dtype=dtype) + snapshots = tuple(v.clone() for v in inputs) + context = make_context(inputs) + baseline = shared_residual_merge_fwd(*inputs, context=context) + with torch.autocast("cpu", dtype=torch.bfloat16): + result = shared_residual_merge_fwd(*inputs, context=context) + assert_bytes(result.token_moe_row_bf16, baseline.token_moe_row_bf16) + for value, saved in zip(inputs, snapshots, strict=True): + assert_bytes(value, saved) + operations = result.receipt["actual_provenance"]["observed_aten_operations"] + assert [op["output_dtype"] for op in operations if op["operator"] == "aten.add.Tensor"] == [ + "torch.float32", + "torch.float32", + ] + assert sum(op["output_dtype"] == "torch.bfloat16" for op in operations) == 1 + + +@pytest.mark.parametrize("width", [1, 17, 129, 4096]) +def test_batch_chunk_padding_permutation(width): + generator = torch.Generator().manual_seed(71) + inputs = tuple(torch.randn(7, width, generator=generator) for _ in range(3)) + expected = shared_residual_merge_fwd(*inputs, context=make_context(inputs)).token_moe_row_bf16 + permutation = [6, 2, 0, 1, 5, 3, 4] + permuted = tuple(v[permutation] for v in inputs) + result = shared_residual_merge_fwd( + *permuted, context=make_context(permuted, tokens=tuple(permutation)) + ) + assert_bytes(result.token_moe_row_bf16, expected[permutation]) + for start, end in ((0, 1), (1, 4), (4, 7)): + chunk = tuple(v[start:end] for v in inputs) + result = shared_residual_merge_fwd( + *chunk, context=make_context(chunk, tokens=tuple(range(start, end))) + ) + assert_bytes(result.token_moe_row_bf16, expected[start:end]) + padded = tuple(torch.cat((v, torch.zeros(3, width))) for v in inputs) + result = shared_residual_merge_fwd(*padded, context=make_context(padded)) + assert_bytes(result.token_moe_row_bf16[:7], expected) + + +def test_duplicate_and_empty_token_batch(): + inputs = tuple(v.repeat(2, 1) for v in make_inputs()) + context = make_context(inputs) + context = replace(context, identity=replace(context.identity, global_token_ids=(1, 1))) + with pytest.raises(MergeContractError, match="AMBIGUOUS_GLOBAL_TOKEN_MAPPING"): + shared_residual_merge_fwd(*inputs, context=context) + empty = tuple(v[:0] for v in inputs) + with pytest.raises(MergeContractError, match="AMBIGUOUS_GLOBAL_TOKEN_MAPPING"): + shared_residual_merge_fwd(*empty, context=make_context(empty)) + + +def test_reference_registry_entry(): + from rl_engine.kernels.ops.pytorch.moe.shared_residual_merge import BACKEND_ID + from rl_engine.kernels.registry import KernelRegistry + from rl_engine.kernels.semantic_registry import ( + OperatorRequirements, + OperatorResolutionError, + OperatorSession, + ) + + session = OperatorSession(KernelRegistry().semantic) + resolution = session.resolve( + semantic_op="shared_residual_merge", + requested_backend=BACKEND_ID, + target="rollout", + requirements=OperatorRequirements(device="cpu", dtype="float32"), + ) + op = session.instantiate(resolution) + inputs = make_inputs() + result = op(*inputs, context=make_context(inputs)) + assert result.receipt["scope"] == "local_forward_reference" + assert resolution.descriptor.determinism_or_alignment_properties["reference_only"] is True + with pytest.raises(OperatorResolutionError): + session.resolve( + semantic_op="shared_residual_merge", + requested_backend=BACKEND_ID, + target="rollout", + requirements=OperatorRequirements( + device="cpu", dtype="float32", alignment_properties={"gpu_launch_observable": True} + ), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA/ROCm GPU: not certified") +@pytest.mark.parametrize("case", CASES, ids=lambda c: c["id"]) +def test_gpu_golden(case): + inputs = make_inputs(case, device="cuda") + result = shared_residual_merge_fwd( + *inputs, context=make_context(inputs, case_id=case["id"]), debug=True + ) + golden = torch.tensor(case["bf16_words"], dtype=torch.uint16).view(torch.bfloat16) + assert_bytes(result.token_moe_row_bf16, golden) + for key, field in zip(BOUNDARIES[:2], ("after_shared", "after_residual"), strict=True): + assert_bytes(result.debug_boundaries[key], torch.tensor(case[field], dtype=torch.float32)) + with torch.autocast("cuda", dtype=torch.bfloat16): + quiet = shared_residual_merge_fwd(*inputs, context=make_context(inputs, case_id=case["id"])) + assert_bytes(quiet.token_moe_row_bf16, result.token_moe_row_bf16) + assert quiet.receipt == result.receipt + + +def test_sealed_reference_evidence(tmp_path): + root = Path(os.environ.get("RL_KERNEL_T04_ARTIFACT_DIR", tmp_path)) + store = ArtifactStore(root) + required = ("t04_receipt.json", "t04_debug.pt") + for case in CASES: + inputs = make_inputs(case) + context = make_context(inputs, case_id=case["id"]) + result = shared_residual_merge_fwd(*inputs, context=context, debug=True) + experiment = { + "context": json.loads(json.dumps(asdict(context))), + "implementation": result.receipt["actual_provenance"], + } + # Identity/build changes cannot silently resume an existing experiment. + store.initialize_experiment(case["id"], experiment=experiment, plan=[]) + attempt = store.create_attempt(case["id"], case["id"]) + store.write_json(attempt, required[0], result.receipt) + store.write_tensor_bundle(attempt, required[1], result.debug_boundaries) + store.complete_attempt( + attempt, + required=required, + summary={ + "schema_version": "cross_config.complete.v1", + "case_id": case["id"], + "attempt_id": attempt.name, + "status": "LOCAL_REFERENCE_ONLY", + }, + ) + store.validate_completed_attempt(attempt, required=required, expected_case_id=case["id"]) + loaded = store.load_tensor_bundle(attempt / required[1])["tensors"] + receipt = json.loads((attempt / required[0]).read_text()) + for boundary in receipt["boundaries"]: + assert tensor_sha256(loaded[boundary["key"]]) == boundary["checksum"] + with pytest.raises(ArtifactError, match="resume metadata differs"): + store.initialize_experiment(case["id"], experiment={"wrong_run": True}, plan=[]) + # Exercise corruption on an isolated copy, never alter a sealed report. + corrupted = tmp_path / "corrupted" / attempt.name + shutil.copytree(attempt, corrupted) + (corrupted / "t04_receipt.json").write_text("{}") + with pytest.raises(ArtifactError, match="artifact hash does not match"): + store.validate_completed_attempt(corrupted, required=required) From 0bf63f7aba510ee55e9b841ce623591d08edaa7e Mon Sep 17 00:00:00 2001 From: HyWang Date: Fri, 11 Sep 2026 07:10:30 +0000 Subject: [PATCH 2/2] Refine MoE merge reference API and validation coverage --- docs/.nav.yml | 2 +- docs/operators/moe-merge.md | 102 +++ docs/operators/shared-residual-merge.md | 151 ----- .../testing/moe_merge.py} | 186 ++---- rl_engine/kernels/ops/pytorch/moe/__init__.py | 22 +- rl_engine/kernels/ops/pytorch/moe/merge.py | 64 ++ rl_engine/kernels/registry.py | 14 +- .../{p6_t04_merge.json => moe_merge.json} | 9 +- tests/test_moe_merge.py | 607 ++++++++++++++++++ tests/test_shared_residual_merge.py | 420 ------------ 10 files changed, 836 insertions(+), 741 deletions(-) create mode 100644 docs/operators/moe-merge.md delete mode 100644 docs/operators/shared-residual-merge.md rename rl_engine/{kernels/ops/pytorch/moe/shared_residual_merge.py => alignment/testing/moe_merge.py} (58%) create mode 100644 rl_engine/kernels/ops/pytorch/moe/merge.py rename tests/fixtures/{p6_t04_merge.json => moe_merge.json} (81%) create mode 100644 tests/test_moe_merge.py delete mode 100644 tests/test_shared_residual_merge.py diff --git a/docs/.nav.yml b/docs/.nav.yml index e50c7524..92cb2256 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -15,7 +15,7 @@ nav: - Operators: - operators/README.md - operators/activation.md - - operators/shared-residual-merge.md + - operators/moe-merge.md - operators/attention.md - operators/fused-logp.md - operators/linear-logp.md diff --git a/docs/operators/moe-merge.md b/docs/operators/moe-merge.md new file mode 100644 index 00000000..396d6803 --- /dev/null +++ b/docs/operators/moe-merge.md @@ -0,0 +1,102 @@ +# MoE Merge + +`shared_residual_merge_fwd` adds the shared-expert output and residual to the combined routed-expert +output. Both additions use FP32 in the fixed order `(routed + shared) + residual`; +only the final result is converted to BF16. Fixing this order prevents different +rounding results between callers that otherwise implement the same formula. + +## Interface + +```python +from rl_engine.kernels.ops.pytorch.moe import MoeMergeOp, shared_residual_merge_fwd + +y = shared_residual_merge_fwd(routed, shared, residual) + +# The registry uses the same nn.Module / forward interface as other native ops. +op = MoeMergeOp() +y = op(routed, shared, residual) +``` + +| Value | Shape and dtype | +| --- | --- | +| `routed` | Contiguous, non-empty `[T, H]` FP32; routing weights and expert combination already applied | +| `shared`, `residual` | Same shape and device as `routed`; FP32 or BF16 | +| `y` | `[T, H]` BF16; no autograd graph | + +The implementation supports CPU and CUDA-device PyTorch eager execution; ROCm +uses PyTorch's CUDA device API. It rejects broadcasting, non-contiguous tensors, +unsupported dtypes, compiled execution and CUDA Graph capture. This is a +forward-only reference. It performs no routing, collective or post-merge mixing. +Token ownership and source provenance are responsibilities of the caller. + +The computation path performs tensor metadata checks and arithmetic. It does not +trace dispatches, hash tensors, build receipts or inspect source files. Values +follow PyTorch arithmetic, including non-finite propagation; the validation path +below rejects non-finite inputs and results. + +## Intermediate results + +The reference exposes the same computation with its two FP32 intermediates: + +```python +after_shared, after_residual, output = op.forward_with_intermediates( + routed, shared, residual, +) +``` + +`forward` and `forward_with_intermediates` share the arithmetic implementation. +The validation helper names these tensors `after_shared`, `after_residual` and +`final_bf16`. Collecting them does not change the output bytes. + +## Registration + +The existing semantic registry resolves `shared_residual_merge` to the explicit +backend `rlkernel.moe.shared_residual_merge.reference.v1`. Its factory constructs +`MoeMergeOp`; implementation provenance comes from `OperatorSession`. + +The descriptor declares deterministic, batch-invariant, forward-only reference +arithmetic. Its dtype is the routed input's FP32 dtype; shared and residual dtypes +are validated by the operator. The empty topology capability requires a local +request. Distributed callers validate their topology and invoke the operator on +each token owner's rows. GPU launch parameters are not observable at this layer. + +## Validation + +Run the focused tests from the repository root: + +```bash +python3 -m pytest -q tests/test_moe_merge.py + +# Retain receipts and tensor snapshots through the existing ArtifactStore. +RL_KERNEL_MOE_MERGE_ARTIFACT_DIR=/tmp/moe-merge-evidence \ + python3 -m pytest -q tests/test_moe_merge.py +``` + +`tests/fixtures/moe_merge.json` contains synthetic inputs, manually specified FP32 +intermediates and exact BF16 words. Tests cover addition order, early rounding, +signed zero, ties, dtype limits, autocast, input immutability and batch/chunk/ +permutation/padding invariance. Comparisons use raw bytes, including signed zero. + +`rl_engine.alignment.testing.moe_merge` holds the identity/count gate and evidence +collection used by these tests. It validates source identities, checksums, token +ownership, weighted routed inputs and application history before numeric checks. +Missing or repeated sources, prior shared/residual application, early downcasts, +and post-merge mixing inputs are rejected. This metadata is not part of the tensor +operator's API. + +The receipt records the validated upstream history, reference schedule, actual +intermediate tensor checksums and the shared registry's implementation provenance. +Application counts describe that reference schedule; they are not measured ATen +or GPU launch counts. Fixed numerical fixtures validate the arithmetic. A receipt +cannot establish that an upstream producer's claims are truthful. + +CPU tests simulate expert-parallel replication at degrees 1, 2, 4 and 8. GPU tests +start one NCCL/RCCL process per device at those degrees, merge uneven owner-local +rows and compare gathered token IDs and output bytes against an independent +integer oracle. Insufficient GPU counts are skipped. Communication exists only in +the test harness; these checks do not establish full model integration. + +Artifact tests use the existing `ArtifactStore` for append-only writes, completion +seals, reload checks and corruption detection. Receipts and boundary tensors use +`merge_receipt.json` and `merge_debug.pt`. Use a new evidence directory after +changing fixtures, metadata or implementation; completed artifacts remain immutable. diff --git a/docs/operators/shared-residual-merge.md b/docs/operators/shared-residual-merge.md deleted file mode 100644 index afd818c8..00000000 --- a/docs/operators/shared-residual-merge.md +++ /dev/null @@ -1,151 +0,0 @@ -# P6 / T04 Shared and Residual Merge Reference - -T04 merges three token-local rows using exactly `(routed + shared) + residual`. -Both additions materialize FP32; only the final output is converted to BF16. -The routed row is already weighted by P5 and summed by T03. This operator never -multiplies route weights, runs a collective, computes gradients, or performs mHC -post-mixing. It is an eager **forward reference**, not a fused production kernel. - -## Interface and scope - -```python -from rl_engine.kernels.ops.pytorch.moe import shared_residual_merge_fwd - -result = shared_residual_merge_fwd( - routed, shared, residual, context=context, debug=True, -) -y = result.token_moe_row_bf16 -receipt = result.receipt -``` - -| Value | Required shape / type | -| --- | --- | -| `routed` | Contiguous `[T,H]` FP32, already combined in canonical slot order | -| `shared`, `residual` | Same shape and device; BF16 or FP32 | -| `context` | `MergeContext`, carrying validated producer receipts and expected identities | -| `y` | `[T,H]` BF16, no autograd graph | - -CPU and CUDA-device PyTorch eager execution are implemented; ROCm uses PyTorch's -CUDA device API. Hardware qualification requires running the named GPU tests. -Ascend, tensor subclasses, non-contiguous tensors, empty batches, broadcasting, -FP16 inputs, compiled execution, and CUDA Graph capture are rejected. A zero-count -EP rank does not call the local merge; its participation belongs to the P4/T08 gate. - -The semantic registry exposes the explicit backend -`rlkernel.moe.shared_residual_merge.reference.v1` for `shared_residual_merge`. -It is marked `reference_only`, `forward_only`, and not production-certified. It -is not inserted into generic fallback priorities or a model's FFN path. A caller -requesting strict GPU launch observability cannot use this reference as proof. - -## Evidence and exactly-once - -`MergeIdentity` binds case/run/pass, checkpoint and weight fingerprints, opaque -RoutePlan/ExchangePlan/CombinePlan fingerprints, fixture checksum, and ordered -global token IDs. `MergeSource` attaches a role, a stable expected source ID, -the same identity, and a checksum over the tensor's dtype/shape/raw bytes. -The caller obtains these from validated producer artifacts. The reference does -not create a RoutePlan, ExchangePlan, CombinePlan, inverse map, or input adapter. - -`MergeContext` requires one source for each of routed/shared/residual. It also -requires an upstream `route_weight` event exactly once, P5 ownership and weighted -input, and no prior shared/residual application or P6 local downcast. That local -downcast history does **not** count BF16 storage conversions inside P5. The fixed -local order is shared, residual, final BF16 cast. Unknown operations and mHC post -input roles fail closed. - -All rows in a call have the same application history and must belong to the -executing token owner. Shared-expert replicas are recorded separately; their -number never scales the shared contribution. `rank` and replica coordinates are -caller-supplied P4 evidence, not measurements of a process group. The reference -cannot prove rank completeness, global uniqueness of ownership across independent -calls, or validate a false upstream claim. T08/live integration must establish -those facts. Repeating a pure reference call for comparison is allowed; the -output receipt's executed events prevent reusing its result as an unmerged input -when receipts are propagated correctly. - -Identity and discrete checks precede tensor checksum and numerical checks. The -function then checks finite values at every boundary, including BF16 overflow. -Receipt counts derive from recorded local events and observed ATen dispatches; -they are not a certificate that upstream metadata was truthful. - -The three boundary keys are: - -- `P6.merge.after_shared` -- `P6.merge.after_residual` -- `P6.merge.final_bf16` - -Every boundary records identity, rank, shape/dtype/layout, event index, backend, -implementation hash and tensor checksum. Slot is not applicable at this token-row -boundary; T03 owns the slot-order hash. `merge_order_hash` fingerprints only T04's -versioned three operations and is not a replacement for P6's full combine hash. -Debug off returns only hashes and the output. Debug on also returns intermediate -tensors; the execution and hashes are identical in both modes. Treat returned -tensors as immutable when associating them with receipts. - -ATen adds and casts are observed through `TorchDispatchMode`. Actual GPU kernel -symbols, tile/warp/stage/vector width are not exposed by this layer and are -explicitly unobserved, never populated with guessed configuration values. The -reference synchronizes and copies tensor bytes to CPU for evidence, so it must -not be used as a production latency baseline. - -## Validation and handoff - -Use the existing pytest entry point from the repository root: - -```bash -python3 -m pytest -q tests/test_shared_residual_merge.py - -# Retain append-only local reference artifacts with the existing ArtifactStore. -RL_KERNEL_T04_ARTIFACT_DIR=/tmp/p6-t04-evidence \ - python3 -m pytest -q tests/test_shared_residual_merge.py \ - --junitxml=/tmp/p6-t04-tests.xml - -# Run the named GPU slice on a GPU host; skips do not constitute a pass. -python3 -m pytest -q tests/test_shared_residual_merge.py -k gpu_golden -``` - -`tests/fixtures/p6_t04_merge.json` contains only synthetic inputs and manually -specified intermediate FP32 values and BF16 bit patterns. Its SHA256 is pinned -in the test. Goldens cover cancellation/order, early-rounding loss, signed zero, -FP32 and BF16 ties, largest finite BF16, and smallest normal BF16. Negative tests -exercise source duplication/missing data, identity drift, ownership, mHC misuse, -invalid types/layouts, non-finite inputs and intermediate/final overflow. -Batch/chunk/permutation/padding checks inspect raw bytes, preserving signed zero. - -EP=1/2/4/8 cases simulate token ownership and shared replication on CPU. They -are local semantic tests, **not** a multi-GPU WS2 certificate. CUDA/ROCm golden -tests compare the same final and intermediate bit patterns when hardware exists. -No T05 fused-forward, T06 backward or live producer claim is made by this suite. - -The artifact test uses the existing `ArtifactStore` and its completion seal, -checksum validation and immutable experiment identity. It writes T04 receipts -and debug bundles for the three synthetic golden cases. Retained files can be -read with `json` and `ArtifactStore.load_tensor_bundle`; each boundary checksum -is validated after reloading. A seal means complete local artifacts, not complete -P6 certification. Reusing an evidence root after a fixture, context or build -change fails resume; use a new experiment root for a new revision. - -The broad `check_operator.py` dtype/tolerance harness is not modified: this -checked reference requires producer receipts and strict raw-byte tests, while -T01 has not yet frozen the P6 runner. T07 should consume these receipts in its -existing comparator/store. T05 should compare all three boundaries, not just the -final BF16 tensor. T03 is the designated reviewer for the T04 handoff. - -## Contract delta for T01 review - -| Field | Proposal | -| --- | --- | -| Current contract / ABI | `p6-task-contract.v1` / `foundation-moe-return.v1` | -| Proposed local receipt version | `p6.t04.merge-receipt.v1-proposal` | -| Change | Add `MergeIdentity`, `MergeSource`, `MergeContext`, `MergeResult` and keyword `context`/`debug` at the reference boundary | -| Compatibility | Does not change arithmetic, weight owner, return layout, gradient ownership or Foundation ABI | -| Affected consumers | T01, T03 reviewer, T05, T07, T08 | -| Input dtype profile | Routed FP32; shared/residual BF16 or FP32; T01 must approve this admitted subset | -| Fixture IDs | `T04-LOCAL-*` are provisional local IDs; T01 assigns canonical catalog IDs/checksums | -| Blocking scope | Formal P6 ABI certification and live integration; local reference can be tested independently | - -No T01 start kit or public CombinePlan implementation was present in the base -checkout. These small receipt dataclasses are an explicit review proposal rather -than a claim to be that missing schema. Map them to the canonical start kit once -available; do not maintain a competing backend-specific protocol. Review, -approved fixture IDs, GPU evidence and live integration remain separate gates. diff --git a/rl_engine/kernels/ops/pytorch/moe/shared_residual_merge.py b/rl_engine/alignment/testing/moe_merge.py similarity index 58% rename from rl_engine/kernels/ops/pytorch/moe/shared_residual_merge.py rename to rl_engine/alignment/testing/moe_merge.py index 87008d8c..fe2b17fe 100644 --- a/rl_engine/kernels/ops/pytorch/moe/shared_residual_merge.py +++ b/rl_engine/alignment/testing/moe_merge.py @@ -1,12 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""P6/T04 eager forward reference; no routing, collective, or mHC arithmetic. +"""Validate MoE merge source identities and collect reference evidence. -The receipt/context interface is a T01 review proposal, not a second CombinePlan. -Upstream validation binds its opaque plan fingerprint and source receipts. This -module validates the local merge boundary; it cannot authenticate dishonest -producer receipts or certify unobserved GPU launch parameters. +This module is for alignment tests. Model code calls the tensor operator directly. +Application counts describe validated source history and the reference schedule; +they are not runtime dispatch measurements or proof of truthful producer data. """ from __future__ import annotations @@ -15,22 +14,21 @@ import json from collections import Counter from dataclasses import asdict, dataclass -from pathlib import Path from typing import Any import torch from torch import Tensor -from torch.utils._python_dispatch import TorchDispatchMode -CONTRACT_VERSION = "p6-task-contract.v1" -FOUNDATION_ABI = "foundation-moe-return.v1" -RECEIPT_VERSION = "p6.t04.merge-receipt.v1-proposal" +from rl_engine.kernels.registry import kernel_registry +from rl_engine.kernels.semantic_registry import OperatorRequirements, OperatorSession + +RECEIPT_VERSION = "moe_merge.receipt.v1" BACKEND_ID = "rlkernel.moe.shared_residual_merge.reference.v1" ORDER = ("shared", "residual", "cast_bf16") BOUNDARIES = ( - "P6.merge.after_shared", - "P6.merge.after_residual", - "P6.merge.final_bf16", + "after_shared", + "after_residual", + "final_bf16", ) @@ -97,12 +95,8 @@ class MergeContext: rank: int token_owner_ranks: tuple[int, ...] shared_replica_ranks: tuple[int, ...] - contract_version: str = CONTRACT_VERSION - foundation_abi: str = FOUNDATION_ABI - receipt_version: str = RECEIPT_VERSION merge_order: tuple[str, ...] = ORDER input_row_weighted: bool = True - route_weight_owner: str = "P5" @dataclass @@ -114,12 +108,6 @@ class MergeResult: def _validate_context(context: MergeContext) -> dict[str, MergeSource]: _require(isinstance(context, MergeContext), "INVALID_COMBINE_PLAN", "context") - for field, expected in ( - ("contract_version", CONTRACT_VERSION), - ("foundation_abi", FOUNDATION_ABI), - ("receipt_version", RECEIPT_VERSION), - ): - _require(getattr(context, field) == expected, "SCHEMA_MISMATCH", field) identity = context.identity _require(isinstance(identity, MergeIdentity), "IDENTITY_DRIFT", "identity") for field, value in asdict(identity).items(): @@ -177,9 +165,9 @@ def _validate_context(context: MergeContext) -> dict[str, MergeSource]: events = context.upstream_events _require( - context.input_row_weighted is True and context.route_weight_owner == "P5", + context.input_row_weighted is True, "INVALID_COMBINE_PLAN", - "weighted row owner", + "weighted routed input", ) _require( isinstance(events, tuple) and all(isinstance(e, str) for e in events), @@ -224,28 +212,7 @@ def _finite(value: Tensor, boundary: str) -> None: _require(bool(torch.isfinite(value).all().item()), "NON_FINITE", boundary) -class _ArithmeticTrace(TorchDispatchMode): - """Observe actual ATen adds/casts without replacing the underlying operations.""" - - def __init__(self): - super().__init__() - self.operations: list[dict[str, Any]] = [] - - def __torch_dispatch__(self, func, types, args=(), kwargs=None): - result = func(*args, **(kwargs or {})) - if func in (torch.ops.aten.add.Tensor, torch.ops.aten._to_copy.default): - self.operations.append( - { - "operator": str(func), - "output_dtype": str(result.dtype), - "shape": list(result.shape), - } - ) - return result - - -@torch.no_grad() -def shared_residual_merge_fwd( +def check_moe_merge( routed: Tensor, shared: Tensor, residual: Tensor, @@ -253,48 +220,19 @@ def shared_residual_merge_fwd( context: MergeContext, debug: bool = False, ) -> MergeResult: - """Merge an owner-local batch; return BF16 output and hash-only evidence. + """Validate producer evidence and collect reference boundaries for acceptance. - CPU and CUDA eager reference only. Hashing synchronizes/copies to CPU; this - is deliberately not a production kernel, CUDA Graph path, or backward op. - ``context`` is required: three bare tensors cannot prove exactly-once. + Identity/discrete failures precede numeric checks. Checksums and finite checks + synchronize GPU tensors; this helper belongs to validation, not model forward. """ - _require(not torch.compiler.is_compiling(), "UNSUPPORTED_CAPABILITY", "compiled reference") sources = _validate_context(context) inputs = {"routed": routed, "shared": shared, "residual": residual} - for role, value in inputs.items(): - _require(type(value) is Tensor, "UNSUPPORTED_CAPABILITY", f"{role}.tensor") - _require( - value.layout == torch.strided and value.ndim == 2 and value.is_contiguous(), - "UNSUPPORTED_CAPABILITY", - f"{role}.layout", - ) - _require(value.device.type in {"cpu", "cuda"}, "UNSUPPORTED_CAPABILITY", f"{role}.device") - _require( - value.shape[0] == len(context.identity.global_token_ids) and value.shape[1] > 0, - "INVALID_COMBINE_PLAN", - f"{role}.shape", - ) - _require( - value.shape == routed.shape and value.device == routed.device, - "INVALID_COMBINE_PLAN", - f"{role}.shape/device", - ) - if role == "routed": - _require(value.dtype == torch.float32, "EARLY_OR_MULTIPLE_DOWNCAST", "routed.dtype") - else: - _require( - value.dtype in {torch.float32, torch.bfloat16}, - "UNSUPPORTED_CAPABILITY", - f"{role}.dtype", - ) - if routed.device.type == "cuda": - with torch.cuda.device(routed.device): - _require( - not torch.cuda.is_current_stream_capturing(), - "UNSUPPORTED_CAPABILITY", - "CUDA Graph reference", - ) + _require( + routed.ndim == 2 and routed.shape[0] == len(context.identity.global_token_ids), + "INVALID_COMBINE_PLAN", + "routed.shape", + ) + _require(routed.dtype == torch.float32, "EARLY_OR_MULTIPLE_DOWNCAST", "routed.dtype") for role, value in inputs.items(): _require( tensor_sha256(value) == sources[role].tensor_checksum, @@ -304,48 +242,29 @@ def shared_residual_merge_fwd( for role, value in inputs.items(): _finite(value, role) - events: list[str] = [] - # Disable an enclosing autocast context; each eager add materializes FP32. - trace = _ArithmeticTrace() - with torch.autocast(device_type=routed.device.type, enabled=False), trace: - after_shared = torch.add(routed, shared.float()) - events.append("shared") - _finite(after_shared, BOUNDARIES[0]) - after_residual = torch.add(after_shared, residual.float()) - events.append("residual") - _finite(after_residual, BOUNDARIES[1]) - output = after_residual.to(torch.bfloat16) - events.append("cast_bf16") - _finite(output, BOUNDARIES[2]) - - actual_adds = [op for op in trace.operations if op["operator"] == "aten.add.Tensor"] - actual_downcasts = [ - op - for op in trace.operations - if op["operator"] == "aten._to_copy.default" and op["output_dtype"] == "torch.bfloat16" - ] - _require( - len(actual_adds) == 2 and all(op["output_dtype"] == "torch.float32" for op in actual_adds), - "ADDITION_ORDER_MISMATCH", - "observed ATen adds", + session = OperatorSession(kernel_registry.semantic) + resolution = session.resolve( + semantic_op="shared_residual_merge", + requested_backend=BACKEND_ID, + target="rollout", + requirements=OperatorRequirements( + device="rocm" if torch.version.hip and routed.is_cuda else routed.device.type, + dtype="float32", + ), ) - _require(len(actual_downcasts) == 1, "EARLY_OR_MULTIPLE_DOWNCAST", "observed ATen casts") - - boundary_tensors = dict(zip(BOUNDARIES, (after_shared, after_residual, output), strict=True)) - implementation_hash = hashlib.sha256(Path(__file__).read_bytes()).hexdigest() + op = session.instantiate(resolution) + values = op.forward_with_intermediates(routed, shared, residual) + boundary_tensors = dict(zip(BOUNDARIES, values, strict=True)) + for key, value in boundary_tensors.items(): + _finite(value, key) + instance = session.instance_provenance(resolution, op) provenance = { - "backend": BACKEND_ID, - "implementation_sha256": implementation_hash, + "operator_instance": instance.to_dict(), "torch_version": str(torch.__version__), "torch_git_version": torch.version.git_version, "device": str(routed.device), - "device_name": ( - torch.cuda.get_device_name(routed.device) if routed.device.type == "cuda" else "cpu" - ), + "device_name": (torch.cuda.get_device_name(routed.device) if routed.is_cuda else "cpu"), "build_runtime": torch.version.hip or torch.version.cuda, - "observed_aten_operations": trace.operations, - "accumulator_dtype": "torch.float32", - "output_dtype": "torch.bfloat16", "launch_parameters": None, "launch_provenance_status": "NOT_OBSERVED_AT_REFERENCE_LEVEL", } @@ -359,40 +278,31 @@ def shared_residual_merge_fwd( "checksum": tensor_sha256(value), "rank": context.rank, "identity": asdict(context.identity), - "backend": BACKEND_ID, - "implementation_sha256": implementation_hash, + "backend": instance.backend_id, + "implementation_fingerprint": instance.implementation_fingerprint, } for index, (key, value) in enumerate(boundary_tensors.items()) ] - applied = Counter(context.upstream_events + tuple(events)) + applied = Counter(context.upstream_events + ORDER) receipt = { "schema_version": RECEIPT_VERSION, - "contract_version": CONTRACT_VERSION, - "foundation_abi": FOUNDATION_ABI, "identity": asdict(context.identity), "sources": [asdict(sources[role]) for role in inputs], "upstream_events": list(context.upstream_events), "input_row_weighted": context.input_row_weighted, - "route_weight_owner": context.route_weight_owner, - "executed_events": events, + "reference_events": list(ORDER), + "count_basis": "validated_upstream_history_and_reference_schedule", "route_weight_applied_count": applied["route_weight"], "shared_applied_count": applied["shared"], "residual_applied_count": applied["residual"], - "local_downcast_count": len(actual_downcasts), - "merge_order_hash": _hash_json({"version": RECEIPT_VERSION, "order": events}), + "local_downcast_count": applied["cast_bf16"], + "merge_order_hash": _hash_json({"version": RECEIPT_VERSION, "order": ORDER}), "rank": context.rank, "token_owner_ranks": list(context.token_owner_ranks), "shared_replica_ranks": list(context.shared_replica_ranks), "identity_discrete_gate": "PASS", "scope": "local_forward_reference", - "certification": "INCOMPLETE: T01 ABI approval and GPU/integration evidence required", "actual_provenance": provenance, "boundaries": boundaries, } - return MergeResult(output, receipt, boundary_tensors if debug else {}) - - -class SharedResidualMergeReferenceOp: - """Explicit reference backend factory for the existing semantic registry.""" - - __call__ = staticmethod(shared_residual_merge_fwd) + return MergeResult(values[-1], receipt, boundary_tensors if debug else {}) diff --git a/rl_engine/kernels/ops/pytorch/moe/__init__.py b/rl_engine/kernels/ops/pytorch/moe/__init__.py index 4d17c436..df73c2d4 100644 --- a/rl_engine/kernels/ops/pytorch/moe/__init__.py +++ b/rl_engine/kernels/ops/pytorch/moe/__init__.py @@ -1,24 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from .shared_residual_merge import ( - MergeContext, - MergeContractError, - MergeIdentity, - MergeResult, - MergeSource, - SharedResidualMergeReferenceOp, - shared_residual_merge_fwd, - tensor_sha256, -) +from .merge import MoeMergeOp, shared_residual_merge_fwd -__all__ = [ - "MergeContext", - "MergeContractError", - "MergeIdentity", - "MergeResult", - "MergeSource", - "SharedResidualMergeReferenceOp", - "shared_residual_merge_fwd", - "tensor_sha256", -] +__all__ = ["MoeMergeOp", "shared_residual_merge_fwd"] diff --git a/rl_engine/kernels/ops/pytorch/moe/merge.py b/rl_engine/kernels/ops/pytorch/moe/merge.py new file mode 100644 index 00000000..c6abd67b --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/moe/merge.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import torch +from torch import Tensor + + +@torch.no_grad() +def _merge(routed: Tensor, shared: Tensor, residual: Tensor) -> tuple[Tensor, Tensor, Tensor]: + # Compilation may reassociate additions; this is the eager FP32 reference. + if torch.compiler.is_compiling(): + raise RuntimeError("shared_residual_merge_fwd requires PyTorch eager execution.") + for name, value in (("routed", routed), ("shared", shared), ("residual", residual)): + if type(value) is not Tensor: + raise TypeError(f"{name} must be a torch.Tensor.") + if value.layout != torch.strided or value.ndim != 2 or not value.is_contiguous(): + raise ValueError(f"{name} must be a contiguous [T, H] tensor.") + if value.shape[0] == 0 or value.shape[1] == 0: + raise ValueError(f"{name} must have non-empty token and hidden dimensions.") + if value.shape != routed.shape: + raise ValueError("routed, shared and residual must share shape.") + if value.device != routed.device: + raise RuntimeError("routed, shared and residual must be on the same device.") + if value.device.type not in {"cpu", "cuda"}: + raise RuntimeError("shared_residual_merge_fwd supports CPU and CUDA/ROCm tensors.") + dtypes = (torch.float32,) if name == "routed" else (torch.float32, torch.bfloat16) + if value.dtype not in dtypes: + raise TypeError(f"{name} must have dtype {dtypes}, got {value.dtype}.") + if routed.is_cuda: + with torch.cuda.device(routed.device): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("shared_residual_merge_fwd does not support CUDA Graph capture.") + + with torch.autocast(device_type=routed.device.type, enabled=False): + after_shared = routed + shared.float() + after_residual = after_shared + residual.float() + output = after_residual.to(torch.bfloat16) + return after_shared, after_residual, output + + +def shared_residual_merge_fwd(routed: Tensor, shared: Tensor, residual: Tensor) -> Tensor: + """Compute ``(routed + shared) + residual`` in FP32, then cast once to BF16. + + Inputs are contiguous [T, H] rows on the same device. Routed rows are already + weighted and combined in FP32; shared and residual may be FP32 or BF16. + This is a forward-only reference. Producer identity and ownership checks + belong to the caller's validation path. + """ + return _merge(routed, shared, residual)[-1] + + +class MoeMergeOp(torch.nn.Module): + """Pure PyTorch reference for the MoE shared/residual merge.""" + + def forward(self, routed: Tensor, shared: Tensor, residual: Tensor) -> Tensor: + return shared_residual_merge_fwd(routed, shared, residual) + + def forward_with_intermediates( + self, routed: Tensor, shared: Tensor, residual: Tensor + ) -> tuple[Tensor, Tensor, Tensor]: + """Return after-shared FP32, after-residual FP32 and final BF16 tensors.""" + return _merge(routed, shared, residual) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index fb1b3716..a8de5f93 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -189,23 +189,25 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: backend_id="rlkernel.moe.shared_residual_merge.reference.v1", supported_targets=frozenset({"rollout", "training"}), supported_devices=frozenset({"cpu", "cuda", "rocm"}), + # Primary (routed) input dtype; shared/residual may also be BF16. supported_dtypes=frozenset({"float32"}), - supported_topologies={"*": "*"}, + # Local merge only; distributed topology requirements belong to the caller. + supported_topologies={}, determinism_or_alignment_properties={ - "algorithm": "fp32_routed_plus_shared_then_residual_final_bf16", + "algorithm": "fixed_order_fp32_merge", + "batch_invariant": True, + "deterministic": True, "reference_only": True, "forward_only": True, - "requires_merge_context": True, "strict_observable": True, "gpu_launch_observable": False, - "certification": "pending_T01_and_GPU_evidence", }, lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, implementation_class_or_factory=( - "rl_engine.kernels.ops.pytorch.moe.SharedResidualMergeReferenceOp" + "rl_engine.kernels.ops.pytorch.moe.merge.MoeMergeOp" ), fallback_policy=OperatorFallbackPolicy.ERROR, - version_or_build_fingerprint="p6.t04.merge-receipt.v1-proposal", + version_or_build_fingerprint="MoeMergeOp-v1", ), OperatorBackendDescriptor( semantic_op="selected_logprob", diff --git a/tests/fixtures/p6_t04_merge.json b/tests/fixtures/moe_merge.json similarity index 81% rename from tests/fixtures/p6_t04_merge.json rename to tests/fixtures/moe_merge.json index 0f9f4eae..1818f200 100644 --- a/tests/fixtures/p6_t04_merge.json +++ b/tests/fixtures/moe_merge.json @@ -1,9 +1,8 @@ { - "schema_version": "p6.t04.synthetic-fixtures.v1-proposal", - "approval": "T01 case IDs and checksum approval pending", + "schema_version": "moe_merge.fixtures.v1", "cases": [ { - "id": "T04-LOCAL-MERGE", + "id": "simple_merge", "routed": [[1.0, 2.0, 3.0]], "shared": [[2.0, 3.0, 4.0]], "residual": [[3.0, 4.0, 5.0]], @@ -12,7 +11,7 @@ "bf16_words": [[16576, 16656, 16704]] }, { - "id": "T04-LOCAL-ORDER-CAST-ZERO", + "id": "rounding_order_and_signed_zero", "routed": [[16777216.0, 256.0, 1.0, -0.0]], "shared": [[1.0, 1.0, 0.000000059604644775390625, -0.0]], "residual": [[-16777216.0, -256.0, 0.0, -0.0]], @@ -21,7 +20,7 @@ "bf16_words": [[0, 16256, 16256, 32768]] }, { - "id": "T04-LOCAL-BF16-TIES", + "id": "bf16_ties_and_limits", "routed": [[1.00390625, 1.01171875, 3.3895313892515355e38, 1.1754943508222875e-38]], "shared": [[0.0, 0.0, 0.0, 0.0]], "residual": [[0.0, 0.0, 0.0, 0.0]], diff --git a/tests/test_moe_merge.py b/tests/test_moe_merge.py new file mode 100644 index 00000000..31718f19 --- /dev/null +++ b/tests/test_moe_merge.py @@ -0,0 +1,607 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MoE merge: fixed goldens, injected faults, and owner-local EP checks.""" + +from __future__ import annotations + +import hashlib +import json +import os +import shutil +import time +from dataclasses import asdict, replace +from datetime import timedelta +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.alignment.cross_config.artifacts import ArtifactError, ArtifactStore +from rl_engine.alignment.testing.moe_merge import ( + BACKEND_ID, + BOUNDARIES, + ORDER, + MergeContext, + MergeContractError, + MergeIdentity, + MergeSource, + check_moe_merge, + tensor_sha256, +) +from rl_engine.kernels.ops.pytorch.moe import MoeMergeOp, shared_residual_merge_fwd + +FIXTURE = Path(__file__).parent / "fixtures" / "moe_merge.json" +FIXTURE_CHECKSUM = hashlib.sha256(FIXTURE.read_bytes()).hexdigest() +CASES = json.loads(FIXTURE.read_text())["cases"] + + +@pytest.fixture( + params=[ + "cpu", + pytest.param( + "cuda", + marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA/ROCm GPU"), + ), + ] +) +def device(request): + return request.param + + +def test_fixture_checksum(): + assert FIXTURE_CHECKSUM == "f92fb938160285f54de88f746dd814471598ff117701103bfac8c55e870f11cc" + + +def make_inputs(case=CASES[0], *, device="cpu", dtype=torch.float32): + return tuple( + torch.tensor( + case[role], + dtype=torch.float32 if role == "routed" else dtype, + device=device, + ) + for role in ("routed", "shared", "residual") + ) + + +def make_context(inputs, *, case_id="simple_merge", rank=0, ep=1, tokens=None): + identity = MergeIdentity( + case_id=case_id, + run_id="synthetic-merge", + pass_id="forward-0", + checkpoint_fingerprint="synthetic-no-checkpoint", + weight_fingerprint="synthetic-no-weights", + route_plan_fingerprint="synthetic-route", + exchange_plan_fingerprint="synthetic-exchange", + combine_plan_fingerprint="synthetic-combine", + fixture_checksum=FIXTURE_CHECKSUM, + global_token_ids=tokens or tuple(range(inputs[0].shape[0])), + ) + sources = tuple( + MergeSource(role, f"synthetic-{role}", identity, tensor_sha256(value)) + for role, value in zip(("routed", "shared", "residual"), inputs, strict=True) + ) + return MergeContext( + identity=identity, + expected_source_ids=tuple(s.source_id for s in sources), + sources=sources, + upstream_events=("route_weight",), + rank=rank, + token_owner_ranks=(rank,) * len(identity.global_token_ids), + shared_replica_ranks=tuple(range(ep)), + ) + + +def assert_bytes(a, b): + assert a.shape == b.shape and a.dtype == b.dtype + assert torch.equal( + a.detach().cpu().contiguous().view(torch.uint8), + b.detach().cpu().contiguous().view(torch.uint8), + ) + + +@pytest.mark.parametrize("case", CASES, ids=lambda c: c["id"]) +def test_frozen_golden_and_debug(case, device): + inputs = make_inputs(case, device=device) + context = make_context(inputs, case_id=case["id"]) + result = check_moe_merge(*inputs, context=context, debug=True) + golden = torch.tensor(case["bf16_words"], dtype=torch.uint16).view(torch.bfloat16) + assert_bytes(result.token_moe_row_bf16, golden) + assert_bytes(shared_residual_merge_fwd(*inputs), golden) + op = MoeMergeOp() + assert_bytes(op(*inputs), golden) + for actual, key in zip(op.forward_with_intermediates(*inputs), BOUNDARIES, strict=True): + assert_bytes(actual, result.debug_boundaries[key]) + for key, name in zip(BOUNDARIES[:2], ("after_shared", "after_residual"), strict=True): + assert_bytes(result.debug_boundaries[key], torch.tensor(case[name], dtype=torch.float32)) + quiet = check_moe_merge(*inputs, context=context) + assert quiet.debug_boundaries == {} + assert quiet.receipt == result.receipt + assert_bytes(quiet.token_moe_row_bf16, golden) + for count in ( + "route_weight_applied_count", + "shared_applied_count", + "residual_applied_count", + "local_downcast_count", + ): + assert result.receipt[count] == 1 + assert result.receipt["reference_events"] == list(ORDER) + + +def test_wrong_arithmetic_is_detectable(device): + r, s, x = make_inputs(CASES[1], device=device) + good = shared_residual_merge_fwd(r, s, x) + wrong_order = (r + (s + x)).bfloat16() + early_cast = ((r + s).bfloat16().float() + x).bfloat16() + assert wrong_order[0, 0].item() == 1 and good[0, 0].item() == 0 + assert early_cast[0, 1].item() == 0 and good[0, 1].item() == 1 + + +@pytest.mark.parametrize( + "event,status", + [ + ("shared", "SHARED_APPLIED_NOT_ONCE"), + ("residual", "RESIDUAL_APPLIED_NOT_ONCE"), + ("cast_bf16", "EARLY_OR_MULTIPLE_DOWNCAST"), + ("route_weight", "ROUTE_WEIGHT_APPLIED_TWICE"), + ("mhc_post", "INVALID_COMBINE_PLAN"), + ], +) +def test_reject_prior_application(event, status, device): + inputs = make_inputs(device=device) + context = make_context(inputs) + context = replace(context, upstream_events=context.upstream_events + (event,)) + with pytest.raises(MergeContractError, match=status): + check_moe_merge(*inputs, context=context) + + +@pytest.mark.parametrize( + "role,status", + [("shared", "SHARED_APPLIED_NOT_ONCE"), ("residual", "RESIDUAL_APPLIED_NOT_ONCE")], +) +@pytest.mark.parametrize("duplicate", [False, True]) +def test_missing_duplicate_source(role, status, duplicate, device): + inputs = make_inputs(device=device) + context = make_context(inputs) + source = next(s for s in context.sources if s.role == role) + sources = ( + context.sources + (source,) + if duplicate + else tuple(s for s in context.sources if s.role != role) + ) + with pytest.raises(MergeContractError, match=status): + check_moe_merge(*inputs, context=replace(context, sources=sources)) + + +@pytest.mark.parametrize( + "field", + [ + "case_id", + "run_id", + "pass_id", + "checkpoint_fingerprint", + "weight_fingerprint", + "route_plan_fingerprint", + "exchange_plan_fingerprint", + "combine_plan_fingerprint", + "fixture_checksum", + "global_token_ids", + ], +) +def test_identity_precedes_numerics(field, device): + inputs = make_inputs(device=device) + context = make_context(inputs) + source = context.sources[1] + identity = replace( + source.identity, **{field: (42,) if field == "global_token_ids" else "wrong"} + ) + context = replace( + context, + sources=( + context.sources[0], + replace(source, identity=identity), + context.sources[2], + ), + ) + inputs[0].fill_(float("nan")) + with pytest.raises(MergeContractError, match="IDENTITY_DRIFT: shared.identity"): + check_moe_merge(*inputs, context=context) + + +@pytest.mark.parametrize( + "mutation,status", + [ + ( + {"merge_order": ("residual", "shared", "cast_bf16")}, + "ADDITION_ORDER_MISMATCH", + ), + ({"upstream_events": ()}, "INVALID_COMBINE_PLAN"), + ({"token_owner_ranks": (1,)}, "SHARED_APPLIED_NOT_ONCE"), + ({"rank": -1}, "MISSING_PROVENANCE"), + ({"shared_replica_ranks": (0, 0)}, "INVALID_COMBINE_PLAN"), + ({"input_row_weighted": False}, "INVALID_COMBINE_PLAN"), + ( + { + "expected_source_ids": ( + "wrong", + "synthetic-shared", + "synthetic-residual", + ) + }, + "IDENTITY_DRIFT", + ), + ], +) +def test_invalid_context(mutation, status, device): + inputs = make_inputs(device=device) + with pytest.raises(MergeContractError, match=status): + check_moe_merge(*inputs, context=replace(make_context(inputs), **mutation)) + + +def test_mhc_post_input_rejected(): + inputs = make_inputs() + context = make_context(inputs) + context = replace( + context, + sources=(*context.sources[:2], replace(context.sources[2], role="mhc_post")), + ) + with pytest.raises(MergeContractError, match="INVALID_COMBINE_PLAN"): + check_moe_merge(*inputs, context=context) + + +def test_checksum_detects_stale_tensor(): + inputs = make_inputs() + context = make_context(inputs) + inputs[1][0, 0] += 1 + with pytest.raises(MergeContractError, match="IDENTITY_DRIFT: shared.tensor_checksum"): + check_moe_merge(*inputs, context=context) + + +@pytest.mark.parametrize("role", range(3)) +@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")]) +def test_nonfinite_inputs(role, value, device): + inputs = make_inputs(device=device) + inputs[role][0, 0] = value + with pytest.raises(MergeContractError, match="NON_FINITE"): + check_moe_merge(*inputs, context=make_context(inputs)) + + +@pytest.mark.parametrize("stage", ["after_shared", "after_residual", "final_bf16"]) +def test_intermediate_and_cast_overflow(stage, device): + maxval = torch.finfo(torch.float32).max + values = { + "after_shared": (maxval, maxval, -maxval), + "after_residual": (maxval, 0.0, maxval), + "final_bf16": (maxval, 0.0, 0.0), + }[stage] + inputs = tuple(torch.tensor([[v]], dtype=torch.float32, device=device) for v in values) + with pytest.raises(MergeContractError, match=f"NON_FINITE: {stage}"): + check_moe_merge(*inputs, context=make_context(inputs)) + + +@pytest.mark.parametrize("kind", ["routed_bf16", "shared_fp16", "broadcast", "strided"]) +def test_unsupported_payloads(kind, device): + inputs = list(make_inputs(device=device)) + if kind == "routed_bf16": + inputs[0] = inputs[0].bfloat16() + elif kind == "shared_fp16": + inputs[1] = inputs[1].half() + elif kind == "broadcast": + inputs[1] = inputs[1][:, :1].contiguous() + else: + inputs[1] = torch.ones(1, 6, device=device)[:, ::2] + with pytest.raises((TypeError, ValueError)): + shared_residual_merge_fwd(*inputs) + if kind == "routed_bf16": + with pytest.raises(MergeContractError, match="EARLY_OR_MULTIPLE_DOWNCAST"): + check_moe_merge(*inputs, context=make_context(inputs)) + + +@pytest.mark.parametrize("ep", [1, 2, 4, 8]) +def test_ep_replication_owner_local_mock(ep): + # No collectives: replicas exist on every mock rank; exactly one owner merges + # each token. Process-group ownership is covered separately on GPUs. + inputs = tuple(v.repeat(16, 1) for v in make_inputs()) + canonical = check_moe_merge(*inputs, context=make_context(inputs)) + gathered = torch.empty_like(canonical.token_moe_row_bf16) + for rank in range(ep): + tokens = tuple(range(rank, 16, ep)) + local = tuple(v[list(tokens)].contiguous() for v in inputs) + context = make_context(local, rank=rank, ep=ep, tokens=tokens) + result = check_moe_merge(*local, context=context) + gathered[list(tokens)] = result.token_moe_row_bf16 + assert result.receipt["shared_applied_count"] == 1 + assert result.receipt["merge_order_hash"] == canonical.receipt["merge_order_hash"] + if ep > 1: + with pytest.raises(MergeContractError, match="SHARED_APPLIED_NOT_ONCE"): + check_moe_merge(*local, context=replace(context, rank=(rank + 1) % ep)) + assert_bytes(gathered, canonical.token_moe_row_bf16) + + +@pytest.mark.parametrize("shared_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("residual_dtype", [torch.float32, torch.bfloat16]) +def test_input_immutability_and_autocast(shared_dtype, residual_dtype, device): + routed, shared, residual = make_inputs(CASES[1], device=device) + inputs = (routed, shared.to(shared_dtype), residual.to(residual_dtype)) + for value in inputs: + value.requires_grad_(True) + snapshots = tuple(v.detach().clone() for v in inputs) + op = MoeMergeOp() + baseline = op(*inputs) + with torch.autocast(device, dtype=torch.bfloat16): + after_shared, after_residual, output = op.forward_with_intermediates(*inputs) + result = op(*inputs) + golden = torch.tensor(CASES[1]["bf16_words"], dtype=torch.uint16).view(torch.bfloat16) + assert_bytes(result, golden) + assert_bytes(result, baseline) + assert_bytes(result, output) + assert after_shared.dtype == after_residual.dtype == torch.float32 + assert not any(v.requires_grad for v in (after_shared, after_residual, output, result)) + for value, saved in zip(inputs, snapshots, strict=True): + assert_bytes(value, saved) + + +@pytest.mark.parametrize("width", [1, 17, 129, 4096]) +def test_batch_chunk_padding_permutation(width, device): + generator = torch.Generator().manual_seed(71) + inputs = tuple(torch.randn(7, width, generator=generator).to(device) for _ in range(3)) + expected = shared_residual_merge_fwd(*inputs) + permutation = [6, 2, 0, 1, 5, 3, 4] + permuted = tuple(v[permutation] for v in inputs) + assert_bytes(shared_residual_merge_fwd(*permuted), expected[permutation]) + for start, end in ((0, 1), (1, 4), (4, 7)): + chunk = tuple(v[start:end] for v in inputs) + assert_bytes(shared_residual_merge_fwd(*chunk), expected[start:end]) + padded = tuple(torch.cat((v, torch.zeros(3, width, device=device))) for v in inputs) + assert_bytes(shared_residual_merge_fwd(*padded)[:7], expected) + + +def test_duplicate_and_empty_token_batch(): + inputs = tuple(v.repeat(2, 1) for v in make_inputs()) + context = make_context(inputs) + context = replace(context, identity=replace(context.identity, global_token_ids=(1, 1))) + with pytest.raises(MergeContractError, match="AMBIGUOUS_GLOBAL_TOKEN_MAPPING"): + check_moe_merge(*inputs, context=context) + empty = tuple(v[:0] for v in inputs) + with pytest.raises(MergeContractError, match="AMBIGUOUS_GLOBAL_TOKEN_MAPPING"): + check_moe_merge(*empty, context=make_context(empty)) + + +@pytest.mark.parametrize("target", ["rollout", "training"]) +def test_reference_registry_entry(target): + from rl_engine.kernels.registry import KernelRegistry + from rl_engine.kernels.semantic_registry import OperatorRequirements, OperatorSession + + session = OperatorSession(KernelRegistry().semantic) + resolution = session.resolve( + semantic_op="shared_residual_merge", + requested_backend=BACKEND_ID, + target=target, + requirements=OperatorRequirements( + device="cpu", + dtype="float32", + alignment_properties={"deterministic": True, "batch_invariant": True}, + ), + ) + op = session.instantiate(resolution) + inputs = make_inputs() + result = op(*inputs) + assert isinstance(op, torch.nn.Module) + assert_bytes(result, shared_residual_merge_fwd(*inputs)) + assert resolution.descriptor.determinism_or_alignment_properties["reference_only"] is True + provenance = session.instance_provenance(resolution, op) + assert provenance.backend_id == BACKEND_ID + assert provenance.concrete_implementation.endswith("MoeMergeOp") + assert provenance.implementation_fingerprint + + +@pytest.mark.parametrize( + "requirements,capability", + [ + ({"dtype": "bfloat16"}, "dtype"), + ({"topology": {"tensor_parallel_size": 2}}, "topology"), + ({"topology": {"expert_parallel_size": 2}}, "topology"), + ({"alignment_properties": {"cross_tp_bitwise": True}}, "alignment_properties"), + ({"alignment_properties": {"gpu_launch_observable": True}}, "alignment_properties"), + ({"alignment_properties": {"reference_only": False}}, "alignment_properties"), + ({"alignment_properties": {"forward_only": False}}, "alignment_properties"), + ], +) +def test_reference_registry_rejects_unsupported_requirements(requirements, capability): + from rl_engine.kernels.registry import KernelRegistry + from rl_engine.kernels.semantic_registry import ( + OperatorRequirements, + OperatorResolutionError, + OperatorSession, + ) + + session = OperatorSession(KernelRegistry().semantic) + requested = {"device": "cpu", "dtype": "float32", **requirements} + with pytest.raises(OperatorResolutionError) as error: + session.resolve( + semantic_op="shared_residual_merge", + requested_backend=BACKEND_ID, + target="rollout", + requirements=OperatorRequirements(**requested), + ) + assert { + decision.capability + for decision in error.value.trace.capability_decisions + if not decision.passed + } == {capability} + + +def test_sealed_reference_evidence(tmp_path, device): + root = Path(os.environ.get("RL_KERNEL_MOE_MERGE_ARTIFACT_DIR", tmp_path)) + store = ArtifactStore(root / device) + required = ("merge_receipt.json", "merge_debug.pt") + for case in CASES: + inputs = make_inputs(case, device=device) + context = make_context(inputs, case_id=case["id"]) + result = check_moe_merge(*inputs, context=context, debug=True) + experiment = { + "context": json.loads(json.dumps(asdict(context))), + "implementation": result.receipt["actual_provenance"], + } + # Identity/build changes cannot silently resume an existing experiment. + store.initialize_experiment(case["id"], experiment=experiment, plan=[]) + attempt = store.create_attempt(case["id"], case["id"]) + store.write_json(attempt, required[0], result.receipt) + store.write_tensor_bundle(attempt, required[1], result.debug_boundaries) + store.complete_attempt( + attempt, + required=required, + summary={ + "schema_version": "cross_config.complete.v1", + "case_id": case["id"], + "attempt_id": attempt.name, + "status": "LOCAL_REFERENCE_ONLY", + }, + ) + store.validate_completed_attempt(attempt, required=required, expected_case_id=case["id"]) + loaded = store.load_tensor_bundle(attempt / required[1])["tensors"] + receipt = json.loads((attempt / required[0]).read_text()) + for boundary in receipt["boundaries"]: + assert tensor_sha256(loaded[boundary["key"]]) == boundary["checksum"] + with pytest.raises(ArtifactError, match="resume metadata differs"): + store.initialize_experiment(case["id"], experiment={"wrong_run": True}, plan=[]) + # Exercise corruption on an isolated copy, never alter a sealed report. + corrupted = tmp_path / "corrupted" / attempt.name + shutil.copytree(attempt, corrupted) + (corrupted / "merge_receipt.json").write_text("{}") + with pytest.raises(ArtifactError, match="artifact hash does not match"): + store.validate_completed_attempt(corrupted, required=required) + + +def _gpu_owner_worker(rank, ep, rendezvous, artifact_root): + torch.cuda.set_device(rank) + torch.cuda.set_per_process_memory_fraction(0.05, rank) + dist.init_process_group( + "nccl", + init_method=rendezvous, + rank=rank, + world_size=ep, + timeout=timedelta(seconds=120), + ) + try: + # Every process holds a complete shared replica and merges only its own + # token rows. Uneven counts and reversed rows catch mapping mistakes. + # Integer-valued inputs have an independent exact oracle. + grid = torch.arange(17 * 17, dtype=torch.float32).reshape(17, 17) + routed = (grid.remainder(31) - 15) * 256 + shared_replica = (grid.remainder(7) + 1).to(rank) + residual = 3 - routed + tokens = tuple(reversed(range(rank, 17, ep))) + local = ( + routed[list(tokens)].to(rank), + shared_replica[list(tokens)].contiguous(), + residual[list(tokens)].to(rank), + ) + context = make_context( + local, + case_id=f"owner-local-ep-{ep}", + rank=dist.get_rank(), + ep=dist.get_world_size(), + tokens=tokens, + ) + result = check_moe_merge(*local, context=context, debug=True) + for field in ( + "shared_applied_count", + "residual_applied_count", + "local_downcast_count", + ): + assert result.receipt[field] == 1 + if ep > 1: + with pytest.raises(MergeContractError, match="SHARED_APPLIED_NOT_ONCE"): + check_moe_merge(*local, context=replace(context, rank=(rank + 1) % ep)) + wrong = (local[0] + ep * local[1] + local[2]).bfloat16() + assert not torch.equal( + wrong.view(torch.int16), result.token_moe_row_bf16.view(torch.int16) + ) + + # This gather belongs solely to the test harness. The merge never calls + # a collective; the caller owns communication and token assignment. + record = { + "rank": dist.get_rank(), + "device": str(result.token_moe_row_bf16.device), + "tokens": tokens, + "shared_replica_checksum": tensor_sha256(shared_replica), + "output_words": result.token_moe_row_bf16.cpu().view(torch.int16).tolist(), + } + gathered = [None] * ep + dist.all_gather_object(gathered, record) + assert [item["rank"] for item in gathered] == list(range(ep)) + assert sorted(t for item in gathered for t in item["tokens"]) == list(range(17)) + assert len({item["shared_replica_checksum"] for item in gathered}) == 1 + expected = (grid.remainder(7) + 4).bfloat16().view(torch.int16) + for item in gathered: + assert_bytes( + torch.tensor(item["output_words"], dtype=torch.int16), + expected[list(item["tokens"])], + ) + + store = ArtifactStore(Path(artifact_root) / f"ep-{ep}") + case_id = f"rank-{rank}" + store.initialize_experiment( + case_id, + experiment={ + "context": json.loads(json.dumps(asdict(context))), + "implementation": result.receipt["actual_provenance"], + "test_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(), + "world_size": dist.get_world_size(), + "test_collective_backend": dist.get_backend(), + }, + plan=[], + ) + attempt = store.create_attempt(case_id, case_id) + required = ("merge_receipt.json", "merge_debug.pt", "rank_gather.json") + store.write_json(attempt, required[0], result.receipt) + store.write_tensor_bundle(attempt, required[1], result.debug_boundaries) + store.write_json(attempt, required[2], {"ranks": gathered}) + store.complete_attempt( + attempt, + required=required, + summary={ + "schema_version": "cross_config.complete.v1", + "case_id": case_id, + "attempt_id": attempt.name, + "status": "LOCAL_REFERENCE_ONLY", + }, + ) + store.validate_completed_attempt(attempt, required=required, expected_case_id=case_id) + loaded = store.load_tensor_bundle(attempt / required[1])["tensors"] + for boundary in result.receipt["boundaries"]: + assert tensor_sha256(loaded[boundary["key"]]) == boundary["checksum"] + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("ep", [1, 2, 4, 8]) +def test_gpu_ep_shared_once(ep, tmp_path): + if torch.cuda.device_count() < ep: + pytest.skip(f"needs {ep} GPUs for real owner ranks") + if not dist.is_nccl_available(): + pytest.skip("NCCL/RCCL unavailable") + root = Path(os.environ.get("RL_KERNEL_MOE_MERGE_ARTIFACT_DIR", tmp_path)) / "distributed" + workers = mp.spawn( + _gpu_owner_worker, + args=(ep, (tmp_path / "rendezvous").as_uri(), str(root)), + nprocs=ep, + join=False, + ) + deadline = time.monotonic() + 180 + try: + while not workers.join(timeout=1): + if time.monotonic() >= deadline: + pytest.fail(f"EP={ep} workers did not finish within 180 seconds") + finally: + for process in workers.processes: + if process.is_alive(): + process.terminate() + for process in workers.processes: + process.join(timeout=5) + if process.is_alive(): + process.kill() + process.join(timeout=5) diff --git a/tests/test_shared_residual_merge.py b/tests/test_shared_residual_merge.py deleted file mode 100644 index 6d7189fe..00000000 --- a/tests/test_shared_residual_merge.py +++ /dev/null @@ -1,420 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -"""T04 reference acceptance; fixed goldens, injected faults, and owner-local EP mocks.""" - -from __future__ import annotations - -import hashlib -import json -import os -import shutil -from dataclasses import asdict, replace -from pathlib import Path - -import pytest -import torch - -from rl_engine.alignment.cross_config.artifacts import ArtifactError, ArtifactStore -from rl_engine.kernels.ops.pytorch.moe import ( - MergeContext, - MergeContractError, - MergeIdentity, - MergeSource, - shared_residual_merge_fwd, - tensor_sha256, -) -from rl_engine.kernels.ops.pytorch.moe.shared_residual_merge import BOUNDARIES, ORDER - -FIXTURE = Path(__file__).parent / "fixtures" / "p6_t04_merge.json" -FIXTURE_CHECKSUM = hashlib.sha256(FIXTURE.read_bytes()).hexdigest() -CASES = json.loads(FIXTURE.read_text())["cases"] - - -def test_fixture_checksum(): - assert FIXTURE_CHECKSUM == "79813b6999c56489b4e08575fd71ed62d6d65cdd7ad61a90d7727a83bceb20d4" - - -def make_inputs(case=CASES[0], *, device="cpu", dtype=torch.float32): - return tuple( - torch.tensor(case[role], dtype=torch.float32 if role == "routed" else dtype, device=device) - for role in ("routed", "shared", "residual") - ) - - -def make_context(inputs, *, case_id="T04-LOCAL-MERGE", rank=0, ep=1, tokens=None): - identity = MergeIdentity( - case_id=case_id, - run_id="synthetic-t04", - pass_id="forward-0", - checkpoint_fingerprint="synthetic-no-checkpoint", - weight_fingerprint="synthetic-no-weights", - route_plan_fingerprint="synthetic-route", - exchange_plan_fingerprint="synthetic-exchange", - combine_plan_fingerprint="synthetic-combine", - fixture_checksum=FIXTURE_CHECKSUM, - global_token_ids=tokens or tuple(range(inputs[0].shape[0])), - ) - sources = tuple( - MergeSource(role, f"synthetic-{role}", identity, tensor_sha256(value)) - for role, value in zip(("routed", "shared", "residual"), inputs, strict=True) - ) - return MergeContext( - identity=identity, - expected_source_ids=tuple(s.source_id for s in sources), - sources=sources, - upstream_events=("route_weight",), - rank=rank, - token_owner_ranks=(rank,) * len(identity.global_token_ids), - shared_replica_ranks=tuple(range(ep)), - ) - - -def assert_bytes(a, b): - assert a.shape == b.shape and a.dtype == b.dtype - assert torch.equal( - a.detach().cpu().contiguous().view(torch.uint8), - b.detach().cpu().contiguous().view(torch.uint8), - ) - - -@pytest.mark.parametrize("case", CASES, ids=lambda c: c["id"]) -def test_frozen_golden_and_debug(case): - inputs = make_inputs(case) - context = make_context(inputs, case_id=case["id"]) - result = shared_residual_merge_fwd(*inputs, context=context, debug=True) - golden = torch.tensor(case["bf16_words"], dtype=torch.uint16).view(torch.bfloat16) - assert_bytes(result.token_moe_row_bf16, golden) - for key, name in zip(BOUNDARIES[:2], ("after_shared", "after_residual"), strict=True): - assert_bytes(result.debug_boundaries[key], torch.tensor(case[name], dtype=torch.float32)) - quiet = shared_residual_merge_fwd(*inputs, context=context) - assert quiet.debug_boundaries == {} - assert quiet.receipt == result.receipt - assert_bytes(quiet.token_moe_row_bf16, golden) - for count in ( - "route_weight_applied_count", - "shared_applied_count", - "residual_applied_count", - "local_downcast_count", - ): - assert result.receipt[count] == 1 - assert result.receipt["executed_events"] == list(ORDER) - - -def test_wrong_arithmetic_is_detectable(): - r, s, x = make_inputs(CASES[1]) - good = shared_residual_merge_fwd(r, s, x, context=make_context((r, s, x))) - wrong_order = (r + (s + x)).bfloat16() - early_cast = ((r + s).bfloat16().float() + x).bfloat16() - assert wrong_order[0, 0].item() == 1 and good.token_moe_row_bf16[0, 0].item() == 0 - assert early_cast[0, 1].item() == 0 and good.token_moe_row_bf16[0, 1].item() == 1 - - -@pytest.mark.parametrize( - "event,status", - [ - ("shared", "SHARED_APPLIED_NOT_ONCE"), - ("residual", "RESIDUAL_APPLIED_NOT_ONCE"), - ("cast_bf16", "EARLY_OR_MULTIPLE_DOWNCAST"), - ("route_weight", "ROUTE_WEIGHT_APPLIED_TWICE"), - ("mhc_post", "INVALID_COMBINE_PLAN"), - ], -) -def test_reject_prior_application(event, status): - inputs = make_inputs() - context = make_context(inputs) - context = replace(context, upstream_events=context.upstream_events + (event,)) - with pytest.raises(MergeContractError, match=status): - shared_residual_merge_fwd(*inputs, context=context) - - -@pytest.mark.parametrize( - "role,status", - [("shared", "SHARED_APPLIED_NOT_ONCE"), ("residual", "RESIDUAL_APPLIED_NOT_ONCE")], -) -@pytest.mark.parametrize("duplicate", [False, True]) -def test_missing_duplicate_source(role, status, duplicate): - inputs = make_inputs() - context = make_context(inputs) - source = next(s for s in context.sources if s.role == role) - sources = ( - context.sources + (source,) - if duplicate - else tuple(s for s in context.sources if s.role != role) - ) - with pytest.raises(MergeContractError, match=status): - shared_residual_merge_fwd(*inputs, context=replace(context, sources=sources)) - - -@pytest.mark.parametrize( - "field", - [ - "case_id", - "run_id", - "pass_id", - "checkpoint_fingerprint", - "weight_fingerprint", - "route_plan_fingerprint", - "exchange_plan_fingerprint", - "combine_plan_fingerprint", - "fixture_checksum", - "global_token_ids", - ], -) -def test_identity_precedes_numerics(field): - inputs = make_inputs() - context = make_context(inputs) - source = context.sources[1] - identity = replace( - source.identity, **{field: (42,) if field == "global_token_ids" else "wrong"} - ) - context = replace( - context, - sources=(context.sources[0], replace(source, identity=identity), context.sources[2]), - ) - inputs[0].fill_(float("nan")) - with pytest.raises(MergeContractError, match="IDENTITY_DRIFT: shared.identity"): - shared_residual_merge_fwd(*inputs, context=context) - - -@pytest.mark.parametrize( - "mutation,status", - [ - ({"merge_order": ("residual", "shared", "cast_bf16")}, "ADDITION_ORDER_MISMATCH"), - ({"upstream_events": ()}, "INVALID_COMBINE_PLAN"), - ({"token_owner_ranks": (1,)}, "SHARED_APPLIED_NOT_ONCE"), - ({"rank": -1}, "MISSING_PROVENANCE"), - ({"shared_replica_ranks": (0, 0)}, "INVALID_COMBINE_PLAN"), - ({"foundation_abi": "wrong"}, "SCHEMA_MISMATCH"), - ({"receipt_version": "wrong"}, "SCHEMA_MISMATCH"), - ({"input_row_weighted": False}, "INVALID_COMBINE_PLAN"), - ({"route_weight_owner": "P6"}, "INVALID_COMBINE_PLAN"), - ( - {"expected_source_ids": ("wrong", "synthetic-shared", "synthetic-residual")}, - "IDENTITY_DRIFT", - ), - ], -) -def test_invalid_context(mutation, status): - inputs = make_inputs() - with pytest.raises(MergeContractError, match=status): - shared_residual_merge_fwd(*inputs, context=replace(make_context(inputs), **mutation)) - - -def test_mhc_post_input_rejected(): - inputs = make_inputs() - context = make_context(inputs) - context = replace( - context, sources=(*context.sources[:2], replace(context.sources[2], role="mhc_post")) - ) - with pytest.raises(MergeContractError, match="INVALID_COMBINE_PLAN"): - shared_residual_merge_fwd(*inputs, context=context) - - -def test_checksum_detects_stale_tensor(): - inputs = make_inputs() - context = make_context(inputs) - inputs[1][0, 0] += 1 - with pytest.raises(MergeContractError, match="IDENTITY_DRIFT: shared.tensor_checksum"): - shared_residual_merge_fwd(*inputs, context=context) - - -@pytest.mark.parametrize("role", range(3)) -@pytest.mark.parametrize("value", [float("nan"), float("inf"), -float("inf")]) -def test_nonfinite_inputs(role, value): - inputs = make_inputs() - inputs[role][0, 0] = value - with pytest.raises(MergeContractError, match="NON_FINITE"): - shared_residual_merge_fwd(*inputs, context=make_context(inputs)) - - -@pytest.mark.parametrize("stage", ["after_shared", "after_residual", "final_bf16"]) -def test_intermediate_and_cast_overflow(stage): - maxval = torch.finfo(torch.float32).max - values = { - "after_shared": (maxval, maxval, -maxval), - "after_residual": (maxval, 0.0, maxval), - "final_bf16": (maxval, 0.0, 0.0), - }[stage] - inputs = tuple(torch.tensor([[v]], dtype=torch.float32) for v in values) - with pytest.raises(MergeContractError, match=f"NON_FINITE: P6.merge.{stage}"): - shared_residual_merge_fwd(*inputs, context=make_context(inputs)) - - -@pytest.mark.parametrize("kind", ["routed_bf16", "shared_fp16", "broadcast", "strided"]) -def test_unsupported_payloads(kind): - inputs = list(make_inputs()) - if kind == "routed_bf16": - inputs[0] = inputs[0].bfloat16() - elif kind == "shared_fp16": - inputs[1] = inputs[1].half() - elif kind == "broadcast": - inputs[1] = inputs[1][:, :1].contiguous() - else: - inputs[1] = torch.ones(1, 6)[:, ::2] - with pytest.raises(MergeContractError): - shared_residual_merge_fwd(*inputs, context=make_context(inputs)) - - -@pytest.mark.parametrize("ep", [1, 2, 4, 8]) -def test_ep_replication_owner_local_mock(ep): - # No collectives: replicas exist on every mock rank; exactly one owner merges - # each token. This is WS2 semantic evidence, not multi-GPU certification. - inputs = tuple(v.repeat(16, 1) for v in make_inputs()) - canonical = shared_residual_merge_fwd(*inputs, context=make_context(inputs)) - gathered = torch.empty_like(canonical.token_moe_row_bf16) - for rank in range(ep): - tokens = tuple(range(rank, 16, ep)) - local = tuple(v[list(tokens)].contiguous() for v in inputs) - context = make_context(local, rank=rank, ep=ep, tokens=tokens) - result = shared_residual_merge_fwd(*local, context=context) - gathered[list(tokens)] = result.token_moe_row_bf16 - assert result.receipt["shared_applied_count"] == 1 - assert result.receipt["merge_order_hash"] == canonical.receipt["merge_order_hash"] - if ep > 1: - with pytest.raises(MergeContractError, match="SHARED_APPLIED_NOT_ONCE"): - shared_residual_merge_fwd(*local, context=replace(context, rank=(rank + 1) % ep)) - assert_bytes(gathered, canonical.token_moe_row_bf16) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_input_immutability_and_autocast(dtype): - inputs = make_inputs(dtype=dtype) - snapshots = tuple(v.clone() for v in inputs) - context = make_context(inputs) - baseline = shared_residual_merge_fwd(*inputs, context=context) - with torch.autocast("cpu", dtype=torch.bfloat16): - result = shared_residual_merge_fwd(*inputs, context=context) - assert_bytes(result.token_moe_row_bf16, baseline.token_moe_row_bf16) - for value, saved in zip(inputs, snapshots, strict=True): - assert_bytes(value, saved) - operations = result.receipt["actual_provenance"]["observed_aten_operations"] - assert [op["output_dtype"] for op in operations if op["operator"] == "aten.add.Tensor"] == [ - "torch.float32", - "torch.float32", - ] - assert sum(op["output_dtype"] == "torch.bfloat16" for op in operations) == 1 - - -@pytest.mark.parametrize("width", [1, 17, 129, 4096]) -def test_batch_chunk_padding_permutation(width): - generator = torch.Generator().manual_seed(71) - inputs = tuple(torch.randn(7, width, generator=generator) for _ in range(3)) - expected = shared_residual_merge_fwd(*inputs, context=make_context(inputs)).token_moe_row_bf16 - permutation = [6, 2, 0, 1, 5, 3, 4] - permuted = tuple(v[permutation] for v in inputs) - result = shared_residual_merge_fwd( - *permuted, context=make_context(permuted, tokens=tuple(permutation)) - ) - assert_bytes(result.token_moe_row_bf16, expected[permutation]) - for start, end in ((0, 1), (1, 4), (4, 7)): - chunk = tuple(v[start:end] for v in inputs) - result = shared_residual_merge_fwd( - *chunk, context=make_context(chunk, tokens=tuple(range(start, end))) - ) - assert_bytes(result.token_moe_row_bf16, expected[start:end]) - padded = tuple(torch.cat((v, torch.zeros(3, width))) for v in inputs) - result = shared_residual_merge_fwd(*padded, context=make_context(padded)) - assert_bytes(result.token_moe_row_bf16[:7], expected) - - -def test_duplicate_and_empty_token_batch(): - inputs = tuple(v.repeat(2, 1) for v in make_inputs()) - context = make_context(inputs) - context = replace(context, identity=replace(context.identity, global_token_ids=(1, 1))) - with pytest.raises(MergeContractError, match="AMBIGUOUS_GLOBAL_TOKEN_MAPPING"): - shared_residual_merge_fwd(*inputs, context=context) - empty = tuple(v[:0] for v in inputs) - with pytest.raises(MergeContractError, match="AMBIGUOUS_GLOBAL_TOKEN_MAPPING"): - shared_residual_merge_fwd(*empty, context=make_context(empty)) - - -def test_reference_registry_entry(): - from rl_engine.kernels.ops.pytorch.moe.shared_residual_merge import BACKEND_ID - from rl_engine.kernels.registry import KernelRegistry - from rl_engine.kernels.semantic_registry import ( - OperatorRequirements, - OperatorResolutionError, - OperatorSession, - ) - - session = OperatorSession(KernelRegistry().semantic) - resolution = session.resolve( - semantic_op="shared_residual_merge", - requested_backend=BACKEND_ID, - target="rollout", - requirements=OperatorRequirements(device="cpu", dtype="float32"), - ) - op = session.instantiate(resolution) - inputs = make_inputs() - result = op(*inputs, context=make_context(inputs)) - assert result.receipt["scope"] == "local_forward_reference" - assert resolution.descriptor.determinism_or_alignment_properties["reference_only"] is True - with pytest.raises(OperatorResolutionError): - session.resolve( - semantic_op="shared_residual_merge", - requested_backend=BACKEND_ID, - target="rollout", - requirements=OperatorRequirements( - device="cpu", dtype="float32", alignment_properties={"gpu_launch_observable": True} - ), - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="no CUDA/ROCm GPU: not certified") -@pytest.mark.parametrize("case", CASES, ids=lambda c: c["id"]) -def test_gpu_golden(case): - inputs = make_inputs(case, device="cuda") - result = shared_residual_merge_fwd( - *inputs, context=make_context(inputs, case_id=case["id"]), debug=True - ) - golden = torch.tensor(case["bf16_words"], dtype=torch.uint16).view(torch.bfloat16) - assert_bytes(result.token_moe_row_bf16, golden) - for key, field in zip(BOUNDARIES[:2], ("after_shared", "after_residual"), strict=True): - assert_bytes(result.debug_boundaries[key], torch.tensor(case[field], dtype=torch.float32)) - with torch.autocast("cuda", dtype=torch.bfloat16): - quiet = shared_residual_merge_fwd(*inputs, context=make_context(inputs, case_id=case["id"])) - assert_bytes(quiet.token_moe_row_bf16, result.token_moe_row_bf16) - assert quiet.receipt == result.receipt - - -def test_sealed_reference_evidence(tmp_path): - root = Path(os.environ.get("RL_KERNEL_T04_ARTIFACT_DIR", tmp_path)) - store = ArtifactStore(root) - required = ("t04_receipt.json", "t04_debug.pt") - for case in CASES: - inputs = make_inputs(case) - context = make_context(inputs, case_id=case["id"]) - result = shared_residual_merge_fwd(*inputs, context=context, debug=True) - experiment = { - "context": json.loads(json.dumps(asdict(context))), - "implementation": result.receipt["actual_provenance"], - } - # Identity/build changes cannot silently resume an existing experiment. - store.initialize_experiment(case["id"], experiment=experiment, plan=[]) - attempt = store.create_attempt(case["id"], case["id"]) - store.write_json(attempt, required[0], result.receipt) - store.write_tensor_bundle(attempt, required[1], result.debug_boundaries) - store.complete_attempt( - attempt, - required=required, - summary={ - "schema_version": "cross_config.complete.v1", - "case_id": case["id"], - "attempt_id": attempt.name, - "status": "LOCAL_REFERENCE_ONLY", - }, - ) - store.validate_completed_attempt(attempt, required=required, expected_case_id=case["id"]) - loaded = store.load_tensor_bundle(attempt / required[1])["tensors"] - receipt = json.loads((attempt / required[0]).read_text()) - for boundary in receipt["boundaries"]: - assert tensor_sha256(loaded[boundary["key"]]) == boundary["checksum"] - with pytest.raises(ArtifactError, match="resume metadata differs"): - store.initialize_experiment(case["id"], experiment={"wrong_run": True}, plan=[]) - # Exercise corruption on an isolated copy, never alter a sealed report. - corrupted = tmp_path / "corrupted" / attempt.name - shutil.copytree(attempt, corrupted) - (corrupted / "t04_receipt.json").write_text("{}") - with pytest.raises(ArtifactError, match="artifact hash does not match"): - store.validate_completed_attempt(corrupted, required=required)