diff --git a/docs/.nav.yml b/docs/.nav.yml index 73281687..92cb2256 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -15,6 +15,7 @@ nav: - Operators: - operators/README.md - operators/activation.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/rl_engine/alignment/testing/moe_merge.py b/rl_engine/alignment/testing/moe_merge.py new file mode 100644 index 00000000..fe2b17fe --- /dev/null +++ b/rl_engine/alignment/testing/moe_merge.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Validate MoE merge source identities and collect reference evidence. + +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 + +import hashlib +import json +from collections import Counter +from dataclasses import asdict, dataclass +from typing import Any + +import torch +from torch import Tensor + +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 = ( + "after_shared", + "after_residual", + "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, ...] + merge_order: tuple[str, ...] = ORDER + input_row_weighted: bool = True + + +@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") + 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, + "INVALID_COMBINE_PLAN", + "weighted routed input", + ) + _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) + + +def check_moe_merge( + routed: Tensor, + shared: Tensor, + residual: Tensor, + *, + context: MergeContext, + debug: bool = False, +) -> MergeResult: + """Validate producer evidence and collect reference boundaries for acceptance. + + Identity/discrete failures precede numeric checks. Checksums and finite checks + synchronize GPU tensors; this helper belongs to validation, not model forward. + """ + sources = _validate_context(context) + inputs = {"routed": routed, "shared": shared, "residual": residual} + _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, + "IDENTITY_DRIFT", + f"{role}.tensor_checksum", + ) + for role, value in inputs.items(): + _finite(value, role) + + 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", + ), + ) + 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 = { + "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.is_cuda else "cpu"), + "build_runtime": torch.version.hip or torch.version.cuda, + "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": instance.backend_id, + "implementation_fingerprint": instance.implementation_fingerprint, + } + for index, (key, value) in enumerate(boundary_tensors.items()) + ] + applied = Counter(context.upstream_events + ORDER) + receipt = { + "schema_version": RECEIPT_VERSION, + "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, + "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": 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", + "actual_provenance": provenance, + "boundaries": boundaries, + } + 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 new file mode 100644 index 00000000..df73c2d4 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/moe/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .merge import MoeMergeOp, shared_residual_merge_fwd + +__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 7070728a..a8de5f93 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -184,6 +184,31 @@ 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"}), + # Primary (routed) input dtype; shared/residual may also be BF16. + supported_dtypes=frozenset({"float32"}), + # Local merge only; distributed topology requirements belong to the caller. + supported_topologies={}, + determinism_or_alignment_properties={ + "algorithm": "fixed_order_fp32_merge", + "batch_invariant": True, + "deterministic": True, + "reference_only": True, + "forward_only": True, + "strict_observable": True, + "gpu_launch_observable": False, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=( + "rl_engine.kernels.ops.pytorch.moe.merge.MoeMergeOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="MoeMergeOp-v1", + ), OperatorBackendDescriptor( semantic_op="selected_logprob", backend_id="rlkernel.reference_logp", diff --git a/tests/fixtures/moe_merge.json b/tests/fixtures/moe_merge.json new file mode 100644 index 00000000..1818f200 --- /dev/null +++ b/tests/fixtures/moe_merge.json @@ -0,0 +1,32 @@ +{ + "schema_version": "moe_merge.fixtures.v1", + "cases": [ + { + "id": "simple_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": "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]], + "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": "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]], + "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_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)