[RL-Kernel] Delegate operator comparison standard

tests/test_rl_kernel_operator_comparison.py

@@ -0,0 +1,119 @@
+import importlib
+import sys
+from types import ModuleType
+
+import pytest
+
+
+def _drop_rl_engine_modules() -> None:
+    for name in list(sys.modules):
+        if name == "rl_engine" or name.startswith("rl_engine."):
+            sys.modules.pop(name, None)
+
+
+def _install_fake_rl_kernel_operator_comparison(monkeypatch) -> ModuleType:
+    rl_engine = ModuleType("rl_engine")
+    alignment = ModuleType("rl_engine.alignment")
+    cross_config = ModuleType("rl_engine.alignment.cross_config")
+    operator_comparison = ModuleType("rl_engine.alignment.cross_config.operator_comparison")
+
+    class FakeOperatorTolerance:
+        def __init__(self, *, atol=1e-6, rtol=1e-6):
+            self.atol = atol
+            self.rtol = rtol
+
+    operator_comparison.RLK_OP_LOGP = "logp"
+    operator_comparison.PHASE4_TARGET_OPERATORS = ("logp",)
+    operator_comparison.OPERATOR_COMPARISON_SPECS = {"logp": "spec"}
+    operator_comparison.OperatorTolerance = FakeOperatorTolerance
+    operator_comparison.iter_operator_comparison_specs = lambda: ("spec",)
+    operator_comparison.compare_operator_outputs = lambda op_name, train, infer, **kwargs: {
+        "op_name": op_name,
+        "train": train,
+        "infer": infer,
+        "kwargs": kwargs,
+    }
+
+    monkeypatch.setitem(sys.modules, "rl_engine", rl_engine)
+    monkeypatch.setitem(sys.modules, "rl_engine.alignment", alignment)
+    monkeypatch.setitem(sys.modules, "rl_engine.alignment.cross_config", cross_config)
+    monkeypatch.setitem(
+        sys.modules,
+        "rl_engine.alignment.cross_config.operator_comparison",
+        operator_comparison,
+    )
+    return operator_comparison
+
+
+@pytest.mark.unit
+def test_operator_comparison_adapter_import_does_not_import_rl_engine():
+    _drop_rl_engine_modules()
+    sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None)
+
+    importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison")
+
+    assert not any(name == "rl_engine" or name.startswith("rl_engine.") for name in sys.modules)
+
+
+@pytest.mark.unit
+def test_operator_comparison_adapter_delegates_to_rl_kernel(monkeypatch):
+    _drop_rl_engine_modules()
+    sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None)
+    fake = _install_fake_rl_kernel_operator_comparison(monkeypatch)
+    adapter = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison")
+
+    assert adapter.RLK_OP_LOGP == "logp"
+    assert adapter.PHASE4_TARGET_OPERATORS == ("logp",)
+    assert adapter.OPERATOR_COMPARISON_SPECS == {"logp": "spec"}
+    assert adapter.iter_operator_comparison_specs() == ("spec",)
+    assert adapter.OperatorTolerance(atol=0.5).atol == pytest.approx(0.5)
+    assert adapter.compare_operator_outputs("logp", "train", "infer") == {
+        "op_name": "logp",
+        "train": "train",
+        "infer": "infer",
+        "kwargs": {},
+    }
+    assert adapter.load_operator_comparison_module() is fake
+
+
+@pytest.mark.unit
+def test_package_reexports_operator_comparison_lazily(monkeypatch):
+    _drop_rl_engine_modules()
+    sys.modules.pop("vime.backends.rl_kernel_utils", None)
+    sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None)
+
+    package = importlib.import_module("vime.backends.rl_kernel_utils")
+
+    assert not any(name == "rl_engine" or name.startswith("rl_engine.") for name in sys.modules)
+
+    _install_fake_rl_kernel_operator_comparison(monkeypatch)
+
+    assert package.RLK_OP_LOGP == "logp"
+    assert package.iter_operator_comparison_specs() == ("spec",)
+
+
+@pytest.mark.unit
+def test_unavailable_operator_comparison_error_is_clear(monkeypatch):
+    _drop_rl_engine_modules()
+    sys.modules.pop("vime.backends.rl_kernel_utils.operator_comparison", None)
+    adapter = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison")
+
+    def fail_import(name, *args, **kwargs):
+        if name == "rl_engine.alignment.cross_config.operator_comparison":
+            raise ModuleNotFoundError("No module named 'rl_engine'")
+        return importlib.import_module(name, *args, **kwargs)
+
+    monkeypatch.setattr(adapter.importlib, "import_module", fail_import)
+
+    with pytest.raises(adapter.RlkOperatorComparisonUnavailable, match="operator comparison standard"):
+        adapter.load_operator_comparison_module()
+
+
+@pytest.mark.unit
+def test_vime_operator_comparison_does_not_carry_reference_implementations():
+    module = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison")
+    source = module.__loader__.get_source(module.__name__)
+
+    assert "rl_engine.alignment.cross_config.operator_comparison" in source
+    assert "def reference_rmsnorm" not in source
+    assert "torch.nn.functional" not in source

tests/utils/test_consistency_audit.py

@@ -77,15 +77,19 @@ def _metadata(sample: Sample | None = None, *, args: Namespace | None = None, **
 
 def _batch(*, metadata=None, layout=None):
     metadata = [_metadata()] if metadata is None else metadata
-    layout = [
-        {
-            "fingerprint": "layout-1",
-            "active_mask_density": 1.0,
-            "dp_rank": 0,
-            "microbatch_id": 0,
-            "microbatch_offset": 0,
-        }
-    ] if layout is None else layout
+    layout = (
+        [
+            {
+                "fingerprint": "layout-1",
+                "active_mask_density": 1.0,
+                "dp_rank": 0,
+                "microbatch_id": 0,
+                "microbatch_offset": 0,
+            }
+        ]
+        if layout is None
+        else layout
+    )
     return {
         "tokens": [torch.tensor([101, 201, 202])],
         "unconcat_tokens": [torch.tensor([101, 201, 202])],

vime/backends/rl_kernel_utils/__init__.py

@@ -1,5 +1,8 @@
 """RL-Kernel adapter and execution utilities owned by vime."""
 
+import importlib
+from typing import Any
+
 from vime.backends.rl_kernel_utils.adapter import (
     RLK_ALL_OPERATORS,
     RLK_OP_LINEAR_LOGP,
@@ -40,6 +43,67 @@
     select_execution_decision,
 )
 
+from vime.backends.rl_kernel_utils.operator_comparison import (
+    RlkOperatorComparisonUnavailable,
+    load_operator_comparison_module,
+)
+
+_OPERATOR_COMPARISON_EXPORTS = frozenset(
+    {
+        "OPERATOR_COMPARISON_SPECS",
+        "PHASE4_TARGET_OPERATORS",
+        "RLK_OP_ATTENTION",
+        "RLK_OP_DPO_FRAGMENT",
+        "RLK_OP_EMBEDDING",
+        "RLK_OP_GRPO_FRAGMENT",
+        "RLK_OP_LM_HEAD",
+        "RLK_OP_LOGP",
+        "RLK_OP_MATMUL_PROJECTION",
+        "RLK_OP_PPO_FRAGMENT",
+        "RLK_OP_RATIO_KL",
+        "RLK_OP_RMSNORM",
+        "RLK_OP_ROPE",
+        "RLK_OP_SWIGLU",
+        "BatchInvarianceCase",
+        "ForwardChainComparisonResult",
+        "ForwardChainStep",
+        "OperatorComparisonResult",
+        "OperatorComparisonSpec",
+        "OperatorPair",
+        "OperatorTolerance",
+        "StrictBackendAdmissionReport",
+        "build_single_card_batch_invariance_cases",
+        "build_strict_backend_admission_report",
+        "compare_batch_invariance",
+        "compare_operator_outputs",
+        "compare_operator_pair",
+        "get_operator_comparison_spec",
+        "iter_operator_comparison_specs",
+        "reference_attention",
+        "reference_embedding",
+        "reference_linear_logp",
+        "reference_lm_head",
+        "reference_matmul_projection",
+        "reference_ppo_fragment",
+        "reference_ratio_kl",
+        "reference_rmsnorm",
+        "reference_rope",
+        "reference_selected_logprobs",
+        "reference_swiglu",
+        "run_deterministic_repeatability_check",
+        "run_forward_chain_comparison",
+        "run_reference_operator",
+    }
+)
+
+
+def __getattr__(name: str) -> Any:
+    if name not in _OPERATOR_COMPARISON_EXPORTS:
+        raise AttributeError(name)
+    module = importlib.import_module("vime.backends.rl_kernel_utils.operator_comparison")
+    return getattr(module, name)
+
+
 __all__ = [
     "BackendCapability",
     "CapabilityQueryResult",
@@ -76,4 +140,49 @@
     "linear_logp_inputs_from_vime",
     "rlk_policy_context_from_args",
     "runtime_batch_metadata_from_vime_batch",
+    "RlkOperatorComparisonUnavailable",
+    "load_operator_comparison_module",
+    "BatchInvarianceCase",
+    "ForwardChainComparisonResult",
+    "ForwardChainStep",
+    "OPERATOR_COMPARISON_SPECS",
+    "OperatorComparisonResult",
+    "OperatorComparisonSpec",
+    "OperatorPair",
+    "OperatorTolerance",
+    "PHASE4_TARGET_OPERATORS",
+    "RLK_OP_ATTENTION",
+    "RLK_OP_DPO_FRAGMENT",
+    "RLK_OP_EMBEDDING",
+    "RLK_OP_GRPO_FRAGMENT",
+    "RLK_OP_LM_HEAD",
+    "RLK_OP_LOGP",
+    "RLK_OP_MATMUL_PROJECTION",
+    "RLK_OP_PPO_FRAGMENT",
+    "RLK_OP_RATIO_KL",
+    "RLK_OP_RMSNORM",
+    "RLK_OP_ROPE",
+    "RLK_OP_SWIGLU",
+    "StrictBackendAdmissionReport",
+    "build_single_card_batch_invariance_cases",
+    "build_strict_backend_admission_report",
+    "compare_batch_invariance",
+    "compare_operator_outputs",
+    "compare_operator_pair",
+    "get_operator_comparison_spec",
+    "iter_operator_comparison_specs",
+    "reference_attention",
+    "reference_embedding",
+    "reference_linear_logp",
+    "reference_lm_head",
+    "reference_matmul_projection",
+    "reference_ppo_fragment",
+    "reference_ratio_kl",
+    "reference_rmsnorm",
+    "reference_rope",
+    "reference_selected_logprobs",
+    "reference_swiglu",
+    "run_deterministic_repeatability_check",
+    "run_forward_chain_comparison",
+    "run_reference_operator",
 ]

vime/backends/rl_kernel_utils/operator_comparison.py

@@ -0,0 +1,84 @@
+"""Lazy re-export of RL-Kernel-owned operator comparison helpers."""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+_RLK_OPERATOR_COMPARISON_MODULE = "rl_engine.alignment.cross_config.operator_comparison"
+
+_OPERATOR_COMPARISON_EXPORTS = frozenset(
+    {
+        "OPERATOR_COMPARISON_SPECS",
+        "PHASE4_TARGET_OPERATORS",
+        "RLK_OP_ATTENTION",
+        "RLK_OP_DPO_FRAGMENT",
+        "RLK_OP_EMBEDDING",
+        "RLK_OP_GRPO_FRAGMENT",
+        "RLK_OP_LM_HEAD",
+        "RLK_OP_LOGP",
+        "RLK_OP_MATMUL_PROJECTION",
+        "RLK_OP_PPO_FRAGMENT",
+        "RLK_OP_RATIO_KL",
+        "RLK_OP_RMSNORM",
+        "RLK_OP_ROPE",
+        "RLK_OP_SWIGLU",
+        "BatchInvarianceCase",
+        "ForwardChainComparisonResult",
+        "ForwardChainStep",
+        "OperatorComparisonResult",
+        "OperatorComparisonSpec",
+        "OperatorPair",
+        "OperatorTolerance",
+        "StrictBackendAdmissionReport",
+        "build_single_card_batch_invariance_cases",
+        "build_strict_backend_admission_report",
+        "compare_batch_invariance",
+        "compare_operator_outputs",
+        "compare_operator_pair",
+        "get_operator_comparison_spec",
+        "iter_operator_comparison_specs",
+        "reference_attention",
+        "reference_embedding",
+        "reference_linear_logp",
+        "reference_lm_head",
+        "reference_matmul_projection",
+        "reference_ppo_fragment",
+        "reference_ratio_kl",
+        "reference_rmsnorm",
+        "reference_rope",
+        "reference_selected_logprobs",
+        "reference_swiglu",
+        "run_deterministic_repeatability_check",
+        "run_forward_chain_comparison",
+        "run_reference_operator",
+    }
+)
+
+
+class RlkOperatorComparisonUnavailable(RuntimeError):
+    """Raised when RL-Kernel's operator comparison standard cannot be imported."""
+
+
+def load_operator_comparison_module() -> Any:
+    try:
+        return importlib.import_module(_RLK_OPERATOR_COMPARISON_MODULE)
+    except Exception as exc:
+        raise RlkOperatorComparisonUnavailable(f"RL-Kernel operator comparison standard is unavailable; install an RL-Kernel build that exposes {_RLK_OPERATOR_COMPARISON_MODULE!r}.") from exc
+
+
+def __getattr__(name: str) -> Any:
+    if name not in _OPERATOR_COMPARISON_EXPORTS:
+        raise AttributeError(name)
+    return getattr(load_operator_comparison_module(), name)
+
+
+def __dir__() -> list[str]:
+    return sorted((*globals(), *_OPERATOR_COMPARISON_EXPORTS))
+
+
+__all__ = [
+    "RlkOperatorComparisonUnavailable",
+    "load_operator_comparison_module",
+    *_OPERATOR_COMPARISON_EXPORTS,
+]

vime/utils/consistency_audit.py

@@ -22,11 +22,7 @@
     stable_fingerprint,
     validate_samples_consistency_metadata,
 )
-from vime.utils.dlogp_diagnostics import (
-    DlogpAuditReport,
-    compute_dlogp_diagnostics,
-    get_rlk_consistency_mode,
-)
+from vime.utils.dlogp_diagnostics import DlogpAuditReport, compute_dlogp_diagnostics, get_rlk_consistency_mode
 
 CONSISTENCY_AUDIT_SCHEMA_VERSION = 1
 AUDIT_REQUIRED_METADATA_FIELDS = (
@@ -289,8 +285,7 @@ def build_consistency_replay_manifest(
                 "total_length": _int_or_none(_sequence_value(_batch_get(batch, "total_lengths"), position)),
                 "response_length": response_length,
                 "active_token_count": _active_token_count(loss_mask, response_length),
-                "has_rollout_log_probs": _batch_sequence_value(batch, "rollout_log_probs", position)
-                is not None,
+                "has_rollout_log_probs": _batch_sequence_value(batch, "rollout_log_probs", position) is not None,
                 "consistency_metadata_fingerprint": None if record is None else record.get("fingerprint"),
                 "batch_layout_fingerprint": _first_present(
                     None if layout is None else layout.get("fingerprint"),
@@ -545,10 +540,7 @@ def _with_audit_required_metadata_issues(
 
     warnings = list(validation.warnings)
     failures = list(validation.failures)
-    seen = {
-        (issue.code, issue.sample_index, issue.rollout_id, issue.field)
-        for issue in (*warnings, *failures)
-    }
+    seen = {(issue.code, issue.sample_index, issue.rollout_id, issue.field) for issue in (*warnings, *failures)}
 
     for position, record in enumerate(records):
         if record is None: