From 0b28e3e7231ce3e6962179ccda7e4bc96320ee03 Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sun, 19 Jul 2026 08:29:04 +0800 Subject: [PATCH 01/33] feat(kernels): add semantic operator catalog --- rl_engine/kernels/registry.py | 143 ++++- rl_engine/kernels/semantic_registry.py | 790 +++++++++++++++++++++++++ 2 files changed, 901 insertions(+), 32 deletions(-) create mode 100644 rl_engine/kernels/semantic_registry.py diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 041ed3e1..5f3719de 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -1,11 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +from __future__ import annotations + import importlib import os from enum import Enum, EnumMeta from typing import Any, Dict, Optional, Set, Type +from rl_engine.kernels.semantic_registry import ( + OperatorBackendDescriptor, + OperatorFallbackPolicy, + OperatorLifecycle, + SemanticOperatorCatalog, +) from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -16,9 +24,11 @@ class _KernelEnumMeta(EnumMeta): def __getitem__(cls, name: str): try: return super().__getitem__(name) - except KeyError as e: + except KeyError as exc: valid_ops = ", ".join(cls.__members__.keys()) - raise ValueError(f"Operator '{name}' not found. Supported backends: {valid_ops}") from e + raise ValueError( + f"Operator '{name}' not found. Supported backends: {valid_ops}" + ) from exc class OpBackend(Enum, metaclass=_KernelEnumMeta): @@ -78,6 +88,57 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_EMBEDDING = "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp" +def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: + return ( + OperatorBackendDescriptor( + semantic_op="selected_logprob", + backend_id="rlkernel.reference_logp", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu", "cuda", "rocm"}), + supported_dtypes=frozenset({"float32", "bfloat16", "float16"}), + supported_topologies={ + "rollout": { + "world_size": (1,), + "tensor_parallel_size": (1,), + "context_parallel_size": (1,), + }, + "training": { + "world_size": (1,), + "sharding": ("unsharded",), + }, + }, + determinism_or_alignment_properties={ + "algorithm": "pytorch.log_softmax_gather", + "batch_invariant": True, + "deterministic": True, + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=( + "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="NativeLogpOp-selected-logprob-v1", + ), + OperatorBackendDescriptor( + semantic_op="selected_logprob", + backend_id="native", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"*"}), + supported_dtypes=frozenset({"*"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "selection": "runtime_native", + "strict_observable": False, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=None, + fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, + version_or_build_fingerprint="runtime-native-unresolved-v1", + ), + ) + + def resolve_logp_op_type( logp_backend: Optional[str] = None, *, @@ -126,14 +187,12 @@ def resolve_logp_op_type( class KernelRegistry: - """ - Central dispatcher for high-performance kernels. - Handles dynamic routing between ROCm and CUDA backends at runtime. - """ + """Legacy hardware dispatcher plus a composed semantic operator catalog.""" def __init__(self): self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() + self.semantic = SemanticOperatorCatalog(_default_semantic_descriptors()) self._priority_map = { "cuda": { @@ -163,23 +222,33 @@ def __init__(self): OpBackend.CUDA_DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE, ], - "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], + "attn": [ + OpBackend.FLASH_ATTN, + OpBackend.TRITON_GENERIC, + OpBackend.PYTORCH_ATTN, + ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], - "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], + "linear_logp": [ + OpBackend.TRITON_LINEAR_LOGP, + OpBackend.PYTORCH_LINEAR_LOGP, + ], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], - # Default dispatch logic for new operators "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], }, "rocm": { - "logp": [OpBackend.ROCM_AITER, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], + "logp": [ + OpBackend.ROCM_AITER, + OpBackend.TRITON_GENERIC, + OpBackend.PYTORCH_NATIVE, + ], "logp_deterministic": [OpBackend.PYTORCH_NATIVE], "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [ @@ -191,7 +260,10 @@ def __init__(self): "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], - "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], + "linear_logp": [ + OpBackend.TRITON_LINEAR_LOGP, + OpBackend.PYTORCH_LINEAR_LOGP, + ], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], @@ -237,14 +309,15 @@ def _adjust_priority_from_env(self): OpBackend.ROCM_FLASH_ATTN, OpBackend.TRITON_GENERIC, ] - elif rocm_attn_backend and rocm_attn_backend not in {"native", "pytorch", "sdpa"}: + elif rocm_attn_backend: logger.warning( "Unknown RL_KERNEL_ROCM_ATTN_BACKEND=%s; using default ROCm attention priority.", rocm_attn_backend, ) def _adjust_priority_for_hardware(self): - """Adjust CUDA priorities for hardware-gated experimental and production kernels.""" + """Adjust CUDA priorities for hardware-gated kernels.""" + if device_ctx.device_type != "cuda": return try: @@ -266,8 +339,6 @@ def _adjust_priority_for_hardware(self): if OpBackend.CUDA_FUSED_LOGP_SM90 not in logp_list: logp_list.insert(0, OpBackend.CUDA_FUSED_LOGP_SM90) - # The fused linear-logp SM90 kernel uses TMA bulk-tensor copies built - # for sm_90a -- gate strictly on cc_major == 9 (Hopper), not >= 9. linear_logp_compiled = _EXT_AVAILABLE and hasattr(_C, "fused_linear_logp_sm90") if linear_logp_compiled and cc_major == 9: ll_list = self._priority_map["cuda"]["linear_logp"] @@ -278,25 +349,26 @@ def _adjust_priority_for_hardware(self): f"SM{cc}: fused linear-logp SM90 kernel not compiled into _C; " "using generic linear-logp backend." ) - except Exception as e: - logger.warning(f"Failed to probe device capability: {e}") + except Exception as exc: + logger.warning(f"Failed to probe device capability: {exc}") def get_op(self, op_type: str) -> Any: - """Core distribution logic: Automatically select the best operator - based on hardware and priority. - """ + """Select the best legacy operator based on hardware and priority.""" + if device_ctx.is_rocm: platform = "rocm" elif device_ctx.device_type == "cuda": platform = "cuda" else: platform = "cpu" - candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) + candidates = self._priority_map.get(platform, {}).get( + op_type, + [OpBackend.PYTORCH_NATIVE], + ) for backend in candidates: if backend.name in self._instance_cache: return self._instance_cache[backend.name] - if backend.name in self._failed_backends: continue @@ -306,8 +378,8 @@ def get_op(self, op_type: str) -> Any: op_instance = op_class() self._instance_cache[backend.name] = op_instance return op_instance - except Exception as e: - logger.error(f"Failed to instantiate {backend.name}: {e}") + except Exception as exc: + logger.error(f"Failed to instantiate {backend.name}: {exc}") self._failed_backends.add(backend.name) else: self._failed_backends.add(backend.name) @@ -315,23 +387,30 @@ def get_op(self, op_type: str) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") def _load_backend(self, backend: OpBackend) -> Optional[Type]: - """Dynamic loading technique: Import modules only when needed - and check environment dependencies. - """ + """Import a legacy backend and distinguish wrapper bugs from absence.""" + module_path, class_name = backend.value.rsplit(".", 1) try: module = importlib.import_module(module_path) return getattr(module, class_name) - except (ImportError, AttributeError, ModuleNotFoundError) as e: - missing_module = str(e.name) if hasattr(e, "name") else "" + except (ImportError, AttributeError, ModuleNotFoundError) as exc: + missing_module = str(exc.name) if hasattr(exc, "name") else "" is_missing_backend = missing_module and ( missing_module == module_path or module_path.startswith(missing_module) ) if missing_module and "rl_engine" in missing_module and not is_missing_backend: - logger.critical(f"Internal wrapper implementation bug in '{module_path}': {e}") - raise e - logger.warning(f"Backend {backend.name} unavailable: {e}. Falling back...") + logger.critical(f"Internal wrapper implementation bug in '{module_path}': {exc}") + raise + logger.warning(f"Backend {backend.name} unavailable: {exc}. Falling back...") return None kernel_registry = KernelRegistry() + + +__all__ = [ + "KernelRegistry", + "OpBackend", + "kernel_registry", + "resolve_logp_op_type", +] diff --git a/rl_engine/kernels/semantic_registry.py b/rl_engine/kernels/semantic_registry.py new file mode 100644 index 00000000..99e152d9 --- /dev/null +++ b/rl_engine/kernels/semantic_registry.py @@ -0,0 +1,790 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Exact semantic-operator catalog with case-local instantiation state.""" + +from __future__ import annotations + +import hashlib +import importlib +import inspect +import json +from dataclasses import dataclass, field, fields, replace +from enum import Enum +from pathlib import Path +from types import CodeType, MappingProxyType +from typing import Any, Callable, Iterable, Mapping, Optional, Sequence, cast + + +class OperatorLifecycle(str, Enum): + REQUEST = "request" + ENGINE_CONSTRUCTION = "engine_construction" + DISTRIBUTED_CONTEXT = "distributed_context" + PROCESS = "process" + + +class OperatorFallbackPolicy(str, Enum): + ERROR = "error" + DECLARED = "declared" + RUNTIME_MANAGED = "runtime_managed" + + +@dataclass(frozen=True) +class OperatorResolutionPolicy: + strict: bool = True + allow_test_backends: bool = False + + +_Policy = Optional[OperatorResolutionPolicy] + + +class _JsonRecord: + def to_dict(self) -> dict[str, Any]: + return { + item.name: _json_value(getattr(self, item.name)) for item in fields(cast(Any, self)) + } + + +@dataclass(frozen=True) +class OperatorRequirements(_JsonRecord): + device: str + dtype: str + topology: Mapping[str, Any] = field(default_factory=dict) + alignment_properties: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "rlkernel.semantic_operator.requirements.v1" + + def __post_init__(self) -> None: + normalized = ( + ("device", _normalize_device(self.device)), + ("dtype", _normalize_dtype(self.dtype)), + ("topology", _freeze(self.topology)), + ("alignment_properties", _freeze(self.alignment_properties)), + ) + for name, value in normalized: + object.__setattr__(self, name, value) + + +@dataclass(frozen=True) +class OperatorBackendDescriptor(_JsonRecord): + semantic_op: str + backend_id: str + supported_targets: frozenset[str] + supported_devices: frozenset[str] + supported_dtypes: frozenset[str] + supported_topologies: Mapping[str, Any] + determinism_or_alignment_properties: Mapping[str, Any] + lifecycle: OperatorLifecycle + implementation_class_or_factory: Optional[str | Callable[..., Any]] + fallback_policy: OperatorFallbackPolicy + version_or_build_fingerprint: str + is_smoke_only: bool = False + schema_version: str = "rlkernel.semantic_operator.backend_descriptor.v1" + + def __post_init__(self) -> None: + values = { + "semantic_op": self.semantic_op.strip(), + "backend_id": self.backend_id.strip(), + "supported_targets": _normalized_values(self.supported_targets, str), + "supported_devices": _normalized_values(self.supported_devices, _normalize_device), + "supported_dtypes": _normalized_values(self.supported_dtypes, _normalize_dtype), + } + for name, value in values.items(): + if not value: + raise ValueError(f"{name} must not be empty") + if not self.version_or_build_fingerprint.strip(): + raise ValueError("version_or_build_fingerprint must not be empty") + values.update( + supported_topologies=_freeze(self.supported_topologies), + determinism_or_alignment_properties=_freeze(self.determinism_or_alignment_properties), + lifecycle=OperatorLifecycle(self.lifecycle), + fallback_policy=OperatorFallbackPolicy(self.fallback_policy), + ) + for name, value in values.items(): + object.__setattr__(self, name, value) + + @property + def implementation_reference(self) -> Optional[str]: + return _reference(self.implementation_class_or_factory) + + @property + def is_strictly_observable(self) -> bool: + return bool( + self.determinism_or_alignment_properties.get( + "strict_observable", self.implementation_class_or_factory is not None + ) + ) + + @property + def descriptor_fingerprint(self) -> str: + return _fingerprint(self.to_dict(include_descriptor_fingerprint=False)) + + def to_dict(self, *, include_descriptor_fingerprint: bool = True) -> dict[str, Any]: + result = super().to_dict() + result["implementation_class_or_factory"] = self.implementation_reference + if include_descriptor_fingerprint: + result["descriptor_fingerprint"] = self.descriptor_fingerprint + return result + + +@dataclass(frozen=True) +class OperatorCapabilityDecision(_JsonRecord): + capability: str + requested: Any + supported: Any + passed: bool + reason: str + + +@dataclass(frozen=True) +class OperatorResolutionTrace(_JsonRecord): + semantic_op: str + requested_backend: str + target: str + strict: bool + status: str + concrete_backend: Optional[str] + implementation_reference: Optional[str] + descriptor_fingerprint: Optional[str] + capability_decisions: tuple[OperatorCapabilityDecision, ...] + fallback_attempts: tuple[str, ...] = () + schema_version: str = "rlkernel.semantic_operator.resolution_trace.v1" + + +@dataclass(frozen=True) +class OperatorResolution(_JsonRecord): + descriptor: OperatorBackendDescriptor + requirements: OperatorRequirements + target: str + strict: bool + trace: OperatorResolutionTrace + schema_version: str = "rlkernel.semantic_operator.resolution.v1" + + +@dataclass(frozen=True) +class OperatorInstanceProvenance(_JsonRecord): + semantic_op: str + backend_id: str + target: str + factory_reference: str + concrete_implementation: str + descriptor_fingerprint: str + implementation_fingerprint: str + instance_fingerprint: str + factory_options: Mapping[str, Any] = field(default_factory=dict) + factory_options_fingerprint: str = "" + schema_version: str = "rlkernel.semantic_operator.instance_provenance.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "factory_options", _freeze(self.factory_options)) + + +@dataclass(frozen=True) +class _InstanceRecord: + instance: Any + descriptor_fingerprint: str + target: str + factory: Callable[..., Any] + factory_options: Mapping[str, Any] + + +class OperatorRegistrationError(ValueError): + pass + + +class OperatorResolutionError(RuntimeError): + def __init__(self, message: str, trace: OperatorResolutionTrace): + super().__init__(message) + self.trace = trace + + +class OperatorInstantiationError(RuntimeError): + pass + + +class SemanticOperatorCatalog: + def __init__(self, descriptors: Iterable[OperatorBackendDescriptor] = ()): + self._descriptors: dict[tuple[str, str], OperatorBackendDescriptor] = {} + for descriptor in descriptors: + self.register_backend(descriptor) + + def register_backend( + self, + descriptor: OperatorBackendDescriptor, + *, + replace: bool = False, + ) -> None: + if not isinstance(descriptor, OperatorBackendDescriptor): + raise TypeError("descriptor must be an OperatorBackendDescriptor") + key = (descriptor.semantic_op, descriptor.backend_id) + if key in self._descriptors and not replace: + raise OperatorRegistrationError(f"operator backend is already registered: {key!r}") + self._descriptors[key] = descriptor + + def backend_descriptor( + self, + semantic_op: str, + backend_id: str, + ) -> Optional[OperatorBackendDescriptor]: + return self._descriptors.get((semantic_op.strip(), backend_id.strip())) + + def backend_descriptors( + self, + semantic_op: Optional[str] = None, + ) -> tuple[OperatorBackendDescriptor, ...]: + values: Iterable[OperatorBackendDescriptor] = self._descriptors.values() + if semantic_op is not None: + normalized = semantic_op.strip() + values = (value for value in values if value.semantic_op == normalized) + return tuple(sorted(values, key=lambda value: (value.semantic_op, value.backend_id))) + + def session(self, policy: _Policy = None) -> OperatorSession: + return OperatorSession(self, policy=policy) + + def _resolve( + self, + *, + semantic_op: str, + requested_backend: str, + target: str, + requirements: OperatorRequirements, + policy: OperatorResolutionPolicy, + ) -> OperatorResolution: + semantic_op = semantic_op.strip() + requested_backend = requested_backend.strip() + target = target.strip().lower() + if not semantic_op or not requested_backend or not target: + raise ValueError("semantic_op, requested_backend, and target must not be empty") + if not isinstance(requirements, OperatorRequirements): + raise TypeError("requirements must be an OperatorRequirements") + + descriptor = self.backend_descriptor(semantic_op, requested_backend) + if descriptor is None: + decision = _decision( + "registration", + requested_backend, + [item.backend_id for item in self.backend_descriptors(semantic_op)], + passed=False, + ) + trace = _trace( + semantic_op, + requested_backend, + target, + policy, + "unsupported", + (decision,), + ) + raise OperatorResolutionError( + f"exact operator backend {requested_backend!r} is not registered; " + "no fallback was attempted", + trace, + ) + + topology_capabilities = _target_topology_capabilities( + descriptor.supported_topologies, + target, + ) + topology_ok = topology_capabilities is not None and _supports_complete_mapping( + topology_capabilities, + requirements.topology, + ) + decisions = ( + _decision("target", target, descriptor.supported_targets), + _decision( + "smoke_opt_in", + descriptor.is_smoke_only, + policy.allow_test_backends, + not descriptor.is_smoke_only or policy.allow_test_backends, + ), + _decision("device", requirements.device, descriptor.supported_devices), + _decision("dtype", requirements.dtype, descriptor.supported_dtypes), + _decision( + "topology", + requirements.topology, + topology_capabilities, + topology_ok, + ), + _decision( + "alignment_properties", + requirements.alignment_properties, + descriptor.determinism_or_alignment_properties, + ), + _decision( + "strict_observability", + policy.strict, + descriptor.is_strictly_observable, + not policy.strict or descriptor.is_strictly_observable, + ), + _decision( + "fallback_policy", + "error" if policy.strict else "declared", + descriptor.fallback_policy.value, + not policy.strict or descriptor.fallback_policy is OperatorFallbackPolicy.ERROR, + ), + ) + failed = tuple(item for item in decisions if not item.passed) + observable = descriptor.is_strictly_observable + status = "unsupported" if failed else ("resolved" if observable else "unobservable") + trace = _trace( + semantic_op, + requested_backend, + target, + policy, + status, + decisions, + descriptor, + ) + if failed: + raise OperatorResolutionError( + f"exact operator backend {requested_backend!r} is unsupported: " + + "; ".join(item.reason for item in failed), + trace, + ) + return OperatorResolution(descriptor, requirements, target, policy.strict, trace) + + +class OperatorSession: + def __init__(self, catalog: SemanticOperatorCatalog, policy: _Policy = None): + if not isinstance(catalog, SemanticOperatorCatalog): + raise TypeError("catalog must be a SemanticOperatorCatalog") + self.catalog = catalog + self.policy = policy or OperatorResolutionPolicy() + self._cache: dict[str, Any] = {} + self._records: dict[int, _InstanceRecord] = {} + + def resolve( + self, + *, + semantic_op: str, + requested_backend: str, + target: str, + requirements: OperatorRequirements, + policy: _Policy = None, + strict: Optional[bool] = None, + ) -> OperatorResolution: + return self.catalog._resolve( + semantic_op=semantic_op, + requested_backend=requested_backend, + target=target, + requirements=requirements, + policy=_resolve_policy(policy or self.policy, strict), + ) + + def instantiate( + self, + resolution: OperatorResolution, + *, + factory_kwargs: Optional[Mapping[str, Any]] = None, + cache: bool = False, + ) -> Any: + if not isinstance(resolution, OperatorResolution): + raise TypeError("resolution must be an OperatorResolution") + descriptor = resolution.descriptor + implementation = descriptor.implementation_class_or_factory + if resolution.trace.status != "resolved" or implementation is None: + raise OperatorInstantiationError( + f"backend {descriptor.backend_id!r} has no exact implementation" + ) + options = dict(factory_kwargs or {}) + cache_key = _fingerprint( + { + "descriptor": descriptor.descriptor_fingerprint, + "target": resolution.target, + "requirements": resolution.requirements.to_dict(), + "options": options, + } + ) + if cache and cache_key in self._cache: + return self._cache[cache_key] + factory = _load_factory(implementation) + try: + instance = factory(**options) + except Exception as exc: + raise OperatorInstantiationError( + f"failed to instantiate backend {descriptor.backend_id!r}: {exc}" + ) from exc + if instance is None: + raise OperatorInstantiationError("operator factory returned None") + self._records[id(instance)] = _InstanceRecord( + instance, + descriptor.descriptor_fingerprint, + resolution.target, + factory, + _freeze(options), + ) + if cache: + self._cache[cache_key] = instance + return instance + + def instance_provenance( + self, + resolution: OperatorResolution, + instance: Any, + ) -> OperatorInstanceProvenance: + descriptor = resolution.descriptor + record = self._records.get(id(instance)) + if ( + record is None + or record.instance is not instance + or record.descriptor_fingerprint != descriptor.descriptor_fingerprint + or record.target != resolution.target + ): + raise OperatorInstantiationError( + "operator instance does not match this session resolution" + ) + factory_reference = descriptor.implementation_reference + concrete = _reference(type(instance)) + if factory_reference is None or concrete is None: + raise OperatorInstantiationError("operator implementation is not observable") + options_fingerprint = _fingerprint(record.factory_options) + implementation_fingerprint = operator_implementation_fingerprint( + record.factory, + instance, + ) + instance_fingerprint = operator_instance_fingerprint( + descriptor_fingerprint=descriptor.descriptor_fingerprint, + factory_reference=factory_reference, + concrete_implementation=concrete, + implementation_fingerprint=implementation_fingerprint, + factory_options_fingerprint=options_fingerprint, + ) + return OperatorInstanceProvenance( + descriptor.semantic_op, + descriptor.backend_id, + resolution.target, + factory_reference, + concrete, + descriptor.descriptor_fingerprint, + implementation_fingerprint, + instance_fingerprint, + record.factory_options, + options_fingerprint, + ) + + def clear_instance_cache(self) -> None: + self._cache.clear() + + +def operator_implementation_fingerprint( + implementation: str | Callable[..., Any], + instance: Any, +) -> str: + return implementation_fingerprint( + implementation, + instance=instance, + entrypoints=("apply_fp32", "__call__"), + ) + + +def implementation_fingerprint( + implementation: str | Callable[..., Any], + *, + instance: Any = None, + entrypoints: Sequence[str] = (), +) -> str: + """Fingerprint executable code, not only its import reference. + + The identity includes source or bytecode for the resolved factory, its + concrete class, the defining modules, and explicitly named runtime entry + points. Module content covers helper functions called by an entry point; + callable identities additionally make in-process replacements observable. + """ + + factory = _load_factory(implementation) + concrete_type = type(instance) if instance is not None else None + runtime_entrypoints = {} + if instance is not None: + for name in sorted(set(entrypoints)): + value = getattr(instance, name, None) + if callable(value): + runtime_entrypoints[name] = _callable_identity(value) + return _fingerprint( + { + "factory": _implementation_identity(factory), + "concrete_type": ( + _implementation_identity(concrete_type) if concrete_type is not None else None + ), + "runtime_entrypoints": runtime_entrypoints, + } + ) + + +def operator_instance_fingerprint(**identity: str) -> str: + return _fingerprint(identity) + + +def _trace( + semantic_op: str, + backend: str, + target: str, + policy: OperatorResolutionPolicy, + status: str, + decisions: tuple[OperatorCapabilityDecision, ...], + descriptor: Optional[OperatorBackendDescriptor] = None, +) -> OperatorResolutionTrace: + observable = descriptor is not None and descriptor.is_strictly_observable + return OperatorResolutionTrace( + semantic_op, + backend, + target, + policy.strict, + status, + ( + descriptor.backend_id + if descriptor is not None and observable and status != "unsupported" + else None + ), + descriptor.implementation_reference if descriptor else None, + descriptor.descriptor_fingerprint if descriptor else None, + decisions, + ) + + +def _decision( + capability: str, + requested: Any, + supported: Any, + passed: Optional[bool] = None, +) -> OperatorCapabilityDecision: + passed = _supports(supported, requested) if passed is None else passed + actionable = { + "smoke_opt_in": "smoke backend use requires explicit opt-in", + "strict_observability": "runtime-native implementation is not exactly observable", + "fallback_policy": "strict resolution forbids declared or runtime fallback", + } + return OperatorCapabilityDecision( + capability, + requested, + supported, + passed, + ( + f"{capability} is supported" + if passed + else actionable.get(capability, f"{capability} is unsupported") + ), + ) + + +def _supports(supported: Any, requested: Any) -> bool: + if isinstance(supported, str) and supported in {"*", "any"}: + return True + if isinstance(supported, Mapping): + if not isinstance(requested, Mapping): + return False + wildcard = supported.get("*") + return all( + _supports(supported.get(key, wildcard), value) + for key, value in requested.items() + if key in supported or wildcard is not None + ) and all(key in supported or wildcard is not None for key in requested) + if isinstance(supported, (set, frozenset, tuple, list)): + if isinstance(requested, (set, frozenset, tuple, list)): + return all(any(_supports(item, value) for item in supported) for value in requested) + return any(_supports(item, requested) for item in supported) + return supported == requested + + +def _target_topology_capabilities(supported: Any, target: str) -> Any: + if not isinstance(supported, Mapping): + return supported + targeted = any(key in supported for key in ("rollout", "training")) + if not targeted: + return supported + return supported.get(target, supported.get("*")) + + +def _supports_complete_mapping(supported: Any, requested: Any) -> bool: + if isinstance(supported, Mapping) and "*" not in supported: + if not isinstance(requested, Mapping) or any(key not in requested for key in supported): + return False + return _supports(supported, requested) + + +def _resolve_policy(policy: _Policy, strict: Optional[bool]) -> OperatorResolutionPolicy: + policy = policy or OperatorResolutionPolicy() + return policy if strict is None else replace(policy, strict=strict) + + +def _load_factory(value: str | Callable[..., Any]) -> Callable[..., Any]: + if callable(value): + return value + try: + module_name, attribute = value.rsplit(".", 1) + factory = getattr(importlib.import_module(module_name), attribute) + except (ValueError, ImportError, AttributeError, ModuleNotFoundError) as exc: + raise OperatorInstantiationError(f"operator factory {value!r} is unavailable") from exc + if not callable(factory): + raise OperatorInstantiationError(f"operator factory {value!r} is not callable") + return factory + + +def _reference(value: Any) -> Optional[str]: + if value is None or isinstance(value, str): + return value + module = getattr(value, "__module__", type(value).__module__) + qualname = getattr(value, "__qualname__", type(value).__qualname__) + return f"{module}.{qualname}" + + +def _normalize_device(value: Any) -> str: + value = str(value).strip().lower() + if value.startswith("torch.device("): + value = value.removeprefix("torch.device(").removesuffix(")").strip("'\"") + if value.startswith("cuda:"): + return "cuda" + return {"gpu": "cuda", "hip": "rocm"}.get(value, value) + + +def _normalize_dtype(value: Any) -> str: + value = str(value).strip().lower().replace("torch.", "") + return { + "fp32": "float32", + "float": "float32", + "bf16": "bfloat16", + "fp16": "float16", + "half": "float16", + }.get(value, value) + + +def _normalized_values(values: Iterable[Any], normalize: Callable[[Any], str]) -> frozenset[str]: + return frozenset(value for item in values if (value := normalize(item).strip().lower())) + + +def _freeze(value: Any) -> Any: + if isinstance(value, Mapping): + return MappingProxyType({str(key): _freeze(item) for key, item in value.items()}) + if isinstance(value, (tuple, list)): + return tuple(_freeze(item) for item in value) + if isinstance(value, (set, frozenset)): + return frozenset(_freeze(item) for item in value) + return value + + +def _json_value(value: Any) -> Any: + if isinstance(value, _JsonRecord): + return value.to_dict() + if isinstance(value, Enum): + return value.value + if isinstance(value, Mapping): + return {str(key): _json_value(item) for key, item in sorted(value.items())} + if isinstance(value, (set, frozenset)): + return sorted((_json_value(item) for item in value), key=repr) + if isinstance(value, (tuple, list)): + return [_json_value(item) for item in value] + if callable(value): + return _reference(value) + return value + + +def _implementation_identity(value: Any) -> Mapping[str, Any]: + reference = _reference(value) + identity: dict[str, Any] = { + "reference": reference, + "kind": "class" if inspect.isclass(value) else "callable", + "callable": _callable_identity(value), + "module": _module_identity(getattr(value, "__module__", None)), + } + if inspect.isclass(value): + identity["members"] = { + name: _callable_identity(member) + for name, raw_member in sorted(vars(value).items()) + if (member := _descriptor_callable(raw_member)) is not None + } + return identity + + +def _descriptor_callable(value: Any) -> Optional[Callable[..., Any]]: + if isinstance(value, (classmethod, staticmethod)): + value = value.__func__ + elif isinstance(value, property): + return None + return value if callable(value) else None + + +def _callable_identity(value: Any) -> Mapping[str, Any]: + if inspect.ismethod(value): + value = value.__func__ + try: + unwrapped = inspect.unwrap(value) + except (TypeError, ValueError): + unwrapped = value + code = getattr(unwrapped, "__code__", None) + try: + source = inspect.getsource(unwrapped) + except (OSError, TypeError): + source = None + identity: dict[str, Any] = { + "reference": _reference(unwrapped), + "source_sha256": ( + hashlib.sha256(source.encode("utf-8")).hexdigest() if source is not None else None + ), + "code_sha256": _code_fingerprint(code) if isinstance(code, CodeType) else None, + } + if isinstance(code, CodeType): + identity["defaults"] = _code_value(getattr(unwrapped, "__defaults__", None)) + identity["keyword_defaults"] = _code_value(getattr(unwrapped, "__kwdefaults__", None)) + return identity + + +def _code_fingerprint(code: CodeType) -> str: + return _fingerprint( + { + "bytecode": code.co_code.hex(), + "constants": tuple(_code_value(value) for value in code.co_consts), + "names": code.co_names, + "variables": code.co_varnames, + "free_variables": code.co_freevars, + "cell_variables": code.co_cellvars, + "positional_arguments": code.co_argcount, + "positional_only_arguments": code.co_posonlyargcount, + "keyword_only_arguments": code.co_kwonlyargcount, + "flags": code.co_flags, + } + ) + + +def _code_value(value: Any) -> Any: + if isinstance(value, CodeType): + return {"nested_code_sha256": _code_fingerprint(value)} + if isinstance(value, bytes): + return {"bytes_sha256": hashlib.sha256(value).hexdigest()} + if isinstance(value, Mapping): + return { + str(key): _code_value(item) + for key, item in sorted(value.items(), key=lambda pair: repr(pair[0])) + } + if isinstance(value, (tuple, list)): + return [_code_value(item) for item in value] + if isinstance(value, (set, frozenset)): + return sorted((_code_value(item) for item in value), key=repr) + if value is None or isinstance(value, (bool, int, float, str)): + return value + return {"type": _reference(type(value)), "repr": repr(value)} + + +def _module_identity(module_name: Optional[str]) -> Optional[Mapping[str, Any]]: + if not module_name: + return None + try: + module = importlib.import_module(module_name) + except (ImportError, ModuleNotFoundError): + return {"name": module_name, "content_sha256": None} + module_file = getattr(module, "__file__", None) + if not module_file: + return {"name": module_name, "content_sha256": None} + path = Path(module_file) + try: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + except OSError: + return {"name": module_name, "content_sha256": None} + return { + "name": module_name, + "content_sha256": digest.hexdigest(), + } + + +def _fingerprint(value: Any) -> str: + encoded = json.dumps(_json_value(value), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() From 2bdf057345d0399ba47f4e9e4092c10477802f5a Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sun, 19 Jul 2026 08:30:48 +0800 Subject: [PATCH 02/33] feat(alignment): define cross-configuration contracts and plans --- rl_engine/alignment/cross_config/_json.py | 44 ++ .../alignment/cross_config/comparison.py | 321 ++++++++++ rl_engine/alignment/cross_config/config.py | 424 +++++++++++++ .../alignment/cross_config/execution_plan.py | 77 +++ rl_engine/alignment/cross_config/planner.py | 573 +++++++++++++++++ rl_engine/alignment/cross_config/schema.py | 583 ++++++++++++++++++ rl_engine/kernels/gtest/tolerance.py | 64 +- 7 files changed, 2085 insertions(+), 1 deletion(-) create mode 100644 rl_engine/alignment/cross_config/_json.py create mode 100644 rl_engine/alignment/cross_config/comparison.py create mode 100644 rl_engine/alignment/cross_config/config.py create mode 100644 rl_engine/alignment/cross_config/execution_plan.py create mode 100644 rl_engine/alignment/cross_config/planner.py create mode 100644 rl_engine/alignment/cross_config/schema.py diff --git a/rl_engine/alignment/cross_config/_json.py b/rl_engine/alignment/cross_config/_json.py new file mode 100644 index 00000000..fe8eb876 --- /dev/null +++ b/rl_engine/alignment/cross_config/_json.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fail-closed JSON decoding shared by configs and artifacts.""" + +from __future__ import annotations + +import json +import math +from typing import Any + + +def strict_json_loads(value: str) -> Any: + """Decode RFC JSON while rejecting duplicate keys and non-finite numbers.""" + + return json.loads( + value, + object_pairs_hook=_unique_object, + parse_constant=_reject_json_constant, + parse_float=_parse_finite_float, + ) + + +def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"duplicate JSON key {key!r}") + result[key] = value + return result + + +def _reject_json_constant(value: str) -> Any: + raise ValueError(f"non-finite JSON constant {value!r} is forbidden") + + +def _parse_finite_float(value: str) -> float: + result = float(value) + if not math.isfinite(result): + raise ValueError(f"non-finite JSON number {value!r} is forbidden") + return result + + +__all__ = ["strict_json_loads"] diff --git a/rl_engine/alignment/cross_config/comparison.py b/rl_engine/alignment/cross_config/comparison.py new file mode 100644 index 00000000..135dbd75 --- /dev/null +++ b/rl_engine/alignment/cross_config/comparison.py @@ -0,0 +1,321 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fixed-contract selected-token comparison for cross-configuration cases.""" + +from __future__ import annotations + +import math +from dataclasses import fields +from typing import Any + +import torch + +from rl_engine.alignment.cross_config.schema import ( + AlignmentResult, + AlignmentStatus, + ScoreArtifact, + ScoreSide, + SemanticIdentitySpec, + TokenComparisonArtifact, +) +from rl_engine.kernels.gtest.tolerance import ( + resolve_logprob_threshold, + tolerance_contract_fingerprint, +) + + +def semantic_identity_errors( + rollout: SemanticIdentitySpec, + training: SemanticIdentitySpec, +) -> tuple[str, ...]: + """Return every logical identity field that differs between the two sides.""" + + return tuple( + item.name + for item in fields(SemanticIdentitySpec) + if getattr(rollout, item.name) != getattr(training, item.name) + ) + + +def recompute_mismatch_mask( + rollout_logprobs: torch.Tensor, + training_logprobs: torch.Tensor, + active_mask: torch.Tensor, + fixed_threshold: float, +) -> torch.Tensor: + """Recompute the sole token mismatch signal from persisted tensors.""" + + if rollout_logprobs.shape != training_logprobs.shape: + raise ValueError("rollout and training logprobs must have identical shapes") + if active_mask.shape != rollout_logprobs.shape: + raise ValueError("active_mask shape must match selected logprobs") + if fixed_threshold < 0.0: + raise ValueError("fixed_threshold must be non-negative") + active = active_mask.to(device=rollout_logprobs.device, dtype=torch.bool) + training = training_logprobs.to(device=rollout_logprobs.device) + return active & (torch.abs(training - rollout_logprobs) > fixed_threshold) + + +class FixedThresholdComparator: + """Compare paired selected logprobs using only the current WS1 contract.""" + + def compare(self, rollout: ScoreArtifact, training: ScoreArtifact) -> AlignmentResult: + contract_fingerprint = tolerance_contract_fingerprint() + artifact_errors = _artifact_errors(rollout, training) + if artifact_errors: + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.INVALID_ARTIFACT, + comparable=False, + passed=False, + active_token_count=0, + mismatch_count=0, + contract_fingerprint=contract_fingerprint, + artifact_errors=artifact_errors, + ) + + identity_errors = list(semantic_identity_errors(rollout.identity, training.identity)) + identity_errors.extend(_artifact_identity_errors(rollout, training)) + if identity_errors: + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.INVALID_IDENTITY, + comparable=False, + passed=False, + active_token_count=0, + mismatch_count=0, + contract_fingerprint=contract_fingerprint, + identity_errors=tuple(dict.fromkeys(identity_errors)), + ) + + threshold, threshold_error = _resolve_fixed_threshold(rollout, training) + if threshold_error is not None: + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.INVALID_ARTIFACT, + comparable=False, + passed=False, + active_token_count=0, + mismatch_count=0, + contract_fingerprint=contract_fingerprint, + artifact_errors=(threshold_error,), + ) + assert threshold is not None + fixed_threshold = threshold + + rollout_logprobs = rollout.selected_logprobs.detach().cpu() + training_logprobs = training.selected_logprobs.detach().cpu() + active_mask = rollout.active_mask.detach().cpu().to(dtype=torch.bool) + active_token_count = int(active_mask.sum().item()) + if active_token_count: + active_rollout = rollout_logprobs[active_mask] + active_training = training_logprobs[active_mask] + if not bool(torch.isfinite(active_rollout).all().item()) or not bool( + torch.isfinite(active_training).all().item() + ): + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.INVALID_ARTIFACT, + comparable=False, + passed=False, + active_token_count=active_token_count, + mismatch_count=0, + contract_fingerprint=contract_fingerprint, + fixed_threshold=fixed_threshold, + artifact_errors=("active selected logprobs must be finite",), + ) + + # Inactive positions are outside the numerical contract. Canonicalize + # them before persistence so an ignored NaN/Inf cannot break strict JSON + # serialization or make resume artifacts non-reproducible. + rollout_logprobs = rollout_logprobs.masked_fill(~active_mask, 0.0) + training_logprobs = training_logprobs.masked_fill(~active_mask, 0.0) + absolute_diff = torch.abs(training_logprobs - rollout_logprobs) + mismatch_mask = recompute_mismatch_mask( + rollout_logprobs, + training_logprobs, + active_mask, + fixed_threshold, + ) + token_artifact = TokenComparisonArtifact( + rollout_logprobs=rollout_logprobs, + training_logprobs=training_logprobs, + active_mask=active_mask, + absolute_diff=absolute_diff, + mismatch_mask=mismatch_mask, + fixed_threshold=fixed_threshold, + ) + if active_token_count == 0: + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.ZERO_ACTIVE_TOKENS, + comparable=False, + passed=False, + active_token_count=0, + mismatch_count=0, + contract_fingerprint=contract_fingerprint, + fixed_threshold=fixed_threshold, + token_artifact=token_artifact, + ) + + mismatch_count = int(mismatch_mask.sum().item()) + passed = mismatch_count == 0 + return AlignmentResult( + case_id=rollout.case_id, + attempt_id=rollout.attempt_id, + status=AlignmentStatus.PASS if passed else AlignmentStatus.FAIL, + comparable=True, + passed=passed, + active_token_count=active_token_count, + mismatch_count=mismatch_count, + contract_fingerprint=contract_fingerprint, + fixed_threshold=fixed_threshold, + diagnostics=_diagnostics( + rollout_logprobs, + training_logprobs, + active_mask, + absolute_diff, + mismatch_count, + ), + token_artifact=token_artifact, + ) + + +def compare_score_artifacts( + rollout: ScoreArtifact, + training: ScoreArtifact, +) -> AlignmentResult: + """Convenience wrapper whose API deliberately exposes no threshold override.""" + + return FixedThresholdComparator().compare(rollout, training) + + +def _resolve_fixed_threshold( + rollout: ScoreArtifact, + training: ScoreArtifact, +) -> tuple[float | None, str | None]: + """Resolve one WS1 threshold, rejecting any mixed-dtype ambiguity.""" + + try: + rollout_threshold = resolve_logprob_threshold(rollout.scorer.dtype) + training_threshold = resolve_logprob_threshold(training.scorer.dtype) + except ValueError as exc: + return None, f"fixed WS1 threshold is unavailable: {exc}" + if rollout_threshold != training_threshold: + return ( + None, + "fixed WS1 threshold is ambiguous for scorer dtypes " + f"rollout={rollout.scorer.dtype!r}, training={training.scorer.dtype!r}", + ) + return rollout_threshold, None + + +def _artifact_errors(rollout: ScoreArtifact, training: ScoreArtifact) -> tuple[str, ...]: + errors: list[str] = [] + if rollout.side is not ScoreSide.ROLLOUT: + errors.append("first artifact side must be rollout") + if training.side is not ScoreSide.TRAINING: + errors.append("second artifact side must be training") + if rollout.case_id != training.case_id: + errors.append("case_id") + if rollout.attempt_id != training.attempt_id: + errors.append("attempt_id") + if rollout.selected_logprobs.shape != training.selected_logprobs.shape: + errors.append("selected_logprobs shape") + for label, artifact in (("rollout", rollout), ("training", training)): + expected_dtype = _score_dtype(artifact.scorer.dtype) + if not artifact.selected_logprobs.is_floating_point(): + errors.append(f"{label}.selected_logprobs must be floating point") + elif expected_dtype is None: + errors.append(f"{label}.scorer dtype is unsupported") + elif artifact.selected_logprobs.dtype != expected_dtype: + errors.append( + f"{label}.selected_logprobs dtype does not match scorer dtype " + f"({artifact.selected_logprobs.dtype} != {expected_dtype})" + ) + return tuple(errors) + + +def _score_dtype(value: str) -> torch.dtype | None: + normalized = str(value).strip().lower().removeprefix("torch.") + return { + "float32": torch.float32, + "fp32": torch.float32, + "float16": torch.float16, + "fp16": torch.float16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + "float64": torch.float64, + }.get(normalized) + + +def _artifact_identity_errors( + rollout: ScoreArtifact, + training: ScoreArtifact, +) -> tuple[str, ...]: + errors: list[str] = [] + rollout_identity_mask = _identity_mask(rollout.identity) + training_identity_mask = _identity_mask(training.identity) + rollout_mask = rollout.active_mask.detach().cpu().to(dtype=torch.bool) + training_mask = training.active_mask.detach().cpu().to(dtype=torch.bool) + if rollout_mask.shape != rollout_identity_mask.shape or not torch.equal( + rollout_mask, rollout_identity_mask + ): + errors.append("rollout.active_mask") + if training_mask.shape != training_identity_mask.shape or not torch.equal( + training_mask, training_identity_mask + ): + errors.append("training.active_mask") + if rollout_mask.shape != training_mask.shape or not torch.equal(rollout_mask, training_mask): + errors.append("active_mask") + return tuple(errors) + + +def _identity_mask(identity: SemanticIdentitySpec) -> torch.Tensor: + return torch.tensor(identity.active_mask, dtype=torch.bool) + + +def _diagnostics( + rollout_logprobs: torch.Tensor, + training_logprobs: torch.Tensor, + active_mask: torch.Tensor, + absolute_diff: torch.Tensor, + mismatch_count: int, +) -> dict[str, Any]: + active_diff = absolute_diff[active_mask].float() + delta = (training_logprobs[active_mask] - rollout_logprobs[active_mask]).float() + worst_active_index = int(torch.argmax(active_diff).item()) + active_coordinates = torch.nonzero(active_mask, as_tuple=False) + worst_coordinate = tuple(int(item) for item in active_coordinates[worst_active_index].tolist()) + approximate_kl = torch.exp(delta.double()) - delta.double() - 1.0 + approximate_kl_mean = _finite_float_or_none(approximate_kl.mean()) + active_count = int(active_diff.numel()) + return { + "mean_abs_diff": _finite_float_or_none(active_diff.mean()), + "p95_abs_diff": _finite_float_or_none(torch.quantile(active_diff, 0.95)), + "p99_abs_diff": _finite_float_or_none(torch.quantile(active_diff, 0.99)), + "max_abs_diff": _finite_float_or_none(active_diff.max()), + "mismatch_ratio": mismatch_count / active_count, + "approximate_kl_mean": approximate_kl_mean, + "approximate_kl_finite": approximate_kl_mean is not None, + "worst_token_index": worst_coordinate, + } + + +def _finite_float_or_none(value: torch.Tensor) -> float | None: + result = float(value.item()) + return result if math.isfinite(result) else None + + +__all__ = [ + "FixedThresholdComparator", + "compare_score_artifacts", + "recompute_mismatch_mask", + "semantic_identity_errors", +] diff --git a/rl_engine/alignment/cross_config/config.py b/rl_engine/alignment/cross_config/config.py new file mode 100644 index 00000000..e02c2198 --- /dev/null +++ b/rl_engine/alignment/cross_config/config.py @@ -0,0 +1,424 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Strict, dependency-free experiment configuration.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from dataclasses import fields as dataclass_fields +from pathlib import Path +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Mapping + +from rl_engine.alignment.cross_config._json import strict_json_loads +from rl_engine.alignment.cross_config.planner import normalize_backend_id +from rl_engine.alignment.cross_config.schema import ( + ExperimentCase, + ExperimentDefinition, + InterventionSpec, + PlanningStrategy, + SemanticIdentitySpec, +) + +if TYPE_CHECKING: + from rl_engine.alignment.cross_config.planner import ExperimentPlan + + +CONFIG_SCHEMA_VERSION = "cross_config.experiment_config.v1" +_FORBIDDEN_THRESHOLD_KEYS = frozenset({"threshold", "fixed_threshold", "tolerance", "atol", "rtol"}) +_TOP_LEVEL_KEYS = frozenset( + { + "schema_version", + "experiment_id", + "scenario_id", + "contract_source", + "contract_version", + "strategy", + "strict_fallback", + "identity", + "baseline", + "interventions", + "pairwise_paths", + "operators", + "scenario", + } +) +_IDENTITY_KEYS = frozenset( + item.name for item in dataclass_fields(SemanticIdentitySpec) if item.name != "schema_version" +) +_INTERVENTION_KEYS = frozenset({"path", "values"}) +_OPERATOR_NAMES = frozenset({"selected_logprob"}) +_OPERATOR_TARGETS = frozenset({"rollout", "training"}) +_OPERATOR_BINDING_KEYS = frozenset({"backend", "options"}) + + +@dataclass(frozen=True) +class OperatorSelection: + """Concrete selected-logprob implementation requested for each scorer side. + + ``logp.backend`` remains the concise both-sides shortcut. This explicit form + is needed only when rollout and training intentionally use different + implementations. + """ + + rollout_backend: str + training_backend: str + rollout_options: Mapping[str, Any] = field(default_factory=dict) + training_options: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + object.__setattr__(self, "rollout_backend", normalize_backend_id(self.rollout_backend)) + object.__setattr__( + self, + "training_backend", + normalize_backend_id(self.training_backend), + ) + object.__setattr__(self, "rollout_options", _freeze_mapping(self.rollout_options)) + object.__setattr__(self, "training_options", _freeze_mapping(self.training_options)) + + def backend_for(self, target: str) -> str: + if target == "rollout": + return self.rollout_backend + if target == "training": + return self.training_backend + raise ValueError("operator target must be 'rollout' or 'training'") + + def options_for(self, target: str) -> Mapping[str, Any]: + if target == "rollout": + return self.rollout_options + if target == "training": + return self.training_options + raise ValueError("operator target must be 'rollout' or 'training'") + + def to_dict(self) -> dict[str, Any]: + return { + "selected_logprob": { + "rollout": { + "backend": self.rollout_backend, + "options": _plain_value(self.rollout_options), + }, + "training": { + "backend": self.training_backend, + "options": _plain_value(self.training_options), + }, + } + } + + +@dataclass(frozen=True) +class ExperimentConfig: + """Loaded experiment plus optional target-specific operator selection.""" + + definition: ExperimentDefinition + source_path: Path + operators: OperatorSelection | None = None + schema_version: str = CONFIG_SCHEMA_VERSION + + def to_dict(self) -> dict[str, Any]: + """Return the normalized, portable experiment-config representation.""" + + payload = self.definition.to_dict() + payload["schema_version"] = self.schema_version + if self.operators is not None: + payload["operators"] = self.operators.to_dict() + return payload + + def plan(self) -> ExperimentPlan: + """Build the deterministic plan without importing a runtime backend.""" + + from rl_engine.alignment.cross_config.planner import Planner + + return Planner().plan(self.definition) + + def operators_for(self, case: ExperimentCase) -> OperatorSelection: + """Resolve the concise ``logp.backend`` shortcut for one planned case.""" + + backend = _case_logp_backend(case) + if self.operators is None: + return OperatorSelection(backend, backend) + if self.operators.rollout_backend != backend: + raise ValueError( + "operators.selected_logprob.rollout must match the planned " + f"logp.backend: {self.operators.rollout_backend!r} != {backend!r}" + ) + return self.operators + + +def load_config(path: str | Path) -> ExperimentConfig: + """Load one versioned JSON experiment with no threshold override surface.""" + + source = Path(path) + try: + raw = strict_json_loads(source.read_text(encoding="utf-8")) + except (OSError, ValueError) as exc: + raise ValueError(f"failed to load cross-configuration config {source}: {exc}") from exc + if not isinstance(raw, dict): + raise ValueError("cross-configuration config must contain a JSON object") + _reject_unknown_keys(raw, _TOP_LEVEL_KEYS, "config") + _reject_threshold_keys(raw) + if raw.get("schema_version") != CONFIG_SCHEMA_VERSION: + raise ValueError( + f"unsupported cross-configuration config schema {raw.get('schema_version')!r}; " + f"expected {CONFIG_SCHEMA_VERSION!r}" + ) + + identity_raw = _required_mapping(raw, "identity") + _reject_unknown_keys(identity_raw, _IDENTITY_KEYS, "identity") + scenario = _optional_mapping(raw, "scenario") + _reject_scenario_controls(scenario) + + interventions_raw = raw.get("interventions", []) + if not isinstance(interventions_raw, list): + raise ValueError("interventions must be a list") + interventions = tuple(_load_intervention(item) for item in interventions_raw) + + pairwise_raw = raw.get("pairwise_paths", []) + if not isinstance(pairwise_raw, list): + raise ValueError("pairwise_paths must be a list") + pairwise_paths = tuple(_load_pair(item) for item in pairwise_raw) + + strict_fallback = raw.get("strict_fallback", True) + if not isinstance(strict_fallback, bool): + raise ValueError("strict_fallback must be a JSON boolean") + + definition = ExperimentDefinition( + experiment_id=_required_string(raw, "experiment_id"), + scenario_id=_required_string(raw, "scenario_id"), + identity=SemanticIdentitySpec(**identity_raw), + baseline=_required_mapping(raw, "baseline"), + interventions=interventions, + scenario=scenario, + strategy=PlanningStrategy(raw.get("strategy", "one_at_a_time")), + strict_fallback=strict_fallback, + pairwise_paths=pairwise_paths, + contract_source=raw.get("contract_source", "ws1"), + contract_version=raw.get("contract_version", "current"), + ) + operators = _load_operators(raw.get("operators")) + if operators is not None: + if any(item.path == "logp.backend" for item in interventions): + raise ValueError( + "explicit operators cannot be combined with logp.backend interventions; " + "use the shortcut or one fixed target mapping" + ) + baseline_backend = _definition_logp_backend(definition) + if operators.rollout_backend != baseline_backend: + raise ValueError( + "operators.selected_logprob.rollout must match baseline logp.backend: " + f"{operators.rollout_backend!r} != {baseline_backend!r}" + ) + + return ExperimentConfig( + definition=definition, + operators=operators, + source_path=source, + ) + + +def bind_operator_selection( + case: ExperimentCase, + selection: OperatorSelection, +) -> ExperimentCase: + """Bind target-specific operators into the execution identity. + + Planning remains semantic-operator agnostic; the immutable binding extends + the case and resume key before any runtime is created. + """ + + requested_backend = _case_logp_backend(case) + if selection.rollout_backend != requested_backend: + raise ValueError( + "rollout operator must match the planned logp.backend: " + f"{selection.rollout_backend!r} != {requested_backend!r}" + ) + binding = selection.to_dict() + payload = { + "base_case_id": case.case_id, + "base_scenario_fingerprint": case.scenario_fingerprint, + "operators": binding, + } + serialized = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8") + operator_fingerprint = hashlib.sha256(serialized).hexdigest() + case_hash = hashlib.sha256( + json.dumps( + {"base_case_id": case.case_id, "operator_fingerprint": operator_fingerprint}, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).hexdigest()[:24] + scenario_fingerprint = hashlib.sha256( + f"{case.scenario_fingerprint}:{operator_fingerprint}".encode("utf-8") + ).hexdigest() + return ExperimentCase( + case_id=f"cross-config-{case_hash}", + experiment_id=case.experiment_id, + scenario_id=case.scenario_id, + identity=case.identity, + requested=case.requested, + execution_binding={"operators": binding}, + changed_paths=case.changed_paths, + contract_fingerprint=case.contract_fingerprint, + scenario_fingerprint=scenario_fingerprint, + ) + + +def _load_intervention(value: Any) -> InterventionSpec: + if not isinstance(value, Mapping): + raise ValueError("each intervention must be an object") + _reject_unknown_keys(value, _INTERVENTION_KEYS, "intervention") + values = value.get("values") + if not isinstance(values, list): + raise ValueError("intervention values must be a list") + return InterventionSpec(path=_required_string(value, "path"), values=tuple(values)) + + +def _load_pair(value: Any) -> tuple[str, str]: + if ( + not isinstance(value, list) + or len(value) != 2 + or not all(isinstance(item, str) for item in value) + ): + raise ValueError("each pairwise_paths entry must contain exactly two string paths") + return value[0], value[1] + + +def _load_operators(value: Any) -> OperatorSelection | None: + if value is None: + return None + if not isinstance(value, Mapping): + raise ValueError("operators must be an object") + _reject_unknown_keys(value, _OPERATOR_NAMES, "operators") + selected = value.get("selected_logprob") + if not isinstance(selected, Mapping): + raise ValueError("operators.selected_logprob must be an object") + _reject_unknown_keys(selected, _OPERATOR_TARGETS, "operators.selected_logprob") + rollout_backend, rollout_options = _load_operator_binding(selected, "rollout") + training_backend, training_options = _load_operator_binding(selected, "training") + return OperatorSelection( + rollout_backend=rollout_backend, + training_backend=training_backend, + rollout_options=rollout_options, + training_options=training_options, + ) + + +def _load_operator_binding( + value: Mapping[str, Any], + target: str, +) -> tuple[str, Mapping[str, Any]]: + binding = value.get(target) + if isinstance(binding, str): + if not binding.strip(): + raise ValueError(f"operators.selected_logprob.{target} must not be empty") + return binding, {} + if not isinstance(binding, Mapping): + raise ValueError(f"operators.selected_logprob.{target} must be a backend string or object") + _reject_unknown_keys(binding, _OPERATOR_BINDING_KEYS, f"{target} operator binding") + return _required_string(binding, "backend"), _optional_mapping(binding, "options") + + +def _case_logp_backend(case: ExperimentCase) -> str: + logp = case.requested.get("logp") + backend = logp.get("backend") if isinstance(logp, Mapping) else None + if not isinstance(backend, str) or not backend: + raise ValueError("planned cases must contain a non-empty string logp.backend") + return normalize_backend_id(backend) + + +def _definition_logp_backend(definition: ExperimentDefinition) -> str: + logp = definition.baseline.get("logp") + backend = logp.get("backend") if isinstance(logp, Mapping) else None + if not isinstance(backend, str) or not backend: + raise ValueError("baseline must contain a non-empty string logp.backend") + return normalize_backend_id(backend) + + +def _reject_scenario_controls(scenario: Mapping[str, Any]) -> None: + behavior_keys = sorted( + set(scenario).intersection( + {"execution", "plan_only", "operator_cases", "expected_status", "allow_smoke_operators"} + ) + ) + if behavior_keys: + raise ValueError( + "scenario is metadata only; move execution and operator policy to the CLI/config: " + f"{behavior_keys}" + ) + + +def _reject_threshold_keys(value: Any, prefix: str = "") -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + normalized = str(key).strip().lower() + path = f"{prefix}.{key}" if prefix else str(key) + if normalized in _FORBIDDEN_THRESHOLD_KEYS: + raise ValueError( + f"{path} is forbidden: the fixed numerical-contract threshold is imported" + ) + _reject_threshold_keys(child, path) + elif isinstance(value, list): + for index, child in enumerate(value): + _reject_threshold_keys(child, f"{prefix}[{index}]") + + +def _reject_unknown_keys( + value: Mapping[str, Any], + allowed: frozenset[str], + label: str, +) -> None: + unknown = sorted(set(value).difference(allowed)) + if unknown: + raise ValueError(f"unknown {label} keys: {unknown}") + + +def _required_mapping(value: Mapping[str, Any], key: str) -> dict[str, Any]: + child = value.get(key) + if not isinstance(child, Mapping): + raise ValueError(f"{key} must be an object") + return dict(child) + + +def _optional_mapping(value: Mapping[str, Any], key: str) -> dict[str, Any]: + child = value.get(key, {}) + if not isinstance(child, Mapping): + raise ValueError(f"{key} must be an object") + return dict(child) + + +def _required_string(value: Mapping[str, Any], key: str) -> str: + child = value.get(key) + if not isinstance(child, str) or not child.strip(): + raise ValueError(f"{key} must be a non-empty string") + return child.strip() + + +def _freeze_mapping(value: Mapping[str, Any]) -> Mapping[str, Any]: + return MappingProxyType({str(key): _freeze_value(child) for key, child in value.items()}) + + +def _freeze_value(value: Any) -> Any: + if isinstance(value, Mapping): + return _freeze_mapping(value) + if isinstance(value, list): + return tuple(_freeze_value(child) for child in value) + return value + + +def _plain_value(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _plain_value(child) for key, child in value.items()} + if isinstance(value, (tuple, list)): + return [_plain_value(child) for child in value] + return value + + +__all__ = [ + "CONFIG_SCHEMA_VERSION", + "ExperimentConfig", + "OperatorSelection", + "bind_operator_selection", + "load_config", +] diff --git a/rl_engine/alignment/cross_config/execution_plan.py b/rl_engine/alignment/cross_config/execution_plan.py new file mode 100644 index 00000000..d3be05b2 --- /dev/null +++ b/rl_engine/alignment/cross_config/execution_plan.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Canonical, runtime-independent execution-plan construction.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Mapping + +from rl_engine.alignment.cross_config.config import ( + ExperimentConfig, + OperatorSelection, + bind_operator_selection, +) +from rl_engine.alignment.cross_config.planner import PlanningIssue +from rl_engine.alignment.cross_config.schema import ExperimentCase + + +@dataclass(frozen=True) +class ExecutionPlanEntry: + """One operator-bound case and its resolved operator selection.""" + + case: ExperimentCase + operators: OperatorSelection + schema_version: str = "cross_config.execution_plan_entry.v1" + + def to_dict(self) -> dict[str, Any]: + """Return the canonical append-only plan row.""" + + return { + "schema_version": self.schema_version, + "case": self.case.to_dict(), + "operators": self.operators.to_dict(), + } + + +@dataclass(frozen=True) +class ExecutionPlan: + """Canonical metadata shared by planning and every runtime adapter.""" + + experiment: Mapping[str, Any] + entries: tuple[ExecutionPlanEntry, ...] + issues: tuple[PlanningIssue, ...] = () + schema_version: str = "cross_config.execution_plan.v1" + + def rows(self) -> tuple[dict[str, Any], ...]: + """Serialize all plan entries in deterministic execution order.""" + + return tuple(entry.to_dict() for entry in self.entries) + + +def build_execution_plan(config: ExperimentConfig) -> ExecutionPlan: + """Plan, resolve operators, and bind them into immutable case identities.""" + + planned = config.plan() + entries: list[ExecutionPlanEntry] = [] + for case in planned.cases: + operators = config.operators_for(case) + entries.append( + ExecutionPlanEntry( + case=bind_operator_selection(case, operators), + operators=operators, + ) + ) + return ExecutionPlan( + experiment=config.to_dict(), + entries=tuple(entries), + issues=planned.issues, + ) + + +__all__ = [ + "ExecutionPlan", + "ExecutionPlanEntry", + "build_execution_plan", +] diff --git a/rl_engine/alignment/cross_config/planner.py b/rl_engine/alignment/cross_config/planner.py new file mode 100644 index 00000000..476da5de --- /dev/null +++ b/rl_engine/alignment/cross_config/planner.py @@ -0,0 +1,573 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed baseline, one-at-a-time, and explicitly bounded pairwise planning.""" + +from __future__ import annotations + +import hashlib +import itertools +import json +from dataclasses import dataclass +from typing import Any, Callable, Mapping, Optional, Sequence + +from rl_engine.alignment.cross_config.schema import ( + ExperimentCase, + ExperimentDefinition, + IsolationScope, + KnobDescriptor, + PlanningStrategy, +) +from rl_engine.kernels.gtest.tolerance import tolerance_contract_fingerprint + +Normalizer = Callable[[Any], Any] +Constraint = Callable[[str, Any, Mapping[str, Any]], Optional["PlanningIssue"]] +MAX_PLAN_CASES = 256 + + +@dataclass(frozen=True) +class PlanningIssue: + """Structured planning rejection that callers can persist or display.""" + + code: str + reason: str + path: Optional[str] = None + value: Any = None + + def to_dict(self) -> dict[str, Any]: + return { + "code": self.code, + "reason": self.reason, + "path": self.path, + "value": self.value, + } + + +class PlanningError(ValueError): + """Raised for an invalid experiment definition with structured issues.""" + + def __init__(self, issues: Sequence[PlanningIssue]): + self.issues = tuple(issues) + message = "; ".join( + f"{issue.code}{f'[{issue.path}]' if issue.path else ''}: {issue.reason}" + for issue in self.issues + ) + super().__init__(message) + + +@dataclass(frozen=True) +class ExperimentPlan: + """A deterministic plan plus non-fatal capability findings.""" + + definition: ExperimentDefinition + cases: tuple[ExperimentCase, ...] + issues: tuple[PlanningIssue, ...] = () + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": "cross_config.experiment_plan.v1", + "experiment_id": self.definition.experiment_id, + "cases": [case.to_dict() for case in self.cases], + "issues": [issue.to_dict() for issue in self.issues], + } + + +def _positive_int(value: Any) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("must be a positive integer") + return value + + +def _strict_bool(value: Any) -> bool: + if not isinstance(value, bool): + raise ValueError("must be a JSON boolean") + return value + + +def _normalize_dtype(value: Any) -> str: + if not isinstance(value, str): + raise ValueError("must be a dtype string") + normalized = value.strip().lower().replace("torch.", "") + aliases = { + "bf16": "bfloat16", + "bfloat16": "bfloat16", + "fp16": "float16", + "half": "float16", + "float16": "float16", + "fp32": "float32", + "float": "float32", + "float32": "float32", + } + try: + return aliases[normalized] + except KeyError as exc: + raise ValueError(f"unsupported dtype {value!r}") from exc + + +def _normalize_choice(*choices: str) -> Normalizer: + allowed = frozenset(choices) + + def normalize(value: Any) -> str: + if not isinstance(value, str): + raise ValueError("must be a string") + normalized = value.strip().lower().replace("-", "_") + if normalized not in allowed: + raise ValueError(f"must be one of {sorted(allowed)}") + return normalized + + return normalize + + +def normalize_backend_id(value: Any) -> str: + """Normalize the public selected-logprob backend shortcut.""" + + if not isinstance(value, str) or not value.strip(): + raise ValueError("must be a non-empty backend ID") + normalized = value.strip().lower().replace("-", "_") + aliases = { + "auto": "native", + "default": "native", + "pytorch": "rlkernel.reference_logp", + "reference": "rlkernel.reference_logp", + } + return aliases.get(normalized, normalized) + + +V1_KNOB_DESCRIPTORS: tuple[KnobDescriptor, ...] = ( + KnobDescriptor("batch.size", IsolationScope.REQUEST, ("rollout", "training")), + KnobDescriptor("rollout.tensor_parallel_size", IsolationScope.PROCESS, ("rollout",)), + KnobDescriptor("rollout.context_parallel_size", IsolationScope.PROCESS, ("rollout",)), + KnobDescriptor("rollout.dtype", IsolationScope.ENGINE_CONSTRUCTION, ("rollout",)), + KnobDescriptor( + "rollout.enable_prefix_caching", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + KnobDescriptor("rollout.enforce_eager", IsolationScope.ENGINE_CONSTRUCTION, ("rollout",)), + KnobDescriptor( + "training.attention_backend", + IsolationScope.ENGINE_CONSTRUCTION, + ("training",), + allowed_values=("flash_attention_2", "sdpa", "eager", "model_default"), + ), + KnobDescriptor("training.compute_dtype", IsolationScope.ENGINE_CONSTRUCTION, ("training",)), + KnobDescriptor( + "logp.backend", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + derived=True, + ), + KnobDescriptor( + "training.sharding", + IsolationScope.PROCESS, + ("training",), + allowed_values=("unsharded", "fsdp"), + ), +) + +V1_KNOBS: Mapping[str, KnobDescriptor] = { + descriptor.path: descriptor for descriptor in V1_KNOB_DESCRIPTORS +} + +_NORMALIZERS: Mapping[str, Normalizer] = { + "batch.size": _positive_int, + "rollout.tensor_parallel_size": _positive_int, + "rollout.context_parallel_size": _positive_int, + "rollout.dtype": _normalize_dtype, + "rollout.enable_prefix_caching": _strict_bool, + "rollout.enforce_eager": _strict_bool, + "training.attention_backend": _normalize_choice( + "flash_attention_2", "sdpa", "eager", "model_default" + ), + "training.compute_dtype": _normalize_dtype, + "logp.backend": normalize_backend_id, + "training.sharding": _normalize_choice("unsharded", "fsdp"), +} + + +class Planner: + """Generate a bounded plan without importing runtime- or operator-specific branches.""" + + def __init__( + self, + *, + knobs: Mapping[str, KnobDescriptor] = V1_KNOBS, + normalizers: Mapping[str, Normalizer] = _NORMALIZERS, + constraints: Sequence[Constraint] = (), + ): + self.knobs = dict(knobs) + self.normalizers = dict(normalizers) + self.constraints = tuple(constraints) + + def plan(self, definition: ExperimentDefinition) -> ExperimentPlan: + issues = self._validate_definition(definition) + if issues: + raise PlanningError(issues) + + baseline = self.normalize_requested(definition.baseline) + requested_cases: list[tuple[dict[str, Any], tuple[str, ...]]] = [(baseline, ())] + intervention_values: dict[str, tuple[Any, ...]] = {} + + def append_requested(requested: dict[str, Any], changed_paths: tuple[str, ...]) -> None: + if len(requested_cases) >= MAX_PLAN_CASES: + raise PlanningError( + ( + PlanningIssue( + code="PLAN_TOO_LARGE", + reason=f"a plan may contain at most {MAX_PLAN_CASES} cases", + value=MAX_PLAN_CASES, + ), + ) + ) + requested_cases.append((requested, changed_paths)) + + for intervention in definition.interventions: + if len(intervention.values) > MAX_PLAN_CASES: + raise PlanningError( + ( + PlanningIssue( + code="PLAN_TOO_LARGE", + reason=(f"an intervention may contain at most {MAX_PLAN_CASES} values"), + path=intervention.path, + value=len(intervention.values), + ), + ) + ) + normalized_values = tuple( + self._normalize_value(intervention.path, value) for value in intervention.values + ) + intervention_values[intervention.path] = normalized_values + baseline_value = _get_path(baseline, intervention.path) + for value in normalized_values: + if value == baseline_value: + continue + requested = _deep_copy_mapping(baseline) + _set_path(requested, intervention.path, value) + append_requested(requested, (intervention.path,)) + + if definition.strategy == PlanningStrategy.PAIRWISE: + for first_path, second_path in definition.pairwise_paths: + first_baseline = _get_path(baseline, first_path) + second_baseline = _get_path(baseline, second_path) + for first_value, second_value in itertools.product( + intervention_values[first_path], intervention_values[second_path] + ): + if first_value == first_baseline or second_value == second_baseline: + continue + requested = _deep_copy_mapping(baseline) + _set_path(requested, first_path, first_value) + _set_path(requested, second_path, second_value) + append_requested(requested, tuple(sorted((first_path, second_path)))) + + contract_fingerprint = tolerance_contract_fingerprint() + scenario_fingerprint = _fingerprint( + {"scenario_id": definition.scenario_id, "scenario": definition.scenario} + ) + cases: list[ExperimentCase] = [] + seen_ids: set[str] = set() + capability_issues: list[PlanningIssue] = [] + for requested, changed_paths in requested_cases: + case_issues = self._apply_constraints(requested, changed_paths) + capability_issues.extend(case_issues) + case_id = self._case_id( + definition, + requested, + contract_fingerprint=contract_fingerprint, + scenario_fingerprint=scenario_fingerprint, + ) + if case_id in seen_ids: + continue + seen_ids.add(case_id) + cases.append( + ExperimentCase( + case_id=case_id, + experiment_id=definition.experiment_id, + scenario_id=definition.scenario_id, + identity=definition.identity, + requested=requested, + changed_paths=changed_paths, + contract_fingerprint=contract_fingerprint, + scenario_fingerprint=scenario_fingerprint, + ) + ) + + return ExperimentPlan( + definition=definition, + cases=tuple(cases), + issues=tuple(capability_issues), + ) + + def normalize_requested(self, requested: Mapping[str, Any]) -> dict[str, Any]: + flattened = _flatten(requested) + issues: list[PlanningIssue] = [] + normalized: dict[str, Any] = {} + for path, value in flattened.items(): + if path not in self.knobs: + code = "DERIVED_KNOB" if path == "logp.tp_layout" else "UNSUPPORTED_PATH" + issues.append( + PlanningIssue( + code=code, + path=path, + value=value, + reason="path is not a user-settable V1 knob", + ) + ) + continue + try: + normalized[path] = self._normalize_value(path, value) + except (TypeError, ValueError) as exc: + issues.append( + PlanningIssue( + code="UNSUPPORTED_VALUE", + path=path, + value=value, + reason=str(exc), + ) + ) + if issues: + raise PlanningError(issues) + result: dict[str, Any] = {} + for path, value in normalized.items(): + _set_path(result, path, value) + return result + + def isolation_for(self, changed_paths: Sequence[str]) -> IsolationScope: + if not changed_paths: + return IsolationScope.REQUEST + order = { + IsolationScope.REQUEST: 0, + IsolationScope.ENGINE_CONSTRUCTION: 1, + IsolationScope.DISTRIBUTED_CONTEXT: 2, + IsolationScope.PROCESS: 3, + } + return max((self.knobs[path].lifecycle for path in changed_paths), key=order.__getitem__) + + def _validate_definition(self, definition: ExperimentDefinition) -> list[PlanningIssue]: + issues: list[PlanningIssue] = [] + try: + baseline = self.normalize_requested(definition.baseline) + except PlanningError as exc: + return list(exc.issues) + baseline_paths = set(_flatten(baseline)) + for path in sorted(set(self.knobs).difference(baseline_paths)): + issues.append( + PlanningIssue( + code="MISSING_BASELINE_VALUE", + path=path, + reason="strict baselines must declare every allowlisted knob", + ) + ) + declared_paths: set[str] = set() + for intervention in definition.interventions: + path = intervention.path + if path not in self.knobs: + issues.append( + PlanningIssue( + code="UNSUPPORTED_PATH", + path=path, + reason="intervention path is not in the V1 allowlist", + ) + ) + continue + if path in declared_paths: + issues.append( + PlanningIssue( + code="DUPLICATE_INTERVENTION", + path=path, + reason="each intervention path must be declared once", + ) + ) + declared_paths.add(path) + if not intervention.values: + issues.append( + PlanningIssue( + code="EMPTY_INTERVENTION", + path=path, + reason="intervention values cannot be empty", + ) + ) + try: + _get_path(baseline, path) + except KeyError: + issues.append( + PlanningIssue( + code="MISSING_BASELINE_VALUE", + path=path, + reason="every intervention path must exist in baseline", + ) + ) + for value in intervention.values: + try: + self._normalize_value(path, value) + except (TypeError, ValueError) as exc: + issues.append( + PlanningIssue( + code="UNSUPPORTED_VALUE", + path=path, + value=value, + reason=str(exc), + ) + ) + + if definition.strategy == PlanningStrategy.ONE_AT_A_TIME and definition.pairwise_paths: + issues.append( + PlanningIssue( + code="PAIRWISE_NOT_ENABLED", + reason="pairwise_paths require strategy='pairwise'", + ) + ) + if definition.strategy == PlanningStrategy.PAIRWISE and not definition.pairwise_paths: + issues.append( + PlanningIssue( + code="PAIRWISE_PATHS_REQUIRED", + reason="pairwise strategy requires at least one explicit path pair", + ) + ) + seen_pairs: set[tuple[str, str]] = set() + for pair in definition.pairwise_paths: + if len(pair) != 2: + issues.append( + PlanningIssue( + code="INVALID_PAIR", + reason="each pairwise entry must contain exactly two paths", + value=pair, + ) + ) + continue + first, second = pair + canonical_pair = (first, second) if first < second else (second, first) + if first == second: + issues.append( + PlanningIssue( + code="INVALID_PAIR", + reason="pairwise paths must be distinct", + value=pair, + ) + ) + elif first not in declared_paths or second not in declared_paths: + issues.append( + PlanningIssue( + code="UNDECLARED_PAIR_PATH", + reason="pairwise paths must both have declared interventions", + value=pair, + ) + ) + elif canonical_pair in seen_pairs: + issues.append( + PlanningIssue( + code="DUPLICATE_PAIR", + reason="pairwise path pair is duplicated", + value=pair, + ) + ) + seen_pairs.add(canonical_pair) + return issues + + def _normalize_value(self, path: str, value: Any) -> Any: + try: + normalizer = self.normalizers[path] + except KeyError as exc: + raise ValueError(f"no normalizer registered for {path}") from exc + return normalizer(value) + + def _apply_constraints( + self, requested: Mapping[str, Any], changed_paths: Sequence[str] + ) -> list[PlanningIssue]: + issues: list[PlanningIssue] = [] + paths = changed_paths or tuple(_flatten(requested)) + for path in paths: + value = _get_path(requested, path) + for constraint in self.constraints: + issue = constraint(path, value, requested) + if issue is not None: + issues.append(issue) + return issues + + @staticmethod + def _case_id( + definition: ExperimentDefinition, + requested: Mapping[str, Any], + *, + contract_fingerprint: str, + scenario_fingerprint: str, + ) -> str: + payload = { + "requested": requested, + "identity": definition.identity.to_dict(), + "contract": { + "source": definition.contract_source, + "version": definition.contract_version, + "fingerprint": contract_fingerprint, + }, + "scenario_id": definition.scenario_id, + "scenario_fingerprint": scenario_fingerprint, + } + return f"cross-config-{_fingerprint(payload)[:20]}" + + +def _flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + flattened: dict[str, Any] = {} + for key, child in value.items(): + if not isinstance(key, str) or not key: + raise PlanningError( + [PlanningIssue(code="INVALID_PATH", reason="configuration keys must be strings")] + ) + path = f"{prefix}.{key}" if prefix else key + if isinstance(child, Mapping): + flattened.update(_flatten(child, path)) + else: + flattened[path] = child + return flattened + + +def _get_path(value: Mapping[str, Any], path: str) -> Any: + current: Any = value + for part in path.split("."): + if not isinstance(current, Mapping) or part not in current: + raise KeyError(path) + current = current[part] + return current + + +def _set_path(value: dict[str, Any], path: str, child: Any) -> None: + current = value + parts = path.split(".") + for part in parts[:-1]: + existing = current.setdefault(part, {}) + if not isinstance(existing, dict): + raise ValueError(f"configuration path collision at {path}") + current = existing + current[parts[-1]] = child + + +def _deep_copy_mapping(value: Mapping[str, Any]) -> dict[str, Any]: + return json.loads(json.dumps(value)) + + +def _fingerprint(value: Mapping[str, Any]) -> str: + payload = json.dumps( + _json_plain(value), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _json_plain(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _json_plain(child) for key, child in value.items()} + if isinstance(value, (tuple, list)): + return [_json_plain(child) for child in value] + return value + + +__all__ = [ + "ExperimentPlan", + "Planner", + "PlanningError", + "PlanningIssue", + "V1_KNOBS", + "V1_KNOB_DESCRIPTORS", + "normalize_backend_id", +] diff --git a/rl_engine/alignment/cross_config/schema.py b/rl_engine/alignment/cross_config/schema.py new file mode 100644 index 00000000..27f668db --- /dev/null +++ b/rl_engine/alignment/cross_config/schema.py @@ -0,0 +1,583 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Stable, versioned domain schema for cross-configuration alignment.""" + +from __future__ import annotations + +import json +import math +from dataclasses import dataclass, field, fields +from enum import Enum +from pathlib import Path +from types import MappingProxyType +from typing import Any, Mapping, Optional, Sequence + +import torch + + +class ScoreSide(str, Enum): + ROLLOUT = "rollout" + TRAINING = "training" + + +class IsolationScope(str, Enum): + REQUEST = "request" + ENGINE_CONSTRUCTION = "engine_construction" + DISTRIBUTED_CONTEXT = "distributed_context" + PROCESS = "process" + + +class PlanningStrategy(str, Enum): + ONE_AT_A_TIME = "one_at_a_time" + PAIRWISE = "pairwise" + + +class MaterializationStatus(str, Enum): + UNSUPPORTED = "unsupported" + APPLIED = "applied" + FALLBACK = "fallback" + UNOBSERVABLE = "unobservable" + ERROR = "error" + + +class AlignmentStatus(str, Enum): + PASS = "pass" + FAIL = "fail" + INVALID_IDENTITY = "invalid_identity" + INVALID_ARTIFACT = "invalid_artifact" + ZERO_ACTIVE_TOKENS = "zero_active_tokens" + + +class SerializableModel: + """Mixin providing a stable JSON-compatible representation.""" + + def to_dict(self) -> dict[str, Any]: + return { + item.name: _serialize_value(getattr(self, item.name)) + for item in fields(self) # type: ignore[arg-type] + } + + def to_json(self, *, indent: Optional[int] = None) -> str: + return json.dumps(self.to_dict(), indent=indent, sort_keys=True) + + +def _serialize_value(value: Any) -> Any: + if isinstance(value, Enum): + return value.value + if isinstance(value, SerializableModel): + return value.to_dict() + if isinstance(value, torch.Tensor): + snapshot = value.detach().cpu() + return { + "dtype": str(snapshot.dtype).replace("torch.", ""), + "shape": list(snapshot.shape), + "values": snapshot.tolist(), + } + if isinstance(value, Mapping): + return {str(key): _serialize_value(item) for key, item in value.items()} + if isinstance(value, (tuple, list)): + return [_serialize_value(item) for item in value] + if isinstance(value, (set, frozenset)): + return [_serialize_value(item) for item in sorted(value, key=repr)] + if isinstance(value, Path): + return str(value) + if isinstance(value, torch.dtype): + return str(value).replace("torch.", "") + return value + + +def _freeze_value(value: Any) -> Any: + if isinstance(value, Mapping): + return MappingProxyType({str(key): _freeze_value(item) for key, item in value.items()}) + if isinstance(value, (tuple, list)): + return tuple(_freeze_value(item) for item in value) + if isinstance(value, (set, frozenset)): + return tuple(sorted((_freeze_value(item) for item in value), key=repr)) + return value + + +def _freeze_mapping(value: Mapping[str, Any]) -> Mapping[str, Any]: + return _freeze_value(value) + + +def _coerce_enum(value: Any, enum_type: type[Enum]) -> Enum: + if isinstance(value, enum_type): + return value + return enum_type(value) + + +def _int_matrix(value: Sequence[Sequence[int]]) -> tuple[tuple[int, ...], ...]: + rows: list[tuple[int, ...]] = [] + for row in value: + normalized: list[int] = [] + for item in row: + if isinstance(item, bool) or not isinstance(item, int): + raise ValueError("integer identity matrices accept JSON integers only") + normalized.append(item) + rows.append(tuple(normalized)) + return tuple(rows) + + +def _bool_matrix(value: Sequence[Sequence[bool]]) -> tuple[tuple[bool, ...], ...]: + rows: list[tuple[bool, ...]] = [] + for row in value: + normalized: list[bool] = [] + for item in row: + if not isinstance(item, bool): + raise ValueError("boolean identity matrices accept JSON booleans only") + normalized.append(item) + rows.append(tuple(normalized)) + return tuple(rows) + + +def _validate_rectangular(name: str, value: tuple[tuple[Any, ...], ...]) -> None: + if not value: + return + width = len(value[0]) + if any(len(row) != width for row in value): + raise ValueError(f"{name} must be rectangular") + + +def _validate_same_matrix_shape( + left_name: str, + left: tuple[tuple[Any, ...], ...], + right_name: str, + right: tuple[tuple[Any, ...], ...], +) -> None: + if left and right and (len(left), len(left[0])) != (len(right), len(right[0])): + raise ValueError(f"{left_name} shape must match {right_name} shape") + + +def _snapshot_tensor(value: torch.Tensor, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + if not isinstance(value, torch.Tensor): + raise TypeError(f"expected torch.Tensor, got {type(value)!r}") + snapshot = value.detach().clone() + return snapshot.to(dtype=dtype) if dtype is not None else snapshot + + +@dataclass(frozen=True) +class SemanticIdentitySpec(SerializableModel): + """Logical inputs that must match before numerical comparison is meaningful.""" + + checkpoint_id: str + model_version: str + tokenizer_policy: str + token_ids: tuple[tuple[int, ...], ...] + selected_token_ids: tuple[tuple[int, ...], ...] + active_mask: tuple[tuple[bool, ...], ...] + pre_update_state: str + tokenizer_id: str = "" + attention_mask: tuple[tuple[bool, ...], ...] = () + position_ids: tuple[tuple[int, ...], ...] = () + cache_metadata: Mapping[str, Any] = field(default_factory=dict) + packing_metadata: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.semantic_identity.v1" + + def __post_init__(self) -> None: + if self.schema_version != "cross_config.semantic_identity.v1": + raise ValueError("unsupported SemanticIdentitySpec schema_version") + if not self.checkpoint_id: + raise ValueError("checkpoint_id must not be empty") + if not self.model_version: + raise ValueError("model_version must not be empty") + if not self.tokenizer_policy: + raise ValueError("tokenizer_policy must not be empty") + if not self.pre_update_state: + raise ValueError("pre_update_state must not be empty") + + object.__setattr__(self, "token_ids", _int_matrix(self.token_ids)) + object.__setattr__(self, "selected_token_ids", _int_matrix(self.selected_token_ids)) + object.__setattr__(self, "active_mask", _bool_matrix(self.active_mask)) + object.__setattr__(self, "attention_mask", _bool_matrix(self.attention_mask)) + object.__setattr__(self, "position_ids", _int_matrix(self.position_ids)) + object.__setattr__(self, "cache_metadata", _freeze_mapping(self.cache_metadata)) + object.__setattr__(self, "packing_metadata", _freeze_mapping(self.packing_metadata)) + + if not self.token_ids or not self.token_ids[0]: + raise ValueError("token_ids must contain at least one token") + if not self.selected_token_ids: + raise ValueError("selected_token_ids must not be empty") + if not self.active_mask: + raise ValueError("active_mask must not be empty") + if not self.attention_mask: + raise ValueError("attention_mask must not be empty") + for name in ( + "token_ids", + "selected_token_ids", + "active_mask", + "attention_mask", + "position_ids", + ): + _validate_rectangular(name, getattr(self, name)) + _validate_same_matrix_shape( + "token_ids", + self.token_ids, + "selected_token_ids", + self.selected_token_ids, + ) + _validate_same_matrix_shape( + "selected_token_ids", self.selected_token_ids, "active_mask", self.active_mask + ) + _validate_same_matrix_shape( + "token_ids", self.token_ids, "attention_mask", self.attention_mask + ) + _validate_same_matrix_shape("token_ids", self.token_ids, "position_ids", self.position_ids) + + +@dataclass(frozen=True) +class ScorerSpec(SerializableModel): + side: ScoreSide + backend_id: str + dtype: str + device: str = "cpu" + world_size: int = 1 + topology: Mapping[str, Any] = field(default_factory=dict) + construction_options: Mapping[str, Any] = field(default_factory=dict) + operator_overrides: Mapping[str, str] = field(default_factory=dict) + schema_version: str = "cross_config.scorer.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "side", _coerce_enum(self.side, ScoreSide)) + if not self.backend_id: + raise ValueError("backend_id must not be empty") + if not self.dtype: + raise ValueError("dtype must not be empty") + if self.world_size < 1: + raise ValueError("world_size must be >= 1") + object.__setattr__(self, "topology", _freeze_mapping(self.topology)) + object.__setattr__(self, "construction_options", _freeze_mapping(self.construction_options)) + object.__setattr__(self, "operator_overrides", _freeze_mapping(self.operator_overrides)) + + +@dataclass(frozen=True) +class KnobDescriptor(SerializableModel): + path: str + lifecycle: IsolationScope + targets: tuple[str, ...] + allowed_values: tuple[Any, ...] = () + derived: bool = False + critical: bool = True + schema_version: str = "cross_config.knob_descriptor.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "lifecycle", _coerce_enum(self.lifecycle, IsolationScope)) + object.__setattr__(self, "targets", tuple(str(target) for target in self.targets)) + object.__setattr__(self, "allowed_values", tuple(_freeze_value(self.allowed_values))) + if not self.path: + raise ValueError("path must not be empty") + if not self.targets: + raise ValueError("targets must not be empty") + + +@dataclass(frozen=True) +class InterventionSpec(SerializableModel): + path: str + values: tuple[Any, ...] + schema_version: str = "cross_config.intervention.v1" + + def __post_init__(self) -> None: + if not self.path: + raise ValueError("path must not be empty") + object.__setattr__(self, "values", tuple(_freeze_value(self.values))) + if not self.values: + raise ValueError("values must not be empty") + + +@dataclass(frozen=True) +class ExperimentDefinition(SerializableModel): + experiment_id: str + scenario_id: str + identity: SemanticIdentitySpec + baseline: Mapping[str, Any] + interventions: tuple[InterventionSpec, ...] = () + scenario: Mapping[str, Any] = field(default_factory=dict) + strategy: PlanningStrategy = PlanningStrategy.ONE_AT_A_TIME + strict_fallback: bool = True + pairwise_paths: tuple[tuple[str, str], ...] = () + contract_source: str = "ws1" + contract_version: str = "current" + schema_version: str = "cross_config.experiment_definition.v1" + + def __post_init__(self) -> None: + if not self.experiment_id: + raise ValueError("experiment_id must not be empty") + if not self.scenario_id: + raise ValueError("scenario_id must not be empty") + if self.contract_source != "ws1": + raise ValueError("Cross-configuration alignment V1 requires contract_source='ws1'") + if self.contract_version != "current": + raise ValueError("Cross-configuration alignment V1 requires contract_version='current'") + object.__setattr__(self, "baseline", _freeze_mapping(self.baseline)) + object.__setattr__(self, "scenario", _freeze_mapping(self.scenario)) + object.__setattr__(self, "interventions", tuple(self.interventions)) + object.__setattr__(self, "strategy", _coerce_enum(self.strategy, PlanningStrategy)) + normalized_pairs: list[tuple[str, str]] = [] + for pair in self.pairwise_paths: + if len(pair) != 2: + raise ValueError("each pairwise_paths entry must contain exactly two paths") + normalized_pairs.append((str(pair[0]), str(pair[1]))) + object.__setattr__(self, "pairwise_paths", tuple(normalized_pairs)) + + +@dataclass(frozen=True) +class ExperimentCase(SerializableModel): + case_id: str + experiment_id: str + scenario_id: str + identity: SemanticIdentitySpec + requested: Mapping[str, Any] + execution_binding: Mapping[str, Any] = field(default_factory=dict) + changed_paths: tuple[str, ...] = () + contract_fingerprint: str = "" + scenario_fingerprint: str = "" + schema_version: str = "cross_config.experiment_case.v1" + + def __post_init__(self) -> None: + if not self.case_id: + raise ValueError("case_id must not be empty") + if not self.experiment_id: + raise ValueError("experiment_id must not be empty") + if not self.scenario_id: + raise ValueError("scenario_id must not be empty") + object.__setattr__(self, "requested", _freeze_mapping(self.requested)) + object.__setattr__(self, "execution_binding", _freeze_mapping(self.execution_binding)) + object.__setattr__(self, "changed_paths", tuple(str(path) for path in self.changed_paths)) + + +@dataclass(frozen=True) +class MaterializedCase(SerializableModel): + case: ExperimentCase + normalized: Mapping[str, Any] + materialized: Mapping[str, Any] + isolation_scope: IsolationScope + construction_fingerprint: str = "" + distributed_context_fingerprint: str = "" + process_fingerprint: str = "" + status: MaterializationStatus = MaterializationStatus.APPLIED + evidence: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.materialized_case.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "normalized", _freeze_mapping(self.normalized)) + object.__setattr__(self, "materialized", _freeze_mapping(self.materialized)) + object.__setattr__( + self, + "isolation_scope", + _coerce_enum(self.isolation_scope, IsolationScope), + ) + object.__setattr__(self, "status", _coerce_enum(self.status, MaterializationStatus)) + object.__setattr__(self, "evidence", _freeze_mapping(self.evidence)) + + +@dataclass(frozen=True) +class CanonicalScoringBatch(SerializableModel): + identity: SemanticIdentitySpec + input_ids: torch.Tensor + selected_token_ids: torch.Tensor + active_mask: torch.Tensor + attention_mask: torch.Tensor + position_ids: Optional[torch.Tensor] = None + metadata: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.canonical_scoring_batch.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "input_ids", _snapshot_tensor(self.input_ids, dtype=torch.long)) + object.__setattr__( + self, "selected_token_ids", _snapshot_tensor(self.selected_token_ids, dtype=torch.long) + ) + object.__setattr__( + self, "active_mask", _snapshot_tensor(self.active_mask, dtype=torch.bool) + ) + object.__setattr__( + self, "attention_mask", _snapshot_tensor(self.attention_mask, dtype=torch.bool) + ) + if self.position_ids is not None: + object.__setattr__( + self, "position_ids", _snapshot_tensor(self.position_ids, dtype=torch.long) + ) + object.__setattr__(self, "metadata", _freeze_mapping(self.metadata)) + + shape = self.input_ids.shape + if self.input_ids.ndim != 2: + raise ValueError("input_ids must have shape [batch, sequence]") + for name in ("selected_token_ids", "active_mask", "attention_mask"): + if getattr(self, name).shape != shape: + raise ValueError(f"{name} shape must match input_ids shape") + if self.position_ids is not None and self.position_ids.shape != shape: + raise ValueError("position_ids shape must match input_ids shape") + + _require_tensor_matches_matrix("input_ids", self.input_ids, self.identity.token_ids) + _require_tensor_matches_matrix( + "selected_token_ids", self.selected_token_ids, self.identity.selected_token_ids + ) + _require_tensor_matches_matrix("active_mask", self.active_mask, self.identity.active_mask) + _require_tensor_matches_matrix( + "attention_mask", self.attention_mask, self.identity.attention_mask + ) + if self.identity.position_ids: + if self.position_ids is None: + raise ValueError("position_ids are required by the semantic identity") + _require_tensor_matches_matrix( + "position_ids", self.position_ids, self.identity.position_ids + ) + elif self.position_ids is not None: + raise ValueError("position_ids were supplied but are absent from semantic identity") + + +def _require_tensor_matches_matrix( + name: str, + tensor: torch.Tensor, + matrix: tuple[tuple[Any, ...], ...], +) -> None: + expected = torch.tensor(matrix, dtype=tensor.dtype, device=tensor.device) + if expected.shape != tensor.shape or not torch.equal(tensor, expected): + raise ValueError(f"{name} does not match the semantic identity") + + +@dataclass(frozen=True) +class RuntimeProvenance(SerializableModel): + requested: Mapping[str, Any] + normalized: Mapping[str, Any] + materialized: Mapping[str, Any] + actual: Mapping[str, Any] + status: MaterializationStatus = MaterializationStatus.APPLIED + construction_fingerprint: str = "" + distributed_context_fingerprint: str = "" + process_fingerprint: str = "" + implementation_fingerprint: str = "" + evidence: Mapping[str, Any] = field(default_factory=dict) + rank: int = 0 + world_size: int = 1 + schema_version: str = "cross_config.runtime_provenance.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "requested", _freeze_mapping(self.requested)) + object.__setattr__(self, "normalized", _freeze_mapping(self.normalized)) + object.__setattr__(self, "materialized", _freeze_mapping(self.materialized)) + object.__setattr__(self, "actual", _freeze_mapping(self.actual)) + object.__setattr__(self, "status", _coerce_enum(self.status, MaterializationStatus)) + object.__setattr__(self, "evidence", _freeze_mapping(self.evidence)) + if self.rank < 0: + raise ValueError("rank must be >= 0") + if self.world_size < 1: + raise ValueError("world_size must be >= 1") + if self.rank >= self.world_size: + raise ValueError("rank must be less than world_size") + + +@dataclass(frozen=True) +class ScoreArtifact(SerializableModel): + case_id: str + attempt_id: str + side: ScoreSide + identity: SemanticIdentitySpec + scorer: ScorerSpec + selected_logprobs: torch.Tensor + active_mask: torch.Tensor + provenance: RuntimeProvenance + schema_version: str = "cross_config.score_artifact.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "side", _coerce_enum(self.side, ScoreSide)) + if not self.case_id: + raise ValueError("case_id must not be empty") + if not self.attempt_id: + raise ValueError("attempt_id must not be empty") + if self.scorer.side is not self.side: + raise ValueError("scorer side must match score artifact side") + object.__setattr__(self, "selected_logprobs", _snapshot_tensor(self.selected_logprobs)) + object.__setattr__( + self, "active_mask", _snapshot_tensor(self.active_mask, dtype=torch.bool) + ) + if self.selected_logprobs.shape != self.active_mask.shape: + raise ValueError("selected_logprobs shape must match active_mask shape") + + +@dataclass(frozen=True) +class TokenComparisonArtifact(SerializableModel): + rollout_logprobs: torch.Tensor + training_logprobs: torch.Tensor + active_mask: torch.Tensor + absolute_diff: torch.Tensor + mismatch_mask: torch.Tensor + fixed_threshold: float + schema_version: str = "cross_config.token_comparison.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "rollout_logprobs", _snapshot_tensor(self.rollout_logprobs)) + object.__setattr__(self, "training_logprobs", _snapshot_tensor(self.training_logprobs)) + object.__setattr__( + self, "active_mask", _snapshot_tensor(self.active_mask, dtype=torch.bool) + ) + object.__setattr__(self, "absolute_diff", _snapshot_tensor(self.absolute_diff)) + object.__setattr__( + self, "mismatch_mask", _snapshot_tensor(self.mismatch_mask, dtype=torch.bool) + ) + shape = self.rollout_logprobs.shape + for name in ( + "training_logprobs", + "active_mask", + "absolute_diff", + "mismatch_mask", + ): + if getattr(self, name).shape != shape: + raise ValueError(f"{name} shape must match rollout_logprobs shape") + if not math.isfinite(self.fixed_threshold) or self.fixed_threshold < 0.0: + raise ValueError("fixed_threshold must be finite and non-negative") + + +@dataclass(frozen=True) +class AlignmentResult(SerializableModel): + case_id: str + attempt_id: str + status: AlignmentStatus + comparable: bool + passed: bool + active_token_count: int + mismatch_count: int + contract_fingerprint: str + fixed_threshold: Optional[float] = None + identity_errors: tuple[str, ...] = () + artifact_errors: tuple[str, ...] = () + diagnostics: Mapping[str, Any] = field(default_factory=dict) + token_artifact: Optional[TokenComparisonArtifact] = None + schema_version: str = "cross_config.alignment_result.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "status", _coerce_enum(self.status, AlignmentStatus)) + object.__setattr__(self, "identity_errors", tuple(self.identity_errors)) + object.__setattr__(self, "artifact_errors", tuple(self.artifact_errors)) + object.__setattr__(self, "diagnostics", _freeze_mapping(self.diagnostics)) + if self.active_token_count < 0: + raise ValueError("active_token_count must be non-negative") + if self.mismatch_count < 0: + raise ValueError("mismatch_count must be non-negative") + if self.mismatch_count > self.active_token_count: + raise ValueError("mismatch_count cannot exceed active_token_count") + if self.status is AlignmentStatus.PASS and not self.passed: + raise ValueError("PASS result must set passed=True") + if self.status is not AlignmentStatus.PASS and self.passed: + raise ValueError("only PASS results may set passed=True") + + +__all__ = [ + "AlignmentResult", + "AlignmentStatus", + "CanonicalScoringBatch", + "ExperimentCase", + "ExperimentDefinition", + "InterventionSpec", + "IsolationScope", + "KnobDescriptor", + "MaterializationStatus", + "MaterializedCase", + "PlanningStrategy", + "RuntimeProvenance", + "ScoreArtifact", + "ScoreSide", + "ScorerSpec", + "SemanticIdentitySpec", + "SerializableModel", + "TokenComparisonArtifact", +] diff --git a/rl_engine/kernels/gtest/tolerance.py b/rl_engine/kernels/gtest/tolerance.py index d0481e83..917373ce 100644 --- a/rl_engine/kernels/gtest/tolerance.py +++ b/rl_engine/kernels/gtest/tolerance.py @@ -3,7 +3,9 @@ from __future__ import annotations +import hashlib import json +import math from pathlib import Path from typing import Any @@ -17,4 +19,64 @@ def load_contract(path: str | Path = _CONTRACT_PATH) -> dict[str, Any]: return json.load(handle) -__all__ = ["load_contract"] +def resolve_logprob_threshold(dtype: Any) -> float: + """Return the fixed WS1 selected-logprob absolute-difference threshold. + + The contract path is intentionally not configurable through this accessor. + Cross-configuration experiment definitions may select a dtype, but they cannot + inject or override a numerical threshold. + """ + + dtype_name = _normalize_dtype_name(dtype) + contract = load_contract() + try: + values = contract["accuracy"]["default"]["logprob"][dtype_name] + raw_threshold = values["atol"] + except (KeyError, TypeError) as exc: + raise ValueError(f"WS1 has no logprob threshold for dtype {dtype_name!r}") from exc + if isinstance(raw_threshold, bool) or not isinstance(raw_threshold, (int, float)): + raise ValueError(f"invalid WS1 logprob threshold for dtype {dtype_name!r}") + threshold = float(raw_threshold) + if not math.isfinite(threshold) or threshold < 0.0: + raise ValueError(f"invalid WS1 logprob threshold for dtype {dtype_name!r}") + return threshold + + +def tolerance_contract_fingerprint() -> str: + """Return a deterministic fingerprint of the current WS1 contract contents.""" + + canonical = json.dumps( + load_contract(), + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + return hashlib.sha256(canonical).hexdigest() + + +def _normalize_dtype_name(dtype: Any) -> str: + normalized = str(dtype).strip().lower().replace("torch.", "").replace("-", "") + aliases = { + "bf16": "bfloat16", + "bfloat16": "bfloat16", + "fp16": "float16", + "float16": "float16", + "half": "float16", + "fp32": "float32", + "float32": "float32", + "float": "float32", + } + try: + return aliases[normalized] + except KeyError as exc: + valid = ", ".join(sorted(set(aliases.values()))) + raise ValueError( + f"unsupported WS1 logprob dtype {dtype!r}; expected one of: {valid}" + ) from exc + + +__all__ = [ + "load_contract", + "resolve_logprob_threshold", + "tolerance_contract_fingerprint", +] From eb6e6bb7d5b5f588b8f9cf3dbf0f0f47d0baab58 Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sun, 19 Jul 2026 08:31:43 +0800 Subject: [PATCH 03/33] feat(alignment): add runtime materialization and scoring bridge --- rl_engine/alignment/cross_config/operators.py | 265 ++++++++++ rl_engine/alignment/cross_config/runtime.py | 465 ++++++++++++++++++ rl_engine/executors/stateless_executor.py | 44 +- 3 files changed, 770 insertions(+), 4 deletions(-) create mode 100644 rl_engine/alignment/cross_config/operators.py create mode 100644 rl_engine/alignment/cross_config/runtime.py diff --git a/rl_engine/alignment/cross_config/operators.py b/rl_engine/alignment/cross_config/operators.py new file mode 100644 index 00000000..e3105a28 --- /dev/null +++ b/rl_engine/alignment/cross_config/operators.py @@ -0,0 +1,265 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Target-specific semantic operator selection for alignment cases.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, replace +from typing import Any, Literal, Mapping, Optional, cast + +import torch + +from rl_engine.kernels.semantic_registry import ( + OperatorInstanceProvenance, + OperatorRequirements, + OperatorResolution, + OperatorResolutionPolicy, + OperatorSession, + SemanticOperatorCatalog, +) + +OperatorTarget = Literal["rollout", "training", "both"] +ConcreteOperatorTarget = Literal["rollout", "training"] + + +@dataclass(frozen=True) +class OperatorOverride: + """Backend overrides for one semantic operator on either or both sides.""" + + semantic_op: str + rollout_backend: Optional[str] = None + training_backend: Optional[str] = None + + def __post_init__(self) -> None: + semantic_op = self.semantic_op.strip() + if not semantic_op: + raise ValueError("semantic_op must not be empty") + rollout_backend = _normalized_optional_backend(self.rollout_backend) + training_backend = _normalized_optional_backend(self.training_backend) + if rollout_backend is None and training_backend is None: + raise ValueError("operator override must select rollout, training, or both") + object.__setattr__(self, "semantic_op", semantic_op) + object.__setattr__(self, "rollout_backend", rollout_backend) + object.__setattr__(self, "training_backend", training_backend) + + @classmethod + def for_target( + cls, + *, + semantic_op: str, + backend_id: str, + target: OperatorTarget, + ) -> OperatorOverride: + """Create a rollout-only, training-only, or dual-side override.""" + + normalized_target = target.strip().lower() + if normalized_target == "rollout": + return cls(semantic_op=semantic_op, rollout_backend=backend_id) + if normalized_target == "training": + return cls(semantic_op=semantic_op, training_backend=backend_id) + if normalized_target == "both": + return cls( + semantic_op=semantic_op, + rollout_backend=backend_id, + training_backend=backend_id, + ) + raise ValueError("target must be 'rollout', 'training', or 'both'") + + def backend_for(self, target: ConcreteOperatorTarget) -> Optional[str]: + normalized_target = _concrete_target(target) + if normalized_target == "rollout": + return self.rollout_backend + return self.training_backend + + def to_dict(self) -> dict[str, Any]: + return { + "semantic_op": self.semantic_op, + "rollout_backend": self.rollout_backend, + "training_backend": self.training_backend, + } + + +@dataclass(frozen=True) +class ResolvedOperatorOverride: + """Target-specific exact resolutions produced from an operator override.""" + + semantic_op: str + rollout: Optional[OperatorResolution] = None + training: Optional[OperatorResolution] = None + + def for_target(self, target: ConcreteOperatorTarget) -> Optional[OperatorResolution]: + normalized_target = _concrete_target(target) + return self.rollout if normalized_target == "rollout" else self.training + + def to_dict(self) -> dict[str, Any]: + return { + "semantic_op": self.semantic_op, + "rollout": None if self.rollout is None else self.rollout.to_dict(), + "training": None if self.training is None else self.training.to_dict(), + } + + +class OperatorBridge: + """Resolve and instantiate semantic operator overrides without planner branches.""" + + def __init__( + self, + catalog: Optional[SemanticOperatorCatalog | OperatorSession] = None, + *, + policy: Optional[OperatorResolutionPolicy] = None, + ): + """Create a bridge backed by one case-local operator session.""" + + if isinstance(catalog, OperatorSession): + self.catalog = catalog.catalog + self.session = catalog + self.policy = policy or catalog.policy + else: + if catalog is None: + # Built-in descriptors are repository integration details; the + # generic semantic catalog itself remains backend-neutral. + from rl_engine.kernels.registry import kernel_registry + + catalog = kernel_registry.semantic + if not isinstance(catalog, SemanticOperatorCatalog): + raise TypeError("catalog must be a SemanticOperatorCatalog or OperatorSession") + self.catalog = catalog + self.policy = policy or OperatorResolutionPolicy() + self.session = self.catalog.session(self.policy) + + def resolve_override( + self, + override: OperatorOverride, + *, + requirements: Mapping[str, OperatorRequirements], + strict: bool = True, + ) -> ResolvedOperatorOverride: + """Resolve only the sides explicitly selected by ``override``.""" + + resolved: dict[ConcreteOperatorTarget, OperatorResolution] = {} + targets: tuple[ConcreteOperatorTarget, ...] = ("rollout", "training") + for target in targets: + backend_id = override.backend_for(target) + if backend_id is None: + continue + target_requirements = requirements.get(target) + if target_requirements is None: + raise ValueError(f"missing operator requirements for target {target!r}") + target_policy = replace(self.policy, strict=strict) + resolved[target] = self.session.resolve( + semantic_op=override.semantic_op, + requested_backend=backend_id, + target=target, + requirements=target_requirements, + policy=target_policy, + ) + return ResolvedOperatorOverride( + semantic_op=override.semantic_op, + rollout=resolved.get("rollout"), + training=resolved.get("training"), + ) + + def instantiate( + self, + resolved: ResolvedOperatorOverride, + *, + target: ConcreteOperatorTarget, + factory_kwargs: Optional[Mapping[str, Any]] = None, + cache: bool = False, + ) -> Any: + """Instantiate one resolved side; rollout and training remain independent.""" + + resolution = resolved.for_target(target) + if resolution is None: + raise ValueError(f"operator override does not select target {target!r}") + return self.session.instantiate( + resolution, + factory_kwargs=factory_kwargs, + cache=cache, + ) + + def instance_provenance( + self, + resolved: ResolvedOperatorOverride, + *, + target: ConcreteOperatorTarget, + instance: Any, + ) -> OperatorInstanceProvenance: + resolution = resolved.for_target(target) + if resolution is None: + raise ValueError(f"operator override does not select target {target!r}") + return self.session.instance_provenance(resolution, instance) + + +def selected_logprobs_with_operator( + operator: Any, + logits: torch.Tensor, + token_ids: torch.Tensor, + *, + active_mask: Optional[torch.Tensor] = None, + temperature: float = 1.0, + output_dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Apply the repository selected-logprob interface with common mask semantics.""" + + if not math.isfinite(temperature) or temperature <= 0.0: + raise ValueError("temperature must be finite and greater than zero") + if logits.shape[:-1] != token_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"token_ids shape {tuple(token_ids.shape)}" + ) + mask: Optional[torch.Tensor] = None + safe_token_ids = token_ids.to(device=logits.device, dtype=torch.long) + if active_mask is not None: + if active_mask.shape != token_ids.shape: + raise ValueError("active_mask shape must match token_ids shape") + mask = active_mask.to(device=logits.device, dtype=torch.bool) + safe_token_ids = safe_token_ids.masked_fill(~mask, 0) + + scaled_logits = logits.float() / float(temperature) + if hasattr(operator, "apply_fp32") and callable(operator.apply_fp32): + selected = operator.apply_fp32(scaled_logits, safe_token_ids) + elif callable(operator): + selected = operator(scaled_logits, safe_token_ids) + else: + raise TypeError("selected-logprob operator must be callable or expose apply_fp32") + if not isinstance(selected, torch.Tensor): + raise TypeError("selected-logprob operator must return a torch.Tensor") + if selected.shape != token_ids.shape: + raise ValueError( + f"selected-logprob output shape {tuple(selected.shape)} must match " + f"token_ids shape {tuple(token_ids.shape)}" + ) + selected = selected.to(device=logits.device, dtype=output_dtype) + if mask is not None: + selected = selected.masked_fill(~mask, 0.0) + return selected + + +def _normalized_optional_backend(value: Optional[str]) -> Optional[str]: + if value is None: + return None + normalized = value.strip() + if not normalized: + raise ValueError("backend_id must not be empty") + return normalized + + +def _concrete_target(target: ConcreteOperatorTarget) -> ConcreteOperatorTarget: + normalized = target.strip().lower() + if normalized not in {"rollout", "training"}: + raise ValueError("target must be 'rollout' or 'training'") + return cast(ConcreteOperatorTarget, normalized) + + +__all__ = [ + "ConcreteOperatorTarget", + "OperatorBridge", + "OperatorOverride", + "OperatorTarget", + "ResolvedOperatorOverride", + "selected_logprobs_with_operator", +] diff --git a/rl_engine/alignment/cross_config/runtime.py b/rl_engine/alignment/cross_config/runtime.py new file mode 100644 index 00000000..3eab0352 --- /dev/null +++ b/rl_engine/alignment/cross_config/runtime.py @@ -0,0 +1,465 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Thin runtime materialization facade for the V1 allowlist.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Iterable, Mapping, Protocol, Sequence + +from rl_engine.alignment.cross_config.planner import V1_KNOBS +from rl_engine.alignment.cross_config.schema import ( + ExperimentCase, + IsolationScope, + KnobDescriptor, + MaterializationStatus, + MaterializedCase, + RuntimeProvenance, + SerializableModel, +) + + +@dataclass(frozen=True) +class KnobApplication(SerializableModel): + """One adapter's requested, materialized, and observed value.""" + + path: str + requested: Any + materialized: Any + actual: Any + lifecycle: IsolationScope + status: MaterializationStatus + evidence: Mapping[str, Any] = field(default_factory=dict) + critical: bool = True + schema_version: str = "cross_config.knob_application.v1" + + def __post_init__(self) -> None: + object.__setattr__(self, "lifecycle", IsolationScope(self.lifecycle)) + object.__setattr__(self, "status", MaterializationStatus(self.status)) + for name in ("requested", "materialized", "actual"): + object.__setattr__(self, name, _freeze_value(getattr(self, name))) + object.__setattr__(self, "evidence", _freeze_mapping(self.evidence)) + + +@dataclass(frozen=True) +class RuntimeBinding: + """Small, backend-neutral handoff from materialization to execution. + + Runtime adapters may construct repository-specific objects internally, but + the core runner sees only the values required to create scorers and validate + lifecycle identity. New vLLM, FSDP, or other adapters therefore do not + change the runner's type surface. + """ + + batch_size: int + side_configs: Mapping[str, Mapping[str, Any]] + topology: Mapping[str, Mapping[str, Any]] + scorer: Mapping[str, Any] + operator_backends: Mapping[str, str] + runtime_kind: str + + def __post_init__(self) -> None: + if isinstance(self.batch_size, bool) or not isinstance(self.batch_size, int): + raise TypeError("batch_size must be an integer") + if self.batch_size < 1: + raise ValueError("batch_size must be greater than zero") + if not isinstance(self.runtime_kind, str) or not self.runtime_kind.strip(): + raise ValueError("runtime_kind must be a non-empty string") + for name, value in ( + ("side_configs", self.side_configs), + ("topology", self.topology), + ): + for target in ("rollout", "training"): + if not isinstance(value.get(target), Mapping): + raise ValueError(f"{name} must define a {target} mapping") + for target in ("rollout", "training"): + world_size = self.topology[target].get("world_size") + if isinstance(world_size, bool) or not isinstance(world_size, int) or world_size < 1: + raise ValueError(f"{target} topology must define a positive integer world_size") + backend = self.operator_backends.get(target) + if not isinstance(backend, str) or not backend.strip(): + raise ValueError(f"operator_backends must define a non-empty {target} backend") + object.__setattr__(self, "side_configs", _freeze_mapping(self.side_configs)) + object.__setattr__(self, "topology", _freeze_mapping(self.topology)) + object.__setattr__(self, "scorer", _freeze_mapping(self.scorer)) + object.__setattr__(self, "operator_backends", _freeze_mapping(self.operator_backends)) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": "cross_config.runtime_binding.v1", + "runtime_kind": self.runtime_kind, + "batch_size": self.batch_size, + "side_configs": _plain_mapping(self.side_configs), + "topology": _plain_mapping(self.topology), + "scorer": _plain_mapping(self.scorer), + "operators": dict(self.operator_backends), + } + + +@dataclass(frozen=True) +class AdapterMaterialization: + """Output of a typed runtime adapter before the facade adds fingerprints.""" + + applications: tuple[KnobApplication, ...] + binding: RuntimeBinding + + def __post_init__(self) -> None: + object.__setattr__(self, "applications", tuple(self.applications)) + + +class RuntimeMaterializer(Protocol): + """Adapter boundary used by the small ``RuntimeTools`` facade. + + The declared implementation fingerprint must deterministically identify the + executable materialization path and change when that implementation changes. + """ + + runtime_kind: str + + @property + def implementation_fingerprint(self) -> str: ... + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: ... + + +@dataclass(frozen=True) +class RuntimeMaterialization: + materialized_case: MaterializedCase + provenance: RuntimeProvenance + applications: tuple[KnobApplication, ...] + binding: RuntimeBinding + + @property + def executable_in_strict_mode(self) -> bool: + return self.materialized_case.status is MaterializationStatus.APPLIED + + +class RuntimeMaterializationError(RuntimeError): + pass + + +class RuntimeTools: + """Materialize cases and compute reuse fingerprints without owning execution.""" + + def __init__(self, descriptors: Mapping[str, KnobDescriptor] = V1_KNOBS): + self.descriptors = dict(descriptors) + + def materialize( + self, + case: ExperimentCase, + adapter: RuntimeMaterializer, + ) -> RuntimeMaterialization: + runtime_kind = _adapter_identity(adapter, "runtime_kind") + adapter_implementation_fingerprint = _adapter_identity( + adapter, + "implementation_fingerprint", + ) + normalized = _plain_mapping(case.requested) + adapter_result = adapter.materialize(normalized, self.descriptors) + if not isinstance(adapter_result, AdapterMaterialization): + raise RuntimeMaterializationError("runtime adapter must return AdapterMaterialization") + if not isinstance(adapter_result.binding, RuntimeBinding): + raise RuntimeMaterializationError("runtime adapter must return a RuntimeBinding") + if adapter_result.binding.runtime_kind != runtime_kind: + raise RuntimeMaterializationError( + "runtime binding kind must match the materializer runtime_kind" + ) + applications = tuple(adapter_result.applications) + _validate_application_contract(normalized, applications, self.descriptors) + materialized = _mapping_from_applications(applications, "materialized") + actual = _mapping_from_applications(applications, "actual") + status = _aggregate_status(application.status for application in applications) + construction_fingerprint = _scope_fingerprint( + runtime_kind, + adapter_implementation_fingerprint, + applications, + scopes=( + IsolationScope.ENGINE_CONSTRUCTION, + IsolationScope.DISTRIBUTED_CONTEXT, + IsolationScope.PROCESS, + ), + ) + distributed_fingerprint = _scope_fingerprint( + runtime_kind, + adapter_implementation_fingerprint, + applications, + scopes=(IsolationScope.DISTRIBUTED_CONTEXT, IsolationScope.PROCESS), + ) + process_fingerprint = _scope_fingerprint( + runtime_kind, + adapter_implementation_fingerprint, + applications, + scopes=(IsolationScope.PROCESS,), + ) + isolation_scope = _strongest_scope( + [ + self.descriptors[path].lifecycle + for path in (case.changed_paths or tuple(_flatten(normalized))) + ] + ) + evidence = { + "runtime_kind": runtime_kind, + "execution_binding": case.execution_binding, + "adapter_implementation_fingerprint": adapter_implementation_fingerprint, + "binding_fingerprint": _fingerprint(adapter_result.binding.to_dict()), + "applications": { + application.path: application.to_dict() for application in applications + }, + } + materialized_case = MaterializedCase( + case=case, + normalized=normalized, + materialized=materialized, + isolation_scope=isolation_scope, + construction_fingerprint=construction_fingerprint, + distributed_context_fingerprint=distributed_fingerprint, + process_fingerprint=process_fingerprint, + status=status, + evidence=evidence, + ) + provenance = RuntimeProvenance( + requested=_plain_mapping(case.requested), + normalized=normalized, + materialized=materialized, + actual=actual, + status=status, + construction_fingerprint=construction_fingerprint, + distributed_context_fingerprint=distributed_fingerprint, + process_fingerprint=process_fingerprint, + implementation_fingerprint=adapter_implementation_fingerprint, + evidence=evidence, + ) + return RuntimeMaterialization( + materialized_case=materialized_case, + provenance=provenance, + applications=applications, + binding=adapter_result.binding, + ) + + @staticmethod + def require_executable( + materialization: RuntimeMaterialization, + *, + strict: bool, + ) -> None: + status = materialization.materialized_case.status + rejected = { + MaterializationStatus.ERROR, + MaterializationStatus.UNSUPPORTED, + MaterializationStatus.UNOBSERVABLE, + } + if strict: + rejected.add(MaterializationStatus.FALLBACK) + if status in rejected: + problems = [ + f"{application.path}={application.status.value}: " + f"{application.evidence.get('reason', 'no evidence')}" + for application in materialization.applications + if application.status is not MaterializationStatus.APPLIED + ] + raise RuntimeMaterializationError( + f"case {materialization.materialized_case.case.case_id} is not executable: " + + "; ".join(problems) + ) + + @staticmethod + def can_reuse(previous: RuntimeMaterialization, current: RuntimeMaterialization) -> bool: + """Reuse only exact semantic and implementation identities with matching state.""" + + previous_case = previous.materialized_case + current_case = current.materialized_case + return ( + previous_case.status is MaterializationStatus.APPLIED + and current_case.status is MaterializationStatus.APPLIED + and previous_case.case.identity == current_case.case.identity + and previous_case.case.execution_binding == current_case.case.execution_binding + and previous.provenance.implementation_fingerprint + == current.provenance.implementation_fingerprint + and previous_case.process_fingerprint == current_case.process_fingerprint + and previous_case.distributed_context_fingerprint + == current_case.distributed_context_fingerprint + and previous_case.construction_fingerprint == current_case.construction_fingerprint + ) + + +def _aggregate_status(statuses: Iterable[MaterializationStatus]) -> MaterializationStatus: + priority = ( + MaterializationStatus.ERROR, + MaterializationStatus.UNSUPPORTED, + MaterializationStatus.UNOBSERVABLE, + MaterializationStatus.FALLBACK, + MaterializationStatus.APPLIED, + ) + status_set = set(statuses) + if not status_set: + return MaterializationStatus.ERROR + return next(status for status in priority if status in status_set) + + +def _validate_application_contract( + normalized: Mapping[str, Any], + applications: tuple[KnobApplication, ...], + descriptors: Mapping[str, KnobDescriptor], +) -> None: + expected = _flatten(normalized) + observed_paths = [application.path for application in applications] + duplicate_paths = sorted(path for path in set(observed_paths) if observed_paths.count(path) > 1) + missing_paths = sorted(set(expected).difference(observed_paths)) + unknown_paths = sorted(set(observed_paths).difference(expected)) + missing_descriptors = sorted(set(expected).difference(descriptors)) + problems: list[str] = [] + if missing_paths: + problems.append(f"missing paths={missing_paths!r}") + if duplicate_paths: + problems.append(f"duplicate paths={duplicate_paths!r}") + if unknown_paths: + problems.append(f"unknown paths={unknown_paths!r}") + if missing_descriptors: + problems.append(f"missing descriptors={missing_descriptors!r}") + for application in applications: + descriptor = descriptors.get(application.path) + if descriptor is None or application.path not in expected: + continue + if _plain_value(application.requested) != _plain_value(expected[application.path]): + problems.append(f"{application.path} requested value differs from normalized case") + if application.lifecycle is not descriptor.lifecycle: + problems.append(f"{application.path} lifecycle differs from descriptor") + if application.critical is not descriptor.critical: + problems.append(f"{application.path} critical flag differs from descriptor") + if application.status is MaterializationStatus.APPLIED: + if _plain_value(application.actual) != _plain_value(application.materialized): + problems.append(f"{application.path} applied actual differs from materialized") + if not descriptor.derived and _plain_value(application.materialized) != _plain_value( + expected[application.path] + ): + problems.append( + f"{application.path} applied materialized value differs from normalized case" + ) + if problems: + raise RuntimeMaterializationError( + "runtime adapter returned invalid V1 knob applications: " + "; ".join(problems) + ) + + +def _mapping_from_applications( + applications: Sequence[KnobApplication], attribute: str +) -> dict[str, Any]: + result: dict[str, Any] = {} + for application in applications: + _set_path(result, application.path, getattr(application, attribute)) + return result + + +def _scope_fingerprint( + runtime_kind: str, + implementation_fingerprint: str, + applications: Sequence[KnobApplication], + *, + scopes: Sequence[IsolationScope], +) -> str: + scope_set = set(scopes) + values = { + application.path: application.materialized + for application in applications + if application.lifecycle in scope_set + } + return _fingerprint( + { + "runtime_kind": runtime_kind, + "implementation_fingerprint": implementation_fingerprint, + "values": values, + } + ) + + +def _strongest_scope(scopes: Sequence[IsolationScope]) -> IsolationScope: + order = { + IsolationScope.REQUEST: 0, + IsolationScope.ENGINE_CONSTRUCTION: 1, + IsolationScope.DISTRIBUTED_CONTEXT: 2, + IsolationScope.PROCESS: 3, + } + return max(scopes, key=order.__getitem__, default=IsolationScope.REQUEST) + + +def _flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + result: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}.{key}" if prefix else key + if isinstance(child, Mapping): + result.update(_flatten(child, path)) + else: + result[path] = child + return result + + +def _set_path(value: dict[str, Any], path: str, child: Any) -> None: + current = value + parts = path.split(".") + for part in parts[:-1]: + current = current.setdefault(part, {}) + current[parts[-1]] = child + + +def _plain_mapping(value: Mapping[str, Any]) -> dict[str, Any]: + return {str(key): _plain_value(item) for key, item in value.items()} + + +def _plain_value(value: Any) -> Any: + if isinstance(value, Mapping): + return {str(key): _plain_value(item) for key, item in value.items()} + if isinstance(value, (tuple, list)): + return [_plain_value(item) for item in value] + if isinstance(value, (set, frozenset)): + return [_plain_value(item) for item in sorted(value, key=repr)] + return value + + +def _freeze_value(value: Any) -> Any: + if isinstance(value, Mapping): + return MappingProxyType({str(key): _freeze_value(item) for key, item in value.items()}) + if isinstance(value, (tuple, list)): + return tuple(_freeze_value(item) for item in value) + if isinstance(value, (set, frozenset)): + return frozenset(_freeze_value(item) for item in value) + return value + + +def _freeze_mapping(value: Mapping[str, Any]) -> Mapping[str, Any]: + return _freeze_value(value) + + +def _fingerprint(value: Mapping[str, Any]) -> str: + payload = json.dumps( + _plain_mapping(value), + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _adapter_identity(adapter: RuntimeMaterializer, attribute: str) -> str: + value = getattr(adapter, attribute, None) + if not isinstance(value, str) or not value.strip(): + raise RuntimeMaterializationError(f"runtime adapter {attribute} must be a non-empty string") + return value.strip() + + +__all__ = [ + "AdapterMaterialization", + "KnobApplication", + "RuntimeBinding", + "RuntimeMaterialization", + "RuntimeMaterializationError", + "RuntimeMaterializer", + "RuntimeTools", +] diff --git a/rl_engine/executors/stateless_executor.py b/rl_engine/executors/stateless_executor.py index 2047218f..70a25f99 100644 --- a/rl_engine/executors/stateless_executor.py +++ b/rl_engine/executors/stateless_executor.py @@ -16,6 +16,7 @@ StatelessForwardMode = Literal["reference", "reward", "both"] StatelessAttentionBackend = Literal["flash_attention_2", "sdpa", "eager", "model_default"] RewardAdapter = Callable[["StatelessForwardOutputs", "StatelessForwardInputs"], torch.Tensor] +SelectedLogprobCallable = Callable[..., torch.Tensor] _MISSING = object() @@ -57,6 +58,7 @@ class StatelessForwardInputs: attention_mask: torch.Tensor completion_mask: torch.Tensor labels: Optional[torch.Tensor] = None + position_ids: Optional[torch.Tensor] = None @dataclass(frozen=True) @@ -105,10 +107,12 @@ def __init__( config: Optional[StatelessForwardConfig] = None, *, reward_adapter: Optional[RewardAdapter] = None, + selected_logprob_fn: Optional[SelectedLogprobCallable] = None, ): self.model = model self.config = config or StatelessForwardConfig() self.reward_adapter = reward_adapter or default_reward_adapter + self.selected_logprob_fn = selected_logprob_fn def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: _validate_inputs(inputs, self.config) @@ -170,6 +174,7 @@ def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult: inputs, temperature=self.config.temperature, output_dtype=self.config.output_dtype, + selected_logprob_fn=self.selected_logprob_fn, ) if self.config.return_token_scores: token_scores = reference_logps @@ -215,8 +220,14 @@ def score_reference_logprobs( *, temperature: float = 1.0, output_dtype: torch.dtype = torch.float32, + selected_logprob_fn: Optional[SelectedLogprobCallable] = None, ) -> torch.Tensor: - """Compute causal next-token selected logprobs aligned to ``[B, S]`` masks.""" + """Compute causal next-token selected logprobs aligned to ``[B, S]`` masks. + + ``selected_logprob_fn`` is an exact injection seam with the same callable + contract as :func:`selected_logprobs_reference`. Leaving it unset preserves + the historical PyTorch-reference behavior. + """ if logits.ndim != 3: raise ValueError(f"reference logits must have shape [B, S, V], got {tuple(logits.shape)}") @@ -237,13 +248,21 @@ def score_reference_logprobs( shifted_logits = logits[:, :-1, :] shifted_labels = labels[:, 1:] shifted_mask = _bool_mask(inputs.completion_mask[:, 1:], device=logits.device) - shifted_logps = selected_logprobs_reference( + scorer = selected_logprob_fn or selected_logprobs_reference + shifted_logps = scorer( shifted_logits, shifted_labels.to(device=logits.device), mask=shifted_mask, temperature=temperature, output_dtype=output_dtype, ) + if not isinstance(shifted_logps, torch.Tensor): + raise TypeError("selected_logprob_fn must return a torch.Tensor") + if shifted_logps.shape != shifted_labels.shape: + raise ValueError( + "selected_logprob_fn output shape must match selected token IDs, got " + f"{tuple(shifted_logps.shape)} and {tuple(shifted_labels.shape)}" + ) result = torch.zeros( inputs.input_ids.shape, device=logits.device, @@ -372,11 +391,20 @@ def _temporarily_configure_stateless_model( config: StatelessForwardConfig, ) -> Iterator[dict[str, float | int | str | bool]]: saved = _model_config_snapshot(model, config) - policy = configure_stateless_model(model, config) + saved_training_modes = tuple((module, module.training) for module in model.modules()) + model.eval() try: + policy = configure_stateless_model(model, config) + policy["model_eval_during_forward"] = True yield policy finally: - _restore_model_config_snapshot(saved) + try: + _restore_model_config_snapshot(saved) + finally: + # Restore each module directly. Calling ``model.train(...)`` would + # flatten intentionally mixed child-module modes. + for module, was_training in saved_training_modes: + module.training = was_training def extract_kv_cache_outputs(raw_outputs: Any) -> Optional[Any]: @@ -450,12 +478,16 @@ def _validate_inputs(inputs: StatelessForwardInputs, config: StatelessForwardCon raise ValueError("completion_mask shape must match input_ids shape") if inputs.labels is not None and inputs.labels.shape != input_ids.shape: raise ValueError("labels shape must match input_ids shape") + if inputs.position_ids is not None and inputs.position_ids.shape != input_ids.shape: + raise ValueError("position_ids shape must match input_ids shape") if attention_mask.device != input_ids.device: raise ValueError("attention_mask device must match input_ids device") if completion_mask.device != input_ids.device: raise ValueError("completion_mask device must match input_ids device") if inputs.labels is not None and inputs.labels.device != input_ids.device: raise ValueError("labels device must match input_ids device") + if inputs.position_ids is not None and inputs.position_ids.device != input_ids.device: + raise ValueError("position_ids device must match input_ids device") if config.max_batch_size is not None and input_ids.shape[0] > config.max_batch_size: raise ValueError( f"batch size {input_ids.shape[0]} exceeds max_batch_size {config.max_batch_size}" @@ -479,6 +511,10 @@ def _run_no_cache_forward( "input_ids": inputs.input_ids, "attention_mask": inputs.attention_mask, } + if inputs.position_ids is not None: + if not _call_accepts_keyword(model, "position_ids"): + raise ValueError("model does not accept the canonical batch position_ids") + kwargs["position_ids"] = inputs.position_ids if _call_accepts_keyword(model, "use_cache"): kwargs["use_cache"] = False return model(**kwargs), True From 18291393a05589e316a78ae5a725d6b165e58ad5 Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sun, 19 Jul 2026 08:32:57 +0800 Subject: [PATCH 04/33] feat(alignment): add execution and artifact primitives --- .../alignment/cross_config/_execution.py | 617 ++++++++++++++++++ rl_engine/alignment/cross_config/artifacts.py | 508 ++++++++++++++ 2 files changed, 1125 insertions(+) create mode 100644 rl_engine/alignment/cross_config/_execution.py create mode 100644 rl_engine/alignment/cross_config/artifacts.py diff --git a/rl_engine/alignment/cross_config/_execution.py b/rl_engine/alignment/cross_config/_execution.py new file mode 100644 index 00000000..a860b951 --- /dev/null +++ b/rl_engine/alignment/cross_config/_execution.py @@ -0,0 +1,617 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Private scoring contracts and child-process supervision.""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import multiprocessing as mp +import os +import tempfile +import time +import traceback +from contextlib import contextmanager +from dataclasses import dataclass, field +from enum import Enum +from pathlib import Path +from typing import Any, Iterator, Mapping, Optional, Protocol, Sequence + +import torch + +from rl_engine.alignment.cross_config._json import strict_json_loads +from rl_engine.alignment.cross_config.schema import CanonicalScoringBatch, ScorerSpec, ScoreSide + + +class PairedRunnerError(RuntimeError): + """Base error for a paired scoring attempt.""" + + +class OperatorExecutionError(PairedRunnerError): + """Raised when exact operator evidence cannot authorize execution.""" + + +class ChildScoringError(PairedRunnerError): + """Raised when a scoring child exits without a valid result.""" + + +class ScoringTimeoutError(PairedRunnerError): + """Raised after all scoring children are stopped at the deadline.""" + + +class RankCompletenessError(PairedRunnerError): + """Raised when rank results are missing, duplicated, or inconsistent.""" + + +class ScorerIdentityError(PairedRunnerError): + """Raised when paired scorer model state is not logically identical.""" + + +@dataclass(frozen=True) +class RankScore: + """One rank's full canonical selected-logprob observation.""" + + rank: int + world_size: int + selected_logprobs: torch.Tensor + metadata: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.rank < 0: + raise ValueError("rank must be >= 0") + if self.world_size < 1: + raise ValueError("world_size must be >= 1") + if self.rank >= self.world_size: + raise ValueError("rank must be less than world_size") + if not isinstance(self.selected_logprobs, torch.Tensor): + raise TypeError("selected_logprobs must be a torch.Tensor") + object.__setattr__( + self, + "selected_logprobs", + self.selected_logprobs.detach().to(device="cpu").clone(), + ) + object.__setattr__(self, "metadata", dict(self.metadata)) + + +class PairedScorer(Protocol): + """Small injection boundary used by the paired-runner control plane.""" + + spec: ScorerSpec + + def score( + self, + batch: CanonicalScoringBatch, + *, + batch_size: int, + operator: Any, + ) -> torch.Tensor | RankScore | Sequence[RankScore]: ... + + +class ChildSupervisor: + """Own the lifecycle of the two isolated scoring children.""" + + def __init__(self, start_method: Optional[str] = None): + available = mp.get_all_start_methods() + resolved = start_method or ("fork" if "fork" in available else "spawn") + if resolved not in available: + raise ValueError(f"multiprocessing start method is unavailable: {resolved}") + self.start_method = resolved + self._active_processes: list[mp.Process] = [] + + @property + def active_child_pids(self) -> tuple[int, ...]: + return tuple( + process.pid + for process in self._active_processes + if process.pid is not None and process.is_alive() + ) + + def run( + self, + attempt_dir: Path, + batch: CanonicalScoringBatch, + *, + batch_size: int, + scorers: Mapping[str, PairedScorer], + specs: Mapping[str, ScorerSpec], + instances: Mapping[str, Any], + timeout_seconds: float, + ) -> dict[str, Mapping[str, Any]]: + context: Any = mp.get_context(self.start_method) + processes: dict[str, mp.Process] = {} + with tempfile.TemporaryDirectory(prefix=".paired-runner-", dir=attempt_dir) as tmp: + temporary_dir = Path(tmp) + try: + for target in ("rollout", "training"): + process = context.Process( + target=_score_child, + name=f"cross-config-{target}", + args=( + temporary_dir / f"{target}.pt", + temporary_dir / f"{target}.error.json", + scorers[target], + specs[target], + batch, + batch_size, + instances[target], + ), + ) + process.start() + processes[target] = process + self._active_processes = list(processes.values()) + self._wait( + processes, + temporary_dir, + timeout_seconds=timeout_seconds, + ) + return { + target: _load_child_result(temporary_dir / f"{target}.pt") + for target in ("rollout", "training") + } + finally: + _stop_processes(tuple(processes.values())) + self._active_processes = [] + + @staticmethod + def _wait( + processes: Mapping[str, mp.Process], + temporary_dir: Path, + *, + timeout_seconds: float, + ) -> None: + deadline = time.monotonic() + timeout_seconds + unfinished = set(processes) + while unfinished: + for target in tuple(unfinished): + process = processes[target] + process.join(timeout=0.01) + if process.is_alive(): + continue + unfinished.remove(target) + if process.exitcode != 0: + detail = _child_error_detail(temporary_dir / f"{target}.error.json") + raise ChildScoringError( + f"{target} scoring child failed with exit code " + f"{process.exitcode}: {detail}" + ) + if unfinished and time.monotonic() >= deadline: + labels = ", ".join(sorted(unfinished)) + raise ScoringTimeoutError( + f"paired scoring exceeded {timeout_seconds:.3f}s; " + f"stopped children: {labels}" + ) + + +def _score_child( + result_path: Path, + error_path: Path, + scorer: PairedScorer, + spec: ScorerSpec, + batch: CanonicalScoringBatch, + batch_size: int, + operator: Any, +) -> None: + try: + with _read_only_scoring_guard(scorer, verify_state=True) as evidence: + with torch.no_grad(): + output = scorer.score(batch, batch_size=batch_size, operator=operator) + ranks = _coerce_rank_scores(output, spec.world_size) + payload = { + "schema_version": 1, + "guard_evidence": evidence, + "ranks": [ + { + "rank": rank.rank, + "world_size": rank.world_size, + "selected_logprobs": rank.selected_logprobs, + "metadata": json_safe(rank.metadata), + } + for rank in ranks + ], + } + temporary = result_path.with_suffix(".tmp") + torch.save(payload, temporary) + os.replace(temporary, result_path) + except BaseException as exc: + error = { + "type": f"{type(exc).__module__}.{type(exc).__qualname__}", + "message": str(exc), + "traceback": traceback.format_exc(), + } + error_path.write_text(json.dumps(error, sort_keys=True), encoding="utf-8") + raise SystemExit(1) from None + + +@contextmanager +def _read_only_scoring_guard( + scorer: PairedScorer, + *, + verify_state: bool, +) -> Iterator[dict[str, Any]]: + model = scorer_model(scorer) + if verify_state and getattr(scorer, "optimizer", None) is not None: + raise ValueError("scorer must not own an active optimizer") + if model is None: + yield { + "model_state_verified": False, + "model_eval": False, + "no_grad": True, + "optimizer_step": False, + } + return + + modes = tuple((module, module.training) for module in model.modules()) + snapshot = _module_tensor_snapshot(model) if verify_state else None + model.eval() + evidence = { + "model_state_verified": verify_state, + "model_eval": True, + "no_grad": True, + "optimizer_step": False, + "model_modes_restored": False, + "model_state_unchanged": False if verify_state else None, + } + try: + yield evidence + finally: + for module, was_training in modes: + module.training = was_training + evidence["model_modes_restored"] = True + if snapshot is not None: + mutations = _module_state_mutations(model, snapshot) + if mutations: + raise RuntimeError( + "read-only scorer mutated model parameters/buffers: " + ", ".join(mutations) + ) + evidence["model_state_unchanged"] = True + + +def _module_tensor_snapshot(model: torch.nn.Module) -> dict[str, torch.Tensor]: + values = { + f"parameter:{name}": tensor.detach().to(device="cpu").clone() + for name, tensor in model.named_parameters() + } + values.update( + { + f"buffer:{name}": tensor.detach().to(device="cpu").clone() + for name, tensor in model.named_buffers() + } + ) + return values + + +def _module_state_mutations( + model: torch.nn.Module, + before: Mapping[str, torch.Tensor], +) -> list[str]: + after = _module_tensor_snapshot(model) + mutations: list[str] = [] + for name in sorted(set(before) | set(after)): + left = before.get(name) + right = after.get(name) + if left is None or right is None: + mutations.append(name) + continue + if left.dtype != right.dtype or left.shape != right.shape or not torch.equal(left, right): + mutations.append(name) + return mutations + + +def scorer_model(scorer: PairedScorer) -> Optional[torch.nn.Module]: + if isinstance(scorer, torch.nn.Module): + return scorer + candidate = getattr(scorer, "model", None) + return candidate if isinstance(candidate, torch.nn.Module) else None + + +def paired_model_state_fingerprints( + rollout_scorer: PairedScorer, + training_scorer: PairedScorer, +) -> dict[str, Optional[str]]: + fingerprints = { + "rollout": _scorer_model_state_fingerprint(rollout_scorer), + "training": _scorer_model_state_fingerprint(training_scorer), + } + if fingerprints["rollout"] is None or fingerprints["training"] is None: + raise ScorerIdentityError( + "rollout and training model state fingerprints must both be observable" + ) + if fingerprints["rollout"] != fingerprints["training"]: + raise ScorerIdentityError( + "rollout and training model state fingerprints differ before scoring" + ) + return fingerprints + + +def _scorer_model_state_fingerprint(scorer: PairedScorer) -> Optional[str]: + declared = getattr(scorer, "model_state_fingerprint", None) + if declared is not None and (not isinstance(declared, str) or not declared): + raise ScorerIdentityError("scorer model_state_fingerprint must be a non-empty string") + model = scorer_model(scorer) + if model is None: + return declared + observed_model = _module_state_fingerprint(model) + if declared is not None and declared != observed_model: + raise ScorerIdentityError( + "declared scorer model_state_fingerprint does not match observed model state" + ) + return observed_model + + +def scorer_implementation_fingerprint(scorer: PairedScorer) -> str: + declared = getattr(scorer, "implementation_fingerprint", None) + if declared is not None and (not isinstance(declared, str) or not declared): + raise ScorerIdentityError("scorer implementation_fingerprint must be a non-empty string") + scorer_type = f"{type(scorer).__module__}.{type(scorer).__qualname__}" + score_source = _source_text(getattr(type(scorer), "score", None)) + return canonical_fingerprint( + { + "declared_implementation": declared, + "scorer_type": scorer_type, + "score_source_fingerprint": hashlib.sha256(score_source.encode("utf-8")).hexdigest(), + } + ) + + +def _module_state_fingerprint(model: torch.nn.Module) -> str: + digest = hashlib.sha256() + digest.update(f"{type(model).__module__}.{type(model).__qualname__}".encode("utf-8")) + digest.update(_source_text(type(model)).encode("utf-8")) + tensors = tuple( + (f"parameter:{name}", tensor) for name, tensor in model.named_parameters() + ) + tuple((f"buffer:{name}", tensor) for name, tensor in model.named_buffers()) + for name, tensor in tensors: + snapshot = tensor.detach().to(device="cpu") + if snapshot.is_sparse: + snapshot = snapshot.to_dense() + snapshot = snapshot.contiguous() + digest.update(name.encode("utf-8")) + digest.update(str(snapshot.dtype).encode("utf-8")) + digest.update(str(tuple(snapshot.shape)).encode("utf-8")) + digest.update(snapshot.reshape(-1).view(torch.uint8).numpy().tobytes()) + return digest.hexdigest() + + +def _source_text(value: Any) -> str: + try: + return inspect.getsource(value) + except (OSError, TypeError): + return repr(value) + + +def _coerce_rank_scores( + output: torch.Tensor | RankScore | Sequence[RankScore], + expected_world_size: int, +) -> tuple[RankScore, ...]: + if isinstance(output, torch.Tensor): + if expected_world_size != 1: + raise RankCompletenessError( + "a bare tensor result is valid only for a world_size=1 scorer" + ) + return (RankScore(rank=0, world_size=1, selected_logprobs=output),) + if isinstance(output, RankScore): + return (output,) + if not isinstance(output, Sequence) or isinstance(output, (str, bytes)): + raise TypeError("scorer must return a tensor, RankScore, or sequence of RankScore") + values = tuple(output) + if not all(isinstance(value, RankScore) for value in values): + raise TypeError("every scorer sequence item must be a RankScore") + return values + + +def _load_child_result(path: Path) -> Mapping[str, Any]: + if not path.is_file(): + raise ChildScoringError(f"scoring child produced no result artifact: {path.name}") + try: + payload = torch.load(path, map_location="cpu", weights_only=True) + except Exception as exc: + raise ChildScoringError(f"failed to load scoring child result {path.name}: {exc}") from exc + if not isinstance(payload, Mapping) or payload.get("schema_version") != 1: + raise ChildScoringError(f"malformed scoring child result: {path.name}") + return payload + + +def validate_rank_outputs( + payload: Mapping[str, Any], + spec: ScorerSpec, + *, + expected_shape: torch.Size, + target: str, +) -> dict[int, RankScore]: + raw_ranks = payload.get("ranks") + if not isinstance(raw_ranks, Sequence): + raise RankCompletenessError(f"{target} child result has no rank sequence") + ranks: dict[int, RankScore] = {} + duplicates: list[int] = [] + for raw in raw_ranks: + if not isinstance(raw, Mapping): + raise RankCompletenessError(f"{target} rank result must be a mapping") + rank_score = RankScore( + rank=int(raw["rank"]), + world_size=int(raw["world_size"]), + selected_logprobs=raw["selected_logprobs"], + metadata=raw.get("metadata", {}), + ) + if rank_score.rank in ranks: + duplicates.append(rank_score.rank) + ranks[rank_score.rank] = rank_score + if duplicates: + raise RankCompletenessError(f"{target} returned duplicate ranks: {sorted(set(duplicates))}") + expected = set(range(spec.world_size)) + actual = set(ranks) + if actual != expected: + missing = sorted(expected - actual) + unexpected = sorted(actual - expected) + raise RankCompletenessError( + f"{target} rank set is incomplete; missing={missing}, unexpected={unexpected}" + ) + for rank_index, value in ranks.items(): + if value.world_size != spec.world_size: + raise RankCompletenessError( + f"{target} rank {rank_index} reported world_size={value.world_size}, " + f"expected {spec.world_size}" + ) + if value.selected_logprobs.shape != expected_shape: + raise RankCompletenessError( + f"{target} rank {rank_index} selected_logprobs shape " + f"{tuple(value.selected_logprobs.shape)} does not match " + f"canonical shape {tuple(expected_shape)}" + ) + expected_dtype = torch_dtype(spec.dtype) + if ( + not value.selected_logprobs.is_floating_point() + or value.selected_logprobs.dtype != expected_dtype + ): + raise RankCompletenessError( + f"{target} rank {rank_index} selected_logprobs dtype " + f"{value.selected_logprobs.dtype} does not match scorer dtype {expected_dtype}" + ) + rank_zero = ranks[0].selected_logprobs + for rank_index, value in ranks.items(): + if rank_index == 0: + continue + if value.selected_logprobs.dtype != rank_zero.dtype or not torch.equal( + value.selected_logprobs, + rank_zero, + ): + raise RankCompletenessError( + f"{target} rank {rank_index} selected_logprobs diverge from rank 0" + ) + return ranks + + +def scorer_spec(scorer: PairedScorer, expected_side: ScoreSide) -> ScorerSpec: + spec = getattr(scorer, "spec", None) + if not isinstance(spec, ScorerSpec): + raise TypeError("paired scorer must expose a ScorerSpec as .spec") + if spec.side is not expected_side: + raise ValueError(f"scorer side {spec.side.value!r} does not match {expected_side.value!r}") + model = scorer_model(scorer) + if model is not None: + _require_module_on_device(model, device_type(spec.device)) + _require_module_float_dtype(model, torch_dtype(spec.dtype)) + return spec + + +def validate_scorer_identity( + specs: Mapping[str, ScorerSpec], + batch: CanonicalScoringBatch, +) -> None: + identity = batch.identity + expected = { + "checkpoint_id": identity.checkpoint_id, + "model_version": identity.model_version, + "pre_update_state": identity.pre_update_state, + } + for target in ("rollout", "training"): + observed = specs[target].construction_options + mismatches = [key for key, value in expected.items() if observed.get(key) != value] + if mismatches: + raise ScorerIdentityError( + f"{target} scorer construction identity differs from canonical identity: " + + ", ".join(mismatches) + ) + + +def _child_error_detail(path: Path) -> str: + try: + value = _read_json_object(path) + except (OSError, ValueError, json.JSONDecodeError): + return "child did not publish structured error evidence" + return f"{value.get('type', 'error')}: {value.get('message', '')}" + + +def _read_json_object(path: Path) -> dict[str, Any]: + value = strict_json_loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object: {path}") + return value + + +def _stop_processes(processes: Iterator[mp.Process] | Sequence[mp.Process]) -> None: + values = tuple(processes) + for process in values: + if process.is_alive(): + process.terminate() + for process in values: + if process.pid is not None: + process.join(timeout=1.0) + for process in values: + if process.is_alive() and hasattr(process, "kill"): + process.kill() + process.join(timeout=1.0) + + +def _require_module_on_device(model: torch.nn.Module, expected: str) -> None: + for name, tensor in tuple(model.named_parameters()) + tuple(model.named_buffers()): + if tensor.device.type != expected: + raise ValueError( + f"scorer model tensor {name!r} is on {tensor.device}; expected {expected}" + ) + + +def _require_module_float_dtype(model: torch.nn.Module, expected: torch.dtype) -> None: + mismatches = [ + f"{name}={tensor.dtype}" + for name, tensor in tuple(model.named_parameters()) + tuple(model.named_buffers()) + if tensor.is_floating_point() and tensor.dtype != expected + ] + if mismatches: + raise ValueError( + f"scorer floating model state must use {expected}: " + ", ".join(mismatches) + ) + + +def device_type(value: str) -> str: + try: + return torch.device(value).type + except (TypeError, RuntimeError) as exc: + raise ValueError(f"invalid scorer device: {value!r}") from exc + + +def normalized_dtype(value: str) -> str: + dtype = torch_dtype(value) + return str(dtype).removeprefix("torch.") + + +def torch_dtype(value: str) -> torch.dtype: + normalized = str(value).strip().lower().replace("torch.", "") + dtypes = { + "float32": torch.float32, + "fp32": torch.float32, + "float16": torch.float16, + "fp16": torch.float16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + } + try: + return dtypes[normalized] + except KeyError as exc: + raise ValueError(f"unsupported stateless scorer dtype: {value!r}") from exc + + +def json_safe(value: Any) -> Any: + if isinstance(value, Enum): + return value.value + if isinstance(value, Mapping): + return {str(key): json_safe(item) for key, item in value.items()} + if isinstance(value, (set, frozenset, tuple, list)): + items = [json_safe(item) for item in value] + if isinstance(value, (set, frozenset)): + return sorted(items, key=lambda item: json.dumps(item, sort_keys=True)) + return items + if isinstance(value, (str, int, float, bool)) or value is None: + return value + return repr(value) + + +def canonical_fingerprint(value: Any) -> str: + serialized = json.dumps( + json_safe(value), + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ) + return hashlib.sha256(serialized.encode("utf-8")).hexdigest() diff --git a/rl_engine/alignment/cross_config/artifacts.py b/rl_engine/alignment/cross_config/artifacts.py new file mode 100644 index 00000000..296c2809 --- /dev/null +++ b/rl_engine/alignment/cross_config/artifacts.py @@ -0,0 +1,508 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Append-only, crash-safe artifacts for cross-configuration runs.""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from pathlib import Path +from typing import Any, Iterable, Mapping, Optional + +import torch + +from rl_engine.alignment.cross_config._json import strict_json_loads + +REQUIRED_CASE_ARTIFACTS = frozenset( + { + "requested.json", + "materialized.json", + "actual.json", + "identity.json", + "score_rollout.pt", + "score_training.pt", + "comparison.json", + "token_diffs.pt", + } +) +_JSON_SCHEMAS = { + "requested.json": "cross_config.requested.v1", + "materialized.json": "cross_config.materialized_envelope.v1", + "actual.json": "cross_config.actual.v1", + "identity.json": "cross_config.identity_envelope.v1", + "comparison.json": "cross_config.alignment_result.v1", +} +_JSON_REQUIRED_KEYS = { + "requested.json": frozenset({"case"}), + "materialized.json": frozenset({"materialized_case"}), + "actual.json": frozenset({"rollout", "training"}), + "identity.json": frozenset({"identity"}), + "comparison.json": frozenset({"status", "comparable", "passed"}), +} +_TENSOR_REQUIRED_KEYS = { + "score_rollout.pt": frozenset({"selected_logprobs", "active_mask"}), + "score_training.pt": frozenset({"selected_logprobs", "active_mask"}), + "token_diffs.pt": frozenset( + { + "rollout_logprobs", + "training_logprobs", + "active_mask", + "absolute_diff", + "mismatch_mask", + } + ), +} + + +class ArtifactError(RuntimeError): + """Raised when an artifact is incomplete, malformed, or would be overwritten.""" + + +class ArtifactStore: + """Persist immutable attempt directories and atomically mark completed attempts.""" + + def __init__(self, root: str | Path): + self.root = Path(root) + + def experiment_dir(self, experiment_id: str) -> Path: + return self.root / _safe_component(experiment_id, "experiment_id") + + def initialize_experiment( + self, + experiment_id: str, + *, + experiment: Mapping[str, Any], + plan: Iterable[Mapping[str, Any]], + ) -> Path: + """Create immutable experiment metadata, or verify an identical resume target.""" + + directory = self.experiment_dir(experiment_id) + directory.mkdir(parents=True, exist_ok=True) + self._write_or_verify_json(directory / "experiment.json", experiment) + plan_text = "".join(_canonical_json(item) + "\n" for item in plan) + self._write_or_verify_text(directory / "plan.jsonl", plan_text) + return directory + + def create_attempt( + self, + experiment_id: str, + case_id: str, + *, + attempt_id: Optional[str] = None, + ) -> Path: + """Allocate an append-only attempt directory for a case.""" + + case_dir = ( + self.experiment_dir(experiment_id) / "cases" / _safe_component(case_id, "case_id") + ) + case_dir.mkdir(parents=True, exist_ok=True) + if attempt_id is not None: + attempt_dir = case_dir / _safe_component(attempt_id, "attempt_id") + try: + attempt_dir.mkdir() + except FileExistsError as exc: + raise ArtifactError(f"attempt already exists: {attempt_dir}") from exc + _fsync_directory(case_dir) + return attempt_dir + + # Another controller can win after _next_attempt_id() observes the + # directory. mkdir is the atomic allocator; retry rather than aliasing or + # overwriting the winning attempt. + while True: + resolved_attempt_id = self._next_attempt_id(case_dir) + attempt_dir = case_dir / resolved_attempt_id + try: + attempt_dir.mkdir() + except FileExistsError: + continue + _fsync_directory(case_dir) + return attempt_dir + + def write_json(self, attempt_dir: str | Path, name: str, value: Mapping[str, Any]) -> Path: + path = self._attempt_path(attempt_dir, name, suffix=".json") + self._write_new_text(path, _canonical_json(value) + "\n") + return path + + def write_tensor_bundle( + self, + attempt_dir: str | Path, + name: str, + tensors: Mapping[str, torch.Tensor], + *, + metadata: Optional[Mapping[str, Any]] = None, + ) -> Path: + """Write CPU tensor payloads that can be loaded with ``weights_only=True``.""" + + path = self._attempt_path(attempt_dir, name, suffix=".pt") + payload: dict[str, Any] = { + "schema_version": 1, + "tensors": { + key: tensor.detach().to(device="cpu").contiguous() + for key, tensor in tensors.items() + }, + "metadata": dict(metadata or {}), + } + self._atomic_torch_save(path, payload) + return path + + def load_tensor_bundle(self, path: str | Path) -> dict[str, Any]: + try: + payload = torch.load(Path(path), map_location="cpu", weights_only=True) + except Exception as exc: + raise ArtifactError(f"failed to load tensor artifact {path}: {exc}") from exc + if not isinstance(payload, dict) or payload.get("schema_version") != 1: + raise ArtifactError(f"unsupported tensor artifact schema: {path}") + tensors = payload.get("tensors") + if not isinstance(tensors, dict) or not all( + isinstance(value, torch.Tensor) for value in tensors.values() + ): + raise ArtifactError(f"malformed tensor payload: {path}") + return payload + + def complete_attempt( + self, + attempt_dir: str | Path, + *, + summary: Mapping[str, Any], + required: Iterable[str] = REQUIRED_CASE_ARTIFACTS, + ) -> Path: + """Validate all payloads before publishing an atomic ``COMPLETE`` marker.""" + + directory = Path(attempt_dir) + required_names = frozenset(required) + missing = sorted(name for name in required_names if not (directory / name).is_file()) + if missing: + raise ArtifactError(f"cannot complete {directory}; missing artifacts: {missing}") + self._validate_machine_artifacts(directory) + marker_value = dict(summary) + marker_value["artifact_sha256"] = { + name: _sha256_file(directory / name) for name in sorted(required_names) + } + self._validate_complete_summary(directory, marker_value) + marker = directory / "COMPLETE" + self._write_new_text(marker, _canonical_json(marker_value) + "\n") + return marker + + def completed_attempt( + self, + experiment_id: str, + case_id: str, + *, + required: Iterable[str] = REQUIRED_CASE_ARTIFACTS, + ) -> Optional[Path]: + """Return the newest valid completed attempt, ignoring partial attempts.""" + + case_dir = ( + self.experiment_dir(experiment_id) / "cases" / _safe_component(case_id, "case_id") + ) + if not case_dir.is_dir(): + return None + for attempt_dir in sorted( + case_dir.iterdir(), + key=_attempt_sort_key, + reverse=True, + ): + if not attempt_dir.is_dir() or not (attempt_dir / "COMPLETE").is_file(): + continue + try: + self.validate_completed_attempt( + attempt_dir, + required=required, + expected_case_id=case_id, + ) + except ArtifactError: + continue + return attempt_dir + return None + + def validate_completed_attempt( + self, + attempt_dir: str | Path, + *, + required: Iterable[str] = REQUIRED_CASE_ARTIFACTS, + expected_case_id: Optional[str] = None, + ) -> None: + directory = Path(attempt_dir) + required_names = frozenset(required) + marker = directory / "COMPLETE" + if not marker.is_file(): + raise ArtifactError(f"missing COMPLETE marker: {directory}") + try: + marker_value = strict_json_loads(marker.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError, ValueError) as exc: + raise ArtifactError(f"malformed COMPLETE marker: {marker}") from exc + if not isinstance(marker_value, dict): + raise ArtifactError(f"COMPLETE marker must contain a JSON object: {marker}") + self._validate_complete_summary( + directory, + marker_value, + expected_case_id=expected_case_id, + ) + missing = sorted(name for name in required_names if not (directory / name).is_file()) + if missing: + raise ArtifactError(f"completed attempt is missing artifacts: {missing}") + self._validate_artifact_hashes(directory, marker_value, required_names) + self._validate_machine_artifacts(directory, expected_case_id=expected_case_id) + + @staticmethod + def _validate_artifact_hashes( + directory: Path, + marker: Mapping[str, Any], + required: frozenset[str], + ) -> None: + recorded = marker.get("artifact_sha256") + if not isinstance(recorded, Mapping) or set(recorded) != set(required): + raise ArtifactError(f"COMPLETE marker has invalid artifact hashes: {directory}") + for name in sorted(required): + expected = recorded.get(name) + if not isinstance(expected, str) or expected != _sha256_file(directory / name): + raise ArtifactError(f"artifact hash does not match COMPLETE: {directory / name}") + + def _validate_machine_artifacts( + self, + directory: Path, + *, + expected_case_id: Optional[str] = None, + ) -> None: + tensor_payloads: dict[str, dict[str, Any]] = {} + for name in ("score_rollout.pt", "score_training.pt", "token_diffs.pt"): + path = directory / name + if path.exists(): + payload = self.load_tensor_bundle(path) + tensor_payloads[name] = payload + tensors = payload["tensors"] + missing_tensor_keys = sorted(_TENSOR_REQUIRED_KEYS[name].difference(tensors)) + if missing_tensor_keys: + raise ArtifactError( + f"tensor artifact {name} is missing keys: {missing_tensor_keys}" + ) + metadata = payload.get("metadata", {}) + if not isinstance(metadata, Mapping): + raise ArtifactError(f"tensor artifact metadata must be an object: {path}") + if expected_case_id is not None and metadata.get("case_id") != expected_case_id: + raise ArtifactError( + f"tensor artifact case_id does not match {expected_case_id!r}: {path}" + ) + if metadata.get("attempt_id") != directory.name: + raise ArtifactError( + f"tensor artifact attempt_id does not match {directory.name!r}: {path}" + ) + json_payloads: dict[str, dict[str, Any]] = {} + for name in ( + "requested.json", + "materialized.json", + "actual.json", + "identity.json", + "comparison.json", + ): + path = directory / name + if not path.exists(): + continue + try: + value = strict_json_loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError, ValueError) as exc: + raise ArtifactError(f"malformed JSON artifact: {path}") from exc + if not isinstance(value, dict): + raise ArtifactError(f"JSON artifact must contain an object: {path}") + json_payloads[name] = value + if expected_case_id is not None and value.get("case_id") != expected_case_id: + raise ArtifactError( + f"JSON artifact case_id does not match {expected_case_id!r}: {path}" + ) + if value.get("schema_version") != _JSON_SCHEMAS[name]: + raise ArtifactError(f"JSON artifact has an unsupported schema: {path}") + missing_json_keys = sorted(_JSON_REQUIRED_KEYS[name].difference(value)) + if missing_json_keys: + raise ArtifactError(f"JSON artifact {name} is missing keys: {missing_json_keys}") + if value.get("attempt_id") != directory.name: + raise ArtifactError( + f"JSON artifact attempt_id does not match {directory.name!r}: {path}" + ) + + if expected_case_id is None and json_payloads: + case_ids = {payload.get("case_id") for payload in json_payloads.values()} + if len(case_ids) != 1 or None in case_ids: + raise ArtifactError("JSON artifacts must declare one consistent case_id") + inferred_case_id = next(iter(case_ids)) + for name, payload in tensor_payloads.items(): + if payload["metadata"].get("case_id") != inferred_case_id: + raise ArtifactError( + f"tensor artifact case_id does not match {inferred_case_id!r}: " + f"{directory / name}" + ) + + @staticmethod + def _validate_complete_summary( + directory: Path, + summary: Mapping[str, Any], + *, + expected_case_id: Optional[str] = None, + ) -> None: + if summary.get("schema_version") != "cross_config.complete.v1": + raise ArtifactError(f"COMPLETE marker has an unsupported schema: {directory}") + case_id = summary.get("case_id") + if not isinstance(case_id, str) or not case_id: + raise ArtifactError(f"COMPLETE marker is missing case_id: {directory}") + if expected_case_id is not None and case_id != expected_case_id: + raise ArtifactError( + f"COMPLETE marker case_id does not match {expected_case_id!r}: {directory}" + ) + if summary.get("attempt_id") != directory.name: + raise ArtifactError( + f"COMPLETE marker attempt_id does not match {directory.name!r}: {directory}" + ) + if not isinstance(summary.get("status"), str): + raise ArtifactError(f"COMPLETE marker is missing status: {directory}") + if not isinstance(summary.get("artifact_sha256"), Mapping): + raise ArtifactError(f"COMPLETE marker is missing artifact hashes: {directory}") + + def _write_or_verify_json(self, path: Path, value: Mapping[str, Any]) -> None: + self._write_or_verify_text(path, _canonical_json(value) + "\n") + + def _write_or_verify_text(self, path: Path, text: str) -> None: + if path.exists(): + self._verify_existing_text(path, text) + return + try: + self._atomic_write_text(path, text) + except ArtifactError: + # A concurrent writer may have atomically published the same immutable + # experiment metadata. Accept only byte-identical content. + if not path.exists(): + raise + self._verify_existing_text(path, text) + + @staticmethod + def _verify_existing_text(path: Path, text: str) -> None: + try: + existing = path.read_text(encoding="utf-8") + except OSError as exc: + raise ArtifactError(f"failed to read existing artifact {path}: {exc}") from exc + if existing != text: + raise ArtifactError(f"resume metadata differs from existing artifact: {path}") + + def _write_new_text(self, path: Path, text: str) -> None: + if path.exists(): + raise ArtifactError(f"refusing to overwrite artifact: {path}") + self._atomic_write_text(path, text) + + def _atomic_write_text(self, path: Path, text: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + temporary = Path(temporary_name) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + handle.write(text) + handle.flush() + os.fsync(handle.fileno()) + _publish_new_file(temporary, path) + except Exception: + temporary.unlink(missing_ok=True) + raise + + def _atomic_torch_save(self, path: Path, payload: Mapping[str, Any]) -> None: + if path.exists(): + raise ArtifactError(f"refusing to overwrite artifact: {path}") + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + os.close(fd) + temporary = Path(temporary_name) + try: + torch.save(dict(payload), temporary) + with temporary.open("rb") as handle: + os.fsync(handle.fileno()) + _publish_new_file(temporary, path) + except Exception: + temporary.unlink(missing_ok=True) + raise + + @staticmethod + def _next_attempt_id(case_dir: Path) -> str: + indices: list[int] = [] + for child in case_dir.iterdir(): + if not child.is_dir() or not child.name.startswith("attempt-"): + continue + suffix = child.name.removeprefix("attempt-") + if suffix.isdigit(): + indices.append(int(suffix)) + return f"attempt-{max(indices, default=0) + 1:04d}" + + @staticmethod + def _attempt_path(attempt_dir: str | Path, name: str, *, suffix: str) -> Path: + directory = Path(attempt_dir) + safe_name = _safe_component(name, "artifact name") + if not safe_name.endswith(suffix): + safe_name += suffix + return directory / safe_name + + +def _canonical_json(value: Mapping[str, Any]) -> str: + try: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + except (TypeError, ValueError) as exc: + raise ArtifactError(f"artifact is not strict JSON: {exc}") from exc + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + try: + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + except OSError as exc: + raise ArtifactError(f"failed to hash artifact {path}: {exc}") from exc + return digest.hexdigest() + + +def _safe_component(value: str, label: str) -> str: + if not value or value in {".", ".."} or Path(value).name != value: + raise ValueError(f"{label} must be a single non-empty path component") + return value + + +def _attempt_sort_key(path: Path) -> tuple[int, int, str]: + """Order standard attempt IDs numerically and retain nonstandard fallbacks.""" + + prefix = "attempt-" + suffix = path.name.removeprefix(prefix) + if path.name.startswith(prefix) and suffix.isdigit(): + return (1, int(suffix), path.name) + return (0, -1, path.name) + + +def _publish_new_file(temporary: Path, destination: Path) -> None: + """Atomically publish without ever replacing an existing artifact.""" + + try: + os.link(temporary, destination) + except FileExistsError as exc: + raise ArtifactError(f"refusing to overwrite artifact: {destination}") from exc + temporary.unlink() + _fsync_directory(destination.parent) + + +def _fsync_directory(directory: Path) -> None: + """Persist directory entry changes where the host filesystem supports it.""" + + try: + descriptor = os.open(directory, os.O_RDONLY) + except OSError: + return + try: + os.fsync(descriptor) + except OSError: + pass + finally: + os.close(descriptor) + + +__all__ = ["ArtifactError", "ArtifactStore", "REQUIRED_CASE_ARTIFACTS"] From 967416ceb8629c121f5339ebd40b8aa43ae30670 Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sun, 19 Jul 2026 08:33:36 +0800 Subject: [PATCH 05/33] feat(alignment): add provenance-aware paired execution --- rl_engine/alignment/cross_config/__init__.py | 40 + .../alignment/cross_config/_provenance.py | 327 ++++++++ rl_engine/alignment/cross_config/_resume.py | 397 ++++++++++ rl_engine/alignment/cross_config/runner.py | 714 ++++++++++++++++++ 4 files changed, 1478 insertions(+) create mode 100644 rl_engine/alignment/cross_config/__init__.py create mode 100644 rl_engine/alignment/cross_config/_provenance.py create mode 100644 rl_engine/alignment/cross_config/_resume.py create mode 100644 rl_engine/alignment/cross_config/runner.py diff --git a/rl_engine/alignment/cross_config/__init__.py b/rl_engine/alignment/cross_config/__init__.py new file mode 100644 index 00000000..26a28e8d --- /dev/null +++ b/rl_engine/alignment/cross_config/__init__.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plan and run cross-configuration alignment experiments. + +The package root is intentionally small and lazily loads execution code. Extension +authors import adapter, artifact, operator, or schema details from their owning +submodule. +""" + +from importlib import import_module +from typing import Any + +from rl_engine.alignment.cross_config.comparison import compare_score_artifacts +from rl_engine.alignment.cross_config.config import ExperimentConfig, load_config +from rl_engine.alignment.cross_config.execution_plan import ExecutionPlan, build_execution_plan +from rl_engine.alignment.cross_config.planner import ExperimentPlan, Planner + + +def __getattr__(name: str) -> Any: + if name in {"PairedRunResult", "PairedRunner"}: + return getattr(import_module("rl_engine.alignment.cross_config.runner"), name) + if name in {"RuntimeMaterializer", "RuntimeTools"}: + return getattr(import_module("rl_engine.alignment.cross_config.runtime"), name) + raise AttributeError(name) + + +__all__ = [ + "ExperimentConfig", + "ExperimentPlan", + "ExecutionPlan", + "PairedRunResult", + "PairedRunner", + "Planner", + "RuntimeMaterializer", + "RuntimeTools", + "build_execution_plan", + "compare_score_artifacts", + "load_config", +] diff --git a/rl_engine/alignment/cross_config/_provenance.py b/rl_engine/alignment/cross_config/_provenance.py new file mode 100644 index 00000000..05817b17 --- /dev/null +++ b/rl_engine/alignment/cross_config/_provenance.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Private execution identity and provenance construction.""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import platform +from dataclasses import replace +from pathlib import Path +from typing import Any, Mapping, Optional + +import torch + +from rl_engine.alignment.cross_config._execution import ( + OperatorExecutionError, + PairedRunnerError, + canonical_fingerprint, + device_type, + json_safe, +) +from rl_engine.alignment.cross_config.runtime import RuntimeMaterialization +from rl_engine.alignment.cross_config.schema import ( + MaterializationStatus, + RuntimeProvenance, + ScoreArtifact, + ScorerSpec, + ScoreSide, +) +from rl_engine.kernels.semantic_registry import OperatorInstanceProvenance, OperatorResolution + +PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT = "cross_config.paired_runner.v2" + + +def effective_runtime_status( + materialization: RuntimeMaterialization, +) -> MaterializationStatus: + """Aggregate runtime status after exact resolution supersedes logp readback.""" + + statuses = [ + application.status + for application in materialization.applications + if not ( + application.path == "logp.backend" + and application.status is MaterializationStatus.UNOBSERVABLE + ) + ] + precedence = ( + MaterializationStatus.ERROR, + MaterializationStatus.UNSUPPORTED, + MaterializationStatus.UNOBSERVABLE, + MaterializationStatus.FALLBACK, + MaterializationStatus.APPLIED, + ) + return next( + (status for status in precedence if status in statuses), + MaterializationStatus.APPLIED, + ) + + +def side_provenance( + base: RuntimeProvenance, + resolution: OperatorResolution, + instance: OperatorInstanceProvenance, + child_payload: Mapping[str, Any], + spec: ScorerSpec, + *, + status: MaterializationStatus, + factory_options: Mapping[str, Any], + model_state_fingerprint: Optional[str], + scorer_implementation_fingerprint: str, +) -> RuntimeProvenance: + payload = base.to_dict() + actual = dict(payload["actual"]) + actual["operators"] = { + "selected_logprob": { + "backend_id": instance.backend_id, + "descriptor_fingerprint": instance.descriptor_fingerprint, + "implementation_fingerprint": instance.implementation_fingerprint, + "instance_fingerprint": instance.instance_fingerprint, + "concrete_implementation": instance.concrete_implementation, + "factory_options": json_safe(factory_options), + "factory_options_fingerprint": factory_options_fingerprint(factory_options), + } + } + actual["model_state_fingerprint"] = model_state_fingerprint + actual["scorer_implementation_fingerprint"] = scorer_implementation_fingerprint + evidence = dict(payload["evidence"]) + evidence.update( + { + "operator_resolution": resolution.to_dict(), + "operator_instance": instance.to_dict(), + "operator_factory_options": json_safe(factory_options), + "scoring_guard": json_safe(child_payload.get("guard_evidence", {})), + "rank_metadata": [ + json_safe(rank.get("metadata", {})) + for rank in child_payload.get("ranks", ()) + if isinstance(rank, Mapping) + ], + "model_state_fingerprint": model_state_fingerprint, + "scorer_implementation_fingerprint": scorer_implementation_fingerprint, + } + ) + implementation_fingerprint = hashlib.sha256( + f"{base.implementation_fingerprint}:{instance.instance_fingerprint}".encode("utf-8") + ).hexdigest() + return RuntimeProvenance( + requested=payload["requested"], + normalized=payload["normalized"], + materialized=payload["materialized"], + actual=actual, + status=status, + construction_fingerprint=base.construction_fingerprint, + distributed_context_fingerprint=base.distributed_context_fingerprint, + process_fingerprint=base.process_fingerprint, + implementation_fingerprint=implementation_fingerprint, + evidence=evidence, + rank=0, + world_size=spec.world_size, + ) + + +def concrete_scorer_spec( + spec: ScorerSpec, + instance: OperatorInstanceProvenance, +) -> ScorerSpec: + overrides = dict(spec.operator_overrides) + overrides["selected_logprob"] = instance.backend_id + return replace(spec, operator_overrides=overrides) + + +def score_metadata(artifact: ScoreArtifact) -> dict[str, Any]: + value = artifact.to_dict() + value.pop("selected_logprobs", None) + value.pop("active_mask", None) + return { + "case_id": artifact.case_id, + "attempt_id": artifact.attempt_id, + "side": artifact.side.value, + "score_artifact": value, + } + + +def execution_fingerprint( + materialization: RuntimeMaterialization, + *, + specs: Mapping[str, ScorerSpec], + instance_provenance: Mapping[str, OperatorInstanceProvenance], + operator_factory_options: Optional[Mapping[str | ScoreSide, Mapping[str, Any]]], + model_state_fingerprints: Mapping[str, Optional[str]], + scorer_implementation_fingerprints: Mapping[str, str], + environment: Mapping[str, Any], +) -> str: + payload = { + "schema_version": "cross_config.execution_identity.v1", + "runner_implementation_fingerprint": PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT, + "environment": environment, + "materialized_case": materialization.materialized_case.to_dict(), + "runtime_provenance": materialization.provenance.to_dict(), + "runtime_binding": materialization.binding.to_dict(), + "applications": [application.to_dict() for application in materialization.applications], + "targets": { + target: { + "scorer": concrete_scorer_spec( + specs[target], + instance_provenance[target], + ).to_dict(), + "operator_instance": instance_provenance[target].to_dict(), + "operator_factory_options": json_safe( + target_factory_options(operator_factory_options, target) + ), + "model_state_fingerprint": model_state_fingerprints[target], + "scorer_implementation_fingerprint": (scorer_implementation_fingerprints[target]), + } + for target in ("rollout", "training") + }, + } + return canonical_fingerprint(payload) + + +def mapping_target(mapping: Mapping[str | ScoreSide, Any], target: str) -> Any: + if target in mapping: + return mapping[target] + side = ScoreSide(target) + return mapping.get(side) + + +def target_factory_options( + options: Optional[Mapping[str | ScoreSide, Mapping[str, Any]]], + target: str, +) -> Mapping[str, Any]: + if options is None: + return {} + value = mapping_target(options, target) + if value is None: + return {} + if not isinstance(value, Mapping): + raise OperatorExecutionError(f"{target} operator factory options must be a mapping") + return dict(value) + + +def factory_options_fingerprint(options: Mapping[str, Any]) -> str: + return canonical_fingerprint(options) + + +def runtime_adapter_fingerprint(materialization: RuntimeMaterialization) -> str: + observed = materialization.provenance.evidence.get("adapter_implementation_fingerprint") + if isinstance(observed, str) and observed: + return observed + return materialization.provenance.implementation_fingerprint + + +def execution_environment_provenance( + specs: Mapping[str, ScorerSpec], + *, + runtime_adapter_fingerprint: str, + operator_implementation_fingerprints: Mapping[str, str], +) -> dict[str, Any]: + source_root = Path(__file__).resolve().parents[3] + try: + package_version = importlib.metadata.version("rl-kernel") + except importlib.metadata.PackageNotFoundError: + package_version = None + torch_config = torch.__config__.show() + execution_devices = {target: device_type(spec.device) for target, spec in sorted(specs.items())} + return { + "schema_version": "cross_config.environment.v1", + "execution_devices": execution_devices, + "python": { + "implementation": platform.python_implementation(), + "version": platform.python_version(), + }, + "torch": { + "version": str(torch.__version__), + "git_version": getattr(torch.version, "git_version", None), + "cuda_build": getattr(torch.version, "cuda", None), + "hip_build": getattr(torch.version, "hip", None), + "debug_build": bool(getattr(torch.version, "debug", False)), + "config_fingerprint": hashlib.sha256(torch_config.encode("utf-8")).hexdigest(), + }, + "host_runtime": { + "platform": platform.platform(), + "machine": platform.machine(), + "processor": platform.processor(), + "mkldnn_available": bool(torch.backends.mkldnn.is_available()), + "mkl_available": bool(torch.backends.mkl.is_available()), + }, + "rl_kernel": { + "package_version": package_version, + "git_revision": _git_revision(source_root), + "source_tree_fingerprint": _cross_config_source_tree_fingerprint( + source_root, + implementation_fingerprints={ + "runtime_adapter": runtime_adapter_fingerprint, + "operators": dict(operator_implementation_fingerprints), + }, + ), + }, + } + + +def _git_revision(source_root: Path) -> Optional[str]: + git_dir = source_root / ".git" + try: + if git_dir.is_file(): + marker = git_dir.read_text(encoding="utf-8").strip() + if not marker.startswith("gitdir: "): + return None + resolved = Path(marker.removeprefix("gitdir: ")) + git_dir = resolved if resolved.is_absolute() else source_root / resolved + head = (git_dir / "HEAD").read_text(encoding="utf-8").strip() + if not head.startswith("ref: "): + return head or None + reference = head.removeprefix("ref: ") + loose_ref = git_dir / reference + if loose_ref.is_file(): + return loose_ref.read_text(encoding="utf-8").strip() or None + packed_refs = git_dir / "packed-refs" + if packed_refs.is_file(): + suffix = f" {reference}" + for line in packed_refs.read_text(encoding="utf-8").splitlines(): + if line.endswith(suffix): + return line.split(" ", 1)[0] + except OSError: + return None + return None + + +def _cross_config_source_tree_fingerprint( + source_root: Path, + *, + implementation_fingerprints: Mapping[str, Any], +) -> str: + paths = list((source_root / "rl_engine/alignment/cross_config").glob("*.py")) + paths.extend( + source_root / relative + for relative in ( + "rl_engine/executors/stateless_executor.py", + "rl_engine/kernels/gtest/tolerance.py", + "rl_engine/kernels/registry.py", + "rl_engine/kernels/semantic_registry.py", + ) + ) + digest = hashlib.sha256() + for path in sorted(set(paths)): + try: + content = path.read_bytes() + except OSError as exc: + raise PairedRunnerError( + f"cannot fingerprint cross-configuration source file {path}: {exc}" + ) from exc + digest.update(str(path.relative_to(source_root)).encode("utf-8")) + digest.update(b"\0") + digest.update(content) + digest.update(b"\0") + digest.update( + json.dumps( + json_safe(implementation_fingerprints), + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + ) + return digest.hexdigest() diff --git a/rl_engine/alignment/cross_config/_resume.py b/rl_engine/alignment/cross_config/_resume.py new file mode 100644 index 00000000..f3d31e70 --- /dev/null +++ b/rl_engine/alignment/cross_config/_resume.py @@ -0,0 +1,397 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Private validation for append-only attempt resume.""" + +from __future__ import annotations + +import hashlib +import json +import math +from pathlib import Path +from typing import Any, Mapping, Optional + +import torch + +from rl_engine.alignment.cross_config._execution import ( + canonical_fingerprint, + json_safe, + torch_dtype, +) +from rl_engine.alignment.cross_config._json import strict_json_loads +from rl_engine.alignment.cross_config._provenance import ( + PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT, + concrete_scorer_spec, + effective_runtime_status, + factory_options_fingerprint, + target_factory_options, +) +from rl_engine.alignment.cross_config.artifacts import REQUIRED_CASE_ARTIFACTS +from rl_engine.alignment.cross_config.comparison import recompute_mismatch_mask +from rl_engine.alignment.cross_config.runtime import RuntimeMaterialization +from rl_engine.alignment.cross_config.schema import ( + CanonicalScoringBatch, + ExperimentCase, + ScorerSpec, + ScoreSide, +) +from rl_engine.kernels.gtest.tolerance import ( + resolve_logprob_threshold, + tolerance_contract_fingerprint, +) +from rl_engine.kernels.semantic_registry import OperatorInstanceProvenance + +_COMPLETE_KEYS = frozenset( + { + "schema_version", + "case_id", + "attempt_id", + "status", + "comparable", + "passed", + "active_token_count", + "mismatch_count", + "worst_token_index", + "max_abs_diff", + "rollout_backend", + "training_backend", + "execution_fingerprint", + "environment_fingerprint", + "runner_implementation_fingerprint", + "artifact_sha256", + } +) + + +def completed_attempt_matches( + attempt_dir: Path, + case: ExperimentCase, + batch: CanonicalScoringBatch, + *, + materialization: RuntimeMaterialization, + specs: Mapping[str, ScorerSpec], + instance_provenance: Mapping[str, OperatorInstanceProvenance], + operator_factory_options: Optional[Mapping[str | ScoreSide, Mapping[str, Any]]], + model_state_fingerprints: Mapping[str, Optional[str]], + scorer_implementation_fingerprints: Mapping[str, str], + environment: Mapping[str, Any], + execution_fingerprint: str, +) -> bool: + try: + identity = read_json_object(attempt_dir / "identity.json") + requested = read_json_object(attempt_dir / "requested.json") + actual = read_json_object(attempt_dir / "actual.json") + marker = read_json_object(attempt_dir / "COMPLETE") + except (OSError, ValueError, json.JSONDecodeError): + return False + if not ( + set(marker) == _COMPLETE_KEYS + and identity.get("schema_version") == "cross_config.identity_envelope.v1" + and requested.get("schema_version") == "cross_config.requested.v1" + and actual.get("schema_version") == "cross_config.actual.v1" + and marker.get("schema_version") == "cross_config.complete.v1" + and actual.get("runner_implementation_fingerprint") + == PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT + and actual.get("environment") == environment + and actual.get("environment_fingerprint") == canonical_fingerprint(environment) + and marker.get("runner_implementation_fingerprint") + == PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT + and marker.get("environment_fingerprint") == canonical_fingerprint(environment) + and isinstance(marker.get("artifact_sha256"), Mapping) + and set(marker["artifact_sha256"]) == set(REQUIRED_CASE_ARTIFACTS) + ): + return False + if not ( + identity.get("case_id") == case.case_id + and identity.get("identity") == batch.identity.to_dict() + and requested.get("case") == case.to_dict() + and marker.get("execution_fingerprint") == execution_fingerprint + and actual.get("execution_fingerprint") == execution_fingerprint + and marker.get("rollout_backend") == instance_provenance["rollout"].backend_id + and marker.get("training_backend") == instance_provenance["training"].backend_id + ): + return False + + runtime = materialization.provenance.to_dict() + effective_status = effective_runtime_status(materialization).value + score_tensors: dict[str, Mapping[str, torch.Tensor]] = {} + for target in ("rollout", "training"): + prior = actual.get(target) + if not isinstance(prior, Mapping): + return False + instance = instance_provenance[target] + options = target_factory_options(operator_factory_options, target) + expected_operator = { + "backend_id": instance.backend_id, + "descriptor_fingerprint": instance.descriptor_fingerprint, + "implementation_fingerprint": instance.implementation_fingerprint, + "instance_fingerprint": instance.instance_fingerprint, + "concrete_implementation": instance.concrete_implementation, + "factory_options": json_safe(options), + "factory_options_fingerprint": factory_options_fingerprint(options), + } + prior_actual = prior.get("actual") + if not isinstance(prior_actual, Mapping): + return False + if any(prior_actual.get(key) != value for key, value in runtime["actual"].items()): + return False + if prior_actual.get("operators", {}).get("selected_logprob") != expected_operator: + return False + if prior_actual.get("model_state_fingerprint") != model_state_fingerprints[target]: + return False + if ( + prior_actual.get("scorer_implementation_fingerprint") + != scorer_implementation_fingerprints[target] + ): + return False + expected_implementation = hashlib.sha256( + ( + f"{materialization.provenance.implementation_fingerprint}:" + f"{instance.instance_fingerprint}" + ).encode("utf-8") + ).hexdigest() + for key, expected in ( + ("requested", runtime["requested"]), + ("normalized", runtime["normalized"]), + ("materialized", runtime["materialized"]), + ("status", effective_status), + ( + "construction_fingerprint", + materialization.provenance.construction_fingerprint, + ), + ( + "distributed_context_fingerprint", + materialization.provenance.distributed_context_fingerprint, + ), + ("process_fingerprint", materialization.provenance.process_fingerprint), + ("implementation_fingerprint", expected_implementation), + ("world_size", specs[target].world_size), + ): + if prior.get(key) != expected: + return False + try: + score_payload = _load_resume_tensor_bundle(attempt_dir / f"score_{target}.pt") + score_artifact = score_payload["metadata"]["score_artifact"] + prior_scorer = score_artifact["scorer"] + except (OSError, KeyError, TypeError, RuntimeError, ValueError): + return False + expected_scorer = concrete_scorer_spec(specs[target], instance).to_dict() + if prior_scorer != expected_scorer: + return False + if ( + score_artifact.get("schema_version") != "cross_config.score_artifact.v1" + or score_artifact.get("case_id") != case.case_id + or score_artifact.get("attempt_id") != attempt_dir.name + or score_artifact.get("side") != target + or score_artifact.get("identity") != batch.identity.to_dict() + or score_artifact.get("provenance") != prior + ): + return False + tensors = score_payload["tensors"] + selected = tensors.get("selected_logprobs") + active_mask = tensors.get("active_mask") + expected_dtype = torch_dtype(specs[target].dtype) + if ( + not isinstance(selected, torch.Tensor) + or not isinstance(active_mask, torch.Tensor) + or selected.shape != batch.input_ids.shape + or active_mask.shape != batch.input_ids.shape + or selected.dtype != expected_dtype + or active_mask.dtype != torch.bool + or not torch.equal(active_mask, batch.active_mask.to(device="cpu")) + ): + return False + metadata = score_payload["metadata"] + if ( + metadata.get("case_id") != case.case_id + or metadata.get("attempt_id") != attempt_dir.name + or metadata.get("side") != target + ): + return False + score_tensors[target] = tensors + return _resume_comparison_matches( + attempt_dir, + case, + batch, + marker, + score_tensors, + specs, + ) + + +def _load_resume_tensor_bundle(path: Path) -> Mapping[str, Any]: + payload = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(payload, Mapping) or payload.get("schema_version") != 1: + raise ValueError(f"invalid tensor bundle schema: {path}") + tensors = payload.get("tensors") + metadata = payload.get("metadata") + if not isinstance(tensors, Mapping) or not all( + isinstance(tensor, torch.Tensor) for tensor in tensors.values() + ): + raise ValueError(f"invalid tensor bundle payload: {path}") + if not isinstance(metadata, Mapping): + raise ValueError(f"invalid tensor bundle metadata: {path}") + return payload + + +def _resume_comparison_matches( + attempt_dir: Path, + case: ExperimentCase, + batch: CanonicalScoringBatch, + marker: Mapping[str, Any], + scores: Mapping[str, Mapping[str, torch.Tensor]], + specs: Mapping[str, ScorerSpec], +) -> bool: + try: + comparison = read_json_object(attempt_dir / "comparison.json") + token_bundle = _load_resume_tensor_bundle(attempt_dir / "token_diffs.pt") + except (OSError, ValueError, RuntimeError, json.JSONDecodeError): + return False + required_token_keys = { + "rollout_logprobs", + "training_logprobs", + "active_mask", + "absolute_diff", + "mismatch_mask", + } + token_tensors = token_bundle["tensors"] + if not required_token_keys.issubset(token_tensors): + return False + rollout = scores["rollout"]["selected_logprobs"] + training = scores["training"]["selected_logprobs"] + active_mask = batch.active_mask.to(device="cpu", dtype=torch.bool) + if not bool(torch.isfinite(rollout[active_mask]).all().item()) or not bool( + torch.isfinite(training[active_mask]).all().item() + ): + return False + rollout_threshold = resolve_logprob_threshold(specs["rollout"].dtype) + training_threshold = resolve_logprob_threshold(specs["training"].dtype) + if rollout_threshold != training_threshold: + return False + fixed_threshold = rollout_threshold + rollout = rollout.masked_fill(~active_mask, 0.0) + training = training.masked_fill(~active_mask, 0.0) + absolute_diff = torch.abs(training - rollout) + mismatch_mask = recompute_mismatch_mask( + rollout, + training, + active_mask, + fixed_threshold, + ) + expected_tensors = { + "rollout_logprobs": rollout, + "training_logprobs": training, + "active_mask": active_mask, + "absolute_diff": absolute_diff, + "mismatch_mask": mismatch_mask, + } + if any( + token_tensors[name].dtype != expected.dtype + or token_tensors[name].shape != expected.shape + or not torch.equal(token_tensors[name], expected) + for name, expected in expected_tensors.items() + ): + return False + token_metadata = token_bundle["metadata"] + active_count = int(active_mask.sum().item()) + if active_count == 0: + return False + mismatch_count = int(mismatch_mask.sum().item()) + passed = mismatch_count == 0 + status = "pass" if passed else "fail" + if ( + token_metadata.get("case_id") != case.case_id + or token_metadata.get("attempt_id") != attempt_dir.name + or token_metadata.get("status") != status + or token_metadata.get("fixed_threshold") != fixed_threshold + ): + return False + diagnostics = _comparison_diagnostics( + rollout, + training, + active_mask, + absolute_diff, + mismatch_count, + ) + expected_comparison = { + "schema_version": "cross_config.alignment_result.v1", + "case_id": case.case_id, + "attempt_id": attempt_dir.name, + "status": status, + "comparable": True, + "passed": passed, + "active_token_count": active_count, + "mismatch_count": mismatch_count, + "contract_fingerprint": tolerance_contract_fingerprint(), + "fixed_threshold": fixed_threshold, + "identity_errors": [], + "artifact_errors": [], + "diagnostics": diagnostics, + "token_artifact": { + **{name: _serialized_tensor(tensor) for name, tensor in expected_tensors.items()}, + "fixed_threshold": fixed_threshold, + "schema_version": "cross_config.token_comparison.v1", + }, + } + if comparison != expected_comparison: + return False + if ( + marker.get("case_id") != case.case_id + or marker.get("attempt_id") != attempt_dir.name + or marker.get("status") != status + or marker.get("comparable") is not True + or marker.get("passed") is not passed + or marker.get("active_token_count") != active_count + or marker.get("mismatch_count") != mismatch_count + or marker.get("max_abs_diff") != diagnostics["max_abs_diff"] + or marker.get("worst_token_index") != diagnostics["worst_token_index"] + ): + return False + return True + + +def _comparison_diagnostics( + rollout: torch.Tensor, + training: torch.Tensor, + active_mask: torch.Tensor, + absolute_diff: torch.Tensor, + mismatch_count: int, +) -> dict[str, Any]: + active_diff = absolute_diff[active_mask].float() + delta = (training[active_mask] - rollout[active_mask]).float() + worst_index = int(torch.argmax(active_diff).item()) + coordinates = torch.nonzero(active_mask, as_tuple=False) + worst_token = [int(item) for item in coordinates[worst_index].tolist()] + approximate_kl = torch.exp(delta.double()) - delta.double() - 1.0 + approximate_kl_mean = _finite_float(approximate_kl.mean()) + active_count = int(active_diff.numel()) + return { + "mean_abs_diff": _finite_float(active_diff.mean()), + "p95_abs_diff": _finite_float(torch.quantile(active_diff, 0.95)), + "p99_abs_diff": _finite_float(torch.quantile(active_diff, 0.99)), + "max_abs_diff": _finite_float(active_diff.max()), + "mismatch_ratio": mismatch_count / active_count, + "approximate_kl_mean": approximate_kl_mean, + "approximate_kl_finite": approximate_kl_mean is not None, + "worst_token_index": worst_token, + } + + +def _finite_float(value: torch.Tensor) -> Optional[float]: + result = float(value.item()) + return result if math.isfinite(result) else None + + +def _serialized_tensor(tensor: torch.Tensor) -> dict[str, Any]: + return { + "dtype": str(tensor.dtype).removeprefix("torch."), + "shape": list(tensor.shape), + "values": tensor.tolist(), + } + + +def read_json_object(path: Path) -> dict[str, Any]: + value = strict_json_loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"expected a JSON object: {path}") + return value diff --git a/rl_engine/alignment/cross_config/runner.py b/rl_engine/alignment/cross_config/runner.py new file mode 100644 index 00000000..1d9418b2 --- /dev/null +++ b/rl_engine/alignment/cross_config/runner.py @@ -0,0 +1,714 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Control plane for paired read-only scoring runs.""" + +from __future__ import annotations + +import importlib +import math +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Mapping, Optional + +import torch + +from rl_engine.alignment.cross_config._execution import ( + ChildScoringError, + ChildSupervisor, + OperatorExecutionError, + PairedRunnerError, + PairedScorer, + RankCompletenessError, + RankScore, + ScorerIdentityError, + ScoringTimeoutError, + canonical_fingerprint, + device_type, + json_safe, + normalized_dtype, + paired_model_state_fingerprints, + scorer_implementation_fingerprint, + scorer_spec, + validate_rank_outputs, + validate_scorer_identity, +) +from rl_engine.alignment.cross_config._provenance import ( + PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT, + concrete_scorer_spec, + effective_runtime_status, + execution_environment_provenance, + execution_fingerprint, + factory_options_fingerprint, + mapping_target, + runtime_adapter_fingerprint, + score_metadata, + side_provenance, + target_factory_options, +) +from rl_engine.alignment.cross_config._resume import completed_attempt_matches, read_json_object +from rl_engine.alignment.cross_config.artifacts import ArtifactStore +from rl_engine.alignment.cross_config.comparison import compare_score_artifacts +from rl_engine.alignment.cross_config.operators import ResolvedOperatorOverride +from rl_engine.alignment.cross_config.runtime import ( + RuntimeMaterialization, + RuntimeMaterializationError, +) +from rl_engine.alignment.cross_config.schema import ( + AlignmentResult, + CanonicalScoringBatch, + ExperimentCase, + MaterializationStatus, + ScoreArtifact, + ScorerSpec, + ScoreSide, +) +from rl_engine.kernels.semantic_registry import ( + OperatorInstanceProvenance, + OperatorResolution, + operator_implementation_fingerprint, + operator_instance_fingerprint, +) + + +@dataclass(frozen=True) +class PairedRunResult: + """Completed attempt or a validated resume hit.""" + + case_id: str + attempt_id: str + attempt_dir: Path + resumed: bool + rollout_score: Optional[ScoreArtifact] = None + training_score: Optional[ScoreArtifact] = None + alignment: Optional[AlignmentResult] = None + summary: Mapping[str, Any] = field(default_factory=dict) + + +class PairedRunner: + """Supervise paired scorers and publish one append-only attempt.""" + + def __init__( + self, + artifact_store: ArtifactStore, + *, + timeout_seconds: float = 30.0, + start_method: Optional[str] = None, + ): + if not math.isfinite(timeout_seconds) or timeout_seconds <= 0.0: + raise ValueError("timeout_seconds must be finite and greater than zero") + self.artifact_store = artifact_store + self.timeout_seconds = float(timeout_seconds) + self._child_supervisor = ChildSupervisor(start_method) + self.start_method = self._child_supervisor.start_method + + @property + def active_child_pids(self) -> tuple[int, ...]: + return self._child_supervisor.active_child_pids + + def run( + self, + case: ExperimentCase, + materialization: RuntimeMaterialization, + batch: CanonicalScoringBatch, + rollout_scorer: PairedScorer, + training_scorer: PairedScorer, + resolved_override: ResolvedOperatorOverride, + operator_instances: Mapping[str | ScoreSide, Any], + operator_instance_provenance: Mapping[ + str | ScoreSide, + OperatorInstanceProvenance, + ], + *, + operator_factory_options: Optional[Mapping[str | ScoreSide, Mapping[str, Any]]] = None, + strict: bool = True, + timeout_seconds: Optional[float] = None, + resume: bool = True, + ) -> PairedRunResult: + """Run both sides against one canonical batch and persist the comparison.""" + + deadline_seconds = self.timeout_seconds if timeout_seconds is None else timeout_seconds + if not math.isfinite(deadline_seconds) or deadline_seconds <= 0.0: + raise ValueError("timeout_seconds must be finite and greater than zero") + self._validate_case_inputs(case, materialization, batch) + rollout_spec = scorer_spec(rollout_scorer, ScoreSide.ROLLOUT) + training_spec = scorer_spec(training_scorer, ScoreSide.TRAINING) + specs = {"rollout": rollout_spec, "training": training_spec} + validate_scorer_identity(specs, batch) + model_state_fingerprints = paired_model_state_fingerprints( + rollout_scorer, + training_scorer, + ) + scorer_implementation_fingerprints = { + "rollout": scorer_implementation_fingerprint(rollout_scorer), + "training": scorer_implementation_fingerprint(training_scorer), + } + resolutions, instances, instance_provenance = _validate_exact_operators( + materialization, + resolved_override, + operator_instances, + operator_instance_provenance, + operator_factory_options=operator_factory_options, + specs=specs, + strict=strict, + ) + environment = execution_environment_provenance( + specs, + runtime_adapter_fingerprint=runtime_adapter_fingerprint(materialization), + operator_implementation_fingerprints={ + target: instance_provenance[target].implementation_fingerprint + for target in ("rollout", "training") + }, + ) + environment_fingerprint = canonical_fingerprint(environment) + _require_materialization_executable(materialization, strict=strict) + current_execution_fingerprint = execution_fingerprint( + materialization, + specs=specs, + instance_provenance=instance_provenance, + operator_factory_options=operator_factory_options, + model_state_fingerprints=model_state_fingerprints, + scorer_implementation_fingerprints=scorer_implementation_fingerprints, + environment=environment, + ) + + if resume: + completed = self.artifact_store.completed_attempt(case.experiment_id, case.case_id) + if completed is not None and completed_attempt_matches( + completed, + case, + batch, + materialization=materialization, + specs=specs, + instance_provenance=instance_provenance, + operator_factory_options=operator_factory_options, + model_state_fingerprints=model_state_fingerprints, + scorer_implementation_fingerprints=scorer_implementation_fingerprints, + environment=environment, + execution_fingerprint=current_execution_fingerprint, + ): + summary = read_json_object(completed / "COMPLETE") + return PairedRunResult( + case_id=case.case_id, + attempt_id=completed.name, + attempt_dir=completed, + resumed=True, + summary=summary, + ) + + attempt_dir = self.artifact_store.create_attempt(case.experiment_id, case.case_id) + attempt_id = attempt_dir.name + self._write_attempt_inputs(attempt_dir, attempt_id, case, materialization, batch) + + child_results = self._child_supervisor.run( + attempt_dir, + batch, + batch_size=materialization.binding.batch_size, + scorers={ + "rollout": rollout_scorer, + "training": training_scorer, + }, + specs=specs, + instances=instances, + timeout_seconds=float(deadline_seconds), + ) + rollout_ranks = validate_rank_outputs( + child_results["rollout"], + rollout_spec, + expected_shape=batch.input_ids.shape, + target="rollout", + ) + training_ranks = validate_rank_outputs( + child_results["training"], + training_spec, + expected_shape=batch.input_ids.shape, + target="training", + ) + + rollout_provenance = side_provenance( + materialization.provenance, + resolutions["rollout"], + instance_provenance["rollout"], + child_results["rollout"], + rollout_spec, + status=effective_runtime_status(materialization), + factory_options=target_factory_options(operator_factory_options, "rollout"), + model_state_fingerprint=model_state_fingerprints["rollout"], + scorer_implementation_fingerprint=scorer_implementation_fingerprints["rollout"], + ) + training_provenance = side_provenance( + materialization.provenance, + resolutions["training"], + instance_provenance["training"], + child_results["training"], + training_spec, + status=effective_runtime_status(materialization), + factory_options=target_factory_options(operator_factory_options, "training"), + model_state_fingerprint=model_state_fingerprints["training"], + scorer_implementation_fingerprint=scorer_implementation_fingerprints["training"], + ) + rollout_artifact = ScoreArtifact( + case_id=case.case_id, + attempt_id=attempt_id, + side=ScoreSide.ROLLOUT, + identity=batch.identity, + scorer=concrete_scorer_spec( + rollout_spec, + instance_provenance["rollout"], + ), + selected_logprobs=rollout_ranks[0].selected_logprobs, + active_mask=batch.active_mask, + provenance=rollout_provenance, + ) + training_artifact = ScoreArtifact( + case_id=case.case_id, + attempt_id=attempt_id, + side=ScoreSide.TRAINING, + identity=batch.identity, + scorer=concrete_scorer_spec( + training_spec, + instance_provenance["training"], + ), + selected_logprobs=training_ranks[0].selected_logprobs, + active_mask=batch.active_mask, + provenance=training_provenance, + ) + alignment = compare_score_artifacts(rollout_artifact, training_artifact) + self._write_attempt_results( + attempt_dir, + rollout_artifact, + training_artifact, + alignment, + execution_fingerprint=current_execution_fingerprint, + environment=environment, + environment_fingerprint=environment_fingerprint, + ) + summary = { + "schema_version": "cross_config.complete.v1", + "case_id": case.case_id, + "attempt_id": attempt_id, + "status": alignment.status.value, + "comparable": alignment.comparable, + "passed": alignment.passed, + "active_token_count": alignment.active_token_count, + "mismatch_count": alignment.mismatch_count, + "worst_token_index": alignment.diagnostics.get("worst_token_index"), + "max_abs_diff": alignment.diagnostics.get("max_abs_diff"), + "rollout_backend": instance_provenance["rollout"].backend_id, + "training_backend": instance_provenance["training"].backend_id, + "execution_fingerprint": current_execution_fingerprint, + "environment_fingerprint": environment_fingerprint, + "runner_implementation_fingerprint": PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT, + } + marker = self.artifact_store.complete_attempt(attempt_dir, summary=summary) + summary = read_json_object(marker) + return PairedRunResult( + case_id=case.case_id, + attempt_id=attempt_id, + attempt_dir=attempt_dir, + resumed=False, + rollout_score=rollout_artifact, + training_score=training_artifact, + alignment=alignment, + summary=summary, + ) + + @staticmethod + def _validate_case_inputs( + case: ExperimentCase, + materialization: RuntimeMaterialization, + batch: CanonicalScoringBatch, + ) -> None: + materialized_case = materialization.materialized_case.case + if materialized_case != case: + raise ValueError("materialization case does not exactly match the requested case") + if batch.identity != case.identity: + raise ValueError("canonical scoring batch identity does not match the case identity") + if batch.input_ids.shape[0] < 1: + raise ValueError("canonical scoring batch must contain at least one sequence") + + def _write_attempt_inputs( + self, + attempt_dir: Path, + attempt_id: str, + case: ExperimentCase, + materialization: RuntimeMaterialization, + batch: CanonicalScoringBatch, + ) -> None: + envelope = {"case_id": case.case_id, "attempt_id": attempt_id} + self.artifact_store.write_json( + attempt_dir, + "requested", + {**envelope, "schema_version": "cross_config.requested.v1", "case": case.to_dict()}, + ) + self.artifact_store.write_json( + attempt_dir, + "materialized", + { + **envelope, + "schema_version": "cross_config.materialized_envelope.v1", + "materialized_case": materialization.materialized_case.to_dict(), + }, + ) + self.artifact_store.write_json( + attempt_dir, + "identity", + { + **envelope, + "schema_version": "cross_config.identity_envelope.v1", + "identity": batch.identity.to_dict(), + }, + ) + + def _write_attempt_results( + self, + attempt_dir: Path, + rollout: ScoreArtifact, + training: ScoreArtifact, + alignment: AlignmentResult, + *, + execution_fingerprint: str, + environment: Mapping[str, Any], + environment_fingerprint: str, + ) -> None: + self.artifact_store.write_json( + attempt_dir, + "actual", + { + "case_id": rollout.case_id, + "attempt_id": rollout.attempt_id, + "schema_version": "cross_config.actual.v1", + "execution_fingerprint": execution_fingerprint, + "environment": environment, + "environment_fingerprint": environment_fingerprint, + "runner_implementation_fingerprint": PAIRED_RUNNER_IMPLEMENTATION_FINGERPRINT, + "operator_source": "exact_resolution_and_instance", + "rollout": rollout.provenance.to_dict(), + "training": training.provenance.to_dict(), + }, + ) + self.artifact_store.write_tensor_bundle( + attempt_dir, + "score_rollout", + { + "selected_logprobs": rollout.selected_logprobs, + "active_mask": rollout.active_mask, + }, + metadata=score_metadata(rollout), + ) + self.artifact_store.write_tensor_bundle( + attempt_dir, + "score_training", + { + "selected_logprobs": training.selected_logprobs, + "active_mask": training.active_mask, + }, + metadata=score_metadata(training), + ) + self.artifact_store.write_json( + attempt_dir, + "comparison", + alignment.to_dict(), + ) + token_artifact = alignment.token_artifact + if token_artifact is None: + empty = torch.empty((0,), dtype=torch.float32) + token_tensors = { + "rollout_logprobs": empty, + "training_logprobs": empty, + "active_mask": torch.empty((0,), dtype=torch.bool), + "absolute_diff": empty, + "mismatch_mask": torch.empty((0,), dtype=torch.bool), + } + else: + token_tensors = { + "rollout_logprobs": token_artifact.rollout_logprobs, + "training_logprobs": token_artifact.training_logprobs, + "active_mask": token_artifact.active_mask, + "absolute_diff": token_artifact.absolute_diff, + "mismatch_mask": token_artifact.mismatch_mask, + } + self.artifact_store.write_tensor_bundle( + attempt_dir, + "token_diffs", + token_tensors, + metadata={ + "case_id": alignment.case_id, + "attempt_id": alignment.attempt_id, + "status": alignment.status.value, + "fixed_threshold": alignment.fixed_threshold, + }, + ) + + +def _validate_exact_operators( + materialization: RuntimeMaterialization, + resolved: ResolvedOperatorOverride, + instances: Mapping[str | ScoreSide, Any], + instance_provenance: Mapping[str | ScoreSide, OperatorInstanceProvenance], + *, + operator_factory_options: Optional[Mapping[str | ScoreSide, Mapping[str, Any]]], + specs: Mapping[str, ScorerSpec], + strict: bool, +) -> tuple[ + dict[str, OperatorResolution], + dict[str, Any], + dict[str, OperatorInstanceProvenance], +]: + if resolved.semantic_op != "selected_logprob": + raise OperatorExecutionError("PairedRunner V1 requires semantic_op='selected_logprob'") + resolutions: dict[str, OperatorResolution] = {} + concrete_instances: dict[str, Any] = {} + provenance: dict[str, OperatorInstanceProvenance] = {} + for target in ("rollout", "training"): + resolution = resolved.for_target(target) # type: ignore[arg-type] + if resolution is None: + raise OperatorExecutionError(f"missing exact {target} operator resolution") + if resolution.target != target: + raise OperatorExecutionError( + f"{target} operator resolution reports target={resolution.target!r}" + ) + if ( + resolution.descriptor.semantic_op != "selected_logprob" + or resolution.trace.semantic_op != "selected_logprob" + ): + raise OperatorExecutionError(f"{target} resolution does not describe selected_logprob") + if resolution.trace.status != "resolved" or resolution.trace.concrete_backend is None: + raise OperatorExecutionError( + f"{target} operator is not exactly observable: {resolution.trace.status}" + ) + if resolution.trace.fallback_attempts: + raise OperatorExecutionError(f"{target} operator resolution attempted fallback") + if strict and not resolution.strict: + raise OperatorExecutionError(f"{target} operator was not resolved in strict mode") + if resolution.trace.concrete_backend != resolution.descriptor.backend_id: + raise OperatorExecutionError(f"{target} resolution backend evidence is inconsistent") + if resolution.trace.descriptor_fingerprint != resolution.descriptor.descriptor_fingerprint: + raise OperatorExecutionError( + f"{target} resolution descriptor fingerprint is inconsistent" + ) + if device_type(resolution.requirements.device) != device_type(specs[target].device): + raise OperatorExecutionError( + f"{target} operator resolution device does not match scorer device" + ) + if normalized_dtype(resolution.requirements.dtype) != normalized_dtype(specs[target].dtype): + raise OperatorExecutionError(f"{target} resolution dtype does not match scorer dtype") + _validate_exact_topology( + materialization, + resolution, + specs[target], + target=target, + ) + instance = mapping_target(instances, target) + if instance is None: + raise OperatorExecutionError(f"missing instantiated {target} operator") + _require_instance_matches_resolution(resolution, instance, target=target) + instance_evidence = mapping_target(instance_provenance, target) + if not isinstance(instance_evidence, OperatorInstanceProvenance): + raise OperatorExecutionError(f"missing sealed {target} operator instance provenance") + _validate_instance_provenance( + resolution, + instance, + instance_evidence, + factory_options=target_factory_options(operator_factory_options, target), + target=target, + ) + if instance_evidence.backend_id != resolution.trace.concrete_backend: + raise OperatorExecutionError(f"{target} instance backend does not match resolution") + declared = materialization.binding.operator_backends.get(target) + if declared is not None and declared != instance_evidence.backend_id: + raise OperatorExecutionError( + f"{target} exact backend {instance_evidence.backend_id!r} does not match " + f"declared override {declared!r}" + ) + resolutions[target] = resolution + concrete_instances[target] = instance + provenance[target] = instance_evidence + requested_logp = materialization.materialized_case.case.requested.get("logp") + requested_backend = ( + requested_logp.get("backend") if isinstance(requested_logp, Mapping) else None + ) + if requested_backend != provenance["rollout"].backend_id: + raise OperatorExecutionError( + "exact rollout operator does not match the public logp.backend request: " + f"{provenance['rollout'].backend_id!r} != {requested_backend!r}" + ) + return resolutions, concrete_instances, provenance + + +def _require_materialization_executable( + materialization: RuntimeMaterialization, + *, + strict: bool, +) -> None: + rejected = { + MaterializationStatus.ERROR, + MaterializationStatus.UNSUPPORTED, + MaterializationStatus.UNOBSERVABLE, + } + if strict: + rejected.add(MaterializationStatus.FALLBACK) + if not materialization.applications and materialization.materialized_case.status in rejected: + raise RuntimeMaterializationError( + f"case {materialization.materialized_case.case.case_id} is not executable: " + f"materialization={materialization.materialized_case.status.value}" + ) + problems = [] + for application in materialization.applications: + if ( + application.path == "logp.backend" + and application.status is MaterializationStatus.UNOBSERVABLE + ): + continue + if application.status in rejected: + problems.append( + f"{application.path}={application.status.value}: " + f"{application.evidence.get('reason', 'no evidence')}" + ) + if problems: + raise RuntimeMaterializationError( + f"case {materialization.materialized_case.case.case_id} is not executable: " + + "; ".join(problems) + ) + + +def _validate_exact_topology( + materialization: RuntimeMaterialization, + resolution: OperatorResolution, + spec: ScorerSpec, + *, + target: str, +) -> None: + bound_topology = mapping_target(materialization.binding.topology, target) + if not isinstance(bound_topology, Mapping): + raise OperatorExecutionError(f"{target} materialized topology is missing") + expected = dict(bound_topology) + topology_paths = { + "rollout": ( + ("rollout.tensor_parallel_size", "tensor_parallel_size"), + ("rollout.context_parallel_size", "context_parallel_size"), + ), + "training": (("training.sharding", "sharding"),), + } + required_keys = {"world_size", *(key for _, key in topology_paths[target])} + missing_keys = sorted(required_keys.difference(expected)) + if missing_keys: + raise OperatorExecutionError( + f"{target} materialized topology is missing required keys: {missing_keys!r}" + ) + if expected.get("world_size") != spec.world_size: + raise OperatorExecutionError( + f"{target} scorer world_size does not match materialized topology" + ) + if dict(resolution.requirements.topology) != expected: + raise OperatorExecutionError( + f"{target} resolution topology does not match materialized topology" + ) + if dict(spec.topology) != expected: + raise OperatorExecutionError( + f"{target} scorer topology does not match materialized topology" + ) + + actual_by_path = { + application.path: application.actual for application in materialization.applications + } + for path, key in topology_paths[target]: + if actual_by_path.get(path) != expected[key]: + raise OperatorExecutionError( + f"{target} actual {path} does not match exact operator topology" + ) + + +def _require_instance_matches_resolution( + resolution: OperatorResolution, + instance: Any, + *, + target: str, +) -> None: + implementation = resolution.descriptor.implementation_class_or_factory + factory = implementation + if isinstance(implementation, str): + try: + module_name, object_name = implementation.rsplit(".", 1) + factory = getattr(importlib.import_module(module_name), object_name) + except (ValueError, ImportError, AttributeError, ModuleNotFoundError) as exc: + raise OperatorExecutionError( + f"{target} exact operator factory cannot be verified: {exc}" + ) from exc + if isinstance(factory, type) and not isinstance(instance, factory): + raise OperatorExecutionError( + f"{target} operator instance type {type(instance).__qualname__!r} " + f"does not match resolved factory {factory.__qualname__!r}" + ) + if not callable(instance) and not callable(getattr(instance, "apply_fp32", None)): + raise OperatorExecutionError( + f"{target} selected-logprob operator instance is not executable" + ) + + +def _validate_instance_provenance( + resolution: OperatorResolution, + instance: Any, + provenance: OperatorInstanceProvenance, + *, + factory_options: Mapping[str, Any], + target: str, +) -> None: + expected_concrete = f"{type(instance).__module__}.{type(instance).__qualname__}" + expected_factory = resolution.descriptor.implementation_reference + if expected_factory is None: + raise OperatorExecutionError(f"{target} resolved operator has no factory reference") + implementation = resolution.descriptor.implementation_class_or_factory + if implementation is None: + raise OperatorExecutionError(f"{target} resolved operator has no implementation") + mismatches = [] + if provenance.semantic_op != resolution.descriptor.semantic_op: + mismatches.append("semantic_op") + if provenance.backend_id != resolution.descriptor.backend_id: + mismatches.append("backend_id") + if provenance.target != target: + mismatches.append("target") + if provenance.factory_reference != expected_factory: + mismatches.append("factory_reference") + if provenance.concrete_implementation != expected_concrete: + mismatches.append("concrete_implementation") + if provenance.descriptor_fingerprint != resolution.descriptor.descriptor_fingerprint: + mismatches.append("descriptor_fingerprint") + observed_implementation_fingerprint = operator_implementation_fingerprint( + implementation, + instance, + ) + if provenance.implementation_fingerprint != observed_implementation_fingerprint: + mismatches.append("implementation_fingerprint") + + if json_safe(provenance.factory_options) != json_safe(factory_options): + mismatches.append("factory_options") + if provenance.factory_options_fingerprint != factory_options_fingerprint(factory_options): + mismatches.append("factory_options_fingerprint") + expected_instance_fingerprint = operator_instance_fingerprint( + descriptor_fingerprint=resolution.descriptor.descriptor_fingerprint, + factory_reference=expected_factory, + concrete_implementation=expected_concrete, + implementation_fingerprint=observed_implementation_fingerprint, + factory_options_fingerprint=factory_options_fingerprint(factory_options), + ) + if provenance.instance_fingerprint != expected_instance_fingerprint: + mismatches.append("instance_fingerprint") + if mismatches: + raise OperatorExecutionError( + f"{target} operator instance provenance is inconsistent: " + ", ".join(mismatches) + ) + + +__all__ = [ + "ChildScoringError", + "OperatorExecutionError", + "PairedRunResult", + "PairedRunner", + "PairedRunnerError", + "PairedScorer", + "RankCompletenessError", + "RankScore", + "ScorerIdentityError", + "ScoringTimeoutError", +] From 4ce611f0d6129dbba089072f8d34c3158d21ca18 Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sun, 19 Jul 2026 08:34:41 +0800 Subject: [PATCH 06/33] feat(alignment): add CPU smoke workflow and CLI --- .gitignore | 3 + examples/cross_config_s0_cpu_smoke.json | 75 ++ .../cross_config_s1_distributed_smoke.json | 98 +++ examples/cross_config_s2_vllm_tp_vs_fsdp.json | 128 ++++ ...cross_config_s3_qwen3_8b_tp4_cp4_bf16.json | 134 ++++ rl_engine/alignment/cross_config/__main__.py | 138 ++++ rl_engine/alignment/testing/__init__.py | 18 + .../alignment/testing/cpu_cross_config.py | 696 ++++++++++++++++++ .../testing/smoke_ops/SMOKE_OPERATORS.md | 30 + .../alignment/testing/smoke_ops/__init__.py | 75 ++ .../smoke_ops/smoke_only_logp_offset.py | 98 +++ .../smoke_ops/smoke_only_logp_reference.py | 113 +++ 12 files changed, 1606 insertions(+) create mode 100644 examples/cross_config_s0_cpu_smoke.json create mode 100644 examples/cross_config_s1_distributed_smoke.json create mode 100644 examples/cross_config_s2_vllm_tp_vs_fsdp.json create mode 100644 examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json create mode 100644 rl_engine/alignment/cross_config/__main__.py create mode 100644 rl_engine/alignment/testing/__init__.py create mode 100644 rl_engine/alignment/testing/cpu_cross_config.py create mode 100644 rl_engine/alignment/testing/smoke_ops/SMOKE_OPERATORS.md create mode 100644 rl_engine/alignment/testing/smoke_ops/__init__.py create mode 100644 rl_engine/alignment/testing/smoke_ops/smoke_only_logp_offset.py create mode 100644 rl_engine/alignment/testing/smoke_ops/smoke_only_logp_reference.py diff --git a/.gitignore b/.gitignore index ae89c0d7..3f4ea52e 100644 --- a/.gitignore +++ b/.gitignore @@ -206,3 +206,6 @@ cython_debug/ marimo/_static/ marimo/_lsp/ __marimo__/ + +# Cross-configuration alignment local run artifacts +/runs/ diff --git a/examples/cross_config_s0_cpu_smoke.json b/examples/cross_config_s0_cpu_smoke.json new file mode 100644 index 00000000..0720c67d --- /dev/null +++ b/examples/cross_config_s0_cpu_smoke.json @@ -0,0 +1,75 @@ +{ + "schema_version": "cross_config.experiment_config.v1", + "experiment_id": "cross-config-s0-cpu-smoke-v1", + "scenario_id": "cross_config.s0.cpu_smoke.v1", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "identity": { + "checkpoint_id": "cross_config.synthetic.cpu_logits.v1", + "model_version": "immutable:cross-config-synthetic-cpu-v1", + "tokenizer_id": "cross_config.synthetic_tokenizer.v1", + "tokenizer_policy": "pretokenized_fixture.v1;add_special_tokens=false;padding_side=right", + "token_ids": [ + [11, 12, 13, 21, 22, 23], + [31, 32, 33, 41, 42, 43] + ], + "selected_token_ids": [ + [11, 12, 13, 21, 22, 23], + [31, 32, 33, 41, 42, 43] + ], + "active_mask": [ + [false, false, false, true, true, true], + [false, false, false, true, true, true] + ], + "attention_mask": [ + [true, true, true, true, true, true], + [true, true, true, true, true, true] + ], + "position_ids": [ + [0, 1, 2, 3, 4, 5], + [0, 1, 2, 3, 4, 5] + ], + "pre_update_state": "synthetic_read_only:no_parameters:no_optimizer", + "cache_metadata": { + "use_cache": false + }, + "packing_metadata": { + "packed": false + } + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "float32", + "enable_prefix_caching": false, + "enforce_eager": true + }, + "training": { + "attention_backend": "eager", + "compute_dtype": "float32", + "sharding": "unsharded" + }, + "logp": { + "backend": "smoke_only.logp_reference" + } + }, + "interventions": [], + "operators": { + "selected_logprob": { + "rollout": "smoke_only.logp_reference", + "training": "smoke_only.logp_reference" + } + }, + "scenario": { + "level": "S0", + "name": "CPU framework smoke", + "device": "cpu", + "hardware_required": false + } +} diff --git a/examples/cross_config_s1_distributed_smoke.json b/examples/cross_config_s1_distributed_smoke.json new file mode 100644 index 00000000..e7112c0d --- /dev/null +++ b/examples/cross_config_s1_distributed_smoke.json @@ -0,0 +1,98 @@ +{ + "schema_version": "cross_config.experiment_config.v1", + "experiment_id": "cross-config-s1-distributed-smoke-v1", + "scenario_id": "cross_config.s1.distributed_smoke.v1", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "identity": { + "checkpoint_id": "Qwen/Qwen3-0.6B", + "model_version": "c1899de289a04d12100db370d81485cdf75e47ca", + "tokenizer_id": "Qwen/Qwen3-0.6B@c1899de289a04d12100db370d81485cdf75e47ca", + "tokenizer_policy": "pretokenized_fixture.v1;add_special_tokens=false;padding_side=left", + "token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "selected_token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "active_mask": [ + [false, false, false, false, true, true, true, true], + [false, false, false, false, true, true, true, true] + ], + "attention_mask": [ + [true, true, true, true, true, true, true, true], + [true, true, true, true, true, true, true, true] + ], + "position_ids": [ + [0, 1, 2, 3, 4, 5, 6, 7], + [0, 1, 2, 3, 4, 5, 6, 7] + ], + "pre_update_state": "checkpoint_revision:c1899de289a04d12100db370d81485cdf75e47ca;optimizer_steps=0", + "cache_metadata": { + "position_policy": "absolute" + }, + "packing_metadata": { + "packed": false + } + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "bfloat16", + "enable_prefix_caching": true, + "enforce_eager": false + }, + "training": { + "attention_backend": "sdpa", + "compute_dtype": "bfloat16", + "sharding": "unsharded" + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "batch.size", + "values": [ + 1 + ] + }, + { + "path": "rollout.tensor_parallel_size", + "values": [ + 2 + ] + }, + { + "path": "rollout.context_parallel_size", + "values": [ + 2 + ] + }, + { + "path": "training.sharding", + "values": [ + "fsdp" + ] + } + ], + "scenario": { + "level": "S1", + "name": "Smallest distributed lifecycle smoke", + "device": "cuda", + "hardware_required": true, + "model_id": "Qwen/Qwen3-0.6B", + "model_revision": "c1899de289a04d12100db370d81485cdf75e47ca", + "rollout_engine": "vllm", + "training_engine": "fsdp_score_only" + } +} diff --git a/examples/cross_config_s2_vllm_tp_vs_fsdp.json b/examples/cross_config_s2_vllm_tp_vs_fsdp.json new file mode 100644 index 00000000..10a44cf3 --- /dev/null +++ b/examples/cross_config_s2_vllm_tp_vs_fsdp.json @@ -0,0 +1,128 @@ +{ + "schema_version": "cross_config.experiment_config.v1", + "experiment_id": "cross-config-s2-vllm-tp-vs-fsdp-v1", + "scenario_id": "cross_config.s2.vllm_tp_vs_fsdp.v1", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "identity": { + "checkpoint_id": "Qwen/Qwen3-8B", + "model_version": "b968826d9c46dd6066d109eabc6255188de91218", + "tokenizer_id": "Qwen/Qwen3-8B@b968826d9c46dd6066d109eabc6255188de91218", + "tokenizer_policy": "pretokenized_fixture.v1;add_special_tokens=false;padding_side=left", + "token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "selected_token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "active_mask": [ + [false, false, false, false, true, true, true, true], + [false, false, false, false, true, true, true, true] + ], + "attention_mask": [ + [true, true, true, true, true, true, true, true], + [true, true, true, true, true, true, true, true] + ], + "position_ids": [ + [0, 1, 2, 3, 4, 5, 6, 7], + [0, 1, 2, 3, 4, 5, 6, 7] + ], + "pre_update_state": "checkpoint_revision:b968826d9c46dd6066d109eabc6255188de91218;optimizer_steps=0", + "cache_metadata": { + "position_policy": "absolute" + }, + "packing_metadata": { + "packed": false + } + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 1, + "dtype": "bfloat16", + "enable_prefix_caching": true, + "enforce_eager": false + }, + "training": { + "attention_backend": "flash_attention_2", + "compute_dtype": "bfloat16", + "sharding": "fsdp" + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "batch.size", + "values": [ + 1 + ] + }, + { + "path": "rollout.tensor_parallel_size", + "values": [ + 1 + ] + }, + { + "path": "rollout.dtype", + "values": [ + "float32" + ] + }, + { + "path": "rollout.enable_prefix_caching", + "values": [ + false + ] + }, + { + "path": "rollout.enforce_eager", + "values": [ + true + ] + }, + { + "path": "training.attention_backend", + "values": [ + "eager" + ] + }, + { + "path": "training.compute_dtype", + "values": [ + "float32" + ] + }, + { + "path": "logp.backend", + "values": [ + "rlkernel.reference_logp" + ] + }, + { + "path": "training.sharding", + "values": [ + "unsharded" + ] + } + ], + "scenario": { + "level": "S2", + "name": "Issue 111 vLLM TP=2 versus training FSDP", + "device": "cuda", + "hardware_required": true, + "model_id": "Qwen/Qwen3-8B", + "model_revision": "b968826d9c46dd6066d109eabc6255188de91218", + "rollout_engine": "vllm", + "training_engine": "fsdp_score_only" + } +} diff --git a/examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json b/examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json new file mode 100644 index 00000000..2dcf79fb --- /dev/null +++ b/examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json @@ -0,0 +1,134 @@ +{ + "schema_version": "cross_config.experiment_config.v1", + "experiment_id": "cross-config-s3-qwen3-8b-tp4-cp4-bf16-v1", + "scenario_id": "cross_config.s3.qwen3_8b_tp4_cp4_bf16.v1", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "identity": { + "checkpoint_id": "Qwen/Qwen3-8B", + "model_version": "b968826d9c46dd6066d109eabc6255188de91218", + "tokenizer_id": "Qwen/Qwen3-8B@b968826d9c46dd6066d109eabc6255188de91218", + "tokenizer_policy": "pretokenized_fixture.v1;add_special_tokens=false;padding_side=left", + "token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "selected_token_ids": [ + [151643, 8948, 198, 11, 12, 13, 14, 15], + [151643, 8948, 198, 21, 22, 23, 24, 25] + ], + "active_mask": [ + [false, false, false, false, true, true, true, true], + [false, false, false, false, true, true, true, true] + ], + "attention_mask": [ + [true, true, true, true, true, true, true, true], + [true, true, true, true, true, true, true, true] + ], + "position_ids": [ + [0, 1, 2, 3, 4, 5, 6, 7], + [0, 1, 2, 3, 4, 5, 6, 7] + ], + "pre_update_state": "checkpoint_revision:b968826d9c46dd6066d109eabc6255188de91218;optimizer_steps=0", + "cache_metadata": { + "position_policy": "absolute" + }, + "packing_metadata": { + "packed": false + } + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 4, + "context_parallel_size": 4, + "dtype": "bfloat16", + "enable_prefix_caching": true, + "enforce_eager": false + }, + "training": { + "attention_backend": "flash_attention_2", + "compute_dtype": "bfloat16", + "sharding": "fsdp" + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "batch.size", + "values": [ + 1 + ] + }, + { + "path": "rollout.tensor_parallel_size", + "values": [ + 1 + ] + }, + { + "path": "rollout.context_parallel_size", + "values": [ + 1 + ] + }, + { + "path": "rollout.dtype", + "values": [ + "float32" + ] + }, + { + "path": "rollout.enable_prefix_caching", + "values": [ + false + ] + }, + { + "path": "rollout.enforce_eager", + "values": [ + true + ] + }, + { + "path": "training.attention_backend", + "values": [ + "eager" + ] + }, + { + "path": "training.compute_dtype", + "values": [ + "float32" + ] + }, + { + "path": "logp.backend", + "values": [ + "rlkernel.reference_logp" + ] + }, + { + "path": "training.sharding", + "values": [ + "unsharded" + ] + } + ], + "scenario": { + "level": "S3", + "name": "Roadmap Qwen3-8B TP=4 CP=4 BF16 milestone", + "device": "cuda", + "hardware_required": true, + "model_id": "Qwen/Qwen3-8B", + "model_revision": "b968826d9c46dd6066d109eabc6255188de91218", + "rollout_engine": "vllm", + "training_engine": "fsdp_score_only" + } +} diff --git a/rl_engine/alignment/cross_config/__main__.py b/rl_engine/alignment/cross_config/__main__.py new file mode 100644 index 00000000..cd047ca3 --- /dev/null +++ b/rl_engine/alignment/cross_config/__main__.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Command-line interface for cross-configuration experiments.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any, Optional, Sequence + +from rl_engine.alignment.cross_config.artifacts import ArtifactStore +from rl_engine.alignment.cross_config.config import ExperimentConfig, load_config +from rl_engine.alignment.cross_config.execution_plan import build_execution_plan + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + + plan = commands.add_parser("plan", help="validate and persist a plan without execution") + _add_common_arguments(plan) + + run = commands.add_parser("run", help="execute a plan with an explicit runtime adapter") + _add_common_arguments(run) + run.add_argument( + "--runtime", + required=True, + choices=("cpu-smoke",), + help="Runtime adapter; only the temporary CPU smoke adapter ships in V1", + ) + run.add_argument( + "--allow-smoke-operators", + action="store_true", + help="Explicitly authorize temporary smoke-only operator backends", + ) + run.add_argument( + "--timeout-seconds", + type=float, + default=30.0, + help="Per paired-scoring attempt deadline", + ) + run.add_argument( + "--no-resume", + action="store_true", + help="Create new attempts even when matching COMPLETE artifacts exist", + ) + return parser + + +def _add_common_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("config", type=Path, help="Versioned experiment JSON") + parser.add_argument( + "--output-root", + type=Path, + default=Path("runs"), + help="Append-only artifact root (default: runs)", + ) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = build_parser().parse_args(argv) + try: + config = load_config(args.config) + if args.command == "plan": + summary = record_plan(config, args.output_root) + else: + summary = _run(config, args) + except Exception as exc: + summary = { + "schema_version": "cross_config.cli_summary.v1", + "status": "error", + "error_type": f"{type(exc).__module__}.{type(exc).__qualname__}", + "error": str(exc), + } + print(json.dumps(summary, sort_keys=True)) + print(f"cross-configuration error: {type(exc).__name__}: {exc}", file=sys.stderr) + return 2 + + print(json.dumps(summary, sort_keys=True)) + if args.command == "plan": + print( + f"planned {summary['planned_case_count']} cases; no runtime was created", + file=sys.stderr, + ) + return 0 + print( + f"CPU smoke: {summary['status']} ({len(summary['cases'])} cases)", + file=sys.stderr, + ) + for case in summary["cases"]: + print( + f" {case['case_id']}: {case['status']}; actual backends " + f"rollout={case['rollout_backend']}, training={case['training_backend']}; " + f"worst sample/token={case['worst_token_index']}; " + f"mismatches={case['mismatch_count']}; resumed={case['resumed']}", + file=sys.stderr, + ) + return 0 if summary["status"] == "pass" else 1 + + +def record_plan(config: ExperimentConfig, output_root: Path) -> dict[str, Any]: + plan = build_execution_plan(config) + store = ArtifactStore(output_root) + experiment_dir = store.initialize_experiment( + config.definition.experiment_id, + experiment=plan.experiment, + plan=plan.rows(), + ) + return { + "schema_version": "cross_config.cli_summary.v1", + "status": "planned", + "experiment_id": config.definition.experiment_id, + "scenario_id": config.definition.scenario_id, + "planned_case_count": len(plan.entries), + "planning_issues": [issue.to_dict() for issue in plan.issues], + "artifact_dir": str(experiment_dir), + } + + +def _run(config: ExperimentConfig, args: argparse.Namespace) -> dict[str, Any]: + if args.runtime != "cpu-smoke": # pragma: no cover - argparse owns choices + raise ValueError(f"unsupported runtime {args.runtime!r}") + from rl_engine.alignment.testing.cpu_cross_config import run_cpu_experiment + + return run_cpu_experiment( + config, + output_root=args.output_root, + allow_smoke_operators=args.allow_smoke_operators, + timeout_seconds=args.timeout_seconds, + resume=not args.no_resume, + ) + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/rl_engine/alignment/testing/__init__.py b/rl_engine/alignment/testing/__init__.py new file mode 100644 index 00000000..2254e26a --- /dev/null +++ b/rl_engine/alignment/testing/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Test-only integration helpers for the alignment framework.""" + +from .smoke_ops import ( + SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + register_smoke_operators, + smoke_operator_descriptors, +) + +__all__ = [ + "SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID", + "SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID", + "register_smoke_operators", + "smoke_operator_descriptors", +] diff --git a/rl_engine/alignment/testing/cpu_cross_config.py b/rl_engine/alignment/testing/cpu_cross_config.py new file mode 100644 index 00000000..9c9e2ce2 --- /dev/null +++ b/rl_engine/alignment/testing/cpu_cross_config.py @@ -0,0 +1,696 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU-only adapters for cross-configuration smoke execution. + +This module is deliberately outside the production framework package. It gives +the CLI and tests a deterministic execution target without implying CUDA, +distributed, vLLM, or training-runtime support. +""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Mapping, Optional + +import torch + +from rl_engine.alignment.cross_config.artifacts import ArtifactStore +from rl_engine.alignment.cross_config.config import ExperimentConfig, OperatorSelection +from rl_engine.alignment.cross_config.execution_plan import build_execution_plan +from rl_engine.alignment.cross_config.operators import ( + OperatorBridge, + OperatorOverride, + selected_logprobs_with_operator, +) +from rl_engine.alignment.cross_config.runner import PairedRunner, PairedRunResult, RankScore +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, + RuntimeTools, +) +from rl_engine.alignment.cross_config.schema import ( + CanonicalScoringBatch, + ExperimentCase, + KnobDescriptor, + MaterializationStatus, + ScorerSpec, + ScoreSide, +) +from rl_engine.executors.stateless_executor import ( + StatelessForwardConfig, + StatelessForwardExecutor, + StatelessForwardInputs, +) +from rl_engine.kernels.registry import kernel_registry +from rl_engine.kernels.semantic_registry import ( + OperatorRequirements, + OperatorResolutionPolicy, + SemanticOperatorCatalog, +) +from rl_engine.kernels.semantic_registry import ( + implementation_fingerprint as fingerprint_implementation, +) + +CPU_SCORER_IMPLEMENTATION_FINGERPRINT = "cross_config.cpu_stateless_scorer.v1" + + +class SyntheticCpuCausalLM(torch.nn.Module): + """Deterministic parameter-free model for the named CPU smoke scenario.""" + + def __init__(self, vocab_size: int): + super().__init__() + self.vocab_axis: torch.Tensor + self.register_buffer( + "vocab_axis", + torch.arange(vocab_size, dtype=torch.float32), + persistent=False, + ) + self.config = SimpleNamespace(use_cache=False, _attn_implementation="eager") + self.generation_config = SimpleNamespace(use_cache=False) + + def forward( + self, + input_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + use_cache: Optional[bool] = None, + ) -> Any: + del attention_mask + if use_cache not in {None, False}: + raise ValueError("CPU smoke scoring forbids KV-cache generation") + if input_ids.device.type != "cpu": + raise ValueError("the synthetic smoke model accepts CPU tensors only") + if position_ids is None: + position_ids = torch.arange(input_ids.shape[1], device="cpu").expand_as(input_ids) + centers = torch.remainder(input_ids + position_ids + 1, self.vocab_axis.numel()).float() + logits = -torch.abs(self.vocab_axis.view(1, 1, -1) - centers.unsqueeze(-1)) * 0.125 + return SimpleNamespace(logits=logits, past_key_values=None) + + +class CpuStatelessScorer: + """Read-only teacher-forcing adapter over ``StatelessForwardExecutor``.""" + + optimizer = None + implementation_fingerprint = CPU_SCORER_IMPLEMENTATION_FINGERPRINT + + def __init__( + self, + model: torch.nn.Module, + spec: ScorerSpec, + config: Optional[StatelessForwardConfig] = None, + ): + if spec.world_size != 1: + raise ValueError("CpuStatelessScorer supports only world_size=1") + if _device_type(spec.device) != "cpu": + raise ValueError("CpuStatelessScorer is explicitly CPU-only") + resolved_config = config or StatelessForwardConfig( + mode="reference", + attention_backend="eager", + output_dtype=_torch_dtype(spec.dtype), + ) + if resolved_config.mode not in {"reference", "both"}: + raise ValueError("CpuStatelessScorer requires reference scoring mode") + expected_dtype = _torch_dtype(spec.dtype) + if resolved_config.output_dtype is not expected_dtype: + raise ValueError("stateless output_dtype must match the scorer dtype") + _require_module_on_cpu(model) + _require_module_float_dtype(model, expected_dtype) + self.model = model + self.spec = spec + self.config = resolved_config + + def score( + self, + batch: CanonicalScoringBatch, + *, + batch_size: int, + operator: Any, + ) -> tuple[RankScore, ...]: + if batch_size < 1: + raise ValueError("batch_size must be greater than zero") + + def selected_logprob_fn( + logits: torch.Tensor, + token_ids: torch.Tensor, + *, + mask: Optional[torch.Tensor] = None, + temperature: float = 1.0, + output_dtype: torch.dtype = torch.float32, + ) -> torch.Tensor: + return selected_logprobs_with_operator( + operator, + logits, + token_ids, + active_mask=mask, + temperature=temperature, + output_dtype=output_dtype, + ) + + executor = StatelessForwardExecutor( + self.model, + self.config, + selected_logprob_fn=selected_logprob_fn, + ) + chunks: list[torch.Tensor] = [] + observed_ranges: list[tuple[int, int]] = [] + for start in range(0, batch.input_ids.shape[0], batch_size): + stop = min(start + batch_size, batch.input_ids.shape[0]) + inputs = StatelessForwardInputs( + input_ids=batch.input_ids[start:stop], + attention_mask=batch.attention_mask[start:stop], + completion_mask=batch.active_mask[start:stop], + labels=batch.selected_token_ids[start:stop], + position_ids=( + None if batch.position_ids is None else batch.position_ids[start:stop] + ), + ) + result = executor.score(inputs) + if result.reference_logps is None: # pragma: no cover - guarded by config mode + raise RuntimeError("stateless scorer returned no selected logprobs") + chunks.append(result.reference_logps.detach().to(device="cpu")) + observed_ranges.append((start, stop)) + selected = torch.cat(chunks, dim=0) + return ( + RankScore( + rank=0, + world_size=1, + selected_logprobs=selected, + metadata={ + "device": "cpu", + "teacher_forcing": True, + "use_cache": False, + "optimizer_step": False, + "batch_ranges": observed_ranges, + }, + ), + ) + + +class CpuSmokeMaterializer: + """Materialize the exact single-process CPU surface used by smoke tests.""" + + runtime_kind = "cpu_smoke" + + @property + def implementation_fingerprint(self) -> str: + """Seal the adapter's concrete class and materialization entry point.""" + + return fingerprint_implementation( + type(self), + instance=self, + entrypoints=("materialize",), + ) + + def __init__( + self, + *, + requested_operator_backends: Optional[Mapping[str, str]] = None, + actual_operator_backends: Optional[Mapping[str, str]] = None, + ): + self.requested_operator_backends = dict(requested_operator_backends or {}) + self.actual_operator_backends = dict(actual_operator_backends or {}) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = _flatten(normalized) + applications = tuple( + self._application(path, value, descriptors[path]) for path, value in flat.items() + ) + batch_size = int(flat["batch.size"]) + requested_logp = str(flat["logp.backend"]) + operator_backends = self.requested_operator_backends or { + "rollout": requested_logp, + "training": requested_logp, + } + return AdapterMaterialization( + applications=applications, + binding=RuntimeBinding( + batch_size=batch_size, + side_configs={ + "rollout": { + "device": "cpu", + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + }, + "training": { + "device": "cpu", + "dtype": "float32", + "attention_backend": "eager", + "sharding": "unsharded", + }, + }, + topology={ + "rollout": { + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + "training": {"world_size": 1, "sharding": "unsharded"}, + }, + scorer={ + "mode": "reference", + "use_cache": False, + "attention_backend": "eager", + "output_dtype": "float32", + }, + operator_backends=operator_backends, + runtime_kind=self.runtime_kind, + ), + ) + + def _application( + self, + path: str, + requested: Any, + descriptor: KnobDescriptor, + ) -> KnobApplication: + fixed_values = { + "rollout.tensor_parallel_size": 1, + "rollout.context_parallel_size": 1, + "rollout.dtype": "float32", + "rollout.enable_prefix_caching": False, + "rollout.enforce_eager": True, + "training.attention_backend": "eager", + "training.compute_dtype": "float32", + "training.sharding": "unsharded", + } + if path == "batch.size": + return _application( + descriptor, + requested, + requested, + requested, + MaterializationStatus.APPLIED, + "canonical batch is partitioned at scorer invocation", + ) + if path == "logp.backend": + requested_backends = self.requested_operator_backends or { + "rollout": requested, + "training": requested, + } + if requested_backends.get("rollout") != requested: + return _application( + descriptor, + requested, + requested_backends, + None, + MaterializationStatus.ERROR, + "rollout operator conflicts with public logp.backend", + ) + actual_backends = { + "rollout": self.actual_operator_backends.get("rollout"), + "training": self.actual_operator_backends.get("training"), + } + if None in actual_backends.values(): + return _application( + descriptor, + requested, + requested_backends, + None, + MaterializationStatus.UNOBSERVABLE, + "operator resolution trace has not been supplied", + ) + status = ( + MaterializationStatus.APPLIED + if actual_backends == requested_backends + else MaterializationStatus.FALLBACK + ) + return _application( + descriptor, + requested, + requested_backends, + actual_backends, + status, + "concrete CPU backends were read from exact resolution traces", + ) + + actual = fixed_values[path] + status = ( + MaterializationStatus.APPLIED + if requested == actual + else MaterializationStatus.UNSUPPORTED + ) + reason = ( + "read back from the single-process CPU scorer" + if status is MaterializationStatus.APPLIED + else f"CPU smoke supports only {path}={actual!r}" + ) + return _application(descriptor, requested, requested, actual, status, reason) + + +def run_cpu_experiment( + config: ExperimentConfig, + *, + output_root: str | Path, + allow_smoke_operators: bool = False, + timeout_seconds: float = 30.0, + resume: bool = True, +) -> dict[str, Any]: + """Run every planned case through the explicit CPU smoke adapter.""" + + scenario_device = str(config.definition.scenario.get("device", "")).strip().lower() + if scenario_device != "cpu": + raise ValueError("the CPU runtime requires scenario.device='cpu'") + plan = build_execution_plan(config) + + store = ArtifactStore(output_root) + experiment_dir = store.initialize_experiment( + config.definition.experiment_id, + experiment=plan.experiment, + plan=plan.rows(), + ) + batch = canonical_cpu_batch(config) + runs = [ + run_cpu_case( + store, + entry.case, + batch, + entry.operators, + allow_smoke_operators=allow_smoke_operators, + strict=config.definition.strict_fallback, + timeout_seconds=timeout_seconds, + resume=resume, + ) + for entry in plan.entries + ] + cases = [ + { + "case_id": run.case_id, + "attempt_id": run.attempt_id, + "status": str(run.summary["status"]), + "rollout_backend": run.summary["rollout_backend"], + "training_backend": run.summary["training_backend"], + "mismatch_count": run.summary.get("mismatch_count"), + "worst_token_index": run.summary.get("worst_token_index"), + "resumed": run.resumed, + "attempt_dir": str(run.attempt_dir), + } + for run in runs + ] + return { + "schema_version": "cross_config.cli_summary.v1", + "status": "pass" if all(item["status"] == "pass" for item in cases) else "fail", + "experiment_id": config.definition.experiment_id, + "scenario_id": config.definition.scenario_id, + "runtime": "cpu-smoke", + "artifact_dir": str(experiment_dir), + "cases": cases, + } + + +def run_cpu_case( + store: ArtifactStore, + case: ExperimentCase, + batch: CanonicalScoringBatch, + selection: OperatorSelection, + *, + allow_smoke_operators: bool, + strict: bool, + timeout_seconds: float, + resume: bool, +) -> PairedRunResult: + """Execute one already-bound CPU case with case-local operator state.""" + + catalog = SemanticOperatorCatalog(kernel_registry.semantic.backend_descriptors()) + if allow_smoke_operators: + from rl_engine.alignment.testing.smoke_ops import register_smoke_operators + + register_smoke_operators(catalog, allow_smoke_operators=True) + bridge = OperatorBridge( + catalog, + policy=OperatorResolutionPolicy( + strict=strict, + allow_test_backends=allow_smoke_operators, + ), + ) + override = OperatorOverride( + semantic_op="selected_logprob", + rollout_backend=selection.rollout_backend, + training_backend=selection.training_backend, + ) + topologies: dict[str, Mapping[str, Any]] = { + "rollout": { + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + "training": {"world_size": 1, "sharding": "unsharded"}, + } + rollout_dtype = str(case.requested["rollout"]["dtype"]) + training_dtype = str(case.requested["training"]["compute_dtype"]) + requirements = { + "rollout": OperatorRequirements( + device="cpu", + dtype=rollout_dtype, + topology=topologies["rollout"], + alignment_properties={"deterministic": True}, + ), + "training": OperatorRequirements( + device="cpu", + dtype=training_dtype, + topology=topologies["training"], + alignment_properties={"deterministic": True}, + ), + } + resolved = bridge.resolve_override(override, requirements=requirements, strict=strict) + options = { + target: _factory_options( + selection.backend_for(target), + selection.options_for(target), + allow_smoke_operators=allow_smoke_operators, + ) + for target in ("rollout", "training") + } + instances = { + target: bridge.instantiate( + resolved, + target=target, # type: ignore[arg-type] + factory_kwargs=options[target], + ) + for target in ("rollout", "training") + } + provenance = { + target: bridge.instance_provenance( + resolved, + target=target, # type: ignore[arg-type] + instance=instances[target], + ) + for target in ("rollout", "training") + } + actual_backends = {target: provenance[target].backend_id for target in provenance} + materialization = RuntimeTools().materialize( + case, + CpuSmokeMaterializer( + requested_operator_backends={ + "rollout": selection.rollout_backend, + "training": selection.training_backend, + }, + actual_operator_backends=actual_backends, + ), + ) + RuntimeTools.require_executable(materialization, strict=strict) + + minimum_token_id = min( + int(batch.input_ids.min().item()), + int(batch.selected_token_ids.min().item()), + ) + if minimum_token_id < 0: + raise ValueError("CPU smoke token IDs must be non-negative") + vocab_size = ( + max( + int(batch.input_ids.max().item()), + int(batch.selected_token_ids.max().item()), + ) + + 17 + ) + scorers = { + "rollout": CpuStatelessScorer( + SyntheticCpuCausalLM(vocab_size), + _scorer_spec( + ScoreSide.ROLLOUT, + rollout_dtype, + selection.rollout_backend, + topologies["rollout"], + case, + ), + ), + "training": CpuStatelessScorer( + SyntheticCpuCausalLM(vocab_size), + _scorer_spec( + ScoreSide.TRAINING, + training_dtype, + selection.training_backend, + topologies["training"], + case, + ), + ), + } + return PairedRunner(store, timeout_seconds=timeout_seconds).run( + case, + materialization, + batch, + scorers["rollout"], + scorers["training"], + resolved, + instances, + provenance, + operator_factory_options=options, + strict=strict, + timeout_seconds=timeout_seconds, + resume=resume, + ) + + +def canonical_cpu_batch(config: ExperimentConfig) -> CanonicalScoringBatch: + """Build the immutable CPU tensors frozen by an experiment identity.""" + + identity = config.definition.identity + position_ids = ( + torch.tensor(identity.position_ids, dtype=torch.long, device="cpu") + if identity.position_ids + else None + ) + return CanonicalScoringBatch( + identity=identity, + input_ids=torch.tensor(identity.token_ids, dtype=torch.long, device="cpu"), + selected_token_ids=torch.tensor( + identity.selected_token_ids, + dtype=torch.long, + device="cpu", + ), + active_mask=torch.tensor(identity.active_mask, dtype=torch.bool, device="cpu"), + attention_mask=torch.tensor( + identity.attention_mask, + dtype=torch.bool, + device="cpu", + ), + position_ids=position_ids, + metadata={"source": "named_json", "device": "cpu"}, + ) + + +def _factory_options( + backend_id: str, + configured: Mapping[str, Any], + *, + allow_smoke_operators: bool, +) -> dict[str, Any]: + options = dict(configured) + if backend_id == "smoke_only.logp_offset": + if not allow_smoke_operators: + raise PermissionError("smoke offset requires explicit test authorization") + options["allow_smoke_operators"] = True + return options + + +def _scorer_spec( + side: ScoreSide, + dtype: str, + backend_id: str, + topology: Mapping[str, Any], + case: ExperimentCase, +) -> ScorerSpec: + identity = case.identity + return ScorerSpec( + side=side, + backend_id="cpu_stateless_teacher_forcing", + dtype=dtype, + device="cpu", + world_size=1, + topology=topology, + construction_options={ + "checkpoint_id": identity.checkpoint_id, + "model_version": identity.model_version, + "pre_update_state": identity.pre_update_state, + "teacher_forcing": True, + "use_cache": False, + }, + operator_overrides={"selected_logprob": backend_id}, + ) + + +def _application( + descriptor: KnobDescriptor, + requested: Any, + materialized: Any, + actual: Any, + status: MaterializationStatus, + reason: str, +) -> KnobApplication: + return KnobApplication( + path=descriptor.path, + requested=requested, + materialized=materialized, + actual=actual, + lifecycle=descriptor.lifecycle, + status=status, + critical=descriptor.critical, + evidence={"reason": reason}, + ) + + +def _flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + result: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}.{key}" if prefix else key + if isinstance(child, Mapping): + result.update(_flatten(child, path)) + else: + result[path] = child + return result + + +def _require_module_on_cpu(model: torch.nn.Module) -> None: + tensors = tuple(model.parameters()) + tuple(model.buffers()) + if any(tensor.device.type != "cpu" for tensor in tensors): + raise ValueError("CPU smoke models must remain on CPU") + + +def _require_module_float_dtype(model: torch.nn.Module, expected: torch.dtype) -> None: + tensors = tuple(model.parameters()) + tuple(model.buffers()) + mismatched = sorted( + { + str(tensor.dtype).replace("torch.", "") + for tensor in tensors + if tensor.is_floating_point() and tensor.dtype is not expected + } + ) + if mismatched: + raise ValueError( + f"CPU smoke model floating dtype must be {expected}; observed {mismatched}" + ) + + +def _device_type(value: str) -> str: + return value.split(":", 1)[0].strip().lower() + + +def _torch_dtype(value: str) -> torch.dtype: + normalized = value.strip().lower().replace("torch.", "") + aliases = {"fp32": "float32", "bf16": "bfloat16", "fp16": "float16"} + normalized = aliases.get(normalized, normalized) + try: + return { + "float32": torch.float32, + "bfloat16": torch.bfloat16, + "float16": torch.float16, + }[normalized] + except KeyError as exc: + raise ValueError(f"unsupported scorer dtype {value!r}") from exc + + +__all__ = [ + "CpuSmokeMaterializer", + "CpuStatelessScorer", + "SyntheticCpuCausalLM", + "canonical_cpu_batch", + "run_cpu_case", + "run_cpu_experiment", +] diff --git a/rl_engine/alignment/testing/smoke_ops/SMOKE_OPERATORS.md b/rl_engine/alignment/testing/smoke_ops/SMOKE_OPERATORS.md new file mode 100644 index 00000000..0ff9e6ea --- /dev/null +++ b/rl_engine/alignment/testing/smoke_ops/SMOKE_OPERATORS.md @@ -0,0 +1,30 @@ +# Cross-configuration smoke-only operators + +These files are temporary test scaffolding. They validate operator selection, +strict resolution, active-token scoring, and provenance; they do not establish +production numerical alignment. + +| File | Backend | Purpose | Replacement owner / issue | +| --- | --- | --- | --- | +| `smoke_only_logp_reference.py` | `smoke_only.logp_reference` | CPU PyTorch `log_softmax` plus gather reference for rollout/training injection tests. | Production selected-logprob operator workstream; roadmap issue #83 / WS1 contract issue #108. | +| `smoke_only_logp_offset.py` | `smoke_only.logp_offset` | Adds an explicit deterministic active-token offset so comparator mismatch detection can be tested. | Test-only fault injection; no production replacement should preserve the offset. | +| `__init__.py` | registration boundary | Keeps registration disabled by default and requires `allow_smoke_operators=True`. | Remove with both smoke implementations. | + +Removal trigger: delete this package once equivalent production RL-Kernel +selected-logprob operators are integrated and the same framework tests pass using +those production backends on both rollout and training sides. + +Exact deletion steps: + +1. Change `tests/test_cross_config_runtime.py` to exercise the production backend + IDs while preserving disabled/unavailable, capability, paired-output, and + provenance coverage. +2. Remove the `smoke_operator` test marker if no other temporary smoke operator + tests use it. +3. Delete `rl_engine/alignment/testing/smoke_ops/` and remove its exports from + `rl_engine/alignment/testing/__init__.py`. +4. Search the repository for `smoke_only.`, `allow_smoke_operators`, and + `RL_KERNEL_ALLOW_SMOKE_OPS`; remove configuration and documentation references + that no longer describe an active test boundary. +5. Run the cross-configuration contract, runtime, runner, and production-backend + tests before merging the deletion. diff --git a/rl_engine/alignment/testing/smoke_ops/__init__.py b/rl_engine/alignment/testing/smoke_ops/__init__.py new file mode 100644 index 00000000..5b6296ae --- /dev/null +++ b/rl_engine/alignment/testing/smoke_ops/__init__.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Opt-in registration for temporary Cross-configuration alignment smoke-only operators.""" + +from __future__ import annotations + +from typing import Any + +from rl_engine.kernels.semantic_registry import OperatorBackendDescriptor, SemanticOperatorCatalog + +SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID = "smoke_only.logp_reference" +SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID = "smoke_only.logp_offset" + + +def smoke_operator_descriptors() -> tuple[OperatorBackendDescriptor, ...]: + """Build smoke descriptors lazily without registering them globally.""" + + from .smoke_only_logp_offset import SMOKE_ONLY_LOGP_OFFSET_DESCRIPTOR + from .smoke_only_logp_reference import SMOKE_ONLY_LOGP_REFERENCE_DESCRIPTOR + + return ( + SMOKE_ONLY_LOGP_REFERENCE_DESCRIPTOR, + SMOKE_ONLY_LOGP_OFFSET_DESCRIPTOR, + ) + + +def register_smoke_operators( + catalog: SemanticOperatorCatalog, + *, + allow_smoke_operators: bool = False, + replace: bool = False, +) -> tuple[OperatorBackendDescriptor, ...]: + """Register every smoke backend after an explicit per-call opt-in. + + Importing this package never mutates a catalog. Resolution independently + requires an ``OperatorResolutionPolicy`` that allows test backends; this + registration guard is the first fail-closed boundary. + """ + + if allow_smoke_operators is not True: + raise PermissionError( + "smoke operator registration requires explicit " "allow_smoke_operators=True" + ) + if not isinstance(catalog, SemanticOperatorCatalog): + raise TypeError("catalog must be a SemanticOperatorCatalog") + + descriptors = smoke_operator_descriptors() + for descriptor in descriptors: + catalog.register_backend(descriptor, replace=replace) + return descriptors + + +def __getattr__(name: str) -> Any: + """Lazily expose implementation classes without default torch imports.""" + + if name == "SmokeOnlyLogpReference": + from .smoke_only_logp_reference import SmokeOnlyLogpReference + + return SmokeOnlyLogpReference + if name == "SmokeOnlyLogpOffset": + from .smoke_only_logp_offset import SmokeOnlyLogpOffset + + return SmokeOnlyLogpOffset + raise AttributeError(name) + + +__all__ = [ + "SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID", + "SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID", + "SmokeOnlyLogpOffset", + "SmokeOnlyLogpReference", + "register_smoke_operators", + "smoke_operator_descriptors", +] diff --git a/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_offset.py b/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_offset.py new file mode 100644 index 00000000..9bbe5e0e --- /dev/null +++ b/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_offset.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""TEMPORARY TEST SCAFFOLD - NOT A PRODUCTION RL-KERNEL OPERATOR""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch + +from rl_engine.kernels.semantic_registry import ( + OperatorBackendDescriptor, + OperatorFallbackPolicy, + OperatorLifecycle, +) + +from . import SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID +from .smoke_only_logp_reference import SmokeOnlyLogpReference + + +class SmokeOnlyLogpOffset(SmokeOnlyLogpReference): + """CPU reference plus a deterministic test-only active-token offset. + + Inputs and output follow :class:`SmokeOnlyLogpReference`. ``offset`` is zero + by default. A non-zero value requires the constructor's explicit + ``allow_smoke_operators=True`` guard. The cross-configuration bridge masks + inactive output positions after invocation, so drift applies only to active + selected tokens. + """ + + def __init__( + self, + offset: float = 0.0, + *, + allow_smoke_operators: bool = False, + ) -> None: + normalized_offset = float(offset) + if not math.isfinite(normalized_offset): + raise ValueError("offset must be finite") + if normalized_offset != 0.0 and allow_smoke_operators is not True: + raise PermissionError( + "a non-zero smoke offset requires explicit " "allow_smoke_operators=True" + ) + self.offset = normalized_offset + + def apply_fp32( + self, + logits: torch.Tensor, + token_ids: torch.Tensor, + active_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + selected = super().apply_fp32(logits, token_ids, active_mask=active_mask) + if self.offset == 0.0: + return selected + if active_mask is None: + return selected + self.offset + mask = active_mask.to(device=selected.device, dtype=torch.bool) + return selected + mask.to(dtype=selected.dtype) * self.offset + + +SMOKE_ONLY_LOGP_OFFSET_DESCRIPTOR = OperatorBackendDescriptor( + semantic_op="selected_logprob", + backend_id=SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu"}), + supported_dtypes=frozenset({"bfloat16", "float16", "float32"}), + supported_topologies={ + "rollout": { + "world_size": (1,), + "tensor_parallel_size": (1,), + "context_parallel_size": (1,), + }, + "training": { + "world_size": (1,), + "sharding": ("unsharded",), + }, + }, + determinism_or_alignment_properties={ + "algorithm": "pytorch.log_softmax_gather_plus_test_offset", + "batch_invariant": True, + "deterministic": True, + "strict_observable": True, + "test_offset_configurable": True, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=SmokeOnlyLogpOffset, + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="cross-config-smoke-only-logp-offset-v1", + is_smoke_only=True, +) + + +__all__ = [ + "SMOKE_ONLY_LOGP_OFFSET_DESCRIPTOR", + "SmokeOnlyLogpOffset", +] diff --git a/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_reference.py b/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_reference.py new file mode 100644 index 00000000..e64d37fa --- /dev/null +++ b/rl_engine/alignment/testing/smoke_ops/smoke_only_logp_reference.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""TEMPORARY TEST SCAFFOLD - NOT A PRODUCTION RL-KERNEL OPERATOR""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from rl_engine.kernels.semantic_registry import ( + OperatorBackendDescriptor, + OperatorFallbackPolicy, + OperatorLifecycle, +) + +from . import SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID + + +class SmokeOnlyLogpReference: + """CPU selected-logprob reference used only to test operator plumbing. + + The semantic inputs are logits shaped ``[..., vocabulary]`` and selected + token IDs shaped ``[...]``. The result is one float32 log probability per + selected token. When supplied, ``active_mask`` has the token-ID shape and + inactive output positions are exactly zero. The cross-configuration bridge + applies the same masking rule when invoking the two-argument interface. + """ + + op_class = "logprob" + + def __call__( + self, + logits: torch.Tensor, + token_ids: torch.Tensor, + active_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.apply_fp32(logits, token_ids, active_mask=active_mask) + + def apply_fp32( + self, + logits: torch.Tensor, + token_ids: torch.Tensor, + active_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Compute CPU log-softmax/gather output with optional active masking.""" + + _validate_inputs(logits, token_ids, active_mask) + selected_ids = token_ids.to(device=logits.device, dtype=torch.long) + mask = None + if active_mask is not None: + mask = active_mask.to(device=logits.device, dtype=torch.bool) + selected_ids = selected_ids.masked_fill(~mask, 0) + + log_probs = torch.log_softmax(logits.float(), dim=-1) + selected = torch.gather(log_probs, dim=-1, index=selected_ids.unsqueeze(-1)).squeeze(-1) + if mask is not None: + selected = selected.masked_fill(~mask, 0.0) + return selected + + +def _validate_inputs( + logits: torch.Tensor, + token_ids: torch.Tensor, + active_mask: Optional[torch.Tensor], +) -> None: + if logits.device.type != "cpu": + raise ValueError("smoke-only logprob operators support CPU tensors only") + if logits.shape[:-1] != token_ids.shape: + raise ValueError( + f"logits leading shape {tuple(logits.shape[:-1])} must match " + f"token_ids shape {tuple(token_ids.shape)}" + ) + if active_mask is not None and active_mask.shape != token_ids.shape: + raise ValueError("active_mask shape must match token_ids shape") + + +SMOKE_ONLY_LOGP_REFERENCE_DESCRIPTOR = OperatorBackendDescriptor( + semantic_op="selected_logprob", + backend_id=SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu"}), + supported_dtypes=frozenset({"bfloat16", "float16", "float32"}), + supported_topologies={ + "rollout": { + "world_size": (1,), + "tensor_parallel_size": (1,), + "context_parallel_size": (1,), + }, + "training": { + "world_size": (1,), + "sharding": ("unsharded",), + }, + }, + determinism_or_alignment_properties={ + "algorithm": "pytorch.log_softmax_gather", + "batch_invariant": True, + "deterministic": True, + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=SmokeOnlyLogpReference, + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="cross-config-smoke-only-logp-reference-v1", + is_smoke_only=True, +) + + +__all__ = [ + "SMOKE_ONLY_LOGP_REFERENCE_DESCRIPTOR", + "SmokeOnlyLogpReference", +] From 381edcea858a3e3a947fbfac20bf22fec239becd Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sun, 19 Jul 2026 08:35:20 +0800 Subject: [PATCH 07/33] test(alignment): add focused cross-configuration coverage --- .github/workflows/ci.yml | 8 + pyproject.toml | 5 + tests/test_cross_config_cli.py | 145 ++++++ tests/test_cross_config_contract.py | 440 ++++++++++++++++ tests/test_cross_config_runner.py | 659 +++++++++++++++++++++++ tests/test_cross_config_runtime.py | 780 ++++++++++++++++++++++++++++ tests/test_stateless_executor.py | 84 +++ tests/test_tolerance_contract.py | 59 ++- 8 files changed, 2179 insertions(+), 1 deletion(-) create mode 100644 tests/test_cross_config_cli.py create mode 100644 tests/test_cross_config_contract.py create mode 100644 tests/test_cross_config_runner.py create mode 100644 tests/test_cross_config_runtime.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..e9e9c91d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -68,6 +68,14 @@ jobs: python -m pytest rl_engine/tests/test_dispatch.py -v PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests/test_attention_correctness.py -q -rs + - name: Run Cross-Configuration Contract Tests (CPU-safe) + run: | + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest -q \ + tests/test_cross_config_*.py \ + tests/test_stateless_executor.py \ + tests/test_tolerance_contract.py \ + tests/test_kernel_registry.py + - name: Run Attention Ground-Truth Tests (CPU-safe) run: | python -m pytest tests/test_attention.py -v -k "not large and not gpu" diff --git a/pyproject.toml b/pyproject.toml index 9c8300ce..d4ed5c57 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,3 +43,8 @@ ignore = [] [tool.mypy] ignore_missing_imports = true follow_imports = "silent" + +[tool.pytest.ini_options] +markers = [ + "smoke_operator: temporary smoke-only operator plumbing tests", +] diff --git a/tests/test_cross_config_cli.py b/tests/test_cross_config_cli.py new file mode 100644 index 00000000..d5a3ce4a --- /dev/null +++ b/tests/test_cross_config_cli.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest +import torch + +import rl_engine.alignment.cross_config.__main__ as cli_main +from rl_engine.alignment.cross_config.artifacts import ArtifactStore + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +_EXAMPLES = _REPOSITORY_ROOT / "examples" +_CPU_RUNTIME_MODULE = "rl_engine.alignment.testing.cpu_cross_config" + + +def _summary(captured: str) -> dict: + summaries = [] + for line in captured.splitlines(): + try: + value = json.loads(line) + except json.JSONDecodeError: + continue + if value.get("schema_version") == "cross_config.cli_summary.v1": + summaries.append(value) + assert len(summaries) == 1 + return summaries[0] + + +def test_run_uses_only_cpu_and_resumes_when_cuda_is_available(tmp_path, capsys, monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + config_path = _EXAMPLES / "cross_config_s0_cpu_smoke.json" + plan_argv = [ + "plan", + str(config_path), + "--output-root", + str(tmp_path), + ] + argv = [ + "run", + str(config_path), + "--runtime", + "cpu-smoke", + "--allow-smoke-operators", + "--output-root", + str(tmp_path), + "--timeout-seconds", + "10", + ] + + assert cli_main.main(plan_argv) == 0 + planned = _summary(capsys.readouterr().out) + experiment_path = Path(planned["artifact_dir"]) / "experiment.json" + plan_path = Path(planned["artifact_dir"]) / "plan.jsonl" + planned_experiment = experiment_path.read_bytes() + planned_cases = plan_path.read_bytes() + stored_config = json.loads(planned_experiment) + stored_row = json.loads(planned_cases) + assert stored_config["schema_version"] == "cross_config.experiment_config.v1" + assert stored_row["schema_version"] == "cross_config.execution_plan_entry.v1" + assert stored_row["case"]["execution_binding"]["operators"] == stored_row["operators"] + + assert cli_main.main(argv) == 0 + captured = capsys.readouterr() + first = _summary(captured.out) + assert first["status"] == "pass" + assert first["runtime"] == "cpu-smoke" + assert "actual backends rollout=smoke_only.logp_reference" in captured.err + assert "training=smoke_only.logp_reference" in captured.err + assert "worst sample/token=[0, 3]" in captured.err + assert first["cases"] + assert all(case["status"] == "pass" for case in first["cases"]) + assert all(case["resumed"] is False for case in first["cases"]) + assert experiment_path.read_bytes() == planned_experiment + assert plan_path.read_bytes() == planned_cases + + store = ArtifactStore(tmp_path) + for case in first["cases"]: + attempt_dir = Path(case["attempt_dir"]) + assert (attempt_dir / "COMPLETE").is_file() + actual = json.loads((attempt_dir / "actual.json").read_text(encoding="utf-8")) + assert actual["environment"]["execution_devices"] == { + "rollout": "cpu", + "training": "cpu", + } + for name in ("score_rollout.pt", "score_training.pt"): + bundle = store.load_tensor_bundle(attempt_dir / name) + assert bundle["tensors"]["selected_logprobs"].device.type == "cpu" + + assert cli_main.main(argv) == 0 + resumed = _summary(capsys.readouterr().out) + assert resumed["status"] == "pass" + assert [case["attempt_id"] for case in resumed["cases"]] == [ + case["attempt_id"] for case in first["cases"] + ] + assert all(case["resumed"] is True for case in resumed["cases"]) + + +@pytest.mark.parametrize( + ("filename", "expected_cases"), + [ + ("cross_config_s1_distributed_smoke.json", 5), + ("cross_config_s2_vllm_tp_vs_fsdp.json", 10), + ("cross_config_s3_qwen3_8b_tp4_cp4_bf16.json", 11), + ], +) +def test_plan_records_named_scenarios_without_loading_a_runtime( + tmp_path, + capsys, + monkeypatch, + filename, + expected_cases, +): + monkeypatch.delitem(sys.modules, _CPU_RUNTIME_MODULE, raising=False) + + def runtime_must_not_run(*args, **kwargs): + raise AssertionError(f"plan unexpectedly invoked the CPU runtime: {args!r}, {kwargs!r}") + + monkeypatch.setattr(cli_main, "_run", runtime_must_not_run) + assert ( + cli_main.main( + [ + "plan", + str(_EXAMPLES / filename), + "--output-root", + str(tmp_path), + ] + ) + == 0 + ) + + summary = _summary(capsys.readouterr().out) + artifact_dir = Path(summary["artifact_dir"]) + assert summary["status"] == "planned" + assert summary["planned_case_count"] == expected_cases + assert (artifact_dir / "experiment.json").is_file() + assert len((artifact_dir / "plan.jsonl").read_text(encoding="utf-8").splitlines()) == ( + expected_cases + ) + assert not list(artifact_dir.glob("cases/*/*")) + assert _CPU_RUNTIME_MODULE not in sys.modules diff --git a/tests/test_cross_config_contract.py b/tests/test_cross_config_contract.py new file mode 100644 index 00000000..817de626 --- /dev/null +++ b/tests/test_cross_config_contract.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +from collections.abc import Mapping +from dataclasses import replace +from pathlib import Path +from typing import Any + +import pytest +import torch + +from rl_engine.alignment.cross_config.comparison import ( + compare_score_artifacts, + recompute_mismatch_mask, +) +from rl_engine.alignment.cross_config.config import ( + CONFIG_SCHEMA_VERSION, + bind_operator_selection, + load_config, +) +from rl_engine.alignment.cross_config.planner import MAX_PLAN_CASES, Planner, PlanningError +from rl_engine.alignment.cross_config.schema import ( + AlignmentStatus, + ExperimentDefinition, + InterventionSpec, + PlanningStrategy, + RuntimeProvenance, + ScoreArtifact, + ScorerSpec, + ScoreSide, + SemanticIdentitySpec, + TokenComparisonArtifact, +) +from rl_engine.kernels.gtest.tolerance import resolve_logprob_threshold + + +def _identity( + *, + checkpoint_id: str = "tiny-checkpoint", + tokenizer_policy: str = "tokenizer-v1:right-padding", + active_mask: tuple[tuple[bool, ...], ...] = ((True, False, True),), +) -> SemanticIdentitySpec: + return SemanticIdentitySpec( + checkpoint_id=checkpoint_id, + model_version="weights-v7", + tokenizer_id="tiny-tokenizer", + tokenizer_policy=tokenizer_policy, + token_ids=((11, 12, 13),), + selected_token_ids=((12, 13, 14),), + active_mask=active_mask, + attention_mask=((True, True, True),), + position_ids=((0, 1, 2),), + pre_update_state="state-before-step-9", + cache_metadata={"use_cache": False}, + packing_metadata={"packed": False}, + ) + + +def _score( + side: ScoreSide, + values: torch.Tensor, + *, + identity: SemanticIdentitySpec | None = None, + active_mask: torch.Tensor | None = None, +) -> ScoreArtifact: + identity = identity or _identity() + backend = f"test.{side.value}.selected_logprob" + return ScoreArtifact( + case_id="case-001", + attempt_id="attempt-001", + side=side, + identity=identity, + scorer=ScorerSpec( + side=side, + backend_id=f"{side.value}-scorer", + dtype="float32", + operator_overrides={"selected_logprob": backend}, + ), + selected_logprobs=values, + active_mask=( + active_mask + if active_mask is not None + else torch.tensor(identity.active_mask, dtype=torch.bool) + ), + provenance=RuntimeProvenance( + requested={"logp": {"backend": backend}}, + normalized={"logp": {"backend": backend}}, + materialized={"logp": {"backend": backend}}, + actual={"logp": {"backend": backend}}, + implementation_fingerprint=f"{side.value}-implementation-v1", + ), + ) + + +def _tensor_from_payload(payload: Mapping[str, Any]) -> torch.Tensor: + return torch.tensor(payload["values"], dtype=getattr(torch, str(payload["dtype"]))).reshape( + payload["shape"] + ) + + +def _baseline() -> dict[str, Any]: + return { + "batch": {"size": 8}, + "rollout": { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + }, + "training": { + "sharding": "unsharded", + "attention_backend": "eager", + "compute_dtype": "float32", + }, + "logp": {"backend": "rlkernel.reference_logp"}, + } + + +def _definition( + *, + strategy: PlanningStrategy = PlanningStrategy.ONE_AT_A_TIME, + pairwise_paths: tuple[tuple[str, str], ...] = (), +) -> ExperimentDefinition: + return ExperimentDefinition( + experiment_id="planner-test", + scenario_id="cpu-contract", + scenario={"model": "synthetic", "device": "cpu"}, + identity=_identity(), + baseline=_baseline(), + interventions=( + InterventionSpec("batch.size", (1, 4)), + InterventionSpec("rollout.dtype", ("bfloat16",)), + InterventionSpec("training.attention_backend", ("sdpa",)), + ), + strategy=strategy, + pairwise_paths=pairwise_paths, + ) + + +def _flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + result: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}.{key}" if prefix else key + if isinstance(child, Mapping): + result.update(_flatten(child, path)) + else: + result[path] = child + return result + + +def _config() -> dict[str, Any]: + return { + "schema_version": CONFIG_SCHEMA_VERSION, + "experiment_id": "cpu-config-test", + "scenario_id": "cpu-smoke", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": True, + "identity": { + "checkpoint_id": "tiny", + "model_version": "weights-v1", + "tokenizer_policy": "synthetic-v1", + "token_ids": [[1, 2, 3]], + "selected_token_ids": [[0, 2, 3]], + "active_mask": [[False, True, True]], + "attention_mask": [[True, True, True]], + "pre_update_state": "iteration-0", + }, + "baseline": { + **_baseline(), + "batch": {"size": 1}, + }, + "interventions": [{"path": "batch.size", "values": [2]}], + "operators": { + "selected_logprob": { + "rollout": "rlkernel.reference_logp", + "training": { + "backend": "smoke_only.logp_offset", + "options": {"offset": 0.1}, + }, + } + }, + "scenario": {"device": "cpu", "workload": "tiny"}, + } + + +def _write_config(tmp_path: Path, value: Mapping[str, Any], name: str = "config.json") -> Path: + path = tmp_path / name + path.write_text(json.dumps(value), encoding="utf-8") + return path + + +def test_fixed_threshold_uses_only_active_tokens_and_is_reproducible_offline(): + threshold = resolve_logprob_threshold("float32") + rollout_values = torch.zeros((1, 3), dtype=torch.float32) + training_values = torch.tensor( + [[threshold * 2.0, 1_000.0, threshold * 0.5]], + dtype=torch.float32, + ) + + result = compare_score_artifacts( + _score(ScoreSide.ROLLOUT, rollout_values), + _score(ScoreSide.TRAINING, training_values), + ) + + assert result.status is AlignmentStatus.FAIL + assert result.active_token_count == 2 + assert result.mismatch_count == 1 + assert result.fixed_threshold == threshold + assert result.token_artifact is not None + assert result.token_artifact.mismatch_mask.tolist() == [[True, False, False]] + + payload = json.loads(json.dumps(result.token_artifact.to_dict())) + offline = recompute_mismatch_mask( + _tensor_from_payload(payload["rollout_logprobs"]), + _tensor_from_payload(payload["training_logprobs"]), + _tensor_from_payload(payload["active_mask"]), + float(payload["fixed_threshold"]), + ) + assert torch.equal(offline, result.token_artifact.mismatch_mask) + assert not recompute_mismatch_mask( + torch.zeros(1, dtype=torch.float64), + torch.tensor([threshold], dtype=torch.float64), + torch.tensor([True]), + threshold, + ).item() + + inactive_nonfinite = training_values.clone() + inactive_nonfinite[0, 1] = float("nan") + sanitized = compare_score_artifacts( + _score(ScoreSide.ROLLOUT, rollout_values), + _score(ScoreSide.TRAINING, inactive_nonfinite), + ) + assert sanitized.status is result.status + assert sanitized.mismatch_count == result.mismatch_count + assert sanitized.token_artifact is not None + assert sanitized.token_artifact.training_logprobs[0, 1].item() == 0.0 + json.dumps(sanitized.to_dict(), allow_nan=False) + + +def test_zero_tokens_identity_mismatch_and_invalid_scores_are_not_numerical_failures(): + empty_identity = _identity(active_mask=((False, False, False),)) + empty = compare_score_artifacts( + _score(ScoreSide.ROLLOUT, torch.zeros((1, 3)), identity=empty_identity), + _score(ScoreSide.TRAINING, torch.zeros((1, 3)), identity=empty_identity), + ) + assert empty.status is AlignmentStatus.ZERO_ACTIVE_TOKENS + assert empty.comparable is False + assert empty.passed is False + + identity = _identity() + changed_identity = replace(identity, tokenizer_policy="different-policy") + mismatched = compare_score_artifacts( + _score(ScoreSide.ROLLOUT, torch.zeros((1, 3)), identity=identity), + _score(ScoreSide.TRAINING, torch.ones((1, 3)), identity=changed_identity), + ) + assert mismatched.status is AlignmentStatus.INVALID_IDENTITY + assert "tokenizer_policy" in mismatched.identity_errors + + invalid = compare_score_artifacts( + _score(ScoreSide.ROLLOUT, torch.zeros((1, 3))), + _score(ScoreSide.TRAINING, torch.tensor([[float("nan"), 0.0, 0.0]])), + ) + assert invalid.status is AlignmentStatus.INVALID_ARTIFACT + assert invalid.comparable is False + + with pytest.raises(ValueError, match="finite and non-negative"): + TokenComparisonArtifact( + rollout_logprobs=torch.zeros(1), + training_logprobs=torch.zeros(1), + active_mask=torch.ones(1, dtype=torch.bool), + absolute_diff=torch.zeros(1), + mismatch_mask=torch.zeros(1, dtype=torch.bool), + fixed_threshold=float("nan"), + ) + + +def test_planner_emits_one_stable_baseline_and_one_change_per_oat_case(): + definition = _definition() + plan = Planner().plan(definition) + baseline = _flatten(plan.cases[0].requested) + + assert len(plan.cases) == 5 + assert sum(not case.changed_paths for case in plan.cases) == 1 + for case in plan.cases[1:]: + requested = _flatten(case.requested) + changed = {path for path, value in requested.items() if value != baseline[path]} + assert changed == set(case.changed_paths) + assert len(changed) == 1 + + reordered = replace( + definition, + experiment_id="same-plan-from-another-run", + baseline={ + "logp": {"backend": "reference"}, + "training": { + "compute_dtype": "fp32", + "attention_backend": "eager", + "sharding": "unsharded", + }, + "rollout": { + "enforce_eager": True, + "enable_prefix_caching": False, + "dtype": "fp32", + "context_parallel_size": 1, + "tensor_parallel_size": 1, + }, + "batch": {"size": 8}, + }, + ) + assert [case.case_id for case in plan.cases] == [ + case.case_id for case in Planner().plan(reordered).cases + ] + + +def test_pairwise_is_explicit_and_planning_errors_remain_structured(): + pairwise = _definition( + strategy=PlanningStrategy.PAIRWISE, + pairwise_paths=(("batch.size", "rollout.dtype"),), + ) + pairwise_cases = [ + case for case in Planner().plan(pairwise).cases if len(case.changed_paths) == 2 + ] + assert len(pairwise_cases) == 2 + assert all(case.changed_paths == ("batch.size", "rollout.dtype") for case in pairwise_cases) + + with pytest.raises(PlanningError) as not_enabled: + Planner().plan( + replace( + _definition(), + pairwise_paths=(("batch.size", "rollout.dtype"),), + ) + ) + assert {issue.code for issue in not_enabled.value.issues} == {"PAIRWISE_NOT_ENABLED"} + + invalid_requests = ( + ({"logp": {"tp_layout": "arbitrary"}}, "DERIVED_KNOB"), + ({"batch": {"size": True}}, "UNSUPPORTED_VALUE"), + ({"rollout": {"unknown": 1}}, "UNSUPPORTED_PATH"), + ) + for requested, expected_code in invalid_requests: + with pytest.raises(PlanningError) as invalid: + Planner().normalize_requested(requested) + assert invalid.value.issues[0].code == expected_code + + incomplete = replace( + _definition(), + baseline={key: value for key, value in _baseline().items() if key != "training"}, + ) + with pytest.raises(PlanningError) as missing: + Planner().plan(incomplete) + assert {issue.path for issue in missing.value.issues} == { + "training.attention_backend", + "training.compute_dtype", + "training.sharding", + } + assert all(issue.code == "MISSING_BASELINE_VALUE" for issue in missing.value.issues) + + oversized = replace( + _definition(), + interventions=(InterventionSpec("batch.size", tuple(range(1, MAX_PLAN_CASES + 2))),), + ) + with pytest.raises(PlanningError) as too_large: + Planner().plan(oversized) + assert too_large.value.issues[0].code == "PLAN_TOO_LARGE" + + +def test_versioned_config_loads_and_binds_target_specific_operators(tmp_path: Path): + loaded = load_config(_write_config(tmp_path, _config())) + base_case = loaded.plan().cases[0] + selection = loaded.operators_for(base_case) + bound = bind_operator_selection(base_case, selection) + + assert loaded.schema_version == CONFIG_SCHEMA_VERSION + assert loaded.definition.strategy is PlanningStrategy.ONE_AT_A_TIME + assert selection.rollout_backend == "rlkernel.reference_logp" + assert selection.training_backend == "smoke_only.logp_offset" + assert selection.training_options == {"offset": 0.1} + assert bound == bind_operator_selection(base_case, selection) + assert bound.case_id != base_case.case_id + assert bound.requested == base_case.requested + assert bound.execution_binding["operators"] == selection.to_dict() + + +def test_config_rejects_schema_escape_hatches_and_incomplete_operator_coverage( + tmp_path: Path, +): + wrong_schema = _config() + wrong_schema["schema_version"] = "cross_config.experiment_config.v999" + + unknown_key = _config() + unknown_key["strict_falback"] = True + + threshold_override = _config() + threshold_override["scenario"]["nested"] = {"threshold": 999.0} + + scenario_policy = _config() + scenario_policy["scenario"]["execution"] = "run" + + conflicting_axis = _config() + conflicting_axis["interventions"].append( + {"path": "logp.backend", "values": ["smoke_only.logp_offset"]} + ) + + incomplete_targets = _config() + del incomplete_targets["operators"]["selected_logprob"]["training"] + + invalid_configs = ( + (wrong_schema, "unsupported cross-configuration config schema"), + (unknown_key, "unknown config keys"), + (threshold_override, "fixed numerical-contract threshold"), + (scenario_policy, "scenario is metadata only"), + (conflicting_axis, "cannot be combined with logp.backend interventions"), + (incomplete_targets, "selected_logprob.training"), + ) + for index, (value, message) in enumerate(invalid_configs): + with pytest.raises(ValueError, match=message): + load_config(_write_config(tmp_path, value, f"invalid-{index}.json")) + + duplicate_key = tmp_path / "duplicate.json" + duplicate_key.write_text( + '{"schema_version":"cross_config.experiment_config.v1",' + '"schema_version":"cross_config.experiment_config.v1"}', + encoding="utf-8", + ) + with pytest.raises(ValueError, match="duplicate JSON key"): + load_config(duplicate_key) + + overflow = tmp_path / "overflow.json" + overflow.write_text( + json.dumps(_config()).replace('"size": 1', '"size": 1e400'), + encoding="utf-8", + ) + with pytest.raises(ValueError, match="non-finite JSON number"): + load_config(overflow) diff --git a/tests/test_cross_config_runner.py b/tests/test_cross_config_runner.py new file mode 100644 index 00000000..21ad8505 --- /dev/null +++ b/tests/test_cross_config_runner.py @@ -0,0 +1,659 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import json +import time +from pathlib import Path + +import pytest +import torch + +from rl_engine.alignment.cross_config.artifacts import ArtifactError, ArtifactStore +from rl_engine.alignment.cross_config.comparison import recompute_mismatch_mask +from rl_engine.alignment.cross_config.config import OperatorSelection, bind_operator_selection +from rl_engine.alignment.cross_config.operators import OperatorBridge, OperatorOverride +from rl_engine.alignment.cross_config.runner import ( + ChildScoringError, + PairedRunner, + RankCompletenessError, + RankScore, + ScoringTimeoutError, +) +from rl_engine.alignment.cross_config.runtime import RuntimeTools +from rl_engine.alignment.cross_config.schema import ( + CanonicalScoringBatch, + ExperimentCase, + ScorerSpec, + ScoreSide, + SemanticIdentitySpec, +) +from rl_engine.alignment.testing.cpu_cross_config import CpuSmokeMaterializer, run_cpu_case +from rl_engine.alignment.testing.smoke_ops import ( + SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, +) +from rl_engine.alignment.testing.smoke_ops.smoke_only_logp_reference import SmokeOnlyLogpReference +from rl_engine.kernels.gtest.tolerance import resolve_logprob_threshold +from rl_engine.kernels.semantic_registry import OperatorRequirements + +_JSON_ARTIFACTS = ( + "requested.json", + "materialized.json", + "actual.json", + "identity.json", + "comparison.json", +) +_TENSOR_ARTIFACTS = ( + "score_rollout.pt", + "score_training.pt", + "token_diffs.pt", +) + + +class FixedRankScorer: + optimizer = None + model_state_fingerprint = "fixed-rank-scorer-state-v1" + + def __init__(self, spec: ScorerSpec, ranks): + self.spec = spec + self.ranks = tuple(ranks) + + def score(self, batch, *, batch_size, operator): + del batch_size, operator + return tuple( + RankScore( + rank=rank, + world_size=self.spec.world_size, + selected_logprobs=torch.zeros_like(batch.input_ids, dtype=torch.float32), + ) + for rank in self.ranks + ) + + +class FailingScorer(FixedRankScorer): + def score(self, batch, *, batch_size, operator): + del batch, batch_size, operator + raise RuntimeError("intentional scorer failure") + + +class SlowScorer(FixedRankScorer): + def score(self, batch, *, batch_size, operator): + time.sleep(2.0) + return super().score(batch, batch_size=batch_size, operator=operator) + + +def _identity() -> SemanticIdentitySpec: + token_ids = ( + (1, 2, 3, 4), + (2, 3, 4, 5), + (3, 4, 5, 6), + ) + selected = ( + (0, 2, 3, 4), + (0, 3, 4, 5), + (0, 4, 5, 6), + ) + active = tuple((False, True, True, True) for _ in token_ids) + attention = tuple((True, True, True, True) for _ in token_ids) + return SemanticIdentitySpec( + checkpoint_id="tiny-cpu-checkpoint", + model_version="weights-v1", + tokenizer_policy="synthetic-tokenizer-v1", + token_ids=token_ids, + selected_token_ids=selected, + active_mask=active, + attention_mask=attention, + pre_update_state="iteration-0", + ) + + +def _batch() -> CanonicalScoringBatch: + identity = _identity() + return CanonicalScoringBatch( + identity=identity, + input_ids=torch.tensor(identity.token_ids, device="cpu"), + selected_token_ids=torch.tensor(identity.selected_token_ids, device="cpu"), + active_mask=torch.tensor(identity.active_mask, device="cpu"), + attention_mask=torch.tensor(identity.attention_mask, device="cpu"), + metadata={"source": "runner-test", "device": "cpu"}, + ) + + +def _requested(*, backend: str = "rlkernel.reference_logp") -> dict[str, object]: + return { + "batch": {"size": 2}, + "rollout": { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + }, + "training": { + "attention_backend": "eager", + "compute_dtype": "float32", + "sharding": "unsharded", + }, + "logp": {"backend": backend}, + } + + +def _case( + *, + case_id: str = "case-runner", + backend: str = "rlkernel.reference_logp", +) -> ExperimentCase: + return ExperimentCase( + case_id=case_id, + experiment_id="runner-test", + scenario_id="S0", + identity=_identity(), + requested=_requested(backend=backend), + contract_fingerprint="contract-sha", + scenario_fingerprint="scenario-sha", + ) + + +def _topology(side: ScoreSide) -> dict[str, object]: + if side is ScoreSide.ROLLOUT: + return { + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + } + return {"world_size": 1, "sharding": "unsharded"} + + +def _requirements(side: ScoreSide) -> OperatorRequirements: + return OperatorRequirements( + device="cpu", + dtype="float32", + topology=_topology(side), + alignment_properties={"deterministic": True}, + ) + + +def _operators(): + bridge = OperatorBridge() + resolved = bridge.resolve_override( + OperatorOverride.for_target( + semantic_op="selected_logprob", + backend_id="rlkernel.reference_logp", + target="both", + ), + requirements={ + "rollout": _requirements(ScoreSide.ROLLOUT), + "training": _requirements(ScoreSide.TRAINING), + }, + strict=True, + ) + instances = { + target: bridge.instantiate(resolved, target=target) for target in ("rollout", "training") + } + provenance = { + target: bridge.instance_provenance( + resolved, + target=target, + instance=instances[target], + ) + for target in ("rollout", "training") + } + return resolved, instances, provenance + + +def _materialization(case: ExperimentCase): + backend = str(case.requested["logp"]["backend"]) + backends = {"rollout": backend, "training": backend} + return RuntimeTools().materialize( + case, + CpuSmokeMaterializer( + requested_operator_backends=backends, + actual_operator_backends=backends, + ), + ) + + +def _spec(side: ScoreSide) -> ScorerSpec: + identity = _identity() + return ScorerSpec( + side=side, + backend_id="fixed_cpu_teacher_forcing", + dtype="float32", + device="cpu", + world_size=1, + topology=_topology(side), + construction_options={ + "checkpoint_id": identity.checkpoint_id, + "model_version": identity.model_version, + "pre_update_state": identity.pre_update_state, + "teacher_forcing": True, + "use_cache": False, + }, + operator_overrides={"selected_logprob": "rlkernel.reference_logp"}, + ) + + +def _bound_smoke_case(scenario: str) -> tuple[ExperimentCase, OperatorSelection]: + threshold_offset = resolve_logprob_threshold("float32") * 4.0 + if scenario == "reference-reference": + rollout_backend = SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID + training_backend = SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID + rollout_options = {} + training_options = {} + elif scenario == "reference-offset": + rollout_backend = SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID + training_backend = SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID + rollout_options = {} + training_options = {"offset": threshold_offset} + elif scenario == "offset-offset": + rollout_backend = SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID + training_backend = SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID + rollout_options = {"offset": threshold_offset} + training_options = {"offset": threshold_offset} + else: # pragma: no cover - test helper contract + raise ValueError(f"unknown scenario: {scenario}") + selection = OperatorSelection( + rollout_backend=rollout_backend, + training_backend=training_backend, + rollout_options=rollout_options, + training_options=training_options, + ) + case = bind_operator_selection( + _case(case_id=f"case-{scenario}", backend=rollout_backend), + selection, + ) + return case, selection + + +def _write_required_artifacts( + store: ArtifactStore, + attempt_dir: Path, + *, + case_id: str = "case-1", + omit: frozenset[str] = frozenset(), + rollout_logprobs: torch.Tensor | None = None, + training_logprobs: torch.Tensor | None = None, + active_mask: torch.Tensor | None = None, + threshold: float = 0.05, +) -> None: + attempt_id = attempt_dir.name + json_values = { + "requested.json": { + "schema_version": "cross_config.requested.v1", + "case_id": case_id, + "attempt_id": attempt_id, + "case": {"case_id": case_id}, + }, + "materialized.json": { + "schema_version": "cross_config.materialized_envelope.v1", + "case_id": case_id, + "attempt_id": attempt_id, + "materialized_case": {"case": {"case_id": case_id}}, + }, + "actual.json": { + "schema_version": "cross_config.actual.v1", + "case_id": case_id, + "attempt_id": attempt_id, + "rollout": {}, + "training": {}, + }, + "identity.json": { + "schema_version": "cross_config.identity_envelope.v1", + "case_id": case_id, + "attempt_id": attempt_id, + "identity": {"checkpoint_id": "tiny"}, + }, + "comparison.json": { + "schema_version": "cross_config.alignment_result.v1", + "case_id": case_id, + "attempt_id": attempt_id, + "status": "pass", + "comparable": True, + "passed": True, + }, + } + for name, value in json_values.items(): + if name not in omit: + store.write_json(attempt_dir, name, value) + + rollout = rollout_logprobs if rollout_logprobs is not None else torch.tensor([-1.0, -2.0, -3.0]) + training = training_logprobs if training_logprobs is not None else rollout.clone() + active = active_mask if active_mask is not None else torch.tensor([True, True, True]) + mismatch = recompute_mismatch_mask(rollout, training, active, threshold) + tensor_values = { + "score_rollout.pt": { + "selected_logprobs": rollout, + "active_mask": active, + }, + "score_training.pt": { + "selected_logprobs": training, + "active_mask": active, + }, + "token_diffs.pt": { + "rollout_logprobs": rollout, + "training_logprobs": training, + "active_mask": active, + "absolute_diff": torch.abs(training - rollout), + "mismatch_mask": mismatch, + }, + } + for name, tensors in tensor_values.items(): + if name not in omit: + store.write_tensor_bundle( + attempt_dir, + name, + tensors, + metadata={ + "case_id": case_id, + "attempt_id": attempt_id, + "artifact": name, + "fixed_threshold": threshold, + }, + ) + + +def _complete_attempt( + store: ArtifactStore, + *, + experiment_id: str = "experiment-1", + case_id: str = "case-1", + **artifact_options, +) -> Path: + attempt_dir = store.create_attempt(experiment_id, case_id) + _write_required_artifacts(store, attempt_dir, case_id=case_id, **artifact_options) + store.complete_attempt( + attempt_dir, + summary={ + "schema_version": "cross_config.complete.v1", + "case_id": case_id, + "attempt_id": attempt_dir.name, + "status": "pass", + }, + ) + return attempt_dir + + +@pytest.mark.smoke_operator +def test_cpu_smoke_cases_preserve_read_only_scoring_and_exact_provenance(tmp_path: Path): + store = ArtifactStore(tmp_path) + batch = _batch() + inputs_before = batch.input_ids.clone() + expected = { + "reference-reference": (True, 0), + "reference-offset": (False, int(batch.active_mask.sum().item())), + "offset-offset": (True, 0), + } + + for scenario, (expected_pass, expected_mismatches) in expected.items(): + case, selection = _bound_smoke_case(scenario) + result = run_cpu_case( + store, + case, + batch, + selection, + allow_smoke_operators=True, + strict=True, + timeout_seconds=5.0, + resume=False, + ) + + assert result.resumed is False + assert result.alignment is not None + assert result.alignment.passed is expected_pass + assert result.alignment.mismatch_count == expected_mismatches + assert result.rollout_score is not None + assert result.training_score is not None + assert result.rollout_score.selected_logprobs.device.type == "cpu" + assert result.training_score.selected_logprobs.device.type == "cpu" + assert result.rollout_score.scorer.device == "cpu" + assert result.training_score.scorer.device == "cpu" + + guard = result.training_score.provenance.evidence["scoring_guard"] + assert guard == { + "model_state_verified": True, + "model_eval": True, + "no_grad": True, + "optimizer_step": False, + "model_modes_restored": True, + "model_state_unchanged": True, + } + assert result.rollout_score.provenance.evidence["scoring_guard"] == guard + assert result.training_score.provenance.evidence["rank_metadata"][0]["batch_ranges"] == ( + (0, 2), + (2, 3), + ) + rollout_state = result.rollout_score.provenance.evidence["model_state_fingerprint"] + training_state = result.training_score.provenance.evidence["model_state_fingerprint"] + assert rollout_state == training_state + + actual = json.loads((result.attempt_dir / "actual.json").read_text(encoding="utf-8")) + assert actual["operator_source"] == "exact_resolution_and_instance" + for target, backend in ( + ("rollout", selection.rollout_backend), + ("training", selection.training_backend), + ): + operator = actual[target]["actual"]["operators"]["selected_logprob"] + assert operator["backend_id"] == backend + assert operator["descriptor_fingerprint"] + assert operator["implementation_fingerprint"] + assert operator["instance_fingerprint"] + complete = result.attempt_dir / "COMPLETE" + assert complete.is_file() + assert result.summary == json.loads(complete.read_text(encoding="utf-8")) + + assert torch.equal(batch.input_ids, inputs_before) + + +@pytest.mark.smoke_operator +def test_runner_resumes_valid_attempt_and_retries_after_identity_or_tensor_change( + tmp_path: Path, + monkeypatch, +): + store = ArtifactStore(tmp_path) + case, selection = _bound_smoke_case("reference-reference") + batch = _batch() + + def run(): + return run_cpu_case( + store, + case, + batch, + selection, + allow_smoke_operators=True, + strict=True, + timeout_seconds=5.0, + resume=True, + ) + + first = run() + resumed = run() + assert first.attempt_id == "attempt-0001" + assert resumed.resumed is True + assert resumed.attempt_id == first.attempt_id + assert resumed.rollout_score is None + assert resumed.summary == first.summary + + token_path = first.attempt_dir / "token_diffs.pt" + payload = torch.load(token_path, map_location="cpu", weights_only=True) + payload["tensors"]["mismatch_mask"] = torch.ones_like(payload["tensors"]["mismatch_mask"]) + torch.save(payload, token_path) + + retried = run() + assert retried.resumed is False + assert retried.attempt_id == "attempt-0002" + assert (retried.attempt_dir / "COMPLETE").is_file() + + original_apply_fp32 = SmokeOnlyLogpReference.apply_fp32 + + def equivalent_apply_fp32(self, logits, token_ids, active_mask=None): + return original_apply_fp32(self, logits, token_ids, active_mask=active_mask) + + monkeypatch.setattr(SmokeOnlyLogpReference, "apply_fp32", equivalent_apply_fp32) + implementation_changed = run() + assert implementation_changed.resumed is False + assert implementation_changed.attempt_id == "attempt-0003" + before = json.loads((retried.attempt_dir / "actual.json").read_text(encoding="utf-8")) + after = json.loads( + (implementation_changed.attempt_dir / "actual.json").read_text(encoding="utf-8") + ) + assert ( + before["rollout"]["actual"]["operators"]["selected_logprob"]["implementation_fingerprint"] + != after["rollout"]["actual"]["operators"]["selected_logprob"]["implementation_fingerprint"] + ) + attempts = sorted(path.name for path in retried.attempt_dir.parent.iterdir()) + assert attempts == ["attempt-0001", "attempt-0002", "attempt-0003"] + + +@pytest.mark.parametrize( + ("mode", "error_type", "message"), + [ + ("failure", ChildScoringError, "intentional scorer failure"), + ("timeout", ScoringTimeoutError, "stopped children"), + ("missing-rank", RankCompletenessError, r"missing=\[0\]"), + ("duplicate-rank", RankCompletenessError, "duplicate ranks"), + ], +) +def test_runner_supervision_fails_closed_and_cleans_children( + tmp_path: Path, + mode: str, + error_type: type[Exception], + message: str, +): + case = _case(case_id=f"case-{mode}") + resolved, instances, provenance = _operators() + if mode == "failure": + rollout = FailingScorer(_spec(ScoreSide.ROLLOUT), ()) + training = FixedRankScorer(_spec(ScoreSide.TRAINING), (0,)) + timeout = 5.0 + elif mode == "timeout": + rollout = SlowScorer(_spec(ScoreSide.ROLLOUT), (0,)) + training = SlowScorer(_spec(ScoreSide.TRAINING), (0,)) + timeout = 0.1 + elif mode == "missing-rank": + rollout = FixedRankScorer(_spec(ScoreSide.ROLLOUT), ()) + training = FixedRankScorer(_spec(ScoreSide.TRAINING), (0,)) + timeout = 5.0 + else: + rollout = FixedRankScorer(_spec(ScoreSide.ROLLOUT), (0, 0)) + training = FixedRankScorer(_spec(ScoreSide.TRAINING), (0,)) + timeout = 5.0 + + runner = PairedRunner(ArtifactStore(tmp_path), timeout_seconds=timeout) + with pytest.raises(error_type, match=message): + runner.run( + case, + _materialization(case), + _batch(), + rollout, + training, + resolved, + instances, + provenance, + timeout_seconds=timeout, + ) + + assert runner.active_child_pids == () + attempt_dir = tmp_path / case.experiment_id / "cases" / case.case_id / "attempt-0001" + assert attempt_dir.is_dir() + assert not (attempt_dir / "COMPLETE").exists() + assert not list(attempt_dir.glob(".paired-runner-*")) + + +def test_artifacts_are_append_only_and_complete_marker_is_published_last(tmp_path: Path): + store = ArtifactStore(tmp_path) + attempt_dir = store.create_attempt("experiment-1", "case-1") + _write_required_artifacts( + store, + attempt_dir, + omit=frozenset({"token_diffs.pt"}), + ) + requested_path = attempt_dir / "requested.json" + rollout_path = attempt_dir / "score_rollout.pt" + requested_before = requested_path.read_bytes() + rollout_before = rollout_path.read_bytes() + + with pytest.raises(ArtifactError, match="refusing to overwrite"): + store.write_json(attempt_dir, "requested", {"case_id": "changed"}) + with pytest.raises(ArtifactError, match="refusing to overwrite"): + store.write_tensor_bundle( + attempt_dir, + "score_rollout", + {"selected_logprobs": torch.tensor([0.0])}, + ) + assert requested_path.read_bytes() == requested_before + assert rollout_path.read_bytes() == rollout_before + + summary = { + "schema_version": "cross_config.complete.v1", + "case_id": "case-1", + "attempt_id": attempt_dir.name, + "status": "pass", + } + with pytest.raises(ArtifactError, match=r"missing artifacts.*token_diffs\.pt"): + store.complete_attempt(attempt_dir, summary=summary) + assert not (attempt_dir / "COMPLETE").exists() + + _write_required_artifacts( + store, + attempt_dir, + omit=frozenset(_JSON_ARTIFACTS + _TENSOR_ARTIFACTS[:-1]), + ) + marker = store.complete_attempt(attempt_dir, summary=summary) + store.validate_completed_attempt(attempt_dir, expected_case_id="case-1") + payload_times = [ + (attempt_dir / name).stat().st_mtime_ns for name in _JSON_ARTIFACTS + _TENSOR_ARTIFACTS + ] + assert marker.stat().st_mtime_ns >= max(payload_times) + marker_value = json.loads(marker.read_text(encoding="utf-8")) + artifact_hashes = marker_value.pop("artifact_sha256") + assert marker_value == summary + assert set(artifact_hashes) == set(_JSON_ARTIFACTS + _TENSOR_ARTIFACTS) + assert all(len(value) == 64 for value in artifact_hashes.values()) + assert not list(attempt_dir.glob(".COMPLETE.*")) + with pytest.raises(ArtifactError, match="refusing to overwrite"): + store.complete_attempt(attempt_dir, summary=summary) + + next_attempt = store.create_attempt("experiment-1", "case-1") + assert next_attempt.name == "attempt-0002" + + +def test_resume_uses_newest_valid_attempt_and_tensors_support_offline_recompute(tmp_path: Path): + store = ArtifactStore(tmp_path) + rollout = torch.tensor([-1.0, -2.0, -3.0]) + training = torch.tensor([-1.01, -2.20, -2.50]) + active = torch.tensor([True, True, False]) + older = _complete_attempt( + store, + rollout_logprobs=rollout, + training_logprobs=training, + active_mask=active, + threshold=0.05, + ) + newer = _complete_attempt(store) + partial = store.create_attempt("experiment-1", "case-1") + store.write_json(partial, "requested", {"case_id": "case-1"}) + + assert store.completed_attempt("experiment-1", "case-1") == newer + (newer / "COMPLETE").write_text("{not-json", encoding="utf-8") + assert store.completed_attempt("experiment-1", "case-1") == older + + token_payload = store.load_tensor_bundle(older / "token_diffs.pt") + tensors = token_payload["tensors"] + recomputed = recompute_mismatch_mask( + tensors["rollout_logprobs"], + tensors["training_logprobs"], + tensors["active_mask"], + token_payload["metadata"]["fixed_threshold"], + ) + assert torch.equal(recomputed, tensors["mismatch_mask"]) + assert torch.equal(recomputed, torch.tensor([False, True, False])) + assert all(tensor.device.type == "cpu" for tensor in tensors.values()) + assert partial.name == "attempt-0003" + + materialized_path = older / "materialized.json" + materialized = json.loads(materialized_path.read_text(encoding="utf-8")) + materialized["materialized_case"]["case"]["case_id"] = "tampered" + materialized_path.write_text(json.dumps(materialized), encoding="utf-8") + assert store.completed_attempt("experiment-1", "case-1") is None diff --git a/tests/test_cross_config_runtime.py b/tests/test_cross_config_runtime.py new file mode 100644 index 00000000..a95ddd52 --- /dev/null +++ b/tests/test_cross_config_runtime.py @@ -0,0 +1,780 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import inspect +from dataclasses import replace +from pathlib import Path + +import pytest +import torch + +from rl_engine.alignment.cross_config.operators import ( + OperatorBridge, + OperatorOverride, + selected_logprobs_with_operator, +) +from rl_engine.alignment.cross_config.planner import V1_KNOBS +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, + RuntimeMaterializationError, + RuntimeTools, +) +from rl_engine.alignment.cross_config.schema import ( + ExperimentCase, + IsolationScope, + MaterializationStatus, + SemanticIdentitySpec, +) +from rl_engine.alignment.testing.cpu_cross_config import CpuSmokeMaterializer +from rl_engine.alignment.testing.smoke_ops import ( + SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + SmokeOnlyLogpOffset, + register_smoke_operators, +) +from rl_engine.kernels.ops.pytorch.loss.logp import NativeLogpOp +from rl_engine.kernels.semantic_registry import ( + OperatorRequirements, + OperatorResolutionError, + OperatorResolutionPolicy, + SemanticOperatorCatalog, +) +from rl_engine.kernels.semantic_registry import ( + implementation_fingerprint as fingerprint_implementation, +) +from rl_engine.testing import selected_logprobs_reference + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +_TEMPORARY_DOCSTRING = "TEMPORARY TEST SCAFFOLD - NOT A PRODUCTION RL-KERNEL OPERATOR" +_TOPOLOGIES = { + "rollout": { + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + }, + "training": {"world_size": 1, "sharding": "unsharded"}, +} + + +def _identity() -> SemanticIdentitySpec: + return SemanticIdentitySpec( + checkpoint_id="tiny-cpu-checkpoint", + model_version="weights-v1", + tokenizer_policy="synthetic-tokenizer-v1", + token_ids=((1, 2, 3),), + selected_token_ids=((0, 2, 3),), + active_mask=((False, True, True),), + attention_mask=((True, True, True),), + pre_update_state="iteration-0", + ) + + +def _requested(**overrides): + requested = { + "batch": {"size": 2}, + "rollout": { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + }, + "training": { + "attention_backend": "eager", + "compute_dtype": "float32", + "sharding": "unsharded", + }, + "logp": {"backend": "rlkernel.reference_logp"}, + } + for path, value in overrides.items(): + current = requested + parts = path.split(".") + for part in parts[:-1]: + current = current[part] + current[parts[-1]] = value + return requested + + +def _case( + *, + case_id: str = "case-1", + changed_paths=(), + requested=None, + execution_binding=None, +) -> ExperimentCase: + return ExperimentCase( + case_id=case_id, + experiment_id="runtime-test", + scenario_id="S0", + identity=_identity(), + requested=requested or _requested(), + changed_paths=changed_paths, + execution_binding=execution_binding or {}, + contract_fingerprint="contract-sha", + scenario_fingerprint="scenario-sha", + ) + + +def _value_at(requested, path: str): + current = requested + for part in path.split("."): + current = current[part] + return current + + +def _readback(requested): + return {path: _value_at(requested, path) for path in V1_KNOBS if path != "batch.size"} + + +class _RuntimeTestAdapter: + """Small observable fake kept beside the lifecycle tests that need it.""" + + runtime_kind = "test_runtime" + + def __init__(self, *, actual_readback=None): + self.actual_readback = dict(actual_readback or {}) + + @property + def implementation_fingerprint(self): + return fingerprint_implementation( + type(self), + instance=self, + entrypoints=("materialize",), + ) + + def materialize(self, normalized, descriptors): + applications = [] + for path, descriptor in descriptors.items(): + requested = _value_at(normalized, path) + materialized = requested + actual = self.actual_readback.get(path) + status = MaterializationStatus.UNOBSERVABLE + reason = "no runtime readback is available" + + unsupported = (path == "rollout.context_parallel_size" and requested != 1) or ( + path == "training.sharding" and requested != "unsharded" + ) + if unsupported: + materialized = actual = None + status = MaterializationStatus.UNSUPPORTED + reason = "the test adapter does not support this topology" + elif path == "batch.size": + actual = requested + status = MaterializationStatus.APPLIED + reason = "batch size is observed at scorer invocation" + elif path in self.actual_readback: + status = ( + MaterializationStatus.APPLIED + if actual == requested + else MaterializationStatus.FALLBACK + ) + reason = "runtime readback was captured" + + applications.append( + KnobApplication( + path=path, + requested=requested, + materialized=materialized, + actual=actual, + lifecycle=descriptor.lifecycle, + status=status, + evidence={"reason": reason}, + critical=descriptor.critical, + ) + ) + + backend = _value_at(normalized, "logp.backend") + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=_value_at(normalized, "batch.size"), + side_configs={"rollout": {}, "training": {}}, + topology={ + "rollout": { + "world_size": 1, + "tensor_parallel_size": _value_at( + normalized, "rollout.tensor_parallel_size" + ), + "context_parallel_size": _value_at( + normalized, "rollout.context_parallel_size" + ), + }, + "training": { + "world_size": 1, + "sharding": _value_at(normalized, "training.sharding"), + }, + }, + scorer={}, + operator_backends={"rollout": backend, "training": backend}, + runtime_kind=self.runtime_kind, + ), + ) + + +def _cpu_materializer() -> CpuSmokeMaterializer: + backends = { + "rollout": "rlkernel.reference_logp", + "training": "rlkernel.reference_logp", + } + return CpuSmokeMaterializer( + requested_operator_backends=backends, + actual_operator_backends=backends, + ) + + +def _requirements( + *, + target: str = "rollout", + device: str = "cpu", +) -> OperatorRequirements: + return OperatorRequirements( + device=device, + dtype="float32", + topology=_TOPOLOGIES[target], + alignment_properties={"deterministic": True}, + ) + + +def _catalog() -> SemanticOperatorCatalog: + """Clone repository descriptors so each test owns registration state.""" + + return SemanticOperatorCatalog(OperatorBridge().catalog.backend_descriptors()) + + +def test_cpu_materialization_records_all_ten_knobs_across_three_stages(): + case = _case() + materialization = RuntimeTools().materialize(case, _cpu_materializer()) + applications = {application.path: application for application in materialization.applications} + + assert len(V1_KNOBS) == 10 + assert set(applications) == set(V1_KNOBS) + assert materialization.materialized_case.status is MaterializationStatus.APPLIED + assert materialization.executable_in_strict_mode + RuntimeTools.require_executable(materialization, strict=True) + + for path, descriptor in V1_KNOBS.items(): + application = applications[path] + assert application.requested == _value_at(case.requested, path) + assert application.lifecycle is descriptor.lifecycle + assert application.status is MaterializationStatus.APPLIED + assert application.evidence["reason"] + + provenance = materialization.provenance + assert provenance.requested == case.requested + assert provenance.normalized == case.requested + assert provenance.materialized["batch"]["size"] == 2 + assert provenance.materialized["rollout"] == { + "tensor_parallel_size": 1, + "context_parallel_size": 1, + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + } + assert provenance.materialized["training"] == { + "attention_backend": "eager", + "compute_dtype": "float32", + "sharding": "unsharded", + } + assert provenance.materialized["logp"]["backend"] == { + "rollout": "rlkernel.reference_logp", + "training": "rlkernel.reference_logp", + } + assert provenance.actual == provenance.materialized + assert provenance.implementation_fingerprint == _cpu_materializer().implementation_fingerprint + assert provenance.evidence["adapter_implementation_fingerprint"] == ( + provenance.implementation_fingerprint + ) + + binding = materialization.binding + assert binding.runtime_kind == "cpu_smoke" + assert binding.side_configs["rollout"]["device"] == "cpu" + assert binding.side_configs["training"]["device"] == "cpu" + assert binding.side_configs["training"]["dtype"] == "float32" + assert binding.side_configs["rollout"] == { + "device": "cpu", + "dtype": "float32", + "enable_prefix_caching": False, + "enforce_eager": True, + } + assert binding.topology["rollout"] == { + "world_size": 1, + "tensor_parallel_size": 1, + "context_parallel_size": 1, + } + assert binding.topology["training"] == {"world_size": 1, "sharding": "unsharded"} + assert binding.scorer == { + "mode": "reference", + "use_cache": False, + "attention_backend": "eager", + "output_dtype": "float32", + } + with pytest.raises(TypeError): + binding.side_configs["rollout"]["device"] = "cuda" + with pytest.raises(TypeError): + applications["batch.size"].evidence["reason"] = "changed after fingerprinting" + + +def test_lifecycle_fingerprints_allow_request_reuse_and_isolate_engine_and_process_changes( + monkeypatch, +): + tools = RuntimeTools() + baseline_case = _case(case_id="baseline") + baseline_adapter = _RuntimeTestAdapter(actual_readback=_readback(baseline_case.requested)) + baseline = tools.materialize( + baseline_case, + baseline_adapter, + ) + + batch_requested = _requested(**{"batch.size": 1}) + batch = tools.materialize( + _case( + case_id="batch", + requested=batch_requested, + changed_paths=("batch.size",), + ), + _RuntimeTestAdapter(actual_readback=_readback(batch_requested)), + ) + assert batch.materialized_case.isolation_scope is IsolationScope.REQUEST + assert batch.materialized_case.construction_fingerprint == ( + baseline.materialized_case.construction_fingerprint + ) + assert batch.materialized_case.distributed_context_fingerprint == ( + baseline.materialized_case.distributed_context_fingerprint + ) + assert batch.materialized_case.process_fingerprint == ( + baseline.materialized_case.process_fingerprint + ) + assert tools.can_reuse(baseline, batch) + + dtype_requested = _requested(**{"rollout.dtype": "bfloat16"}) + dtype = tools.materialize( + _case( + case_id="dtype", + requested=dtype_requested, + changed_paths=("rollout.dtype",), + ), + _RuntimeTestAdapter(actual_readback=_readback(dtype_requested)), + ) + assert dtype.materialized_case.isolation_scope is IsolationScope.ENGINE_CONSTRUCTION + assert dtype.materialized_case.construction_fingerprint != ( + baseline.materialized_case.construction_fingerprint + ) + assert dtype.materialized_case.process_fingerprint == ( + baseline.materialized_case.process_fingerprint + ) + assert not tools.can_reuse(baseline, dtype) + + topology_requested = _requested(**{"rollout.tensor_parallel_size": 2}) + topology = tools.materialize( + _case( + case_id="topology", + requested=topology_requested, + changed_paths=("rollout.tensor_parallel_size",), + ), + _RuntimeTestAdapter(actual_readback=_readback(topology_requested)), + ) + assert topology.materialized_case.isolation_scope is IsolationScope.PROCESS + assert topology.materialized_case.process_fingerprint != ( + baseline.materialized_case.process_fingerprint + ) + assert not tools.can_reuse(baseline, topology) + + rebound_case = replace( + baseline_case, + case_id="rebound", + execution_binding={"operator_case": {"rollout_options": {"offset": 0.1}}}, + ) + rebound = tools.materialize( + rebound_case, + _RuntimeTestAdapter(actual_readback=_readback(rebound_case.requested)), + ) + assert rebound.materialized_case.construction_fingerprint == ( + baseline.materialized_case.construction_fingerprint + ) + assert not tools.can_reuse(baseline, rebound) + + changed_identity_case = replace( + baseline_case, + case_id="changed-identity", + identity=replace(baseline_case.identity, pre_update_state="iteration-1"), + ) + changed_identity = tools.materialize(changed_identity_case, baseline_adapter) + assert changed_identity.materialized_case.construction_fingerprint == ( + baseline.materialized_case.construction_fingerprint + ) + assert not tools.can_reuse(baseline, changed_identity) + + original_materialize = _RuntimeTestAdapter.materialize + + def materialize_with_same_result(self, normalized, descriptors): + return original_materialize(self, normalized, descriptors) + + monkeypatch.setattr( + _RuntimeTestAdapter, + "materialize", + materialize_with_same_result, + ) + changed_adapter = tools.materialize( + baseline_case, + _RuntimeTestAdapter(actual_readback=_readback(baseline_case.requested)), + ) + assert changed_adapter.provenance.actual == baseline.provenance.actual + assert ( + changed_adapter.provenance.implementation_fingerprint + != baseline.provenance.implementation_fingerprint + ) + assert not tools.can_reuse(baseline, changed_adapter) + + +def test_materialization_fails_closed_for_fallback_unobservable_and_unsupported_paths(): + tools = RuntimeTools() + + incomplete_descriptors = { + path: descriptor for path, descriptor in V1_KNOBS.items() if path != "batch.size" + } + with pytest.raises(RuntimeMaterializationError, match="missing descriptors"): + RuntimeTools(incomplete_descriptors).materialize( + _case(case_id="missing-descriptor"), + _RuntimeTestAdapter(actual_readback=_readback(_requested())), + ) + + fallback_requested = _requested(**{"training.attention_backend": "flash_attention_2"}) + fallback_readback = _readback(fallback_requested) + fallback_readback["training.attention_backend"] = "eager" + fallback = tools.materialize( + _case( + case_id="fallback", + requested=fallback_requested, + changed_paths=("training.attention_backend",), + ), + _RuntimeTestAdapter(actual_readback=fallback_readback), + ) + assert fallback.materialized_case.status is MaterializationStatus.FALLBACK + with pytest.raises(RuntimeMaterializationError, match=r"training\.attention_backend"): + tools.require_executable(fallback, strict=True) + tools.require_executable(fallback, strict=False) + + unobservable = tools.materialize( + _case(case_id="unobservable"), + _RuntimeTestAdapter(), + ) + assert unobservable.materialized_case.status is MaterializationStatus.UNOBSERVABLE + with pytest.raises(RuntimeMaterializationError, match="no runtime readback"): + tools.require_executable(unobservable, strict=False) + + unsupported_requested = _requested( + **{ + "rollout.context_parallel_size": 4, + "training.sharding": "fsdp", + } + ) + unsupported = tools.materialize( + _case( + case_id="unsupported", + requested=unsupported_requested, + changed_paths=("rollout.context_parallel_size", "training.sharding"), + ), + _RuntimeTestAdapter(), + ) + unsupported_paths = { + application.path + for application in unsupported.applications + if application.status is MaterializationStatus.UNSUPPORTED + } + assert unsupported_paths == {"rollout.context_parallel_size", "training.sharding"} + assert unsupported.materialized_case.status is MaterializationStatus.UNSUPPORTED + with pytest.raises(RuntimeMaterializationError, match=r"rollout\.context_parallel_size"): + tools.require_executable(unsupported, strict=False) + + cpu_only_requested = _requested(**{"rollout.tensor_parallel_size": 2}) + cpu_only = tools.materialize( + _case( + case_id="cpu-only", + requested=cpu_only_requested, + changed_paths=("rollout.tensor_parallel_size",), + ), + _cpu_materializer(), + ) + assert cpu_only.materialized_case.status is MaterializationStatus.UNSUPPORTED + assert cpu_only.binding.side_configs["rollout"]["device"] == "cpu" + assert cpu_only.binding.side_configs["training"]["device"] == "cpu" + + +def test_operator_binding_selects_rollout_training_and_both_without_side_leakage(): + bridge = OperatorBridge() + requirements = { + "rollout": _requirements(target="rollout"), + "training": _requirements(target="training"), + } + assert requirements["rollout"].to_dict()["schema_version"] == ( + "rlkernel.semantic_operator.requirements.v1" + ) + + for target, expected_targets in ( + ("rollout", {"rollout"}), + ("training", {"training"}), + ("both", {"rollout", "training"}), + ): + resolved = bridge.resolve_override( + OperatorOverride.for_target( + semantic_op="selected_logprob", + backend_id="rlkernel.reference_logp", + target=target, + ), + requirements=requirements, + strict=True, + ) + selected_targets = { + side for side in ("rollout", "training") if resolved.for_target(side) is not None + } + assert selected_targets == expected_targets + + instances = {} + for side in expected_targets: + resolution = resolved.for_target(side) + assert resolution is not None + assert resolution.to_dict()["schema_version"] == ( + "rlkernel.semantic_operator.resolution.v1" + ) + assert resolution.descriptor.to_dict()["schema_version"] == ( + "rlkernel.semantic_operator.backend_descriptor.v1" + ) + assert resolution.trace.to_dict()["schema_version"] == ( + "rlkernel.semantic_operator.resolution_trace.v1" + ) + instance = bridge.instantiate(resolved, target=side) + instances[side] = instance + assert isinstance(instance, NativeLogpOp) + provenance = bridge.instance_provenance( + resolved, + target=side, + instance=instance, + ) + assert provenance.backend_id == "rlkernel.reference_logp" + assert provenance.target == side + assert provenance.instance_fingerprint + assert provenance.to_dict()["schema_version"] == ( + "rlkernel.semantic_operator.instance_provenance.v1" + ) + with pytest.raises(TypeError): + provenance.factory_options["unexpected"] = True + if target == "both": + assert instances["rollout"] is not instances["training"] + + catalog = _catalog() + descriptor = catalog.backend_descriptor("selected_logprob", "rlkernel.reference_logp") + assert descriptor is not None + catalog.register_backend( + replace(descriptor, supported_topologies={"*": "*"}), + replace=True, + ) + asymmetric = OperatorBridge(catalog).resolve_override( + OperatorOverride.for_target( + semantic_op="selected_logprob", + backend_id="rlkernel.reference_logp", + target="both", + ), + requirements={ + "rollout": OperatorRequirements( + device="cpu", + dtype="float32", + topology={"world_size": 2, "tensor_parallel_size": 2}, + ), + "training": OperatorRequirements( + device="cpu", + dtype="float32", + topology={"world_size": 1, "sharding": "fsdp"}, + ), + }, + ) + assert asymmetric.rollout is not None and asymmetric.training is not None + assert asymmetric.rollout.requirements.topology != asymmetric.training.requirements.topology + + strict_session = _catalog().session() + with pytest.raises(OperatorResolutionError, match="topology"): + strict_session.resolve( + semantic_op="selected_logprob", + requested_backend="rlkernel.reference_logp", + target="rollout", + requirements=OperatorRequirements(device="cpu", dtype="float32", topology={}), + strict=True, + ) + session = catalog.session() + with pytest.raises(OperatorResolutionError, match="not registered") as unsupported: + session.resolve( + semantic_op="selected_logprob", + requested_backend="missing.backend", + target="rollout", + requirements=_requirements(), + strict=True, + ) + assert unsupported.value.trace.status == "unsupported" + assert unsupported.value.trace.fallback_attempts == () + + native_requirements = OperatorRequirements( + device="cpu", + dtype="float32", + topology=_TOPOLOGIES["training"], + ) + with pytest.raises(OperatorResolutionError, match="not exactly observable"): + session.resolve( + semantic_op="selected_logprob", + requested_backend="native", + target="training", + requirements=native_requirements, + strict=True, + ) + native = session.resolve( + semantic_op="selected_logprob", + requested_backend="native", + target="training", + requirements=native_requirements, + strict=False, + ) + assert native.trace.status == "unobservable" + assert native.trace.concrete_backend is None + + +@pytest.mark.smoke_operator +def test_smoke_package_is_temporary_cpu_only_and_disabled_without_two_explicit_opt_ins(): + from rl_engine.alignment.testing.smoke_ops import ( + smoke_only_logp_offset, + smoke_only_logp_reference, + ) + + assert inspect.getdoc(smoke_only_logp_reference) == _TEMPORARY_DOCSTRING + assert inspect.getdoc(smoke_only_logp_offset) == _TEMPORARY_DOCSTRING + + manifest = ( + _REPOSITORY_ROOT / "rl_engine/alignment/testing/smoke_ops/SMOKE_OPERATORS.md" + ).read_text(encoding="utf-8") + for required_text in ( + "temporary test scaffolding", + "smoke_only_logp_reference.py", + "smoke_only_logp_offset.py", + "allow_smoke_operators=True", + "delete this package", + ): + assert required_text in manifest + + catalog = _catalog() + for backend_id in ( + SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + ): + assert catalog.backend_descriptor("selected_logprob", backend_id) is None + + with pytest.raises(PermissionError, match="allow_smoke_operators=True"): + register_smoke_operators(catalog) + + descriptors = register_smoke_operators(catalog, allow_smoke_operators=True) + assert {descriptor.backend_id for descriptor in descriptors} == { + SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + SMOKE_ONLY_LOGP_OFFSET_BACKEND_ID, + } + for descriptor in descriptors: + assert descriptor.supported_devices == frozenset({"cpu"}) + assert descriptor.is_smoke_only is True + disabled_session = catalog.session( + OperatorResolutionPolicy(strict=True, allow_test_backends=False) + ) + with pytest.raises(OperatorResolutionError, match="explicit opt-in"): + disabled_session.resolve( + semantic_op="selected_logprob", + requested_backend=descriptor.backend_id, + target="training", + requirements=_requirements(target="training"), + ) + enabled_session = catalog.session( + OperatorResolutionPolicy(strict=True, allow_test_backends=True) + ) + with pytest.raises(OperatorResolutionError) as error: + enabled_session.resolve( + semantic_op="selected_logprob", + requested_backend=descriptor.backend_id, + target="training", + requirements=_requirements(target="training", device="cuda"), + ) + failed = { + decision.capability + for decision in error.value.trace.capability_decisions + if not decision.passed + } + assert failed == {"device"} + + assert SmokeOnlyLogpOffset().offset == 0.0 + with pytest.raises(PermissionError, match="allow_smoke_operators=True"): + SmokeOnlyLogpOffset(offset=0.01) + + +@pytest.mark.smoke_operator +def test_explicit_smoke_opt_in_runs_both_sides_on_cpu_with_sealed_provenance(): + catalog = _catalog() + register_smoke_operators(catalog, allow_smoke_operators=True) + bridge = OperatorBridge( + catalog, + policy=OperatorResolutionPolicy(strict=True, allow_test_backends=True), + ) + resolved = bridge.resolve_override( + OperatorOverride.for_target( + semantic_op="selected_logprob", + backend_id=SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID, + target="both", + ), + requirements={ + "rollout": _requirements(target="rollout"), + "training": _requirements(target="training"), + }, + strict=True, + ) + instances = { + target: bridge.instantiate(resolved, target=target) for target in ("rollout", "training") + } + logits = torch.tensor( + [[[1.0, 2.0, -1.0], [0.0, 3.0, 1.0], [2.0, 0.0, 4.0]]], + device="cpu", + ) + token_ids = torch.tensor([[1, -100, 2]], device="cpu") + active_mask = torch.tensor([[True, False, True]], device="cpu") + expected = selected_logprobs_reference(logits, token_ids, mask=active_mask) + + outputs = {} + for target, instance in instances.items(): + output = selected_logprobs_with_operator( + instance, + logits, + token_ids, + active_mask=active_mask, + ) + outputs[target] = output + assert output.device.type == "cpu" + torch.testing.assert_close(output, expected, atol=0.0, rtol=0.0) + assert torch.count_nonzero(output[~active_mask]).item() == 0 + + provenance = bridge.instance_provenance( + resolved, + target=target, + instance=instance, + ) + assert provenance.backend_id == SMOKE_ONLY_LOGP_REFERENCE_BACKEND_ID + assert provenance.target == target + assert provenance.concrete_implementation.endswith("SmokeOnlyLogpReference") + assert provenance.descriptor_fingerprint + assert provenance.instance_fingerprint + + for invalid_temperature in (float("nan"), float("inf"), float("-inf")): + with pytest.raises(ValueError, match="finite and greater than zero"): + selected_logprobs_with_operator( + instances["rollout"], + logits, + token_ids, + active_mask=active_mask, + temperature=invalid_temperature, + ) + + assert instances["rollout"] is not instances["training"] + torch.testing.assert_close(outputs["rollout"], outputs["training"], atol=0.0, rtol=0.0) diff --git a/tests/test_stateless_executor.py b/tests/test_stateless_executor.py index fc6d9b3a..851402f7 100644 --- a/tests/test_stateless_executor.py +++ b/tests/test_stateless_executor.py @@ -172,6 +172,90 @@ def test_executor_runs_full_sequence_forward_with_use_cache_false_and_detaches_o assert not hasattr(model.generation_config, "attn_implementation") +def test_executor_exact_selected_logprob_callable_is_optional_and_injected(): + inputs = _inputs() + logits = _logits_for(inputs) + calls = [] + + def selected_logprob_fn( + shifted_logits, + shifted_labels, + *, + mask, + temperature, + output_dtype, + ): + calls.append((shifted_logits, shifted_labels, mask, temperature, output_dtype)) + reference = selected_logprobs_reference( + shifted_logits, + shifted_labels, + mask=mask, + temperature=temperature, + output_dtype=output_dtype, + ) + return reference + mask.to(dtype=output_dtype) * 0.25 + + default = StatelessForwardExecutor( + FakeReferenceModel(logits), + StatelessForwardConfig(mode="reference"), + ).score(inputs) + injected = StatelessForwardExecutor( + FakeReferenceModel(logits), + StatelessForwardConfig(mode="reference"), + selected_logprob_fn=selected_logprob_fn, + ).score(inputs) + + assert default.reference_logps is not None + assert injected.reference_logps is not None + assert len(calls) == 1 + assert torch.equal(calls[0][2], inputs.completion_mask[:, 1:]) + torch.testing.assert_close( + injected.reference_logps[inputs.completion_mask], + default.reference_logps[inputs.completion_mask] + 0.25, + ) + assert torch.equal( + injected.reference_logps[~inputs.completion_mask], + torch.zeros_like(injected.reference_logps[~inputs.completion_mask]), + ) + + +def test_executor_uses_eval_no_grad_and_restores_mixed_module_modes_read_only(): + inputs = _inputs() + + class ReadOnlyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(2.0)) + self.register_buffer("counter", torch.tensor(3.0)) + self.child = torch.nn.Linear(1, 1) + + def forward(self, input_ids, attention_mask=None, use_cache=None): + del attention_mask + assert self.training is False + assert self.child.training is False + assert torch.is_grad_enabled() is False + assert use_cache is False + batch, sequence = input_ids.shape + logits = torch.zeros(batch, sequence, 8) + return {"logits": logits + self.weight * 0.0 + self.counter * 0.0} + + model = ReadOnlyModel() + model.train() + model.child.eval() + state_before = {name: value.detach().clone() for name, value in model.state_dict().items()} + + result = StatelessForwardExecutor( + model, + StatelessForwardConfig(mode="reference", attention_backend="eager"), + ).score(inputs) + + assert result.reference_logps is not None + assert result.metrics["model_eval_during_forward"] is True + assert model.training is True + assert model.child.training is False + assert all(torch.equal(state_before[name], value) for name, value in model.state_dict().items()) + + def test_executor_falls_back_for_models_without_use_cache_argument(): inputs = _inputs() executor = StatelessForwardExecutor( diff --git a/tests/test_tolerance_contract.py b/tests/test_tolerance_contract.py index fb429d81..2dff672e 100644 --- a/tests/test_tolerance_contract.py +++ b/tests/test_tolerance_contract.py @@ -3,7 +3,19 @@ from __future__ import annotations -from rl_engine.kernels.gtest.tolerance import load_contract +import hashlib +import inspect +import json + +import pytest +import torch + +from rl_engine.kernels.gtest import tolerance as tolerance_module +from rl_engine.kernels.gtest.tolerance import ( + load_contract, + resolve_logprob_threshold, + tolerance_contract_fingerprint, +) def test_load_contract_contains_expected_operator_classes(): @@ -27,3 +39,48 @@ def test_logprob_bfloat16_tolerance_covers_observed_reference_drift(): tolerance = contract["accuracy"]["default"]["logprob"]["bfloat16"] assert tolerance["atol"] >= 5.0e-2 assert tolerance["rtol"] == 0.0 + + +@pytest.mark.parametrize( + ("dtype", "dtype_name"), + ( + (torch.float32, "float32"), + ("fp32", "float32"), + (torch.bfloat16, "bfloat16"), + ("bf16", "bfloat16"), + (torch.float16, "float16"), + ("fp16", "float16"), + ), +) +def test_resolve_logprob_threshold_reads_current_ws1_absolute_tolerance(dtype, dtype_name): + expected = load_contract()["accuracy"]["default"]["logprob"][dtype_name]["atol"] + + assert resolve_logprob_threshold(dtype) == expected + + +def test_resolve_logprob_threshold_has_no_contract_or_value_override_parameter(): + assert tuple(inspect.signature(resolve_logprob_threshold).parameters) == ("dtype",) + + +def test_resolve_logprob_threshold_rejects_dtype_outside_ws1_contract(monkeypatch): + with pytest.raises(ValueError, match="unsupported WS1 logprob dtype"): + resolve_logprob_threshold(torch.float64) + + for invalid in (True, -1.0, float("nan"), float("inf"), "0.1"): + contract = load_contract() + contract["accuracy"]["default"]["logprob"]["float32"]["atol"] = invalid + monkeypatch.setattr(tolerance_module, "load_contract", lambda value=contract: value) + with pytest.raises(ValueError, match="invalid WS1 logprob threshold"): + resolve_logprob_threshold("float32") + + +def test_tolerance_contract_fingerprint_is_canonical_content_sha256(): + canonical = json.dumps( + load_contract(), + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + + assert tolerance_contract_fingerprint() == hashlib.sha256(canonical).hexdigest() + assert len(tolerance_contract_fingerprint()) == 64 From 1c4457e16866328b55f7642bbea060146b6184ad Mon Sep 17 00:00:00 2001 From: CyberSecurityErial <2710555967@qq.com> Date: Sun, 19 Jul 2026 08:35:49 +0800 Subject: [PATCH 08/33] docs(alignment): document cross-configuration workflow --- docs/assets/ws2-cross-config-before-after.png | Bin 93735 -> 0 bytes .../cross_config_implementation_report.md | 185 ++++ .../cross_config_logprob_drift_contract.md | 308 ++++++ ...ws2_cross_config_logprob_drift_contract.md | 874 ------------------ 4 files changed, 493 insertions(+), 874 deletions(-) delete mode 100644 docs/assets/ws2-cross-config-before-after.png create mode 100644 docs/design/cross_config_implementation_report.md create mode 100644 docs/design/cross_config_logprob_drift_contract.md delete mode 100644 docs/design/ws2_cross_config_logprob_drift_contract.md diff --git a/docs/assets/ws2-cross-config-before-after.png b/docs/assets/ws2-cross-config-before-after.png deleted file mode 100644 index c83db1cea0624b150234ca0899f789e3bba3cda8..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 93735 zcmeEv2{@Ho+drZVl_;}hNHS!br^=W)DzmU<*oLsT*)~%&nF&deqRgbsV+tW5^Gq_d z%`%Vwwb{E)=RNQHo$veVJH6-so~vu0y`N`2>sjlb*1higcQ3aus>&ZArX|L~!8xF) za83;eXDm%x1o6C-QbZUf%kE5-IQ#{+udVU-nSmEvWU7ToR8(sZ{s6C39{%Lom`?_Yp@jSW)X-pJNM1z`%a20fZ0ce~)@K{r4T+|(5_YF>VU z-7Dty*cV}Dg1(?`>4u#i=0S5OOH-Hw=HcBZ9T5m?M@!r9FPb3WaF_{(c$glI?ClXQ z-#0TuSYxJ#X=7_-j(y{9@8}Ea|3GH+Rc%XCM~mH&i0}!2YpnpYG`GOMnO8`3_l}Jb zwlQWV4i-kH2$$XGcP{;wQ@e!3lG7f60R8=jk-xo-ZDfO#fmx%ejin>zJ*<3k|NUpb zQo|kw$4upa)18oxg`xve&5ln)TUhXvsMNU&79yv3fgqq|0cQFSB>(nS%s?CvPWC1k z#{6+-yL~x0x?&Xq-K{$2C&J#*0%4AT8(GVKyDDvuaDtn{(EJCgA^ok*c?7~1T;v58 zufiN1T`^K<LI62Js|n6;6k zrStc)Z-h}6^RMmD{Q5Hj1=?Wb@n4SxZR>s~wtqOSKafYPWPIcJH-4g7=VD>$2vfH; z+7TNUV2-}$9k8|5G6-wLuC3%ZGcpx5MRU=?(H>z1!`^`-;D993yTo@kvUb7{h?P&s z8n8jy3<1=GkH>8HE#Ru16IxZJ&pVk|ni_#SG6=W>0*K_-JK)+kA;mUFcXP^N$Jn85 zHLtKOkR9+=?0xXf9Nn#{B`|4xJixe_S(<~3J7WN}-W|qn8_ZZ?_Kq;L0D_-8l=|0v z_(Pen_>CbG!l=K?H>1akeh4iPnAh-O+hA-v9}njKFC<3%yCeP~CTB<0e#=e1|36L6 z4o&|D;`NK091MGHjs9^l`Y|TQl*a@H0}k4sG&x%5&#R~#+1Ogc?153)(SNk^Isf&F z&i@N4I>xK~A&QPM8sE(Jud3*~#_hlK5q4@0&=)lKgjWD;hM;QMysxBb%%c(jqCID$s_+an< zI?2)9Jn;|mdoi*2Z;btwb0$A?ariTRf6KmKL|EGD>b{h9?Vnd zbl`st&cMEG2cq#60rIbh$q%q$KNiq|k&mC}IRE>`<$k|K{uSeL7*G1=vCjFpD{{@3Ve3;zjkJ0sin?Urgf%`8(Onf`xkALm9VqINOj_?)1%=3#UtoeRq%5&GS z{Z$y%UkJv;|97?UKRp5Q-?Z?*3XbU??E5bmZ~2kI{#|qUS84h;kfJU?C}H*fzXB}& z6d)7-4psk#>jEa+|MPsgf8U77@0I-*jF|9Y0$xAFnZqVBfBop{Pu{rk@9bdze#(wz z@Xu5B|6RW3|Ek{Y-vU>_%lA7>(=WJ}T zGSBqa!4~{o_n=q?|Ei1CH!tb$3NGN~6a3}T1=w>we#{DlneI2%{s&$Fa##QIANr_; z-@yxbd9e`5za73{cWM5aS@%=xIt)XoWuulIi!}W7&0spSaWq}`=-T8~9 z<$rYWE_PY_jg0;84@&#_81g@}f?^S`Qug*npbCZuyK%Ec*OdX2`{xdR`E4NikHW^n zzY9Tq-8O=a?4LSW=D&D>DEdCS%mHI^e|G<0t9Jf-~e zwF9jA{Pu?>emkR`yV$>Ov3mpa0Q!$=>@G@m{6$_kIE*-o=cF`{(8)fck^_s|HFJ%* zhnd&oC>2SYRd=g7Y7{up)qoHVPS#$$!cwn3?7Ejdw9$9)1z$KO7w%9?o2HND>DZ|2H2z?Ac64 z57IACTiIJuIs{0Tg(oFRU%#PNsN$aQc-xE72M3Q3{Sia_h76kR^7Z%5MZ$Yzj4yef zeA1yIxU5OvT(A&-s}5c)(=2ryniBX-wQnW%7 z&yOmTW8Xum@2K3%656Fnm0Ntarj2r%U0vsFSDtZUx(anP{AgwwVL+Ha*)CTq_JYCC zg<9f$$D-NCl5*Mxchd~yQ}Zlms)m}zwy9P`jUL!h$f2mT4YS=bL-*;yCxrBBk*!_F z&~ukG1?R3fZ_<8tv!8f$DLh}{QuIP7{2cQmI8$x`qQfyQRZl+UxAZ^{L0H3A{P*;D zJ05jO-)Sg^S+@Bh1Y;4t{f?lFp{hQ!(lZaQfsZvlvE7!A zhrqiM?gZj~$GU7CFe&4U(Kf-el`w*Y)0nxTJGl1ndoD{dKExUF?>|idJ{DmZqKq>F zvdK`;?7>oZXV6ywiT$e{are*<`PUOmyI|BA-2uUAg72F8NQwb|mbwswkBR#6-Ih&K zVEi{Z`1WGOc1Kt%j6ur)DT6%GN|I;_ziA8yb~)HL1yJp-#iOUmm-)((DD+?<4c7!K`pFnQ#=rZ{+?fdDV`vxARJ$P<~JQAzFt zqdJ=Umijj~3n)tf0U?BIuT$43Aw*y6P!U$*(H(TMf8P}U0})_&sXF*DFqFHw!%Cln zU{qonVZ7fAoD#Wrh^jMNK2bc1Za) zZ+qe7<2y!K6G;?=kP?kkYr96wM*<9NL)S$v=sS^lpk7fo`|(qX`Uk7GJgcTd#$AQx z7qa_OqqU>PPP>G!!^=x>lnGeaHY+Wy114?F1%~RB0+;|}V@HnT}Wljew zZ&t&M6==V=(>;OGv1x}T2)n56)qa=qyQg%eM^&lmw2*dFo+`)ZDsW>2tQBlaS>!|N z%$1{S0~^aiIk`*EQC?8wsXR~_uw*pF#ubPq%Nr%1S!S&qYpU;tGn?(BIV2nMN(NRv zX85+Fvh1 zs9m$ueCxt55pI=L6h8a-z>uuKI>Q9>uBQ7q0iL((=CFPz%93=roRuVtleg>g2Uum? zm+z{ciZ!=*v2pWQQe!sRyxcZb(i(;4r|Y8ebqlNU_4R$lN?>@Vi&^E0H%CcpA2IZO zO2LIFr+lpX)TMg#TF?D*)0n`R2aFP%TgCczpOZdh>Z*6;S-HHmdGekDx>5G|LJFmK z{CPTlrTYQ}EU#xjzApS?+OMrV5_`<+uE*+pU65dASpU)P+9WEd#e&HwAj67f3vBX{ z@BBOE>#(Z(@t-~tk0e!<9!T>1a`U+M`v=wbE)wmkmxbZSPcKb|vB!IG^teog%enM& zAJr4!uVz4|q-n+-T}tKJg8S2-bw1xZl^bW>qBLEW1hM5a!KubPO6{h^26}#An_rn$gK-4-zlw+9zi-Gq&0f&RNQB7y)s>n zq&pP+f*MY3ZwB-GqrQBh3gw7~GUxXpuLF5eW4KeUhsw2JrKpdLaJd=hIa>8BfKvIz zvm&K2ty3eu#nX3R1L>KyJAC}YbZ-u7%qhoVB){5zj2I%P`oPqWjnlk=2hZ2sM#JTU9!nDFPmK)w0asf(gIer|E1_AqIFXn7tyGXLrM$g}I@BOL@^k`xl8tEVzh zLhd@mqY{mTmUj$neeJ9=5=?69-!NSdv`f|A;!IQQQez!y8e*9r-=h-flc*f4RW0YlxH4eO@F?rqp0jX$o3OP=p_1f*oZKU&kFw8udpg=aA7)d1yQ*nLy}jLftenOh z+OdLrJ0KuJQD)nv=Ol&WyDj^K^Kqie9EKiY9*eKs!%ykuIS%;~v!cz~n&Za6)LkJ& z&xuI5&6P9mPoy5kUK^@f&pIe99aS;$iu*$S(}q~sREcP{AwgvTOL^u2XwY2!F}kGc z(xB_wUL_(|UzV(H?VDfDsgmO%N1he-NIbu48oZHf{ruCrFqVQ-PQ@AFvH>?rsmP#y zD%Hdok-72$ED;TZ0=q^SYYLE36?bMoCUp$l)$V8Dyy>;d{>hb6kgg9>4(~l;s-UVG zMU=Hayj_Ks;B~+o!L-B**D?Q#Ne3apMAwKRf^HZm^{a_+rJZf%{Tj+&1}0Yf6i~!;2`sjjsOr1(uBK zfIODOqm9R`xk=2WGnk5!JJy$<;{L9m1&~tC9m9SHO0;U|)Dd3Nt?aqpWlm%0kKFj%7C8U48 z=i2J4F`2}p<-|LOt?F*$R_ZU+1i~blJDTBBh1WGVff}mdDS# zl%}Z8uKD=hc?O~51rLgHHKz3oKSZ+FAfM00dw=r3^kk$()=)74k03CON24Gc{`%Ow zQseZ8(kYWyv&=T}Z4NwQ%XEh(ZxYkTU2VB6cc$ZXVjr<^AG1xYpdcJVK`6r4bD^J( z_xydqk{X@qsEMoy$Uy9oW^P#|85tet?t1Uz69f#Zka5eA?>vq{3SdwJY)-q>Cgt}t zAr4Hkg^1gF-;ilMS#ptduz#s<-m(7KLCDh700efmBTRPYLX?(Xs9EX0YpB#>+PbQddYU06uq<8+NSkGJLe71B+nnP)6~Yu7JKMaUr9`A~ zNPBpuEt1(&%sCVuB0-z;fEdy(=&cYp)2oc->nT5oOoJ(R9Dk8+!w;D1~Nn29L&vEtZyj1Cdr*b%>{kx z??Ki3n8y+4r0QK9$*|9z+uoxseU28IuD|}G)iv3F?o428=Gk@n1B+{;HEoaM^2q{b zAIl6-!9LeT&aYobs?CDk}Fn=?G?!w|~12s+`~K;i1QXBQE? zajr(AbtJOeC_fk2-~^(M!16{Ld0VJN#D4OgzD9o8^Qq`! z7mCGgi7gM`lEqgVt5)?|pSKqxUiHZ7I*n9CP4$e_ogNO#PDG5UR%=o#Yo*_ncVCR^ zKiX3tk)cOp!cx+-(ooZcal|TsZ+u1mcNjdQN%D1l*TRbTDf zk7W~WTJ|Lwym+csbdG2^OaA=5lUZ!%!%o#lCr+rNdgKVUo^OwLz!nnmMO{&MP3xmn z?hnvGIu`PGgtQh#NW{U;IO-Fyqz2Vd86F3+2ss+(ed)4b5ri7L z;63GZiInET=6w;*r;eW`E(ULxylr_8v9Wn4zFld3 zEG}ANYpbSUjag6s!I}OADVo)d7q9lkQ)N15?4M|<*MWpnutB}CE1ZW#r`v{kNfm* zOdwV!5W~(YFr}!_f`Wp3FTXyOtGl>&|Zp+ai+T&&?f)zV_fg@nI zoSnR%=sQ`urFCVXF%9P0vBVV1uJ+kJrEjXxLbzd#8;G>66xlc)#tJ@A2ixbPSb@jr z8so!?=LV-ur2?7^hsHS7G_JZ#aJmK@k$pFe_eCsEN%-1c^)V$7-8w7IR}_L-jEm5r z05!$x@AvmQT#qVAD2x-mB6mVAM5WBHKbm#v2NLiVp`-=7R?^yd;V}(g{kc)FbiC7h ze*8NVJ;;d(L~XVN6{>k0PB}8rsZCk5U&wCQO_$f+wrFU!gv5a3oyhL!o8-ZcJpM8t z`SCNb_fHdwB9nH1ft(ngj zSaH@m?vY_uAWScph^!T;*(MHG%jF)>a(*V+d{r><{ad@~N%~Uk5^4zcn=3^d`x(CD z>D$Yaed1RquDRk<>#Dn6CAvm(Au{0oZs&uifq5K=e1Q9vmOFC@BTDAv$>X{gJH*e! z-1GQ!$C||#RoC%Kwdu(R1f_M_IWILVnl5tWTOwYycj9BoeuxgYGN7{{NAo)pN-FzS zXi|&aq*jSa57S-DY|ixZqwJ2ni;Ia_OW!E;)&>5o%|>tU9a|Xj1H4 zoyioN@3G2q)O6F-Gj1*cPOACx>Va4a%Y)dN-U9I)W-*C&IgGW%458^fBR99Inwpvb z9mVUfp{WvJU57H1PErFXM9eGogro$>kTe%s9VWlqo&go zKsyjN<21i(w+Art;G>qWav3TIw8r3fK_?tv4bW`zAa?#fU%=1z9X<@-2Sx}B(9t>B zsz;odhJ$)&)@T{85w9{ZhgqA? zG$P-(laxabj_N;y`$`s{A4R!1%8Q2cb6DNo6-BQgY}4`ZV8K&8SsGIGhnW~TCpASDJ8ZdkKQD8gM)duy2W(H=rZ5UZ>a z43O%g-+3U5U`2n?Y4r*1nITyXkMyt6AjXPX!gU}E<0mTYRPyieG3>mW4YYPr8|bk) z3`G`uEqpaBw)#F1WHri?`wmHhpk=TX<@FrVod*m_W(!dh^YVho%7Ek*9gLF$Zh`ho zKv|nqe1`ACDPU)IBubSSv`#F`?>SL2>%2LB7E9#80vtA$;Ao$NV1`wA^c`G#VGcVF zlu-^kv~`SxT%Gl!x-|Ek5Bq=*Xg>+&`Y0C6wU46?xntBH909ErJrVgT$f2pVW_K(# zE9`Jp<>(=}M|Z-QBAL4T>qd7T@Hka5mYP+6tZk48^7x|(4Eq2MXx|CuDk%)+TBLb* zZbvso*g@-vwSIW)Y-P{lljW~UGK}{iPy~u62FwK|iLL$%>bng)#%|@ZPot(`v_$Hw z!_ae8HFekra6tQPQZQFIn5$#j$>AMKWvd2Svn`hP2@Uz@Psq=Fbth{GL}=_v@Ew7v zrP1I*9Y)LXpD}5x%T0UWPaAVRJ{|c9ZUqR zRZ;H6HkBlZj*$MZW+a%zmhU5yXdDQG4p)U;gYiY7ooKb8-U~=FiesE# zA0V{ls}u)wEh4hjZ^?KZxDRcVD-=QNhpo%bJ$_XEna4`yv4)wRYS8q}9C;IH1Y+~U z_p2UR?^yAQ4j{^FMH`n6#LY5uh--bb_uJ}|Nk!0Q&_aRA{iscO< zeTRlcbGf(oaCz&{jjDzhyXuGkl5)^PYGS_H(bkWOcLX;4wW}@A)wrRF+%2DJyjWpG z)Ex>=rMNS$f!hb^MXu}@l@F#5zkByM*bo&R<`LhOd*c%otJ3c*atC;Z{FhkGe6W2< zS_1*RodMAUJaxyH3#4X3DFv{|2m1-U)udWlX(8S*G-Y@XAme&HIw66x%}SFB#EAsN52Oa!5S^UM<4G03>6 z$~h)K_O6z+7>d6vFZSUqo2p_BMR{GiTgSN~qV^VjpF=WV@c2R2Gh5&13qif2=ZjbG zUI%A@$G(@Rzd_jgilQDmY8tk$>JQ`D71}HJx~{V7`CcQBX3Jt*JI_WYIhc-9V}O_@ zqIW6P6?Qucw)b+z!*yC26pRJg?GYJFJ{7vSc>8xK%V-5SpCe>^JCHyc>;W!6ScNpN zFb72o6KL3mW!RFjycq;MnWh4B*rJ8C$pjbc63xwc9Rjl;)m&Pl&gZ_znRe4`Yhs)^ zlACJr7z7ZNF$3g`OYOPr;*Z1i)m00c5*0<0wUsN8PWK|9*${+$$v5g0WTy8tw&0yCXq!qeY z$_1AB=(K0v1N_CP$Ikg<1>Mry^Ze5ITvVOd1b4B%r;(^29Z;bld=bN@dKmK_Eu|c<6)U^ezS(q$I_ab{*2^g+|!Hd(LQb4>&oa6Xn$&m_+%2 zF3%Mp;&PIG(pT*UGCvDFg&$3l-Vfc;>;&+xZ25B>hoU`JE-K9jd$SnCFY0rLmTKrS z^iR#xHlN)ldLU}A1K6YI5?s929u)=yUJ;@DO+B!WgdNc~D9F9MK!S@PN~abT>60m$ z2&83`O?yguzEqIUEdagzL4bZPKK1G{)}QxYR|{XBUaCiW3adKxoOGf)3`Q}?;LE>$ z*nc?F9cAQnL1g;!sk)^S;7Q%0O8wL*P7YmpF&#*Ydm8`yl>5V@8!v%>&Tx3p#0_~20kmhtdQ%*u1kXbQGdb*xhW!K>;m>~ZgtnJ5_s$U+$=mOAGU9hJz zwRuAhSSBe6Ox2RV0`P5_Bfq?L^SC`4Cf%%Jh)xb0z=Mm?iEJ1UXhQcJdZ113WtugA zpJni^#85^K_!mpS@?43}g`Wa(G_7WV;UQ%K^g40|o!+Utv`herXcDSeNpiI(I@6NftwA#pEN|4!Y5$ zdyd|?A38MX;p@qy26pZWBWH3L3?I8DUUY2H1})s8V}~=*KnNN@i^{kmEg(|)S~t9y z20g*InT}90;IVoFCG0fM&{OuKD$*2aK(3Ihpx31fU=YY0>!+t1>^cM}NUk0SIw4op zYJJj=>ORHYA;v*w&;=F6^-R#k1c8X5t>zl1%s^#a5(|N(_0+} z@b@51nT3~hm4uK%~py|0M##e=oRFX6h*V!Jj4Fg>S+>z1- zT`U4(6qKFIoO*O>sA^dvBa);(G7===iYTQ!(6h>ngy};m0gvq$bEZ$%^nsB}O8~uF zNK4^nVJ>M(3J+xA_cf`x@V0w#A_+ATIZ{&M>YH;3XsVp`eXvt(L^FMrTucZ@uQC8n zTQl0o(@pVNR0FJjX^MOx@A{dpqR5wst&2K)B-}O%HgbkYQ#4C2>r8YN&UO!U)Yl`* zP8o7<%K;eLr_Py)!~rXlPcM?{O>0}1S-s0VoSo%Y$Qyt!AX9NyGW~Uc&M{7tQcv|; z+|XOYnE?SZ9l#;1Rd{F7WqTKh$0fY``lW6jEUPvxCw2f4L1zl>(b?D}uc^x)(mYSI ztLOODX87lK#_*e)cBLLmgl0@!=ah4JPXnDWkA}An5e;uZW|fpwg|^ItH5gFaX$Z+Y{nd66V>?h8&U+A>cU1<>J^V}L!OTHqdyIA|~-bScZD*%mjcFv)* z=dbdfh>i)I9&N4PoUf5vof&*JXc+WhHq)bQ(n@z_>r-fhW|n>7k$3SD;V1YEZq8P! zz0VesLy}Qak+AsqGy?UQT|IcHYSxAn@_0N=z`93jex!Z^HOssm*=DMc#qc3%@pbe5 zwj8Hb@&-O7AwM4&-t&VHW6$)!#idi%oh@YOU5;LRrYn!$Hr|b|Ze#*RyK&miGh*LJ zYCT7-pY|3N6#sG@Itk;J5dSV-1PU+zytn?efXBwUOGZ+t*~)pj%HF}W!~K!CCjFr@ z6YYM5W-9@Qp?(OnHd;R_EyWxF)|$Nz7~_>}Z(VI+iA;dAesCYR)!pbL7WMj=*y|>9 zil)2xnMO&W*KFP;R8Bt{^-wDIGNKsOPfy==_rmDg0?U?>{GiyxbIp%G%(lDnWlyw( zpVsD%S^KO{BVkkIQM@+!X{MUCL!rukVMw`Bd52tm_kgfMFuiD^^Ws%L z)T()s#@xqOi|F*G=Z&CHgqQxQ+v}v%5pXBvwCV$JQtlR)HvogT)n1;Ex8K!0J(=iZ zLx%xpNbkV8>EcbHVJ)4Ig-H{x_@a=i07~9#Q7-21?h@3O6A>BUX`$0>GV%xO4*-Y= z(^+?g_#u{0(?w264VyuyQbY_N>sG#%JCv=#XoF1fI-0DCdKS#U&Qh%+5KOg?^jLUo zW~mnRbC(J?mtMO?IY>i}T$3KF=R8|{b78D@3!)fI-<34lrv31XPKp;rr}x6dd>=~C z9$Mz1wt0txFVWnkut0isaH$`Oxbsf8K*6g&TZ%(7{bc?C^15Asgi~+eCmoNDCpJmu z>v=hrU(8$)?QGTYK;$DGJtvz(Av~%N#ZRH4BkHv`8PnP`A_D3ng@a#g#(OEvk2IJp z%%D?{<^3Lo?%JPuk+;&f7KJ}IZ7>ZhT=2JPA$Edx%mrbS;ordLdOd1(t z=4+qjcy@4GI@U~mcrMFta`4OgQ;8tXUiEnIltpox^{68vGbN5=51QZYSvn?myz!uW zn@HnXx5;~U?oZizUWOMVQev5GwTAq~@;yc}Tuk3Tu9{zN_Z|r+;Z%6igNOhaA^%>t z`gis;{#0Z48@pb}uRNisB=K&Ixy0<_BM24)o#FAC3>`}7gELM_(-xb_lim4^E97Oy zg*nxw^3mxl(6Y;OrDVp*=FsD8Y|8x{+Vv~1Ap6Q69h-kIGF*5F4GR*6x(9J>dPOGl z4EOiIBWqvZyyX`BCN;i8F?wrH5N&Dm_9PXOwAab}si0nweWYDdb%tI~mn7@== zHa;KT=BHNL|4>}LHCh*?Hkca8Szlww#j^$V#7#Nz0w9`%xdXgTZo!A^E0-l=SiDc^ znsPaau`8lcMI0ioN%Ag*$sVF69ua0Rr6HS{RtK}ehKX2G74|kp&dO}k@ zR+*kZuzxT09=USv!%rJ|D(_fjCL6!(GkJS4#d&pPeoilo&g7^N{Mvr62^&!fZ`>67 z&#d-6LBa}M;jBrY&$`swWH69X_?-D-L04Gd;mQ)~1%IR9TDZFpmposR^!~AW}E&$$U6pKlUXFK#}HK5asm;NQELO zm#dYFZ{DvN!;&vWzI@J zBnf&TqS)@eS~wT8{(%lbzEX|UGF~?SbVnHyX($XQU%Xz}EtpzA2?$(tPnX*U8P?{( zz38Ag*>!rs#(!3Npqq~m?D?QzuP~vqnIl_Y8NfP_b$71~1Ap9vEi6e%N@8&qVYeRZ~&phG>PW9&xA-G=&Th2~2CpmJgNG!WT3lH%m8UK`2( zjBDNY4gf4P=oqRTYJ20>u7{vD9859f^cL_g+sJjan+=f6eQ;l-Omt*VSGI-kTvAWz zlo-G+K!huA+9xT~^r#uBB(1qnmouoDkd(CAJOR+19zqCS*2|%)jx;ZxGK+P)_gusE zbrlm+q~dp9r_j8f=f@5qh0sF(iTi+WNlYGd3Icb1z(K z=o*Qx_C1MyiCDV$mco_jSYy<|#`2Tf8LO{!wEd|X%f*!3^4xn1x|wIyw6Bzipaf}z z$45eM`ii%&iw<)HBA$9%aM6y}&kb}Cc(-*HsGG-Zx2eUfj#9JmXhU+}N4iwg4Y>xm z)6te6@?1-j%K5U6NV@n6ubK}P5SZr~M*{VYb~6Vd|0%tmjy+?hZ#qU?7D2>)lD$#z z09}0+oiMFL1z|}mDPC&`k!su-*v8w(K&R-okHC?3pE`1@T_3uo98+s8$bXJOalT$D z4usQ1Hj|DrKiv0QXq>&5?~5-ui)v%3sC_Qpr}V7-SW^rEmeI^>%l{?vY4O}F+)3WI zLfODIZPb>dZ=ZA8U=200B&Dqtn!+A$I(g_!2 zH)Gax7zjJ_J@fVNP!%#3u#YSP2w>3goWwRfqI+3bP+YiM(2B#QU!U#xamC`phoLL; zHdE_T4tyWi4JvYR{HT`h%K2|`Y&4ELZ9hElsfz71owjSLcEONrw?j61cQjUfhOC5H zymywl*VTvEOasa1QhA!`^}6fk-r;sF%SMwnX&?*T;yv1H{o2dB z$+|rSVhatpn{OuVEnf$8kx<@8ukb)7l;no0phxjIxTIbr~f z?OfN&64-+FF(h=g^aUDjKRo05f$)sFq2Qw^MfY@#{D$JS=^JBjE)88Hrt#8Q4oKc- z_1-Y&+(CLVYyF@mu+Lx`%IPw(8@%W2nVXuk@g#+5Y9Bn3g78GE%}Q&=llG0J;v*2X zH-s_XJuho@#)@N((KN;r>br(5Deun>T$BBjc6dM#&YDo208dNu!cB$xjH=Lrzy4B18GPyCaC~Nbu+RDbY`L=yCG_lCfi^8zVg_zG?osZzG z>049d2no9Sltvq;33(??ca*k&EIcf);r?F43p!OKS?w7{+bTJ;VQ<8nY{xIY4uGG; z8y>FnrbzB&`0EG!<_Cvs?98I;$786>Mje(WY25NFBX`R%K=|Ry$G{QW+m0c|r3Hze z?rN=6md#~xu9l&5O)t?1&K&kd&Q9Zw_PT{S5-0a$hU+)$js;NSBNWLyeY&B|B;15s zwC!lw(B?tuQV#DeH|_nvvZ$=`C?cYf+9oo^pk@X3{J!l|uM|j0E5sUaLIPAAC@LxO zF>FPxB~=FWxi9+%cAZD^8dshhPH?aqYp87kK+EpZuthT4e6R%hyB%^_&I|E(gc`&= zWw76z3bd>qNK5{}QT+`;AVYgPuX3o&D=9rXG!iCYj}4iv19^?I%U#EJp3tG#l-wtEo7A&9CA| zF_rM-20RZ%$45czFX)C0*g^0%F1`z>maYw5PD)cNoBZ#<9^ z219QOz|1o_Il{eM2jH135jRSifmEC2E;->h-u`#9KB?x!FGYO%N7+mjmAzwVC=Tr+v z?Tbz!$}Y+(a`U?B=5xOdywgcDPHUeD|dM{djP zy~&l9m%H601>nD$MkwTZ@8P4@>yvvRDUY=K;3xp|_jBUXBH8Daac`+^?_wYc1zs{I zUn>R0^?nb%rL7gu*|_$r1{A%?dL(FXQEw&XqPbLM@hh*ih{fJWV640!P1%;eyr>OL z?jcPb&T%vQu+U*h;tEwLrOXKA+^cTr`K+6zWVU)e6nYsX2XBGWiUxGW)dW{#ZW7R4 zUvtG%Zr(Y)9tJAIFAO}Zw@V$L&-G?ha85mV^?8<=hjhGYSys2mK53(D@I|OYb;zja zi6?Aq)Douk@eo78#zIh=f@nPG-G;2zS>#N`?@)FxhP|jyIske1er&Rfe*q4h)}1Oe zW{NA&2hZub!=^8#o9iMIrY}IxmW62)WnWwP-1@4Z>g~Hms#`@H+M8EJUI3F$js_@C zcD|55allhHGgDtzT({t&T%fnM(vefl0nd)(rzMt|sz>ytJdv-1Q9uOX2g19Q{YkS_ zlnyzOuHMLRkJg-MuT8!kaZZJJ52>EIpQrnsE?RiCjeqvCh0-(U~q;JprnWo_gfihTUbqBpu~dt`=v}#<#+1^s@M*;bU3O zw`(8)5~_>5v^ZZ^?clZ(8%ZTztltvuk2Tg`JVa=xiNGUG)7U;eTjhm7ick-3o^rt*?dY2JDJAN*4`X zwI)I%f^H!w55lz_uihh{wf>OfOjyixZ z=Y!X#W%_A1z6A^YMTvU_iPEk#i_@mVy(*=ZrSrzqAF~1*KjrUT;;#%-=FlMQFG6k% zMx%p5K~rFt)S2U1+&mDQnZe|F+J!H(QYXha=}dQ8z8&i%fNCenh-)JgGPw&szAsL4 zoclOG-K3qURz(7~x7z}lLYfk^tFBwm(4nwljwAekM($3wO^LQrZKA|SP?RS?3@eP( zkdFOSk+X0am-;gGfirtp7SGL>U*9(}C*D|Vt{kN}m>`7m`bhNp=8X+n^UZi~n*8O9 zh6*V?Qb9)rcv^^4o%aU1zDDZYUz5#ri{vz@>wJ12xffcmwaI9+zCapd%>lVomOe)WQFGx`JgCA@~q$G($;8_sr53(mN9Scu8$wj zo!w>w)lAjAmkf?RJ3%+r?ImZ7qCA8uBrCV>JjpM<-rM<9tbNIPKnrYQnF^FTqW9I` zjDOEBzULslIM;G(u`Cp|^pfn(J|JiHr5%T#CKk)*t?j>*@a5x-7WHzdxbsK9IoBs} zha~Cnic@LKn;)c|)C`XDeVTgyICD}x{%~CuDDU}Hs(n*A;TnPTjrVrK&ND~zCZQnq z->XMgz8pAP8FzM2f5;#8N<}4>F8(2k-=MTra~-9tiMk29n%TS%SA1^smdV?mVbk?H zFY<3{#P^`&?NYMdYkOy9Drq`Imu-6|B^MPzA4YqYEm*M|JT9=-S@BOA^1rIyH=0sf zR9`W$O%3XEH9C!X_6vJv#n!7^}G_Hj$?6W!;ynN`p(U=Ov}L` zPp96s_mx5}32kiV_i7B^4^<9n%L;1CFlf!V(3){29SO206nCum;93#zkwCP(#q=_Z z(!(b=XCy>DJ|*jGlYE;gdfLMY1E{OyPHD;Y+^N*|lO)y`2I>r zf$@S_AzmeIb0&uXgL7VnU{5L_)FeJ`x1NKu+|-|uSu&pUNZaqSxQ2h z;u`BBTzFc2plns9S_~DCLnT*sg5^XBr@NO96Tz1;;oOXL%8$w~)BRRl@F%I>a+JJH zm7i&hW3h>!;si0gBlb?0wI)R8W>Q|0uywmURtYXvO>TMG0z4@Z1KjRLkO#fyxd#{D z7+`0~ICb*^eDZ-D09c?qsk7k>z^|m%R`4Qjx$&$xX=Vv|`IlT-6%H@(E(i5d8nk_O zv9QvQ$^!K#yAI(+5uKxhwj#Q^3#7Rg_oSYhcu+;hT#}6j$qlOEs(#SFOkKnZK&SX1 zI|;IP0Xnrg2KXhXDoXAFn4Np}y4wZh&=nQ*?u5&rP-(00c}qc3TVH-_B#6?OT84nC zFdT_9rhPW-9E>~hvrE+m5*C~8Le-NghX(1YUU(LJnwz`vhYjL zP2aT8b_gW6*egX~`248~McVc2e~^QgFPgAnJV7@LbJI#t$?EE^pDaP14h`jU|M z%X45B5C1cZzX>mt^ykaYve({r&!ygfT>%2%B9jz^1?aOrJO!WU$kV%Iw3_Nc$p()% za7NA>N)SOdv>oIEeGKp<*a9QGwYf+wGv#Ks)XJn}ii;${aSMaeyVBZw>TF+9!Kdw| zq~kMvQ=g@pR^I`n3nx<-F_kYN8Wdz}=Ka{V*XfcJ*EDgB0RSvNV8a(t{BkbFqpQ|a zA}eadpyy^_DL7Ij3E-SOKEHZICRwrT=y#Fhi4pNYPLwoRdON??`3U0Ecf((98VWD? zi@%X%bpLE0e#+aI+B!ys*$33~vCET^)buGV_LegMFjdDbhRcu&#rum-t($3mL9xNx zYm^D)*WPO(d_Yqh7IYHS^@xd$%)WrO{%N=gUT;YtL-A0~mrzTr0NPgJ36NZU+J`ya-q>{z%jONq z0*h3lBe=cbI&rNq$IW8faZlzhI&Fy*pG6i2R~esNI5Qf3u1F7MV{x`$bP(g9>kjMU z25uV&;jok_h}!9gFHe#$vO&(`f1DqUn5Vu{6T(oV7j)F_VqaHbJSBMWle5F+|-MKntf(zwrPT`UTa z4BR!c79I0bW&jHpwDoRC#@=iIekq@SXvEHU5@5;Z$ip)#RvC^F_Q^{kaD$6!+oFdv z<*H0ih7sC^Q{JYX=wuJ=YNJXe1sFR}f<&@qX z{ALpHo>C^gTcs5@{OoeXjjSi>9s{`kY$UcYr;L{Wx-5?GyPZ{Xb}L(p@p#F_SS`|N zafex-M0)X)?23x*dDZR4Cmu17(?ys2m~tY;^Y7e=;iHpd6aw5c8K#}cu^@t z$6>RCrbciw$<25V^Ce_>Z|5M;~A`K$4Be7kS>D z+aH+N0-&G;={;dC=AUz5_R4)lVTcMFE3^6&kFN#j7qW6|bLTGxZ-Ed6HF)Lg3VZuY zVBtr#jHr4?K(O}%xA^P20f=Eq^$(rfr21EYFdiO1jZ}Z zaT5FwycO{Eso8ZWJgA%v0323dn_LWi>cKc`#V3_4H$GZ?KJr1&_g+IVTnbbAp3;^A z^5M?RhfoG$&-{yzwtcoflpA`J?LAd9>^w@<%8& zNTOl!=J@wYM2A(=2s_y;)_YZLx&gT6G(g99C7h=NNHrfXJP$8`O3T+ieVVzF#C+2d zeS}97)TcIf2MA3#5%n@9e1HM-@EZf`)=VD#wWs-n4p02nhE6_sIHe zd7nm7klI@k*nTOE1n^*fHGoK~Aj3U@2+nRfiI+z_E1n-2TMe%51=R{dj4f!K*P{fS z_+x%lwA1au2XD_-Uh<+v1M(NBuFOC64x(YgD{1WdeSO0c;q0muK#IhT#%q16cM@T) z_|!n=Oe|cuYT0N3?u_@<1_d;W8%NslNwkXg$Eq1(;M1Iv0T_oFN2w1cwXy!*aD{lV#l7~)~im`n|{%76qdtU}c z-WM`x*a2k*ZvkS6x`J3mELnkGI(4hS_ z5+2u<^6J1tfC&Fkc;b$aW@@(ig)bL4{iwE51kv^tpPuZ>{3zk+Am%G+I5Q48Klzhnl z$%a5MTG1dt(Y9#;bQl@`;`g;ElVoV7f<0=3Lx4{((DST zFG$k5(OybCxPIFXqvzCIy~sfU{Ey|@IE9Lbm(hKAoBm% zd+VsE!>xZ*5R{M<7^FKC8A3`@TIoh&=#uU(rBhk~L0UjsI;D{kK{|&9iJ_bOjOU#7 zo_DQx-QWG=U3cBP*6|OPFf&i=XYc)d_WpkMp52fq-3Q>wN74_Zb$<0j0A0_Vz0{==59ZT0Dmq{7;#g6XPOHAA~$B zkN!-lHMndV2ayKyMW!|#?JX>vL0#*X);|SSApL|UhQn;Nc$_WG;!}|9X)lKSPvd^X zbnxMKAjiiG5~&{@wa@K;o~>pSnpI7c>o=IFY{mZ zIv5VW(e3~8ng4+|erOT^ZY&r}D`eU9%=|0Th$ZR4SY*QE0LnOnW%ha90yeodeR8k1 z*d^%jP_BWcQ)R9&Rcmh$hN`8%UxPFnW8`OeOKpfzYOw%~+Ruz%F~q)`4EQFoX8>j7 zjF{-7;q#~5=}Gx|iubqL1Cu%w@@B(IO}VsT@hNiJPNdZT*7?{0#w|nT8;A6$zm(rQ zApZo&@Q-`EFlgQD_R`Q|E%ZwJq^hT4lB9SGdFor>54WB z7|61wzZdx*{iQ(PL5;(5n2lDCIR@qO4*yjZG=24ojIA&Rdbl>2;2UnScY%ZzIMQZw z_!oE^vEi=i8S|>3f}O#j?^k?T=agub*Ev&@GX7)qovEq0~jp(;XZ>kL>FhM3$ z^EO}x@Q3Nj@O%1tceEcOmm~=`#_;CrJHW#JXJ5V7&|V}xzU>zK%YODhIzL{@EBJuq z9ML!9u84LGNRtfmz`cQnPA|(p!Qcl>egaQ)enn}!o>Zh%8GuyjRZJ++xF9s&4EKOj zv_u~wHZ~=LF)%XrkI)yaj1_{?ASNl*?44k_aZ#X9M32>b8JojP$sOrHW1{qxp-&PA zD}$`{POD*d@4iNZEE99WaIb%XVqe9I>C0ambp(u&ds-8lb=#N5(wLt`YAbuXUf!e& z>=oo)cc+}D;tnOm`6=?4|A>XO}S+!KHWbZJh0Qn<{ zKvV%|*>IuyJeT<IV6MFTpeieXXVm#C#%3K)#vsZlf0edCXey^_urmd$;(;vcYn;BZkMMo zMEj6gDfrhe29&DCXBrGkzfz^1-z|((Z_pd_N7g6>e*16TGx@i+Kxc$6S{zNv%pr|# zl!*zV3b~<(AsfNMU7pGnZwG5BMma)5*m1+toE|9^s+zVXySZ>9a$% z#eS^yZ?D2k0`tsj;a9%49w^?$y?%sm*IO|QT?#*%pwv3#)?eZ4HZuPMq&6bld45L$A zOh2M^t@WLrBLTKvO?Utv`WSS!lvS^{PA@a%6b&>-;SfEZ_j|7q`!Um;v9nd?N!!)4 ze>kx*)Bz>#d1|hIl{mev#Ye>mdRlJV<@7p@gMH2`uoxQk1nRcSPvh-IjqAPqbb9fvL%)8K0C2;w@7&H0*id6- zZNDLo1xn0$Hq&*XpxjA)zfIX`ZGb&R{W1F%7udQss{_}06x%{$Y^bwt_>UR@!46D^ z*)gwJ z-TTq;O;Qo5_+|u1pwJx95IhEN(I35j*Wju?8+hHL%6nUd$ESe)Se>oCqmieiB%x#4 z^D5};tJytMh$Vn7&xpUg&gH<{J2A8cPXuwP_}K@PgzM-({NXxOOHgJbNxpM9T3)A3 z3Mz;24gD@ezed?T>e7H1{gGz_4G%tI{adw^BM|qisfz16E~!|F>}s+R#L&;18>Go< zsQ#1!k6WquSJV%`fl4Ye>@VPfXaAw6r(E0CzVTj0NBh;7?xh%ivzZ@R3`u_~v7z4^ z&4|^%*l4XF2(q-xjYL*U*f!~Wf^!{>Q2WsraJOBd50&2Q-}h^K$ujU+RJMFrU&mjD zdS`1>ZLNez$gIRl2r>6aFy<6vct|f+zD3V<`hJzM38fT%V(98oVLd@>Y+6ify1-C(RmwWt_}z7I)Jr z8mib~d&JD4uriu!nYLzB@RwTN^{q!s!F=0m>SZ4@pEH)ZHx4KrBU@*J`e)=u6nFi< zMF$ul8%C&ZK3iT{`B@MgizpWmh~nf(lWFS#gLqY)fx@7#X==^pep`u6^E+IkzmX4G zP?&)R5W@KDZ%24h{RlwWh+G~f)JMn4J9Mxwt9kxgNI)hUzrc}sHK_f6%~#+Bs@(tc z@&AbK|HEo?Jb!C9%`#PKmcj)wu(|G@GS)_YS)kM((|Q|^8jLNs1ZQm5NeRw@EUs{P zNif%d!oHve_Zjev%Wh|O)_JX9Tr2dv@HjutF7Dha*J=D~&sE}nd#Z4ssP6eqLA8@n zeWhxprONGt)Hvy#qW3Hg}vSSGF!#Q?zq{C9eXac2&*=@yzWcm%LKGW5g%DB zA{Nq-8h(MOu;`>p{QZl*rT*A&$(hMdabmaxSo_WvrkqzW9-S}7TJ)PujFp@DZ!A>S zr)!meKJxLl$=|uIv&?F`vLs(#s4^S2yDDtqVDorE!0J09-Xq0Zq*t~b+|NU6>iss>kFf?t#{{2es zX!in1(;{kpYkz{%-trCuyKapGpOx3S#O*b+n#_X)`!&DEJ-F1#DftR1#%6$WtofAj zrlLmL_WA(6@-r>zL=C{4qdy`)uEx^#G(LF`UX5LBafuGjCZ4kb+_T^~KMDVvy`()# zum0B!e|KLQT0HAmdsS~1iD9j2cp3+j2tg7{ zR1dd`jUTIk4vk!2m7tqrmGwuDc*Q}tP$I6FKU0-2L1R|Xz%@kpN&W+ewA=19Z742f zBDj%~Yhp!N`R25Tx!w=(x|n==%yqe3mg|9tH>y4!UTe+WJ5=hQZ*I&p{C;TGpWM2d zm0yzQT?>gaP?I@$?s*X8@@Gs{l|d>5&sTTXR$WWer;+xn0%)q))z!{~-sI<0=XvcV z;AH6U@U3dWQ&-c!S{s%j!Y#R&U0fRP!y8;+u88ZeU8j6cc)S~fJ|c;g{7>w>QMkKCrlp!e;>B*N)9n(rfMze}e- z>~fc`((uP6*KTl&>fX)$m9=h#+;+31)&0P`l}_Iw2MF>qxuor~l$Hd-sB#vPvP~0M zf`G|FT!8^OeQ+csr4nL57vAXh!}*ji*?XoY<$i|GW0JuJr!-JS$|jfmsgfQ5!XH=9 z`YGekLE;uD%mMd*Qy99XU|-h;P@0E)d@oPN%P_eSmm{-^EFOmGtuKhQ`iv+w5HEGfpDe&lnH%C9<&zS8^2DFcXIl;Ox6mcL!_xHb)E0cAGwO*?9P$du>8N zjbQ1AH!_4U3Slov_+E(#8rxG{A@d_Br(>qjfu%|PUQ2%Wg3Q5%?=|$msfqK&zqkNn zqp?4-q5g^77V%~irJuTqDVIRE+4d3R*Hx%6Q2}|up{czlM(P++_R-s`9=5>M{$OV( ze!s5A%Ub7!1{U!NdDl1-a&fA8rW?%8(f)wLTQ}0~G|R?YXED?ItKN>%Y9uG++3t~K zLtbktA}(rkTqA>uM`ufT-%!4xU@Mc^KxC8K&se@@FE1L9U<-9_w~Hy4&<{uV+p*UQ z=6#?e+r0ocNZX}Mr4A^S@bNTCtreK({dpa8zImGtxJ^^*>QAco>>npVTeO*QyczX# z9zRgdYWru;=3gy}A8lA|JvcA>r7Aazp`#_uppvZqV8~Z)1+XFLD_Ox~-cVySA*)pwYxh;}*CwH~!x!@xr`v0DPwBK6nwY^qugvmPeGo2PK) z&EEx8QNG?q6`njdN^3r*HrpZ7u2R8bCA!s0wYP+@jf-SAH$6W=0rlFXR8I46KCO$+ zQ`KgXiKW^qZ{}QUg?^s*E$*Zx`w6Kf`#+b_IrdbXr~IjOb9HV}U*&3Tb!}siqrElp z+Cbh@spbX|O2n=y|7Lz(jB0mI#_7hNhMZsE>&%MZ`O2qSv;L4{)#{-^C+|ZN)#lx& z-Mpw;FjFGf7k}#c!gn^8hf0r{9V#lCelK;6i}LoSHG944`7U8PbXHNh)ep@xPI}B` zxyGh)x^!*&NAO_g^oi-Rl69_WF!t&Xo$5TPVOYGNJuFAm|GH9<>Bw8ftC@S;Be^ri z|0Zr^ScsAGUNfCRQ>Md-Cd|=Mwbt=r{d%0tyWzAj7vbOeZq*qBk^o>|a@my0>_h1)SENmF5aKEq?xbYm`(m*KZB8xGIdoB7Vv= zEc{cFYU3bx?2EtkzCu}hjs5x%Pq&*@b0zV4x3i2y>fq)HMuxkacGpj&?f6&NkO2pvlsp-rmYL$O+XKpU(TPdtUe$RMA`bS)7~< z_`0uss0op`BF^w|T;3I4DN?1}wosHPIapm>-j1}AOIbdvKje2`#p23Ud}n-qXzWp2 z>JUz>hU>p{;p%x{yc1O@o59any3ibX@-EcqWIyY`G^kGVRk|O3fq>*xt zq#yAe(@t-3Ab2#I|DxlK#z@TzVM%wlWZNgzHt+IztQXI>C$yVw?`Jgogp;lQ$oB2H zID6{NYjR61B+O|DClZJqhQZEwtVvv~#?g{q5v91&8Jk7;Pcbo|NWp_G=d;vIRtD8D zN_5=fn>9*c4?aCVNq!qp2$$-570h9Ffcj^$LM4JG`vK`BIxzT-_gTU#eb|n+>czaS zt8%vy#6@yh^r!ac1l;Cr{m_bX|1w>YxdvRZm_PL#o?EiH@^L!THKv;sY1tyv5rQ^>KM=eLZq>lw^57?(V@tu;J35 zKWt^u&t2h6H#e8xh{U_2GGb(6sSJO7Q~A}MCBYO#nh8gzzCTX0XqncrgtNW;{n%Rd zgKUBT_f4}a_LH9I7vnzX9F|sX{*!#eAER-&^S?9N3FY>sW7bCyhw2^>vFJ44;~h*5 z=Sg7t{HkB`VD;#?T7iSdp)?VT0gABsK>SYVSn<2!0!4U1dg~8^TWrq7w$$HsjyDuc zJF`!3sWLp4jyLIb*3NPAw#OyXR7^}WlS!a<250xLi?pN;TsH9s`yMgA`JL=Z^TvHE zsyc!BrFXD)93my%<9dZNIff$U!IZOMTYSKBmX>0A+s}Kh6^mn!7!*F;i?y2-J3TNn zHU7hOOh%(rs9NW?_RQ_6{#%#VRgcOLT-9@tGO;uVei-CMwbrqdj?3^dD#zw=tJ!83 z0ux?4>65p?KkH37uQan6ba2QCHp}Ep$2*4@BUgD|TX9*u$TJ^&&!+E{KGmp&#dGX^ zB7RNP?BLRG?sI*B?tFP1XX2jD@$#muSD#KYQ)g2(r$5o^S6Mbj{#)k+ly~oW>_wV9 zPl){GGb6NOsp)N#zv6r9e1p{qjOE|ArG;^bT+7{rg&+!fTj{?|U!aj@a*3GvkRdXH z5rtO6blnA_bt+99wm2VuJr+WUue`^vXBe`ZVlils{lJMG$!C*jEVcV*YJX#s`-n;u zrM}X*ww+V)QpmvQVGkJFo79!YKUBtox0eM(=6y#1n*^Fs~aS+nFt zrsG47a<%)3PgF_m=l$C@rkrzG%Dp$~IvpLqYf)1tT2G@w7h1Z`M4ID7Tpv>iyHkZK zWyYCc5=g8xVJJ0mOdj9$7hrXtt&z{yGdYZr4^X|b) zZ_Mi_S}$ptE{8>Pd~+^(KQQ6cPzwd2Z;zYErYMsB5?GWD-uQUekH{cT{iTb|EdD9A zX!q=IYsAHlpx#U?dAOeb>TkWSZ9eHH{tC{|5`qkhA2`}>&?C!DuSi+H2KD?PTBuam zKqHETh#TVf9Cr2JG@jaJSyhqo`)p(r5y%rmC5Ixp3+!@yx=<&pwu#a|<@DdSTK706 zvVH1jfrY7|-N*gYp5Z&=u=#7j!O2OyDu3r^zN?mJyOf<=(jm76xIj;UdANW-D zDgpv$=!)8uV+?-o+TUlLM#lFhnpue|%dg`Vzvop55p=%U`<7Smnq{-J+jWhtDFrGq)XgCZX%mJMyHC$d=3wU^d#7j>#awoY_77M+>V3|s zS?_(~cN_e0n{8eazcFr0L&zi>#y+P5-_OE6;8t<(Ezg*&*kS4Hi#DS7>fdgZBy@P^ zL8};_tCT{(Q*^&+fQ-M6+*mR;&`5$ggftD9KPE7L*)$2ncJ52c8I##$+irs{?0iTJ zMRH$aut!{(&1+WJS29tO>gR4;88!IFKy{-}g0VpaNMgzBP(DH_2!XaoROD0~?L~62 zy6pWRUM*YiYGsyBRUbdRIl{EXT9Y^j0gHn%J~W6Up278vi2*z@gnxuKC{DzAonAAY z|A5;%93uYtwexj8=k2?4=Qk6IUw+nEy?=-r%6kmGC4E>QTsYbI9#hZ8SjJ~|&LRi`>FgYiT_T>Xavx_vKmbkg1z zOo(6JNV+MI@`g5>4=2$jXGPhMVhvLZwGBvwM*Cy$=-hK8);2`6uWxHT8_JNvf@tVg z^VazKtn)S`vxhHIlpyA6r7~wJrr%bs_LL`id_PM>6s~GPy~nUox(lL|9%V{5Ss3}e znSL0H*u=9Vg39YkNW}Q8L*?ymJq)>k@l?bu+TOGx7|$rsMO`?(1uBJiIIQ><0}z_7O&NetPfr zcFze4Z?x1izTk_uEpc0@w7izno2(}m;4J0dF5tAqRwHbfynD^6<1f8%|l)ku#W;I zYLy0gJDV9VNm$R(-f}1@9}~@Hu>eR4V8^cZ-&C4Cp|6pnxTRb)u`{JLEjjB!ft$N%H43NGW-I@#i2<{N;obeI^rm_?#$8g(cJslmC-fo}T zzKfnNE_HZ!C*Y3e$9^A7)2eXzenCU!aenZqRIi!E|3$Os_H(w^-@c!Drq~Ut6q%x% z4(xM>#QuqVnuc2HNuY(|vN=t0d76Qv*7RlKGu_28CGG1lI?4Q*T6^WD#fNR?XcSh? zETu%x-5o?V^;-ZEAdNVK@eRdkz);4FX3=syDt6Z>b3ACaA`Hfrw#U(zYj?$xfnDuf zU;8*-wxj3vB|dW0(~pLC)i4jgJ5x!gkbTdOdGS=wh6pO|>@qM}G4e&Jl?GZ}!U%gt z{*07oCUgLY_Rqoj0s0=SrIQ)YU8v5B`0(CQ$5`|I_p~xG98zIkbKco6%I3SR9w~Z+ z%~F0mT;Z($WXP=8dXyn3icwjez$DJlOTrATTp7Gi3;Ra#oFYjVZG|o3%RfxU+I}84 z{!AZ3e>-2kJ^%hKS!hI*hvViMXV^pPDm-e5qD((5WAo@SRu-7#3csJ725Jo<6nS%1 zZ;6|)r!uvj5jidNojqIOGWWSK8E!ut`__}LiZBKb^DIZQgRR~bHfb~;^{GJTVz+&Q z>OU5ZiRVA&(XgvO8R*x`7ormJBdIu53qgoN zX1EzS-!h^v=fzffqYt{&7B2=ZENBhI@r|iRExYuHRG|mrB(o~)bRm{c>h>XoJ-Iy` z|A{1L7(}GD?EWLz$_oA!_#&Klvl`Z~j7-+ypZ9P&t?(Ro#zu~<@p)bmvBJ>`OHA;!;zy{n~mXU{84)|otuK5LHz2-XfFFQKpwnUjt#MqJQIB{G#~i@Xav?xFE{ z#JP5(mll4uc-UF98@(rv`hp*2xzChlu3{32pu25%i|-7l#mOtr12|CO?(MYphfvA& zW^WvxA&aH#N0W6+0V+-S*RWFmTY^zS_Uo|DMN(u`jCREo`Bn*R3Q{27PW%t zWP`0&#MYnPGfyV#m+KJRjyMxhR&x8B3Cz+lyRi!7xIc#J=wZ)n)L&;X5!64E{3vXZ zEssBx0NJ`B-0FKSGfjb8q^6reM%v)yiUElz&d3%JFgi*U9#dQPcU{>+4^hZsgY?9Goo__T8rd1TZs^cxHGaYwx3D~F9LHOKK!Kj7J*kgN3F+0{c z&QkGP2(uLzTVeQF zGy`WxJ+DlB$ea{7M0j>ynnrjy9&M5@43mS$in!6t{&(5-(<9pr{r_`p4uKXaON zibdTXnBS{eX$wy{*!qG!VNdfiTT+%(v}ItFCW5xmU7HzN^YnxDzz)q;WN}Kv18U+m z$u3t#xK2UDQyqH6Ka-RE-qPo5#CnPNo*Y4jCTETK&<0-Xh`zg*=oOr{afFuot))${A2bd@JWELl z%s-w~1bGmoC+%pyGl?Rq@j%R8UlyliWU1oQb%EewAAH>tE@cxhv`^J|Z{GGptzl}@&`SV8*9qI2(4 z-4OB3CxW`M!C1;zu=F~y*BR+UjK24caxFP8(v@C@XL1n08|!JmrA%a(%5Jc|@_+}Q z>8;$2*Sfztz`8@wr33H!(;*YLRgf92Mo|C6OQV$K@cFZ?8EQ^_y%h_YTD$MxN&=gjV^f&M&+d$o!ZzE*c;~@cEyv8-FmQF~@C7x& zaHh&wfZ=5l*u}Wx z6wLF;ijc{XAUBdjyhSQ~=uM^Lff1ph4;BY95*hg$iQuV3K>jNdX-bC#UN$;y*Z5Q{ zv>4ru`ZG27LD_q5NZ)e2(I5!kl_V^_D}03U z#@M4Xhs`~jlczT(IYiKw`39!j3F|{8@aBM<)vmLh*CLDe&*!b4`M---X1z z*JY^U4THd|Vm_tzLiXT~SyWIB$Qet+385RyJXZ8_iDZ6$o>b60q^XFt&V%+(tMmH3dw83HGPmUG5I|YFtP=tB!&!}=hG>}gGH-hNr)X!Q3A1qPD!S*M z1aTNfMIIHtOFlN~0(YHbC{j@4Sj59D^=UvXKg z;j3%nmjIlyoIP%^@&bc}Tm3xQ@in5K;4bxAQ+fAjjIVTZz|$^lB;clRWPd$efx#hy zPMUdx575RvpWAn3nB4{W!Vp&2)uJ?_{Tq+ztKwdOfQxcDo!6EI7c8PYDLXv|FOPbr zDFwl;9GfXBv&KbSZg?d7tmK&9J*O-)Q^*$K-F2YKGSzTPGyQ-fY9{ogzjGp_(&&`x z+3pmxDo8OU&r0<^}GM~4z?I$TLy=HKNbw-op4W_xbh2rhKXS0gnD2|kzD>G=7yf_l= ze@IL~OdIBdIFU|$wi(q|k7FD3Y1lwWDs^lVpI&+(tO4uLInGBYSMi*FG!9f;k9oxoXT#E%58XTN83VMd7@R$MM=ymwi-%{vnWYdNc}Equ=ZFqU~b z#ZU>$Bg75=Xf;2cZ}nnTlFaQ-rYtz!{XXvLyttK23g4nvVWWk8&EUXoOo1k}4nTM@ z5p;@IXJ2MYWWLNtWwH>cnw}xBu~l9+eK6qjB;=j-Ce4nfg-GouK=_6;djfQRkGS86 z6BVgHg5F+3Liz04!okV3-h}^9 zU(_{%$nXcRWf<2nIgMw-5!ckk)b(Y#(Z#dGQZ4G8o_k?uarzIF=P4p8Vr}MwJ~n&N z_}2V!QZ$yhfR9LYD!f;YKo~3PKwt5N>SUI|{?z7yh7#scN*(#bXJZ*p5=^jBk-^?Q zkw*h(&yIw#m(@lyQeG{q#T~w#kp9@6TH=Qv_SM62t6eJ5@~1Y_=hD}uY*a#iZ#8CW zOqMfzxjbDZC!~W<#9NksE=@MMi+X)~Lz`fSy_9|1*~jm)^{jMf=Bbirh;;ZK}!djo?fSr%OL$w(z-Ji`A2ge976qA9+)9xMsW_FQt8P+J$eu zoX{v}e)K1#_4Z2b=D>g`ZaJ%GV9G&`fF?snys6IN3>x)jp6FdNdx`*SkOyrk8CG%P zm+MAE1gpfS{Z4E|c9Z|$)`3^vB4&W3?*q3@Px=}@!jqA;aC#8CU$i)ldqEVbsvo6Fu2FXN^VPSt|@OH{*Jtkp@V6U z8bm3TyiD;GME&oHgHdHGpI3?*Mu zU28~H!o&DF#`cCywUKfe=#kZ^4i>d=@Ts1Z>h=?S z_8k}l}43RBUjkvB6P_`o_A$O@b6rjFr_nCA_R zJI<7yyxKy5Z!As0ogjeHqlrs!Aq>*AaS8MWX4>1HJZ(M)-HzpKS}n zcuKpoam*BOCRxLG1Rd0H;1tf_#{yv-bUIh{ew1rC8?mmnyS>Gj?^2*Jm};Bs`)A=o zln-H??g}h{41`sYdgNABra0|dU*mYu53|%id09?2#>TU2yw(DTo0Q-nmcs)3u;w8dV9F%)>J*Q(D1io=;y4f3{GP~^t$aV|+-o60I zd#RJ%?PYqUdY|KicxLD~&F3j>F(50ReL8zCJqo1|h>BAn>r;ckJEO`+IPu6S^bKV;d!y!cHsm{cRXAANvqBe?T4PL*xAeV%clw4d(H&2a9{6c*d_% z6lGJN^Kh;AMS7<|mv4SQIcs)?cs6Tjw+5*kroWxZcX9rbSHC%>83FM+DP3LfJ{vob zpUSy$g8eMGU}!s2PxI3t%T#UpnV)IJgZVzd7gM>Z9O~Z~Jv5j#_jB9JFapxxR#wWy zT&n)u-B6TDuQ5c4YFD+KMK7=DacMOd@weVu$K6E~A^$V6A&C|AzSnhOg>!}NZ1bJ} z;sTU@`Zbuf#kBbOo(fXVY5ca`1x#;EsvJN}xQaT01uR+zF@`6Y)LlRX{kCqRBag~p z=nqcj&)(fe?Y6-z{qk);WjkEgJ!8f9Xn&-|HrIaIg*YwJps~OlUa-)DO<+D9k7fho z#4&bxF^HkQ!eoh7Vui6V4qSlW7~0YyxQ+Mf1YEQybY0=I6gtFEL1vUtrq3?U8$Zj{ zpLs?=9y^8HlQ;4yLSeT0DcM~>=0c{xgc6Of{8uz#CRR*9G({M|;}^{u37XUie516M zAq!8FkU8L8e}!7Z2!^uM-R9TPMy|>B|H+yX{;bMXkEaFNC@AApqE5jEd!(H$>29tL%;B{bysl>_78>t99Mf1j%Cc7h=ep zRGio@hXB|N8NGvNm$;CCuTa=;&kG!x#ECMYR#E@8(u&QWv%Q3}aleSe6FQ@9T*v7H zWeMs|5!|KT9@0hZPvY(XOCdy^TJGskax{5j)!y)!#?T4@AKZ>L^eIB^(dk3$F3?2X zUYz3+aotC6+e`=`Mq2AJ8?34M3NPl=l03*QTmH4mmvylxUIF^RP4c4#a&%$ixtHzX zQ8&4w-NO~=aoL>1U*JiW82Xe@Llu}}&;tK6b2?5QAdGxFTi{hVA&uwP54eKat#&=9VTX6i>(`4m|o#PJ`(gYh|0Q^tZL;1lLAluRY%Q+WO^DL~A! zfW33U`sNrf+voP?`ZCgD*W>zv-g7f5ppN|&MZX0tiZwp|gmn1wIUCXo8nU9o8(19K zGNrr$iwlMUL{MFeOJ~BlJld;lOsNKD=SiZN0S6c))s}dsz6bAGz7|9b2QAUQH!D)b z1_|aBV>%Q<8WbtC`MLV=QMH~BOZo%-Tvh`fv8MflgYQ+$)HF9KQq*Y1Ju!iah57_f zG+lhW`i1Mx&!>fpO5?$jC6s-K)K@i8smi?wfz#Bsk6>U4J-|9t#xiVTjvtae@U zq`TnauyESE^Gw~ADodE*vn^_VcYzlK+F2O6EIbp;?@QbC6zzc{);ld(EAO3AFT!$k zls$KQv$wl~jWt2f)A?O#BZ#1#w~p?o?U5E#_&rCjzbUl!v689gI|W;#LTr_S#)~uB zocEwyjQ{uKRH_XeG6`PIL@j3H$PgLmO^_YmvYPL?+F%v9 zBFy)G@)G{UMRA6!HAh0)BaJ?_;Ae(YkX8^f^4^rWH7Uxy3_!Dy5=*bazYWrO01jy# z@E5N7L-Al9mlyO(hLbDbjkuozWM{7@`>~apTNWclGs!_qQ8844hxmPX2|W?v>g*LB8tR>O>INOD z$g}$0C@O#VQRbTocdznvL?KmEufY2NTTJc8BCbRTaXso|k6U&u(i;7IifA+!1)4es zS|)PwCJvyJGU1IOzP?vWw+i!Yd67;>y#z#16>*3&K|&Y_4_3B}yZFF|p)i8cmwbnC z$!{;)`WFl`+VCf2&(UXPd&E)sy^b&XA zN}d8rfJST?y$R+&lv|eo$m_|kX|hm6Va=79t@TsN%;ymvM*ETcfr!1EPg!eiQtni)RMk!aUPNf5T`UxiK7on(9{E-cOb2I6 zjix{8!1)F&I3Lc6`K=*LtikaUD!&RC*3awIHo!!YA%dbF5{QRP6hl;h0w*Po`Fxxa zh0p;X74F4SYj7y}riKVg7aX31^a=rJ@{9m56|PLipna1-4(cQIKMo3kCYBrPBc}=r zyVDgJhJkQK{z8LNg#i{tBErjrqJY16IxEB%m+5&mnTi}bk#Nyji|9}5sBIsLu!}slFXh?xfCf~_ytXUl#k!_p@}aefLAP8T@A?D#1L=c=&k5~Czb&dPZNlQGNF(Y zAXoi!#ezz-h6%Z$JhZb}jaN~}N%)^VvVusqPQ0iM7Nd5}&rfCSiq@*0l_a}vVX7U^_2Qh?#F5t@4cyLB%v!n<;0 zV&9ASdUZxUQoC8+-DMBw7ZCV9JdC%{>c36~u0HVDavI!1H8YB!Uc;e=}2mh570qwcJ~oj)L*m-HXDlE z>`xQjf8VSaLYxre<_R`fiq=1tB8Obcf!Q9zXC|-|`NMme3Hbm{{}opf7$HC2MDTxm zgmzi6s~_1)nIemY*nUV9D%-{jPOTcpB1Rnbgh>q1?t-^3VD((Eh5;#iHH`gwYk)yi z{{1>2f<~crma+XRP{KpX^{?z~kaC@ISYu%@v`22{^bXiem5JQn%>*MiqijO{@0&5q z_5i%h*x^i{#pUJuxAqIJ9V2sdl%=Jm4Ih<2N(Y}lo z>_B*u6V2Y9y3T+9ZDS*!Ps@b>6+!e6zq!0cXC&ky;$dA|mlWC`ou_4YX3-qT0~$4% zmjc`xQF<|GxC2&FKiuA@HV`^J=`MyEAVyM#pDXZ3ONYFqsZo2z5KD|*sAo*qQpNo; zE?^K-Tek^77v~fmQ?c#>lJC8W@_-mJX6j-ii%LSX*o(|2VYAV=y+Gbaw*(k;FpLrS z@e?Hv?@Pk>(wG=abdJC399&&T79RvzI^Juu#+g=3xQhY)E*ijAbnjsH5_{M03eBBb zWLtfQg{LtGT=2nax(h^Xh&Tp}8w)D#m~-32k*-a~pg zXcU|x;K^r5k_QA9My9jO-8*gw1Op-p;PMNo8I)ofS4vwg8-X+Tg z*9RA0jPFzy#&@Yf?Prg!FrcFT8q+d4NcZBHp8lQ@MuA1T_SV!}3hgT+63DIG)6g>a zDSC6tAK)0JMA_Dto*jY7uY(L%yCUQA0VK1G2>_5^<%;>XWj!2wupVmjt_wP=TA+a8 zsAq63rAJ&H%FI5t$AV50VNj=mN)xkPToAUM#qQ$(-hNycl&l>da2Bk)3;L4=CnlX$ z&b&wailmkPUI283BGJq@*3v*RR-rn@e5f%*$aLAv)~XNK^PiVu5fErb z?^{U$2@U-65I~PpCoZk7BZyLZ8k0!*m<}Yi2-r25P_Q35B%US0gM&fJh#o*cOH>sf z-|Loz=EKy!de^!KfQGqLApepkX0y094ANKpF!qlQ^8&Qq^4!6dm6&wRM{I4-fByFo*pyh`^4D-ORhr)CF6z z^Sl@qS&(WbTnZZs2R;FYmJQVY;$~w>*E@RHu|%H6_CoTp(!B5eHq9q$hNma=%KAKg zjZ*QuhkGKEc<3@rD1!LCnJg&rJrP!+aRRwH1m)j6mX>yLz&!d=QA)?#y0!yRfV_Jp z`Q-Rj8v$Gm(;>jtgfS#NuYVj~9$M_4kfxbT2Z%)^mdd6}4o-xm-`=ha`A-F) zg_h$5#@}WRGqm~w?p_fHe)Y(^#8G1;X+o`qOrsvo#N`jvPlpHObw%=(iGo1tQ@YhV zBaocE!xf4VdS_XU9gnb_8AT=xqzosWgB@*vV`>IA4iD}*)q)Bg@j{2zMNwZS^vG12 zBGd<#(E<`4*Tm+3l=;jNK)2>%}B)U|??0l(>!Q4245jcQ=)5!@rO5MZ8)J{Q8 zn*tU7em2%1Qf09hrL^-h)8uZHQazSzTLu~hEifRXVdCxzaK9c&17JNbL}nv0OR+Wk zy(p84FSN*=0B4K`MP`y>YdVlDm7)t9M?eleA6R1GLyBpTOTpdoQ=H%Z=Rl`s(szMY zx=4b7V$vT?s{@8~{W}d&VLM!nJvy>4)~GS%0%O&8d7cSGK**euvE)H=8kL#<8z!); z>Fe~SIZS?W0Q7+u)Wf^s>-6b2;n6L>HA&7-%7D*62BY6?ro(8(A(g;?ug+i z#K40?((_Q&9%UA_VdC1N0`82PD6_pgY79Q`v*njuIVj+2drk*1lFU%9N0zh28z-PY zZ+_W6)8YH8KjaoxK6F5TcE!lklo0LRZ3226Ah}h~@c%yp+q7^#J6MqMIEADCTJJ?k z;diuM!x1$w00}o77k$#nwHd>^N^~bwT8*SSBUX(m1Yf)jVo|MQ@H!@oARRFYvKk%Ko%mU6AoFL| z??JI9w7nzrl*6SrsCS4~wmGFNWBFj!U%5q3YU!%a`#QO02{u(91^88=;o-OIq*m;&EHIW4rpYX^c-)%?C)0t{FQ* z;n&4l8{GCUf?O3!9Bvk0L|IL-jjaV2Yf5C+A1drjGwlpfO~&SRgc^M>)W<6smSTh(#3ZeDywHFwr2)Ok9i-q0jFpQv2gl{k2dxgfDjz-3N8_xy+|L_ zwwG0{n^8*O_?Dpb-33O6fIELe&!S5Cg-4|NHx;Mqd-*Sth;}qfo4*MeBmD9&9wXm8 zUDE}fhxx~fwV?anbCs4e`Q)k3>>lPnWMhS;3%e6?n@4p`^E9q{h&O=>3K7%Z{i->V z7m>-;BiT|iY2l24%x{zs@ZLCjkauGV!67HX_3pRFr51IQLBO&1FZWsvf91a7MR+?P zjBjr=S+qH#DQ0G7Agle*;9_g}Ka&I(*zmBZtBc=w_VYf{Me2x-uU;EJg^#vtA^x{> zjM=+3XP%SrU<_HUuiAm%tVXjv57#tn`^E;BK0mKs+FA=?SIei(mkZ$cHryc_pq`|z zx_ur&^oYyQ?zZJ@Yc-gP+d}3I6=fvk#o!60O3jN7jO`iKCT-8n$HrYaUPTOV##qAE zRw%MQoqWYl_2Be4e|4qvO&g;sJw5%8+X1{|Tgx(Too_(U?-iC&4GZx2FFu-JpuXyi z?Crb{HSmJ}#>T5!!U7pss}0ge|EMfFX#88sFm=0hUUnB8M8X58NcFM%Pt_aTdAu%c z>o~Jk;Iv|7+lt`9ccvSaymRG0k_k89wU1iFavGZ&%+@)gZ;svu1zyY0>(-Ln;-*ph1`+A*l#mojm2Qx3 zq$Q=K;mn15|IYcJc+PV^pY!s0v$uPD&s=NPteUmHYX?Wihee%90vPtQ9&gpN?yonhD#`HAQ3tio1J^AxR(XEAB=WbnV`I2&AsBc4WTmQETNsv*ZXbfr?tcK{yCO2`=S&d0ie*A0+?zw7l>_f z@XbyHaM+f20J;A8wO)#FE%h`B$JfseSDglR*mWwYTxFb7LN>RrPB(GAn!F_6NnJl~ z+!2j>@KQtR*9_i^<87Jm=3VMvWS(f*jox2;9gNGmW7y!g{Od=y$Fyq&xB43d$-Fgs=^PW-FSnK8FeLcHef|h@=0QHp3_C| zmxE46%cBkd0EOsG^?9HDd7NUsh)WyopOtHJ%WT>deRuB0IpfKMi_{( zD5zWN_xkEx29-q2qdJ#hp>6!JnOGkkkbwrN|2%v(I%8=ju}=W~x!MW!HR;}gFmCtL z;LyK77l(tFb=hu4p`)twHt6t>4oQdV0}{&^x`qux9w%WGv-9)RI%USvyG`3X-W$)B zyEN6(z84sma9NMqA~|Y#_s?`qG##N+iNgsOrwYEkZ4cDD9xJgS6*NPCXeY$dady5t zobbqwU%O=HUE{f{>1h}|fzO8V1CugW8-gqKoL1q3xh&fv=Eo0=JuS1BqDA8zp`eHD zYe&bON=wzRs}a76)9gB&Z?rjuGhY{m={CH6IM?jUzE!E_bcytZp#*E6*9e}Iwv?-Rc-72_2lRBa(L6F zaF)`qFpWbNLw<6Nd=*5UOH!8N%2bHGzwpjBt1hr?ejFmG&miu_qb(_m%j+>Mpm(hA$w$lHtqwBTU z>Hd5d>=C9`fuD8BbQVAxlGb?Kbt2|^1A3#)wwD#GCi2T{uNS5S4$IVuK9xzHyw5oI`bxkMD#_Uks$@fAs zS?V;GQbF@hO2gD6hf=~?+0ywC0^yqHvXS=dLT32;@$7eeO~gYqg<_McC4Lq?y(Jy1 zPXFLl7w$!dLO58*_~QNW_vm={{B!?8vu~9-LI`btIwu9FWuW_H zq-n3PvH9(4ySQ0(zPlM`#TKX`e#OJk)oQ?naleVDlTqv(ZK7f}nU%&4_tIX>008ke=`O1Z`Jyp1;rw2XTxB~rZdy?94S z&g1Y+3-o`OE&pa^(JCOVu{C18)@yVME~b|b<6YTbJQ=cCpR25HxGGd{(cZ6o>$@Y_ zO<4G9>m*aH4Ur~#oBntC#UD!Yu3^d`g8SzzW|EyurU|7&6-aDp%<^5dgk<>`R zU5Piz12%!}m9w*?KZ^73{S{LN#buNzgDmUO755|SPrZ7y(r;G#1Q+{L=d5NU=GoSz%9V zhQt<^JLsM$i|YGgV&D zn%Ec#B<%t3u2<$WUnPh=4MU-Fn5UC1*4@_%aF~Kb=ef|BVR-4T@OS-~=mb7&W~3w5{h=D7oqw^>hV{I8BT_kXpUH+S*@E&_OogfX=AK zHlgWSb#nD+n0Z~lj={QW)35QPc82WqpMuFu`io1iArUt&c!^HQJLGxlDCZoDx%mec z13nF&r&9ecZGk-E(`}8cfeulbm=~%XKlwRP@fL;$8DnYcCTw~B2Q_}XK9^P zneKRGip4h=iYR(hbs!zAUT4dku!`nuLBpmC!KoStKRv$qmf;r4`c;-JljS)TZ^|aa?FSe%L?mR8xkoz7Wj#?Z^ z-3&nN`ry#Vx-n{isCPS;Cl7ERsykw0{8b6;OY?QXRW}a4iZz8kf`e5Y61ABXV$6p$sX%+)P z`6NkTIz-gdzDs&xwvI^l@~oNeRX>*FGxOD6O z7*xgXM^fW_T_5uO*49S$yd9Mx>*43x`WK+o@U(Bcoo;(=cOulh87<^)9-*(e_$m82 zRxr-eV9?RY?=yjR+ZnntK}^oIw&;4dSC;d}4PJ-n*B3`57nhfxHWa_=<8Xkx3`lWG zjbzI23q)AXHhN(IBy_9}16tEVwye4!lJmtwxh7tVsAfa>bB*7>#rQ%9qFr|$8b>gM z<8jjm`q8MUP)t^kIY(lGdlp3MeNEC(61`7I(`K84U(VKKrmGuH52fPW_+REO{m5@c}vr!$v2SdH*d` ziTQp#zEocwYjmzIf!9u6vBAGu6|xq{7)C3WAmm7#B^;Q###-^5%n4F4A^5iw{)^a_ z6MplQo)i8GRr28t9HrAA4lrKr`HnN<8NJ+-%I80<3Sv7qYi%ca@9HUO-e5ga zQ8~DbQB@{5^JYrDZ=Gf}n|{mM_xYJB{El29gyfVjkMo;OdKqmi!owvwNbkBO*o~CW z9=FE!5ySaZ%_YvoUI%8eLA_i#xWqwBDV>KHbeQcI6@q(~pLP9!)p~4aH86Fn%{9Oo zoz-6N4k;QlJ}uFwJRJ%|teJ};L_I(MybN~ZjSOaa8~hqbQ|REwwf8j4+~9tQ3JbPO zr(FGtxinD&$CGtrrKAS~Y>R=7s>UHOCU8paEZd%)nHdCkXIuCb+I{1 zg+voq+pPIbyuy&oa$2!)KX+x3-ZxX9YdrYC{im&3{aYKN#!hf)W0b1G*mN(;>y)uT zm#4CUNs*A6+SlE~$3NL%3*FfmU47y3mv6bJ8&<&RKq6n{P<97b_WS)61odJo8C%(B zOs{<@!|Rv_nFtC`;J`GB@t7aC8e(JAcSHz>SH4K^w8_F=IdxV&#S)GhrvKzlD^{j$ zhqf_tibgL>@qM)~A_6AX>uY-{1n2SmW9xOgFJnx?7^Zv2VXlG4H$?uAX?GXa96~o)G1r{C~2RqU|}(^R#wr_H>TMZmNO7%gh`UT)_d=11Lt#z+ET(*-?m!;T&m1HIx&HihEK;->WM4G%739% zI4J7IB5v1L8@c(D&&hozi47Y_AAG|4k)sfKZ`(5qT*$Jo&)~SR@a=6qt_}<3gd10H zz^vJ*@6`pX^*Ag2fhNO~I21}@LI2?MJFLbb7bH>>xXh%|crfnPmyuO3`EDmZFV^Ma z&2;UX0=+*BB3|yQVKj>|?R^G1Lz&N?+s?V;c74cgook$a;P8!qYr<8FX;sY2y;+I) zYI|C6q|~Z4(Z|m^uE8zkyaQQ;YI*qen&J#f@ z1j%^Z7lu}WCumsxvCMp}E{ah^-0hL#Z)WQqx>E3Yo@j`G}90kv*Yi9Yvx&7 z*>cmg#Xhp{ID)asm`vxJ;#iDDZZWBfF-5<))BI4;Q21$)mTZ-6qRwgo?ibG1mvOjfkJ?6Gm%BRlRuE-qW#RZ~=!pm4 zaWldMcF)VKu?GJ`63#~LsKuKxk|g> z6dTc-1a&+~P;+HeQ3@eB2YsS)?&&fN#vwLo0hczVv(o(ee$jHlf_AN2VvV1`#ToD; z^;dg~$<5+PS4p>Iik=3iN8&tEO=D9v3^9aqNi9EoxfIJy2Y0KNF%YiL+q-@4GCyj4 zs`!%E;@%#L@A>yrUA~ueW3vr*y`preAw3=MSn;`~7Z11?5k1R`_3@mNJ!UN{FR>g2 z5^(jvr5oW=MC=O5(rJjdn`3mUzj%eJJ>!)WsoXD)hVwe%L@Xxi#@rUQ`-)Ym6X3Q) zE;fDu7Zqpx5Jj#M8*)bH79wxRwV~PYK$zd`4jU7$x|gZl->&NRe59peYS_&!85Uk1 zm|erU(c{?wmLRa)@M+vltt}H@85GDtrSK%uZ?p}uM-a3)cR8)w@EgZ=wub6?`#tGA zrAcs#;5*!bX)%>I!8i3oqDdYHE1_35eQGo~juUcT71GF&XN>kYY4*_rU9hNuxDP8V zF46e%xMLRs%gP+kgK@|o;l*=b(-Sxjq?9XN@z(jgRb(-$qnrA?NekZN)P+$OYUBof z{w&!3F42Fcc_sc5v00iX0uNa4=TWP)^ah=Rl)H0{G9Ur_doP=LW(J@Bd63fn88a=g zPZS^}^L07&U}Z*zD-s|JXv+=oTSEy*hVwZ=quHJd&|Q=2Q2G;b6@65w82 z@>f?E1cnV0chCRK#$9t6FN?z$-#IR^b>_^FaETy5wS8p;p-$2i1{LPu#$kY3c}L=^lLwl_(Yyu}SMWq-l@62mbePPDVR+!y z27kRq&#dvSFT8s-?p+rr{dU+Ul~i&jZ%&lhbmpnX1GzezIOZu zE@9f+niEaD5N;Tt0~IC_Po^>rYjJxN#qkp6=x7`pHJnh*T&YO1O7rhKz{4fvgS$-Z zv4%>uQtNo}{_GFYtJhi!gxI7?lTBrSaOrOCBL#GEKF_1~DBlim1uu7nFL3{E`xSJr z4Kus1b4bG-S1CY(X8$Fs7<|ZS{nbSZyG{sN6rK0&RyXJ9d|{EfhNy9wP)(lXwf^?v z$8zR7Vc?FrCysRo%d8=rNk|5({om$dd&<{);qJay;owfGtnKykeQAVzg=nq4-C<(0 zT9|S~#S2N0?;Lu)Kr{2@6>@fxK)Y)56Hbn{KV`yfeSVcB zx)2?X+%8+}VAZw6q+lCXvdDd*kW_SOLh2jr5#*>|N?bi`&qv0QgpNx9I3XxjJ?=KB zVul+GkjOyW)FeQ`P{zFD;^LGbt=Szj>NA~ecmm!-tg`X(96ABSRcy1Ho_qqmJZ6qrQVtIMpV5tiU_TzALepTm5UUU zF#Hw6_~e;p68SXhF1Vmd+_3RXW3~ME(UD+Uk^7aGT8N``TG;vVJgxI;tA;!eVU&$* zFQ5Gk>Avcsc~|9&G^j3-C5O`~gR756C%-YW^y;USP9o}z6Jw6M+4?kYMxJsq@uHZT z4T)wemRYnA#)C>7x0$Oo>Q+Z^4Z?(`OiYwA`O;!k3+X5G#NH=# z_!c`xUr<6d^FD<3^a$97kI~C^6L0~1-ldl5^&dLi#FMWYKpNMeto+NtYH)E9?i#OH z#Fl1jl?Uiae`iFN{>akHdQ&v(<-B3dO-Bph{CA}9#M~SrY5}u} zbuj+d96m(cK1b1B~Ts&9_-Re4P$yTg%n#zNz_7O7ieRgBZc zzQa7Qh2(JLi*}BN2vm3G0>NF8pi>`Sdw@|^x~`K_17$H)Wl_0)dK1%1NMH)m;lbZ9 zD6AlSo)AT08!=oG7lS*5dTf$ft5tqR>`e}9t{&oE;iA0FMSS`>`SuQI+a>#TC5{t} zD}q7-00j*`XIfNcIB5kZh?aP@>Z5k35eQ?R3ViY){>x(E5y$fZ22)Xh@F?ksgBrfx`%!`4 z4rc4z>H{5_H}Dw%middr0gZ5=qxJ}(g4U-RMaYUBj;&2WBd)*#LsG$WFonyOp-RLJ z_*Vizrie0%3ZI0LSnmhXPkrypr1wa5*2Y1$zc;#srJofJpimYKRocmWg<~Rlf}UjQ z8-3X+ku?DGG#b}z#d!nkPX<5o^!(brBCtfeL!4i9ZJ=KSt171c!m9i~@8L z_GrIAXdbhSt$s`JlL_o7G|4TDiD-c6@%*ti0m{!%cM8-{%O^wZ zd^6l21vDKH%XQyVkU$2jBKrF5=~D=x!{(+tLkiHP;BSaPqfF61UwR-`{XorC4my84 zVg~S=wR@mH`F}-S0kUk4H#`7(>-DxV)7IE$Pk{%*pndigzVZkUfGIq3<;s_ZCeWq< z+lr|+I>Ok7B>->1v-`O^h}YN;`xp#J+So(mREMx{yvF(;`%biMsD_gp0WNS=yB`<0(5P+dN)0XJi!_Au0iBM_-tX-00f5-*pr)@ud;Z38 zR1>E;)O+pZYRwF=l-pEU-ZyOF0&E#s!9f@R*j8Iq#bX&TS}On%tMK(ZG!Izx3qJ4@ z+n_0~TRdXbSg!Z1L8o*K!XE&@z&45ga)uy+_=8m7V^ar-bcioE!#P#_{%#jTh96iJ zrfGRm5rj#P73Q$VUef%#kpRy6&zAtSn3M&=bJ8HEf+m_w=R07R{#SGyK<`z1x`+a- zUY#p9f$QdUJ0ghH6@>d^dIrKPSSq}VMxbdstUE&M2mGi3nMs7baXy0HnqhXGjcUiy8@F z%f^0iTM5|XMfu50Ixw2;ds{bysTueawCGAQ@RNs+?c)7O)2Gmc?+Ad4lDgX@Q2qcm zkE9*K9zwcg6)6!z!Y$I&?C~>GwCAN`4p0J8R;6Nd0ce)HoWe~0^0*7D5AfgX%EUqn z^MjL^FKn-{O|61H8Yraz203|@Um*BXHLkqtO+NRY6ixT0PTG9BxU<{)UytVaHFI9e zEHXGg71?t}y7zHHEXk6}{{zx~@_1mN{`~m~*|?|(=LmT{EwUYMCzUuAK;>WH5o|J#%OU@Hg)EAadw;A;<=l~ z$ow0pA|wt;RXX3(@;qyOU6Fhj4~L_aV;a>DMZdMUO1p(ff)WYiso4p$xOfWuoGShD z;;9)<1P!tcLCEG-z#DsvpZJ_f)xpoM8B85f#Ke369Nn?FX#M?OA=}KtS6KkCG+lGv zV50#VhNhB%aid{8DC$kVpeePD&#l4dQvG?}@JtAL$Yxe6=CwHQzgU2BucLT+r2e9L zCH;c;e|uJSRPjK=>SnesD0JQYzTAr#e~=9f&jXmD_qN2vi2>=Er8JHdNPn$zX0 z@|nOocDqRapX4+iV9x_icXo{3_j+$FCWs+wFoeK8Qr%jmuap4n+x3>^YqH1d8Q8IKgc++SaPCyQN4)Aa>)`Ngm*{=VvSNO-fV4DmGTFtdp;}~cS z9<96c60}D`msUvN?nx!k{qOChO13s8$I?Ep|Io|?9p^@L30`T zQ(h`Y$P8N|lDX#WQ*|?cJRNz_qy2WQU2dgCyS34lpMA`@eER%3PNAXvTwJ)JxBw#W z0NMWGY+#{&30~>~#~|LctW*U&T_JI!Z1PRA*c%npTW%A8K7++lGZSjx6W?gTC%>{N zM4eFEp2Q5c8MZQH$wa(aA4vK?&uxU>HibM?-XOC-c&|~sGR5Zt@mn; zmO;-afev^b9z~-BnSAB+;2>stb+`^%Q%lp=N(1*2W{e2;=KWaIY2N)rm{fg{Mp!zT zfHn+$=mHAzv}ayLcg25=D*{maQOSxWxM#zO&b}6t4Yu{0vgk4pg#7-iV%E8 z%&R8{nd7WZ#y8r+2D++>x+(cV1C~S}DrmRL)U?6}VNEHm8@H(@XJ z^uAI^NK1S`522Ay+MW5H=y2PePjuJi_+!ykv$M#Kl-p6syxs0EpW~;3DCq7C5(nS2 z&3TEyt7%_lfp@j*#pD4G<0_+!C?r=%a)8$)QIR+7@hf^MpjT>9y>K(&KOYk(-U3eq zo0dc21#|Eq7>bF!kwf&71|1bw{ymb5pk$Y25qNQ*pF zr1s*A1cT=`TQSw#2$DURp!Xzc77~#9usZN%YoW0=z)#f0@hpNL)uI(dJ0hG#06BGl zoL?`H4xopfabL;!tp>6u(y2hpRR_zZRHiVa_mi+i8R?oRfPwe(XiQNH>Qm5fh)1ub z1>A`c!DECzRSZ_!pZD%7o&*gO@`kz6AjNGk+j@3j9~J@*LWPz=Rd)+nA{)dEh`nr+)CA}gfn8wH(rHx-`dBMBUkNqCK zl%om(ueN@m2YRAXng~744eobF0po-2QL<4!di~yYdn%Dh*7!h*+qlkk+xm)iV&*q{ zV=pA0Atrux@mU*xZDg*0^h-#a0;v}(;TqnE85HYe~d<1(o7%vXS`#Nw3_hTye*6*_i5S6O%&JF-jN1wdUQkz9VH8}bqWjehJ`w<((z z0^TSESQ`7EFvQpC^%(L;{9+9Bp8H?0rZ)eJ-QYLyM};Vj>kLs+7a}{yhgIf~R>q zb?J{SO+SOFb}mpu-#sFe4z8;lL**=9({ec;EV9gN!?vLQY9)pYB29tdK_4Q0%+}M8kh&T5a8KQ(=uCjw@~s~? z%7~E(ubw;Ah>`d~cX0(h@L}&fN91x~}u1yTjZZ)wr&_!F7ChcDb^8WKMzBi=rgggsHi0cgZD0skya`m=`l(sWD&niMjlk zpW)^E3><*{T3Nw*Vg1!=T;dIJxGdGy-u64k?9mTK&HEj%HyRp*6BA#IJ;W-Wqmi1e zhCNPflb!ABjcmkI13r9y4cX<#?gg)~AT-x1r`5>FY=(wyFy!A%lO+msn^#u3^q(8( zwi3koi9J+|58ao9o?wnEA6xi0wz4QaJ7eovDaX*ps-#t zcj6u-abF+$YlxS0^%?Ca+nM31>Y`K^g^ZNqk6DGkMM5z{=eSZ1s3Q-`Nqr7h+l!_C zg;65)HMU75;@`T5E9lv{^noYaT$je3la1Jj23&LBV8-MC-O~c zKmh~?sY(V58HM%!g%h~`F+hi!Rqy>53j*V7EEtyEYvnP>Vt(SHOZp|Z;No&py{WTt zbVBbV1Ki~+RAg7%FxR);3XsS_8JU%bSBeR_2yU{{AORAP{Q1$WjB_y7y6p{fYCo_p zDt!PUhtG@Tx=lBqL?aOpqqwJ3x)5%d*DS-hPzu2g`J_{h7&RaCY(dhK4o07(7*Ub0 zz+Ds5K0?t@(ic-lXJ^vh&=Oo0MW7GIr;@3M!p*(?%3%4)aIE_RE`kMw)Bk7~h(~I| z11NAvSZUw#QlQX62XK_Ew#RC2oGSM+VL*OgGd+EYrkJXPh)Bf^(E`?-)nV*T{*I+6 z8Y-luH1G3}`Vjm-4JRWIggU%3DGK0bs=-}bO}EQ_?n;Cl2qz+)Yu-tv_Qr{H$k0>+ zUD~aDd$T-BFnXa9M04qtJuAYLL^QN1alad_;%~?SP*|5(wsg1UOp=a%mWL5jrv{05&#PD$gE) zednSAk*i4mG_-NiG;&Yc zg(NU}TDUqec4$$MlSVf;p;=mOW2RP#eL?Ll5qDK3F?La>xIyLqsp{pt`a9vH6uK9J zi@$g}1=uI+_!rX5nV2 zguU{<-@0wnyYvvlE<;3fHnRu1LSOI`Tn2tUiBF3c zOLd4VG%1PMUt!Nlac1hn618vf?&+@D^@%oOdwOq6@qq3D5jmZoL;yz2iE3}PNP@Gbhi_J{r3jk9*yzmB_dwU)LLrT{o$yz}7@w&hODyDn6SsDC;q~C#;x8d7No` z@%(Z3<-G6tU3i^?4(rpRm=y0h3tlhN`DZ%DBG={)@V&_ni{Wp)N{Y+u`{BMI(2T2VqG$c`;ptzX;2m!=eOv%|ar?%*{#GzSlPrEP@9% zF!T1f(W?rM6hXOuYsXOjM|;?X^XsTxg|&#ELnhT}o-@c!qsqsftZNrff<)#twM%um zywfQS3gC`Qj4!%(JjXsf{QUX@8Fdr=t#j9?@%({!lQ`~JgbGZFiu8fvpF+=!*Uj73 zu+uGKC)r9w&Z+pHcy$L~wmVc_0af*f#i={ZKaX90RQ>dE+O{)V{zM-cC1zwFHOj6- zIN6YJF@@m$q)Gx`{xYuIB1O)aF90*qpM8~|Xmc6FoIrZ}c$0RRshMf%1xJj7*m>E^ z(5lFnoj;{NmfvXeF7FvvTpd^cSlat(nFeT!la5+6n6o$Ezv;1uCj;!l5-Ule)ju&Pu zUx=qN%*>fI3xzj4{Z-gkb;DB%lFnLqx!L8eW%ajK9l7rObT6x>-_2Ruls{HiRV>lk zKhodsa$d?OXR@P6F>0=9I5=X1TbxK7^drUu=NL)SoobcNUKeT=>i5~Y?Y)%iEJ;1@ z)=BrFGwHu-YiL%5r3$tBbhVhJ3A$?|6>q&f8>c~-_??dk8zs& zr{__9Sav!5kAA)$E*^}W;}AP;ea!oN#cjEw@_uJ+WRsw8)2k~(*zU>OUfnaLMPs+m zu^87z=GyM-Y*DKnW6=tg@!bXAF3m4`TFjT1^(>pUmc%dj>+|4y7XhlvgY=OPr)z|E zB)Cwb@p4Nc!E)dWD)vL;FcY1!m~pYR z7gUuweF#nRbneep_H6aqAEU%a>Z+R+I<6^MD)XATN3FNGws(F#4XIra=C-_EdmD4^ z%i~<9RigcPYHqS+tJ!qU&FNAmpJ=jSFi1LI?CtXIwb0~N_FV#_r;g*dF8d7AC`1{> ziM`Fo3O-vVuLg>}Q(Q|qh8+|%E*B1It}MAR6MOK$i}p-B)}FvV`f$R1`}K#Bh6aqr zjacNUeZveV&Vs5`dH#2y@%Tq*oPf~80dM+OXc(E9%xD(MZ-i#=SKr8Ap&^D8k)7*s zNVZNtVtG3y-q)CK`t5R%)LNIJca(V$*h!$jgvpW*tJt-XjA24qry^7S76EVPp;2N zc8A%!lP7DEu`3b34wMh1sPE?PUJXWy-QVn`x9J|?_Kg~s!{HohkQcb}_P(7QOET3^ zQ2T_Dub(lWn$?JkjP4Hg5Y4t`_^Qm*>#Yz8^S^e#us$j^0Xy7J2KdD9CRvW&)iR5|G^lJCXkx!qU>uZn3BIgiC@>hWC*Z2DcNmtB`-GPhaJ;B#?h@TKBc`KmYZ6wP2P z{y_Mpg!Nbv{b$2TrG7hjHQcwo{5&=~v@$N?*oN@)_3-n#Sbfi3=2`pZsRqsL0w5S0 zipfN0X%WDUJyebeVou>=P8rflqONZ~q~=k-f3d_s_D72vUL_ZN(INYb4a3fO9jn8I zsJ&Kwp!3PD&(rl;s*PxCFK_0chuk~x-wH#=ChPRIj|@GVl#;3$d>LgQt=;cKydIgj z<8Oubk2%j46(*xZ76#@#uvvQPY~T6V!m}ZvH8KXzKz(tpIIij8<<1u2sg$9~+>x0;_HEAF4KDYzuzPT_{Zh;rMh1a9SW2nvojys!5<#XLQj ztlh&oKxgw#{IlAZY~Hil=d>e!?|o<`aZgoOSy!jzrTJHcL(hl6gWNw&M`4DX<&}F_ zw`;t*=n5-rjxIXn&@dmh$?|ja4m~#ksi4lo)%RedBUXFOwNSV4GE{aFuu4_8Vd9~v z$kaF1lzPBsHQ>gYHc7DbVI7l&rTfo4*r*p#5z0=aFm8^5bBx$~j>HI4ucyy9xNqbZ zw&4&S5c6|CyS?@bB(l+$Rv&ydOPLR-+w{UrL_&E9j@oaB)R}M2BItzR^v*da|2bpl zw(~Rgzp>wU{0z7{N>{Ep=X)rO2^8-YV1JNg{5lujRoj%mbdY!*mq${=_04L}7DRef zFW$Y8PyhvJ>{6N8yY_8y=qDh{jh&{vVg8D}Mj8o<&1V^T<9Iuqn zIOo$>16MHU8l<5EhTBwOJ#->!1Lw2w4w!*lri$g0v?V%;1G++xk`XLSd@%<__UBM? zdcZ7-Z-K7u0F_Ua8SO-)zpD!E9thtP!tUb9vqG``@_!ZEXTZ%2WToB(BVyi5lHK#i zg@{KUTBI$ghsppcU>6=m7geF;s^qn`tDs8*a}^OIxcOGotuhS-1xn%J%4_hMhIQ(S~`L zSjcL_-jwa3G)*oz9_>cIee?^Qms69$(MAyLr}c)t8fy)7I>QACtkpm}G7K0Mmxk1B z3n5e~sGy80jR?-z(NCd}a+U>WWTm1Xm@$tMiOGd&aX}m-s4-WJqcYp)frQaP+ZG?ayOT$kwDVwLb-d;r^c zC_Mc)k+29hv3)6}BwLNSl5xNhPaevgoo|72E;#W)X`mAhkS18$+Mi$~@K8w2g{h~PT%G)8(xC@bm)vZDXbITXdqau}G5n{HEZ zn)v>nIh$H$ai_n4v{Sd`Bm8qvEK5Kk{~Q9&!i3c?I5tW}ztG@=#AEc|#A7s+c>Drn z#zoi{WLpn9dOH#CS&g2|lZ*QV1O%vZqhgP(@)M&;vb&;YbpCMG5&0qbM>2skJp^%t|;tRPV_}Wct{K@ zV-^Vu%=c=kuQdT31p_K(M7{o7u>l3gX#3^ZN57!@i$4uuCynCdd{ABTB?T_>&FcA4 zSd&SHD7enE(T@_t=-z*NJ&SmuGPPXozZQNce_z##ts1}=}i!{=`*q?H|*He-i(&?$-m zz_VIfb5k&Y$YKCW8Bw1*LtlOja%Mq(TQ}?XlTZR)@!6Zz2VlH}23jpl$$B*r)WuPW zGF1M_E7DMa0YdM>-#{MNPhtcN$Rf7%f~NWjj4#s#s_7ZX$rK~GROr$@guI*-ped-k z%k(cz(oBFRUj5rQ^&d5mP=nz>I*|oPN}+{jYyhlqZ-m{4rbj~gg*n(Duk-Hm<_v@a z+#_j7XaxAyH^3!S=JZ0a8f{Q!(lcAj_&a}$E^w)#T@z4+O(KH62n7~>IW{|1n~86dx8MeBxxm5m1NTrL4AAzlqaj1Nu>Ybk~-j$koe6S z=zxoOpkaI)6xMJ+tWL!c#R{SkP2ipLf5?2gp$G${O`{f;`CK5T-UBRJQd=bj$KVR! z&sVnagi-$0<<{_V1?dC&JatV5rXMuZSMsFumL7;eoEw$^nQLVVEW)_IZ(j^ zEDZ&5SyaJmK0Ro9A82~PwGc2pr~>}@TLr8I(VUBTa4)nz2^um${D-I)mZ9Kd=E`8T z#Y4`fpr!{7?Z4!+(80nRfQorUD>B`zcHs-i+71x(9zxq{i}lJmVH^7ukdIGBpm^&U z6oLct!oV_|#i_3XI|sh|5&YhSucrZG&l3oN$bc8%kAyM?$S+a8>kr@;`G9;?f}X$G zlLi8iR5nsfyg?_#gNVU(<*gN{Scm~z;SHmu196c7h!jA~10~U++A2TyMHeFU-xP&O^x*)o&Vj&9vG>775>@d;wDcuxJzup4e{JSeiC=UVoR?e;=cac$TL7H>w zJ=!y{CFt|+Lw^a9m!ok6hwS$gcaROxPIB`~5^1>r`)<}olZaxN&cwHS z0|dDFUm;S!DBL(wB_IU<;#>ay9@s7K`!{y$AI6}fpbs{n`?WHmPq2Y4{hyitha3M0 z77&pCkqyYh|6etQZm`5a@E6Er`Up+WQ`lc0*R(2A+}OyQfB0a6aPFUI_{J2Pz6GX` zyg~vS9hH#u(d&Yl*Tc$A9h7bfH{c+L`}O~TgV2RIl5a}BkTn3|R2C8Q}T3{e`zU-D4#NSa-h_AzsT&0-(ffB&B@2(`Xf1mNC7kwmSz zb|L@B^`@Zw?mvGC%$*h(Aj{bEpBW^aK*^jz@VIxz%zxW=w!H`|I9lOv%m5|?Rujzs zUOMlFOVB_55A&lPj@yZP{vp&hLM3|sxBLv&^Oee<@kv~gsbzl>Bm8Bex&H>llvp$-4&hacmCzYZxi5cu3?rzX_)I}(;K57 znTzXw^M6yq`sW7?goCneCm0Uh883p|Vd5stxys?Uk2ZU%?Mm@4=e_>rA3!icAzmYo z+$TxXVWudTP3fzX$L5>bQ8R|*S=~ZGqaSsiR9lboQgtLDcgNs$b>S82s^@sYZNAgt z{%2b>k&}oMF{jGKm9d5QB?$jv`+IC48(p2wPA^W+cj-QA+nvjilX%VG4z9|K7SWA< zL}(Q)%l1k`vD1xs<=_ss%S^j&H^M~-Ly0TT7Q3z@%Q4?Z|KrqvV^;%yB%l^M)xtlP zT_%o;NUvF}$N$wy6}dgW>j({1X1N&2epU$bo=m&o^G8`rz!zCg}5$vcYMv(`e7 z_^w-wFM1U9Dx7kAQ~2CEE{4ATzD6|-`j_aZfM~sywS_1QtEFdM_)&Ms>4Z^lbueT^ zMoxJ;MBSem>0s!;0ce76(MbF693S)U*0P)o<;kq|aP!;%qO5LZpZ``D=I2S`7n6JB zR@V`sfD0i*Z^Yw2NK|ChLv{6owaP_0Bv4#i#|*D?`3G3^&ksZV#JGyuPDzNuY80}A z@*KZ)h$7pWhNoYJVkbS|5qJ0Y(ppnFR(yLV_eGEv07C!SdJeFoX7D%Y+;{!ceWw1@ zvWWR$tS@D6jO^K{+(7*x1JXb|;WHXoTIF7G8DTsydEB*ai+1}Xx7WaF(ija_|oZ|555cAaM8u z4i@-9Un?x=5X6Bz_|GR;G9n{>bo_7aU3oavU(|2kY-yXHh*DCK5h6=Q)=CJ4#Eh~f z%Vf>akV-_dZzCc5HX=(HMr55CWXrBFWBn0h$zbNa!_@Cxp67ku-~aFYHP7>%x%b?2 z&OP^@&-vW@)tK8zNSFf`g|tSyOqnOoE8?cA=U+sMD_|h3gl|KE$2oul*d*o$)@}=g z1d86d@A&jZRH#y}u^@P5uhb6tjh!mfK#qIZy{!D%J*=+Mld$$ricOrG?>O>d%)cBL z|3$jy3-3qU$7PRVqz!8yS#bUdJf}2(4?8)LUzB>3`w50Xh;FHY|CM`tgwBzd^JD8*tFp zPkHjcl5HC7ldEULpHY07Od6JSW`4kDn8D>9XmO8bTeBk=Uo7N3GqYzZ0t4AAzZ-;T z(H}P;ta(tt1hz${`J^Pk)6lQR7sN?yB5g|hSuNhgZ4rS6$T*?cFc2~8;O>na7X)Ue z8klhHb7n`B(Mba-s0oSBantT0m%vl6^tb;;HlJg_7QfJXOikd48>BWM^fchd-vfkQ zK~d1+y=%R)oM3!Z6Kpg#d^d zS#3^T40(lGy_pMdnEDe?Lp2k<6$c%Cerqxi04^YoJ?XIdCvYDLP1EsdK6bb3tE9I$@7&kq(EOn zRAZ$UrlF_)m@rat$tdY~3!bWtvYfoN;?;kt$9uNoc7vrZt_xvIktRYc5-t7`a zhB`(#94}88PVkPJyq(qXnC(ad2LyfwjjXuc6d`u(?(2L3e~Q}lD1f6j_0o3GKO5yZ z6t2^>-lA|4fnO`&GbwS%|4gA!t~0(SlWEkJXeUd;k#h9luDc};k#H5i>k0R2Ebzuf z7N6Tv^**X|@bngCOUbSMtQz@Ma02A<5{(M1DO+U)Ou?jbUKb0wLLGkU+!8r?JG-gG zyzXkzx!c5i@1YM@mXL!c26?6;h>C@ zvGPMq*|`8S{G_mU#9XlAxz&!GvvSQNP@LIOVJh5HE6X73kh`dWbjqKD987z3ivf>W zV&_IQ?vQu@e-Qqje8;+xC+O(4?9*CA$NC%b=b#nq#W+QK%oF->ARs(YGz1z|RQ#H7 z30`3bd8IwkP`$e@%+WrhqG} zhOyXr(PSbO}vsSgQ?8MH8Z&R>{$X-5;$9 z1pTy_+kDAG9-5rJj}fS1XY0(n<;AjTb&1r@xGpQvju7EL*ok)hrDS8QCbR{a~+Y~G)z#7APQu(IvTj$Xjcsh1`r zU@J?JhV#O)leCZIie-!~crnM?&mL>b2`GMgp)5Z!3R{-pIGIW5EpCJ{Vi|_L3KKF9 zBi5!O^ob8&H!|oY{YyV<-e+TUkR4Zd^`lGf#3y|VqFH+c$=OqGIFr^}lHY{M3!rZm zQz%U%$K)DpICz`PfoZYOp@fT?iu7qon&ST0*ELbL%@xmJrDs!ebG6|7H z+^YN(I|j0P(f*t{ec9V`d@7Id$xlF9J=$&h%-B1Zu)>+%&Kz8ZVRMEHS+Phbr_nk2*?t@TMasN?)Om826=@W6H@-vG9MAdcoAU} z=b*y!U-wHFrazJ=9R`R&qak9ZHW-Ekei7_X}1|0odDrV%`8@`t2n&ROv zm&Z@Q#O?@cIaJkB641nxIoJpN-D{?~1COJ%?R#S` zl(X}lp(4Bo_(XEKNP&gj-FsP4WfW%H0_bP`Q92X!h14-l=th1R*V_E`jf-DN9H$N zsc5{^wIg{(C-7KPTL-sAZUe{wy5#NJGM+GC2)DsNYCIOiIRwz5BR?PryQM}?_Gsyc z3l)6U{GN6mW5-B=3;o@{)2{mv|1wn4xA#TOa++4`0uEj`obRM->v;xWPFW_&KzeW) zG2!8i?u-^nvleQgd+p0XDoHUx-tY}>^gz$guNr!lCZ+MllSO&h9bc_{y%zG*JmXPm zEBA;8lQjtS3EmD^$;D(XPIUi7ZZ!NMmWS8GEjA}snj{5v5o&1I%`e_(qi^1to!2A+ z1};$j!2G>ktb~MX9lMPO;rwZSd$`oY0plp`cXnQ9q}->CKDI23y^}N_ZyN8lPB@=E zh^k+EbS4M*?1xi4e>V4^uuFaRu?FhH9tSCvwDOnUEvB&U@)Z%s;XWZmH5qz4@`ivw zC25ZAe#2v4S@n%q6Ld_#mxXhc_d&vzP1%-(CoKA2h$6*c+bX3Bc)~R$}ALTlVu){#XAM?8Rxw#ocLNpRD;PyRc7c4 zTue+TBVif?D?DQ3NR{9Ubrricr8e}XN1 zoFKS7=&VB(k-O{=1T%~q|Ao)JE4*Fajp!|E_)13+WB{klcW*BI`y1XUb&vw{XDdY@ z{it-lm1um--pR$g6MCtm^)iFvFvYW<=1aaWT7Rb0ep(*XL=9wodyGmb!ZXBfYb)rc zaU8Yq`lA0cIs1ugsxLi6BX6rK*BFZYlFxsjD;JtQ+^ zv}sdf6^=kHqQ=F<9I4i^?E>JP&rkGw5L!pF6urR5gx;L9RR zy-lU3;*ff8qO(wQ{k%wjMM1O9@>vY!%AiJqbB;_Ft(mPNG+y4#jy`nEWe_ESMIzcC zmCpF=Ca%ga=-;BgpQ`eaaQ?FAW4u^V9*Pn>Pw2>OW7I;(H9D|J-OHZWgR*X?YAw%L z-C>yK-Ua+O6o@4Jj^TG5LJuSH_6Df&v1o*MWYb|0Kq$b|v(2ShYwL%DJ{fTkkBW3e zU4nP74P~50-;1_L$2zp`dM*uhb6d}xH6eK3Tkcm2)Pu|LmoezRR;k&K?#kF3XUOI| zk&b5e6=tseq)kcKt+}rZug+TjL@&p#dl^@n95wMOz2ln=T?H2hncSi(FQb1I_Ywc1 zJvf6-sKxtnUd+CDwLhpeu!-+q@W9IUZy^YA_pq@F850D&0`=)odk!$6G_nIF2sj;&pvxf>NAuG zR1X|gQ!uU06Oo(ortnZX_u0c=V=i4FGuOVd0_Eu5Df!dO-W&2bRKXXN-7cL*CRU=Jp zlJGy}DFGrxd9h+JrSYxps3{J+z;K-?r1zTcn~M4C)T#19@5Q<%Hay+al4SyS(fyXx z`-@0?tcLnTu=->dHa^dN1t<}sKef;A*0%)OOrcXfzli2(?TE{^9am8UpU@9q>cp6N z{hC+Dd(AJaZ~~Y`2J={NyeI=J)m0CR?<&p(8=Bd)*2f;2*Kk|5>^pQ#^~x>7ar2Y~ zp`qTJdwOy1BhR@kmKs7SbLq4*qCs;}n$XCXU{Vz{hLeKL7PxN@?B*JrXks&cl@olk zQ($WR;#kA0gQ-*PI)ukN;a2>0HX2kxcC|WIdb@Q;0CTGF3^S(++fLclv^a}(0q;?m z^w~U=ZpHP()SWmnm2b;RCOtW^3KC8!s61cky}&XaFb1xF0FR%>FAt_AC2EdADpu0$ z263`v_aUgUMN}=uHq9xmbS-l5d~9WcGC2_3xR^tXJ5wON-;l zmRG-FnVHSUQf$L9Iwg*h`hzwTkdLOtOnnUl{*<5X0it0d&EEj+_p*|p>}@1D)rU+4 z`^&bqle(9xA;dr?t>n}6792V*$`iTff>|bebfyKeT z1y}CTjXQ%1_B-b;8iZ_qwzE!Twst`1RcdST|J0>gYwRJ{^~`LUc`WI=6Bx0Z$KS@* zEtgY0Im`A5tECVie=Y_fmDa>qmdL;@Sf#OmP08eh5!IS zdk>z8YecsEZ4EMe*5Nn+Y$zNsBl~JJ_BY-FPgztjri5-xcU&gmoMBIm0AOaz~}akAOBLp{k&SzcFel=m9rQ zfMpp~$BMB6tOYo@uV&CK`@jOIO)$CWz8MB$km3WF2^&aJqWe|@_#x!zJfLryH1p@P z;Jl|I;MdGUP`f}0otxO$nG0&a7-o5xxpP162%fHK%l<3bc0UBIY6=_KJk23yaU zErbY(7*cecrenV9Wp}~7bTAwAzYQrafO|IldEVBjiGFY#DW=rY7P$rdf0#ouZeAI! z+76)U`!K29f4aVY?t?|6o{b}2&DX6x5>mG9d2ru0$uVMNaqwl`2M5|=7|K$LKR-** z{U1Pu#@F}Si4^vo3;DXVUz??e_u%_K6(@}yWkW-nbvo{nFk}8>_SOPVPDW!6Ck(+E9X z4?7-t@~i%~)~iF^UW>kh9a*|HQsV9}BIMIwR1PWgLUtOOJXpvndnD|qx1oH!N|kgu*a<4axI^1|KaKZKz-?|a&b(&H z0US$^8BoCu4aTQIVXN)28IZ`~a3iTmRH^#HHp?de2P1;Q)_t&l%%-!W&?gBx%fCR} zuVWR_Z)4I&IaiuSy4|?bY?F-xdj8KGx!+C?JYS3bam!YF<8gU0m(HDkN_1O%#z9S0 z6-ng&M=J2mM|m4+wkk%}Hfn_$UvY4I!8F^0^}g$$xWF6 diff --git a/docs/design/cross_config_implementation_report.md b/docs/design/cross_config_implementation_report.md new file mode 100644 index 00000000..7239f9d3 --- /dev/null +++ b/docs/design/cross_config_implementation_report.md @@ -0,0 +1,185 @@ +# Cross-Configuration Alignment Implementation Report + +Status: V1 framework snapshot, 2026-07-19 + +Related work: + +- [Roadmap #83](https://github.com/RL-Align/RL-Kernel/issues/83) +- [Cross-configuration alignment #111](https://github.com/RL-Align/RL-Kernel/issues/111) +- [Numerical contract #108](https://github.com/RL-Align/RL-Kernel/issues/108) +- [V1 contract](cross_config_logprob_drift_contract.md) + +## Result and claim boundary + +This change provides a small framework for planning and executing paired +rollout/training logprob comparisons across controlled configuration changes. It +includes strict configuration loading, bounded case planning, exact semantic +operator selection, lifecycle-aware runtime materialization, paired read-only +scoring, fixed-contract comparison, append-only artifacts, and validated resume. + +The included executable path is deliberately CPU-only. It validates framework +plumbing with a synthetic model and temporary selected-logprob backends; it does +not claim production vLLM, FSDP, TP, CP, accelerator, or distributed numerical +alignment. The S1, S2, and S3 examples are plans, not execution evidence. + +## Architecture + +The implementation keeps configuration, operator resolution, runtime ownership, +and execution separate: + +```text +JSON -> ExperimentConfig -> Planner -> ExperimentPlan + | + build_execution_plan + | + operator-bound ExecutionPlan + / \ + operator session RuntimeMaterializer + | + RuntimeBinding + | + ArtifactStore <- PairedRunner -> fixed comparator +``` + +| Boundary | Responsibility | +| --- | --- | +| `config.py` and `planner.py` | Load strict, versioned JSON; normalize the ten supported knobs; emit an `ExperimentPlan` containing a baseline plus declared OAT or explicit pairwise cases under a fixed 256-case cap; compute stable semantic case IDs without importing a runtime. | +| `build_execution_plan` | Resolve rollout/training selections, bind them into immutable case identity, and emit canonical operator-bound `ExecutionPlan` rows shared by planning and execution. | +| `SemanticOperatorCatalog` | Store immutable backend descriptors and their target, device, dtype, per-target required topology, lifecycle, factory, and observability constraints. | +| `OperatorSession` | Resolve and instantiate exact rollout/training implementations for one case, cache only within that case, and produce concrete provenance. | +| `RuntimeMaterializer` | Apply each normalized knob through an owning adapter and report requested, materialized, and actual values with status and lifecycle evidence. | +| `RuntimeBinding` | Carry only backend-neutral batch, side-configuration, topology, scorer, operator-backend, and runtime-kind mappings. Runtime-specific engine objects stay behind the adapter boundary. | +| `PairedRunner` | Supervise isolated rollout/training scoring children, enforce timeout and read-only model state, validate ranks and exact operator instances, compare selected logprobs, and coordinate resume/publication. | +| `ArtifactStore` | Publish immutable attempt directories and write `COMPLETE` last with SHA-256 seals for every required payload. Resume accepts only an attempt whose identity, execution, provenance, tensors, and comparison still validate. | + +Runtime bindings deep-freeze their execution handoff. Rollout topology contains +only rollout-owned world/TP/CP state, while training topology contains only +training-owned world/sharding state; neither side is padded with fields owned by +the other. + +The fixed-threshold comparator applies the repository numerical contract only to +active selected tokens. Identity violations, invalid artifacts, non-finite +scores, and zero active tokens cannot become passes. Diagnostics remain separate +from the pass/fail rule. + +## Configuration and operator selection + +Execution controls do not live in experiment JSON. `scenario` is metadata; the +CLI chooses planning versus execution, the runtime adapter, temporary-operator +authorization, timeout, and resume policy. + +`logp.backend` is the concise choice when rollout and training use the same +selected-logprob implementation. The optional top-level `operators` mapping is +the extension point for independent sides and per-backend options: + +```json +{ + "baseline": { + "logp": {"backend": "rlkernel.reference_logp"} + }, + "operators": { + "selected_logprob": { + "rollout": "rlkernel.reference_logp", + "training": { + "backend": "vendor.training_logp", + "options": {"mode": "exact"} + } + } + } +} +``` + +The explicit rollout backend must agree with the baseline `logp.backend`. +Explicit operators cannot be combined with `logp.backend` interventions because +that would make the planned knob differ from the implementation actually used. +Unknown fields, threshold overrides, hidden execution controls, duplicate JSON +keys, and non-finite values are rejected. + +## CLI + +Planning validates and records a plan without constructing a runtime: + +```bash +python -m rl_engine.alignment.cross_config plan \ + examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json +``` + +The only shipped execution adapter is the explicit CPU smoke runtime: + +```bash +python -m rl_engine.alignment.cross_config run \ + examples/cross_config_s0_cpu_smoke.json \ + --runtime cpu-smoke \ + --allow-smoke-operators +``` + +Both commands accept `--output-root`. `run` also exposes a per-attempt timeout +and `--no-resume`; these policies are intentionally absent from the JSON schema. + +## CPU testing adapter + +`rl_engine.alignment.testing.cpu_cross_config` owns the synthetic causal model, +stateless CPU scorer, `CpuSmokeMaterializer`, canonical batch construction, and +CPU experiment helpers. Keeping these objects outside the core package prevents +test hardware and model assumptions from becoming runtime abstractions. + +Temporary selected-logprob implementations live under +`rl_engine/alignment/testing/smoke_ops`. They advertise CPU as their only device, +are not registered by default, and require explicit policy authorization. The +reference backend performs `log_softmax` plus gather; the offset backend is used +only by focused mismatch tests. The shipped S0 example selects reference on both +sides and contains one baseline case. + +## Scenario evidence + +| Scenario | Planned cases | Current evidence | +| --- | ---: | --- | +| S0 CPU framework smoke | 1 | One reference/reference case passes on CPU; a second invocation validates and resumes the same complete attempt. | +| S1 distributed smoke | 5 | Configuration loading and planning only. | +| S2 vLLM TP versus FSDP | 10 | Configuration loading and planning only. | +| S3 Qwen3-8B TP=4, CP=4, BF16 | 11 | Configuration loading and planning only. | + +Planning success does not imply that the requested production topology can be +materialized. Unsupported or unobservable runtime settings fail strict execution +instead of silently falling back. + +## Validation snapshot + +```text +Focused contract/runtime/runner/CLI and existing regression tests: +60 passed in 2.20s + +CPU-collectable repository suite: +418 passed, 242 skipped in 15.83s + +Named scenario checks: +S0 run: 1 pass, then 1 validated resume +S1/S2/S3 plan: 5 / 10 / 11 cases, no runtime constructed +``` + +The final review also checks JSON syntax, formatting, static typing, strict +documentation build, and `git diff --check`. + +The full CPU command excludes `test_grpo_loss.py` and `test_ratio_kl.py` +because those modules require Triton during collection. It ran outside the +restricted sandbox so Gloo and POSIX shared-memory tests could use host +resources. + +## Extension path and known gaps + +A production backend extends the semantic catalog with a descriptor and factory, +then supplies injection/read-back hooks through its runtime adapter. The planner +and runner do not need backend-specific branches. Operator correctness remains +owned by the operator implementation; the framework verifies exact selection, +materialization evidence, paired identity, comparison, and provenance. + +Production execution still requires: + +- verified selected-logprob injection and read-back for the rollout engine; +- a read-only pre-update training scorer for FSDP and distributed rank evidence; +- context-parallel application/read-back and process-group orchestration; +- accelerator-backed lifecycle and cleanup tests; and +- the production kernels tracked by their owning workstreams. + +Until those adapters exist, S1-S3 remain reproducible planning inputs and the CPU +smoke remains a framework claim only. diff --git a/docs/design/cross_config_logprob_drift_contract.md b/docs/design/cross_config_logprob_drift_contract.md new file mode 100644 index 00000000..3f06bc18 --- /dev/null +++ b/docs/design/cross_config_logprob_drift_contract.md @@ -0,0 +1,308 @@ +# Cross-Configuration Logprob Drift Contract + +Status: V1 implementation contract + +Related work: + +- [Roadmap #83](https://github.com/RL-Align/RL-Kernel/issues/83) +- [Cross-configuration alignment #111](https://github.com/RL-Align/RL-Kernel/issues/111) +- [Numerical contract #108](https://github.com/RL-Align/RL-Kernel/issues/108) + +## Goal and boundary + +This framework isolates configuration changes that can make rollout-selected +log probabilities differ from training-side recomputation. It provides typed +plans, lifecycle-aware runtime materialization, exact semantic-operator +selection, paired read-only scoring, append-only artifacts, and safe resume. + +It does not implement production AG, RS, GEMM, attention, logprob, TP-invariant, +CP-aware, or deterministic collective kernels. Those implementations remain +owned by their operator workstreams and integrate through the semantic operator +catalog described below. + +## The only pass/fail rule + +For every active selected response/action token: + +```text +abs(training_logprob - rollout_logprob) > fixed_threshold +``` + +The fixed threshold is loaded from the repository numerical contract. It is not +a config field, CLI flag, experiment axis, workload policy, or operator option. +Equality with the threshold is not a mismatch. + +```python +mismatch_mask = active_mask & ( + torch.abs(training_logprobs - rollout_logprobs) > fixed_threshold +) +``` + +Mean, percentiles, maximum absolute difference, mismatch ratio, worst-token +location, and approximate KL are diagnostics only. They never change pass/fail. +The token-level artifact persists both logprob tensors, the active mask, and the +resolved threshold so the mask can be recomputed offline. + +Zero active tokens produce `ZERO_ACTIVE_TOKENS`, never a pass. Non-finite or +non-floating active scores produce `INVALID_ARTIFACT`. + +## Identity before numerics + +A comparison is valid only when both scorers use the same logical input: + +- immutable checkpoint and model version; +- tokenizer ID and tokenization policy; +- generated token IDs and selected-token IDs; +- active and attention masks; +- pre-update model state; +- required position, cache, and packing metadata. + +The training scorer teacher-forces the already generated sequence. It cannot +generate replacement tokens, use a KV cache when the frozen identity forbids +one, own an optimizer, update parameters or buffers, or leave the model in a +different mode. An identity violation is `INVALID_IDENTITY`, not numerical +drift. + +## V1 configuration + +The user supplies one explicit baseline plus declared interventions. Lists do +not imply a Cartesian product. + +```json +{ + "schema_version": "cross_config.experiment_config.v1", + "experiment_id": "qwen3-8b-alignment", + "scenario_id": "qwen3-8b-tp4-cp4-bf16", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "identity": {"...": "frozen scoring identity"}, + "baseline": { + "batch": {"size": 8}, + "rollout": { + "tensor_parallel_size": 4, + "context_parallel_size": 4, + "dtype": "bfloat16", + "enable_prefix_caching": true, + "enforce_eager": false + }, + "training": { + "attention_backend": "flash_attention_2", + "compute_dtype": "bfloat16", + "sharding": "fsdp" + }, + "logp": {"backend": "native"} + }, + "interventions": [ + {"path": "batch.size", "values": [1]}, + {"path": "logp.backend", "values": ["rlkernel.reference_logp"]} + ], + "scenario": {"level": "S3", "device": "cuda"} +} +``` + +`scenario` is metadata. Execution mode and authorization policy belong to the +CLI, so a config cannot hide `plan_only`, expected test outcomes, or permission +to activate temporary operators. + +### Exact knob allowlist + +| Knob | Minimum lifecycle | Meaning | +|---|---|---| +| `batch.size` | request | Canonical sample chunking only. | +| `rollout.tensor_parallel_size` | process | Rollout TP world. | +| `rollout.context_parallel_size` | process | Rollout CP world. | +| `rollout.dtype` | engine construction | Rollout numerical dtype. | +| `rollout.enable_prefix_caching` | engine construction | Engine cache policy. | +| `rollout.enforce_eager` | engine construction | Eager versus optimized/graph path. | +| `training.attention_backend` | engine construction | Training scorer attention implementation. | +| `training.compute_dtype` | engine construction | Training scorer compute dtype. | +| `logp.backend` | engine construction | Both-sides selected-logprob shortcut. | +| `training.sharding` | process | Training topology, such as unsharded or FSDP. | + +TP/vocabulary layout is derived and recorded, not user-settable. Tokenization, +masks, positions, checkpoint identity, and pre-update state are invariants, not +ordinary knobs. Quantization, FP8, MoE, speculative decoding, pipeline +parallelism, and arbitrary runtime fields are deferred. + +### Planning + +`one_at_a_time` emits one baseline and cases that change exactly one declared +path. `pairwise` is opt-in and expands only explicitly listed path pairs. The +planner normalizes aliases, validates the allowlist and capability constraints, +and reports structured issues without creating engines. A fixed 256-case +framework cap stops OAT or pairwise expansion before unbounded accumulation. + +Stable case IDs hash normalized requested values, identity, contract version, +and scenario definition. Runtime readback never rewrites a case ID. A retry gets +a new attempt ID under the same case. + +## Architecture and extension points + +The core has one-way responsibilities: + +```text +strict config -> Planner -> ExperimentPlan -> build_execution_plan + | + operator-bound ExecutionPlan + | + runtime adapter -> RuntimeBinding + | + paired runner + / \ + artifact store comparator +``` + +- `config.py` owns the external schema and strict JSON loading. +- `schema.py` owns immutable, versioned domain records. +- `planner.py` owns normalization, the knob catalog, OAT, and pairwise cases. +- `execution_plan.py` binds the selected rollout/training operators into each + immutable case and produces canonical rows shared by planning and execution. +- `runtime.py` owns the adapter protocol, three-stage materialization, lifecycle + fingerprints, and a backend-neutral execution binding. +- `comparison.py` owns identity validation and the fixed-threshold result. +- `runner.py` coordinates paired execution and atomic publication; private + execution, provenance, and resume modules isolate process supervision and + validation details. +- `artifacts.py` owns append-only attempts and resume discovery. +- `semantic_registry.py` owns generic operator descriptors and case-local + resolution sessions; it is shared by future alignment features. + +The package root exposes only the common planning and execution facade. Runtime, +artifact, schema, and operator internals remain in their owning modules. + +### Runtime adapters + +A runtime adapter receives the normalized case and returns one application +record per knob: + +```text +requested -> materialized -> actual +``` + +Each record includes status (`applied`, `fallback`, `unsupported`, +`unobservable`, or `error`), evidence, and lifecycle. The facade derives +construction, distributed-context, and process fingerprints. Reuse is allowed +only when all relevant fingerprints and operator bindings match. + +Adapters may construct repository-native vLLM, training, or stateless config +objects internally. The core runner receives only backend-neutral batch, side +configuration, topology, scorer, operator-backend, and runtime-kind mappings, +so adding a runtime does not add branches to the planner or runner. + +Strict execution rejects fallback, ignored settings, unobservable critical +values, stale registry state, and incompatible reuse. Fallback is measurable +only when it is itself the declared intervention. + +## Semantic operator selection + +The first semantic operator is `selected_logprob`. Rollout and training can +select implementations independently: + +```json +{ + "operators": { + "selected_logprob": { + "rollout": "rlkernel.reference_logp", + "training": { + "backend": "rlkernel.reference_logp", + "options": {} + } + } + } +} +``` + +When `operators` is absent, `logp.backend` selects the same backend on both +sides. An explicit mapping is bound into execution identity before a runtime is +created. A `logp.backend` intervention cannot be combined with a fixed explicit +mapping because that would create a knob that no longer changes execution. + +Each backend descriptor declares: + +- semantic operation and backend ID; +- supported target tags, devices, dtypes, and per-target required topology values; +- alignment properties and lifecycle; +- implementation factory and version/build fingerprint; +- explicit fallback policy and temporary-test marker. + +`SemanticOperatorCatalog` stores immutable descriptors. Each case creates an +`OperatorSession` for resolution, instantiation, caching, and provenance. Failed +or cached state cannot leak into the next case. Strict resolution never invokes +legacy priority fallback. + +Adding a production implementation requires its existing semantic interface, +one descriptor, runtime injection hooks where needed, operator-owned correctness +tests, and one framework case. It does not require a planner change. + +## Artifacts and resume + +Attempts are append-only: + +```text +runs// + experiment.json + plan.jsonl + cases/// + requested.json + materialized.json + actual.json + identity.json + score_rollout.pt + score_training.pt + comparison.json + token_diffs.pt + COMPLETE +``` + +`COMPLETE` is published last and seals every required payload with a SHA-256 +digest. Resume accepts only a complete attempt whose case, identity, +materialization, scorer, operator, environment, comparison, and tensor artifacts +match the current execution key. Partial, malformed, or tampered attempts are +ignored; an older valid attempt may still be reused. Existing files are never +overwritten. + +## CPU smoke boundary + +The only executable adapter delivered here is under +`rl_engine.alignment.testing.cpu_cross_config`. It is explicitly CPU-only and +uses a deterministic synthetic model plus read-only stateless scoring. Named +distributed and accelerator scenarios are configuration/plan coverage only. + +Temporary selected-logprob backends live together under +`rl_engine/alignment/testing/smoke_ops`: + +- `smoke_only.logp_reference`: PyTorch `log_softmax` plus gather; +- `smoke_only.logp_offset`: the same result with an authorized deterministic + offset used to prove mismatch detection. + +They are CPU-only, marked `is_smoke_only`, unregistered by default, and require +both explicit registration and execution policy authorization. Their exact +removal procedure is in `SMOKE_OPERATORS.md`. Remove them when equivalent +production operators pass the same framework cases, then remove the opt-in flag +and temporary test marker. + +## Scenario levels and claims + +| Level | Purpose | Current claim | +|---|---|---| +| S0 | Local CPU framework smoke | Executable: config, planner, operator selection, paired scoring, comparison, artifacts, resume. | +| S1 | Small distributed lifecycle smoke | Plan only until suitable hardware/runtime adapters exist. | +| S2 | Named vLLM TP versus training FSDP comparison | Plan only. | +| S3 | Qwen3-8B TP=4, CP=4, BF16 milestone | Plan only; no production alignment claim. | + +Run the shipped examples with: + +```bash +python -m rl_engine.alignment.cross_config plan \ + examples/cross_config_s3_qwen3_8b_tp4_cp4_bf16.json + +python -m rl_engine.alignment.cross_config run \ + examples/cross_config_s0_cpu_smoke.json \ + --runtime cpu-smoke \ + --allow-smoke-operators +``` + +Passing S0 proves framework plumbing only. It does not prove accelerator, +distributed, production-operator, or roadmap numerical alignment. diff --git a/docs/design/ws2_cross_config_logprob_drift_contract.md b/docs/design/ws2_cross_config_logprob_drift_contract.md deleted file mode 100644 index 3d7fac06..00000000 --- a/docs/design/ws2_cross_config_logprob_drift_contract.md +++ /dev/null @@ -1,874 +0,0 @@ -# WS2 Cross-Config Logprob Drift Contract - -Status: RFC - -Tracking issues: - -- [#111: WS2 cross-config alignment](https://github.com/RL-Align/RL-Kernel/issues/111) -- [#108: WS1 numerical contract](https://github.com/RL-Align/RL-Kernel/issues/108) - -## Motivation - -WS2 covers rollout and training paths that use different parallelism strategies, such as -rollout tensor parallelism and training FSDP. The alignment problem is not a single-op -accuracy check. It is end-to-end floating-point drift across tokenizer, masks, serving, -rollout, and training recomputation before any optimizer update. - -For PPO, GRPO, and related RL post-training algorithms, the most direct pre-update signal -is selected-token log probability drift. If rollout-side `old_logprobs` and train-side -recomputed log probabilities disagree for the same checkpoint, same token ids, same masks, -and same model version, classify the failure as infrastructure, precision, mask, -tokenizer, or serving-path drift. Do not classify that failure as an algorithm or reward -problem until pre-update logprob alignment is clean. - -Aggregate KL-style diagnostics are useful but not sufficient as the primary WS2 contract. -In training-inference mismatch cases, KL estimates can stay flat or fail to expose the -early failure phase, because the first-order issue is token-level rollout-vs-training -probability disagreement before the optimizer update, not necessarily a large aggregate -policy-space shift. - -## Framework Upgrade Overview - -The upgrade moves cross-config validation from two separately configured execution paths -that require a manual comparison into a shared runtime flow. `RuntimeTools` coordinates -the rollout and training configurations, while `PairedRunner` collects their selected -logprobs and performs the comparison automatically. This keeps both sides aligned on the -same inputs and makes drift visible as part of the run rather than as a follow-up manual -check. - -![Before and after framework upgrade](../assets/ws2-cross-config-before-after.png) - -The left side shows the pre-upgrade flow, where `VLLMSamplerConfig` and -`TorchRLTrainingConfig` feed independent executors and the results are manually compared. -The right side shows the upgraded flow, where `RuntimeTools` and `PairedRunner` connect the -two paths and produce an automatic comparison while preserving the shared -`KernelRegistry` contract. - -## Scope - -This RFC defines what WS2 cross-config alignment measures, how failures are classified, -and the modular implementation roadmap for making that contract executable. It does not -itself add a test harness, distributed tests, runtime gates, layer-wise probes, or -distributed fixes. - -Out of scope for this document: - -- Implementing multi-GPU test infrastructure. -- Adding runtime pass/fail gates. -- Adding automatic layer-wise drift probes. -- Fixing TP, FSDP, SP, cache, mask, tokenizer, or serving-path bugs. -- Defining a second numerical tolerance table. -- Reimplementing work owned by the adjacent TP, SP, collective, training-integration, or - layer-probe issues referenced by the roadmap below. - -## Measurement Contract - -The primary metric is selected-token logprob drift: - -```text -dlogp = train_recomputed_logp - rollout_old_logp -``` - -Compute `dlogp` only on active response/action tokens. Prompt tokens, padding tokens, and -masked-out response positions are excluded from every aggregate metric. - -The comparison must use teacher-forcing scoring on the training side. The scored sequence -is the already-sampled rollout sequence; the training path must not resample or regenerate -tokens for this contract. - -The rollout and training values are comparable only when they share the same logical -inputs: - -- Same checkpoint and same model version. -- Same input token ids. -- Same selected response/action token ids. -- Same attention mask and action mask. -- Same tokenizer version and tokenization policy. -- Same padding layout semantics, including left-padding or right-padding behavior. -- Same pre-update state, before any optimizer step, weight sync, or policy mutation that - belongs to the next training step. - -If the implementation has explicit position ids, cache-position metadata, sequence ids, or -packed-sequence metadata, those inputs are part of the comparison contract as well. - -## Primary Failure Signal - -The pass/fail decision starts from `dlogp` over active tokens. Reward, gradnorm, -weightnorm, and update norm are downstream symptoms. They are useful for debugging and -triage, but they are not the primary contract for cross-config alignment. - -The zero-update expectation is: - -```text -train_recomputed_logp ~= rollout_old_logp -ratio0 ~= 1 -approx_kl0 ~= 0 -``` - -The acceptable meaning of `~=` is defined by the WS1 per-dtype numerical threshold table -from [#108](https://github.com/RL-Align/RL-Kernel/issues/108). This RFC defines the -measurement surface and classification rules only. - -## Diagnostics - -All diagnostics are computed on active response/action tokens only. - -| Metric | Definition | Purpose | -| --- | --- | --- | -| `ratio0` | `exp(dlogp)` | Zero-update policy ratio implied by train-vs-rollout logprob drift. | -| `clipfrac0` | Mean indicator that `ratio0` falls outside the configured PPO/GRPO clip range. | Detects whether drift alone would trigger clipping before any update. | -| `approx_kl0` | Masked mean of `exp(dlogp) - 1 - dlogp`. | Zero-update approximate KL implied by logprob drift. | -| `mean_abs_dlogp` | Mean of `abs(dlogp)`. | Average selected-token drift. | -| `p95_abs_dlogp` | 95th percentile of `abs(dlogp)`. | Tail drift below outliers. | -| `p99_abs_dlogp` | 99th percentile of `abs(dlogp)`. | High-tail drift. | -| `max_abs_dlogp` | Maximum of `abs(dlogp)`. | Worst selected-token mismatch. | - -When the run is distributed, report optional per-rank versions of the same metrics. The -per-rank view should preserve enough metadata to identify the rollout rank, training rank, -parallelism mode, dtype, padding side, cache mode, and local active-token count for that -rank. - -## Tolerance Source - -This RFC does not define a separate numerical tolerance table. The single source of truth -for acceptable numerical drift is the per-dtype threshold table owned by -[#108](https://github.com/RL-Align/RL-Kernel/issues/108). - -For WS2, acceptable numerical drift means that `max_abs_dlogp` over active -response/action tokens satisfies the WS1 per-dtype threshold from #108. If the #108 table -changes, WS2 inherits that policy without editing this document or maintaining a second -table. - -## Tolerance Interpretation and Effect-Based Validation - -Numerical tolerances in this RFC are infrastructure contract thresholds, not a universal -statement of algorithmic harmlessness. There is no model-independent scale that proves a -given train-vs-rollout logprob difference is harmless for every algorithm, reward model, -prompt distribution, sequence length, or optimization schedule. Any hand-written threshold -encodes a prior about acceptable numerical error. WS2 therefore does not introduce an -additional algorithmic noise budget, nor does it define a new estimator for tolerable -logprob noise. - -The #108 threshold defines whether rollout and training paths are numerically aligned -enough to continue debugging the failure as an algorithmic or reward problem. It does not -prove that all smaller drift is behaviorally irrelevant, and it does not imply that all -larger drift is the only cause of downstream failure. - -When downstream model-effect validation is available, such as reward trajectory, train KL, -eval win rate, collapse rate, policy regression tests, or task-specific success metrics, -use it as a severity and root-cause prioritization signal. It must not replace the -pre-update selected-token logprob contract. A run can be numerically out of contract even -if a short downstream run appears healthy, and a run can be numerically in contract while -still failing because of algorithmic tuning, reward hacking, insufficient KL control, or -data issues. - -The intended interpretation is: - -```text -#108 per-dtype threshold: - numerical infrastructure contract - -selected-token dlogp: - primary WS2 train-vs-rollout drift surface - -downstream model effect: - practical severity and algorithmic relevance signal - -KL / ratio / percentile diagnostics: - debugging and triage signals, not replacement pass/fail criteria -``` - -## Drift Source Taxonomy - -Before treating train-vs-rollout drift as generic algorithmic noise, WS2 should classify -likely sources of mismatch. At minimum, the following source classes should be considered -separately. - -### Arithmetic Schedule Drift - -Arithmetic schedule drift comes from different floating-point operation order between -rollout and training. This includes different kernels, fused vs unfused implementations, -compiler-generated graph rewrites, attention implementation differences, matmul epilogue -differences, accumulation dtype differences, and changes introduced by advanced compilers -or graph optimizers. - -This class answers the question: - -```text -Do rollout and training compute mathematically equivalent expressions using different -floating-point schedules? -``` - -Examples include: - -- Fused attention vs unfused attention. -- Different FlashAttention or SDPA backends. -- Fused RMSNorm or LayerNorm vs decomposed normalization. -- Compiler-reordered graph segments. -- Different matmul epilogues or activation fusion. -- Different accumulation precision in otherwise equivalent kernels. - -### Reduction and Collective Drift - -Reduction drift comes from operations whose floating-point result depends on reduction -order, parallel topology, or concurrent execution. This includes local reductions, -cross-rank reductions, all-reduce, reduce-scatter, gather/scatter patterns, sharded logits -or loss computation, tensor-parallel collectives, FSDP reductions, and nondeterministic -reduction scheduling. - -This class answers the question: - -```text -Does the mismatch appear because rollout and training aggregate partial results in -different orders or across different rank topologies? -``` - -Examples include: - -- TP logits produced through a different collective path from the training path. -- FSDP reduce-scatter or all-gather changing accumulation order. -- Per-rank partial reductions with different shard boundaries. -- Loss or logprob reductions performed before vs after cross-rank communication. -- Nondeterministic collective algorithms or concurrent reductions. - -### Quantization and Dequantization Drift - -Quantization drift comes from representing weights, activations, KV cache, logits, or -intermediate tensors with different quantization policies between rollout and training. -Quantization is not merely a floating-point ordering issue; it introduces representation -noise through scales, zero points, clipping, grouping, calibration, and dequantization -paths. - -This class answers the question: - -```text -Does the mismatch appear because rollout and training use different numerical -representations or quantization policies? -``` - -Examples include: - -- Rollout uses weight-only quantization while training recomputation uses bf16/fp16 - weights. -- Different quantization group sizes. -- Different activation quantization or KV-cache quantization policy. -- Different scale computation or calibration data. -- Different dequantization placement relative to fused kernels. -- Serving-path quantization that is absent from the training path. - -### Logical Input and Metadata Drift - -Logical input mismatch must be ruled out before interpreting any result as numerical -drift. This class includes tokenizer version, tokenization policy, attention mask, action -mask, padding side, explicit position ids, cache positions, sequence ids, packed-sequence -metadata, and serving-path request formatting. - -This class answers the question: - -```text -Are rollout and training actually scoring the same logical sequence under the same masking -and positional semantics? -``` - -If this class is not clean, the comparison is invalid rather than merely noisy. - -## Decision Rule - -Use this order when classifying a cross-config failure: - -1. If pre-update selected-token logprobs do not match under the same checkpoint, same token - ids, same masks, and same model version, treat the failure as infrastructure, - precision, mask, tokenizer, or serving-path drift. -2. If `max_abs_dlogp` violates the #108 threshold but downstream metrics look healthy in a - short run, keep the issue classified as infrastructure drift. Short-horizon model - health does not prove the drift is safe. -3. If KL or ratio diagnostics move before gradnorm or update norm moves, treat the failure - as likely infrastructure or logprob plumbing. -4. If gradnorm or update norm moves first and KL moves later, treat the failure as more - likely algorithmic tuning, such as learning rate, KL beta, reward scale, or advantage - outliers. -5. If only some ranks drift, treat the failure as distributed infrastructure until rank - placement, shard boundaries, collective algorithms, local active-token counts, masks, - and cache-position issues are ruled out. -6. If reward rises and then collapses while pre-update logprob alignment is clean, treat - the failure as more likely algorithmic, reward hacking, data-related, or insufficient - KL constraint. - -This classification does not prove root cause by itself. It defines the first branch in -the debugging tree so WS2 bugs do not get misfiled as reward or algorithm regressions -before the zero-update logprob contract is satisfied. - -## Layered Ablation Strategy - -WS2 should not treat train-vs-rollout mismatch as a single undifferentiated error source. -Later tests should use a layered ablation strategy that changes one source class at a time -whenever the implementation allows it. - -The minimum useful ablation structure is: - -```text -A0. Fully aligned reference - Same checkpoint, same dtype policy, same kernels where possible, same reduction - topology where possible, same quantization policy, same tokenizer, same masks, same - padding, same cache/position metadata. - -A1. Arithmetic-schedule-only mismatch - Keep logical inputs, reduction topology, and quantization policy aligned. Allow only - kernel, fusion, compiler, or graph execution differences. - -A2. Reduction-topology-only mismatch - Keep logical inputs, kernel policy, and quantization policy aligned. Allow only - reduction order, collective topology, sharding, or rank placement differences. - -A3. Quantization-only mismatch - Keep logical inputs, kernel policy, and reduction topology aligned. Allow only - quantization, dequantization, scale, group size, or representation differences. - -A4. Pairwise mismatches - Enable two mismatch classes at a time: - arithmetic + reduction - arithmetic + quantization - reduction + quantization - -A5. Full production mismatch - Use the real rollout and training configurations, including all production - differences. -``` - -Each ablation should collect the same primary and diagnostic metrics: - -```text -primary: - dlogp over active response/action tokens - max_abs_dlogp - -diagnostics: - mean_abs_dlogp - p95_abs_dlogp - p99_abs_dlogp - ratio0 - clipfrac0 - approx_kl0 - per-rank versions when distributed - -metadata: - dtype - kernel/backend choices - fusion/compiler mode - reduction/collective topology - quantization policy - padding side - cache mode - position/cache-position metadata - active-token count -``` - -When downstream model-effect validation is available, the same ablations should also -record practical training outcomes, for example reward trajectory, training KL, entropy, -clip fraction, update norm, collapse rate, and task-specific evaluation metrics. These -downstream metrics are not the WS2 pass/fail contract, but they help rank which numerical -mismatch class matters most for the workload. - -## Ablation Interpretation Rules - -Use these rules when reading the ablation matrix: - -1. If the fully aligned reference fails, the issue is not a cross-config mismatch yet. - First debug the base scoring path, masks, tokenizer, position metadata, checkpoint - identity, or implementation correctness. -2. If a single-source ablation fails the `max_abs_dlogp` contract, that source class is - sufficient to create unacceptable train-vs-rollout drift under the tested workload. For - example, if only quantization is misaligned and the run fails, quantization is a - dominant source candidate for that task and configuration. -3. If all single-source ablations pass, but pairwise or full-production mismatches fail, - the failure is likely an interaction effect. Identify the minimal failing pair before - attributing the issue to any single subsystem. -4. If one single-source ablation passes the numerical contract but shows materially worse - downstream model effect, record it as behaviorally sensitive even if it remains - numerically in contract. This is a signal that the #108 infrastructure tolerance may be - sufficient for numerical alignment but not necessarily predictive of algorithmic - robustness for that workload. -5. If pre-update logprob alignment is clean but downstream training still collapses, - classify the failure as more likely algorithmic, reward-related, data-related, or - KL-control-related rather than cross-config numerical drift. - -## Minimal and Layered Alignment Principle - -The governing principle of WS2 is **minimal alignment**: - -> Keep rollout and training semantically identical, then align only the smallest numerical -> layer needed to satisfy the selected-token logprob contract. - -WS2 does not require every internal tensor, kernel, reduction, or execution schedule to be -identical. If the production rollout and training paths already satisfy the #108 -`logprob` tolerance, no numerical alignment change is required. Different engines are -allowed to keep different high-performance implementations. - -Minimal alignment does not relax logical correctness. Checkpoint/version, token ids, -masks, tokenizer semantics, and required position metadata must match exactly. A logical -input mismatch invalidates the experiment; it is not acceptable numerical drift. - -### Alignment Ladder - -Use the following ladder in order and stop at the first level that satisfies the contract: - -| Level | Action | Production implication | -| --- | --- | --- | -| L0: semantic identity | Make logical inputs and model version exactly comparable. | Mandatory for every case. | -| L1: observable contract | Keep both production paths unchanged and compare selected-token logprobs. | Stop here if #108 passes. | -| L2: source isolation | Change one declared knob at a time to locate the smallest sufficient drift source. | Diagnostic only; do not change production yet. | -| L3: local alignment | Align or fix one operator, collective, metadata field, or representation policy. | Preferred production fix when L1 fails. | -| L4: layered alignment | Align the smallest interacting pair or contiguous layer boundary that is required. | Use only when no single local change is sufficient. | -| L5: full/bitwise alignment | Force broad identical paths or reference implementations. | Diagnostic fallback, not the default WS2 exit criterion. | - -The chosen fix should minimize, in order: - -1. semantic scope changed; -2. number of aligned knobs; -3. performance and memory overhead; -4. engine-specific intrusion; -5. maintenance burden. - -A fix is incomplete if it proves only that the fully aligned reference passes. It must -also show that unrelated rollout/training differences can remain enabled. Conversely, WS2 -must not reject a configuration merely because internal tensors are not bitwise equal when -the selected-token contract passes. - -## Controller-Centered Design - -The central feature is an ablation controller, not a hard-coded list of distributed -tests. It separates experiment planning from engine-specific knob application. - -```mermaid -flowchart LR - Definition["ExperimentDefinition
identity + baseline + axes + constraints"] - Planner["GridPlanner
product / one-at-a-time / pairwise"] - Isolation["IsolationValidator
declared deltas only"] - Definition --> Planner --> Isolation - - Isolation --> Cases["ExperimentCase[]
stable ids + provenance"] - - subgraph Materializers["Knob materializers"] - Rollout["vLLM adapter"] - Training["stateless / FSDP adapter"] - Kernel["kernel policy adapter"] - Environment["process/build environment adapter"] - end - - Cases --> Materializers - Materializers --> Runner["isolated paired runner"] - Runner --> Samples["canonical alignment samples"] - Samples --> Comparator["identity validator + dlogp comparator"] - Comparator --> Cube["result cube
axes + per-rank reports + cost"] - Cube --> Analyzer["minimal sufficient alignment analyzer"] -``` - -### Core Objects - -The implementation should expose a small typed model rather than passing more loose -dictionaries through the current executors: - -- `SemanticIdentitySpec`: checkpoint/weight version, tokenizer, fixed token sequences, - masks, and position metadata that must match. -- `ScorerSpec`: rollout or training engine, world size, device/dtype, and immutable engine - construction settings. -- `KnobDefinition`: one controllable source of variation. -- `ExperimentDefinition`: baseline scorers plus axes, constraints, and measurement policy. -- `ExperimentCase`: one fully materialized grid point with a stable content-derived id. -- `AlignmentSample`: logical tensors, selected logprobs, and actual runtime provenance. -- `AlignmentResult`: global/per-rank drift, pass/fail, actual applied knobs, and optional - cost metrics. -- `ResultCube`: results indexed by normalized knob values, independent of execution order. - -Every `KnobDefinition` must declare: - -```text -name: - stable dotted name, for example rollout.tensor_parallel_size - -source_class: - logical-layout | arithmetic | reduction | representation | execution - -lifecycle: - request | engine-construction | process-start | build - -targets: - rollout | training | both | kernel - -domain: - allowed typed values - -capability: - how an adapter proves that a value is supported - -constraints: - incompatible or conditional combinations - -apply: - engine-specific materialization hook - -provenance: - how the actual applied value is read back and reported -``` - -The controller must compare requested and actual provenance. A silent runtime fallback is -an invalid ablation unless the fallback itself is the declared knob under test. - -### Grid Composition - -`GridPlanner` should support the following modes over the same typed axes: - -- `product`: full Cartesian grid; -- `one_at_a_time`: baseline plus one changed factor per case; -- `pairwise`: covering pairs without requiring the full Cartesian product; -- `zip`: paired values such as compatible model/dtype artifacts; -- fixed overrides and named slices; -- capability and compatibility constraints; -- deterministic case ids, filtering, resume, and retry. - -A normal workflow starts with `one_at_a_time`, expands to `pairwise` only when single -factors do not explain the failure, and uses `product` for an explicit grid search. CI -runs a named slice of the same definition rather than maintaining a separate handwritten -test matrix. - -For every generated case, `IsolationValidator` compares its normalized spec with the -baseline and rejects undeclared changes. This is what makes an arithmetic-only, -reduction-only, or quantization-only claim trustworthy. - -### Minimal Sufficient Alignment Analysis - -The analyzer treats "align this knob between rollout and training" as an intervention. It -reports the smallest passing intervention set found by the executed grid: - -```text -production mismatch: - fail - -align attention backend only: - fail - -align logp reduction only: - pass - -minimal sufficient alignment candidate: - {logp.reduction_policy} - -unrelated differences left enabled: - attention backend, cache policy, TP/FSDP topology -``` - -This result is evidence for the smallest effective intervention, not automatic proof of -root cause. A later fix PR still needs the smallest reproducer and a local regression. - -## Mapping to Current Code - -The controller should initially map to existing configuration surfaces instead of -introducing a second execution stack. - -| High-level knob | Current code path | Required adapter behavior | -| --- | --- | --- | -| `rollout.tensor_parallel_size` | `VLLMSamplerConfig.engine_kwargs` | Materialize `tensor_parallel_size` before vLLM engine construction and read it back from runtime metadata. | -| `rollout.dtype` | `VLLMSamplerConfig.engine_kwargs["dtype"]` | Normalize string/torch dtype and record the actual engine dtype. | -| `sampling.temperature` | `VLLMSamplerConfig.sampling_params` | Apply per request; require the same scoring semantics on both sides. | -| `execution.prefix_cache` | `VLLMSamplerConfig.enable_prefix_caching` | Treat as engine-construction-time, not a request toggle. | -| `training.attention_backend` | `StatelessForwardConfig.attention_backend` | Apply before forward and report requested backend plus any actual fallback. | -| `training.output_dtype` | `StatelessForwardConfig.output_dtype` | Keep observation dtype separate from model compute dtype. | -| `training.compute_dtype` | `TorchRLTrainingConfig.dtype` and FSDP model construction | Materialize before wrapping/sharding the model. | -| `logp.backend` | `RolloutExecutor` / `TorchRLTrainingConfig.logp_backend` | Reuse `resolve_logp_op_type()` aliases and report the resolved op type and concrete backend class. | -| `logp.deterministic` | `require_batch_invariant_logp` | Express policy intent; do not hard-code a CUDA implementation in the controller. | -| `training.sharding` | new score-only FSDP adapter | Materialize world size and sharding strategy before process-group/model construction. | -| `logp.tp_layout` | `linear_logp` `tp_group`, `vocab_start_index`, `global_vocab_size` | Record shard boundaries and reject incomplete ownership metadata. | -| `kernel.fast_math` | `KERNEL_ALIGN_USE_FAST_MATH` | Treat as build-time and bind the case to a distinct built artifact. | -| `kernel.sm90_path` | `KERNEL_ALIGN_FORCE_SM90` and compiled extension | Capability-gate by architecture and build artifact; never switch it after import. | - -The existing `KernelRegistry` caches instances and resolves priority maps during -initialization. vLLM TP, dtype, and prefix caching also belong to engine construction. -Therefore the runner must not mutate these values in a long-lived process and assume the -next case is isolated. - -Cases may share a worker only when their engine-construction and process-start -fingerprints are identical. Request-time knobs may reuse that worker. Build-time knobs -always select a prebuilt artifact and a separate process. The artifact id and extension -build metadata are part of result provenance. - -## Kernel Integration Contract - -Kernel work may require a new or rewritten implementation, but the ablation controller -must not know CUDA/Triton class names or kernel launch details. - -The kernel boundary should expose a backend descriptor with: - -- stable backend id and semantic operator name; -- supported device architectures, dtypes, shapes, and parallel layouts; -- determinism/alignment properties; -- required TP/SP metadata and collectives; -- configuration lifecycle, including build-time flags; -- concrete implementation selected at runtime; -- fallback behavior; -- version/build fingerprint. - -The controller requests a policy such as `production`, `reference`, `deterministic`, or a -stable backend id. The kernel adapter resolves that policy through `KernelRegistry` and -records the concrete implementation. Strict WS2 cases reject an undeclared fallback. - -A rewritten kernel integrates cleanly by: - -1. implementing the existing operator semantic interface; -2. registering a new stable backend descriptor; -3. passing #108 operator accuracy and batch-invariance checks; -4. declaring TP/SP metadata and supported lifecycle knobs; -5. adding one isolated end-to-end controller case; -6. reporting performance/memory overhead against the production backend. - -It should not require a new branch in `GridPlanner`. If a framework cannot inject the -kernel through a supported hook, its engine adapter reports the knob as unsupported; it -must not claim that the ablation ran. - -## Repository Fit - -The current repository already provides useful pieces: - -- #108 owns `tolerance_contract.json`. -- `VLLMSamplerConfig` exposes loose `engine_kwargs`, `sampling_params`, and prefix-cache - configuration. -- `StatelessForwardConfig` exposes attention backend, temperature, and output dtype. -- `TorchRLTrainingConfig` exposes compute dtype, `logp_backend`, and the deterministic - requirement. -- `resolve_logp_op_type()` already separates user-facing logp policy from registry op type. -- TP `linear_logp` already accepts explicit process group and vocab-shard metadata. -- `RolloutStageResult` and the weight bridge carry iteration/weight version. -- `StatelessForwardExecutor` is a reusable no-update teacher-forcing scorer. - -The missing pieces are the typed experiment model, actual-value provenance, strict scoring -payload, FSDP score-only adapter, lifecycle-aware knob materializers, grid planner, and -result cube. - -`DeepSpeedTrainingWorker.train()` still performs backward/step and constructs its current -objective's `old_logps` from recomputed values. It is not a WS2 comparator. A later -DeepSpeed scorer must be a separate read-only adapter. - -## Ownership Boundaries - -| Issue | Boundary | -| --- | --- | -| [#108](https://github.com/RL-Align/RL-Kernel/issues/108) | Owns numerical thresholds. | -| [#109](https://github.com/RL-Align/RL-Kernel/issues/109) | Owns deterministic TP reduction implementations. | -| [#110](https://github.com/RL-Align/RL-Kernel/issues/110) | Owns SP-aware operators and reductions. | -| [#112](https://github.com/RL-Align/RL-Kernel/issues/112) | Owns deterministic collective implementations. | -| [#113](https://github.com/RL-Align/RL-Kernel/issues/113) | Owns the later distributed forward/backward chain gate. | -| [#116](https://github.com/RL-Align/RL-Kernel/issues/116) | Shares the tolerance/report foundation implemented by B1. | -| [#127](https://github.com/RL-Align/RL-Kernel/issues/127) | Owns the pinned multi-GPU dual-engine environment. | -| [#130](https://github.com/RL-Align/RL-Kernel/issues/130) | Owns full FSDP/Megatron training integration and backward. | -| [#131](https://github.com/RL-Align/RL-Kernel/issues/131) | Owns the later production cross-benchmark command. | -| [#136](https://github.com/RL-Align/RL-Kernel/issues/136) | Owns automatic layer-wise probes. | - -## Revised Modular PR Roadmap - -The identifiers below are roadmap labels, not existing GitHub PR numbers. The former -Phases A, B, and C are consolidated because they jointly form the baseline infrastructure. - -### Phase 1: Baseline Infrastructure - -#### B1 — Alignment contract, comparator, and report - -**Scope:** Expose #108 tolerance lookup; add canonical identity/provenance/sample types, -logical comparability validation, active-token drift metrics, and one JSON/human report. - -**Acceptance:** CPU tests cover identity mismatch, masks, percentiles, zero active tokens, -worst-token metadata, and dtype-specific pass/fail without copying threshold values. - -**Why one PR:** These types form one public contract and cannot provide useful independent -behavior when landed separately. - -#### B2 — Exact rollout and teacher-forcing scoring adapters - -**Scope:** Normalize vLLM sampled-token logprobs and rollout provenance; add strict -rollout-to-teacher-forcing collation; define the read-only scorer protocol and adapt -`StatelessForwardExecutor`. - -**Acceptance:** A fixture round trip preserves prompt/generated ids, masks, selected -logprobs, weight version, and available position metadata. Missing identity data or -undeclared backend fallback fails explicitly. Repeated scoring does not change model state. - -**Non-goal:** No FSDP, subprocess runner, or grid planner. - -#### B3 — Score-only FSDP adapter and baseline controls - -**Scope:** Add a PyTorch FSDP scorer with no optimizer/backward, then add A0 identical -stateless scoring and unsharded-vs-FSDP controls. - -**Acceptance:** A0 passes on CPU; a labeled two-GPU/NCCL control proves FSDP recomputation -is clean and model state is unchanged. - -**Non-goal:** Full training integration remains in #130. - -#### B4 — Paired runner, artifacts, and rank aggregation - -**Scope:** Launch rollout/training scorers with independent world sizes; write versioned -canonical artifacts; enforce timeout/cleanup; aggregate deterministic per-rank/global -reports. - -**Acceptance:** CPU fixtures cover child failure, timeout, malformed artifact, duplicate -or missing ranks, weight-version mismatch, and global worst-token selection. - -**Design requirement:** The runner accepts separate construction/process/build -fingerprints so the later controller can isolate cases correctly. - -### Phase 2: Composable Ablation Controller - -#### C1 — Typed experiment model, knob registry, and grid planner - -**Scope:** Implement `ExperimentDefinition`, typed knob descriptors, constraints, -capability declarations, stable case ids, and `product`, `one_at_a_time`, `pairwise`, and -`zip` planners. - -**Acceptance:** Pure CPU tests generate deterministic grids, reject invalid combinations, -resume by case id, and prove each one-at-a-time case changes exactly one declared knob. - -**Non-goal:** Do not launch engines in this PR. - -#### C2 — Runtime knob materializers and capability checks - -**Scope:** Map controller knobs to current vLLM, stateless, and FSDP configuration -surfaces. Separate request-time, engine-construction, and process-start application. Read -back actual values and construction fingerprints. - -**Acceptance:** Fake engine adapters prove every requested value is either applied and -reported or rejected as unsupported. No silent fallback is accepted in strict cases. - -#### C3 — Kernel policy bridge - -**Scope:** Add the backend descriptor and kernel materializer boundary described above. -Adapt existing logp policy aliases and TP metadata without changing kernel math. - -**Acceptance:** The same experiment definition can select production/reference/ -deterministic logp policies and report the concrete registry backend. A fake rewritten -kernel registers without a controller code change. - -**Non-goal:** Kernel rewrites discovered later remain one-root-cause fix PRs. - -#### C4 — Grid executor and result cube - -**Scope:** Execute C1 cases through B4, pool only workers with identical lifecycle -fingerprints, select build artifacts, persist results, and expose filtering/resume plus a -machine-readable result cube. - -**Acceptance:** An interrupted fake grid resumes without rerunning completed cases; -requested and actual provenance are queryable for every axis. - -### Phase 3: Core Scenario and Minimal Alignment - -#### M1 — TP=2 rollout versus FSDP diagnostic grid - -**Scope:** Define the first real experiment using the controller: fixed model/tokenizer/ -tokens, vLLM TP=2 rollout, FSDP recomputation, bf16, and production defaults. Generate the -production point plus one-at-a-time alignment interventions. - -**Acceptance:** Execution and reports succeed on the pinned #127 environment. Numerical -failure is recorded without weakening #108. - -#### M-FIX-N — One minimal root cause per PR - -Each fix PR consumes the smallest controller case that exposes one problem. A kernel -rewrite, collective change, metadata fix, or adapter fix remains separate. - -A fix must show: - -- the failing production or isolated case; -- the smallest intervention that makes it pass; -- one local implementation change; -- A0 and unrelated-knob regressions; -- actual backend provenance; -- performance/memory cost when applicable. - -#### M2 — Promote the minimally aligned core case to a gate - -**Scope:** After required M-FIX PRs, gate TP=2/FSDP using the smallest passing alignment -set, not a fully reference configuration. - -**Acceptance:** The report names which knobs were aligned, which differences remained -enabled, and why a broader alignment level was unnecessary. - -### Phase 4: Grid Coverage and Ablation Closure - -#### G1 — Required composable grid - -**Scope:** Add the required batch-size, padding/layout, dtype, and cache/position axes as -declarative knob values and constraints. Allow full product, named slices, and -one-at-a-time views from the same definition. - -**Acceptance:** A user can request, for example: - -```text -batch_size = [1, 8] -padding_side = [left, right] -dtype = [fp32, bf16, fp16] -prefix_cache = [off, on] -logp.backend = [production, deterministic] -``` - -without writing a new test function. Unsupported combinations are capability-filtered -with explicit reasons, and every result is indexed in the same cube. - -#### G2 — A0-A5 profile and minimal-alignment wrapper - -**Scope:** Express the RFC's A0-A5 ablations as presets over C1 rather than separate test -implementations: - -- A0: fully aligned diagnostic reference; -- A1: arithmetic one-at-a-time; -- A2: reduction/topology one-at-a-time; -- A3: representation/quantization one-at-a-time; -- A4: pairwise expansion only when needed; -- A5: production mismatch. - -Add a CLI/config wrapper that selects profiles, axes, filters, and output location. - -**Acceptance:** Phase F behavior is only a planner/profile layer over G1. It adds no -engine-specific branching. - -#### G3 — Targeted GPU CI and downstream handoff - -**Scope:** Run a curated named slice of the same grid in labeled GPU CI, upload the result -cube, and expose fixtures/reports to #113/#131. - -**Acceptance:** CI distinguishes launch/environment/numerical failure, always cleans up, -and does not maintain a second handwritten matrix. - -## PR Sizing Rules - -The consolidated roadmap uses fewer baseline PRs, but later numerical fixes remain small: - -1. B1-B4 may each land one cohesive baseline subsystem. -2. C1-C4 each own one controller layer: planning, runtime materialization, kernel policy, - or execution/results. -3. Adding a new ordinary knob changes one descriptor and one engine adapter, not the - planner. -4. Adding a rewritten kernel changes the kernel implementation and its backend descriptor, - not the controller. -5. M-FIX PRs contain one root cause only. -6. G1/G2 add declarative grids/profiles and must not include numerical fixes. -7. No PR adds or copies a tolerance value. - -## Completion Criteria for #111 - -#111 is complete when: - -1. semantic identity validation is strict and independent of numerical alignment; -2. the controller can compose, filter, resume, and report a multidimensional configuration - grid; -3. requested knobs are verified against actual runtime/kernel provenance; -4. the TP=2/FSDP production case is brought into #108 contract using the documented - smallest sufficient alignment set; -5. batch, padding/layout, dtype, and cache/position axes are available through G1 without - new test functions; -6. at least one production, one arithmetic, and one reduction/topology slice run through - the same controller/report path; -7. a rewritten kernel can register through C3 without changing grid-planner code; -8. the stable core and selected grid slice run in targeted GPU CI; -9. forward fixtures and result cubes are reusable by #113 and #131. - -SP, additional TP sizes, quantization variants, exhaustive product grids, and downstream -training effects remain extensions unless maintainers promote named grid slices into the -required gate. Full/bitwise internal alignment is not a completion criterion unless a -separate contract explicitly requires it. From dd2ce92b723b24f590692ad6307158aac921ae6d Mon Sep 17 00:00:00 2001 From: vensen Date: Mon, 17 Aug 2026 09:04:53 +0000 Subject: [PATCH 09/33] feat(rocm): add multi-architecture ROCm build support --- csrc/deterministic_logp_kernel.cu | 17 ++- csrc/fused_logp_kernel.cu | 27 +++-- csrc/ops.cpp | 6 +- docker/Dockerfile.rocm | 26 +++++ docker/Dockerfile.rocm_base | 44 ++++++++ setup.py | 172 +++++++++++++++++++----------- 6 files changed, 214 insertions(+), 78 deletions(-) create mode 100644 docker/Dockerfile.rocm create mode 100644 docker/Dockerfile.rocm_base diff --git a/csrc/deterministic_logp_kernel.cu b/csrc/deterministic_logp_kernel.cu index 1ca23287..c238e11d 100644 --- a/csrc/deterministic_logp_kernel.cu +++ b/csrc/deterministic_logp_kernel.cu @@ -15,6 +15,15 @@ constexpr int kDeterministicLogpMediumVocabLimit = 4096; constexpr int kDeterministicLogpWarpSize = 32; constexpr float kDeterministicLogpNegInf = -3.4028234663852886e38F; +template +__device__ __forceinline__ T deterministic_logp_shfl_down_32(T value, unsigned int delta) { +#if defined(__HIPCC__) || defined(__HIP_PLATFORM_AMD__) + return __shfl_down(value, delta, kDeterministicLogpWarpSize); +#else + return __shfl_down_sync(0xffffffffu, value, delta, kDeterministicLogpWarpSize); +#endif +} + template struct DeterministicLogpBlockTraits { static_assert( @@ -36,7 +45,7 @@ __device__ __forceinline__ float deterministicBlockReduceMax(float val) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { - val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset)); + val = fmaxf(val, deterministic_logp_shfl_down_32(val, offset)); } if (lane == 0) { @@ -50,7 +59,7 @@ __device__ __forceinline__ float deterministicBlockReduceMax(float val) { if (wid == 0) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { - val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset)); + val = fmaxf(val, deterministic_logp_shfl_down_32(val, offset)); } } return val; @@ -66,7 +75,7 @@ __device__ __forceinline__ float deterministicBlockReduceSum(float val) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { - val += __shfl_down_sync(0xffffffff, val, offset); + val += deterministic_logp_shfl_down_32(val, offset); } if (lane == 0) { @@ -80,7 +89,7 @@ __device__ __forceinline__ float deterministicBlockReduceSum(float val) { if (wid == 0) { #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { - val += __shfl_down_sync(0xffffffff, val, offset); + val += deterministic_logp_shfl_down_32(val, offset); } } return val; diff --git a/csrc/fused_logp_kernel.cu b/csrc/fused_logp_kernel.cu index b620b047..679a6a30 100644 --- a/csrc/fused_logp_kernel.cu +++ b/csrc/fused_logp_kernel.cu @@ -5,6 +5,17 @@ #include #include +constexpr int kFusedLogpLogicalWarpSize = 32; + +template +__device__ __forceinline__ T fused_logp_shfl_down_32(T value, unsigned int delta) { +#if defined(__HIPCC__) || defined(__HIP_PLATFORM_AMD__) + return __shfl_down(value, delta, kFusedLogpLogicalWarpSize); +#else + return __shfl_down_sync(0xffffffffu, value, delta, kFusedLogpLogicalWarpSize); +#endif +} + template __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) { static __shared__ float shared[32]; @@ -14,7 +25,7 @@ __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) { float f_val = static_cast(val); for (int offset = 16; offset > 0; offset /= 2) - f_val = max(f_val, __shfl_down_sync(0xffffffff, f_val, offset)); + f_val = max(f_val, fused_logp_shfl_down_32(f_val, offset)); if (lane == 0) shared[wid] = f_val; __syncthreads(); @@ -22,7 +33,7 @@ __device__ __forceinline__ scalar_t blockReduceMax(scalar_t val) { f_val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : -1e20f; if (wid == 0) { for (int offset = 16; offset > 0; offset /= 2) - f_val = max(f_val, __shfl_down_sync(0xffffffff, f_val, offset)); + f_val = max(f_val, fused_logp_shfl_down_32(f_val, offset)); } return static_cast(f_val); } @@ -36,7 +47,7 @@ __device__ __forceinline__ scalar_t blockReduceSum(scalar_t val) { float f_val = static_cast(val); for (int offset = 16; offset > 0; offset /= 2) - f_val += __shfl_down_sync(0xffffffff, f_val, offset); + f_val += fused_logp_shfl_down_32(f_val, offset); if (lane == 0) shared[wid] = f_val; __syncthreads(); @@ -44,7 +55,7 @@ __device__ __forceinline__ scalar_t blockReduceSum(scalar_t val) { f_val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f; if (wid == 0) { for (int offset = 16; offset > 0; offset /= 2) - f_val += __shfl_down_sync(0xffffffff, f_val, offset); + f_val += fused_logp_shfl_down_32(f_val, offset); } return static_cast(f_val); } @@ -82,8 +93,8 @@ __device__ __forceinline__ LogSumExpState blockReduceLogSumExp(LogSumExpState st for (int offset = 16; offset > 0; offset /= 2) { LogSumExpState other{ - __shfl_down_sync(0xffffffff, state.max_val, offset), - __shfl_down_sync(0xffffffff, state.sum_exp, offset)}; + fused_logp_shfl_down_32(state.max_val, offset), + fused_logp_shfl_down_32(state.sum_exp, offset)}; state = merge_logsumexp_state(state, other); } @@ -100,8 +111,8 @@ __device__ __forceinline__ LogSumExpState blockReduceLogSumExp(LogSumExpState st if (wid == 0) { for (int offset = 16; offset > 0; offset /= 2) { LogSumExpState other{ - __shfl_down_sync(0xffffffff, state.max_val, offset), - __shfl_down_sync(0xffffffff, state.sum_exp, offset)}; + fused_logp_shfl_down_32(state.max_val, offset), + fused_logp_shfl_down_32(state.sum_exp, offset)}; state = merge_logsumexp_state(state, other); } } diff --git a/csrc/ops.cpp b/csrc/ops.cpp index eee328a4..97b753d3 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -253,6 +253,7 @@ std::vector deterministic_attention_backward( // Prefix-Shared Attention Declarations & Wrappers +#if !defined(USE_ROCM) void prefix_shared_attention_forward( const __nv_bfloat16 *Q, // [bs, G, len_q, DIM] const __nv_bfloat16 *K, // [bs, len_kv, DIM] @@ -296,6 +297,7 @@ at::Tensor prefix_shared_attention( return O; } #endif +#endif // PyBind11 Module Registration PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { @@ -355,8 +357,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("deterministic_logp_forward_indexed_out", &deterministic_logp_forward_indexed_out, "Batch-invariant deterministic logp indexed out"); m.def("deterministic_logp_forward_indexed_fp32", &deterministic_logp_forward_indexed_fp32, "Batch-invariant deterministic logp indexed fp32"); - // registry Prefix-Shared Attention + // Prefix-shared attention uses NVIDIA PTX and falls back to PyTorch SDPA on ROCm. +#if !defined(USE_ROCM) m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO"); +#endif // registry Batch-Invariant Deterministic GEMM m.def("det_gemm_fwd", &det_gemm_fwd, "Batch-invariant deterministic GEMM forward (C=A@B)"); diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm new file mode 100644 index 00000000..f83905fe --- /dev/null +++ b/docker/Dockerfile.rocm @@ -0,0 +1,26 @@ +# docker/Dockerfile.rocm +# base: docker build -f docker/Dockerfile.rocm_base -t rl-kernel:rocm-dev . +# build: docker build -f docker/Dockerfile.rocm -t /rl-kernel-ci:rocm . +# push: docker push /rl-kernel-ci:rocm + +FROM rl-kernel:rocm-dev + +# Build a portable extension by default. The list follows the multi-architecture +# ROCm profile used by vLLM: MI200 (gfx90a), MI300/MI325 (gfx942), MI350/MI355 +# (gfx950), plus supported RDNA 3/4 targets. Override it at build time with +# --build-arg PYTORCH_ROCM_ARCH=, or at run time with `docker run -e`. +# Newer GPU targets need no setup.py change: pass the gfx target supported by the +# installed PyTorch/ROCm pair. +ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1150;gfx1151;gfx1200;gfx1201 +ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV MAX_JOBS=8 + +WORKDIR /opt/rl-kernel + +COPY pyproject.toml setup.py* requirements*.txt ./ + +RUN pip install --no-cache-dir -U pip \ + && pip install --no-cache-dir -r requirements.txt \ + && pip install --no-cache-dir pytest + +WORKDIR /workspace/RL-Kernel diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base new file mode 100644 index 00000000..e4334bc0 --- /dev/null +++ b/docker/Dockerfile.rocm_base @@ -0,0 +1,44 @@ +# docker/Dockerfile.rocm_base +# Build: docker build -f docker/Dockerfile.rocm_base -t rl-kernel:rocm-dev . +# +# To build only the targets deployed in a particular image, override the default: +# docker build -f docker/Dockerfile.rocm_base \ +# --build-arg PYTORCH_ROCM_ARCH='gfx942;gfx950' -t rl-kernel:rocm-dev . + +ARG BASE_IMAGE=rocm/dev-ubuntu-22.04:7.2.3-complete +FROM ${BASE_IMAGE} + +# This is a build-target list, not a hardware allow-list. It covers the target +# families supported by the vLLM ROCm 7.2 reference image: MI200 (gfx90a), +# MI300/MI325 (gfx942), MI350/MI355 (gfx950), and RDNA 3/4. For future GPUs, +# pass the target accepted by the selected ROCm/PyTorch toolchain at build time. +ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1150;gfx1151;gfx1200;gfx1201 +# Keep this wheel index aligned with BASE_IMAGE's ROCm release when overriding it. +ARG PYTORCH_INDEX_URL=https://download.pytorch.org/whl/rocm7.2 +ARG PYTORCH_VERSION=2.12.1 + +ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:${PATH} +ENV ROCM_PATH=/opt/rocm +ENV LD_LIBRARY_PATH=/opt/rocm/lib:/usr/local/lib:${LD_LIBRARY_PATH} +ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} +ENV MAX_JOBS=8 +ENV DEBIAN_FRONTEND=noninteractive + +RUN apt-get update -y \ + && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + cmake \ + git \ + ninja-build \ + pkg-config \ + python3 \ + python3-dev \ + python3-pip \ + python3-venv \ + && python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel \ + && python3 -m pip install --no-cache-dir --index-url "${PYTORCH_INDEX_URL}" "torch==${PYTORCH_VERSION}" \ + && python3 -c "import torch; assert torch.version.hip is not None, torch.__version__" \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /workspace/RL-Kernel diff --git a/setup.py b/setup.py index a17ecb40..53c6bdfc 100644 --- a/setup.py +++ b/setup.py @@ -26,14 +26,12 @@ def _load_torch_extension_tools(): import torch from torch.utils.cpp_extension import BuildExtension, CUDAExtension except ImportError: - return None, None, None, None + return None, None, None - try: - from torch.utils.cpp_extension import ROCMExtension - except ImportError: - ROCMExtension = None - - return torch, BuildExtension, CUDAExtension, ROCMExtension + # CUDAExtension is also the supported extension entry point for ROCm + # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when + # torch.version.hip is set. + return torch, BuildExtension, CUDAExtension def _cuda_define_from_env(name: str, macro: str) -> list[str]: @@ -46,8 +44,42 @@ def _cuda_define_from_env(name: str, macro: str) -> list[str]: return [f"-D{macro}={parsed}"] +_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( + "-Xfatbin", + "-compress-all", + "-gencode", + "--generate-code", + "--expt-", + "-lineinfo", + "-allow-unsupported-compiler", + "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", +) +_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { + "-Xfatbin", + "-gencode", + "--generate-code", +} + + +def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: + """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" + filtered_flags = [] + skip_next = False + for flag in flags: + if skip_next: + skip_next = False + continue + if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: + skip_next = True + continue + if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): + continue + filtered_flags.append(flag) + return filtered_flags + + def get_extensions(): - torch, _, CUDAExtension, ROCMExtension = _load_torch_extension_tools() + torch, _, CUDAExtension = _load_torch_extension_tools() if torch is None: return [] @@ -56,47 +88,49 @@ def get_extensions(): torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = torch.version.hip is not None + is_rocm = getattr(torch.version, "hip", None) is not None - if is_rocm and ROCMExtension is not None: - extensions.append( - ROCMExtension( - name="rl_engine._C", - sources=[ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cpp", - ], - extra_compile_args={ - "cxx": ["-O3", "-std=c++17"], - "hipcc": ["-O3", "--use_fast_math", "-Xhipcc", "-compress-all"], - }, - extra_link_args=list(torch_rpath), - ) + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, + # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also + # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add + # --offload-arch. Do not require a visible GPU when a ROCm target was + # explicitly selected. + no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() + if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: + raise RuntimeError( + "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " + "Set one or more ';'-separated targets, for example " + "PYTORCH_ROCM_ARCH='gfx942;gfx950'." ) - elif torch.cuda.is_available(): + + if is_rocm or torch.cuda.is_available(): cuda_sources = [ "csrc/ops.cpp", "csrc/fused_logp_kernel.cu", "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/attention/prefix_shared_attention.cu", "csrc/cuda/gemm/det_gemm_kernel.cu", "csrc/cuda/rmsnorm.cu", "csrc/cuda/activation.cu", "csrc/cuda/attention/deterministic_attention.cu", ] + if not is_rocm: + # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). + # The ROCm dispatcher falls back to PyTorch SDPA for this operator. + cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): nvcc_flags.append("--use_fast_math") - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") + if not is_rocm: + cc_major, cc_minor = torch.cuda.get_device_capability() + enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + if not enable_sm90: + # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. + nvcc_flags.append( + f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" + ) + nvcc_flags.append("--expt-relaxed-constexpr") + nvcc_flags.append("--expt-extended-lambda") nvcc_flags.extend( _cuda_define_from_env( "FUSED_LOGP_TWOPASS_BLOCK_SIZE", @@ -139,44 +173,52 @@ def get_extensions(): "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", ) ) - if envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): nvcc_flags.append("-lineinfo") - if os.name == "nt" and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC): + if ( + not is_rocm + and os.name == "nt" + and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) + ): nvcc_flags.append("-allow-unsupported-compiler") nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] extra_link_args = list(torch_rpath) - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - "csrc/cuda/embedding_lm_head_sm90.cu", # single-card batch-invariant embedding/lm-head - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - if enable_sm90 and present_sm90: - tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant - cuda_sources.extend(present_sm90) - nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: + if not is_rocm: + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + "csrc/cuda/embedding_lm_head_sm90.cu", # single-card batch-invariant embedding/lm-head + ] + enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + if enable_sm90 and present_sm90: + tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant + cuda_sources.extend(present_sm90) + nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp + # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in + # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. + enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" + if enable_det_gemm_sm90: + tma_arch = f"{cc_major}{cc_minor}a" + arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" + if arch_flag not in nvcc_flags: + nvcc_flags.append(arch_flag) + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + if is_rocm: + nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) extensions.append( CUDAExtension( @@ -194,7 +236,7 @@ def get_extensions(): def get_cmdclass(): - _, BuildExtension, _, _ = _load_torch_extension_tools() + _, BuildExtension, _ = _load_torch_extension_tools() if BuildExtension is None: return {} return {"build_ext": BuildExtension} From 0dd1b7300aad6d2b095a39578450493a6cf3117a Mon Sep 17 00:00:00 2001 From: vensen Date: Mon, 17 Aug 2026 09:23:22 +0000 Subject: [PATCH 10/33] style: satisfy flake8 line length --- setup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 53c6bdfc..94777289 100644 --- a/setup.py +++ b/setup.py @@ -192,7 +192,8 @@ def get_extensions(): "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - "csrc/cuda/embedding_lm_head_sm90.cu", # single-card batch-invariant embedding/lm-head + # Single-card batch-invariant embedding/lm-head. + "csrc/cuda/embedding_lm_head_sm90.cu", ] enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] From 898ed55c70e2b50883ac83fdda07dee730103066 Mon Sep 17 00:00:00 2001 From: vensen Date: Mon, 17 Aug 2026 10:10:15 +0000 Subject: [PATCH 11/33] fix(rocm): enforce native extension builds --- README.md | 8 ++++-- docker/Dockerfile.rocm | 2 ++ docker/Dockerfile.rocm_base | 8 ++++++ docs/getting_started/faq.md | 16 ++++++++---- docs/getting_started/installation.md | 16 ++++++------ envs.py | 1 + rl_engine/__init__.py | 4 +++ setup.py | 39 ++++++++++++++++++++++++++-- 8 files changed, 77 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index f1ab58d8..0be7704e 100644 --- a/README.md +++ b/README.md @@ -117,8 +117,12 @@ RL-Kernel sits between high-level alignment libraries and low-level GPU kernels, git clone https://github.com/RL-Align/RL-Kernel.git cd RL-Kernel -# Install core dependencies (CUDA 12.4+ recommended) -pip install -e . +# CPU-only / pure-Python fallback +python -m pip install -e . + +# Native CUDA or ROCm extension (install a matching PyTorch build first) +RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e . +python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)" ``` ### Contributions diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index f83905fe..93ec25f8 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -15,6 +15,7 @@ ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1150;gfx1151;gfx12 ENV PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} ENV MAX_JOBS=8 +USER root WORKDIR /opt/rl-kernel COPY pyproject.toml setup.py* requirements*.txt ./ @@ -23,4 +24,5 @@ RUN pip install --no-cache-dir -U pip \ && pip install --no-cache-dir -r requirements.txt \ && pip install --no-cache-dir pytest +USER rlkernel WORKDIR /workspace/RL-Kernel diff --git a/docker/Dockerfile.rocm_base b/docker/Dockerfile.rocm_base index e4334bc0..ab577175 100644 --- a/docker/Dockerfile.rocm_base +++ b/docker/Dockerfile.rocm_base @@ -16,6 +16,9 @@ ARG PYTORCH_ROCM_ARCH=gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1150;gfx1151;gfx12 # Keep this wheel index aligned with BASE_IMAGE's ROCm release when overriding it. ARG PYTORCH_INDEX_URL=https://download.pytorch.org/whl/rocm7.2 ARG PYTORCH_VERSION=2.12.1 +ARG RL_KERNEL_USER=rlkernel +ARG RL_KERNEL_UID=10001 +ARG RL_KERNEL_GID=10001 ENV PATH=/opt/rocm/llvm/bin:/opt/rocm/bin:${PATH} ENV ROCM_PATH=/opt/rocm @@ -39,6 +42,11 @@ RUN apt-get update -y \ && python3 -m pip install --no-cache-dir --upgrade pip setuptools wheel \ && python3 -m pip install --no-cache-dir --index-url "${PYTORCH_INDEX_URL}" "torch==${PYTORCH_VERSION}" \ && python3 -c "import torch; assert torch.version.hip is not None, torch.__version__" \ + && groupadd --gid "${RL_KERNEL_GID}" "${RL_KERNEL_USER}" \ + && useradd --uid "${RL_KERNEL_UID}" --gid "${RL_KERNEL_GID}" --create-home --shell /bin/bash "${RL_KERNEL_USER}" \ + && install -d --owner "${RL_KERNEL_USER}" --group "${RL_KERNEL_USER}" /workspace/RL-Kernel \ && rm -rf /var/lib/apt/lists/* +ENV HOME=/home/${RL_KERNEL_USER} WORKDIR /workspace/RL-Kernel +USER ${RL_KERNEL_USER} diff --git a/docs/getting_started/faq.md b/docs/getting_started/faq.md index 878c8c14..27cdd0c1 100644 --- a/docs/getting_started/faq.md +++ b/docs/getting_started/faq.md @@ -13,8 +13,8 @@ change more often than RL-Kernel's public API. | --- | --- | --- | --- | | Read docs or edit docs | `pip install -r requirements-docs.txt` | No | Use `mkdocs build --strict -f mkdocs.yaml` before opening a PR. | | Run CPU/mock tests | `pip install -e ".[dev]"` | No | Matches the default CI style: fallback and mocked integration coverage. | -| Run CUDA operators | `pip install -e ".[cuda]"` | Yes, NVIDIA | Requires a CUDA-enabled PyTorch wheel and a working CUDA toolchain for source builds. | -| Run ROCm operators | `pip install -e ".[rocm]"` | Yes, AMD | Requires a ROCm-enabled PyTorch wheel and ROCm compiler/runtime environment. | +| Run CUDA operators | `RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[cuda]"` | Yes, NVIDIA | Requires a CUDA-enabled PyTorch wheel and a working CUDA toolchain. | +| Run ROCm operators | `RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[rocm]"` | Yes, AMD | Requires a ROCm-enabled PyTorch wheel and ROCm compiler/runtime environment. | | Run real vLLM rollout | `pip install -e ".[vllm]"` | Runtime-dependent | Core tests do not need vLLM; install this only where real vLLM is used. | Do not install every optional extra by default. Install the smallest environment @@ -84,7 +84,13 @@ The important rule is that PyTorch must match your runtime: ```bash git clone https://github.com/RL-Align/RL-Kernel.git cd RL-Kernel + +# CPU-only / pure-Python fallback pip install -e . + +# Native CUDA or ROCm extension (after installing a matching PyTorch build) +RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e . +python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)" ``` The examples on this page use `python3` for system-level commands. Inside an @@ -93,9 +99,9 @@ activated virtual environment, `python` is also fine. ### Which optional extras exist? ```bash -pip install -e ".[cuda]" -pip install -e ".[rocm]" -pip install -e ".[vllm]" +RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[cuda]" +RL_KERNEL_REQUIRE_EXT=1 pip install --no-build-isolation -e ".[rocm]" +pip install --no-build-isolation -e ".[vllm]" pip install -e ".[dev]" ``` diff --git a/docs/getting_started/installation.md b/docs/getting_started/installation.md index 71fe227d..0e8bad46 100644 --- a/docs/getting_started/installation.md +++ b/docs/getting_started/installation.md @@ -14,15 +14,15 @@ git clone https://github.com/RL-Align/RL-Kernel.git cd RL-Kernel # Optional: pin the compile target. If unset, the build targets your GPU's arch. # export TORCH_CUDA_ARCH_LIST="9.0+PTX" # e.g. Hopper; or "8.6+PTX", "12.0+PTX" -pip install --no-build-isolation -e . +RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e . ``` -Without `--no-build-isolation`, PyTorch is invisible to the isolated build -environment, the extension is silently skipped, and the library falls back to the -slower pure-PyTorch kernels. Confirm the compiled extension is present with: +`RL_KERNEL_REQUIRE_EXT=1` makes the build fail if `_C` cannot be compiled. Without +`--no-build-isolation`, PyTorch is invisible to the isolated build environment. +Confirm the compiled extension is present with: ```bash -python -c "from rl_engine import _C; print('compiled extension OK')" +python -c "import rl_engine._C as _C; assert hasattr(_C, 'fused_logp'); print(_C.__file__)" ``` A CPU-only install (plain `pip install -e .` on a machine with no GPU) remains @@ -34,15 +34,15 @@ The extras add optional dependencies on top of the compiled package, so they use the same `--no-build-isolation` flag as the source build above. ```bash -pip install --no-build-isolation -e ".[cuda]" +RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e ".[cuda]" ``` ```bash -pip install --no-build-isolation -e ".[rocm]" +RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e ".[rocm]" ``` ```bash -pip install --no-build-isolation -e ".[vllm]" +python -m pip install --no-build-isolation -e ".[vllm]" ``` Install the vLLM extra only on rollout or benchmark environments that need the diff --git a/envs.py b/envs.py index 34aded7c..03c65ef2 100644 --- a/envs.py +++ b/envs.py @@ -26,3 +26,4 @@ def env_flag(name: str, default: bool = False) -> bool: KERNEL_ALIGN_NCU_LINEINFO = "KERNEL_ALIGN_NCU_LINEINFO" KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC = "KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC" KERNEL_ALIGN_FORCE_SM90 = "KERNEL_ALIGN_FORCE_SM90" +RL_KERNEL_REQUIRE_EXT = "RL_KERNEL_REQUIRE_EXT" diff --git a/rl_engine/__init__.py b/rl_engine/__init__.py index e69de29b..306d6813 100644 --- a/rl_engine/__init__.py +++ b/rl_engine/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import torch # noqa: F401 # Load torch shared libraries before importing rl_engine._C. diff --git a/setup.py b/setup.py index 94777289..6a35bcb3 100644 --- a/setup.py +++ b/setup.py @@ -3,6 +3,7 @@ import importlib.util import os +import warnings from pathlib import Path from setuptools import find_packages, setup @@ -24,16 +25,29 @@ def _load_envs_module(): def _load_torch_extension_tools(): try: import torch - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - except ImportError: + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise return None, None, None + from torch.utils.cpp_extension import BuildExtension, CUDAExtension + # CUDAExtension is also the supported extension entry point for ROCm # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when # torch.version.hip is set. return torch, BuildExtension, CUDAExtension +def _native_extension_required() -> bool: + """Whether the caller explicitly requested a native extension build.""" + return ( + envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) + or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or envs.env_flag("FORCE_CUDA") + ) + + def _cuda_define_from_env(name: str, macro: str) -> list[str]: value = os.environ.get(name) if value is None: @@ -81,6 +95,19 @@ def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: def get_extensions(): torch, _, CUDAExtension = _load_torch_extension_tools() if torch is None: + message = ( + "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " + "CUDA/ROCm PyTorch build first, then run " + "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." + ) + if _native_extension_required(): + raise RuntimeError(message) + warnings.warn( + f"{message} Continuing with the pure-Python fallback because no native extension " + "was explicitly requested.", + RuntimeWarning, + stacklevel=2, + ) return [] extensions = [] @@ -233,6 +260,14 @@ def get_extensions(): extra_link_args=extra_link_args, ) ) + + if _native_extension_required() and not extensions: + raise RuntimeError( + "rl_engine._C was requested but no CUDA/ROCm build environment is available. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " + "PYTORCH_ROCM_ARCH to the target architecture." + ) + return extensions From 9b41f740716885afb250d6c8bb0c611b64675988 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 19 Aug 2026 15:13:30 +0800 Subject: [PATCH 12/33] feat(attention): integrate cross-config wrapper into PR230 --- .github/workflows/ci.yml | 11 + .../ws2-attention-cross-config-integration.md | 185 ++ docs/design/ws2-cp-attention-contract.md | 227 +++ ...config_qwen3_8b_megatron_tp2_cp2_vllm.json | 89 + pyproject.toml | 1 + .../cross_config/adapters/__init__.py | 38 + .../cross_config/adapters/_common.py | 281 +++ .../alignment/cross_config/adapters/knobs.py | 161 ++ .../cross_config/adapters/megatron.py | 481 ++++++ .../alignment/cross_config/adapters/vllm.py | 540 ++++++ .../cross_config/attention_binding.py | 1407 +++++++++++++++ .../alignment/cross_config/determinism.py | 305 ++++ rl_engine/kernels/attention_contract.py | 1532 +++++++++++++++++ rl_engine/kernels/attention_preprocess.py | 275 +++ rl_engine/kernels/attention_projection.py | 254 +++ .../kernels/ops/pytorch/attention/__init__.py | 13 +- .../kernels/ops/pytorch/attention/ablation.py | 650 +++++++ .../ops/pytorch/attention/cp_attention.py | 1182 +++++++++++++ rl_engine/kernels/registry.py | 224 +++ tests/test_attention_ablation.py | 262 +++ tests/test_attention_contract.py | 683 ++++++++ tests/test_attention_cross_config_binding.py | 1295 ++++++++++++++ tests/test_attention_preprocess.py | 116 ++ tests/test_attention_projection.py | 97 ++ tests/test_cp_attention.py | 712 ++++++++ tests/test_cp_attention_transformer_engine.py | 48 + 26 files changed, 11068 insertions(+), 1 deletion(-) create mode 100644 docs/design/ws2-attention-cross-config-integration.md create mode 100644 docs/design/ws2-cp-attention-contract.md create mode 100644 examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json create mode 100644 rl_engine/alignment/cross_config/adapters/__init__.py create mode 100644 rl_engine/alignment/cross_config/adapters/_common.py create mode 100644 rl_engine/alignment/cross_config/adapters/knobs.py create mode 100644 rl_engine/alignment/cross_config/adapters/megatron.py create mode 100644 rl_engine/alignment/cross_config/adapters/vllm.py create mode 100644 rl_engine/alignment/cross_config/attention_binding.py create mode 100644 rl_engine/alignment/cross_config/determinism.py create mode 100644 rl_engine/kernels/attention_contract.py create mode 100644 rl_engine/kernels/attention_preprocess.py create mode 100644 rl_engine/kernels/attention_projection.py create mode 100644 rl_engine/kernels/ops/pytorch/attention/ablation.py create mode 100644 rl_engine/kernels/ops/pytorch/attention/cp_attention.py create mode 100644 tests/test_attention_ablation.py create mode 100644 tests/test_attention_contract.py create mode 100644 tests/test_attention_cross_config_binding.py create mode 100644 tests/test_attention_preprocess.py create mode 100644 tests/test_attention_projection.py create mode 100644 tests/test_cp_attention.py create mode 100644 tests/test_cp_attention_transformer_engine.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e9e9c91d..dc4c773e 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -80,6 +80,17 @@ jobs: run: | python -m pytest tests/test_attention.py -v -k "not large and not gpu" + - name: Run WS2 Attention Cross-Configuration Tests (CPU-safe) + run: | + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest -q \ + tests/test_attention_ablation.py \ + tests/test_attention_contract.py \ + tests/test_attention_cross_config_binding.py \ + tests/test_attention_preprocess.py \ + tests/test_attention_projection.py \ + tests/test_cp_attention.py \ + tests/test_cp_attention_transformer_engine.py + - name: Run KV-Cache Attention Ground-Truth Tests (CPU-safe) run: | python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu" diff --git a/docs/design/ws2-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md new file mode 100644 index 00000000..6ddd45a0 --- /dev/null +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -0,0 +1,185 @@ +# WS2 Attention Cross-Configuration Integration + +Implements PR4 of [#235](https://github.com/RL-Align/RL-Kernel/issues/235): wiring the +CP attention path into the cross-configuration planner/runtime for the Qwen3-8B +TP=2 CP=2 BF16 target. + +Builds on [#236](https://github.com/RL-Align/RL-Kernel/pull/236) (attention contract +and dispatch metadata), [#238](https://github.com/RL-Align/RL-Kernel/pull/238) +(deterministic CP reference) and [#230](https://github.com/RL-Align/RL-Kernel/pull/230) +(cross-configuration framework). + +## What "bind to the same contract" means here + +The PR4 acceptance criteria say rollout and training descriptors must "bind to the +same semantic attention contract". Under the frozen deployment the two sides can +never produce identical `AttentionContract` instances: + +| | training (Megatron) | rollout (vLLM) | +| --- | --- | --- | +| mode | full-sequence prefill | chunked prefill, later decode | +| CP | `context_parallel_size`, whole forward | `prefill_context_parallel_size`, prefill only | +| KV | no paging | paged KV with a block table | +| backend vocabulary | `AttnBackend{flash,fused,unfused,local,auto}` | `AttentionBackendEnum` | + +Read literally, the criterion is unsatisfiable. It is therefore implemented as three +tiers, in `rl_engine/alignment/cross_config/attention_binding.py`: + +| tier | fields | rule | failure | +| --- | --- | --- | --- | +| `IDENTICAL` | checkpoint/model/token identity plus complete TP/CP GQA head and sequence ownership | equal bit for bit | `comparable=False`; no drift number from the pair means anything | +| `SEMANTIC` | reduction contract, dtype, exported LSE, first-class Split-KV request, complete actual Split-KV plan sets, CUDA QK-Norm/RoPE identity, cross-side determinism | both sides equal **and** equal to the WS2 mandate | fail closed | +| `RECORDED` | `mode`, `backend_id`, `reduction.engine`, RoPE materialization state and fusion boundary, KV-cache paging | free to differ | recorded into provenance and measured | + +Two placements are load-bearing: + +* **`reduction.engine` is `RECORDED`, not `SEMANTIC`.** Training may run the in-op + deterministic reference while rollout runs a Transformer Engine merge oracle. + Forcing them equal would defeat the oracle comparison that #235 PR2/PR3/PR5/PR6 + depend on. +* **TP/CP topology is `IDENTICAL`, not `RECORDED`.** TP selects local Qwen3 GQA head + ownership and CP selects local sequence ownership. Different topology is a + different local attention problem, not a backend detail. +* **`reduction.order`, `reduction.acc_dtype`, and actual Split-KV schedules are + `SEMANTIC`.** Both runtimes must export the complete batch x TP x CP x KV-owner + plan set, including logical boundaries, merge order, FP32 accumulation, final + downcast, and fallback state. Configured policy alone never passes strict binding. + +`comparable` and `passed` are separate flags. A pair with mismatched identity is not +comparable. A pair that is comparable but violates the reduction mandate is still +rejected -- the drift would be real but attributable to the wrong thing. + +## H100 Attention input and projection boundary + +Megatron/TE and vLLM/FlashInfer are the first-choice implementations. +`H100AttentionPreprocessor` runs a same-input H100 bitwise probe against the +deterministic RL-Kernel path. An unavailable native callable, a native exception, +or a failed probe switches both sides to `RMSNormCudaOp` + `RoPESM90Op` and +records the fallback reason and probe ID. The launcher passes the returned +backend evidence into `AttentionRuntimeReadback`: + +```python +from rl_engine.kernels.attention_preprocess import H100AttentionPreprocessor + +prepared = H100AttentionPreprocessor(device)( + q, k, q_norm_weight, k_norm_weight, position_ids +) +readback = AttentionRuntimeReadback( + # contract, knobs, Split-KV plan set, source, and scope fields omitted here + **prepared.readback_fields(), +) +``` + +Strict binding rejects a missing or unknown backend and rejects mixed native / +fallback execution. If both sides fall back, they must report the same deterministic +backend IDs and policy ID. Printing a configured backend without executing the +probe is not evidence. + +The Attention boundary includes QKV projection, Q/K RMSNorm, RoPE, core +attention, KV-cache access, CP `(Out, LSE)` communication/merge, and o_proj. +`AttentionProjectionOp` freezes QKV/o_proj to BF16 input and output, FP32 +accumulation, ascending-K reduction, and Split-K disabled. Native projection +callables are accepted only after a bitwise probe against `DetGemmOp`; otherwise +both sides use the deterministic fallback. Its collective contract records QKV +column-parallel plus backward TP all-reduce, o_proj row-parallel partial output, +and the SP all-gather/reduce-scatter directions. The model input RMSNorm and +residual add remain outside this Attention experiment. + +## Determinism is not one thing + +`rl_engine/alignment/cross_config/determinism.py` probes both sides and compares +them, because the two frameworks mean different things by "deterministic": + +| | Megatron `deterministic_mode` | vLLM `VLLM_BATCH_INVARIANT` | +| --- | --- | --- | +| `NCCL_ALGO` | asserts membership in a five-value set | hard-sets `allreduce:tree` | +| `NCCL_PROTO`, channels, threads | not managed | hard-set (`Simple`, `1`, `1`) | +| TF32 | **not managed at all** | disabled (`fp32_precision="ieee"`) | +| BF16 reduced-precision reduction | not managed | disabled | +| cuBLAS workspace / BLAS library | not managed | `:4096:8`, cuBLASLt | +| GEMM | cuBLAS / TE | Triton `matmul_persistent` | +| FlashAttention | forbidden | permitted | + +`NCCL_ALGO`, `NCCL_PROTO` and `CUBLAS_WORKSPACE_CONFIG` change arithmetic, so a +mismatch there is blocking. The remaining differences -- including the TF32 and +BF16-reduction asymmetry, which under a pure BF16 GEMM path does not fire -- are +recorded so the asymmetry appears in every artifact rather than being invisible. + +## Runtime adapters + +Before this PR the only `RuntimeMaterializer` in the repository was +`CpuSmokeMaterializer` over a synthetic CPU model, and every named scenario +(`S1`/`S2`/`S3`) was planning-only. This PR adds the first two framework-shaped +adapters: + +* `adapters/megatron.py` -- `MegatronProvenanceAdapter` (construction and + distributed-context fingerprints, determinism probe, frozen-scope assertions) and + `MegatronAttentionMaterializer`. +* `adapters/vllm.py` -- `VllmProvenanceAdapter` (including diagnostic vLLM split + limits) and `VllmRolloutMaterializer`. +* `AttentionRuntimeReadback` -- the explicit handoff from an executed engine. It + carries the reconstructed actual contract, actual knob values, frozen-scope + verification, executed CUDA QK-Norm/RoPE identities and fallback state, and the + complete Split-KV runtime plan set. + +Constructing a contract is not runtime verification. Without a readback, adapter +applications are `UNOBSERVABLE`; only matching values reconstructed from a real +Megatron or vLLM execution are `APPLIED`. `bind_attention_runtime_readbacks` is the +strict public entry point used after both framework launchers collect that evidence. + +Neither module imports `megatron` or `vllm`; configs are duck-typed, so the binding +rules are exercised on CPU in CI rather than only on a 2-node cluster. + +## Fail closed, never substitute + +`unsupported_reduction_reason` rejects requests that #236 cannot express, instead of +collapsing them onto the supported value: + +| request | status | why | +| --- | --- | --- | +| `attention.reduction_order=arrival` | `UNSUPPORTED` | the control group must stay distinguishable from the treatment | +| `attention.reduction_downcast_at=per_block` | `UNSUPPORTED` | `DowncastPoint` declares only `final_write` | +| `attention.reduction_engine=te_oracle` | `UNSUPPORTED` | the TE merge oracle lands in #235 PR2/PR3; PR4's TE plan is provenance only | +| `attention.reduction_acc_dtype=bf16` | `UNSUPPORTED` | the CP `(out, lse)` merge accumulates in FP32 | +| configured contract without runtime readback | `UNOBSERVABLE` | requested values do not prove what executed | +| `rollout.context_parallel_size>1` with effective decode CP=1 | `ERROR`/`FALLBACK` | strict TP=2/CP=2 acceptance rejects the topology change | +| missing/mismatched/fallback Split-KV plan set | binding failure | Split-KV provenance must cover every batch/TP/CP/owner coordinate | +| missing/unknown QK-Norm or RoPE backend, or mixed native/fallback sides | binding failure | both sides must execute the same verified native policy or the common RL-Kernel CUDA fallback | + +## Knobs + +`adapters/knobs.py` extends `V1_KNOBS` additively. Added: training-side +`tensor_parallel_size` / `context_parallel_size` / `deterministic_mode` / +`cp_comm_type`, `rollout.batch_invariant` / `rollout.kv_block_size`, and the +reduction axis (`acc_dtype`, `order`, `downcast_at`, `engine`) plus +`attention.fusion_boundary` and `attention.split_kv_policy`. + +`training.attention_backend` keeps its path but its value domain is replaced with +Megatron's `AttnBackend`; the HuggingFace names have no Megatron counterpart, so this +is a replacement rather than a mapping. + +Not done here, because both change `V1_KNOBS` itself and would break existing +cross-config tests: removing `training.sharding` (Megatron has no such concept, and +DP=1 makes it moot) and renaming `rollout.context_parallel_size` to reflect that it +binds to `prefill_context_parallel_size`. + +## Scenario + +`examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json` supersedes +`cross_config_s1_distributed_smoke.json` and +`cross_config_s3_qwen3_8b_tp4_cp4_bf16.json`, whose training sides used `sdpa` / +`flash_attention_2` and `sharding: fsdp` -- none of which exist under Megatron -- and +whose TP=4/CP=4 topology does not match the target. +`cross_config_s2_vllm_tp_vs_fsdp.json` has no Megatron-only counterpart and should be +retired rather than rewritten. + +## Out of scope + +Deliberately not in this PR: + +* launching `torchrun`, initializing process groups, or executing core attention; +* pre-attention model RMSNorm and residual add; +* decode-mode materialization, which needs the validated `KVCacheSpec` from #235 PR6 + and is refused with that reference rather than stubbed; +* distributed drift benchmarks and report artifacts (#235 PR5); +* fused production backend alignment (#235 PR7) and backward (#235 PR8). diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md new file mode 100644 index 00000000..95320101 --- /dev/null +++ b/docs/design/ws2-cp-attention-contract.md @@ -0,0 +1,227 @@ +# WS2 CP-Aware Attention Contract + +Status: PR1 contract and dispatch metadata + +Tracking and shared contracts: + +- [#235: CP-aware deterministic Attention](https://github.com/RL-Align/RL-Kernel/issues/235) +- [#83: WS2 roadmap](https://github.com/RL-Align/RL-Kernel/issues/83) +- [#108: WS1 numerical contract](https://github.com/RL-Align/RL-Kernel/issues/108) +- [#111: WS2 cross-config alignment](https://github.com/RL-Align/RL-Kernel/issues/111) +- [#207: cross-config logprob drift contract](https://github.com/RL-Align/RL-Kernel/issues/207) + +## Scope + +This contract describes the logical inputs and deterministic reduction semantics for standard +softmax Attention under tensor parallelism (TP) and context parallelism (CP). It lets runtime +dispatch reject a backend whose numerical semantics do not match the requested layout. + +This PR1 layer does not shard tensors, launch a collective, merge CP partial states, or implement +a fused kernel. The deterministic CP reference implementation and its distributed numerical tests +belong to later work in #235. + +## Contract Objects + +`rl_engine.kernels.attention_contract` defines: + +- `AttentionContract`: role, mode, dtype, causal metadata, sharding, reduction, and optional cache + identity; +- `ShardingSpec`: TP-local head ownership and CP block-to-token ownership; +- `ReductionSpec`: fixed `(out, lse)` merge semantics; +- `KVCacheSpec`: decode replay cache identity; +- `RoPESpec`: Qwen3 RoPE state, position identity, and fused/unfused boundary metadata; +- `AttentionBackendCapability`: the layouts and semantics a backend explicitly supports. + +Construction performs validation immediately. A structurally valid contract means that the +request is complete and internally consistent; it does not mean that an installed backend can +materialize it. + +`AttentionContract.batch_size` is the logical sequence count. For packed varlen input it must +equal `len(packed_sequence_offsets) - 1`; it is not the physical leading dimension of a flattened +token tensor. + +For full `prefill`, `query_sequence_length` equals the local sequence length described by +`ShardingSpec`. Chunked prefill and decode may use shorter query lengths than their available KV +context. + +## Qwen3-8B TP=2 CP=2 Example + +```python +from rl_engine.kernels.attention_contract import ( + AttentionContract, + ReductionSpec, + RoPESpec, + ShardingSpec, +) + +sharding = ShardingSpec( + tp_rank=0, + tp_world_size=2, + cp_rank=0, + cp_world_size=2, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=0, + local_q_heads=16, + local_kv_head_start=0, + local_kv_heads=4, + global_sequence_length=4096, + local_sequence_length=2048, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 2048), +) + +contract = AttentionContract( + role="infer", + mode="prefill", + dtype="bf16", + batch_size=1, + query_sequence_length=2048, + head_dim=128, + causal=True, + causal_offsets=(0,), + sharding=sharding, + reduction=ReductionSpec(), + rope=RoPESpec( + q_state="post_rope", + k_state="post_rope", + k_cache_state="post_rope", + theta=1.0e6, + rotary_dim=128, + query_position_offsets=(0,), + key_position_offsets=(0,), + cast_at="after_rope", + output_dtype="bf16", + fusion_boundary="unfused_rope_attention", + ), +) +``` + +The TP fields preserve the global Qwen3 GQA mapping: each rank owns 16 of 32 query heads and 4 of +8 KV heads. The CP fields map local tensor slices to stable logical global block ids. A rank that +owns non-contiguous blocks uses one global token start per block and one extra local boundary: + +```python +global_block_indices=(0, 3) +global_block_token_starts=(0, 3072) +local_block_offsets=(0, 1024, 2048) +``` + +This metadata is sufficient for a later implementation to restore logical global order without +using ring arrival order. + +## RoPE / Position Semantics + +RoPE is part of the attention contract because rollout can materialize +`RoPE+Attention` as a fused or cache-aware path while training may materialize +`RoPE -> Attention` as separate operators. PR1 does not execute the RoPE kernel, +but it records the metadata required to prove both materializations use the same +model semantics. + +`RoPESpec` records: + +- whether Q, K, and cached K are `pre_rope` or `post_rope`; +- `theta`, optional `rope_scaling`, and `rotary_dim`; +- dense `position_ids` or per-sequence `query_position_offsets` / + `key_position_offsets`; +- the RoPE cast point and output dtype; +- `fusion_boundary`, either `unfused_rope_attention` or `fused_rope_attention`. + +When RoPE metadata is present, construction validates that rotary dimensions fit +the attention head dimension and that offset metadata matches the logical batch +shape. Backends must declare RoPE support through `AttentionBackendCapability`; +a backend that cannot consume RoPE/position metadata or cannot support a fused +RoPE+Attention boundary is rejected before dispatch. + +## Reduction Semantics + +The only PR1 reduction contract is: + +```text +partial state: (out, attention-domain lse) +merge: online_softmax_lse +acc_dtype: fp32 +order: global_block_index +downcast_at: final_write +engine: in_op_reference +``` + +CP output is not a plain sum. A backend that cannot export attention-domain LSE or cannot merge +partial states in fixed logical order is incompatible with this contract. + +The acceptable output and selected-logprob drift thresholds remain owned by #108. This contract +does not introduce another tolerance table. When connected to the rollout/training chain, the +selected-token metric remains the #207 convention: + +```text +dlogp = training-side recomputed logp - rollout-side old logp +``` + +## Mode-Specific Metadata + +All causal calls provide `causal_offsets`. Packed varlen calls provide one causal offset per +packed sequence and validated `packed_sequence_offsets`. + +Decode additionally requires `KVCacheSpec` with: + +- one cache position and KV sequence length per logical sequence; +- a block/page table; +- the physical page size; +- global token positions for every logical cached token; +- a prefix-cache key and explicit shared-prefix page count when prefix caching is enabled. + +Within each logical sequence, global token positions must be strictly increasing. Block-table +padding must be trailing, the active page count must match `ceil(kv_seq_len / page_size)`, and a +sequence cannot repeat one physical page id. Different sequences may share physical pages for an +equivalent prefix only when those pages are declared by `shared_prefix_page_count`, use the same +leading page ids and logical positions, and are fully populated. Declared shared prefix pages are +read-only; all suffix pages are exclusive to one sequence, providing the contract boundary needed +for copy-on-write before divergent decode. When prefix caching is disabled, no active page may be +shared across sequences. Missing or inconsistent decode cache identity is an error at contract +construction time. + +Each `cache_positions` entry is the terminal logical position already present in that sequence's +KV cache, so it must equal the final corresponding `global_token_positions` entry. It is not the +next position to be written. + +## Contract-Aware Dispatch + +Legacy callers continue to use `KernelRegistry.get_op()`. WS2 callers use: + +```python +result = kernel_registry.get_attention_op(contract) +op = result.op +provenance = result.provenance +``` + +Dispatch considers only backends with an `AttentionBackendCapability`. It checks role, attention +mode, dtype, TP/CP degree, LSE export, deterministic CP merge, packed varlen, and KV-cache support. +When RoPE metadata is present, dispatch also checks whether the backend explicitly supports +RoPE/position metadata and fused RoPE+Attention boundaries. +An undeclared or incompatible backend is skipped with an explicit rejection reason. + +The current WS1 PyTorch Attention implementations support local reference math but do not export +attention-domain LSE or materialize deterministic CP merge. Strict WS2 requests therefore fail +clearly today. A later deterministic backend becomes selectable by registering a capability that +truthfully declares those features; no grid-planner branch or silent fallback is required. + +Successful dispatch provenance records: + +- requested and actual backend ids; +- platform and fallback status; +- prior candidate rejection reasons; +- the complete requested contract; +- the selected backend capability descriptor. + +## Validation + +Contract and dispatch behavior are covered by: + +```bash +python -m pytest tests/test_attention_contract.py -q +``` + +The tests include Qwen3 TP=2/CP=2 construction, GQA ownership errors, non-contiguous CP blocks, +packed varlen metadata, decode cache identity, undeclared backend rejection, no incompatible +fallback, RoPE metadata validation, and JSON-compatible provenance. diff --git a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json new file mode 100644 index 00000000..11f9cef1 --- /dev/null +++ b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json @@ -0,0 +1,89 @@ +{ + "experiment_id": "ws2-qwen3-8b-attention-tp2-cp2", + "scenario_id": "qwen3_8b_megatron_tp2_cp2_vllm", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "scenario": { + "issue": "https://github.com/RL-Align/RL-Kernel/issues/235", + "pull_request": "PR4 -- cross-config integration", + "model": "Qwen3-8B dense", + "training_framework": "megatron", + "rollout_framework": "vllm", + "topology": "2 nodes x 2 GPUs, TP=2 CP=2 PP=1 DP=1, BF16, SM90", + "notes": [ + "Supersedes cross_config_s1_distributed_smoke.json and", + "cross_config_s3_qwen3_8b_tp4_cp4_bf16.json, whose training side used", + "HuggingFace attention backends and FSDP sharding. Neither exists in", + "Megatron, and DP=1 makes the sharding knob meaningless.", + "cross_config_s2_vllm_tp_vs_fsdp.json has no Megatron-only counterpart at", + "all and should be retired rather than rewritten.", + "rollout.context_parallel_size binds to vLLM", + "ParallelConfig.prefill_context_parallel_size and therefore applies to", + "prefill only; strict PR4 acceptance covers CP=2 prefill/chunked prefill.", + "A decode request that becomes CP=1 is a blocking fallback and is tested", + "separately by the PR6 logical KV replay harness." + ] + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "dtype": "bfloat16", + "enable_prefix_caching": false, + "enforce_eager": true, + "batch_invariant": true, + "kv_block_size": 16 + }, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "attention_backend": "unfused", + "compute_dtype": "bfloat16", + "deterministic_mode": true, + "cp_comm_type": "p2p", + "sharding": "unsharded" + }, + "attention": { + "reduction_acc_dtype": "fp32", + "reduction_order": "global_block_index", + "reduction_downcast_at": "final_write", + "reduction_engine": "in_op_reference", + "fusion_boundary": "unfused_rope_attention", + "split_kv_policy": 32 + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "training.context_parallel_size", + "values": [1, 2] + }, + { + "path": "training.tensor_parallel_size", + "values": [1, 2] + }, + { + "path": "attention.fusion_boundary", + "values": ["unfused_rope_attention", "fused_rope_attention"] + }, + { + "path": "training.cp_comm_type", + "values": ["p2p", "all_gather"] + }, + { + "path": "attention.reduction_order", + "values": ["global_block_index", "arrival"] + }, + { + "path": "attention.reduction_acc_dtype", + "values": ["fp32", "bf16"] + } + ] +} diff --git a/pyproject.toml b/pyproject.toml index 2507e591..7ffe391c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,4 +47,5 @@ follow_imports = "silent" [tool.pytest.ini_options] markers = [ "smoke_operator: temporary smoke-only operator plumbing tests", + "unit: CPU-safe unit tests", ] diff --git a/rl_engine/alignment/cross_config/adapters/__init__.py b/rl_engine/alignment/cross_config/adapters/__init__.py new file mode 100644 index 00000000..f02b9b38 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime adapters for the WS2 Qwen3-8B Megatron + vLLM cross-config target.""" + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + AttentionRuntimeReadback, + Qwen3ModelSpec, +) +from rl_engine.alignment.cross_config.adapters.knobs import ( + MEGATRON_ATTENTION_BACKENDS, + WS2_ATTENTION_KNOB_DESCRIPTORS, + WS2_ATTENTION_KNOBS, + WS2_ATTENTION_NORMALIZERS, +) +from rl_engine.alignment.cross_config.adapters.megatron import ( + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, +) +from rl_engine.alignment.cross_config.adapters.vllm import ( + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "AttentionRuntimeReadback", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", + "QWEN3_8B", + "Qwen3ModelSpec", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] diff --git a/rl_engine/alignment/cross_config/adapters/_common.py b/rl_engine/alignment/cross_config/adapters/_common.py new file mode 100644 index 00000000..9a3c5384 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/_common.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared pieces for the Megatron and vLLM WS2 attention adapters.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from rl_engine.alignment.cross_config.attention_binding import AttentionRuntimeReadback +from rl_engine.alignment.cross_config.runtime import KnobApplication +from rl_engine.alignment.cross_config.schema import ( + IsolationScope, + KnobDescriptor, + MaterializationStatus, +) +from rl_engine.kernels.attention_contract import ( + AttentionDType, + AttentionMerge, + DowncastPoint, + ReductionEngine, + ReductionOrder, + ReductionSpec, + ShardingSpec, + SplitKVSpec, +) + +__all__ = [ + "QWEN3_8B", + "Qwen3ModelSpec", + "AttentionRuntimeReadback", + "application", + "attention_dtype", + "build_reduction_spec", + "build_sharding_spec", + "causal_offsets_for", + "flatten", + "split_kv_spec", + "unsupported_reduction_reason", +] + + +@dataclass(frozen=True) +class Qwen3ModelSpec: + """Architecture constants for the frozen dense target. + + These are *not* knobs. #235/#239/#241 all fix Qwen3-8B dense, so they belong to + the scenario, and both sides must agree on them or the comparison is void. + """ + + name: str = "qwen3-8b" + hidden_size: int = 4096 + ffn_hidden_size: int = 12288 + num_layers: int = 36 + q_heads: int = 32 + kv_heads: int = 8 + head_dim: int = 128 + real_vocab_size: int = 151936 + rope_theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + qk_layernorm: bool = True + + def identity_fields(self) -> dict[str, Any]: + """The subset of :data:`IDENTITY_FIELDS` this spec is responsible for.""" + + return { + "q_heads": self.q_heads, + "kv_heads": self.kv_heads, + "head_dim": self.head_dim, + "rope_theta": self.rope_theta, + "rope_scaling": self.rope_scaling, + "rotary_dim": self.rotary_dim, + "qk_layernorm": self.qk_layernorm, + } + + +QWEN3_8B = Qwen3ModelSpec() + + +#: The planner normalizes dtype knobs to torch spellings (``bfloat16``), while +#: :class:`AttentionDType` uses short spellings (``bf16``). Passing a normalized knob +#: straight into the enum raises, so every adapter must translate here rather than +#: each inventing its own mapping. +_DTYPE_ALIASES: Mapping[str, AttentionDType] = { + "bf16": AttentionDType.BF16, + "bfloat16": AttentionDType.BF16, + "fp16": AttentionDType.FP16, + "float16": AttentionDType.FP16, + "half": AttentionDType.FP16, + "fp32": AttentionDType.FP32, + "float32": AttentionDType.FP32, + "float": AttentionDType.FP32, +} + + +def attention_dtype(value: Any, *, field: str) -> AttentionDType: + """Translate a normalized knob dtype into an :class:`AttentionDType`.""" + + if isinstance(value, AttentionDType): + return value + key = str(value).strip().lower().replace("torch.", "") + try: + return _DTYPE_ALIASES[key] + except KeyError as exc: + raise ValueError( + f"{field}={value!r} is not a supported attention dtype; " + f"expected one of {sorted(set(_DTYPE_ALIASES))}" + ) from exc + + +def split_kv_spec(flat: Mapping[str, Any]) -> SplitKVSpec: + """Build the first-class logical Split-KV request. + + The integer is a fixed logical KV chunk size in tokens. It is intentionally + not vLLM's ``flash_attn_max_num_splits_for_cuda_graph``: that setting is only + an upper bound and cannot prove which runtime boundaries executed. + """ + + split_size = flat.get("attention.split_kv_policy") + if split_size is None: + return SplitKVSpec.disabled() + return SplitKVSpec.fixed(int(split_size)) + + +def flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + """Flatten nested knob mappings into dotted paths.""" + + flat: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}{key}" + if isinstance(child, Mapping): + flat.update(flatten(child, f"{path}.")) + else: + flat[path] = child + return flat + + +def application( + descriptor: KnobDescriptor, + requested: Any, + materialized: Any, + actual: Any, + status: MaterializationStatus, + reason: str, + **evidence: Any, +) -> KnobApplication: + return KnobApplication( + path=descriptor.path, + requested=requested, + materialized=materialized, + actual=actual, + lifecycle=descriptor.lifecycle, + status=status, + evidence={"reason": reason, **evidence}, + critical=descriptor.critical, + ) + + +def unsupported_reduction_reason(flat: Mapping[str, Any]) -> str | None: + """Return why the requested reduction cannot be materialized, if it cannot. + + #236 declares single-member enums for merge order, downcast point and reduction + engine, so the alternative knob values exist only as control groups. Requesting + one must fail loudly rather than quietly collapse onto the supported value -- + silently substituting ``global_block_index`` for a requested ``arrival`` would + make the control group indistinguishable from the treatment. + """ + + order = flat.get("attention.reduction_order") + if order is not None and order != ReductionOrder.GLOBAL_BLOCK_INDEX.value: + return ( + f"attention.reduction_order={order!r} has no backend; #236 ReductionOrder " + f"declares only {ReductionOrder.GLOBAL_BLOCK_INDEX.value!r}" + ) + downcast = flat.get("attention.reduction_downcast_at") + if downcast is not None and downcast != DowncastPoint.FINAL_WRITE.value: + return ( + f"attention.reduction_downcast_at={downcast!r} has no backend; #236 " + f"DowncastPoint declares only {DowncastPoint.FINAL_WRITE.value!r}" + ) + engine = flat.get("attention.reduction_engine") + if engine is not None and engine != ReductionEngine.IN_OP_REFERENCE.value: + return ( + f"attention.reduction_engine={engine!r} has no backend; the Transformer " + "Engine merge oracle lands in #235 PR2/PR3, not here" + ) + acc_dtype = flat.get("attention.reduction_acc_dtype") + if ( + acc_dtype is not None + and attention_dtype(acc_dtype, field="attention.reduction_acc_dtype") + is not AttentionDType.FP32 + ): + return ( + f"attention.reduction_acc_dtype={acc_dtype!r} violates the WS2 mandate; " + "the CP (out, lse) merge accumulates in fp32" + ) + return None + + +def build_reduction_spec(flat: Mapping[str, Any]) -> ReductionSpec: + """Build the reduction spec, having already rejected unsupported requests.""" + + return ReductionSpec( + merge=AttentionMerge.ONLINE_SOFTMAX_LSE, + acc_dtype=AttentionDType.FP32, + order=ReductionOrder.GLOBAL_BLOCK_INDEX, + downcast_at=DowncastPoint.FINAL_WRITE, + engine=ReductionEngine.IN_OP_REFERENCE, + ) + + +def build_sharding_spec( + *, + model: Qwen3ModelSpec, + tp_rank: int, + tp_world_size: int, + cp_rank: int, + cp_world_size: int, + global_sequence_length: int, +) -> ShardingSpec: + """Build a CP/TP sharding spec for one rank of the frozen layout. + + TP splits heads, CP splits the sequence. The #239 rank layout fixes + ``rank = cp_rank * tp_world_size + tp_rank`` for a 2-node x 2-GPU deployment, + but nothing here depends on that mapping: ownership is derived from the ranks + themselves so the same builder serves CP=1 baselines. + """ + + if model.q_heads % tp_world_size or model.kv_heads % tp_world_size: + raise ValueError( + f"Qwen3 GQA heads ({model.q_heads}/{model.kv_heads}) must divide evenly " + f"across tp_world_size={tp_world_size}" + ) + if global_sequence_length % cp_world_size: + raise ValueError( + f"global_sequence_length={global_sequence_length} must divide evenly " + f"across cp_world_size={cp_world_size}" + ) + + local_q_heads = model.q_heads // tp_world_size + local_kv_heads = model.kv_heads // tp_world_size + local_sequence_length = global_sequence_length // cp_world_size + + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=model.q_heads, + global_kv_heads=model.kv_heads, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=global_sequence_length, + local_sequence_length=local_sequence_length, + # One contiguous CP block per rank. The merge order key is the global block + # index, never the arrival order of the CP exchange. + global_block_indices=(cp_rank,), + global_block_token_starts=(cp_rank * local_sequence_length,), + local_block_offsets=(0, local_sequence_length), + ) + + +def causal_offsets_for(sharding: ShardingSpec, batch_size: int) -> tuple[int, ...]: + """Causal offsets for one CP shard, one entry per batch entry. + + Under CP the local query block does not start at global position zero, so the + causal mask has to be shifted by the number of preceding global tokens. Taking + that from ``global_block_token_starts`` rather than recomputing + ``cp_rank * local_sequence_length`` keeps uneven CP splits correct. + """ + + offset = sharding.global_block_token_starts[0] + return (offset,) * batch_size + + +_PROCESS_SCOPES = (IsolationScope.PROCESS, IsolationScope.DISTRIBUTED_CONTEXT) diff --git a/rl_engine/alignment/cross_config/adapters/knobs.py b/rl_engine/alignment/cross_config/adapters/knobs.py new file mode 100644 index 00000000..33c29820 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/knobs.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 attention knobs for the Qwen3-8B TP=2 CP=2 Megatron + vLLM target. + +``V1_KNOBS`` was written against a HuggingFace/FSDP rollout-vs-training pair. Three +of its entries do not survive contact with the frozen Megatron + vLLM target: + +* ``training.sharding`` takes ``unsharded``/``fsdp``, neither of which exists in + Megatron, and is meaningless at DP=1 anyway; +* ``training.attention_backend`` takes HuggingFace names + (``flash_attention_2``/``sdpa``/``eager``/``model_default``) while Megatron's + ``AttnBackend`` is ``flash``/``fused``/``unfused``/``local``/``auto``; +* there is no training-side ``tensor_parallel_size`` or ``context_parallel_size`` + at all, so the target configuration cannot even be expressed. + +This module is deliberately **additive**: it extends ``V1_KNOBS`` rather than +editing it, and overrides only the normalizer for ``training.attention_backend``. +Deleting the two dead knobs changes ``V1_KNOBS`` itself and would break existing +cross-config tests, so it is left to a follow-up on the framework PR. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +from rl_engine.alignment.cross_config.planner import ( + _NORMALIZERS, + V1_KNOBS, + Normalizer, + _normalize_choice, + _positive_int, + _strict_bool, +) +from rl_engine.alignment.cross_config.schema import IsolationScope, KnobDescriptor + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] + + +#: ``megatron.core.transformer.enums.AttnBackend``. +MEGATRON_ATTENTION_BACKENDS: tuple[str, ...] = ( + "flash", + "fused", + "unfused", + "local", + "auto", +) + + +WS2_ATTENTION_KNOB_DESCRIPTORS: tuple[KnobDescriptor, ...] = ( + # -- training-side parallelism: the target configuration itself ------------ + KnobDescriptor( + "training.tensor_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "training.context_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + # -- determinism switches, one per framework ------------------------------ + KnobDescriptor( + "training.deterministic_mode", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "rollout.batch_invariant", + IsolationScope.PROCESS, + ("rollout",), + ), + # -- reduction knobs: the "turn the noise sources on and off" axis --------- + # These are what make drift attributable. ``reduction.order=arrival`` in + # particular is a control group, not a supported production value. + KnobDescriptor( + "attention.reduction_acc_dtype", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("fp32", "bf16"), + ), + KnobDescriptor( + "attention.reduction_order", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("global_block_index", "arrival"), + ), + KnobDescriptor( + "attention.reduction_downcast_at", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("final_write", "per_block"), + ), + KnobDescriptor( + "attention.reduction_engine", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("in_op_reference", "te_oracle"), + ), + # -- materialization knobs: differences the experiment measures ------------ + KnobDescriptor( + "attention.fusion_boundary", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("unfused_rope_attention", "fused_rope_attention"), + ), + KnobDescriptor( + # Shared logical KV chunk size. Runtime adapters must separately report + # the actual per-owner boundaries; a configured value is not evidence. + "attention.split_kv_policy", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + ), + KnobDescriptor( + # vLLM: CacheConfig.block_size -> AttentionContract.kv_cache.page_size + "rollout.kv_block_size", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + # The CP communication group cannot be reconfigured once built; it is bound to + # the distributed context, not merely to engine construction. + KnobDescriptor( + "training.cp_comm_type", + IsolationScope.DISTRIBUTED_CONTEXT, + ("training",), + allowed_values=("p2p", "all_gather", "a2a", "a2a+p2p"), + ), +) + + +WS2_ATTENTION_KNOBS: Mapping[str, KnobDescriptor] = { + **V1_KNOBS, + **{descriptor.path: descriptor for descriptor in WS2_ATTENTION_KNOB_DESCRIPTORS}, +} + + +WS2_ATTENTION_NORMALIZERS: Mapping[str, Normalizer] = { + **_NORMALIZERS, + # Replace, not map: the HuggingFace names have no Megatron counterpart. + "training.attention_backend": _normalize_choice(*MEGATRON_ATTENTION_BACKENDS), + "training.tensor_parallel_size": _positive_int, + "training.context_parallel_size": _positive_int, + "training.deterministic_mode": _strict_bool, + "rollout.batch_invariant": _strict_bool, + # AttentionDType values, not torch dtype names -- these feed ReductionSpec directly. + "attention.reduction_acc_dtype": _normalize_choice("fp32", "bf16"), + "attention.reduction_order": _normalize_choice("global_block_index", "arrival"), + "attention.reduction_downcast_at": _normalize_choice("final_write", "per_block"), + "attention.reduction_engine": _normalize_choice("in_op_reference", "te_oracle"), + "attention.fusion_boundary": _normalize_choice( + "unfused_rope_attention", "fused_rope_attention" + ), + "attention.split_kv_policy": _positive_int, + "rollout.kv_block_size": _positive_int, + "training.cp_comm_type": _normalize_choice("p2p", "all_gather", "a2a", "a2a+p2p"), +} diff --git a/rl_engine/alignment/cross_config/adapters/megatron.py b/rl_engine/alignment/cross_config/adapters/megatron.py new file mode 100644 index 00000000..55903cc8 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Training-side (Megatron) runtime adapter for WS2 attention cross-config. + +Two things live here: + +``MegatronProvenanceAdapter`` + Read-only. Turns a Megatron config object into the construction and + distributed-context fingerprints the cross-config framework already expects, + plus the determinism probe. It never imports ``megatron`` -- every accessor is + duck-typed -- so this module is importable and testable on a laptop. + +``MegatronAttentionMaterializer`` + Implements the ``RuntimeMaterializer`` protocol. Before this PR the only + implementation in the repository was ``CpuSmokeMaterializer`` over a synthetic + CPU model, so nothing had ever materialized a real distributed runtime. + +Scope boundary: materialization builds and validates the training-side +:class:`AttentionContract` and reports what would be constructed. Without an +``AttentionRuntimeReadback`` it reports ``UNOBSERVABLE``, never ``APPLIED``. It +does not launch ``torchrun``, initialize process groups, or execute attention; +the 2-node x 2-GPU launcher must inject readback collected after execution. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + AttentionRuntimeReadback, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + split_kv_spec, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import ( + DeterminismProbe, + megatron_probe_from_config, +) +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "MEGATRON_CONSTRUCTION_KEYS", + "MEGATRON_DISTRIBUTED_KEYS", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", +] + + +#: ``TransformerConfig`` fields that change attention arithmetic. Hashed into the +#: construction fingerprint. Deliberately excludes MoE, Mamba, MLA and sparse +#: attention fields: the frozen target is Qwen3-8B dense, and those are asserted +#: off rather than recorded. +MEGATRON_CONSTRUCTION_KEYS: tuple[str, ...] = ( + "attention_backend", + "attention_softmax_in_fp32", + "apply_query_key_layer_scaling", + "apply_rope_fusion", + "masked_softmax_fusion", + "bias_activation_fusion", + "bias_dropout_fusion", + "gradient_accumulation_fusion", + "cross_entropy_loss_fusion", + "cross_entropy_fusion_impl", + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + "recompute_modules", + "rotary_base", + "rotary_percent", + "rotary_interleaved", + "rotary_scaling_factor", + "qk_layernorm", + "hidden_dropout", + "attention_dropout", + "params_dtype", + "bf16", + "fp16", + "fp8", + "deterministic_mode", +) + + +#: ``ModelParallelConfig`` fields that define the distributed context. +MEGATRON_DISTRIBUTED_KEYS: tuple[str, ...] = ( + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "virtual_pipeline_model_parallel_size", + "context_parallel_size", + "hierarchical_context_parallel_sizes", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "sequence_parallel", + "cp_comm_type", + "tp_comm_overlap", + "use_te_rng_tracker", +) + + +#: Fields that must hold these values for the frozen dense target. A mismatch is a +#: hard stop, not a recorded difference -- see the exclusion list in the WS2 scope. +MEGATRON_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "pipeline_model_parallel_size": 1, + "expert_model_parallel_size": 1, + "sequence_parallel": False, + "fp8": None, + "hidden_dropout": 0.0, + "attention_dropout": 0.0, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class MegatronProvenanceAdapter: + """Extract fingerprints and determinism evidence from a Megatron config. + + ``config`` may be a real ``TransformerConfig``/``ModelParallelConfig``, a merged + namespace, or a test double. Missing attributes read as ``None`` and are + recorded as such rather than raising: an absent field is itself provenance. + """ + + framework = "megatron" + + def __init__(self, config: Any, *, env: Optional[Mapping[str, str]] = None): + self.config = config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_CONSTRUCTION_KEYS} + + def distributed_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_DISTRIBUTED_KEYS} + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return megatron_probe_from_config(self.config, env=self.env) + + def frozen_scope_violations(self) -> tuple[str, ...]: + """Return the frozen-scope assertions this config violates.""" + + violations: list[str] = [] + for name, expected in MEGATRON_FROZEN_ASSERTIONS.items(): + actual = _value(self.config, name) + if actual is None: + # Not declared. Treated as unknown rather than as satisfied, because + # a silently-absent MoE or FP8 setting is exactly the case that would + # otherwise slip past a dense-only claim. + violations.append(f"{name} is not declared (expected {expected!r})") + elif actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "determinism": self.determinism_probe().to_dict(), + } + + +class MegatronAttentionMaterializer: + """Materialize the training-side attention runtime for the WS2 target.""" + + runtime_kind = "megatron_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + backend_id: str = "rlkernel.cp_attention_reference", + provenance: Optional[MegatronProvenanceAdapter] = None, + runtime_readback: Optional[AttentionRuntimeReadback] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.backend_id = backend_id + self.provenance = provenance + self.runtime_readback = runtime_readback + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + """Build the training-side contract. Raises on an unusable request.""" + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=attention_dtype( + flat.get("training.compute_dtype", "bf16"), field="training.compute_dtype" + ), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + split_kv=split_kv_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "training" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + if contract is None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + f"attention contract is unavailable: {blocked}", + ) + ) + continue + applications.append( + self._runtime_application( + descriptor, + requested, + contract=contract, + scope_violations=scope_violations, + ) + ) + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "megatron", + "attention_backend": flat.get("training.attention_backend"), + "compute_dtype": flat.get("training.compute_dtype"), + "deterministic_mode": flat.get("training.deterministic_mode"), + "cp_comm_type": flat.get("training.cp_comm_type"), + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "runtime_readback": ( + None if self.runtime_readback is None else self.runtime_readback.to_dict() + ), + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"training": side_config, "rollout": {}}, + topology={ + "training": { + "world_size": tp_world_size * cp_world_size, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": cp_world_size, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "rollout": {"world_size": 1}, + }, + scorer={ + "mode": "teacher_forcing", + "framework": "megatron", + "export_lse": True, + }, + operator_backends={ + "training": self.backend_id, + "rollout": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) + + def _runtime_application( + self, + descriptor: KnobDescriptor, + requested: Any, + *, + contract: AttentionContract, + scope_violations: tuple[str, ...], + ) -> KnobApplication: + readback = self.runtime_readback + if readback is None: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + ( + "configured in the training contract, but no Megatron runtime " + "readback was supplied" + ), + frozen_scope_violations=list(scope_violations), + ) + if scope_violations or not readback.frozen_scope_verified: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.UNOBSERVABLE, + "Megatron frozen-scope assertions were not all verified by runtime readback", + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) + if readback.split_kv_fallback: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "Megatron runtime reported a Split-KV fallback", + runtime_readback_source=readback.source, + ) + if readback.contract != contract: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "Megatron runtime contract differs from the requested contract", + runtime_readback_source=readback.source, + ) + if descriptor.path not in readback.actual_knobs: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "Megatron runtime readback does not expose this knob", + runtime_readback_source=readback.source, + ) + actual = readback.actual_knobs[descriptor.path] + status = ( + MaterializationStatus.APPLIED if actual == requested else MaterializationStatus.FALLBACK + ) + reason = ( + "verified from the executed Megatron runtime" + if status is MaterializationStatus.APPLIED + else "Megatron runtime value differs from the requested value" + ) + return application( + descriptor, + requested, + requested, + actual, + status, + reason, + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) diff --git a/rl_engine/alignment/cross_config/adapters/vllm.py b/rl_engine/alignment/cross_config/adapters/vllm.py new file mode 100644 index 00000000..f89ab232 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -0,0 +1,540 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rollout-side (vLLM) runtime adapter for WS2 attention cross-config. + +Mirrors :mod:`.megatron`, with three differences that come straight from what vLLM +actually is: + +* vLLM's context parallelism is ``prefill_context_parallel_size`` -- it applies to + prefill only, so a decode-mode contract must declare ``cp_world_size == 1`` + regardless of what the prefill knob says. +* ``CacheConfig.block_size`` is the paged-KV page size, and it feeds + ``KVCacheSpec.page_size`` directly rather than being invented here. +* Determinism comes from the ``VLLM_BATCH_INVARIANT`` environment variable rather + than from a config field, because vLLM applies it inside + ``init_batch_invariance()`` at worker startup. + +Like the Megatron adapter, nothing here imports ``vllm``; configs are duck-typed so +the module is importable anywhere. Configured-only values remain ``UNOBSERVABLE``; +``APPLIED`` requires an explicit post-execution ``AttentionRuntimeReadback``. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + AttentionRuntimeReadback, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + split_kv_spec, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import DeterminismProbe, vllm_probe_from_env +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "VLLM_ATTENTION_KEYS", + "VLLM_CACHE_KEYS", + "VLLM_FROZEN_ASSERTIONS", + "VLLM_MODEL_KEYS", + "VLLM_PARALLEL_KEYS", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", +] + + +VLLM_MODEL_KEYS: tuple[str, ...] = ( + "dtype", + "seed", + "quantization", + "enforce_eager", + "max_logprobs", + "disable_cascade_attn", + "max_model_len", +) + +VLLM_CACHE_KEYS: tuple[str, ...] = ( + "block_size", + "cache_dtype", + "enable_prefix_caching", + "prefix_caching_hash_algo", + "calculate_kv_scales", + "sliding_window", +) + +VLLM_ATTENTION_KEYS: tuple[str, ...] = ( + "backend", + "flash_attn_version", + "use_prefill_decode_attention", + "flash_attn_max_num_splits_for_cuda_graph", + "use_cudnn_prefill", + "disable_flashinfer_prefill", + "use_non_causal", +) + +VLLM_PARALLEL_KEYS: tuple[str, ...] = ( + "tensor_parallel_size", + "pipeline_parallel_size", + "prefill_context_parallel_size", + "data_parallel_size", +) + + +#: Frozen dense-target assertions on the rollout side. ``cache_dtype`` must stay +#: ``auto`` because an FP8 KV cache is a representation-drift problem tracked +#: separately, and ``disable_cascade_attn`` must stay ``True`` because cascade +#: attention changes the block-merge structure the contract pins down. +VLLM_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "quantization": None, + "cache_dtype": "auto", + "calculate_kv_scales": False, + "disable_cascade_attn": True, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "sliding_window": None, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class VllmProvenanceAdapter: + """Extract fingerprints and determinism evidence from vLLM configs.""" + + framework = "vllm" + + def __init__( + self, + *, + model_config: Any = None, + cache_config: Any = None, + attention_config: Any = None, + parallel_config: Any = None, + env: Optional[Mapping[str, str]] = None, + ): + self.model_config = model_config + self.cache_config = cache_config + self.attention_config = attention_config + self.parallel_config = parallel_config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + view: dict[str, Any] = {} + for prefix, config, keys in ( + ("model", self.model_config, VLLM_MODEL_KEYS), + ("cache", self.cache_config, VLLM_CACHE_KEYS), + ("attention", self.attention_config, VLLM_ATTENTION_KEYS), + ): + for name in keys: + view[f"{prefix}.{name}"] = _value(config, name) + return view + + def distributed_view(self) -> dict[str, Any]: + return { + f"parallel.{name}": _value(self.parallel_config, name) for name in VLLM_PARALLEL_KEYS + } + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return vllm_probe_from_env(self.env, model_config=self.model_config) + + def frozen_scope_violations(self) -> tuple[str, ...]: + sources = { + "quantization": self.model_config, + "disable_cascade_attn": self.model_config, + "cache_dtype": self.cache_config, + "calculate_kv_scales": self.cache_config, + "sliding_window": self.cache_config, + "pipeline_parallel_size": self.parallel_config, + "data_parallel_size": self.parallel_config, + } + violations: list[str] = [] + for name, expected in VLLM_FROZEN_ASSERTIONS.items(): + config = sources.get(name) + if config is None: + continue + actual = _value(config, name) + if actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + @property + def kv_page_size(self) -> Optional[int]: + """vLLM's paged-KV block size, which is the contract's ``page_size``.""" + + block_size = _value(self.cache_config, "block_size") + return int(block_size) if block_size is not None else None + + @property + def split_kv_policy(self) -> Optional[int]: + """Diagnostic vLLM maximum split count, not the logical chunk-size contract.""" + + splits = _value(self.attention_config, "flash_attn_max_num_splits_for_cuda_graph") + return int(splits) if splits is not None else None + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "kv_page_size": self.kv_page_size, + "flash_attn_max_num_splits_for_cuda_graph": self.split_kv_policy, + "determinism": self.determinism_probe().to_dict(), + } + + +class VllmRolloutMaterializer: + """Materialize the rollout-side attention runtime for the WS2 target.""" + + runtime_kind = "vllm_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + mode: AttentionMode = AttentionMode.CHUNKED_PREFILL, + backend_id: str = "vllm.flash_attn", + provenance: Optional[VllmProvenanceAdapter] = None, + runtime_readback: Optional[AttentionRuntimeReadback] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.mode = mode + self.backend_id = backend_id + self.provenance = provenance + self.runtime_readback = runtime_readback + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def effective_cp_world_size(self, flat: Mapping[str, Any]) -> int: + """CP applies to prefill only; decode always runs at CP=1.""" + + requested = int(flat.get("rollout.context_parallel_size", 1)) + if self.mode is AttentionMode.DECODE: + return 1 + return requested + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + if self.mode is AttentionMode.DECODE: + # Decode replay needs a validated KVCacheSpec (cache positions, page + # ownership, prefix-cache identity). That is #235 PR6's contract surface, + # and inventing a placeholder here would let an unvalidated decode case + # look bound. Fail instead. + raise AttentionContractError( + "decode-mode materialization requires KV-cache identity from #235 PR6; " + "this adapter covers prefill and chunked prefill" + ) + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + cp_world_size = self.effective_cp_world_size(flat) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank if cp_world_size > 1 else 0, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.FUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + # vLLM stores post-RoPE K in the cache; recorded, not asserted equal to + # the training side, because it is a materialization fact. + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.INFER, + mode=self.mode, + dtype=attention_dtype(flat.get("rollout.dtype", "bf16"), field="rollout.dtype"), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + split_kv=split_kv_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + requested_cp = int(flat.get("rollout.context_parallel_size", 1)) + effective_cp = self.effective_cp_world_size(flat) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "rollout" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + if contract is None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + f"attention contract is unavailable: {blocked}", + ) + ) + continue + if path == "rollout.context_parallel_size" and effective_cp != requested_cp: + applications.append( + application( + descriptor, + requested, + effective_cp, + effective_cp, + MaterializationStatus.FALLBACK, + ( + "vLLM context parallelism covers prefill only; a decode-mode " + f"contract runs at cp_world_size=1, not {requested_cp}" + ), + vllm_field="ParallelConfig.prefill_context_parallel_size", + ) + ) + continue + applications.append( + self._runtime_application( + descriptor, + requested, + contract=contract, + scope_violations=scope_violations, + ) + ) + + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "vllm", + "dtype": flat.get("rollout.dtype"), + "enforce_eager": flat.get("rollout.enforce_eager"), + "enable_prefix_caching": flat.get("rollout.enable_prefix_caching"), + "batch_invariant": flat.get("rollout.batch_invariant"), + "kv_block_size": flat.get("rollout.kv_block_size"), + "split_kv_policy": flat.get("attention.split_kv_policy"), + "attention_mode": self.mode.value, + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "runtime_readback": ( + None if self.runtime_readback is None else self.runtime_readback.to_dict() + ), + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"rollout": side_config, "training": {}}, + topology={ + "rollout": { + "world_size": tp_world_size * effective_cp, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": effective_cp, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "training": {"world_size": 1}, + }, + scorer={ + "mode": "rollout_logprob", + "framework": "vllm", + "export_lse": True, + }, + operator_backends={ + "rollout": self.backend_id, + "training": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) + + def _runtime_application( + self, + descriptor: KnobDescriptor, + requested: Any, + *, + contract: AttentionContract, + scope_violations: tuple[str, ...], + ) -> KnobApplication: + readback = self.runtime_readback + if readback is None: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "configured in the rollout contract, but no vLLM runtime readback was supplied", + frozen_scope_violations=list(scope_violations), + ) + if scope_violations or not readback.frozen_scope_verified: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.UNOBSERVABLE, + "vLLM frozen-scope assertions were not all verified by runtime readback", + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) + if readback.split_kv_fallback: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "vLLM runtime reported a Split-KV fallback", + runtime_readback_source=readback.source, + ) + if readback.contract != contract: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "vLLM runtime contract differs from the requested contract", + runtime_readback_source=readback.source, + ) + if descriptor.path not in readback.actual_knobs: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "vLLM runtime readback does not expose this knob", + runtime_readback_source=readback.source, + ) + actual = readback.actual_knobs[descriptor.path] + status = ( + MaterializationStatus.APPLIED if actual == requested else MaterializationStatus.FALLBACK + ) + reason = ( + "verified from the executed vLLM runtime" + if status is MaterializationStatus.APPLIED + else "vLLM runtime value differs from the requested value" + ) + return application( + descriptor, + requested, + requested, + actual, + status, + reason, + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py new file mode 100644 index 00000000..13f87000 --- /dev/null +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -0,0 +1,1407 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Three-tier binding between rollout-side and training-side attention contracts. + +Issue #235 PR4 requires that "rollout and training descriptors bind to the same +semantic attention contract". Under the frozen Megatron + vLLM deployment the two +sides can never produce *identical* :class:`AttentionContract` instances: training +runs full-sequence prefill over a CP-sharded sequence, while rollout runs vLLM +paged-KV chunked prefill and decode. Taking "same contract" literally would make +the target configuration permanently unbindable. + +This module therefore splits binding into three tiers: + +``IDENTICAL`` + Logical identity. Both sides must agree bit for bit, otherwise the pair is not + comparable at all and no drift number from it means anything. + +``SEMANTIC`` + The WS2 numerical claim: merge semantics, accumulation dtype, reduction order + and downcast point are decided by the contract, not by the implementation. + Both sides must carry the same values *and* those values must match the WS2 + mandate, otherwise the comparison fails closed. + +``RECORDED`` + Materialization facts that the two sides are expected to differ on -- attention + mode, RoPE fusion boundary, KV-cache paging, backend id, reduction engine. These + differences are exactly what the experiment measures, so they are recorded into + provenance rather than rejected. + +Deliberately *not* in ``SEMANTIC``: ``engine``. Training may run the in-op +deterministic reference while rollout runs a Transformer Engine merge oracle; forcing +those equal would defeat the purpose of the oracle comparison in #235 PR2/3/5/6. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from typing import Any, Optional + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMerge, + AttentionRole, + DowncastPoint, + ReductionOrder, + SplitKVRuntimePlanSet, + validate_split_kv_plan_set_alignment, +) +from rl_engine.kernels.attention_preprocess import ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS, + PREPROCESS_POLICY_ID, +) +from rl_engine.kernels.attention_projection import ( + O_PROJ_COLLECTIVE_CONTRACT, + PROJECTION_POLICY_ID, + QKV_COLLECTIVE_CONTRACT, +) + +__all__ = [ + "ATTENTION_LSE_DOMAIN", + "AttentionBindingError", + "AttentionBindingResult", + "AttentionRuntimeReadback", + "BindingErrorCode", + "BindingIssue", + "BindingTier", + "IDENTITY_FIELDS", + "NULLABLE_IDENTITY_FIELDS", + "RECORDED_FIELDS", + "SEMANTIC_CONTRACT_FIELDS", + "SEMANTIC_REDUCTION_FIELDS", + "TOPOLOGY_FIELDS", + "WS2_ATTENTION_REDUCTION_MANDATE", + "bind_attention_contracts", + "bind_attention_runtime_readbacks", + "first_blocking_issue", + "identity_fingerprint", + "summarize_binding", +] + + +class AttentionBindingError(ValueError): + """Raised when a caller supplies structurally unusable binding inputs.""" + + +class BindingTier(str, Enum): + """Which rule a field is governed by.""" + + IDENTICAL = "identical" + SEMANTIC = "semantic" + RECORDED = "recorded" + + +class BindingErrorCode(str, Enum): + """Stable, machine-readable reasons a binding is rejected. + + Callers branch on these; they are part of the artifact schema and must not be + renamed without a schema version bump. + """ + + IDENTITY_MISSING = "IDENTITY_MISSING" + IDENTITY_MISMATCH = "IDENTITY_MISMATCH" + REDUCTION_SEMANTIC_MISMATCH = "REDUCTION_SEMANTIC_MISMATCH" + REDUCTION_MANDATE_VIOLATION = "REDUCTION_MANDATE_VIOLATION" + LSE_NOT_EXPORTED = "LSE_NOT_EXPORTED" + ROLE_COLLISION = "ROLE_COLLISION" + DETERMINISM_INCOMPATIBLE = "DETERMINISM_INCOMPATIBLE" + TOPOLOGY_MISMATCH = "TOPOLOGY_MISMATCH" + SPLIT_KV_RUNTIME_MISSING = "SPLIT_KV_RUNTIME_MISSING" + SPLIT_KV_MISMATCH = "SPLIT_KV_MISMATCH" + SPLIT_KV_FALLBACK = "SPLIT_KV_FALLBACK" + ATTENTION_PREPROCESS_MISSING = "ATTENTION_PREPROCESS_MISSING" + ATTENTION_PREPROCESS_MISMATCH = "ATTENTION_PREPROCESS_MISMATCH" + ATTENTION_PREPROCESS_FALLBACK = "ATTENTION_PREPROCESS_FALLBACK" + ATTENTION_PROJECTION_MISSING = "ATTENTION_PROJECTION_MISSING" + ATTENTION_PROJECTION_MISMATCH = "ATTENTION_PROJECTION_MISMATCH" + ATTENTION_CORE_MISSING = "ATTENTION_CORE_MISSING" + ATTENTION_CORE_MISMATCH = "ATTENTION_CORE_MISMATCH" + ATTENTION_CORE_SCHEDULE = "ATTENTION_CORE_SCHEDULE" + ATTENTION_NATIVE_ARITHMETIC = "ATTENTION_NATIVE_ARITHMETIC" + ATTENTION_CORE_SPLIT_K = "ATTENTION_CORE_SPLIT_K" + ATTENTION_BACKEND_MISSING = "ATTENTION_BACKEND_MISSING" + ATTENTION_NOT_PRODUCTION_READY = "ATTENTION_NOT_PRODUCTION_READY" + + +@dataclass(frozen=True) +class AttentionRuntimeReadback: + """Actual attention contract and all-rank Split-KV evidence from one engine.""" + + contract: AttentionContract + actual_knobs: Mapping[str, Any] + split_kv_plan_set: SplitKVRuntimePlanSet + source: str + frozen_scope_verified: bool + preprocess_backends: Mapping[str, str] = field(default_factory=dict) + preprocess_fallback: bool = False + preprocess_fallback_reason: str | None = None + preprocess_probe_id: str = "" + preprocess_policy_id: str = PREPROCESS_POLICY_ID + projection_plans: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + strict_mode: bool = False + strict_core_id: str | None = None + strict_schedule: str | None = None + native_attention_arithmetic: bool = True + strict_split_kv_policy: str | None = None + actual_backend: str | None = None + communication_backend: str | None = None + production_ready: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.contract, AttentionContract): + raise TypeError("runtime readback contract must be an AttentionContract") + if not isinstance(self.actual_knobs, Mapping): + raise TypeError("runtime readback actual_knobs must be a mapping") + if not isinstance(self.split_kv_plan_set, SplitKVRuntimePlanSet): + raise TypeError("runtime readback requires a complete SplitKVRuntimePlanSet") + if not isinstance(self.source, str) or not self.source.strip(): + raise ValueError("runtime readback source must be a non-empty string") + if not isinstance(self.frozen_scope_verified, bool): + raise TypeError("frozen_scope_verified must be a bool") + if not isinstance(self.preprocess_backends, Mapping): + raise TypeError("runtime readback preprocess_backends must be a mapping") + for name, backend in self.preprocess_backends.items(): + if not isinstance(name, str) or not name.strip(): + raise ValueError("preprocess backend names must be non-empty strings") + if not isinstance(backend, str) or not backend.strip(): + raise ValueError("preprocess backend IDs must be non-empty strings") + if not isinstance(self.preprocess_fallback, bool): + raise TypeError("preprocess_fallback must be a bool") + if self.preprocess_fallback and not self.preprocess_fallback_reason: + raise ValueError( + "preprocess_fallback_reason is required when preprocess_fallback is true" + ) + if not isinstance(self.preprocess_probe_id, str): + raise TypeError("preprocess_probe_id must be a string") + if not isinstance(self.preprocess_policy_id, str) or not self.preprocess_policy_id.strip(): + raise ValueError("preprocess_policy_id must be a non-empty string") + if not isinstance(self.projection_plans, Mapping): + raise TypeError("projection_plans must be a mapping") + if not isinstance(self.strict_mode, bool): + raise TypeError("strict_mode must be a bool") + if self.strict_core_id is not None and ( + not isinstance(self.strict_core_id, str) or not self.strict_core_id.strip() + ): + raise ValueError("strict_core_id must be a non-empty string when provided") + if self.strict_schedule is not None and ( + not isinstance(self.strict_schedule, str) or not self.strict_schedule.strip() + ): + raise ValueError("strict_schedule must be a non-empty string when provided") + if not isinstance(self.native_attention_arithmetic, bool): + raise TypeError("native_attention_arithmetic must be a bool") + if self.strict_split_kv_policy is not None and self.strict_split_kv_policy not in { + "disabled", + "fixed", + "auto", + }: + raise ValueError("strict_split_kv_policy must be disabled, fixed, or auto") + for name, value in ( + ("actual_backend", self.actual_backend), + ("communication_backend", self.communication_backend), + ): + if value is not None and (not isinstance(value, str) or not value.strip()): + raise ValueError(f"{name} must be a non-empty string when provided") + if not isinstance(self.production_ready, bool): + raise TypeError("production_ready must be a bool") + normalized_projection_plans: dict[str, Mapping[str, Any]] = {} + for name, plan in self.projection_plans.items(): + if not isinstance(name, str) or not isinstance(plan, Mapping): + raise TypeError("projection_plans must map projection names to mappings") + normalized_projection_plans[name] = MappingProxyType(dict(plan)) + + plan_error = _split_kv_plan_contract_error(self.contract, self.split_kv_plan_set) + if plan_error is not None: + raise ValueError(plan_error) + object.__setattr__(self, "actual_knobs", MappingProxyType(dict(self.actual_knobs))) + object.__setattr__( + self, + "preprocess_backends", + MappingProxyType(dict(self.preprocess_backends)), + ) + object.__setattr__( + self, + "projection_plans", + MappingProxyType(normalized_projection_plans), + ) + + @property + def split_kv_fallback(self) -> bool: + return bool(_split_kv_fallbacks(self.split_kv_plan_set)) + + def to_dict(self) -> dict[str, Any]: + return { + "source": self.source, + "frozen_scope_verified": self.frozen_scope_verified, + "contract": self.contract.to_dict(), + "actual_knobs": dict(self.actual_knobs), + "attention_preprocess": { + "backends": dict(self.preprocess_backends), + "fallback": self.preprocess_fallback, + "fallback_reason": self.preprocess_fallback_reason, + "probe_id": self.preprocess_probe_id, + "policy_id": self.preprocess_policy_id, + }, + "attention_projections": { + name: dict(plan) for name, plan in self.projection_plans.items() + }, + "strict_attention": { + "enabled": self.strict_mode, + "core_id": self.strict_core_id, + "schedule": self.strict_schedule, + "native_attention_arithmetic": self.native_attention_arithmetic, + "split_kv_policy": self.strict_split_kv_policy, + }, + "runtime_backend": { + "actual_backend": self.actual_backend, + "communication_backend": self.communication_backend, + "production_ready": self.production_ready, + }, + "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), + } + + +#: Attention exports attention-domain LSE, never vocab-logprob LSE (#235). +#: Recorded explicitly so a future ``LogprobContract`` binding cannot be confused +#: with this one purely because both set ``export_lse=True``. +ATTENTION_LSE_DOMAIN = "attention" + + +#: Fields both sides must agree on bit for bit before any comparison is meaningful. +#: Sourced from #235 "Numerical Contract" preconditions plus the vime-owned rollout +#: provenance (weight version, sampling, padding) that the issue assumes but does +#: not enumerate. +IDENTITY_FIELDS: tuple[str, ...] = ( + "checkpoint_id", + "model_version", + "weight_version", + "tokenizer_fingerprint", + "token_ids_fingerprint", + "active_mask_fingerprint", + "position_ids_fingerprint", + "padding_side", + "pre_update_state", + # model semantics that decide what attention *means* + "q_heads", + "kv_heads", + "head_dim", + "rope_theta", + "rope_scaling", + "rotary_dim", + "qk_layernorm", + # batch composition: batch-invariance is a claim about results not changing with + # batch makeup, so two sides scoring different batches are not comparable at all + "batch_size", + # decode replay identity (#235 PR6) + "global_token_positions_fingerprint", + "kv_seq_lens_fingerprint", +) + + +#: Reduction fields that decide the numerical result. Both sides must carry the +#: same value, and that value must satisfy :data:`WS2_ATTENTION_REDUCTION_MANDATE`. +SEMANTIC_REDUCTION_FIELDS: tuple[str, ...] = ( + "merge", + "acc_dtype", + "order", + "downcast_at", +) + + +#: Contract fields outside ``ReductionSpec`` that still decide the numerical result. +#: ``dtype`` is here rather than in :data:`RECORDED_FIELDS` because comparing a BF16 +#: rollout against an FP16 training pass produces a real drift number attributable to +#: nothing. #235 PR5 does sweep BF16 against an FP32 reference; that sweep opts in via +#: ``allow_dtype_difference`` instead of loosening the default. +SEMANTIC_CONTRACT_FIELDS: tuple[str, ...] = ("dtype",) + + +#: Sharding fields that determine local GQA head and sequence ownership. These are +#: comparison preconditions, not harmless backend provenance: a TP/CP mismatch +#: means the two ranks did not evaluate the same local attention problem. +TOPOLOGY_FIELDS: tuple[str, ...] = ( + "tp_rank", + "tp_world_size", + "cp_rank", + "cp_world_size", + "global_q_heads", + "global_kv_heads", + "local_q_head_start", + "local_q_heads", + "local_kv_head_start", + "local_kv_heads", + "global_sequence_length", + "local_sequence_length", + "global_block_indices", + "global_block_token_starts", + "local_block_offsets", + "packed_sequence_offsets", +) + + +#: The WS2 mandate itself. ``#236`` currently declares single-member enums for +#: ``merge`` / ``order`` / ``downcast_at``, so those checks are tautological today; +#: they are written out anyway so that widening any of those enums later fails here +#: instead of silently admitting a non-conforming backend. +WS2_ATTENTION_REDUCTION_MANDATE: Mapping[str, str] = { + "merge": AttentionMerge.ONLINE_SOFTMAX_LSE.value, + "acc_dtype": AttentionDType.FP32.value, + "order": ReductionOrder.GLOBAL_BLOCK_INDEX.value, + "downcast_at": DowncastPoint.FINAL_WRITE.value, +} + + +#: Materialization facts the two sides are expected to differ on. Recorded into +#: provenance; never a rejection reason. +RECORDED_FIELDS: tuple[str, ...] = ( + "mode", + "backend_id", + "reduction.engine", + "rope.fusion_boundary", + "rope.q_state", + "rope.k_state", + "rope.k_cache_state", + "rope.cast_at", + "rope.output_dtype", + "preprocess.qk_rmsnorm", + "preprocess.rope", + "preprocess.fallback", + "kv_cache.page_size", + "kv_cache.prefix_cache_enabled", + "kv_cache.block_table_shape", +) + + +@dataclass(frozen=True) +class BindingIssue: + """One reason a binding is not comparable or not admissible.""" + + code: BindingErrorCode + tier: BindingTier + field: str + rollout: Any = None + training: Any = None + message: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "code": self.code.value, + "tier": self.tier.value, + "field": self.field, + "rollout": self.rollout, + "training": self.training, + "message": self.message, + } + + +@dataclass(frozen=True) +class AttentionBindingResult: + """Outcome of binding one rollout contract to one training contract. + + ``comparable`` and ``passed`` are deliberately separate. A pair whose identity + does not match is *not comparable* -- reporting a drift number for it would be + meaningless. A pair that is comparable but violates the reduction mandate *is* + comparable yet must still fail closed, because the whole WS2 claim is that + reduction order and accumulation precision come from the contract. + """ + + comparable: bool + passed: bool + issues: tuple[BindingIssue, ...] = () + identity_fingerprint: str = "" + reduction_fingerprint: str = "" + binding_fingerprint: str = "" + recorded_differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + provenance: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.attention_binding.v3" + + def issues_by_code(self, code: BindingErrorCode) -> tuple[BindingIssue, ...]: + return tuple(issue for issue in self.issues if issue.code is code) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "comparable": self.comparable, + "passed": self.passed, + "issues": [issue.to_dict() for issue in self.issues], + "identity_fingerprint": self.identity_fingerprint, + "reduction_fingerprint": self.reduction_fingerprint, + "binding_fingerprint": self.binding_fingerprint, + "recorded_differences": { + key: dict(value) for key, value in self.recorded_differences.items() + }, + "provenance": dict(self.provenance), + } + + +def _canonical_fingerprint(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def identity_fingerprint(identity: Mapping[str, Any]) -> str: + """Fingerprint only the declared :data:`IDENTITY_FIELDS`, in a fixed order. + + Extra keys in ``identity`` are ignored on purpose: callers pass whole + provenance bundles, and the fingerprint must not drift when an unrelated + diagnostic field is added. + """ + + return _canonical_fingerprint({name: identity.get(name) for name in IDENTITY_FIELDS}) + + +def _reduction_view(contract: AttentionContract) -> dict[str, Any]: + reduction = contract.reduction + return { + "merge": reduction.merge.value, + "acc_dtype": reduction.acc_dtype.value, + "order": reduction.order.value, + "downcast_at": reduction.downcast_at.value, + "engine": reduction.engine.value, + } + + +def _recorded_view( + contract: AttentionContract, + extra: Optional[Mapping[str, Any]] = None, +) -> dict[str, Any]: + rope = contract.rope + kv_cache = contract.kv_cache + view: dict[str, Any] = { + "mode": contract.mode.value, + "backend_id": None, + "reduction.engine": contract.reduction.engine.value, + } + if rope is not None: + view.update( + { + "rope.fusion_boundary": rope.fusion_boundary.value, + "rope.q_state": rope.q_state.value, + "rope.k_state": rope.k_state.value, + "rope.k_cache_state": rope.k_cache_state.value, + "rope.cast_at": rope.cast_at.value, + "rope.output_dtype": rope.output_dtype.value, + } + ) + if kv_cache is not None: + view.update( + { + "kv_cache.page_size": kv_cache.page_size, + "kv_cache.prefix_cache_enabled": kv_cache.prefix_cache_enabled, + "kv_cache.block_table_shape": [ + len(kv_cache.block_table), + max((len(row) for row in kv_cache.block_table), default=0), + ], + } + ) + if extra: + view.update(extra) + return view + + +def _topology_view(contract: AttentionContract) -> dict[str, Any]: + sharding = contract.sharding + return {name: getattr(sharding, name) for name in TOPOLOGY_FIELDS} + + +def _split_kv_fallbacks(plan_set: SplitKVRuntimePlanSet) -> list[dict[str, Any]]: + return [ + entry.to_dict() + for entry in plan_set.entries + if entry.execution.fallback + or entry.execution.actual_mode is None + or entry.execution.actual_mode is not entry.execution.requested_mode + or entry.execution.actual_split_size != entry.execution.requested_split_size + ] + + +def _split_kv_plan_contract_error( + contract: AttentionContract, + plan_set: SplitKVRuntimePlanSet, +) -> str | None: + sharding = contract.sharding + expected_topology = ( + contract.batch_size, + sharding.tp_world_size, + sharding.cp_world_size, + ) + actual_topology = ( + plan_set.batch_size, + plan_set.tp_world_size, + plan_set.cp_world_size, + ) + if actual_topology != expected_topology: + return ( + "Split-KV plan-set batch/TP/CP topology does not match the attention " + f"contract: actual={actual_topology}, expected={expected_topology}" + ) + if contract.mode.value in {"prefill", "chunked_prefill"}: + expected_totals = (sharding.global_sequence_length,) * contract.batch_size + if plan_set.total_kv_tokens != expected_totals: + return ( + "Split-KV plan-set KV lengths do not match the prefill attention " + f"contract: actual={plan_set.total_kv_tokens}, expected={expected_totals}" + ) + elif contract.kv_cache is not None: + expected_totals = contract.kv_cache.kv_seq_lens + if plan_set.total_kv_tokens != expected_totals: + return ( + "Split-KV plan-set KV lengths do not match decode KV-cache lengths: " + f"actual={plan_set.total_kv_tokens}, expected={expected_totals}" + ) + for entry in plan_set.entries: + execution = entry.execution + if ( + execution.requested_mode is not contract.split_kv.mode + or execution.requested_split_size != contract.split_kv.fixed_split_size + ): + return ( + "Split-KV runtime request does not match the first-class attention " + f"contract at {entry.coordinate}" + ) + return None + + +#: Identity fields where ``None`` is a real value rather than an omission. Qwen3-8B +#: applies no RoPE scaling, so ``rope_scaling=None`` must not read as "undeclared" -- +#: both sides still have to agree on it, which the equality pass below handles. +NULLABLE_IDENTITY_FIELDS: frozenset[str] = frozenset({"rope_scaling"}) + + +def _missing_identity_fields(identity: Mapping[str, Any]) -> tuple[str, ...]: + return tuple( + name + for name in IDENTITY_FIELDS + if name not in NULLABLE_IDENTITY_FIELDS and identity.get(name) is None + ) + + +def bind_attention_contracts( + *, + rollout_contract: AttentionContract, + training_contract: AttentionContract, + rollout_identity: Mapping[str, Any], + training_identity: Mapping[str, Any], + rollout_backend_id: str, + training_backend_id: str, + determinism_issues: Sequence[BindingIssue] = (), + require_full_identity: bool = True, + allow_dtype_difference: bool = False, + rollout_split_kv_plan_set: Optional[SplitKVRuntimePlanSet] = None, + training_split_kv_plan_set: Optional[SplitKVRuntimePlanSet] = None, + rollout_recorded_extra: Optional[Mapping[str, Any]] = None, + training_recorded_extra: Optional[Mapping[str, Any]] = None, +) -> AttentionBindingResult: + """Bind a rollout attention contract to a training attention contract. + + ``determinism_issues`` is threaded in from + :mod:`rl_engine.alignment.cross_config.determinism` rather than computed here, + so that this module stays free of framework probing and remains testable + without Megatron or vLLM present. + + ``require_full_identity`` exists for the single-GPU harness in #235 PR2, which + legitimately has no KV-cache or decode identity to declare. Distributed callers + must leave it at ``True``. + + ``allow_dtype_difference`` exists for the #235 PR5 sweep that deliberately scores + a BF16 path against an FP32 reference. It must stay ``False`` everywhere else. + + Strict binding requires complete actual Split-KV plan sets from both runtimes. + A configured policy is insufficient because auto-selection, graph capture, and + backend fallbacks can change the executed boundaries. The plan sets cover the + complete batch x TP x CP x KV-owner Cartesian product. + + ``rollout_recorded_extra`` / ``training_recorded_extra`` are diagnostic-only + backend facts. They can never make a semantic mismatch admissible. + """ + + if rollout_contract.role is not AttentionRole.INFER: + raise AttentionBindingError( + f"rollout_contract.role must be {AttentionRole.INFER.value!r}, " + f"got {rollout_contract.role.value!r}" + ) + if training_contract.role is not AttentionRole.TRAIN: + raise AttentionBindingError( + f"training_contract.role must be {AttentionRole.TRAIN.value!r}, " + f"got {training_contract.role.value!r}" + ) + + issues: list[BindingIssue] = [] + + # ---- tier 1: identity, bit for bit ------------------------------------- + if require_full_identity: + for side, identity in (("rollout", rollout_identity), ("training", training_identity)): + for name in _missing_identity_fields(identity): + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISSING, + tier=BindingTier.IDENTICAL, + field=f"{side}.{name}", + message=f"{side} identity does not declare {name!r}", + ) + ) + + for name in IDENTITY_FIELDS: + rollout_value = rollout_identity.get(name) + training_value = training_identity.get(name) + if rollout_value != training_value: + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISMATCH, + tier=BindingTier.IDENTICAL, + field=name, + rollout=rollout_value, + training=training_value, + message=( + f"{name!r} differs between sides; the pair is not comparable " + "and any drift computed from it is meaningless" + ), + ) + ) + + comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + + rollout_topology = _topology_view(rollout_contract) + training_topology = _topology_view(training_contract) + for name in TOPOLOGY_FIELDS: + if rollout_topology[name] != training_topology[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.TOPOLOGY_MISMATCH, + tier=BindingTier.IDENTICAL, + field=f"sharding.{name}", + rollout=rollout_topology[name], + training=training_topology[name], + message=( + f"sharding.{name} changes TP/CP ownership; the pair is not " + "the same local attention problem" + ), + ) + ) + + comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + + # ---- tier 2: reduction semantics, and the WS2 mandate ------------------- + rollout_reduction = _reduction_view(rollout_contract) + training_reduction = _reduction_view(training_contract) + + for name in SEMANTIC_REDUCTION_FIELDS: + if rollout_reduction[name] != training_reduction[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"reduction.{name!r} must be decided by the contract, not by the " + "backend; the two sides disagree" + ), + ) + ) + mandated = WS2_ATTENTION_REDUCTION_MANDATE[name] + for side, view in (("rollout", rollout_reduction), ("training", training_reduction)): + if view[name] != mandated: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_MANDATE_VIOLATION, + tier=BindingTier.SEMANTIC, + field=f"{side}.reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"WS2 requires reduction.{name} == {mandated!r}; " + f"{side} declares {view[name]!r}" + ), + ) + ) + + if not allow_dtype_difference and rollout_contract.dtype is not training_contract.dtype: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field="dtype", + rollout=rollout_contract.dtype.value, + training=training_contract.dtype.value, + message=( + "the two sides compute in different dtypes; the resulting drift is " + "not attributable. Pass allow_dtype_difference=True only for a " + "deliberate precision sweep" + ), + ) + ) + + if rollout_contract.split_kv != training_contract.split_kv: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field="split_kv", + rollout=rollout_contract.split_kv.to_dict(), + training=training_contract.split_kv.to_dict(), + message="training and rollout must request the same first-class Split-KV policy", + ) + ) + + for side, plan_set in ( + ("rollout", rollout_split_kv_plan_set), + ("training", training_split_kv_plan_set), + ): + if plan_set is None: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_RUNTIME_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + message=( + f"{side} did not report a complete actual Split-KV plan set; " + "configured policy alone is not runtime evidence" + ), + ) + ) + continue + contract = rollout_contract if side == "rollout" else training_contract + contract_error = _split_kv_plan_contract_error(contract, plan_set) + if contract_error is not None: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + rollout=plan_set.to_dict() if side == "rollout" else None, + training=plan_set.to_dict() if side == "training" else None, + message=contract_error, + ) + ) + fallbacks = _split_kv_fallbacks(plan_set) + if fallbacks: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_FALLBACK, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + rollout=fallbacks if side == "rollout" else None, + training=fallbacks if side == "training" else None, + message=f"{side} Split-KV runtime used an unknown or fallback plan", + ) + ) + + if rollout_split_kv_plan_set is not None and training_split_kv_plan_set is not None: + try: + validate_split_kv_plan_set_alignment( + training_split_kv_plan_set, + rollout_split_kv_plan_set, + ) + except AttentionContractError as exc: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field="split_kv_runtime_plan_set", + rollout=rollout_split_kv_plan_set.to_dict(), + training=training_split_kv_plan_set.to_dict(), + message=str(exc), + ) + ) + + for side, contract in (("rollout", rollout_contract), ("training", training_contract)): + if not contract.export_lse: + issues.append( + BindingIssue( + code=BindingErrorCode.LSE_NOT_EXPORTED, + tier=BindingTier.SEMANTIC, + field=f"{side}.export_lse", + message=( + "attention-domain LSE must be exported; without it the deterministic " + "CP merge cannot be validated" + ), + ) + ) + + if rollout_backend_id == training_backend_id and rollout_backend_id: + # Not an error, but worth surfacing: an identical backend on both sides means + # the experiment is not actually measuring a cross-implementation difference. + pass + + issues.extend(determinism_issues) + + # ---- tier 3: recorded differences -------------------------------------- + rollout_recorded = _recorded_view(rollout_contract, rollout_recorded_extra) + rollout_recorded["backend_id"] = rollout_backend_id + training_recorded = _recorded_view(training_contract, training_recorded_extra) + training_recorded["backend_id"] = training_backend_id + + recorded_differences: dict[str, dict[str, Any]] = {} + for name in RECORDED_FIELDS: + rollout_value = rollout_recorded.get(name) + training_value = training_recorded.get(name) + if rollout_value != training_value: + recorded_differences[name] = { + "rollout": rollout_value, + "training": training_value, + } + + identity_fp = identity_fingerprint(training_identity if comparable else rollout_identity) + reduction_fp = _canonical_fingerprint( + {name: training_reduction[name] for name in SEMANTIC_REDUCTION_FIELDS} + ) + passed = comparable and not any(issue.tier is BindingTier.SEMANTIC for issue in issues) + + provenance = { + "lse_domain": ATTENTION_LSE_DOMAIN, + "dtype": training_contract.dtype.value, + "split_kv_runtime": { + "rollout": ( + None if rollout_split_kv_plan_set is None else rollout_split_kv_plan_set.to_dict() + ), + "training": ( + None if training_split_kv_plan_set is None else training_split_kv_plan_set.to_dict() + ), + }, + "rollout": { + "contract": rollout_contract.to_dict(), + "backend_id": rollout_backend_id, + "recorded": rollout_recorded, + }, + "training": { + "contract": training_contract.to_dict(), + "backend_id": training_backend_id, + "recorded": training_recorded, + }, + } + + return AttentionBindingResult( + comparable=comparable, + passed=passed, + issues=tuple(issues), + identity_fingerprint=identity_fp, + reduction_fingerprint=reduction_fp, + binding_fingerprint=_canonical_fingerprint( + { + "identity": identity_fp, + "reduction": reduction_fp, + "topology": training_topology, + "split_kv": provenance["split_kv_runtime"], + "lse_domain": ATTENTION_LSE_DOMAIN, + "rollout_backend": rollout_backend_id, + "training_backend": training_backend_id, + "attention_preprocess": { + "rollout": { + name: rollout_recorded.get(f"preprocess.{name}") + for name in ( + *MANDATED_ATTENTION_PREPROCESS_BACKENDS, + "fallback", + ) + }, + "training": { + name: training_recorded.get(f"preprocess.{name}") + for name in ( + *MANDATED_ATTENTION_PREPROCESS_BACKENDS, + "fallback", + ) + }, + }, + } + ), + recorded_differences=recorded_differences, + provenance=provenance, + ) + + +def bind_attention_runtime_readbacks( + *, + rollout: AttentionRuntimeReadback, + training: AttentionRuntimeReadback, + rollout_identity: Mapping[str, Any], + training_identity: Mapping[str, Any], + rollout_backend_id: str, + training_backend_id: str, + determinism_issues: Sequence[BindingIssue] = (), +) -> AttentionBindingResult: + """Strict public handoff from executed framework runtimes to PR4 binding. + + The Megatron/vLLM launchers remain environment-owned. Once both launchers have + reconstructed their actual contracts and all-rank Split-KV reports, this entry + point performs the complete comparison without accepting configured-only data. + """ + + missing_scope_evidence = [] + for side, readback in (("rollout", rollout), ("training", training)): + if not readback.frozen_scope_verified: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"{side}.frozen_scope_verified", + message=f"{side} runtime did not verify the frozen attention scope", + ) + ) + missing_scope_evidence.extend(_attention_preprocess_issues(side, readback)) + missing_scope_evidence.extend(_attention_projection_issues(side, readback)) + missing_scope_evidence.extend(_strict_attention_core_issues(side, readback)) + missing_scope_evidence.extend(_strict_attention_core_pair_issues(rollout, training)) + if rollout.preprocess_fallback != training.preprocess_fallback: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field="preprocess.fallback", + rollout=rollout.preprocess_fallback, + training=training.preprocess_fallback, + message=( + "both runtimes must either pass the native H100 bitwise probe or " + "use the same deterministic preprocess fallback" + ), + ) + ) + if rollout.preprocess_policy_id != training.preprocess_policy_id: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field="preprocess.policy_id", + rollout=rollout.preprocess_policy_id, + training=training.preprocess_policy_id, + message="QK-Norm/RoPE policy IDs differ between runtimes", + ) + ) + if rollout.preprocess_fallback and training.preprocess_fallback: + for name in MANDATED_ATTENTION_PREPROCESS_BACKENDS: + rollout_backend = rollout.preprocess_backends.get(name) + training_backend = training.preprocess_backends.get(name) + if rollout_backend != training_backend: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"preprocess.{name}", + rollout=rollout_backend, + training=training_backend, + message="fallback sides must use the same deterministic preprocess backend", + ) + ) + for projection in ("qkv", "o_proj"): + rollout_plan = rollout.projection_plans.get(projection, {}) + training_plan = training.projection_plans.get(projection, {}) + rollout_fallback = bool(rollout_plan.get("fallback", False)) + training_fallback = bool(training_plan.get("fallback", False)) + if rollout_fallback != training_fallback: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"projection.{projection}.fallback", + rollout=rollout_fallback, + training=training_fallback, + message=( + "QKV/o_proj must use the same native-or-deterministic " "path on both sides" + ), + ) + ) + if rollout_fallback and training_fallback: + for field_name in ("backend_id", "policy_id", "split_k", "reduction_order"): + if rollout_plan.get(field_name) != training_plan.get(field_name): + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"projection.{projection}.{field_name}", + rollout=rollout_plan.get(field_name), + training=training_plan.get(field_name), + message=( + "deterministic projection fallback evidence differs " + "between sides" + ), + ) + ) + return bind_attention_contracts( + rollout_contract=rollout.contract, + training_contract=training.contract, + rollout_identity=rollout_identity, + training_identity=training_identity, + rollout_backend_id=rollout_backend_id, + training_backend_id=training_backend_id, + determinism_issues=tuple(determinism_issues) + tuple(missing_scope_evidence), + rollout_split_kv_plan_set=rollout.split_kv_plan_set, + training_split_kv_plan_set=training.split_kv_plan_set, + rollout_recorded_extra={ + **{ + f"preprocess.{name}": backend + for name, backend in rollout.preprocess_backends.items() + }, + "preprocess.fallback": rollout.preprocess_fallback, + "preprocess.fallback_reason": rollout.preprocess_fallback_reason, + "preprocess.probe_id": rollout.preprocess_probe_id, + "preprocess.policy_id": rollout.preprocess_policy_id, + "strict.enabled": rollout.strict_mode, + "strict.core_id": rollout.strict_core_id, + "strict.schedule": rollout.strict_schedule, + "strict.native_attention_arithmetic": rollout.native_attention_arithmetic, + "strict.split_kv_policy": rollout.strict_split_kv_policy, + "runtime.actual_backend": rollout.actual_backend, + "runtime.communication_backend": rollout.communication_backend, + "runtime.production_ready": rollout.production_ready, + **{ + f"projection.{projection}": dict(plan) + for projection, plan in rollout.projection_plans.items() + }, + }, + training_recorded_extra={ + **{ + f"preprocess.{name}": backend + for name, backend in training.preprocess_backends.items() + }, + "preprocess.fallback": training.preprocess_fallback, + "preprocess.fallback_reason": training.preprocess_fallback_reason, + "preprocess.probe_id": training.preprocess_probe_id, + "preprocess.policy_id": training.preprocess_policy_id, + "strict.enabled": training.strict_mode, + "strict.core_id": training.strict_core_id, + "strict.schedule": training.strict_schedule, + "strict.native_attention_arithmetic": training.native_attention_arithmetic, + "strict.split_kv_policy": training.strict_split_kv_policy, + "runtime.actual_backend": training.actual_backend, + "runtime.communication_backend": training.communication_backend, + "runtime.production_ready": training.production_ready, + **{ + f"projection.{projection}": dict(plan) + for projection, plan in training.projection_plans.items() + }, + }, + ) + + +def _strict_attention_core_issues( + side: str, + readback: AttentionRuntimeReadback, +) -> list[BindingIssue]: + if not readback.strict_mode: + return [] + issues = [] + if readback.actual_backend != "rlkernel.cuda.deterministic_attention": + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_BACKEND_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.runtime.actual_backend", + rollout=readback.actual_backend if side == "rollout" else None, + training=readback.actual_backend if side == "training" else None, + message=(f"{side} strict Attention did not execute the CUDA deterministic core"), + ) + ) + if readback.communication_backend != "self_owned_cuda_ag_rs": + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_BACKEND_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.runtime.communication_backend", + rollout=readback.communication_backend if side == "rollout" else None, + training=readback.communication_backend if side == "training" else None, + message=( + f"{side} strict CP Attention did not execute the self-owned CUDA AG/RS path" + ), + ) + ) + if not readback.production_ready: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_NOT_PRODUCTION_READY, + tier=BindingTier.SEMANTIC, + field=f"{side}.runtime.production_ready", + rollout=False if side == "rollout" else None, + training=False if side == "training" else None, + message=(f"{side} evidence is reference-only and cannot close the production gate"), + ) + ) + if readback.strict_core_id != STRICT_ATTENTION_CORE_ID: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.core_id", + message=( + f"{side} strict Attention did not execute the shared core " + f"{STRICT_ATTENTION_CORE_ID!r}" + ), + ) + ) + if readback.strict_schedule != STRICT_ATTENTION_SCHEDULE_ID: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_SCHEDULE, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.schedule", + message=( + f"{side} strict Attention did not execute the canonical schedule " + f"{STRICT_ATTENTION_SCHEDULE_ID!r}" + ), + ) + ) + if readback.native_attention_arithmetic: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_NATIVE_ARITHMETIC, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.native_attention_arithmetic", + message=f"{side} strict Attention entered native TE/FlashInfer arithmetic", + ) + ) + if ( + readback.strict_split_kv_policy != "disabled" + or readback.contract.split_kv.mode.value != "disabled" + ): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_SPLIT_K, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.split_kv_policy", + message=( + f"{side} strict Attention did not prove Split-KV disabled " + "in both runtime evidence and AttentionContract" + ), + ) + ) + return issues + + +def _strict_attention_core_pair_issues( + rollout: AttentionRuntimeReadback, + training: AttentionRuntimeReadback, +) -> list[BindingIssue]: + if rollout.strict_mode != training.strict_mode: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="strict.enabled", + rollout=rollout.strict_mode, + training=training.strict_mode, + message="training and rollout must use the same strict Attention mode", + ) + ] + if rollout.strict_mode and rollout.strict_core_id != training.strict_core_id: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="strict.core_id", + rollout=rollout.strict_core_id, + training=training.strict_core_id, + message="training and rollout executed different Attention cores", + ) + ] + if rollout.strict_mode and rollout.strict_schedule != training.strict_schedule: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_SCHEDULE, + tier=BindingTier.SEMANTIC, + field="strict.schedule", + rollout=rollout.strict_schedule, + training=training.strict_schedule, + message="training and rollout executed different strict Attention schedules", + ) + ] + if rollout.strict_mode and rollout.actual_backend != training.actual_backend: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="runtime.actual_backend", + rollout=rollout.actual_backend, + training=training.actual_backend, + message="training and rollout executed different Attention backends", + ) + ] + if rollout.strict_mode and rollout.communication_backend != training.communication_backend: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="runtime.communication_backend", + rollout=rollout.communication_backend, + training=training.communication_backend, + message="training and rollout executed different Attention communication backends", + ) + ] + return [] + + +def _attention_preprocess_issues( + side: str, + readback: AttentionRuntimeReadback, +) -> list[BindingIssue]: + issues: list[BindingIssue] = [] + for name, mandated in MANDATED_ATTENTION_PREPROCESS_BACKENDS.items(): + actual = readback.preprocess_backends.get(name) + if actual is None: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.preprocess.{name}", + message=( + f"{side} did not report the executed {name} backend; " + "runtime-native execution cannot validate the Attention input boundary" + ), + ) + ) + elif actual != mandated and not actual.startswith("native."): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.preprocess.{name}", + rollout=actual if side == "rollout" else None, + training=actual if side == "training" else None, + message=( + f"{side} executed {actual!r}; " + f"the H100 experiment requires {mandated!r} or a verified native backend" + ), + ) + ) + if readback.preprocess_fallback and any( + readback.preprocess_backends.get(name) != mandated + for name, mandated in MANDATED_ATTENTION_PREPROCESS_BACKENDS.items() + ): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_FALLBACK, + tier=BindingTier.SEMANTIC, + field=f"{side}.preprocess.fallback", + message=( + f"{side} reported a fallback but did not use the common deterministic " + "QK-Norm/RoPE backends" + ), + ) + ) + return issues + + +def _attention_projection_issues( + side: str, + readback: AttentionRuntimeReadback, +) -> list[BindingIssue]: + issues: list[BindingIssue] = [] + expected_collectives = { + "qkv": QKV_COLLECTIVE_CONTRACT.to_dict(), + "o_proj": O_PROJ_COLLECTIVE_CONTRACT.to_dict(), + } + fixed_fields = { + "input_dtype": "torch.bfloat16", + "weight_dtype": "torch.bfloat16", + "output_dtype": "torch.bfloat16", + "accumulation_dtype": "torch.float32", + "reduction_order": "k_ascending", + "split_k": False, + "policy_id": PROJECTION_POLICY_ID, + } + for projection, expected_collective in expected_collectives.items(): + plan = readback.projection_plans.get(projection) + if plan is None: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}", + message=f"{side} did not report the executed {projection} projection plan", + ) + ) + continue + for field_name, expected in fixed_fields.items(): + actual = plan.get(field_name) + if actual != expected: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.{field_name}", + rollout=actual if side == "rollout" else None, + training=actual if side == "training" else None, + message=f"{side} {projection} {field_name} must be {expected!r}", + ) + ) + collective = plan.get("collective") + if not isinstance(collective, Mapping) or dict(collective) != expected_collective: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.collective", + message=f"{side} {projection} TP/SP collective directions are invalid", + ) + ) + backend_id = plan.get("backend_id") + if not isinstance(backend_id, str) or not backend_id.strip(): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.backend_id", + message=f"{side} {projection} backend identity is missing", + ) + ) + if not isinstance(plan.get("probe_id"), str) or not plan.get("probe_id"): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.probe_id", + message=f"{side} {projection} bitwise probe identity is missing", + ) + ) + if plan.get("fallback"): + if backend_id != "rlkernel.cuda.det_gemm" or not plan.get("fallback_reason"): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.fallback", + message=( + f"{side} {projection} fallback must execute DetGemmOp and record why" + ), + ) + ) + return issues + + +def summarize_binding(result: AttentionBindingResult) -> str: + """One-line human summary for CLI output and failure messages.""" + + if result.passed: + return ( + f"attention binding OK " + f"(identity={result.identity_fingerprint[:12]}, " + f"{len(result.recorded_differences)} recorded difference(s))" + ) + if not result.comparable: + fields = ", ".join( + sorted({issue.field for issue in result.issues if issue.tier is BindingTier.IDENTICAL}) + ) + return f"attention binding NOT COMPARABLE; identity problems: {fields}" + fields = ", ".join( + sorted({issue.field for issue in result.issues if issue.tier is BindingTier.SEMANTIC}) + ) + return f"attention binding FAILED CLOSED; semantic problems: {fields}" + + +def first_blocking_issue( + result: AttentionBindingResult, +) -> Optional[BindingIssue]: + """Return the issue a caller should report, preferring identity over semantics.""" + + for tier in (BindingTier.IDENTICAL, BindingTier.SEMANTIC): + for issue in result.issues: + if issue.tier is tier: + return issue + return None diff --git a/rl_engine/alignment/cross_config/determinism.py b/rl_engine/alignment/cross_config/determinism.py new file mode 100644 index 00000000..be81654e --- /dev/null +++ b/rl_engine/alignment/cross_config/determinism.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-side determinism probing for the Megatron + vLLM cross-config target. + +Both frameworks ship a "make this deterministic" switch, but they mean different +things by it, and neither knows the other exists: + +``Megatron`` ``ModelParallelConfig.deterministic_mode`` + Asserts ``NCCL_ALGO`` is one of five values, forbids FlashAttention and fused + cross-entropy, calls ``torch.use_deterministic_algorithms(True)``, and requires + ``NVTE_ALLOW_NONDETERMINISTIC_ALGO == 0``. It does **not** touch TF32, BF16 + reduced-precision reduction, cuBLAS workspace, NCCL protocol, or NCCL channel + counts. + +``vLLM`` ``VLLM_BATCH_INVARIANT`` + Replaces ``aten::mm/addmm/matmul/linear/bmm``, ``log_softmax``/``softmax``, + ``mean.dim`` and ``rms_norm`` with Triton kernels, disables TF32 and BF16/FP16 + reduced-precision reduction, pins cuBLAS workspace and the BLAS library, and + hard-sets ten NCCL environment variables. + +So a run can have both switches on and still be comparing two different notions of +determinism. This module makes that difference explicit and, where it changes the +numerics, blocking. It never imports Megatron or vLLM: probes are built from plain +mappings so the logic is testable on any machine. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, Optional + +from rl_engine.alignment.cross_config.attention_binding import ( + BindingErrorCode, + BindingIssue, + BindingTier, +) + +__all__ = [ + "COMPARED_NCCL_KEYS", + "DeterminismProbe", + "DeterminismReport", + "compare_determinism", + "megatron_probe_from_config", + "vllm_probe_from_env", +] + + +#: Environment keys whose value can change a reduction result. Compared across +#: sides; a difference is reported, and a difference in the *arithmetic* subset is +#: blocking. Ordering is fixed so the fingerprint is stable. +COMPARED_NCCL_KEYS: tuple[str, ...] = ( + "NCCL_ALGO", + "NCCL_PROTO", + "NCCL_MIN_NCHANNELS", + "NCCL_MAX_NCHANNELS", + "NCCL_NTHREADS", + "NCCL_SOCKET_NTHREADS", + "NCCL_COLLNET_ENABLE", + "NCCL_NVLS_ENABLE", + "NCCL_P2P_NET_DISABLE", + "NCCL_LAUNCH_MODE", + "CUBLAS_WORKSPACE_CONFIG", +) + + +#: The subset above that changes arithmetic rather than only scheduling. A mismatch +#: here fails the binding closed; a mismatch in the remainder is recorded only. +_ARITHMETIC_NCCL_KEYS: frozenset[str] = frozenset( + {"NCCL_ALGO", "NCCL_PROTO", "CUBLAS_WORKSPACE_CONFIG"} +) + + +@dataclass(frozen=True) +class DeterminismProbe: + """What one side actually has switched on. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are tri-state on + purpose: ``None`` means "the framework does not manage this", which is exactly + Megatron's situation and is itself the finding. + """ + + side: str + framework: str + mode_flag: str + enabled: bool + env: Mapping[str, Any] = field(default_factory=dict) + tf32_disabled: Optional[bool] = None + bf16_reduced_precision_reduction: Optional[bool] = None + forbids_flash_attention: Optional[bool] = None + evidence: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_probe.v1" + + def __post_init__(self) -> None: + if self.side not in ("rollout", "training"): + raise ValueError("side must be 'rollout' or 'training'") + if not self.framework: + raise ValueError("framework must not be empty") + object.__setattr__(self, "env", dict(self.env)) + object.__setattr__(self, "evidence", dict(self.evidence)) + + @property + def env_fingerprint(self) -> str: + payload = {key: self.env.get(key) for key in COMPARED_NCCL_KEYS} + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "side": self.side, + "framework": self.framework, + "mode_flag": self.mode_flag, + "enabled": self.enabled, + "env": {key: self.env.get(key) for key in COMPARED_NCCL_KEYS}, + "env_fingerprint": self.env_fingerprint, + "tf32_disabled": self.tf32_disabled, + "bf16_reduced_precision_reduction": self.bf16_reduced_precision_reduction, + "forbids_flash_attention": self.forbids_flash_attention, + "evidence": dict(self.evidence), + } + + +@dataclass(frozen=True) +class DeterminismReport: + """Cross-side comparison result.""" + + rollout: DeterminismProbe + training: DeterminismProbe + issues: tuple[BindingIssue, ...] = () + differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_report.v1" + + @property + def compatible(self) -> bool: + return not self.issues + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "compatible": self.compatible, + "rollout": self.rollout.to_dict(), + "training": self.training.to_dict(), + "issues": [issue.to_dict() for issue in self.issues], + "differences": {key: dict(value) for key, value in self.differences.items()}, + } + + +def megatron_probe_from_config( + config: Any, + env: Optional[Mapping[str, str]] = None, +) -> DeterminismProbe: + """Build a training-side probe from a Megatron config object. + + ``config`` is duck-typed (anything exposing ``deterministic_mode`` and + optionally ``attention_backend`` / ``cross_entropy_loss_fusion``) so this works + against a real ``ModelParallelConfig``, a test double, or a plain namespace, + and so importing this module never requires Megatron. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are reported as + ``None`` because Megatron does not manage them -- a ``grep`` for ``allow_tf32`` + and ``fp32_precision`` across ``megatron/`` returns nothing. That asymmetry + against vLLM is the point of :func:`compare_determinism`. + """ + + environ = dict(env or {}) + enabled = bool(getattr(config, "deterministic_mode", False)) + return DeterminismProbe( + side="training", + framework="megatron", + mode_flag="deterministic_mode", + enabled=enabled, + env={key: environ.get(key) for key in COMPARED_NCCL_KEYS}, + tf32_disabled=None, + bf16_reduced_precision_reduction=None, + forbids_flash_attention=enabled, + evidence={ + "nvte_allow_nondeterministic_algo": environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO"), + "cross_entropy_loss_fusion": getattr(config, "cross_entropy_loss_fusion", None), + "attention_backend": _enum_value(getattr(config, "attention_backend", None)), + "tensor_model_parallel_size": getattr(config, "tensor_model_parallel_size", None), + "context_parallel_size": getattr(config, "context_parallel_size", None), + "sequence_parallel": getattr(config, "sequence_parallel", None), + "manages_tf32": False, + "manages_bf16_reduced_precision_reduction": False, + }, + ) + + +def vllm_probe_from_env( + env: Mapping[str, str], + *, + model_config: Any = None, +) -> DeterminismProbe: + """Build a rollout-side probe from the vLLM process environment. + + ``VLLM_BATCH_INVARIANT`` is read from ``env`` rather than ``vllm.envs`` so the + probe can be constructed from a remote worker's reported environment, which is + how vime's Ray actors expose it. + """ + + enabled = str(env.get("VLLM_BATCH_INVARIANT", "0")).strip() in ("1", "true", "True") + return DeterminismProbe( + side="rollout", + framework="vllm", + mode_flag="VLLM_BATCH_INVARIANT", + enabled=enabled, + env={key: env.get(key) for key in COMPARED_NCCL_KEYS}, + # vLLM sets both to "ieee"/disabled inside init_batch_invariance(). + tf32_disabled=enabled or None, + bf16_reduced_precision_reduction=(False if enabled else None), + forbids_flash_attention=False, + evidence={ + "vllm_allreduce_use_symm_mem": env.get("VLLM_ALLREDUCE_USE_SYMM_MEM"), + "vllm_use_aot_compile": env.get("VLLM_USE_AOT_COMPILE"), + "enforce_eager": getattr(model_config, "enforce_eager", None), + "disable_cascade_attn": getattr(model_config, "disable_cascade_attn", None), + "quantization": getattr(model_config, "quantization", None), + "manages_tf32": True, + "manages_bf16_reduced_precision_reduction": True, + }, + ) + + +def _enum_value(value: Any) -> Any: + return getattr(value, "value", value) + + +def compare_determinism( + *, + rollout: DeterminismProbe, + training: DeterminismProbe, +) -> DeterminismReport: + """Compare two probes and produce blocking issues plus recorded differences.""" + + if rollout.side != "rollout" or training.side != "training": + raise ValueError("compare_determinism expects one rollout probe and one training probe") + + issues: list[BindingIssue] = [] + differences: dict[str, dict[str, Any]] = {} + + for probe in (rollout, training): + if not probe.enabled: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"{probe.side}.{probe.mode_flag}", + rollout=rollout.enabled, + training=training.enabled, + message=( + f"{probe.framework} {probe.mode_flag} is not enabled; the " + f"{probe.side} side is not batch-invariant and cannot anchor a " + "cross-config comparison" + ), + ) + ) + + for key in COMPARED_NCCL_KEYS: + rollout_value = rollout.env.get(key) + training_value = training.env.get(key) + if rollout_value == training_value: + continue + differences[key] = {"rollout": rollout_value, "training": training_value} + if key in _ARITHMETIC_NCCL_KEYS: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"env.{key}", + rollout=rollout_value, + training=training_value, + message=( + f"{key} differs between sides; the two sides would reduce with " + "different arithmetic and the resulting drift is not attributable" + ), + ) + ) + + # Megatron reports None for these because it does not manage them at all. That is + # recorded rather than blocking: under a pure BF16 GEMM path TF32 does not fire, and + # forcing Megatron to manage it is out of scope for this PR. It is surfaced so the + # asymmetry appears in every artifact instead of being invisible. + for name in ("tf32_disabled", "bf16_reduced_precision_reduction"): + rollout_value = getattr(rollout, name) + training_value = getattr(training, name) + if rollout_value != training_value: + differences[name] = { + "rollout": rollout_value, + "training": training_value, + "note": ( + "megatron does not manage this setting; vllm sets it inside " + "init_batch_invariance()" + ), + } + + return DeterminismReport( + rollout=rollout, + training=training, + issues=tuple(issues), + differences=differences, + ) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py new file mode 100644 index 00000000..8b5f588b --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,1532 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for context-parallel standard softmax attention. + +The objects in this module describe a distributed attention invocation. They +do not shard tensors, launch collectives, or implement the ``(out, lse)`` +merge. Keeping description and materialization separate lets dispatch reject +an incompatible backend before any numerically different path is launched. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Iterable, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" +STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks" + + +class AttentionContractError(ValueError): + """Raised when attention metadata does not describe a valid invocation.""" + + +class AttentionRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class AttentionMode(str, Enum): + PREFILL = "prefill" + CHUNKED_PREFILL = "chunked_prefill" + DECODE = "decode" + + +class AttentionDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class AttentionMerge(str, Enum): + ONLINE_SOFTMAX_LSE = "online_softmax_lse" + + +class ReductionOrder(str, Enum): + GLOBAL_BLOCK_INDEX = "global_block_index" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +class SplitKVMode(str, Enum): + DISABLED = "disabled" + FIXED = "fixed" + AUTO = "auto" + + +class RoPEState(str, Enum): + PRE_ROPE = "pre_rope" + POST_ROPE = "post_rope" + + +class RoPECastPoint(str, Enum): + NONE = "none" + BEFORE_ROPE = "before_rope" + AFTER_ROPE = "after_rope" + + +class RoPEFusionBoundary(str, Enum): + UNFUSED_ROPE_ATTENTION = "unfused_rope_attention" + FUSED_ROPE_ATTENTION = "fused_rope_attention" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise AttentionContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise AttentionContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise AttentionContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _integer_tuple(values: Iterable[int], field: str) -> tuple[int, ...]: + try: + result = tuple(values) + except TypeError as exc: + raise AttentionContractError(f"{field} must be an iterable of integers") from exc + for index, value in enumerate(result): + if isinstance(value, bool) or not isinstance(value, int): + raise AttentionContractError(f"{field}[{index}] must be an integer; got {value!r}") + return result + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical TP/CP ownership for one attention invocation. + + TP head shards are currently required to be equal and contiguous. CP + sequence ownership may be uneven, but every local block must carry a stable + logical global index so a later implementation can merge by logical order + instead of collective arrival order. + """ + + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + global_q_heads: int + global_kv_heads: int + local_q_head_start: int + local_q_heads: int + local_kv_head_start: int + local_kv_heads: int + global_sequence_length: int + local_sequence_length: int + global_block_indices: tuple[int, ...] + global_block_token_starts: tuple[int, ...] + local_block_offsets: tuple[int, ...] + packed_sequence_offsets: tuple[int, ...] | None = None + + def __post_init__(self) -> None: + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + if tp_rank >= tp_world_size: + raise AttentionContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + if cp_rank >= cp_world_size: + raise AttentionContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" + ) + + global_q_heads = _positive_int(self.global_q_heads, "global_q_heads") + global_kv_heads = _positive_int(self.global_kv_heads, "global_kv_heads") + if global_q_heads % global_kv_heads != 0: + raise AttentionContractError( + f"global_q_heads={global_q_heads} must be divisible by " + f"global_kv_heads={global_kv_heads} for GQA" + ) + if global_q_heads % tp_world_size != 0 or global_kv_heads % tp_world_size != 0: + raise AttentionContractError( + "global Q/KV heads must be evenly divisible by tp_world_size; " + f"got Hq={global_q_heads}, Hkv={global_kv_heads}, TP={tp_world_size}" + ) + + local_q_heads = _positive_int(self.local_q_heads, "local_q_heads") + local_kv_heads = _positive_int(self.local_kv_heads, "local_kv_heads") + expected_q_heads = global_q_heads // tp_world_size + expected_kv_heads = global_kv_heads // tp_world_size + if local_q_heads != expected_q_heads or local_kv_heads != expected_kv_heads: + raise AttentionContractError( + "local TP head counts do not preserve the global Q/KV mapping; " + f"expected ({expected_q_heads}, {expected_kv_heads}), got " + f"({local_q_heads}, {local_kv_heads})" + ) + + local_q_head_start = _non_negative_int(self.local_q_head_start, "local_q_head_start") + local_kv_head_start = _non_negative_int(self.local_kv_head_start, "local_kv_head_start") + expected_q_start = tp_rank * expected_q_heads + expected_kv_start = tp_rank * expected_kv_heads + if local_q_head_start != expected_q_start or local_kv_head_start != expected_kv_start: + raise AttentionContractError( + "local TP head starts do not match contiguous rank ownership; " + f"expected ({expected_q_start}, {expected_kv_start}), got " + f"({local_q_head_start}, {local_kv_head_start})" + ) + + global_sequence_length = _positive_int( + self.global_sequence_length, "global_sequence_length" + ) + local_sequence_length = _positive_int(self.local_sequence_length, "local_sequence_length") + + block_indices = _integer_tuple(self.global_block_indices, "global_block_indices") + if not block_indices: + raise AttentionContractError("global_block_indices must not be empty") + if any(index < 0 for index in block_indices): + raise AttentionContractError("global_block_indices must be non-negative") + if any( + left >= right for left, right in zip(block_indices, block_indices[1:], strict=False) + ): + raise AttentionContractError( + "global_block_indices must be unique and strictly increasing" + ) + object.__setattr__(self, "global_block_indices", block_indices) + + block_token_starts = _integer_tuple( + self.global_block_token_starts, "global_block_token_starts" + ) + local_block_offsets = _integer_tuple(self.local_block_offsets, "local_block_offsets") + if len(block_token_starts) != len(block_indices): + raise AttentionContractError( + "global_block_token_starts must contain one entry per global_block_indices entry" + ) + if any(start < 0 for start in block_token_starts): + raise AttentionContractError("global_block_token_starts must be non-negative") + if len(local_block_offsets) != len(block_indices) + 1: + raise AttentionContractError( + "local_block_offsets must contain one boundary more than global_block_indices" + ) + if local_block_offsets[0] != 0 or local_block_offsets[-1] != local_sequence_length: + raise AttentionContractError( + "local_block_offsets must start at 0 and end at local_sequence_length" + ) + if any( + left >= right + for left, right in zip(local_block_offsets, local_block_offsets[1:], strict=False) + ): + raise AttentionContractError("local_block_offsets must be strictly increasing") + + previous_global_end = 0 + for index, (global_start, local_start, local_end) in enumerate( + zip( + block_token_starts, + local_block_offsets[:-1], + local_block_offsets[1:], + strict=True, + ) + ): + global_end = global_start + (local_end - local_start) + if global_end > global_sequence_length: + raise AttentionContractError( + f"global block {block_indices[index]} exceeds global_sequence_length" + ) + if index > 0 and global_start < previous_global_end: + raise AttentionContractError( + "global block token ranges must be non-overlapping and ordered" + ) + previous_global_end = global_end + object.__setattr__(self, "global_block_token_starts", block_token_starts) + object.__setattr__(self, "local_block_offsets", local_block_offsets) + + if self.packed_sequence_offsets is not None: + offsets = _integer_tuple(self.packed_sequence_offsets, "packed_sequence_offsets") + if len(offsets) < 2 or offsets[0] != 0: + raise AttentionContractError( + "packed_sequence_offsets must start at 0 and contain an end offset" + ) + if any(left >= right for left, right in zip(offsets, offsets[1:], strict=False)): + raise AttentionContractError("packed_sequence_offsets must be strictly increasing") + if offsets[-1] != local_sequence_length: + raise AttentionContractError( + "the final packed_sequence_offsets value must equal local_sequence_length; " + f"got {offsets[-1]} and {local_sequence_length}" + ) + object.__setattr__(self, "packed_sequence_offsets", offsets) + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic CP ``(out, lse)`` merge semantics.""" + + merge: AttentionMerge = AttentionMerge.ONLINE_SOFTMAX_LSE + acc_dtype: AttentionDType = AttentionDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + + def __post_init__(self) -> None: + object.__setattr__(self, "merge", _enum_value(AttentionMerge, self.merge, "merge")) + object.__setattr__( + self, "acc_dtype", _enum_value(AttentionDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError( + f"CP attention accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class SplitKVExecutionPlan: + """Actual backend-local Split-KV schedule emitted by a runtime. + + CP ownership is intentionally separate from this plan. ``boundaries`` + describe the canonical logical KV ranges reduced by one backend invocation; + CP may transport those partial states between ranks but must not change the + FP32 merge contract recorded here. + """ + + requested_mode: SplitKVMode + requested_split_size: int | None + actual_mode: SplitKVMode | None + actual_split_size: int | None + boundaries: tuple[tuple[int, int], ...] + merge_order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + acc_dtype: AttentionDType = AttentionDType.FP32 + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + backend: str = "reference" + source: str = "contract_exact" + fallback: bool = False + fallback_reason: str | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "requested_mode", + _enum_value(SplitKVMode, self.requested_mode, "requested_mode"), + ) + if self.actual_mode is not None: + object.__setattr__( + self, + "actual_mode", + _enum_value(SplitKVMode, self.actual_mode, "actual_mode"), + ) + for field_name in ("requested_split_size", "actual_split_size"): + value = getattr(self, field_name) + if value is not None: + _positive_int(value, field_name) + if self.requested_mode is SplitKVMode.FIXED: + if self.requested_split_size is None: + raise AttentionContractError( + "requested fixed Split-KV mode requires requested_split_size" + ) + elif self.requested_split_size is not None: + raise AttentionContractError( + "requested_split_size is only valid for requested fixed Split-KV mode" + ) + try: + boundaries = tuple(tuple(boundary) for boundary in self.boundaries) + except TypeError as exc: + raise AttentionContractError( + "Split-KV boundaries must be an iterable of (start, end) pairs" + ) from exc + previous_end = 0 + for index, boundary in enumerate(boundaries): + if len(boundary) != 2: + raise AttentionContractError( + f"Split-KV boundary {index} must contain exactly start and end" + ) + start, end = boundary + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + ): + raise AttentionContractError("Split-KV boundaries must satisfy 0 <= start < end") + if index > 0 and start != previous_end: + raise AttentionContractError( + "Split-KV boundaries must be contiguous and in logical KV order" + ) + previous_end = end + if self.actual_mode is None and boundaries: + raise AttentionContractError("unknown actual Split-KV plan cannot declare boundaries") + if self.actual_mode is not None and not boundaries: + raise AttentionContractError("actual Split-KV plan must declare logical boundaries") + if self.actual_mode is SplitKVMode.FIXED: + if self.actual_split_size is None: + raise AttentionContractError( + "actual fixed Split-KV mode requires actual_split_size" + ) + widths = tuple(end - start for start, end in boundaries) + if any(width != self.actual_split_size for width in widths[:-1]) or ( + widths and widths[-1] > self.actual_split_size + ): + raise AttentionContractError( + "fixed Split-KV boundaries must use actual_split_size except " + "for a shorter final split" + ) + elif self.actual_split_size is not None: + raise AttentionContractError( + "actual_split_size is only valid for actual fixed Split-KV mode" + ) + if self.actual_mode is SplitKVMode.DISABLED and len(boundaries) != 1: + raise AttentionContractError( + "disabled Split-KV execution must contain exactly one boundary" + ) + object.__setattr__(self, "boundaries", boundaries) + object.__setattr__( + self, + "merge_order", + _enum_value(ReductionOrder, self.merge_order, "merge_order"), + ) + object.__setattr__( + self, + "acc_dtype", + _enum_value(AttentionDType, self.acc_dtype, "acc_dtype"), + ) + object.__setattr__( + self, + "downcast_at", + _enum_value(DowncastPoint, self.downcast_at, "downcast_at"), + ) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError("Split-KV partial states must be merged in fp32") + if not isinstance(self.backend, str) or not self.backend.strip(): + raise AttentionContractError("Split-KV backend must be a non-empty string") + if not isinstance(self.source, str) or not self.source.strip(): + raise AttentionContractError("Split-KV plan source must be a non-empty string") + if not isinstance(self.fallback, bool): + raise AttentionContractError("Split-KV fallback must be a bool") + if self.fallback and not self.fallback_reason: + raise AttentionContractError("Split-KV fallback_reason is required for a fallback") + if not self.fallback and self.fallback_reason is not None: + raise AttentionContractError( + "Split-KV fallback_reason must be None when fallback=False" + ) + if not self.fallback and self.actual_mode is not None: + if self.actual_mode is not self.requested_mode: + raise AttentionContractError( + "actual Split-KV mode may differ from requested mode only for a fallback" + ) + if self.actual_split_size != self.requested_split_size: + raise AttentionContractError( + "actual Split-KV size may differ from requested size only for a fallback" + ) + + @property + def actual_split_count(self) -> int | None: + return None if self.actual_mode is None else len(self.boundaries) + + def to_dict(self) -> dict[str, Any]: + return { + "requested_split_kv_policy": self.requested_mode.value, + "requested_split_kv_size": self.requested_split_size, + "actual_split_kv_policy": ( + None if self.actual_mode is None else self.actual_mode.value + ), + "actual_split_kv_size": self.actual_split_size, + "actual_split_kv_count": self.actual_split_count, + "actual_split_boundaries": [list(boundary) for boundary in self.boundaries], + "split_kv_merge_order": self.merge_order.value, + "split_kv_accum_dtype": self.acc_dtype.value, + "split_kv_downcast_at": self.downcast_at.value, + "split_kv_backend": self.backend, + "split_kv_plan_source": self.source, + "split_kv_fallback": self.fallback, + "split_kv_fallback_reason": self.fallback_reason, + } + + +@dataclass(frozen=True) +class SplitKVSpec: + """Requested Split-KV policy shared by training and rollout paths.""" + + mode: SplitKVMode = SplitKVMode.DISABLED + fixed_split_size: int | None = None + strict_consistency: bool = True + + @classmethod + def disabled(cls, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls(mode=SplitKVMode.DISABLED, strict_consistency=strict_consistency) + + @classmethod + def fixed(cls, split_size: int, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls( + mode=SplitKVMode.FIXED, + fixed_split_size=split_size, + strict_consistency=strict_consistency, + ) + + @classmethod + def auto(cls, *, strict_consistency: bool = False) -> "SplitKVSpec": + return cls(mode=SplitKVMode.AUTO, strict_consistency=strict_consistency) + + def __post_init__(self) -> None: + object.__setattr__(self, "mode", _enum_value(SplitKVMode, self.mode, "split_kv.mode")) + if not isinstance(self.strict_consistency, bool): + raise AttentionContractError("split_kv.strict_consistency must be a bool") + if self.mode is SplitKVMode.FIXED: + if self.fixed_split_size is None: + raise AttentionContractError("fixed Split-KV policy requires fixed_split_size") + _positive_int(self.fixed_split_size, "split_kv.fixed_split_size") + elif self.fixed_split_size is not None: + raise AttentionContractError("fixed_split_size is only valid for fixed Split-KV policy") + if self.strict_consistency and self.mode is SplitKVMode.AUTO: + raise AttentionContractError( + "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency" + ) + + def resolve(self, total_kv_tokens: int, *, backend: str) -> SplitKVExecutionPlan: + """Resolve policies whose logical schedule is fully known by contract.""" + + total_kv_tokens = _positive_int(total_kv_tokens, "total_kv_tokens") + if self.mode is SplitKVMode.AUTO: + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=None, + actual_mode=None, + actual_split_size=None, + boundaries=(), + backend=backend, + source="runtime_required", + ) + split_size = total_kv_tokens if self.mode is SplitKVMode.DISABLED else self.fixed_split_size + assert split_size is not None + boundaries = tuple( + (start, min(start + split_size, total_kv_tokens)) + for start in range(0, total_kv_tokens, split_size) + ) + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=self.fixed_split_size, + actual_mode=self.mode, + actual_split_size=self.fixed_split_size, + boundaries=boundaries, + backend=backend, + source="contract_exact", + ) + + def to_dict(self) -> dict[str, Any]: + return { + "mode": self.mode.value, + "fixed_split_size": self.fixed_split_size, + "strict_consistency": self.strict_consistency, + } + + +def validate_split_kv_alignment( + training: SplitKVExecutionPlan, + rollout: SplitKVExecutionPlan, +) -> None: + """Fail closed unless train and rollout executed the same logical plan.""" + + if training.actual_mode is None or rollout.actual_mode is None: + raise AttentionContractError( + "strict Split-KV alignment requires actual runtime plans from both sides" + ) + fields = ( + "requested_mode", + "requested_split_size", + "actual_mode", + "actual_split_size", + "boundaries", + "merge_order", + "acc_dtype", + "downcast_at", + "fallback", + ) + mismatches = [ + field_name + for field_name in fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if mismatches: + raise AttentionContractError( + "training/rollout Split-KV execution plans differ: " + ", ".join(mismatches) + ) + + +@dataclass(frozen=True, order=True) +class SplitKVRuntimeCoordinate: + """Identity of one batch/rank/owner Split-KV runtime plan.""" + + batch_index: int + tp_rank: int + cp_rank: int + owner_cp_rank: int + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + ) -> None: + batch_index = _non_negative_int(self.batch_index, "batch_index") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + owner_cp_rank = _non_negative_int(self.owner_cp_rank, "owner_cp_rank") + if batch_index >= batch_size: + raise AttentionContractError("Split-KV batch_index is outside batch_size") + if tp_rank >= tp_world_size: + raise AttentionContractError("Split-KV tp_rank is outside tp_world_size") + if cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV cp_rank is outside cp_world_size") + if owner_cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV owner_cp_rank is outside cp_world_size") + + def to_dict(self) -> dict[str, int]: + return { + "batch_index": self.batch_index, + "tp_rank": self.tp_rank, + "cp_rank": self.cp_rank, + "owner_cp_rank": self.owner_cp_rank, + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanEntry: + """Actual plan for one batch/TP/CP consumer and logical KV owner.""" + + coordinate: SplitKVRuntimeCoordinate + expected_kv_range: tuple[int, int] + execution: SplitKVExecutionPlan + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + total_kv_tokens: int, + ) -> None: + self.coordinate.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + ) + try: + start, end = self.expected_kv_range + except (TypeError, ValueError) as exc: + raise AttentionContractError( + "expected_kv_range must contain exactly (start, end)" + ) from exc + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + or end > total_kv_tokens + ): + raise AttentionContractError( + "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens" + ) + if self.execution.actual_mode is None: + raise AttentionContractError("complete Split-KV plan sets require actual runtime plans") + if self.execution.boundaries[0][0] != start or self.execution.boundaries[-1][1] != end: + raise AttentionContractError( + "Split-KV execution boundaries must exactly cover expected_kv_range" + ) + if any( + boundary_start < start or boundary_end > end + for boundary_start, boundary_end in self.execution.boundaries + ): + raise AttentionContractError("Split-KV execution boundary escapes expected_kv_range") + + def to_dict(self) -> dict[str, Any]: + return { + **self.coordinate.to_dict(), + "expected_kv_range": list(self.expected_kv_range), + **self.execution.to_dict(), + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanSet: + """Complete actual Split-KV plans across batch, TP, CP, and KV owners. + + Every CP consumer must report the plan used for every CP-owned KV range. + This duplicates owner plans across consumers intentionally: it detects one + rank silently choosing a different Split-K schedule or merge policy. + """ + + batch_size: int + tp_world_size: int + cp_world_size: int + total_kv_tokens: tuple[int, ...] + entries: tuple[SplitKVRuntimePlanEntry, ...] + + def __post_init__(self) -> None: + batch_size = _positive_int(self.batch_size, "batch_size") + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + totals = _integer_tuple(self.total_kv_tokens, "total_kv_tokens") + if len(totals) != batch_size or any(total <= 0 for total in totals): + raise AttentionContractError( + "total_kv_tokens must contain one positive length per batch item" + ) + object.__setattr__(self, "total_kv_tokens", totals) + entries = tuple(self.entries) + object.__setattr__(self, "entries", entries) + expected_coordinates = { + SplitKVRuntimeCoordinate(batch_index, tp_rank, cp_rank, owner_cp_rank) + for batch_index in range(batch_size) + for tp_rank in range(tp_world_size) + for cp_rank in range(cp_world_size) + for owner_cp_rank in range(cp_world_size) + } + actual_coordinates = [entry.coordinate for entry in entries] + if len(set(actual_coordinates)) != len(actual_coordinates): + raise AttentionContractError("Split-KV runtime plan set contains duplicate coordinates") + missing = expected_coordinates.difference(actual_coordinates) + extra = set(actual_coordinates).difference(expected_coordinates) + if missing or extra: + raise AttentionContractError( + "Split-KV runtime plan set coordinate coverage is incomplete; " + f"missing={_format_split_kv_coordinates(missing)}, " + f"extra={_format_split_kv_coordinates(extra)}" + ) + for entry in entries: + entry.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals[entry.coordinate.batch_index], + ) + self._validate_owner_coverage() + self._validate_rank_invariance() + + def _validate_owner_coverage(self) -> None: + for batch_index, total in enumerate(self.total_kv_tokens): + for tp_rank in range(self.tp_world_size): + ranges = [] + for owner_cp_rank in range(self.cp_world_size): + matches = [ + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.tp_rank == tp_rank + and entry.coordinate.cp_rank == 0 + and entry.coordinate.owner_cp_rank == owner_cp_rank + ] + ranges.append(matches[0].expected_kv_range) + previous_end = 0 + for start, end in ranges: + if start != previous_end: + raise AttentionContractError( + "Split-KV owner ranges must be gap-free in CP owner order" + ) + previous_end = end + if previous_end != total: + raise AttentionContractError( + "Split-KV owner ranges do not cover total_kv_tokens" + ) + + def _validate_rank_invariance(self) -> None: + for batch_index in range(self.batch_size): + for owner_cp_rank in range(self.cp_world_size): + entries = sorted( + ( + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.owner_cp_rank == owner_cp_rank + ), + key=lambda entry: ( + entry.coordinate.tp_rank, + entry.coordinate.cp_rank, + ), + ) + reference = entries[0] + for entry in entries[1:]: + if entry.expected_kv_range != reference.expected_kv_range: + raise AttentionContractError( + "Split-KV owner range differs across TP/CP consumers" + ) + try: + validate_split_kv_alignment( + reference.execution, + entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + "Split-KV plan differs across TP/CP consumers for " + f"batch={batch_index}, owner_cp={owner_cp_rank}: {exc}" + ) from exc + + def to_dict(self) -> dict[str, Any]: + return { + "batch_size": self.batch_size, + "tp_world_size": self.tp_world_size, + "cp_world_size": self.cp_world_size, + "total_kv_tokens": list(self.total_kv_tokens), + "entries": [ + entry.to_dict() + for entry in sorted(self.entries, key=lambda entry: entry.coordinate) + ], + "coverage": "complete_batch_tp_cp_owner_cartesian_product", + } + + +def validate_split_kv_plan_set_alignment( + training: SplitKVRuntimePlanSet, + rollout: SplitKVRuntimePlanSet, +) -> None: + """Fail closed unless complete train/rollout runtime plan sets align.""" + + topology_fields = ( + "batch_size", + "tp_world_size", + "cp_world_size", + "total_kv_tokens", + ) + topology_mismatches = [ + field_name + for field_name in topology_fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if topology_mismatches: + raise AttentionContractError( + "training/rollout Split-KV plan-set topology differs: " + ", ".join(topology_mismatches) + ) + training_by_coordinate = {entry.coordinate: entry for entry in training.entries} + rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries} + if training_by_coordinate.keys() != rollout_by_coordinate.keys(): + raise AttentionContractError("training/rollout Split-KV plan-set coordinates differ") + for coordinate in sorted(training_by_coordinate): + train_entry = training_by_coordinate[coordinate] + rollout_entry = rollout_by_coordinate[coordinate] + if train_entry.expected_kv_range != rollout_entry.expected_kv_range: + raise AttentionContractError( + f"training/rollout expected KV range differs at {coordinate}" + ) + try: + validate_split_kv_alignment( + train_entry.execution, + rollout_entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + f"training/rollout Split-KV plan differs at {coordinate}: {exc}" + ) from exc + + +def build_split_kv_runtime_plan_set( + total_kv_tokens: Iterable[int], + *, + tp_world_size: int, + cp_world_size: int, + split_kv: SplitKVSpec, + backend: str = "contract_reference", +) -> SplitKVRuntimePlanSet: + """Build a complete owner-local plan set for contract tests and adapters.""" + + totals = _integer_tuple(total_kv_tokens, "total_kv_tokens") + if not totals or any(total < cp_world_size for total in totals): + raise AttentionContractError( + "contract plan sets require at least one KV token per CP owner" + ) + tp_world_size = _positive_int(tp_world_size, "tp_world_size") + cp_world_size = _positive_int(cp_world_size, "cp_world_size") + if not isinstance(split_kv, SplitKVSpec): + raise AttentionContractError("split_kv must be a SplitKVSpec") + + entries: list[SplitKVRuntimePlanEntry] = [] + for batch_index, total in enumerate(totals): + base = total // cp_world_size + remainder = total % cp_world_size + owner_ranges: list[tuple[int, int]] = [] + start = 0 + for owner_cp_rank in range(cp_world_size): + end = start + base + (1 if owner_cp_rank < remainder else 0) + owner_ranges.append((start, end)) + start = end + for tp_rank in range(tp_world_size): + for cp_rank in range(cp_world_size): + for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + local_total = owner_end - owner_start + local = split_kv.resolve(local_total, backend=backend) + execution = SplitKVExecutionPlan( + requested_mode=local.requested_mode, + requested_split_size=local.requested_split_size, + actual_mode=local.actual_mode, + actual_split_size=local.actual_split_size, + boundaries=tuple( + (owner_start + start, owner_start + end) + for start, end in local.boundaries + ), + merge_order=local.merge_order, + acc_dtype=local.acc_dtype, + downcast_at=local.downcast_at, + backend=local.backend, + source=local.source, + fallback=local.fallback, + fallback_reason=local.fallback_reason, + ) + entries.append( + SplitKVRuntimePlanEntry( + coordinate=SplitKVRuntimeCoordinate( + batch_index=batch_index, + tp_rank=tp_rank, + cp_rank=cp_rank, + owner_cp_rank=owner_cp_rank, + ), + expected_kv_range=(owner_start, owner_end), + execution=execution, + ) + ) + return SplitKVRuntimePlanSet( + batch_size=len(totals), + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals, + entries=tuple(entries), + ) + + +def _format_split_kv_coordinates( + coordinates: Iterable[SplitKVRuntimeCoordinate], +) -> list[dict[str, int]]: + return [coordinate.to_dict() for coordinate in sorted(coordinates)] + + +@dataclass(frozen=True) +class KVCacheSpec: + """Logical identity of the paged/block KV cache used for replay.""" + + cache_positions: tuple[int, ...] + kv_seq_lens: tuple[int, ...] + block_table: tuple[tuple[int, ...], ...] + global_token_positions: tuple[int, ...] + page_size: int + prefix_cache_enabled: bool = False + prefix_cache_key: str | None = None + shared_prefix_page_count: int = 0 + + def __post_init__(self) -> None: + cache_positions = _integer_tuple(self.cache_positions, "cache_positions") + kv_seq_lens = _integer_tuple(self.kv_seq_lens, "kv_seq_lens") + page_size = _positive_int(self.page_size, "page_size") + global_token_positions = _integer_tuple( + self.global_token_positions, "global_token_positions" + ) + if not cache_positions or any(position < 0 for position in cache_positions): + raise AttentionContractError("cache_positions must contain non-negative positions") + if not kv_seq_lens or any(length <= 0 for length in kv_seq_lens): + raise AttentionContractError("kv_seq_lens must contain positive sequence lengths") + if len(cache_positions) != len(kv_seq_lens): + raise AttentionContractError( + "cache_positions must contain one entry per kv_seq_lens entry" + ) + if not global_token_positions or any(position < 0 for position in global_token_positions): + raise AttentionContractError( + "global_token_positions must contain non-negative positions" + ) + if len(global_token_positions) != sum(kv_seq_lens): + raise AttentionContractError( + "global_token_positions must describe every logical cached token; " + f"expected {sum(kv_seq_lens)}, got {len(global_token_positions)}" + ) + token_offset = 0 + sequence_position_rows: list[tuple[int, ...]] = [] + for sequence_index, sequence_length in enumerate(kv_seq_lens): + sequence_positions = global_token_positions[ + token_offset : token_offset + sequence_length + ] + if any( + left >= right + for left, right in zip(sequence_positions, sequence_positions[1:], strict=False) + ): + raise AttentionContractError( + "global_token_positions must be strictly increasing within each sequence; " + f"sequence {sequence_index} is invalid" + ) + sequence_position_rows.append(sequence_positions) + token_offset += sequence_length + + for sequence_index, (cache_position, sequence_positions) in enumerate( + zip(cache_positions, sequence_position_rows, strict=True) + ): + terminal_position = sequence_positions[-1] + if cache_position != terminal_position: + raise AttentionContractError( + "cache_positions must equal the terminal global token position for each " + f"sequence; sequence {sequence_index} expected {terminal_position}, " + f"got {cache_position}" + ) + + try: + block_table = tuple(tuple(row) for row in self.block_table) + except TypeError as exc: + raise AttentionContractError( + "block_table must be a two-dimensional integer table" + ) from exc + if len(block_table) != len(kv_seq_lens) or any(not row for row in block_table): + raise AttentionContractError( + "block_table must contain one non-empty row per kv_seq_lens entry" + ) + active_block_rows: list[tuple[int, ...]] = [] + for row_index, (row, sequence_length) in enumerate( + zip(block_table, kv_seq_lens, strict=True) + ): + row_active_blocks: list[int] = [] + saw_padding = False + for column_index, block in enumerate(row): + if isinstance(block, bool) or not isinstance(block, int) or block < -1: + raise AttentionContractError( + "block_table entries must be integer block ids or -1 padding; " + f"got block_table[{row_index}][{column_index}]={block!r}" + ) + if block == -1: + saw_padding = True + continue + if saw_padding: + raise AttentionContractError( + "block_table -1 padding must be trailing; " + f"row {row_index} contains an active block after padding" + ) + row_active_blocks.append(block) + + expected_blocks = (sequence_length + page_size - 1) // page_size + if len(row_active_blocks) != expected_blocks: + raise AttentionContractError( + "block_table active page count must match kv_seq_lens and page_size; " + f"row {row_index} expected {expected_blocks}, got {len(row_active_blocks)}" + ) + if len(set(row_active_blocks)) != len(row_active_blocks): + raise AttentionContractError( + f"block_table row {row_index} contains duplicate active page ids" + ) + active_block_rows.append(tuple(row_active_blocks)) + + if not isinstance(self.prefix_cache_enabled, bool): + raise AttentionContractError("prefix_cache_enabled must be a bool") + shared_prefix_page_count = _non_negative_int( + self.shared_prefix_page_count, "shared_prefix_page_count" + ) + if self.prefix_cache_enabled and not self.prefix_cache_key: + raise AttentionContractError( + "prefix_cache_key is required when prefix_cache_enabled=True" + ) + if not self.prefix_cache_enabled and self.prefix_cache_key is not None: + raise AttentionContractError( + "prefix_cache_key must be None when prefix_cache_enabled=False" + ) + if not self.prefix_cache_enabled and shared_prefix_page_count != 0: + raise AttentionContractError( + "shared_prefix_page_count must be 0 when prefix_cache_enabled=False" + ) + + shared_prefix_pages: tuple[int, ...] = () + if shared_prefix_page_count > 0: + if any(len(row_blocks) < shared_prefix_page_count for row_blocks in active_block_rows): + raise AttentionContractError( + "shared_prefix_page_count exceeds an active block-table row" + ) + shared_prefix_token_count = shared_prefix_page_count * page_size + if any(length < shared_prefix_token_count for length in kv_seq_lens): + raise AttentionContractError( + "shared prefix pages must be fully populated and read-only" + ) + + shared_prefix_pages = active_block_rows[0][:shared_prefix_page_count] + shared_prefix_positions = sequence_position_rows[0][:shared_prefix_token_count] + for sequence_index, (row_blocks, positions) in enumerate( + zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 + ): + if row_blocks[:shared_prefix_page_count] != shared_prefix_pages: + raise AttentionContractError( + "shared prefix page ids must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + if positions[:shared_prefix_token_count] != shared_prefix_positions: + raise AttentionContractError( + "shared prefix token positions must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + + exclusive_page_owners: dict[int, int] = {} + shared_prefix_page_ids = set(shared_prefix_pages) + for sequence_index, row_blocks in enumerate(active_block_rows): + for page_id in row_blocks[shared_prefix_page_count:]: + if page_id in shared_prefix_page_ids: + raise AttentionContractError( + "a writable suffix page cannot alias a read-only shared prefix page" + ) + previous_owner = exclusive_page_owners.get(page_id) + if previous_owner is not None: + raise AttentionContractError( + "active pages may be shared across sequences only when declared as " + "read-only prefix pages; " + f"page {page_id} is used by sequences {previous_owner} and " + f"{sequence_index}" + ) + exclusive_page_owners[page_id] = sequence_index + + object.__setattr__(self, "cache_positions", cache_positions) + object.__setattr__(self, "kv_seq_lens", kv_seq_lens) + object.__setattr__(self, "block_table", block_table) + object.__setattr__(self, "global_token_positions", global_token_positions) + + +@dataclass(frozen=True) +class RoPESpec: + """Qwen3 RoPE identity for attention inputs and KV-cache replay. + + The CP attention reference consumes attention inputs. This object records + whether those inputs are pre- or post-RoPE and pins the position/cache + metadata needed to compare fused ``RoPE+Attention`` and unfused + ``RoPE -> Attention`` materializations. + """ + + q_state: RoPEState = RoPEState.POST_ROPE + k_state: RoPEState = RoPEState.POST_ROPE + k_cache_state: RoPEState = RoPEState.POST_ROPE + theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + position_ids: tuple[int, ...] | None = None + query_position_offsets: tuple[int, ...] | None = None + key_position_offsets: tuple[int, ...] | None = None + cast_at: RoPECastPoint = RoPECastPoint.AFTER_ROPE + output_dtype: AttentionDType = AttentionDType.BF16 + fusion_boundary: RoPEFusionBoundary = RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION + + def __post_init__(self) -> None: + object.__setattr__(self, "q_state", _enum_value(RoPEState, self.q_state, "q_state")) + object.__setattr__(self, "k_state", _enum_value(RoPEState, self.k_state, "k_state")) + object.__setattr__( + self, "k_cache_state", _enum_value(RoPEState, self.k_cache_state, "k_cache_state") + ) + if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + theta = float(self.theta) + if theta <= 0.0: + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + object.__setattr__(self, "theta", theta) + _positive_int(self.rotary_dim, "rotary_dim") + if self.rope_scaling is not None and ( + not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() + ): + raise AttentionContractError("rope_scaling must be a non-empty string when provided") + for position_field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, position_field) + if values is None: + continue + normalized = _integer_tuple(values, position_field) + if not normalized or any(value < 0 for value in normalized): + raise AttentionContractError( + f"{position_field} must contain non-negative positions" + ) + object.__setattr__(self, position_field, normalized) + object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) + object.__setattr__( + self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") + ) + object.__setattr__( + self, + "fusion_boundary", + _enum_value(RoPEFusionBoundary, self.fusion_boundary, "fusion_boundary"), + ) + + +@dataclass(frozen=True) +class AttentionContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: AttentionRole + mode: AttentionMode + dtype: AttentionDType + batch_size: int + query_sequence_length: int + head_dim: int + causal: bool + causal_offsets: tuple[int, ...] | None + sharding: ShardingSpec + reduction: ReductionSpec + split_kv: SplitKVSpec = field(default_factory=SplitKVSpec.disabled) + kv_cache: KVCacheSpec | None = None + rope: RoPESpec | None = None + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role")) + object.__setattr__(self, "mode", _enum_value(AttentionMode, self.mode, "mode")) + object.__setattr__(self, "dtype", _enum_value(AttentionDType, self.dtype, "dtype")) + batch_size = _positive_int(self.batch_size, "batch_size") + query_sequence_length = _positive_int(self.query_sequence_length, "query_sequence_length") + _positive_int(self.head_dim, "head_dim") + if not isinstance(self.sharding, ShardingSpec): + raise AttentionContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise AttentionContractError("reduction must be a ReductionSpec") + if not isinstance(self.split_kv, SplitKVSpec): + raise AttentionContractError("split_kv must be a SplitKVSpec") + if ( + self.mode is AttentionMode.PREFILL + and query_sequence_length != self.sharding.local_sequence_length + ): + raise AttentionContractError( + "prefill query_sequence_length must equal sharding.local_sequence_length; " + f"got {query_sequence_length} and {self.sharding.local_sequence_length}" + ) + if self.sharding.packed_sequence_offsets is not None: + packed_sequence_count = len(self.sharding.packed_sequence_offsets) - 1 + if packed_sequence_count != batch_size: + raise AttentionContractError( + "packed sequence count must equal logical batch_size; " + f"got {packed_sequence_count} packed sequences and batch_size={batch_size}" + ) + if not isinstance(self.causal, bool): + raise AttentionContractError("causal must be a bool") + if self.causal: + if self.causal_offsets is None: + raise AttentionContractError("causal_offsets are required for causal attention") + if self.causal_offsets is not None: + causal_offsets = _integer_tuple(self.causal_offsets, "causal_offsets") + if not causal_offsets or any(offset < 0 for offset in causal_offsets): + raise AttentionContractError("causal_offsets must contain non-negative offsets") + if self.sharding.packed_sequence_offsets is not None: + expected_causal_offsets = batch_size + offset_owner = "packed sequence" + else: + expected_causal_offsets = batch_size + offset_owner = "batch entry" + if len(causal_offsets) != expected_causal_offsets: + raise AttentionContractError( + f"causal_offsets must contain one entry per {offset_owner}" + ) + object.__setattr__(self, "causal_offsets", causal_offsets) + + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise AttentionContractError( + "export_lse must be True for the WS2 attention-domain LSE contract" + ) + + if self.mode is AttentionMode.DECODE and self.kv_cache is None: + raise AttentionContractError("kv_cache metadata is required for decode attention") + if self.kv_cache is not None and not isinstance(self.kv_cache, KVCacheSpec): + raise AttentionContractError("kv_cache must be a KVCacheSpec when provided") + if self.rope is not None: + if not isinstance(self.rope, RoPESpec): + raise AttentionContractError("rope must be a RoPESpec when provided") + if self.rope.rotary_dim > self.head_dim: + raise AttentionContractError( + f"rotary_dim={self.rope.rotary_dim} must not exceed head_dim={self.head_dim}" + ) + if self.rope.position_ids is not None and len(self.rope.position_ids) not in { + query_sequence_length, + self.sharding.local_sequence_length, + }: + raise AttentionContractError( + "position_ids must describe the local query sequence or full local " + "sequence length" + ) + for position_field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, position_field) + if offsets is not None and len(offsets) != batch_size: + raise AttentionContractError( + f"{position_field} must contain one entry per logical batch entry" + ) + if self.mode is AttentionMode.DECODE and self.kv_cache is not None: + if len(self.kv_cache.kv_seq_lens) != batch_size: + raise AttentionContractError( + "decode kv_seq_lens must contain one entry per batch entry" + ) + if len(self.kv_cache.cache_positions) != batch_size: + raise AttentionContractError( + "decode cache_positions must contain one entry per batch entry" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "global_q_heads": self.sharding.global_q_heads, + "global_kv_heads": self.sharding.global_kv_heads, + "local_q_head_start": self.sharding.local_q_head_start, + "local_q_heads": self.sharding.local_q_heads, + "local_kv_head_start": self.sharding.local_kv_head_start, + "local_kv_heads": self.sharding.local_kv_heads, + "global_sequence_length": self.sharding.global_sequence_length, + "local_sequence_length": self.sharding.local_sequence_length, + "global_block_indices": list(self.sharding.global_block_indices), + "global_block_token_starts": list(self.sharding.global_block_token_starts), + "local_block_offsets": list(self.sharding.local_block_offsets), + "packed_sequence_offsets": ( + list(self.sharding.packed_sequence_offsets) + if self.sharding.packed_sequence_offsets is not None + else None + ), + } + reduction = { + "merge": self.reduction.merge.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + } + kv_cache = None + if self.kv_cache is not None: + kv_cache = { + "cache_positions": list(self.kv_cache.cache_positions), + "kv_seq_lens": list(self.kv_cache.kv_seq_lens), + "block_table": [list(row) for row in self.kv_cache.block_table], + "global_token_positions": list(self.kv_cache.global_token_positions), + "page_size": self.kv_cache.page_size, + "prefix_cache_enabled": self.kv_cache.prefix_cache_enabled, + "prefix_cache_key": self.kv_cache.prefix_cache_key, + "shared_prefix_page_count": self.kv_cache.shared_prefix_page_count, + } + rope = None + if self.rope is not None: + rope = { + "q_state": self.rope.q_state.value, + "k_state": self.rope.k_state.value, + "k_cache_state": self.rope.k_cache_state.value, + "theta": self.rope.theta, + "rotary_dim": self.rope.rotary_dim, + "rope_scaling": self.rope.rope_scaling, + "position_ids": ( + list(self.rope.position_ids) if self.rope.position_ids is not None else None + ), + "query_position_offsets": ( + list(self.rope.query_position_offsets) + if self.rope.query_position_offsets is not None + else None + ), + "key_position_offsets": ( + list(self.rope.key_position_offsets) + if self.rope.key_position_offsets is not None + else None + ), + "cast_at": self.rope.cast_at.value, + "output_dtype": self.rope.output_dtype.value, + "fusion_boundary": self.rope.fusion_boundary.value, + } + return { + "semantic_operator": "standard_softmax_attention", + "role": self.role.value, + "mode": self.mode.value, + "dtype": self.dtype.value, + "batch_size": self.batch_size, + "query_sequence_length": self.query_sequence_length, + "head_dim": self.head_dim, + "causal": self.causal, + "causal_offsets": ( + list(self.causal_offsets) if self.causal_offsets is not None else None + ), + "export_lse": self.export_lse, + "lse_domain": "attention", + "sharding": sharding, + "reduction": reduction, + "split_kv": self.split_kv.to_dict(), + "kv_cache": kv_cache, + "rope": rope, + } + + +@dataclass(frozen=True) +class AttentionBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[AttentionRole] + modes: frozenset[AttentionMode] + dtypes: frozenset[AttentionDType] + cp_world_sizes: tuple[int, ...] + tp_world_sizes: tuple[int, ...] | None = None + exports_attention_lse: bool = False + deterministic_cp_merge: bool = False + supports_packed_varlen: bool = False + supports_kv_cache: bool = False + supports_rope_metadata: bool = False + supports_fused_rope_attention: bool = False + supports_split_kv_disabled: bool = True + supports_split_kv_fixed: bool = False + supports_split_kv_auto: bool = False + reports_actual_split_kv_plan: bool = False + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise AttentionContractError("backend_id must be a non-empty string") + roles = frozenset(_enum_value(AttentionRole, value, "roles") for value in self.roles) + modes = frozenset(_enum_value(AttentionMode, value, "modes") for value in self.modes) + dtypes = frozenset(_enum_value(AttentionDType, value, "dtypes") for value in self.dtypes) + if not roles or not modes or not dtypes: + raise AttentionContractError("backend roles, modes, and dtypes must not be empty") + cp_world_sizes = _integer_tuple(self.cp_world_sizes, "cp_world_sizes") + if not cp_world_sizes or any(size <= 0 for size in cp_world_sizes): + raise AttentionContractError("cp_world_sizes must contain positive values") + if len(set(cp_world_sizes)) != len(cp_world_sizes): + raise AttentionContractError("cp_world_sizes must not contain duplicates") + tp_world_sizes = None + if self.tp_world_sizes is not None: + tp_world_sizes = _integer_tuple(self.tp_world_sizes, "tp_world_sizes") + if not tp_world_sizes or any(size <= 0 for size in tp_world_sizes): + raise AttentionContractError("tp_world_sizes must contain positive values") + if len(set(tp_world_sizes)) != len(tp_world_sizes): + raise AttentionContractError("tp_world_sizes must not contain duplicates") + for capability_field in ( + "exports_attention_lse", + "deterministic_cp_merge", + "supports_packed_varlen", + "supports_kv_cache", + "supports_rope_metadata", + "supports_fused_rope_attention", + "supports_split_kv_disabled", + "supports_split_kv_fixed", + "supports_split_kv_auto", + "reports_actual_split_kv_plan", + ): + if not isinstance(getattr(self, capability_field), bool): + raise AttentionContractError(f"{capability_field} must be a bool") + if self.implementation_kind not in {"production", "reference", "deterministic"}: + raise AttentionContractError( + "implementation_kind must be production, reference, or deterministic" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "modes", modes) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + + def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.mode not in self.modes: + reasons.append(f"mode={contract.mode.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if contract.export_lse and not self.exports_attention_lse: + reasons.append("attention-domain LSE export is unsupported") + if cp_size > 1 and not self.deterministic_cp_merge: + reasons.append("deterministic CP (out, lse) merge is unsupported") + if ( + contract.sharding.packed_sequence_offsets is not None + and not self.supports_packed_varlen + ): + reasons.append("packed varlen layout is unsupported") + if contract.kv_cache is not None and not self.supports_kv_cache: + reasons.append("KV-cache identity materialization is unsupported") + if contract.rope is not None and not self.supports_rope_metadata: + reasons.append("RoPE/position metadata is unsupported") + if ( + contract.rope is not None + and contract.rope.fusion_boundary is RoPEFusionBoundary.FUSED_ROPE_ATTENTION + and not self.supports_fused_rope_attention + ): + reasons.append("fused RoPE+Attention boundary is unsupported") + split_support = { + SplitKVMode.DISABLED: self.supports_split_kv_disabled, + SplitKVMode.FIXED: self.supports_split_kv_fixed, + SplitKVMode.AUTO: self.supports_split_kv_auto, + } + if not split_support[contract.split_kv.mode]: + reasons.append(f"Split-KV policy={contract.split_kv.mode.value} is unsupported") + if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan: + reasons.append("actual Split-KV execution-plan provenance is unsupported") + return tuple(reasons) + + def supports(self, contract: AttentionContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "modes": sorted(mode.value for mode in self.modes), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes), + "exports_attention_lse": self.exports_attention_lse, + "deterministic_cp_merge": self.deterministic_cp_merge, + "supports_packed_varlen": self.supports_packed_varlen, + "supports_kv_cache": self.supports_kv_cache, + "supports_rope_metadata": self.supports_rope_metadata, + "supports_fused_rope_attention": self.supports_fused_rope_attention, + "supports_split_kv_disabled": self.supports_split_kv_disabled, + "supports_split_kv_fixed": self.supports_split_kv_fixed, + "supports_split_kv_auto": self.supports_split_kv_auto, + "reports_actual_split_kv_plan": self.reports_actual_split_kv_plan, + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class AttentionDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: AttentionBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "AttentionContract", + "AttentionContractError", + "AttentionBackendCapability", + "AttentionDispatchResult", + "AttentionDType", + "AttentionMerge", + "AttentionMode", + "AttentionRole", + "DowncastPoint", + "KVCacheSpec", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "RoPECastPoint", + "RoPEFusionBoundary", + "RoPESpec", + "RoPEState", + "ShardingSpec", + "build_split_kv_runtime_plan_set", + "SplitKVExecutionPlan", + "SplitKVMode", + "SplitKVRuntimeCoordinate", + "SplitKVRuntimePlanEntry", + "SplitKVRuntimePlanSet", + "SplitKVSpec", + "STRICT_ATTENTION_CORE_ID", + "STRICT_ATTENTION_SCHEDULE_ID", + "validate_split_kv_alignment", + "validate_split_kv_plan_set_alignment", +] diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py new file mode 100644 index 00000000..e97b365f --- /dev/null +++ b/rl_engine/kernels/attention_preprocess.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Bitwise-bound QK-Norm and RoPE handoff for WS2 Attention. + +Megatron/TE and vLLM/FlashInfer remain the first-choice implementations. They +are admitted only after the same-input H100 probe is bitwise identical to the +deterministic RL-Kernel path. A failed probe (or an unavailable native backend) +selects the common RL-Kernel path for both sides and records why the fallback +was taken. The caller must pass this readback to the cross-config binder. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Callable, Mapping + +import torch +from torch import Tensor + +QK_RMSNORM_BACKEND_ID = "rlkernel.cuda.rmsnorm" +ROPE_BACKEND_ID = "rlkernel.cuda.rope_sm90" +NATIVE_QK_RMSNORM_BACKEND_ID = "native.qk_rmsnorm" +NATIVE_ROPE_BACKEND_ID = "native.rope" +PREPROCESS_POLICY_ID = "ws2.attention.preprocess.v2" +MANDATED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( + { + "qk_rmsnorm": QK_RMSNORM_BACKEND_ID, + "rope": ROPE_BACKEND_ID, + } +) + + +@dataclass(frozen=True) +class AttentionPreprocessResult: + """Post-QK-Norm, post-RoPE tensors plus executed backend evidence.""" + + q: Tensor + k: Tensor + backend_ids: Mapping[str, str] + fallback: bool + device_capability: tuple[int, int] + fallback_reason: str | None = None + probe_id: str = "" + policy_id: str = PREPROCESS_POLICY_ID + + def __post_init__(self) -> None: + object.__setattr__(self, "backend_ids", MappingProxyType(dict(self.backend_ids))) + + def evidence(self) -> dict[str, Any]: + return { + "backends": dict(self.backend_ids), + "fallback": self.fallback, + "device_capability": list(self.device_capability), + } + + def readback_fields(self) -> dict[str, Any]: + """Keyword fields consumed by ``AttentionRuntimeReadback``.""" + + return { + "preprocess_backends": dict(self.backend_ids), + "preprocess_fallback": self.fallback, + "preprocess_fallback_reason": self.fallback_reason, + "preprocess_probe_id": self.probe_id, + "preprocess_policy_id": self.policy_id, + } + + +class H100AttentionPreprocessor: + """Apply native QK-Norm/RoPE when the H100 probe passes. + + ``native_qk_norm`` and ``native_rope`` are framework-owned callables. They + are intentionally injected instead of importing TE/vLLM here, so the same + policy can be used by both runtimes. The deterministic callables default to + RL-Kernel's CUDA operators and are always run to establish the probe oracle. + """ + + def __init__( + self, + device: torch.device | str | int | None = None, + *, + native_qk_norm: Callable[..., Tensor] | None = None, + native_rope: Callable[..., Tensor] | None = None, + native_qk_norm_backend_id: str = NATIVE_QK_RMSNORM_BACKEND_ID, + native_rope_backend_id: str = NATIVE_ROPE_BACKEND_ID, + policy_id: str = PREPROCESS_POLICY_ID, + ) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("H100AttentionPreprocessor requires an available CUDA runtime") + + current_device = torch.cuda.current_device() + self.device = torch.device("cuda", current_device) + if device is not None: + self.device = ( + torch.device("cuda", device) if isinstance(device, int) else torch.device(device) + ) + if self.device.type != "cuda": + raise RuntimeError(f"H100AttentionPreprocessor requires CUDA, got {self.device}") + if self.device.index is None: + self.device = torch.device("cuda", current_device) + + capability = torch.cuda.get_device_capability(self.device) + self.device_capability: tuple[int, int] = (int(capability[0]), int(capability[1])) + if self.device_capability[0] != 9: + raise RuntimeError( + "H100AttentionPreprocessor requires Hopper SM90; " + f"got sm_{self.device_capability[0]}{self.device_capability[1]}" + ) + + # Import only after the hardware gate so CPU tools can inspect the module. + from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op + + self.rmsnorm = RMSNormCudaOp() + self.rope = RoPESM90Op() + if not isinstance(policy_id, str) or not policy_id.strip(): + raise ValueError("policy_id must be a non-empty string") + self.native_qk_norm = native_qk_norm + self.native_rope = native_rope + self.native_qk_norm_backend_id = native_qk_norm_backend_id + self.native_rope_backend_id = native_rope_backend_id + self.policy_id = policy_id + + def __call__( + self, + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + *, + eps: float = 1.0e-6, + theta: float = 1_000_000.0, + ) -> AttentionPreprocessResult: + return self.forward( + q, + k, + q_weight, + k_weight, + positions, + eps=eps, + theta=theta, + ) + + def forward( + self, + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + *, + eps: float = 1.0e-6, + theta: float = 1_000_000.0, + ) -> AttentionPreprocessResult: + _validate_inputs(q, k, q_weight, k_weight, positions, self.device) + q_norm_det = self.rmsnorm(q, q_weight, eps=eps) + k_norm_det = self.rmsnorm(k, k_weight, eps=eps) + q_det = self.rope(q_norm_det, positions, theta=theta) + k_det = self.rope(k_norm_det, positions, theta=theta) + + native_qk_norm = self.native_qk_norm + native_rope = self.native_rope + if native_qk_norm is not None and native_rope is not None: + try: + q_norm_native = native_qk_norm(q, q_weight, eps=eps) + k_norm_native = native_qk_norm(k, k_weight, eps=eps) + q_native = native_rope(q_norm_native, positions, theta=theta) + k_native = native_rope(k_norm_native, positions, theta=theta) + probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) + if ( + torch.equal(q_norm_native, q_norm_det) + and torch.equal(k_norm_native, k_norm_det) + and torch.equal(q_native, q_det) + and torch.equal(k_native, k_det) + ): + return AttentionPreprocessResult( + q=q_native, + k=k_native, + backend_ids=MappingProxyType( + { + "qk_rmsnorm": self.native_qk_norm_backend_id, + "rope": self.native_rope_backend_id, + } + ), + fallback=False, + device_capability=self.device_capability, + probe_id=probe_id, + policy_id=self.policy_id, + ) + fallback_reason = "native_preprocess_bitwise_probe_failed" + except Exception as exc: # framework backend failures must fail over together + fallback_reason = f"native_preprocess_unavailable:{type(exc).__name__}" + probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) + else: + fallback_reason = "native_preprocess_not_supplied" + probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) + return AttentionPreprocessResult( + q=q_det, + k=k_det, + backend_ids=MANDATED_ATTENTION_PREPROCESS_BACKENDS, + fallback=True, + device_capability=self.device_capability, + fallback_reason=fallback_reason, + probe_id=probe_id, + policy_id=self.policy_id, + ) + + +def _probe_id( + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + eps: float, + theta: float, +) -> str: + payload = { + "q_shape": list(q.shape), + "k_shape": list(k.shape), + "q_dtype": str(q.dtype), + "k_dtype": str(k.dtype), + "weight_dtype": str(q_weight.dtype), + "positions_shape": list(positions.shape), + "positions_sha256": hashlib.sha256(positions.detach().cpu().numpy().tobytes()).hexdigest(), + "eps": float(eps), + "theta": float(theta), + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16] + + +def _validate_inputs( + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + device: torch.device, +) -> None: + if q.dim() != 4 or k.dim() != 4: + raise ValueError("q and k must use [B, H, S, D] layout") + if q.shape[0] != k.shape[0] or q.shape[-2:] != k.shape[-2:]: + raise ValueError("q and k must have the same batch, sequence, and head dimensions") + if q.dtype is not torch.bfloat16 or k.dtype is not torch.bfloat16: + raise TypeError("the frozen H100 Attention experiment requires BF16 q and k") + if q.device != device or k.device != device: + raise ValueError(f"q and k must both be on the configured device {device}") + for name, weight in (("q_weight", q_weight), ("k_weight", k_weight)): + if weight.shape != (q.shape[-1],): + raise ValueError(f"{name} must have shape ({q.shape[-1]},)") + if weight.device != device or weight.dtype is not torch.bfloat16: + raise ValueError(f"{name} must be BF16 on {device}") + if positions.device != device: + raise ValueError(f"positions must be on {device}") + if positions.dtype not in (torch.int32, torch.int64): + raise TypeError("positions must use int32 or int64 global token indices") + expected = (q.shape[-2],) if positions.dim() == 1 else (q.shape[0], q.shape[-2]) + if positions.dim() not in (1, 2) or tuple(positions.shape) != expected: + raise ValueError(f"positions must have shape [S] or [B, S], expected {expected}") + + +__all__ = [ + "AttentionPreprocessResult", + "H100AttentionPreprocessor", + "MANDATED_ATTENTION_PREPROCESS_BACKENDS", + "QK_RMSNORM_BACKEND_ID", + "ROPE_BACKEND_ID", + "NATIVE_QK_RMSNORM_BACKEND_ID", + "NATIVE_ROPE_BACKEND_ID", + "PREPROCESS_POLICY_ID", +] diff --git a/rl_engine/kernels/attention_projection.py b/rl_engine/kernels/attention_projection.py new file mode 100644 index 00000000..a7a91bfe --- /dev/null +++ b/rl_engine/kernels/attention_projection.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""QKV and output-projection boundaries for the WS2 Attention experiment. + +The framework still owns the native TE/vLLM implementation. This wrapper only +freezes the semantics that must be shared by training and inference: BF16 I/O, +FP32 accumulation, ascending K reduction, and no Split-K. A native callable is +accepted only when its result is bitwise equal to the deterministic fallback. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Callable, Mapping + +import torch +from torch import Tensor + +from rl_engine.kernels.ops.cuda.matmul.det_gemm import DetGemmOp + +ProjectionCallable = Callable[[Tensor, Tensor], Tensor] + +QKV_PROJECTION = "qkv" +O_PROJ_PROJECTION = "o_proj" +PROJECTION_POLICY_ID = "ws2.attention.projection.v1" + + +@dataclass(frozen=True) +class ProjectionCollectiveContract: + """TP/SP directions fixed by the Attention table.""" + + projection: str + tp_forward: str + tp_backward: str + sp_forward: str + sp_backward: str + reduction_forward: str + reduction_backward: str + + def __post_init__(self) -> None: + if self.projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: + raise ValueError(f"unsupported projection {self.projection!r}") + + def to_dict(self) -> dict[str, str]: + return { + "projection": self.projection, + "tp_forward": self.tp_forward, + "tp_backward": self.tp_backward, + "sp_forward": self.sp_forward, + "sp_backward": self.sp_backward, + "reduction_forward": self.reduction_forward, + "reduction_backward": self.reduction_backward, + } + + +QKV_COLLECTIVE_CONTRACT = ProjectionCollectiveContract( + projection=QKV_PROJECTION, + tp_forward="column_parallel", + tp_backward="all_reduce", + sp_forward="all_gather", + sp_backward="reduce_scatter", + reduction_forward="none", + reduction_backward="none", +) +O_PROJ_COLLECTIVE_CONTRACT = ProjectionCollectiveContract( + projection=O_PROJ_PROJECTION, + tp_forward="row_parallel", + tp_backward="none", + sp_forward="reduce_scatter", + sp_backward="all_gather", + reduction_forward="all_reduce", + reduction_backward="none", +) + + +@dataclass(frozen=True) +class ProjectionPlan: + projection: str + backend_id: str + fallback: bool + fallback_reason: str | None + probe_id: str + input_dtype: str = "torch.bfloat16" + weight_dtype: str = "torch.bfloat16" + output_dtype: str = "torch.bfloat16" + accumulation_dtype: str = "torch.float32" + reduction_order: str = "k_ascending" + split_k: bool = False + policy_id: str = PROJECTION_POLICY_ID + collective: Mapping[str, str] = field(default_factory=dict) + + def __post_init__(self) -> None: + if self.projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: + raise ValueError(f"unsupported projection {self.projection!r}") + if self.input_dtype != "torch.bfloat16" or self.weight_dtype != "torch.bfloat16": + raise ValueError("Attention projections require BF16 input and weight") + if self.output_dtype != "torch.bfloat16" or self.accumulation_dtype != "torch.float32": + raise ValueError("Attention projections require FP32 accumulation and BF16 output") + if self.reduction_order != "k_ascending" or self.split_k: + raise ValueError( + "Attention projections require ascending K reduction with Split-K disabled" + ) + object.__setattr__(self, "collective", MappingProxyType(dict(self.collective))) + + def to_dict(self) -> dict[str, Any]: + return { + "projection": self.projection, + "backend_id": self.backend_id, + "fallback": self.fallback, + "fallback_reason": self.fallback_reason, + "probe_id": self.probe_id, + "input_dtype": self.input_dtype, + "weight_dtype": self.weight_dtype, + "output_dtype": self.output_dtype, + "accumulation_dtype": self.accumulation_dtype, + "reduction_order": self.reduction_order, + "split_k": self.split_k, + "policy_id": self.policy_id, + "collective": dict(self.collective), + } + + +@dataclass(frozen=True) +class ProjectionResult: + output: Tensor + plan: ProjectionPlan + + def to_readback(self) -> dict[str, Any]: + return self.plan.to_dict() + + +class AttentionProjectionOp: + """Native-first projection wrapper with a deterministic common fallback.""" + + def __init__( + self, + projection: str, + *, + native: ProjectionCallable | None = None, + native_backend_id: str | None = None, + deterministic: ProjectionCallable | None = None, + policy_id: str = PROJECTION_POLICY_ID, + ) -> None: + if projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: + raise ValueError(f"unsupported projection {projection!r}") + self.projection = projection + self.native = native + self.native_backend_id = native_backend_id or f"native.{projection}" + self.deterministic = deterministic or DetGemmOp() + self.policy_id = policy_id + self.collective = ( + QKV_COLLECTIVE_CONTRACT if projection == QKV_PROJECTION else O_PROJ_COLLECTIVE_CONTRACT + ) + + def __call__(self, x: Tensor, weight: Tensor) -> ProjectionResult: + _validate_projection_inputs(x, weight) + deterministic_out = self.deterministic(x, weight) + if deterministic_out.dtype is not torch.bfloat16: + deterministic_out = deterministic_out.to(torch.bfloat16) + probe_id = _probe_id(x, weight) + + if self.native is not None: + try: + native_out = self.native(x, weight) + if native_out.dtype is not torch.bfloat16: + native_out = native_out.to(torch.bfloat16) + if torch.equal(native_out, deterministic_out): + return ProjectionResult( + native_out, + ProjectionPlan( + projection=self.projection, + backend_id=self.native_backend_id, + fallback=False, + fallback_reason=None, + probe_id=probe_id, + policy_id=self.policy_id, + collective=self.collective.to_dict(), + ), + ) + reason = "native_projection_bitwise_probe_failed" + except Exception as exc: # framework backend failure: use common fallback + reason = f"native_projection_unavailable:{type(exc).__name__}" + else: + reason = "native_projection_not_supplied" + + return ProjectionResult( + deterministic_out, + ProjectionPlan( + projection=self.projection, + backend_id="rlkernel.cuda.det_gemm", + fallback=True, + fallback_reason=reason, + probe_id=probe_id, + policy_id=self.policy_id, + collective=self.collective.to_dict(), + ), + ) + + +def split_qkv( + projected_qkv: Tensor, + q_heads: int, + kv_heads: int, + head_dim: int, +) -> tuple[Tensor, Tensor, Tensor]: + """Split a [Q, K, V] projection in the fixed contiguous Q/K/V order.""" + + if projected_qkv.dim() != 2: + raise ValueError("projected QKV must be [tokens, features]") + q_width = q_heads * head_dim + kv_width = kv_heads * head_dim + expected = q_width + kv_width + kv_width + if projected_qkv.shape[-1] != expected: + raise ValueError(f"projected QKV width must be {expected}, got {projected_qkv.shape[-1]}") + q, k, v = projected_qkv.split((q_width, kv_width, kv_width), dim=-1) + return q, k, v + + +def _validate_projection_inputs(x: Tensor, weight: Tensor) -> None: + if x.dim() != 2 or weight.dim() != 2: + raise ValueError("projection inputs must be [tokens, K] and [K, N]") + if x.shape[-1] != weight.shape[0]: + raise ValueError("projection K dimensions must match") + if x.dtype is not torch.bfloat16 or weight.dtype is not torch.bfloat16: + raise TypeError("Attention projections require BF16 inputs and weights") + if x.device != weight.device: + raise ValueError("projection inputs must be on the same device") + + +def _probe_id(x: Tensor, weight: Tensor) -> str: + payload = { + "x_shape": list(x.shape), + "weight_shape": list(weight.shape), + "device": str(x.device), + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16] + + +__all__ = [ + "AttentionProjectionOp", + "O_PROJ_COLLECTIVE_CONTRACT", + "O_PROJ_PROJECTION", + "PROJECTION_POLICY_ID", + "ProjectionCollectiveContract", + "ProjectionPlan", + "ProjectionResult", + "QKV_COLLECTIVE_CONTRACT", + "QKV_PROJECTION", + "split_qkv", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/__init__.py b/rl_engine/kernels/ops/pytorch/attention/__init__.py index d2454e67..977ab14c 100644 --- a/rl_engine/kernels/ops/pytorch/attention/__init__.py +++ b/rl_engine/kernels/ops/pytorch/attention/__init__.py @@ -4,6 +4,12 @@ import torch import torch.nn.functional as F +from rl_engine.kernels.ops.pytorch.attention.ablation import ( + AttentionAblationConfig, + AttentionAblationOp, + AttentionAblationResult, +) + class NativeAttentionOp: """PyTorch SDPA fallback for FlashAttention-layout tensors.""" @@ -46,4 +52,9 @@ def __call__( return out.transpose(1, 2) -__all__ = ["NativeAttentionOp"] +__all__ = [ + "AttentionAblationConfig", + "AttentionAblationOp", + "AttentionAblationResult", + "NativeAttentionOp", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py new file mode 100644 index 00000000..dd463906 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -0,0 +1,650 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Unified Attention entry point for the PR230 cross-configuration matrix. + +The matrix needs one stable callable shape even though training and rollout may +materialize different Attention backends. This adapter owns the common +contract checks and provenance only; numerical work remains in the existing +deterministic Attention implementations. +""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import math +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Callable, Mapping + +import torch +from torch import Tensor + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + AttentionContract, + AttentionContractError, + AttentionDType, + SplitKVMode, +) + +BACKEND_ID = "rlkernel.attention.deterministic.v1" +REFERENCE_BACKEND_ID = "rlkernel.attention.reference.v1" + +_TORCH_DTYPES = { + AttentionDType.BF16: torch.bfloat16, + AttentionDType.FP16: torch.float16, + AttentionDType.FP32: torch.float32, +} + + +@dataclass(frozen=True) +class AttentionAblationConfig: + """Per-invocation settings materialized by the ablation runner.""" + + backend: str = "auto" + deterministic: bool = True + communication_backend: str = "none" + return_lse: bool = True + return_gradients: bool = False + strict_core_id: str = STRICT_ATTENTION_CORE_ID + strict_schedule: str = STRICT_ATTENTION_SCHEDULE_ID + validate: bool = True + + def __post_init__(self) -> None: + if not isinstance(self.backend, str) or not self.backend.strip(): + raise AttentionContractError("Attention backend must be a non-empty string") + if not isinstance(self.deterministic, bool): + raise AttentionContractError("deterministic must be a bool") + if ( + not isinstance(self.communication_backend, str) + or not self.communication_backend.strip() + ): + raise AttentionContractError("communication_backend must be a non-empty string") + for name in ("return_lse", "return_gradients", "validate"): + if not isinstance(getattr(self, name), bool): + raise AttentionContractError(f"{name} must be a bool") + if not isinstance(self.strict_core_id, str) or not self.strict_core_id.strip(): + raise AttentionContractError("strict_core_id must be a non-empty string") + if not isinstance(self.strict_schedule, str) or not self.strict_schedule.strip(): + raise AttentionContractError("strict_schedule must be a non-empty string") + object.__setattr__(self, "backend", self.backend.strip().lower()) + object.__setattr__(self, "communication_backend", self.communication_backend.strip()) + + +@dataclass(frozen=True) +class AttentionAblationResult: + """Standardized Attention result consumed by cross-config artifacts.""" + + out: Tensor + lse: Tensor | None + dq: Tensor | None = None + dk: Tensor | None = None + dv: Tensor | None = None + backend_id: str = BACKEND_ID + deterministic: bool = True + provenance: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.out, Tensor): + raise TypeError("Attention result out must be a torch.Tensor") + if self.lse is not None and not isinstance(self.lse, Tensor): + raise TypeError("Attention result lse must be a torch.Tensor or None") + for name in ("dq", "dk", "dv"): + value = getattr(self, name) + if value is not None and not isinstance(value, Tensor): + raise TypeError(f"Attention result {name} must be a torch.Tensor or None") + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise ValueError("Attention result backend_id must be non-empty") + if not isinstance(self.deterministic, bool): + raise TypeError("Attention result deterministic must be a bool") + if not isinstance(self.provenance, Mapping): + raise TypeError("Attention result provenance must be a mapping") + object.__setattr__(self, "provenance", MappingProxyType(dict(self.provenance))) + + @property + def out_lse(self) -> tuple[Tensor, Tensor | None]: + """Compatibility tuple for callers that consume ``(out, lse)``.""" + + return self.out, self.lse + + def readback(self) -> dict[str, Any]: + """Return JSON-compatible execution evidence for PR230 artifacts.""" + + return { + "backend_id": self.backend_id, + "deterministic": self.deterministic, + "out_shape": list(self.out.shape), + "out_dtype": str(self.out.dtype).replace("torch.", ""), + "lse_shape": None if self.lse is None else list(self.lse.shape), + "lse_dtype": None if self.lse is None else str(self.lse.dtype).replace("torch.", ""), + "gradients": { + "dq": self.dq is not None, + "dk": self.dk is not None, + "dv": self.dv is not None, + }, + "provenance": dict(self.provenance), + } + + +class AttentionAblationOp: + """PR230/PR314-style unified Attention wrapper. + + ``core`` and ``reference`` are injectable so the wrapper is usable by the + semantic operator session without importing CUDA at construction time. + ``core`` should expose ``forward_with_lse``; ``reference`` is the existing + pure-PyTorch CP reference with the same method. + """ + + op_class = "attention" + is_batch_invariant = True + backend_id = BACKEND_ID + + def __init__( + self, + *, + core: Any | None = None, + reference: Any | None = None, + native: Any | None = None, + cp_backend: Any | None = None, + communication_backend: str = "none", + ) -> None: + if not isinstance(communication_backend, str) or not communication_backend.strip(): + raise AttentionContractError("communication_backend must be a non-empty string") + self.core = core + self.reference = reference + self.native = native + # CP production execution is injected by the runtime adapter. Keeping + # it separate from the single-device core prevents an accidental + # fallback to the PyTorch reference when AG/RS is required. + self.cp_backend = cp_backend + self.communication_backend = communication_backend.strip() + + def __call__( + self, + q: Tensor, + k: Tensor, + v: Tensor, + *, + contract: AttentionContract, + config: AttentionAblationConfig | Mapping[str, Any] | None = None, + backend: str | Callable[..., Any] | None = None, + deterministic: bool | None = None, + return_lse: bool | None = None, + return_gradients: bool | None = None, + dout: Tensor | None = None, + communication_backend: str | None = None, + validate: bool | None = None, + **kwargs: Any, + ) -> AttentionAblationResult: + return self.apply( + q, + k, + v, + contract=contract, + config=config, + backend=backend, + deterministic=deterministic, + return_lse=return_lse, + return_gradients=return_gradients, + dout=dout, + communication_backend=communication_backend, + validate=validate, + **kwargs, + ) + + def apply( + self, + q: Tensor, + k: Tensor, + v: Tensor, + *, + contract: AttentionContract, + config: AttentionAblationConfig | Mapping[str, Any] | None = None, + backend: str | Callable[..., Any] | None = None, + deterministic: bool | None = None, + return_lse: bool | None = None, + return_gradients: bool | None = None, + dout: Tensor | None = None, + communication_backend: str | None = None, + validate: bool | None = None, + **kwargs: Any, + ) -> AttentionAblationResult: + if not isinstance(contract, AttentionContract): + raise AttentionContractError("contract must be an AttentionContract") + backend_request = ( + backend + if callable(backend) + or hasattr(backend, "forward_with_lse") + or hasattr(backend, "apply") + else None + ) + cfg = _resolve_config( + config, + backend=backend if isinstance(backend, str) else None, + deterministic=deterministic, + return_lse=return_lse, + return_gradients=return_gradients, + communication_backend=( + communication_backend + if communication_backend is not None + else self.communication_backend + ), + validate=validate, + ) + if cfg.validate: + self._validate_inputs(q, k, v, contract) + if cfg.deterministic and contract.split_kv.mode is SplitKVMode.AUTO: + raise AttentionContractError( + "deterministic Attention cannot use runtime-dependent Split-KV=auto" + ) + if cfg.deterministic and ( + cfg.strict_core_id != STRICT_ATTENTION_CORE_ID + or cfg.strict_schedule != STRICT_ATTENTION_SCHEDULE_ID + ): + raise AttentionContractError( + "strict deterministic Attention requires the canonical core and schedule" + ) + if cfg.deterministic and not cfg.return_lse: + raise AttentionContractError("strict deterministic Attention must return LSE") + if cfg.return_gradients and dout is None: + raise AttentionContractError("dout is required when return_gradients=True") + if dout is not None and dout.shape != q.shape: + raise AttentionContractError("dout must have the same shape as q") + + requested = backend_request if backend_request is not None else cfg.backend + selected, selected_id = self._select_backend( + requested, + q, + contract, + communication_backend=cfg.communication_backend, + ) + if cfg.deterministic and selected_id == "native": + raise AttentionContractError( + "deterministic=True cannot execute an unverified native Attention backend" + ) + selected_core_id = getattr(selected, "core_id", None) + selected_schedule = getattr(selected, "strict_schedule", None) + if cfg.deterministic and selected_id not in {BACKEND_ID, REFERENCE_BACKEND_ID}: + if selected_core_id != cfg.strict_core_id or selected_schedule != cfg.strict_schedule: + raise AttentionContractError( + "deterministic Attention requires the shared strict core and schedule" + ) + + call_kwargs = dict(kwargs) + call_kwargs.setdefault("causal", contract.causal) + call_kwargs.setdefault("scale", 1.0 / math.sqrt(contract.head_dim)) + call_kwargs.setdefault("cp_world_size", contract.sharding.cp_world_size) + if contract.split_kv.mode is SplitKVMode.FIXED: + call_kwargs.setdefault("kv_chunk_size", contract.split_kv.fixed_split_size) + out, lse, backend_provenance = self._invoke(selected, q, k, v, call_kwargs, contract) + if cfg.validate: + self._validate_outputs(out, lse, q, contract) + _validate_runtime_provenance( + selected, + selected_id, + backend_provenance, + cfg, + ) + + dq = dk = dv = None + if cfg.return_gradients: + dq, dk, dv = self._backward(selected, q, k, v, out, dout, call_kwargs) + + provenance = { + "schema_version": "rlkernel.attention.ablation_result.v1", + "semantic_operator": "attention", + "backend_id": selected_id, + "deterministic": cfg.deterministic, + "strict_core_id": (cfg.strict_core_id if cfg.deterministic else None), + "strict_schedule": cfg.strict_schedule if cfg.deterministic else None, + "core_id": cfg.strict_core_id if cfg.deterministic else selected_id, + "backend_deterministic": cfg.deterministic, + "native_attention_arithmetic": False if cfg.deterministic else selected_id == "native", + "communication_backend": cfg.communication_backend, + "communication_executed": bool(getattr(selected, "communication_executed", False)), + "split_kv": contract.split_kv.to_dict(), + "actual_split_kv": _actual_split_provenance( + contract, + total_kv_tokens=k.size(2), + backend=selected_id, + ), + "reduction": _reduction_provenance(contract), + "contract_fingerprint": _contract_fingerprint(contract), + "return_lse": cfg.return_lse, + "return_gradients": cfg.return_gradients, + } + provenance.update(backend_provenance) + provenance.setdefault("actual_backend", selected_id) + provenance.setdefault("communication_backend", cfg.communication_backend) + provenance.setdefault("production_ready", False) + return AttentionAblationResult( + out=out, + lse=lse if cfg.return_lse else None, + dq=dq, + dk=dk, + dv=dv, + backend_id=selected_id, + deterministic=cfg.deterministic, + provenance=provenance, + ) + + def apply_fp32(self, *args: Any, **kwargs: Any) -> AttentionAblationResult: + """Stable fingerprint entry point used by ``OperatorSession``.""" + + return self.apply(*args, **kwargs) + + def _select_backend( + self, + requested: str | Callable[..., Any], + q: Tensor, + contract: AttentionContract, + *, + communication_backend: str, + ) -> tuple[Any, str]: + if ( + callable(requested) + or hasattr(requested, "forward_with_lse") + or hasattr(requested, "apply") + ): + return requested, _callable_backend_id(requested) + normalized = str(requested).strip().lower() + if normalized in {"native", "te", "flashinfer"}: + if self.native is None: + raise AttentionContractError( + "native Attention backend was requested but no native callable was injected" + ) + return self.native, "native" + if normalized in {"reference", "pytorch_reference"}: + return self._reference_backend(), REFERENCE_BACKEND_ID + if normalized not in {"auto", "deterministic", "rlkernel"}: + raise AttentionContractError(f"unsupported Attention backend {requested!r}") + if contract.sharding.cp_world_size > 1: + if communication_backend == "self_owned_cuda_ag_rs": + if self.cp_backend is None: + raise AttentionContractError( + "CP production Attention requires an injected AG/RS backend" + ) + return self.cp_backend, _callable_backend_id(self.cp_backend) + return self._reference_backend(), REFERENCE_BACKEND_ID + if q.device.type != "cuda" or torch.version.hip is not None: + return self._reference_backend(), REFERENCE_BACKEND_ID + if self.core is None: + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionOp, + ) + + self.core = DeterministicAttentionOp() + return self.core, BACKEND_ID + + def _reference_backend(self) -> Any: + if self.reference is None: + from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + DeterministicCPAttentionReferenceOp, + ) + + self.reference = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + return self.reference + + @staticmethod + def _validate_inputs(q: Tensor, k: Tensor, v: Tensor, contract: AttentionContract) -> None: + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise AttentionContractError("q, k, and v must use [B, H, S, D] layout") + expected_dtype = _TORCH_DTYPES[contract.dtype] + if ( + q.dtype is not expected_dtype + or k.dtype is not expected_dtype + or v.dtype is not expected_dtype + ): + raise AttentionContractError( + f"q, k, and v must match contract dtype {contract.dtype.value}" + ) + if q.device != k.device or q.device != v.device: + raise AttentionContractError("q, k, and v must be on the same device") + batch, q_heads, q_seq, dim = q.shape + if batch != contract.batch_size: + raise AttentionContractError( + f"q batch={batch} does not match contract batch_size={contract.batch_size}" + ) + if q_seq != contract.query_sequence_length: + raise AttentionContractError( + "q sequence length does not match AttentionContract query_sequence_length" + ) + sharding = contract.sharding + if q_heads != sharding.local_q_heads or k.shape[1] != sharding.local_kv_heads: + raise AttentionContractError("q/k head counts do not match TP sharding in contract") + if ( + k.shape[0] != batch + or v.shape[:3] != k.shape[:3] + or k.shape[-1] != dim + or v.shape[-1] != dim + ): + raise AttentionContractError("q, k, and v shapes are inconsistent") + if dim != contract.head_dim: + raise AttentionContractError("tensor head_dim does not match AttentionContract") + + @staticmethod + def _validate_outputs(out: Tensor, lse: Tensor, q: Tensor, contract: AttentionContract) -> None: + if out.shape != q.shape: + raise AttentionContractError( + f"Attention output shape {tuple(out.shape)} does not match q {tuple(q.shape)}" + ) + expected_lse = (q.shape[0], q.shape[1], q.shape[2]) + if lse.shape != expected_lse: + raise AttentionContractError( + f"attention-domain LSE shape {tuple(lse.shape)} does not match {expected_lse}" + ) + if lse.dtype is not torch.float32: + raise AttentionContractError("attention-domain LSE must remain fp32") + expected_dtype = _TORCH_DTYPES[contract.dtype] + if out.dtype is not expected_dtype: + raise AttentionContractError( + f"Attention output must be written in {contract.dtype.value}, got {out.dtype}" + ) + + @staticmethod + def _invoke( + backend: Any, + q: Tensor, + k: Tensor, + v: Tensor, + kwargs: Mapping[str, Any], + contract: AttentionContract, + ) -> tuple[Tensor, Tensor, dict[str, Any]]: + method = getattr(backend, "forward_with_lse", None) + if not callable(method): + method = getattr(backend, "apply", None) + if not callable(method): + method = backend if callable(backend) else None + if method is None: + raise AttentionContractError( + "Attention backend must expose forward_with_lse, apply, or __call__" + ) + accepted = _accepted_kwargs(method, kwargs) + if contract.sharding.cp_world_size > 1 and "cp_world_size" not in accepted: + raise AttentionContractError( + "CP>1 requires an Attention backend that explicitly accepts cp_world_size" + ) + result = method(q, k, v, **accepted) + backend_provenance: dict[str, Any] = {} + if isinstance(result, AttentionAblationResult): + out, lse = result.out, result.lse + elif isinstance(result, tuple) and len(result) == 2: + out, lse = result + else: + out = getattr(result, "out", None) + lse = getattr(result, "lse", None) + raw_provenance = getattr(result, "provenance", {}) + if isinstance(raw_provenance, Mapping): + backend_provenance = dict(raw_provenance) + if out is None or lse is None: + raise AttentionContractError( + "Attention backend must return (out, lse), AttentionAblationResult, " + "or an object with out/lse/provenance" + ) + if isinstance(result, AttentionAblationResult): + backend_provenance = dict(result.provenance) + if not isinstance(out, Tensor) or not isinstance(lse, Tensor): + raise AttentionContractError("Attention backend returned non-tensor output or LSE") + return out, lse, backend_provenance + + @staticmethod + def _backward( + backend: Any, + q: Tensor, + k: Tensor, + v: Tensor, + out: Tensor, + dout: Tensor | None, + kwargs: Mapping[str, Any], + ) -> tuple[Tensor, Tensor, Tensor]: + backward = getattr(backend, "backward_reference", None) + if callable(backward): + result = backward(q, k, v, dout, **_accepted_kwargs(backward, kwargs)) + gradients = getattr(result, "gradients", None) + if gradients is not None: + return gradients.dq, gradients.dk, gradients.dv + if not out.requires_grad: + raise AttentionContractError( + "Attention backend did not retain an autograd graph for gradients" + ) + return torch.autograd.grad( + out, + (q, k, v), + grad_outputs=dout.to(dtype=out.dtype), + allow_unused=False, + retain_graph=True, + ) + + +def _resolve_config( + config: AttentionAblationConfig | Mapping[str, Any] | None, + **overrides: Any, +) -> AttentionAblationConfig: + if config is None: + values: dict[str, Any] = {} + elif isinstance(config, AttentionAblationConfig): + values = { + name: getattr(config, name) + for name in ( + "backend", + "deterministic", + "communication_backend", + "return_lse", + "return_gradients", + "strict_core_id", + "strict_schedule", + "validate", + ) + } + elif isinstance(config, Mapping): + values = dict(config) + else: + raise AttentionContractError("config must be AttentionAblationConfig, mapping, or None") + values.update({name: value for name, value in overrides.items() if value is not None}) + return AttentionAblationConfig(**values) + + +def _accepted_kwargs(method: Callable[..., Any], kwargs: Mapping[str, Any]) -> dict[str, Any]: + try: + signature = inspect.signature(method) + except (TypeError, ValueError): + return dict(kwargs) + parameters = signature.parameters.values() + if any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters): + return dict(kwargs) + return {name: value for name, value in kwargs.items() if name in signature.parameters} + + +def _callable_backend_id(value: Any) -> str: + explicit = getattr(value, "backend_id", None) or getattr( + value, "__attention_backend_id__", None + ) + if isinstance(explicit, str) and explicit.strip(): + return explicit.strip() + return f"injected.{type(value).__module__}.{type(value).__qualname__}" + + +def _validate_runtime_provenance( + selected: Any, + selected_id: str, + runtime: Mapping[str, Any], + config: AttentionAblationConfig, +) -> None: + """Fail closed when a production strict backend did not prove its identity.""" + + if not config.deterministic: + return + if selected_id == "native": + raise AttentionContractError( + "deterministic Attention cannot execute native Attention arithmetic" + ) + + actual_core = runtime.get("strict_core_id", getattr(selected, "core_id", None)) + actual_schedule = runtime.get("strict_schedule", getattr(selected, "strict_schedule", None)) + if selected_id != REFERENCE_BACKEND_ID and ( + actual_core != config.strict_core_id or actual_schedule != config.strict_schedule + ): + raise AttentionContractError( + "deterministic Attention backend did not prove the shared strict core and schedule" + ) + + if config.communication_backend != "self_owned_cuda_ag_rs": + return + + expected = { + "strict_core_id": config.strict_core_id, + "strict_schedule": config.strict_schedule, + "actual_backend": "rlkernel.cuda.deterministic_attention", + "communication_backend": "self_owned_cuda_ag_rs", + "production_ready": True, + "native_attention_arithmetic": False, + "fallback": False, + } + mismatches = [name for name, value in expected.items() if runtime.get(name) != value] + if mismatches: + raise AttentionContractError( + "CP production Attention runtime provenance is incomplete or mismatched: " + + ", ".join(mismatches) + ) + + +def _actual_split_provenance( + contract: AttentionContract, + *, + total_kv_tokens: int, + backend: str = BACKEND_ID, +) -> dict[str, Any]: + plan = contract.split_kv.resolve( + total_kv_tokens, + backend=backend, + ) + return plan.to_dict() + + +def _reduction_provenance(contract: AttentionContract) -> dict[str, Any]: + reduction = contract.reduction + return { + "merge": reduction.merge.value, + "acc_dtype": reduction.acc_dtype.value, + "order": reduction.order.value, + "downcast_at": reduction.downcast_at.value, + "engine": reduction.engine.value, + } + + +def _contract_fingerprint(contract: AttentionContract) -> str: + payload = json.dumps(contract.to_dict(), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +__all__ = [ + "AttentionAblationConfig", + "AttentionAblationOp", + "AttentionAblationResult", + "BACKEND_ID", + "REFERENCE_BACKEND_ID", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py new file mode 100644 index 00000000..194c80d5 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -0,0 +1,1182 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic context-parallel attention reference. + +This module is the correctness-first WS2 reference for CP-aware standard +softmax attention. It intentionally stays in PyTorch and uses fp32 partial +states so fused CUDA/Triton backends can validate their CP/LSE merge semantics +against a small, inspectable implementation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional, Sequence + +import torch + +from rl_engine.kernels.attention_contract import ( + SplitKVExecutionPlan, + SplitKVMode, + SplitKVRuntimeCoordinate, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + +@dataclass(frozen=True) +class AttentionPartialState: + """One KV block's attention state before deterministic LSE merge. + + ``out`` is already normalized within the local KV block and has shape + ``[B, Hq, Sq, D]``. ``lse`` is the local attention-domain log-sum-exp with + shape ``[B, Hq, Sq]``. ``block_start`` / ``block_end`` are logical global KV + positions and define the canonical merge order. + """ + + out: torch.Tensor + lse: torch.Tensor + block_start: int + block_end: int + + def __post_init__(self) -> None: + if self.out.ndim != 4: + raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") + if self.lse.shape != self.out.shape[:3]: + raise ValueError("partial attention lse must have shape [B, Hq, Sq]") + if self.out.device != self.lse.device: + raise ValueError("partial attention out/lse must be on the same device") + if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: + raise ValueError("partial attention out/lse must remain FP32 before merge") + if self.block_start < 0: + raise ValueError("block_start must be non-negative") + if self.block_end < self.block_start: + raise ValueError("block_end must be >= block_start") + + +@dataclass(frozen=True) +class AttentionBackwardGradients: + """Training-side gradients emitted by the CP attention backward reference.""" + + dq: torch.Tensor + dk: torch.Tensor + dv: torch.Tensor + + +@dataclass(frozen=True) +class AttentionBackwardPathResult: + """One materialized CP attention backward path.""" + + name: str + out: torch.Tensor + lse: torch.Tensor + gradients: AttentionBackwardGradients + provenance: dict[str, object] + + +@dataclass(frozen=True) +class GradientDriftStats: + """Shape-aware absolute drift summary for backward validation reports.""" + + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + def to_dict(self) -> dict[str, object]: + return { + "max_abs": self.max_abs, + "mean_abs": self.mean_abs, + "p95_abs": self.p95_abs, + "p99_abs": self.p99_abs, + "active_count": self.active_count, + } + + +@dataclass(frozen=True) +class AttentionBackwardRankDrift: + """Backward drift for one logical CP rank's sequence ownership.""" + + rank: int + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + + def to_dict(self) -> dict[str, object]: + return { + "rank": self.rank, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + } + + +@dataclass(frozen=True) +class AttentionBackwardPathDrift: + """Candidate-vs-reference backward drift for one CP path.""" + + candidate_name: str + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + out: GradientDriftStats + lse: GradientDriftStats + per_rank: tuple[AttentionBackwardRankDrift, ...] + provenance: dict[str, object] + + def to_dict(self) -> dict[str, object]: + return { + "candidate_name": self.candidate_name, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + "out": self.out.to_dict(), + "lse": self.lse.to_dict(), + "per_rank": [item.to_dict() for item in self.per_rank], + "provenance": self.provenance, + } + + +@dataclass(frozen=True) +class AttentionBackwardComparisonReport: + """Structured PR8 report for CP attention gradient drift validation.""" + + reference_name: str + drifts: tuple[AttentionBackwardPathDrift, ...] + + def to_dict(self) -> dict[str, object]: + return { + "reference_name": self.reference_name, + "drifts": [drift.to_dict() for drift in self.drifts], + } + + +def merge_attention_partial_states( + states: Sequence[AttentionPartialState], +) -> AttentionPartialState: + """Merge CP/chunk partial states in logical block order. + + The merge is the online-softmax/LSE merge used by attention, not a plain + sum. The input order is deliberately ignored: states are sorted by logical + ``block_start`` so the result depends on global block indices rather than + arrival order. + """ + + if not states: + raise ValueError("at least one attention partial state is required") + + ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) + _validate_merge_shapes_and_ranges(ordered) + + merged = ordered[0] + merged_out = merged.out.float() + merged_lse = merged.lse.float() + for state in ordered[1:]: + merged_out, merged_lse = _merge_two_states( + merged_out, + merged_lse, + state.out.float(), + state.lse.float(), + ) + + return AttentionPartialState( + out=merged_out, + lse=merged_lse, + block_start=ordered[0].block_start, + block_end=ordered[-1].block_end, + ) + + +class DeterministicCPAttentionReferenceOp: + """Correctness-first CP attention reference for prefill and chunked prefill. + + The reference consumes attention-ready Q/K. For Qwen3 WS2 this means Q/K + have already passed QK-Norm and RoPE unless an outer contract explicitly + marks them as pre-RoPE. RoPE is intentionally kept outside this CP merge + implementation so fused and unfused ``RoPE+Attention`` paths can compare the + same post-RoPE Q/K boundary before validating CP communication. + + The op emulates CP by splitting query and KV sequence dimensions into + logical CP shards. Each query shard computes one partial attention state per + KV block, then merges those states in fixed global-block order using fp32 + LSE arithmetic. ``forward`` returns the input dtype after the final write; + ``forward_fp32`` keeps the fp32 merged output. + """ + + op_class = "attention" + + def __init__(self, *, strict_bitwise: bool = False) -> None: + if not isinstance(strict_bitwise, bool): + raise TypeError("strict_bitwise must be a bool") + self.strict_bitwise = strict_bitwise + + @staticmethod + def split_kv_execution_plans( + total_kv_tokens: int, + *, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> list[dict[str, object]]: + """Export the actual logical Split-KV plan before execution.""" + + return split_kv_execution_plan_provenance( + total_kv_tokens, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ) + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + return self.forward( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and final input-dtype write.""" + + out, _ = self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=q.dtype, + ) + return out + + def forward_fp32( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and fp32 output.""" + + out, _ = self.forward_fp32_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(out, lse)`` for the CP reference path. + + ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 + until the final write, then downcast to ``output_dtype``. When omitted, + ``output_dtype`` defaults to the input dtype. + """ + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) + if self.strict_bitwise: + out, lse = self._forward_strict_bitwise( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + kv_chunk_size=kv_chunk_size, + ) + else: + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + out = out.to(resolved_output_dtype) + return out, lse + + def forward_fp32_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return fp32 ``(out, lse)`` for the CP reference path.""" + + return self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=torch.float32, + ) + + def _forward_strict_bitwise( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Execute one batch/CP-independent arithmetic schedule.""" + + _validate_qkv(q, k, v) + _validate_scale(scale) + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + kv_bounds = _kv_block_bounds(skv, 1, kv_chunk_size) + out_rows: list[torch.Tensor] = [] + lse_rows: list[torch.Tensor] = [] + for batch_index in range(batch): + q_batch = q[batch_index : batch_index + 1].contiguous() + k_batch = k[batch_index : batch_index + 1].contiguous() + v_batch = v[batch_index : batch_index + 1].contiguous() + pad_batch = ( + None + if key_padding_mask is None + else key_padding_mask[batch_index : batch_index + 1].contiguous() + ) + query_offset = query_offsets[batch_index : batch_index + 1] + key_offset = key_offsets[batch_index : batch_index + 1] + query_rows: list[torch.Tensor] = [] + lse_query_rows: list[torch.Tensor] = [] + for query_index in range(sq): + q_row = q_batch[:, :, query_index : query_index + 1, :].contiguous() + states = [ + self.local_partial_state( + q_row, + k_batch[:, :, key_start:key_end, :].contiguous(), + v_batch[:, :, key_start:key_end, :].contiguous(), + q_start=query_index, + k_start=key_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None + if pad_batch is None + else pad_batch[:, key_start:key_end].contiguous() + ), + query_position_offsets=query_offset, + key_position_offsets=key_offset, + ) + for key_start, key_end in kv_bounds + if key_start != key_end + ] + merged = merge_attention_partial_states(states) + query_rows.append(merged.out) + lse_query_rows.append(merged.lse) + out_rows.append(torch.cat(query_rows, dim=2)) + lse_rows.append(torch.cat(lse_query_rows, dim=2)) + return torch.cat(out_rows, dim=0), torch.cat(lse_rows, dim=0) + + def backward_reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, + name: Optional[str] = None, + ) -> AttentionBackwardPathResult: + """Run the deterministic training-side backward validation path. + + The semantic backward input is ``dout`` plus the forward attention state + produced from the same Q/K/V, masks, position offsets, CP world, and KV + block order. The reference keeps the softmax/merge math in fp32 and + records the final-write dtype in provenance; decode backward is + intentionally out of scope for PR8. + """ + + _validate_qkv(q, k, v) + if dout.shape != q.shape: + raise ValueError("dout must have shape [B, Hq, Sq, D], matching q") + if not torch.is_floating_point(dout) or torch.is_complex(dout): + raise ValueError("dout must be a real floating-point tensor") + if dout.device != q.device: + raise ValueError("dout must be on the same device as q, k, and v") + if dout.dtype != q.dtype: + raise ValueError("dout must have the same dtype as q") + q_leaf = q.detach().clone().requires_grad_(True) + k_leaf = k.detach().clone().requires_grad_(True) + v_leaf = v.detach().clone().requires_grad_(True) + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) + out, lse = self.forward_with_lse( + q_leaf, + k_leaf, + v_leaf, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=resolved_output_dtype, + ) + torch.autograd.backward(out, dout.to(dtype=out.dtype)) + if q_leaf.grad is None or k_leaf.grad is None or v_leaf.grad is None: + raise RuntimeError("CP attention backward did not produce dq/dk/dv") + + return AttentionBackwardPathResult( + name=name + or _backward_path_name(cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size), + out=out.detach(), + lse=lse.detach(), + gradients=AttentionBackwardGradients( + dq=q_leaf.grad.detach(), + dk=k_leaf.grad.detach(), + dv=v_leaf.grad.detach(), + ), + provenance={ + "attention_mode": "prefill" if kv_chunk_size is None else "chunked_prefill", + "gradient_mode": "training_backward", + "gradient_inputs": ["q", "k", "v"], + "gradient_outputs": ["out"], + "saved_forward_state": [ + "out", + "attention_lse", + "causal_mask", + "key_padding_mask", + "query_position_offsets", + "key_position_offsets", + "global_block_index", + ], + "cp_world_size": cp_world_size, + "kv_chunk_size": kv_chunk_size, + "requested_split_kv_policy": ("disabled" if kv_chunk_size is None else "fixed"), + "requested_split_kv_size": kv_chunk_size, + "actual_split_kv_plans": split_kv_execution_plan_provenance( + k.size(2), + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_backward_reference", + ), + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "output_dtype": str(resolved_output_dtype).replace("torch.", ""), + "q_dtype": str(q.dtype).replace("torch.", ""), + "k_dtype": str(k.dtype).replace("torch.", ""), + "v_dtype": str(v.dtype).replace("torch.", ""), + "dout_dtype": str(dout.dtype).replace("torch.", ""), + "te_backward_oracle": "not_used", + "decode_backward": "not_supported", + }, + ) + + def local_partial_state( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + q_start: int, + k_start: int, + total_kv_len: int, + total_query_len: Optional[int] = None, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + ) -> AttentionPartialState: + """Compute one query shard against one logical KV block. + + ``query_position_offsets`` and ``key_position_offsets`` are optional + per-batch-row base positions. They let the reference express varlen or + packed metadata while retaining the dense [B, H, S, D] tensor layout. + For post-RoPE Q/K, these offsets must describe the same absolute token + positions used when RoPE was applied. + """ + + _validate_qkv(q, k, v) + _validate_scale(scale) + if q_start < 0 or k_start < 0: + raise ValueError("q_start and k_start must be non-negative") + if total_kv_len < k_start + k.size(2): + raise ValueError("total_kv_len must cover the local KV block") + if total_query_len is None: + total_query_len = q.size(2) + if total_query_len < q_start + q.size(2): + raise ValueError("total_query_len must cover the local query block") + if key_padding_mask is not None: + if key_padding_mask.shape != (q.size(0), k.size(2)): + raise ValueError("local key_padding_mask must have shape [B, local_skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("local key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + q.size(0), + q.device, + default=total_kv_len - total_query_len, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + q.size(0), + q.device, + default=0, + name="key_position_offsets", + ) + + ctx = NativeAttentionOp._strict_fp32_math(q.device.type) + with ctx: + qf = q.float() + kf = k.float() + vf = v.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hkv != hq: + repeat = hq // hkv + kf = kf.repeat_interleave(repeat, dim=1) + vf = vf.repeat_interleave(repeat, dim=1) + + if skv == 0: + zero_dep = _zero_dependency(qf, kf, vf) + return AttentionPartialState( + out=torch.zeros(q.size(0), hq, sq, dim, device=q.device, dtype=torch.float32) + + zero_dep, + lse=torch.full( + (q.size(0), hq, sq), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep, + block_start=k_start, + block_end=k_start, + ) + + scale_value = scale if scale is not None else (1.0 / math.sqrt(dim)) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * scale_value + if causal: + query_base = query_offsets[:, None] + q_start + key_base = key_offsets[:, None] + k_start + q_pos = torch.arange(sq, device=q.device, dtype=torch.long) + query_base + k_pos = torch.arange(skv, device=q.device, dtype=torch.long) + key_base + causal_mask = k_pos[:, None, :] > q_pos[:, :, None] + scores = scores.masked_fill(causal_mask[:, None, :, :], float("-inf")) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + + lse = torch.logsumexp(scores, dim=-1) + finite_lse = torch.isfinite(lse) + weights = torch.exp(scores - lse.unsqueeze(-1)) + weights = torch.where(finite_lse.unsqueeze(-1), weights, torch.zeros_like(weights)) + out = torch.matmul(weights, vf) + return AttentionPartialState( + out=out, + lse=lse, + block_start=k_start, + block_end=k_start + skv, + ) + + def _forward_impl( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + cp_world_size: int, + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + _validate_qkv(q, k, v) + _validate_scale(scale) + if ( + isinstance(cp_world_size, bool) + or not isinstance(cp_world_size, int) + or cp_world_size < 1 + ): + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and ( + isinstance(kv_chunk_size, bool) + or not isinstance(kv_chunk_size, int) + or kv_chunk_size < 1 + ): + raise ValueError("kv_chunk_size must be >= 1 when provided") + + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + + q_bounds = _split_bounds(sq, cp_world_size) + kv_bounds = _kv_block_bounds(skv, cp_world_size, kv_chunk_size) + out_chunks: list[torch.Tensor] = [] + lse_chunks: list[torch.Tensor] = [] + for q_start, q_end in q_bounds: + if q_start == q_end: + continue + q_block = q[:, :, q_start:q_end, :] + states = [ + self.local_partial_state( + q_block, + k[:, :, k_start:k_end, :], + v[:, :, k_start:k_end, :], + q_start=q_start, + k_start=k_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None if key_padding_mask is None else key_padding_mask[:, k_start:k_end] + ), + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + ) + for k_start, k_end in kv_bounds + if k_start != k_end + ] + if states: + merged = merge_attention_partial_states(states) + out_chunks.append(merged.out) + lse_chunks.append(merged.lse) + else: + zero_dep = _zero_dependency(q_block.float(), k.float(), v.float()) + out_chunks.append( + torch.zeros(batch, hq, q_end - q_start, dim, device=q.device) + zero_dep + ) + lse_chunks.append( + torch.full( + (batch, hq, q_end - q_start), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep + ) + + if not out_chunks: + zero_dep = _zero_dependency(q.float(), k.float(), v.float()) + return ( + torch.empty(batch, hq, 0, dim, device=q.device, dtype=torch.float32) + zero_dep, + torch.empty(batch, hq, 0, device=q.device, dtype=torch.float32) + zero_dep, + ) + return torch.cat(out_chunks, dim=2), torch.cat(lse_chunks, dim=2) + + +def compare_cp_attention_backward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + candidate_cp_world_size: int = 2, + candidate_kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, +) -> AttentionBackwardComparisonReport: + """Compare CP=1 backward with a CP/chunked-prefill candidate. + + The report includes whole-tensor ``dq/dk/dv`` drift and per-logical-CP-rank + slices. It is a validation/reporting helper, not a separate production + backward kernel. + """ + + op = DeterministicCPAttentionReferenceOp() + reference = op.backward_reference( + q, + k, + v, + dout, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=1, + kv_chunk_size=None, + output_dtype=output_dtype, + name="cp1_backward_reference", + ) + candidate = op.backward_reference( + q, + k, + v, + dout, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=candidate_cp_world_size, + kv_chunk_size=candidate_kv_chunk_size, + output_dtype=output_dtype, + ) + return AttentionBackwardComparisonReport( + reference_name=reference.name, + drifts=(_compare_backward_path(candidate, reference),), + ) + + +def _compare_backward_path( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, +) -> AttentionBackwardPathDrift: + cp_world_size = _provenance_int(candidate.provenance, "cp_world_size") + return AttentionBackwardPathDrift( + candidate_name=candidate.name, + dq=_drift_stats(candidate.gradients.dq, reference.gradients.dq), + dk=_drift_stats(candidate.gradients.dk, reference.gradients.dk), + dv=_drift_stats(candidate.gradients.dv, reference.gradients.dv), + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + per_rank=_per_rank_backward_drifts(candidate, reference, cp_world_size), + provenance=candidate.provenance, + ) + + +def _per_rank_backward_drifts( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, + cp_world_size: int, +) -> tuple[AttentionBackwardRankDrift, ...]: + q_bounds = _split_bounds(candidate.gradients.dq.size(2), cp_world_size) + kv_bounds = _split_bounds(candidate.gradients.dk.size(2), cp_world_size) + per_rank = [] + for rank, ((q_start, q_end), (kv_start, kv_end)) in enumerate(zip(q_bounds, kv_bounds)): + per_rank.append( + AttentionBackwardRankDrift( + rank=rank, + dq=_drift_stats( + candidate.gradients.dq[:, :, q_start:q_end, :], + reference.gradients.dq[:, :, q_start:q_end, :], + ), + dk=_drift_stats( + candidate.gradients.dk[:, :, kv_start:kv_end, :], + reference.gradients.dk[:, :, kv_start:kv_end, :], + ), + dv=_drift_stats( + candidate.gradients.dv[:, :, kv_start:kv_end, :], + reference.gradients.dv[:, :, kv_start:kv_end, :], + ), + ) + ) + return tuple(per_rank) + + +def _drift_stats(candidate: torch.Tensor, reference: torch.Tensor) -> GradientDriftStats: + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match " + f"reference shape {tuple(reference.shape)}" + ) + diff = (candidate.float() - reference.float()).abs().reshape(-1) + active_count = int(diff.numel()) + if active_count == 0: + return GradientDriftStats(0.0, 0.0, 0.0, 0.0, 0) + return GradientDriftStats( + max_abs=float(diff.max().item()), + mean_abs=float(diff.mean().item()), + p95_abs=float(torch.quantile(diff, 0.95).item()), + p99_abs=float(torch.quantile(diff, 0.99).item()), + active_count=active_count, + ) + + +def _backward_path_name(*, cp_world_size: int, kv_chunk_size: Optional[int]) -> str: + prefix = f"cp{cp_world_size}" + if kv_chunk_size is None: + return f"{prefix}_backward" + return f"{prefix}_chunked_backward" + + +def _provenance_int(provenance: dict[str, object], key: str) -> int: + value = provenance[key] + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"provenance field {key!r} must be an int") + return value + + +def _merge_two_states( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + merged_lse = torch.logaddexp(lse_a, lse_b) + finite = torch.isfinite(merged_lse) + weight_a = torch.where(finite, torch.exp(lse_a - merged_lse), torch.zeros_like(merged_lse)) + weight_b = torch.where(finite, torch.exp(lse_b - merged_lse), torch.zeros_like(merged_lse)) + merged_out = weight_a.unsqueeze(-1) * out_a + weight_b.unsqueeze(-1) * out_b + return merged_out, merged_lse + + +def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) -> None: + first = states[0] + previous_end = first.block_end + for state in states[1:]: + if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: + raise ValueError("all partial states must have matching out/lse shapes") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be gap-free and non-overlapping") + previous_end = state.block_end + + +def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q, k, and v must have shape [B, H, S, D]") + if k.shape != v.shape: + raise ValueError("k and v must have the same shape") + if q.size(0) != k.size(0) or q.size(3) != k.size(3): + raise ValueError("q, k, and v must share batch size and head dim") + if q.size(1) < 1 or k.size(1) < 1 or q.size(3) < 1: + raise ValueError("q, k, and v must have positive head counts and head dim") + if not all(torch.is_floating_point(tensor) for tensor in (q, k, v)) or any( + torch.is_complex(tensor) for tensor in (q, k, v) + ): + raise ValueError("q, k, and v must be real floating-point tensors") + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError("q, k, and v must have the same dtype") + if q.device != k.device or q.device != v.device: + raise ValueError("q, k, and v must be on the same device") + if q.size(1) % k.size(1) != 0: + raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") + + +def _validate_scale(scale: Optional[float]) -> None: + if scale is None: + return + if isinstance(scale, bool) or not isinstance(scale, (int, float)): + raise ValueError("scale must be a positive finite number") + if not math.isfinite(float(scale)) or float(scale) <= 0: + raise ValueError("scale must be a positive finite number") + + +def _validate_output_dtype(output_dtype: torch.dtype) -> None: + if not isinstance(output_dtype, torch.dtype): + raise ValueError("output_dtype must be a real floating-point torch dtype") + probe = torch.empty((), dtype=output_dtype) + if not torch.is_floating_point(probe) or torch.is_complex(probe): + raise ValueError("output_dtype must be a real floating-point torch dtype") + + +def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: + total = torch.tensor(0.0, device=tensors[0].device) + for tensor in tensors: + total = total + tensor.sum() + return total * 0.0 + + +def _normalize_position_offsets( + offsets: Optional[torch.Tensor], + batch: int, + device: torch.device, + *, + default: int, + name: str, +) -> torch.Tensor: + if offsets is None: + return torch.full((batch,), default, dtype=torch.long, device=device) + if offsets.ndim != 1 or offsets.numel() != batch: + raise ValueError(f"{name} must have shape [B]") + if torch.is_floating_point(offsets) or torch.is_complex(offsets) or offsets.dtype == torch.bool: + raise ValueError(f"{name} must contain integer positions") + return offsets.to(device=device, dtype=torch.long) + + +def _split_bounds(length: int, parts: int) -> list[tuple[int, int]]: + base, extra = divmod(length, parts) + bounds: list[tuple[int, int]] = [] + start = 0 + for index in range(parts): + width = base + (1 if index < extra else 0) + end = start + width + bounds.append((start, end)) + start = end + return bounds + + +def _kv_block_bounds( + length: int, + cp_world_size: int, + kv_chunk_size: Optional[int], +) -> list[tuple[int, int]]: + bounds: list[tuple[int, int]] = [] + for start, end in _split_bounds(length, cp_world_size): + if kv_chunk_size is None: + bounds.append((start, end)) + continue + cursor = start + while cursor < end: + chunk_end = min(cursor + kv_chunk_size, end) + bounds.append((cursor, chunk_end)) + cursor = chunk_end + return bounds + + +def split_kv_execution_plan_provenance( + length: int, + *, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str, +) -> list[dict[str, object]]: + """Return the actual backend-local Split-KV plan for every CP owner.""" + + if length < 1: + raise ValueError("Split-KV sequence length must be >= 1") + if cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + result: list[dict[str, object]] = [] + for owner_cp_rank, (rank_start, rank_end) in enumerate(_split_bounds(length, cp_world_size)): + if rank_start == rank_end: + continue + boundaries: tuple[tuple[int, int], ...] + if kv_chunk_size is None: + boundaries = ((rank_start, rank_end),) + mode = SplitKVMode.DISABLED + else: + boundaries = tuple( + (start, min(start + kv_chunk_size, rank_end)) + for start in range(rank_start, rank_end, kv_chunk_size) + ) + mode = SplitKVMode.FIXED + plan = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + result.append({"owner_cp_rank": owner_cp_rank, **plan.to_dict()}) + return result + + +def build_reference_split_kv_runtime_plan_set( + total_kv_tokens: Sequence[int], + *, + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str = "deterministic_cp_reference", +) -> SplitKVRuntimePlanSet: + """Build complete per-batch/TP/CP/owner plans for the reference path.""" + + totals = tuple(total_kv_tokens) + if not totals or any(total < cp_world_size for total in totals): + raise ValueError("reference runtime plan sets require at least one KV token per CP owner") + if tp_world_size < 1 or cp_world_size < 1: + raise ValueError("TP and CP world sizes must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + + entries: list[SplitKVRuntimePlanEntry] = [] + for batch_index, total in enumerate(totals): + owner_ranges = _split_bounds(total, cp_world_size) + for tp_rank in range(tp_world_size): + for cp_rank in range(cp_world_size): + for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + boundaries: tuple[tuple[int, int], ...] + if kv_chunk_size is None: + mode = SplitKVMode.DISABLED + boundaries = ((owner_start, owner_end),) + else: + mode = SplitKVMode.FIXED + boundaries = tuple( + (start, min(start + kv_chunk_size, owner_end)) + for start in range(owner_start, owner_end, kv_chunk_size) + ) + execution = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + entries.append( + SplitKVRuntimePlanEntry( + coordinate=SplitKVRuntimeCoordinate( + batch_index=batch_index, + tp_rank=tp_rank, + cp_rank=cp_rank, + owner_cp_rank=owner_cp_rank, + ), + expected_kv_range=(owner_start, owner_end), + execution=execution, + ) + ) + return SplitKVRuntimePlanSet( + batch_size=len(totals), + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals, + entries=tuple(entries), + ) + + +CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp + +__all__ = [ + "AttentionBackwardComparisonReport", + "AttentionBackwardGradients", + "AttentionBackwardPathDrift", + "AttentionBackwardPathResult", + "AttentionBackwardRankDrift", + "AttentionPartialState", + "build_reference_split_kv_runtime_plan_set", + "CPAttentionReferenceOp", + "DeterministicCPAttentionReferenceOp", + "GradientDriftStats", + "compare_cp_attention_backward", + "merge_attention_partial_states", + "split_kv_execution_plan_provenance", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 2d5c470f..e9a7d075 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -8,6 +8,15 @@ from enum import Enum, EnumMeta from typing import Any, Dict, Optional, Set, Type +from rl_engine.kernels.attention_contract import ( + AttentionBackendCapability, + AttentionContract, + AttentionContractError, + AttentionDispatchResult, + AttentionDType, + AttentionMode, + AttentionRole, +) from rl_engine.kernels.semantic_registry import ( OperatorBackendDescriptor, OperatorFallbackPolicy, @@ -107,6 +116,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_KV_CACHE_ATTN = ( "rl_engine.kernels.ops.pytorch.attention.kv_cache.NativeKVCacheAttnOp" ) + PYTORCH_CP_ATTENTION = ( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ) # WS1 pure-PyTorch ground-truth linear ops PYTORCH_NATIVE_LM_HEAD = "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp" # WS1 pure-PyTorch ground-truth embedding ops @@ -161,6 +174,45 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, version_or_build_fingerprint="runtime-native-unresolved-v1", ), + OperatorBackendDescriptor( + semantic_op="attention", + backend_id="rlkernel.attention.deterministic.v1", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu", "cuda", "rocm"}), + supported_dtypes=frozenset({"float32", "bfloat16", "float16"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "algorithm": "standard_softmax_attention", + "batch_invariant": True, + "deterministic": True, + "split_kv": "contract_bound", + "reduction_order": "global_block_index", + "strict_schedule": "single_batch_single_query_global_kv_blocks", + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=( + "rl_engine.kernels.ops.pytorch.attention.ablation.AttentionAblationOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="AttentionAblationOp-bitwise-v2", + ), + OperatorBackendDescriptor( + semantic_op="attention", + backend_id="native", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"*"}), + supported_dtypes=frozenset({"*"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "selection": "runtime_native", + "strict_observable": False, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=None, + fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, + version_or_build_fingerprint="runtime-native-attention-unresolved-v1", + ), ) @@ -219,6 +271,54 @@ def __init__(self): self._failed_backends: Set[str] = set() self.semantic = SemanticOperatorCatalog(_default_semantic_descriptors()) + common_roles = frozenset({AttentionRole.TRAIN, AttentionRole.INFER}) + common_dtypes = frozenset({AttentionDType.BF16, AttentionDType.FP16, AttentionDType.FP32}) + self._attention_capabilities = { + OpBackend.PYTORCH_NATIVE_ATTENTION: AttentionBackendCapability( + backend_id="pytorch-native-attention-ws1", + roles=common_roles, + modes=frozenset({AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL}), + dtypes=common_dtypes, + cp_world_sizes=(1,), + exports_attention_lse=False, + deterministic_cp_merge=False, + supports_packed_varlen=False, + supports_kv_cache=False, + implementation_kind="reference", + ), + OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN: AttentionBackendCapability( + backend_id="pytorch-native-kv-cache-attention-ws1", + roles=common_roles, + modes=frozenset({AttentionMode.DECODE}), + dtypes=common_dtypes, + cp_world_sizes=(1,), + exports_attention_lse=False, + deterministic_cp_merge=False, + supports_packed_varlen=False, + supports_kv_cache=False, + implementation_kind="reference", + ), + OpBackend.PYTORCH_CP_ATTENTION: AttentionBackendCapability( + backend_id="pytorch-deterministic-cp-attention-reference", + roles=common_roles, + modes=frozenset({AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL}), + dtypes=common_dtypes, + tp_world_sizes=(1, 2), + cp_world_sizes=(1, 2), + exports_attention_lse=True, + deterministic_cp_merge=True, + supports_packed_varlen=False, + supports_kv_cache=False, + supports_rope_metadata=False, + supports_fused_rope_attention=False, + supports_split_kv_disabled=True, + supports_split_kv_fixed=True, + supports_split_kv_auto=False, + reports_actual_split_kv_plan=True, + implementation_kind="deterministic", + ), + } + self._priority_map = { "cuda": { "logp": [ @@ -252,6 +352,12 @@ def __init__(self): OpBackend.CUDA_DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION, ], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], + "ws2_attention": [ + OpBackend.PYTORCH_CP_ATTENTION, + OpBackend.CUDA_DETERMINISTIC_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "linear_logp": [ @@ -291,6 +397,11 @@ def __init__(self): OpBackend.TRITON_GENERIC, ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], + "ws2_attention": [ + OpBackend.PYTORCH_CP_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], @@ -315,6 +426,11 @@ def __init__(self): "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], + "ws2_attention": [ + OpBackend.PYTORCH_CP_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], @@ -432,6 +548,114 @@ def get_op(self, op_type: str) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + def get_attention_op( + self, + contract: AttentionContract, + *, + requested_backend: str = "deterministic", + ) -> AttentionDispatchResult: + """Resolve only a backend that explicitly supports the WS2 contract.""" + + if not isinstance(contract, AttentionContract): + raise AttentionContractError("contract must be an AttentionContract") + if not isinstance(requested_backend, str) or not requested_backend.strip(): + raise AttentionContractError("requested_backend must be a non-empty string") + requested_backend = requested_backend.strip().lower() + + platform = self._platform() + candidates = self._priority_map.get(platform, {}).get("ws2_attention", []) + rejected: list[str] = [] + + for backend in candidates: + capability = self._attention_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no AttentionBackendCapability declared") + continue + incompatibilities = list(capability.incompatibilities(contract)) + policy_mismatch = self._attention_policy_mismatch(requested_backend, capability) + if policy_mismatch is not None: + incompatibilities.append(policy_mismatch) + if incompatibilities: + rejected.append(f"{backend.name}: " + "; ".join(incompatibilities)) + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + continue + + return AttentionDispatchResult( + op=op, + capability=capability, + provenance={ + "requested_backend": requested_backend, + "actual_backend": capability.backend_id, + "backend_enum": backend.name, + "platform": platform, + "fallback": bool(rejected), + "prior_rejections": list(rejected), + "contract": contract.to_dict(), + "capability": capability.to_dict(), + }, + ) + + details = " | ".join(rejected) if rejected else "no candidates registered" + requested = contract.to_dict() + raise RuntimeError( + "No attention backend supports the requested WS2 contract on " + f"{platform}: role={requested['role']}, mode={requested['mode']}, " + f"dtype={requested['dtype']}, TP={contract.sharding.tp_world_size}, " + f"CP={contract.sharding.cp_world_size}. Rejections: {details}" + ) + + @staticmethod + def _attention_policy_mismatch( + requested_backend: str, + capability: AttentionBackendCapability, + ) -> Optional[str]: + if requested_backend == "auto": + return None + if requested_backend in {"production", "reference", "deterministic"}: + if capability.implementation_kind == requested_backend: + return None + return ( + f"implementation_kind={capability.implementation_kind} does not satisfy " + f"requested_backend={requested_backend}" + ) + if capability.backend_id == requested_backend: + return None + return ( + f"backend_id={capability.backend_id} does not match " + f"requested_backend={requested_backend}" + ) + + @staticmethod + def _platform() -> str: + if device_ctx.is_rocm: + return "rocm" + if device_ctx.device_type == "cuda": + return "cuda" + return "cpu" + + def _get_or_create_backend(self, backend: OpBackend) -> Optional[Any]: + if backend.name in self._instance_cache: + return self._instance_cache[backend.name] + if backend.name in self._failed_backends: + return None + + op_class = self._load_backend(backend) + if op_class is None: + self._failed_backends.add(backend.name) + return None + try: + op = op_class() + except Exception as exc: + logger.error(f"Failed to instantiate {backend.name}: {exc}") + self._failed_backends.add(backend.name) + return None + self._instance_cache[backend.name] = op + return op + def _load_backend(self, backend: OpBackend) -> Optional[Type]: """Import a legacy backend and distinguish wrapper bugs from absence.""" diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py new file mode 100644 index 00000000..70dbd2ca --- /dev/null +++ b/tests/test_attention_ablation.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMode, + AttentionRole, + ReductionSpec, + ShardingSpec, + SplitKVSpec, +) +from rl_engine.kernels.ops.pytorch.attention.ablation import ( + BACKEND_ID, + REFERENCE_BACKEND_ID, + AttentionAblationOp, +) + + +def _contract(*, split_kv: SplitKVSpec | None = None) -> AttentionContract: + sharding = ShardingSpec( + tp_rank=0, + tp_world_size=1, + cp_rank=0, + cp_world_size=1, + global_q_heads=2, + global_kv_heads=1, + local_q_head_start=0, + local_q_heads=2, + local_kv_head_start=0, + local_kv_heads=1, + global_sequence_length=4, + local_sequence_length=4, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 4), + ) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=4, + head_dim=4, + causal=True, + causal_offsets=(0,), + sharding=sharding, + reduction=ReductionSpec(), + split_kv=split_kv or SplitKVSpec.disabled(), + ) + + +def _qkv(): + torch.manual_seed(0) + return ( + torch.randn(1, 2, 4, 4, dtype=torch.bfloat16), + torch.randn(1, 1, 4, 4, dtype=torch.bfloat16), + torch.randn(1, 1, 4, 4, dtype=torch.bfloat16), + ) + + +def test_attention_wrapper_has_unified_result_and_provenance(): + q, k, v = _qkv() + result = AttentionAblationOp()(q, k, v, contract=_contract()) + + assert result.backend_id == REFERENCE_BACKEND_ID + assert result.deterministic + assert result.out.shape == q.shape + assert result.lse is not None + assert result.lse.dtype is torch.float32 + assert result.provenance["semantic_operator"] == "attention" + assert result.provenance["split_kv"]["mode"] == "disabled" + assert result.provenance["strict_core_id"] == STRICT_ATTENTION_CORE_ID + assert result.provenance["strict_schedule"] == STRICT_ATTENTION_SCHEDULE_ID + assert result.readback()["out_shape"] == list(q.shape) + + +def test_attention_wrapper_supports_explicit_injected_backend(): + q, k, v = _qkv() + + class FakeBackend: + backend_id = "test.attention.backend" + + def forward_with_lse(self, q, k, v, *, causal, scale): + del k, v, causal, scale + return q.clone(), torch.zeros(q.shape[:3], dtype=torch.float32) + + result = AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + backend=FakeBackend(), + deterministic=False, + ) + assert result.backend_id == "test.attention.backend" + assert torch.equal(result.out, q) + + +def test_cp_production_configuration_fails_closed_without_ag_rs_backend(): + q, k, v = _qkv() + cp_sharding = ShardingSpec( + tp_rank=0, + tp_world_size=1, + cp_rank=0, + cp_world_size=2, + global_q_heads=2, + global_kv_heads=1, + local_q_head_start=0, + local_q_heads=2, + local_kv_head_start=0, + local_kv_heads=1, + global_sequence_length=4, + local_sequence_length=2, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 2), + ) + contract = AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=2, + head_dim=4, + causal=True, + causal_offsets=(0,), + sharding=cp_sharding, + reduction=ReductionSpec(), + split_kv=SplitKVSpec.disabled(), + ) + with pytest.raises(AttentionContractError, match="injected AG/RS backend"): + AttentionAblationOp(communication_backend="self_owned_cuda_ag_rs")( + q[:, :, :2], k[:, :, :2], v[:, :, :2], contract=contract + ) + + +def test_cp_production_wrapper_preserves_runtime_backend_provenance(): + q, k, v = _qkv() + + class StrictCPBackend: + backend_id = "injected.strict_cp_backend" + core_id = STRICT_ATTENTION_CORE_ID + strict_schedule = STRICT_ATTENTION_SCHEDULE_ID + + def __call__(self, q, k, v, *, causal, scale): + del k, v, causal, scale + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:3], dtype=torch.float32), + provenance={ + "strict_core_id": STRICT_ATTENTION_CORE_ID, + "strict_schedule": STRICT_ATTENTION_SCHEDULE_ID, + "actual_backend": "rlkernel.cuda.deterministic_attention", + "communication_backend": "self_owned_cuda_ag_rs", + "production_ready": True, + "native_attention_arithmetic": False, + "fallback": False, + }, + ) + + result = AttentionAblationOp( + cp_backend=StrictCPBackend(), + communication_backend="self_owned_cuda_ag_rs", + )( + q, + k, + v, + contract=_contract(), + backend=StrictCPBackend(), + ) + assert result.provenance["actual_backend"] == "rlkernel.cuda.deterministic_attention" + assert result.provenance["communication_backend"] == "self_owned_cuda_ag_rs" + assert result.provenance["production_ready"] is True + + +def test_deterministic_attention_rejects_runtime_split_kv_auto(): + q, k, v = _qkv() + contract = _contract(split_kv=SplitKVSpec.auto(strict_consistency=False)) + with pytest.raises(AttentionContractError, match="Split-KV=auto"): + AttentionAblationOp()(q, k, v, contract=contract) + + +def test_deterministic_native_backend_requires_explicit_native_callable(): + q, k, v = _qkv() + with pytest.raises(AttentionContractError, match="native Attention backend"): + AttentionAblationOp()(q, k, v, contract=_contract(), backend="native") + + +def test_attention_wrapper_can_return_dq_dk_dv_from_reference_backend(): + q, k, v = _qkv() + result = AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + return_gradients=True, + dout=torch.ones_like(q), + ) + + assert result.dq is not None and result.dq.shape == q.shape + assert result.dk is not None and result.dk.shape == k.shape + assert result.dv is not None and result.dv.shape == v.shape + + +def test_attention_backend_is_registered_for_pr230_semantic_resolution(): + from rl_engine.kernels.registry import kernel_registry + from rl_engine.kernels.semantic_registry import OperatorRequirements + + session = kernel_registry.semantic.session() + resolution = session.resolve( + semantic_op="attention", + requested_backend=BACKEND_ID, + target="training", + requirements=OperatorRequirements( + device="cpu", + dtype="bfloat16", + topology={"world_size": 1, "tensor_parallel_size": 1, "context_parallel_size": 1}, + alignment_properties={"deterministic": True}, + ), + ) + instance = session.instantiate(resolution) + assert isinstance(instance, AttentionAblationOp) + provenance = session.instance_provenance(resolution, instance) + assert provenance.backend_id == BACKEND_ID + + +def test_strict_wrapper_is_bitwise_invariant_to_batch_shape(): + q, k, v = _qkv() + noise_q, noise_k, noise_v = _qkv() + contract = _contract() + batch_contract = AttentionContract( + role=contract.role, + mode=contract.mode, + dtype=contract.dtype, + batch_size=2, + query_sequence_length=contract.query_sequence_length, + head_dim=contract.head_dim, + causal=contract.causal, + causal_offsets=(0, 0), + sharding=contract.sharding, + reduction=contract.reduction, + split_kv=contract.split_kv, + ) + single = AttentionAblationOp()(q, k, v, contract=contract) + batched = AttentionAblationOp()( + torch.cat((q, noise_q), dim=0), + torch.cat((k, noise_k), dim=0), + torch.cat((v, noise_v), dim=0), + contract=batch_contract, + ) + assert torch.equal(single.out[0], batched.out[0]) + assert torch.equal(single.lse[0], batched.lse[0]) diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py new file mode 100644 index 00000000..350545a1 --- /dev/null +++ b/tests/test_attention_contract.py @@ -0,0 +1,683 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 Attention CP contract and contract-aware dispatch tests (issue #235).""" + +from __future__ import annotations + +import json +from dataclasses import replace + +import pytest + +from rl_engine.kernels.attention_contract import ( + AttentionBackendCapability, + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMode, + AttentionRole, + KVCacheSpec, + ReductionSpec, + RoPEFusionBoundary, + RoPESpec, + ShardingSpec, + SplitKVExecutionPlan, + SplitKVRuntimePlanSet, + SplitKVSpec, + build_split_kv_runtime_plan_set, + validate_split_kv_alignment, + validate_split_kv_plan_set_alignment, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + + +def _sharding( + *, + tp_rank: int = 0, + tp_world_size: int = 2, + cp_rank: int = 0, + cp_world_size: int = 2, + global_sequence_length: int = 4096, + local_sequence_length: int = 2048, + global_block_indices: tuple[int, ...] = (0,), + global_block_token_starts: tuple[int, ...] = (0,), + local_block_offsets: tuple[int, ...] = (0, 2048), + packed_sequence_offsets: tuple[int, ...] | None = None, +) -> ShardingSpec: + local_q_heads = 32 // tp_world_size + local_kv_heads = 8 // tp_world_size + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=global_sequence_length, + local_sequence_length=local_sequence_length, + global_block_indices=global_block_indices, + global_block_token_starts=global_block_token_starts, + local_block_offsets=local_block_offsets, + packed_sequence_offsets=packed_sequence_offsets, + ) + + +def _contract( + *, + role: str = "infer", + mode: str = "prefill", + sharding: ShardingSpec | None = None, + kv_cache: KVCacheSpec | None = None, + causal_offsets: tuple[int, ...] = (0,), + batch_size: int = 1, + query_sequence_length: int | None = None, + rope: RoPESpec | None = None, +) -> AttentionContract: + resolved_sharding = sharding or _sharding() + return AttentionContract( + role=role, + mode=mode, + dtype="bf16", + batch_size=batch_size, + query_sequence_length=( + query_sequence_length + if query_sequence_length is not None + else (1 if mode == "decode" else resolved_sharding.local_sequence_length) + ), + head_dim=128, + causal=True, + causal_offsets=causal_offsets, + sharding=resolved_sharding, + reduction=ReductionSpec(), + kv_cache=kv_cache, + rope=rope, + ) + + +def _declared_cp_backend() -> AttentionBackendCapability: + return AttentionBackendCapability( + backend_id="test-deterministic-cp-attention", + roles=frozenset({AttentionRole.TRAIN, AttentionRole.INFER}), + modes=frozenset( + {AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL, AttentionMode.DECODE} + ), + dtypes=frozenset({AttentionDType.BF16}), + tp_world_sizes=(2,), + cp_world_sizes=(1, 2), + exports_attention_lse=True, + deterministic_cp_merge=True, + supports_packed_varlen=True, + supports_kv_cache=True, + supports_split_kv_fixed=True, + reports_actual_split_kv_plan=True, + implementation_kind="deterministic", + ) + + +def test_qwen3_tp2_cp2_contract_is_representable_and_serializable(): + contract = _contract() + + assert contract.sharding.local_q_heads == 16 + assert contract.sharding.local_kv_heads == 4 + assert contract.sharding.tp_world_size == 2 + assert contract.sharding.cp_world_size == 2 + assert contract.reduction.acc_dtype is AttentionDType.FP32 + assert contract.to_dict()["reduction"] == { + "merge": "online_softmax_lse", + "acc_dtype": "fp32", + "order": "global_block_index", + "downcast_at": "final_write", + "engine": "in_op_reference", + } + json.dumps(contract.to_dict()) + + +def test_rope_metadata_is_part_of_attention_contract_provenance(): + rope = RoPESpec( + q_state="post_rope", + k_state="post_rope", + k_cache_state="post_rope", + theta=1.0e6, + rotary_dim=128, + query_position_offsets=(0,), + key_position_offsets=(0,), + cast_at="after_rope", + output_dtype="bf16", + fusion_boundary="unfused_rope_attention", + ) + + contract = _contract(rope=rope) + payload = contract.to_dict() + + assert payload["rope"] == { + "q_state": "post_rope", + "k_state": "post_rope", + "k_cache_state": "post_rope", + "theta": 1.0e6, + "rotary_dim": 128, + "rope_scaling": None, + "position_ids": None, + "query_position_offsets": [0], + "key_position_offsets": [0], + "cast_at": "after_rope", + "output_dtype": "bf16", + "fusion_boundary": "unfused_rope_attention", + } + json.dumps(payload) + + +def test_rope_position_metadata_is_validated_against_contract_shape(): + with pytest.raises(AttentionContractError, match="rotary_dim=256"): + _contract(rope=RoPESpec(rotary_dim=256)) + + with pytest.raises(AttentionContractError, match="query_position_offsets"): + _contract(batch_size=2, causal_offsets=(0, 0), rope=RoPESpec(query_position_offsets=(0,))) + + with pytest.raises(AttentionContractError, match="position_ids"): + _contract(rope=RoPESpec(position_ids=(0, 1, 2))) + + valid = _contract(rope=RoPESpec(position_ids=tuple(range(2048)))) + assert valid.rope is not None + assert valid.rope.position_ids == tuple(range(2048)) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("tp_rank", 2, "tp_rank=2"), + ("cp_rank", 2, "cp_rank=2"), + ("global_block_indices", (), "must not be empty"), + ("global_block_indices", (1, 0), "strictly increasing"), + ], +) +def test_invalid_rank_and_cp_order_metadata_fail_loudly(field, value, message): + values = { + "tp_rank": 0, + "tp_world_size": 2, + "cp_rank": 0, + "cp_world_size": 2, + "global_q_heads": 32, + "global_kv_heads": 8, + "local_q_head_start": 0, + "local_q_heads": 16, + "local_kv_head_start": 0, + "local_kv_heads": 4, + "global_sequence_length": 4096, + "local_sequence_length": 2048, + "global_block_indices": (0,), + "global_block_token_starts": (0,), + "local_block_offsets": (0, 2048), + } + values[field] = value + + with pytest.raises(AttentionContractError, match=message): + ShardingSpec(**values) + + +def test_tp_local_heads_must_preserve_global_gqa_mapping(): + with pytest.raises(AttentionContractError, match="local TP head counts"): + replace(_sharding(), local_q_heads=7) + + with pytest.raises(AttentionContractError, match="head starts"): + replace(_sharding(tp_rank=1), local_q_head_start=0) + + +def test_sequence_range_and_packed_offsets_are_validated(): + with pytest.raises(AttentionContractError, match="exceeds global_sequence_length"): + _sharding(global_block_token_starts=(4000,)) + + with pytest.raises(AttentionContractError, match="final packed_sequence_offsets"): + _sharding(packed_sequence_offsets=(0, 512)) + + sharding = _sharding(packed_sequence_offsets=(0, 512, 2048)) + assert sharding.packed_sequence_offsets == (0, 512, 2048) + + +def test_non_contiguous_cp_blocks_have_explicit_global_and_local_offsets(): + sharding = _sharding( + global_block_indices=(0, 3), + global_block_token_starts=(0, 3072), + local_block_offsets=(0, 1024, 2048), + ) + + assert sharding.global_block_indices == (0, 3) + assert sharding.global_block_token_starts == (0, 3072) + assert sharding.local_block_offsets == (0, 1024, 2048) + + with pytest.raises(AttentionContractError, match="non-overlapping and ordered"): + _sharding( + global_block_indices=(0, 1), + global_block_token_starts=(0, 512), + local_block_offsets=(0, 1024, 2048), + ) + + +def test_reduction_requires_fp32_accumulation(): + with pytest.raises(AttentionContractError, match="must be fp32"): + ReductionSpec(acc_dtype="bf16") + + +def test_split_kv_policy_is_a_first_class_strict_contract(): + contract = _contract() + assert contract.to_dict()["split_kv"] == { + "mode": "disabled", + "fixed_split_size": None, + "strict_consistency": True, + } + + fixed = replace(contract, split_kv=SplitKVSpec.fixed(128)) + assert fixed.to_dict()["split_kv"]["mode"] == "fixed" + assert fixed.to_dict()["split_kv"]["fixed_split_size"] == 128 + + with pytest.raises(AttentionContractError, match="auto Split-KV"): + SplitKVSpec.auto(strict_consistency=True) + + +def test_split_kv_execution_plan_records_actual_logical_schedule(): + plan = SplitKVSpec.fixed(4).resolve(10, backend="training-reference") + + assert plan.actual_split_count == 3 + assert plan.to_dict()["actual_split_boundaries"] == [[0, 4], [4, 8], [8, 10]] + assert plan.to_dict()["split_kv_merge_order"] == "global_block_index" + assert plan.to_dict()["split_kv_accum_dtype"] == "fp32" + assert plan.to_dict()["split_kv_downcast_at"] == "final_write" + + +def test_strict_split_kv_alignment_rejects_unknown_or_mismatched_actual_plan(): + training = SplitKVSpec.fixed(4).resolve(10, backend="training") + rollout = SplitKVSpec.fixed(4).resolve(10, backend="rollout") + validate_split_kv_alignment(training, rollout) + + unknown = SplitKVSpec.auto().resolve(10, backend="rollout") + with pytest.raises(AttentionContractError, match="actual runtime plans"): + validate_split_kv_alignment(training, unknown) + + mismatched = SplitKVSpec.fixed(5).resolve(10, backend="rollout") + with pytest.raises(AttentionContractError, match="differ"): + validate_split_kv_alignment(training, mismatched) + + with pytest.raises(AttentionContractError, match="contiguous"): + SplitKVExecutionPlan( + requested_mode="fixed", + requested_split_size=4, + actual_mode="fixed", + actual_split_size=4, + boundaries=((0, 4), (5, 10)), + ) + + +def test_complete_split_kv_plan_set_covers_batch_tp_cp_and_owner_coordinates(): + plan_set = build_split_kv_runtime_plan_set( + (8, 10), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(2), + backend="training-reference", + ) + + assert len(plan_set.entries) == 16 + assert plan_set.to_dict()["coverage"] == ("complete_batch_tp_cp_owner_cartesian_product") + assert { + tuple(entry["expected_kv_range"]) + for entry in plan_set.to_dict()["entries"] + if entry["batch_index"] == 0 + } == {(0, 4), (4, 8)} + + +def test_split_kv_plan_set_alignment_rejects_missing_and_mismatched_rank_plans(): + training = build_split_kv_runtime_plan_set( + (8,), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(2), + backend="training", + ) + rollout = build_split_kv_runtime_plan_set( + (8,), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(2), + backend="rollout", + ) + validate_split_kv_plan_set_alignment(training, rollout) + + with pytest.raises(AttentionContractError, match="coordinate coverage is incomplete"): + SplitKVRuntimePlanSet( + batch_size=training.batch_size, + tp_world_size=training.tp_world_size, + cp_world_size=training.cp_world_size, + total_kv_tokens=training.total_kv_tokens, + entries=training.entries[:-1], + ) + + mismatched = build_split_kv_runtime_plan_set( + (8,), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(1), + backend="rollout", + ) + with pytest.raises(AttentionContractError, match="plan differs"): + validate_split_kv_plan_set_alignment(training, mismatched) + + +def test_backend_must_support_policy_and_actual_plan_provenance(): + fixed = replace(_contract(), split_kv=SplitKVSpec.fixed(128)) + capability = replace( + _declared_cp_backend(), + supports_split_kv_fixed=False, + reports_actual_split_kv_plan=False, + ) + + assert capability.incompatibilities(fixed)[-2:] == ( + "Split-KV policy=fixed is unsupported", + "actual Split-KV execution-plan provenance is unsupported", + ) + + +def test_causal_attention_requires_explicit_offset(): + contract = _contract() + with pytest.raises(AttentionContractError, match="causal_offsets are required"): + replace(contract, causal_offsets=None) + + +def test_full_prefill_query_length_must_match_local_sequence_length(): + with pytest.raises(AttentionContractError, match="prefill query_sequence_length must equal"): + _contract(mode="prefill", query_sequence_length=1024) + + chunked = _contract(mode="chunked_prefill", query_sequence_length=512) + decode = _contract( + mode="decode", + query_sequence_length=1, + kv_cache=KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + ), + ) + assert chunked.query_sequence_length == 512 + assert decode.query_sequence_length == 1 + + +def test_decode_requires_complete_kv_cache_identity(): + with pytest.raises(AttentionContractError, match="kv_cache metadata is required"): + _contract(mode="decode") + + cache = KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1, -1),), + global_token_positions=tuple(range(17)), + page_size=16, + prefix_cache_enabled=True, + prefix_cache_key="prefix:sample-0", + ) + contract = _contract(mode="decode", kv_cache=cache) + assert contract.to_dict()["kv_cache"]["block_table"] == [[0, 1, -1]] + + +def test_prefix_cache_key_is_required_only_when_prefix_cache_is_enabled(): + with pytest.raises(AttentionContractError, match="prefix_cache_key is required"): + KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + prefix_cache_enabled=True, + ) + + +def test_cache_positions_must_match_kv_sequence_count(): + with pytest.raises(AttentionContractError, match="one entry per kv_seq_lens"): + KVCacheSpec( + cache_positions=(1,), + kv_seq_lens=(2, 2), + block_table=((0,), (1,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + ) + + +def test_cache_position_must_match_terminal_global_token_position(): + with pytest.raises(AttentionContractError, match="terminal global token position"): + KVCacheSpec( + cache_positions=(999,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + ) + + +@pytest.mark.parametrize("positions", [(7, 6), (7, 7)]) +def test_kv_cache_positions_must_be_strictly_increasing_per_sequence(positions): + with pytest.raises(AttentionContractError, match="strictly increasing"): + KVCacheSpec( + cache_positions=(7,), + kv_seq_lens=(2,), + block_table=((0,),), + global_token_positions=positions, + page_size=2, + ) + + +@pytest.mark.parametrize( + ("block_table", "message"), + [ + ((0, -1, 1), "padding must be trailing"), + ((0, 0, -1), "duplicate active page ids"), + ((0, -1, -1), "active page count"), + ], +) +def test_kv_cache_block_table_page_mapping_is_validated(block_table, message): + with pytest.raises(AttentionContractError, match=message): + KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=(block_table,), + global_token_positions=tuple(range(17)), + page_size=16, + ) + + +def test_prefix_pages_may_be_shared_across_sequences(): + cache = KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + assert cache.block_table == ((3,), (3,)) + assert cache.shared_prefix_page_count == 1 + + +def test_non_prefix_cache_rejects_cross_sequence_page_sharing(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=False, + ) + + +def test_prefix_cache_requires_explicit_shared_page_count(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=0, + ) + + +def test_prefix_cache_rejects_shared_writable_suffix_pages(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(3, 3), + kv_seq_lens=(4, 4), + block_table=((3, 4), (3, 4)), + global_token_positions=(0, 1, 2, 3, 0, 1, 2, 3), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + +def test_shared_prefix_identity_must_match_pages_and_positions(): + with pytest.raises(AttentionContractError, match="page ids must match"): + KVCacheSpec( + cache_positions=(3, 3), + kv_seq_lens=(4, 4), + block_table=((3, 4), (5, 6)), + global_token_positions=(0, 1, 2, 3, 0, 1, 2, 3), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + with pytest.raises(AttentionContractError, match="token positions must match"): + KVCacheSpec( + cache_positions=(3, 13), + kv_seq_lens=(4, 4), + block_table=((3, 4), (3, 5)), + global_token_positions=(0, 1, 2, 3, 10, 11, 12, 13), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + +def test_shared_prefix_pages_must_be_fully_populated(): + with pytest.raises(AttentionContractError, match="fully populated and read-only"): + KVCacheSpec( + cache_positions=(0, 0), + kv_seq_lens=(1, 1), + block_table=((3,), (3,)), + global_token_positions=(0, 0), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="partial-prefix-page", + shared_prefix_page_count=1, + ) + + +def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION] + + with pytest.raises(RuntimeError) as exc_info: + registry.get_attention_op(_contract()) + + message = str(exc_info.value) + assert "CP=2 is unsupported" in message + assert "attention-domain LSE export is unsupported" in message + assert "deterministic CP (out, lse) merge is unsupported" in message + + +def test_undeclared_backend_capability_is_never_selected(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_ATTN] + + with pytest.raises(RuntimeError, match="no AttentionBackendCapability declared"): + registry.get_attention_op(_contract()) + + +def test_declared_compatible_backend_resolves_and_records_provenance(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION] + registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend() + + result = registry.get_attention_op(_contract(), requested_backend="deterministic") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-cp-attention" + assert result.provenance["requested_backend"] == "deterministic" + assert result.provenance["actual_backend"] == "test-deterministic-cp-attention" + assert result.provenance["fallback"] is False + assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 + assert result.provenance["contract"]["sharding"]["cp_world_size"] == 2 + json.dumps(result.provenance) + + +def test_requested_stable_backend_id_is_enforced(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION] + registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend() + + with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"): + registry.get_attention_op(_contract(), requested_backend="another-backend") + + result = registry.get_attention_op( + _contract(), requested_backend="test-deterministic-cp-attention" + ) + assert result.provenance["actual_backend"] == "test-deterministic-cp-attention" + + +def test_packed_layout_requires_declared_backend_support(): + capability = replace(_declared_cp_backend(), supports_packed_varlen=False) + contract = _contract( + sharding=_sharding(packed_sequence_offsets=(0, 512, 2048)), + causal_offsets=(0, 0), + batch_size=2, + ) + + assert capability.incompatibilities(contract) == ("packed varlen layout is unsupported",) + + +def test_rope_contract_requires_declared_backend_support(): + contract = _contract(rope=RoPESpec()) + capability = _declared_cp_backend() + + assert capability.incompatibilities(contract) == ("RoPE/position metadata is unsupported",) + + supported = replace(capability, supports_rope_metadata=True) + assert supported.incompatibilities(contract) == () + + +def test_fused_rope_attention_boundary_requires_declared_backend_support(): + contract = _contract(rope=RoPESpec(fusion_boundary=RoPEFusionBoundary.FUSED_ROPE_ATTENTION)) + capability = replace(_declared_cp_backend(), supports_rope_metadata=True) + + assert capability.incompatibilities(contract) == ( + "fused RoPE+Attention boundary is unsupported", + ) + + supported = replace(capability, supports_fused_rope_attention=True) + assert supported.incompatibilities(contract) == () + + +def test_packed_sequence_count_must_match_logical_batch_size(): + sharding = _sharding(packed_sequence_offsets=(0, 512, 2048)) + + with pytest.raises(AttentionContractError, match="must equal logical batch_size"): + _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=1) + + contract = _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=2) + assert contract.batch_size == 2 diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py new file mode 100644 index 00000000..c1c8d35d --- /dev/null +++ b/tests/test_attention_cross_config_binding.py @@ -0,0 +1,1295 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for #235 PR4: rollout/training attention contract binding. + +Every test here runs on CPU without Megatron or vLLM installed. That is the point: +the binding rules are contract logic, and contract logic that can only be exercised +on a 2-node x 2-GPU cluster would never be exercised. +""" + +from __future__ import annotations + +import json +from dataclasses import replace +from enum import Enum +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from rl_engine.alignment.cross_config.adapters import ( + QWEN3_8B, + WS2_ATTENTION_KNOBS, + AttentionRuntimeReadback, + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) +from rl_engine.alignment.cross_config.attention_binding import ( + ATTENTION_LSE_DOMAIN, + AttentionBindingError, + BindingErrorCode, + BindingTier, + bind_attention_contracts, + bind_attention_runtime_readbacks, + first_blocking_issue, + identity_fingerprint, + summarize_binding, +) +from rl_engine.alignment.cross_config.determinism import ( + compare_determinism, + megatron_probe_from_config, + vllm_probe_from_env, +) +from rl_engine.alignment.cross_config.schema import MaterializationStatus +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + AttentionContractError, + AttentionMode, + AttentionRole, + KVCacheSpec, + SplitKVExecutionPlan, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, + SplitKVSpec, + build_split_kv_runtime_plan_set, +) +from rl_engine.kernels.attention_preprocess import MANDATED_ATTENTION_PREPROCESS_BACKENDS +from rl_engine.kernels.attention_projection import ( + O_PROJ_COLLECTIVE_CONTRACT, + QKV_COLLECTIVE_CONTRACT, + ProjectionPlan, +) + +pytestmark = pytest.mark.unit + + +TRAINING_KNOBS = { + "batch.size": 2, + "training.tensor_parallel_size": 2, + "training.context_parallel_size": 2, + "training.compute_dtype": "bf16", + "attention.split_kv_policy": 32, +} + +ROLLOUT_KNOBS = { + "batch.size": 2, + "rollout.tensor_parallel_size": 2, + "rollout.context_parallel_size": 2, + "rollout.dtype": "bf16", + "attention.split_kv_policy": 32, +} + + +def _identity(**overrides): + identity = { + "checkpoint_id": "qwen3-8b", + "model_version": "v1", + "weight_version": 7, + "tokenizer_fingerprint": "tokenizer-abc", + "token_ids_fingerprint": "tokens-abc", + "active_mask_fingerprint": "mask-abc", + "position_ids_fingerprint": "pos-abc", + "padding_side": "right", + "pre_update_state": "pre_update", + "batch_size": 2, + "global_token_positions_fingerprint": "gtp-abc", + "kv_seq_lens_fingerprint": "kvlen-abc", + } + identity.update(QWEN3_8B.identity_fields()) + identity.update(overrides) + return identity + + +def _contracts(): + training = MegatronAttentionMaterializer().build_contract(TRAINING_KNOBS) + rollout = VllmRolloutMaterializer().build_contract(ROLLOUT_KNOBS) + return rollout, training + + +def _plan_set(contract, *, backend): + return build_split_kv_runtime_plan_set( + (contract.sharding.global_sequence_length,) * contract.batch_size, + tp_world_size=contract.sharding.tp_world_size, + cp_world_size=contract.sharding.cp_world_size, + split_kv=contract.split_kv, + backend=backend, + ) + + +def _bind(rollout_identity=None, training_identity=None, **kwargs): + rollout, training = _contracts() + rollout = kwargs.pop("rollout_contract", rollout) + training = kwargs.pop("training_contract", training) + return bind_attention_contracts( + rollout_contract=rollout, + training_contract=training, + rollout_identity=rollout_identity if rollout_identity is not None else _identity(), + training_identity=training_identity if training_identity is not None else _identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + rollout_split_kv_plan_set=kwargs.pop( + "rollout_split_kv_plan_set", _plan_set(rollout, backend="vllm.readback") + ), + training_split_kv_plan_set=kwargs.pop( + "training_split_kv_plan_set", _plan_set(training, backend="megatron.readback") + ), + **kwargs, + ) + + +# -------------------------------------------------------------------------- +# tier 1: identity +# -------------------------------------------------------------------------- + + +def test_matching_identity_and_topology_bind_despite_different_materialization(): + """The core claim of PR4: same identity + same reduction, different runtimes.""" + + result = _bind() + + assert result.comparable + assert result.passed + assert result.issues == () + # Attention mode is a framework materialization difference, while both sides + # execute the same TP=2/CP=2 local ownership and Split-K schedule. + assert "mode" in result.recorded_differences + assert result.recorded_differences["mode"] == { + "rollout": "chunked_prefill", + "training": "prefill", + } + + +def test_weight_version_mismatch_is_not_comparable(): + result = _bind(rollout_identity=_identity(weight_version=6)) + + assert not result.comparable + assert not result.passed + codes = {issue.code for issue in result.issues} + assert BindingErrorCode.IDENTITY_MISMATCH in codes + blocking = first_blocking_issue(result) + assert blocking is not None and blocking.tier is BindingTier.IDENTICAL + assert "NOT COMPARABLE" in summarize_binding(result) + + +def test_rope_theta_mismatch_is_not_comparable(): + """RoPE math constants are identity, not materialization.""" + + result = _bind(training_identity=_identity(rope_theta=10000.0)) + + assert not result.comparable + assert any(issue.field == "rope_theta" for issue in result.issues) + + +def test_null_rope_scaling_is_a_value_not_an_omission(): + """Qwen3-8B applies no RoPE scaling; ``None`` must not read as undeclared.""" + + result = _bind() + + assert not result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + + +def test_missing_identity_field_is_reported_per_side(): + identity = _identity() + del identity["padding_side"] + result = _bind(rollout_identity=identity, training_identity=identity) + + missing = result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + assert {issue.field for issue in missing} == { + "rollout.padding_side", + "training.padding_side", + } + assert not result.comparable + + +def test_single_gpu_harness_may_waive_full_identity(): + """#235 PR2 has no KV-cache identity to declare; it opts out explicitly.""" + + identity = _identity() + del identity["global_token_positions_fingerprint"] + del identity["kv_seq_lens_fingerprint"] + + strict = _bind(rollout_identity=identity, training_identity=identity) + waived = _bind( + rollout_identity=identity, + training_identity=identity, + require_full_identity=False, + ) + + assert not strict.comparable + assert waived.comparable and waived.passed + + +def test_identity_fingerprint_ignores_undeclared_extra_keys(): + base = _identity() + decorated = dict(base, diagnostic_note="added later") + + assert identity_fingerprint(base) == identity_fingerprint(decorated) + + +# -------------------------------------------------------------------------- +# tier 2: reduction semantics +# -------------------------------------------------------------------------- + + +def test_reduction_semantics_are_bound_and_fingerprinted(): + result = _bind() + + reduction = result.provenance["training"]["contract"]["reduction"] + assert reduction["merge"] == "online_softmax_lse" + assert reduction["acc_dtype"] == "fp32" + assert reduction["order"] == "global_block_index" + assert reduction["downcast_at"] == "final_write" + assert result.reduction_fingerprint + + +def test_reduction_engine_difference_is_recorded_not_rejected(): + """A TE merge oracle on one side must not fail the binding.""" + + from rl_engine.alignment.cross_config.attention_binding import ( + RECORDED_FIELDS, + SEMANTIC_REDUCTION_FIELDS, + ) + + assert "reduction.engine" in RECORDED_FIELDS + assert "engine" not in SEMANTIC_REDUCTION_FIELDS + + +def test_lse_domain_is_recorded_as_attention_domain(): + """#235: attention exports attention-domain LSE, not vocab-logprob LSE.""" + + result = _bind() + + assert result.provenance["lse_domain"] == ATTENTION_LSE_DOMAIN == "attention" + + +def test_mixed_dtypes_fail_closed(): + """BF16 rollout against FP16 training produces an unattributable number.""" + + rollout = VllmRolloutMaterializer().build_contract( + {**ROLLOUT_KNOBS, "rollout.dtype": "float16"} + ) + result = _bind(rollout_contract=rollout) + + assert result.comparable # identity is fine + assert not result.passed + assert any(issue.field == "dtype" for issue in result.issues) + + +def test_precision_sweep_may_opt_into_mixed_dtypes(): + """#235 PR5 sweeps BF16 against an FP32 reference; it says so explicitly.""" + + training = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "float32"} + ) + result = _bind(training_contract=training, allow_dtype_difference=True) + + assert result.passed + assert result.provenance["dtype"] == "fp32" + + +def test_batch_size_mismatch_is_not_comparable(): + """Batch invariance is a claim about batch makeup, so it belongs to identity.""" + + result = _bind(rollout_identity=_identity(batch_size=4)) + + assert not result.comparable + assert any(issue.field == "batch_size" for issue in result.issues) + + +def test_missing_split_kv_runtime_evidence_fails_closed(): + result = _bind(rollout_split_kv_plan_set=None) + + assert result.comparable + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_RUNTIME_MISSING) + + +def test_split_kv_requested_policy_mismatch_fails_closed(): + rollout, training = _contracts() + rollout = replace(rollout, split_kv=SplitKVSpec.fixed(16)) + result = _bind(rollout_contract=rollout) + + assert result.comparable + assert not result.passed + assert any(issue.field == "split_kv" for issue in result.issues) + + +def test_split_kv_runtime_boundary_mismatch_fails_closed(): + rollout, training = _contracts() + mismatched_rollout = build_split_kv_runtime_plan_set( + (rollout.sharding.global_sequence_length,) * rollout.batch_size, + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(16), + backend="vllm.readback", + ) + result = _bind(rollout_split_kv_plan_set=mismatched_rollout) + + assert not result.passed + issues = result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + assert any(issue.field == "split_kv_runtime_plan_set" for issue in issues) + + +def test_split_kv_plan_set_must_match_its_own_contract_topology(): + rollout, _ = _contracts() + wrong_topology = build_split_kv_runtime_plan_set( + (rollout.sharding.global_sequence_length,) * rollout.batch_size, + tp_world_size=1, + cp_world_size=2, + split_kv=rollout.split_kv, + backend="vllm.readback", + ) + + result = _bind(rollout_split_kv_plan_set=wrong_topology) + + assert not result.passed + assert any( + issue.field == "rollout.split_kv_runtime_plan_set" + for issue in result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + ) + + +@pytest.mark.parametrize( + ("field_name", "corrupt_value"), + [ + ("merge_order", Enum("BadOrder", {"ARRIVAL": "arrival"}).ARRIVAL), + ("acc_dtype", Enum("BadDType", {"BF16": "bf16"}).BF16), + ("downcast_at", Enum("BadDowncast", {"PER_BLOCK": "per_block"}).PER_BLOCK), + ], +) +def test_split_kv_runtime_merge_semantic_corruption_fails_closed(field_name, corrupt_value): + rollout, _ = _contracts() + corrupted = _plan_set(rollout, backend="vllm.readback") + # Runtime reports are deserialized at this boundary. Simulate a corrupted + # report after construction to prove binding compares the actual fields. + object.__setattr__(corrupted.entries[0].execution, field_name, corrupt_value) + + result = _bind(rollout_split_kv_plan_set=corrupted) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + + +def test_split_kv_runtime_fallback_fails_closed(): + rollout, _ = _contracts() + plan_set = _plan_set(rollout, backend="vllm.readback") + fallback_entries = [] + for entry in plan_set.entries: + execution = entry.execution + fallback_entries.append( + SplitKVRuntimePlanEntry( + coordinate=entry.coordinate, + expected_kv_range=entry.expected_kv_range, + execution=SplitKVExecutionPlan( + requested_mode=execution.requested_mode, + requested_split_size=execution.requested_split_size, + actual_mode=execution.actual_mode, + actual_split_size=execution.actual_split_size, + boundaries=execution.boundaries, + merge_order=execution.merge_order, + acc_dtype=execution.acc_dtype, + downcast_at=execution.downcast_at, + backend=execution.backend, + source="runtime_fallback", + fallback=True, + fallback_reason="backend substituted a runtime plan", + ), + ) + ) + fallback = SplitKVRuntimePlanSet( + batch_size=plan_set.batch_size, + tp_world_size=plan_set.tp_world_size, + cp_world_size=plan_set.cp_world_size, + total_kv_tokens=plan_set.total_kv_tokens, + entries=tuple(fallback_entries), + ) + + result = _bind(rollout_split_kv_plan_set=fallback) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_FALLBACK) + + +@pytest.mark.parametrize( + "rollout_overrides", + [ + {"rollout.tensor_parallel_size": 1}, + {"rollout.context_parallel_size": 1}, + ], +) +def test_tp_or_cp_topology_mismatch_is_not_comparable(rollout_overrides): + rollout = VllmRolloutMaterializer().build_contract({**ROLLOUT_KNOBS, **rollout_overrides}) + result = _bind(rollout_contract=rollout) + + assert not result.comparable + assert result.issues_by_code(BindingErrorCode.TOPOLOGY_MISMATCH) + + +# -------------------------------------------------------------------------- +# role and input validation +# -------------------------------------------------------------------------- + + +def test_swapped_roles_are_rejected_outright(): + rollout, training = _contracts() + + with pytest.raises(AttentionBindingError): + bind_attention_contracts( + rollout_contract=training, + training_contract=rollout, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="a", + training_backend_id="b", + ) + + +# -------------------------------------------------------------------------- +# determinism cross-check +# -------------------------------------------------------------------------- + + +def _megatron_env(): + return {"NCCL_ALGO": "Tree", "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0"} + + +def _vllm_env(**overrides): + env = { + "VLLM_BATCH_INVARIANT": "1", + "NCCL_ALGO": "allreduce:tree", + "NCCL_PROTO": "Simple", + "NCCL_MIN_NCHANNELS": "1", + "NCCL_MAX_NCHANNELS": "1", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + } + env.update(overrides) + return env + + +def test_nccl_algo_mismatch_blocks_the_binding(): + """Megatron asserts NCCL_ALGO; vLLM hard-sets a different value.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert not report.compatible + fields = {issue.field for issue in report.issues} + assert "env.NCCL_ALGO" in fields + assert "env.NCCL_PROTO" in fields + + +def test_matching_nccl_settings_are_compatible(): + shared = {"NCCL_ALGO": "allreduce:tree", "NCCL_PROTO": "Simple"} + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=dict(shared) + ) + rollout = vllm_probe_from_env({**_vllm_env(**shared), "CUBLAS_WORKSPACE_CONFIG": None}) + + report = compare_determinism(rollout=rollout, training=training) + + assert report.compatible, [issue.to_dict() for issue in report.issues] + + +def test_determinism_switch_off_on_either_side_blocks(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=False), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env(VLLM_BATCH_INVARIANT="0")) + + report = compare_determinism(rollout=rollout, training=training) + + fields = {issue.field for issue in report.issues} + assert "training.deterministic_mode" in fields + assert "rollout.VLLM_BATCH_INVARIANT" in fields + + +def test_tf32_asymmetry_is_recorded_not_blocking(): + """Megatron does not manage TF32 at all; vLLM disables it. Record the gap.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert training.tf32_disabled is None + assert rollout.tf32_disabled is True + assert "tf32_disabled" in report.differences + assert not any(issue.field == "tf32_disabled" for issue in report.issues) + + +def test_determinism_issues_flow_into_the_binding(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + report = compare_determinism(rollout=rollout, training=training) + + result = _bind(determinism_issues=report.issues) + + assert result.comparable # identity is fine + assert not result.passed # but the reduction environment is not + assert result.issues_by_code(BindingErrorCode.DETERMINISM_INCOMPATIBLE) + assert "FAILED CLOSED" in summarize_binding(result) + + +# -------------------------------------------------------------------------- +# sharding derived from the frozen #239 layout +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_cp_shards_cover_the_global_sequence_without_overlap(cp_rank): + contract = MegatronAttentionMaterializer( + cp_rank=cp_rank, global_sequence_length=4096 + ).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert sharding.local_sequence_length == 2048 + assert sharding.global_block_indices == (cp_rank,) + assert sharding.global_block_token_starts == (cp_rank * 2048,) + # The causal offset must be the number of preceding *global* tokens, otherwise + # rank 1 would mask as if its shard started at position zero. + assert contract.causal_offsets == (cp_rank * 2048, cp_rank * 2048) + + +def test_tp_head_shards_split_qwen3_gqa_evenly(): + contract = MegatronAttentionMaterializer(tp_rank=1).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert (sharding.global_q_heads, sharding.global_kv_heads) == (32, 8) + assert (sharding.local_q_heads, sharding.local_kv_heads) == (16, 4) + assert (sharding.local_q_head_start, sharding.local_kv_head_start) == (16, 4) + + +@pytest.mark.parametrize("tp_world_size", [2, 4, 8]) +def test_supported_tp_degrees_shard_qwen3_gqa(tp_world_size): + """Qwen3-8B has 32 Q heads and 8 KV heads, so TP in {2, 4, 8} all divide.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": tp_world_size} + ) + + sharding = contract.sharding + assert sharding.local_q_heads == 32 // tp_world_size + assert sharding.local_kv_heads == 8 // tp_world_size + + +@pytest.mark.parametrize( + ("knob_value", "expected"), + [("bfloat16", "bf16"), ("float16", "fp16"), ("float32", "fp32"), ("fp16", "fp16")], +) +def test_planner_normalized_dtypes_reach_the_contract(knob_value, expected): + """The planner emits torch spellings; AttentionDType uses short ones.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": knob_value} + ) + + assert contract.dtype.value == expected + + +def test_unknown_dtype_is_rejected_with_the_offending_field(): + with pytest.raises(ValueError, match="training.compute_dtype"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "int8"} + ) + + +def test_indivisible_tp_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": 3} + ) + + +def test_indivisible_cp_sequence_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer(global_sequence_length=4097).build_contract(TRAINING_KNOBS) + + +# -------------------------------------------------------------------------- +# materialization: fail closed rather than silently substitute +# -------------------------------------------------------------------------- + + +def _statuses(materialization, path): + return [app.status for app in materialization.applications if app.path == path] + + +def _readback(materializer, flat, *, source): + contract = materializer.build_contract(flat) + projection_plans = { + name: ProjectionPlan( + projection=name, + backend_id="rlkernel.cuda.det_gemm", + fallback=True, + fallback_reason="native projection probe failed", + probe_id=f"{source}-{name}", + collective=collective.to_dict(), + ).to_dict() + for name, collective in ( + ("qkv", QKV_COLLECTIVE_CONTRACT), + ("o_proj", O_PROJ_COLLECTIVE_CONTRACT), + ) + } + return AttentionRuntimeReadback( + contract=contract, + actual_knobs=dict(flat), + split_kv_plan_set=_plan_set(contract, backend=source), + source=source, + frozen_scope_verified=True, + preprocess_backends=MANDATED_ATTENTION_PREPROCESS_BACKENDS, + preprocess_fallback=False, + projection_plans=projection_plans, + actual_backend=f"reference.{source}", + communication_backend="none", + production_ready=False, + ) + + +def _strict_readback(materializer, flat, *, source): + readback = _readback(materializer, flat, source=source) + contract = replace(readback.contract, split_kv=SplitKVSpec.disabled()) + actual_knobs = dict(readback.actual_knobs) + actual_knobs["attention.split_kv_policy"] = "disabled" + return replace( + readback, + contract=contract, + actual_knobs=actual_knobs, + split_kv_plan_set=_plan_set(contract, backend=source), + strict_mode=True, + strict_core_id=STRICT_ATTENTION_CORE_ID, + strict_schedule=STRICT_ATTENTION_SCHEDULE_ID, + native_attention_arithmetic=False, + strict_split_kv_policy="disabled", + actual_backend="rlkernel.cuda.deterministic_attention", + communication_backend="self_owned_cuda_ag_rs", + production_ready=True, + ) + + +def test_configured_contract_without_runtime_readback_is_unobservable(): + normalized = { + "batch": {"size": 2}, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "compute_dtype": "bf16", + }, + "attention": {"split_kv_policy": 32}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert materialization.applications + assert {app.status for app in materialization.applications} == { + MaterializationStatus.UNOBSERVABLE + } + assert all(app.actual is None for app in materialization.applications) + + +@pytest.mark.parametrize( + ("materializer_type", "flat", "source"), + [ + (MegatronAttentionMaterializer, TRAINING_KNOBS, "megatron.runtime_readback"), + (VllmRolloutMaterializer, ROLLOUT_KNOBS, "vllm.runtime_readback"), + ], +) +def test_runtime_readback_can_verify_materialized_knobs(materializer_type, flat, source): + configured = materializer_type() + readback = _readback(configured, flat, source=source) + materializer = materializer_type(runtime_readback=readback) + normalized = {} + for path, value in flat.items(): + section, key = path.split(".", 1) + normalized.setdefault(section, {})[key] = value + + materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) + + assert materialization.applications + assert {app.status for app in materialization.applications} == {MaterializationStatus.APPLIED} + side = "training" if materializer_type is MegatronAttentionMaterializer else "rollout" + assert materialization.binding.side_configs[side]["runtime_readback"]["source"] == source + + +def test_runtime_readback_mismatch_is_a_fallback(): + configured = MegatronAttentionMaterializer() + readback = _readback(configured, TRAINING_KNOBS, source="megatron.runtime_readback") + actual = dict(readback.actual_knobs) + actual["training.context_parallel_size"] = 1 + mismatched = AttentionRuntimeReadback( + contract=readback.contract, + actual_knobs=actual, + split_kv_plan_set=readback.split_kv_plan_set, + source=readback.source, + frozen_scope_verified=True, + preprocess_backends=readback.preprocess_backends, + preprocess_fallback=readback.preprocess_fallback, + ) + normalized = { + "batch": {"size": 2}, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "compute_dtype": "bf16", + }, + "attention": {"split_kv_policy": 32}, + } + materialization = MegatronAttentionMaterializer(runtime_readback=mismatched).materialize( + normalized, WS2_ATTENTION_KNOBS + ) + + assert _statuses(materialization, "training.context_parallel_size") == [ + MaterializationStatus.FALLBACK + ] + + +def test_decode_split_kv_plan_set_must_match_kv_cache_lengths(): + rollout, _ = _contracts() + kv_cache = KVCacheSpec( + cache_positions=(3, 5), + kv_seq_lens=(4, 6), + block_table=((0, 1, -1), (2, 3, 4)), + global_token_positions=tuple(range(4)) + tuple(range(6)), + page_size=2, + ) + decode = replace( + rollout, + role=AttentionRole.INFER, + mode=AttentionMode.DECODE, + query_sequence_length=1, + causal_offsets=(3, 5), + kv_cache=kv_cache, + ) + wrong_lengths = build_split_kv_runtime_plan_set( + (4, 8), + tp_world_size=2, + cp_world_size=2, + split_kv=decode.split_kv, + backend="vllm.decode.readback", + ) + + with pytest.raises(ValueError, match="KV-cache lengths"): + AttentionRuntimeReadback( + contract=decode, + actual_knobs={}, + split_kv_plan_set=wrong_lengths, + source="vllm.decode.readback", + frozen_scope_verified=True, + ) + + +def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): + rollout_materializer = VllmRolloutMaterializer() + training_materializer = MegatronAttentionMaterializer() + rollout = _readback(rollout_materializer, ROLLOUT_KNOBS, source="vllm.runtime_readback") + training = _readback( + training_materializer, + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.passed + assert result.provenance["split_kv_runtime"]["rollout"]["coverage"] == ( + "complete_batch_tp_cp_owner_cartesian_product" + ) + assert result.provenance["rollout"]["recorded"]["preprocess.qk_rmsnorm"] == ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS["qk_rmsnorm"] + ) + + +@pytest.mark.parametrize( + ("changes", "expected_code"), + [ + ( + {"native_attention_arithmetic": True}, + BindingErrorCode.ATTENTION_NATIVE_ARITHMETIC, + ), + ( + {"strict_core_id": "different.core"}, + BindingErrorCode.ATTENTION_CORE_MISSING, + ), + ( + {"strict_schedule": "different_schedule"}, + BindingErrorCode.ATTENTION_CORE_SCHEDULE, + ), + ( + {"strict_split_kv_policy": "fixed"}, + BindingErrorCode.ATTENTION_CORE_SPLIT_K, + ), + ], +) +def test_strict_runtime_readback_rejects_non_shared_attention_arithmetic(changes, expected_code): + rollout = replace( + _strict_readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + **changes, + ) + training = _strict_readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="rlkernel.cuda.deterministic_attention", + training_backend_id="rlkernel.cuda.deterministic_attention", + ) + + assert not result.passed + assert result.issues_by_code(expected_code) + + +def test_strict_runtime_readback_accepts_shared_no_split_k_core(): + rollout = _strict_readback( + VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback" + ) + training = _strict_readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="rlkernel.cuda.deterministic_attention", + training_backend_id="rlkernel.cuda.deterministic_attention", + ) + + assert result.passed + assert result.provenance["rollout"]["recorded"]["strict.core_id"] == (STRICT_ATTENTION_CORE_ID) + assert result.provenance["rollout"]["recorded"]["runtime.actual_backend"] == ( + "rlkernel.cuda.deterministic_attention" + ) + assert result.provenance["rollout"]["recorded"]["runtime.communication_backend"] == ( + "self_owned_cuda_ag_rs" + ) + assert result.provenance["rollout"]["recorded"]["runtime.production_ready"] is True + + +@pytest.mark.parametrize( + ("field", "value", "error_code"), + [ + ( + "actual_backend", + "rlkernel.attention.cp_reference", + BindingErrorCode.ATTENTION_BACKEND_MISSING, + ), + ( + "communication_backend", + "p2p_nccl_reference", + BindingErrorCode.ATTENTION_BACKEND_MISSING, + ), + ("production_ready", False, BindingErrorCode.ATTENTION_NOT_PRODUCTION_READY), + ], +) +def test_strict_runtime_readback_rejects_reference_only_evidence(field, value, error_code): + rollout = replace( + _strict_readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + **{field: value}, + ) + training = _strict_readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="rlkernel.cuda.deterministic_attention", + training_backend_id="rlkernel.cuda.deterministic_attention", + ) + assert not result.passed + assert result.issues_by_code(error_code) + + +def test_strict_runtime_readback_accepts_common_deterministic_preprocess_fallback(): + rollout = replace( + _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + preprocess_fallback=True, + preprocess_fallback_reason="native probe failed", + preprocess_probe_id="rollout-probe", + ) + training = replace( + _readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ), + preprocess_fallback=True, + preprocess_fallback_reason="native probe failed", + preprocess_probe_id="training-probe", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.passed + + +def test_strict_runtime_readback_accepts_distinct_verified_native_preprocess_backends(): + rollout = replace( + _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + preprocess_backends={"qk_rmsnorm": "native.vllm.rmsnorm", "rope": "native.vllm.rope"}, + preprocess_probe_id="rollout-probe", + ) + training = replace( + _readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ), + preprocess_backends={"qk_rmsnorm": "native.te.rmsnorm", "rope": "native.te.rope"}, + preprocess_probe_id="training-probe", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="megatron.te", + ) + + assert result.passed + + +@pytest.mark.parametrize( + ("backends", "fallback", "expected_code"), + [ + ( + {"rope": MANDATED_ATTENTION_PREPROCESS_BACKENDS["rope"]}, + False, + BindingErrorCode.ATTENTION_PREPROCESS_MISSING, + ), + ( + {**MANDATED_ATTENTION_PREPROCESS_BACKENDS, "qk_rmsnorm": "unknown.backend"}, + False, + BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + ), + ( + dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), + True, + BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + ), + ], +) +def test_strict_runtime_readback_rejects_unverified_preprocess_backend( + backends, fallback, expected_code +): + rollout = _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback") + rollout = replace( + rollout, + preprocess_backends=backends, + preprocess_fallback=fallback, + preprocess_fallback_reason=("test fallback" if fallback else None), + ) + training = _readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.comparable + assert not result.passed + assert result.issues_by_code(expected_code) + + +def test_strict_runtime_readback_rejects_projection_split_k_or_missing_plan(): + rollout = _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback") + plans = {name: dict(plan) for name, plan in rollout.projection_plans.items()} + plans["qkv"]["split_k"] = True + plans.pop("o_proj") + rollout = replace(rollout, projection_plans=plans) + training = _readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="megatron.te", + ) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.ATTENTION_PROJECTION_MISMATCH) + assert result.issues_by_code(BindingErrorCode.ATTENTION_PROJECTION_MISSING) + + +def test_strict_runtime_readback_entrypoint_rejects_unverified_frozen_scope(): + rollout_materializer = VllmRolloutMaterializer() + training_materializer = MegatronAttentionMaterializer() + rollout = _readback(rollout_materializer, ROLLOUT_KNOBS, source="vllm.runtime_readback") + rollout = replace(rollout, frozen_scope_verified=False) + training = _readback( + training_materializer, + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.comparable + assert not result.passed + assert any(issue.field == "rollout.frozen_scope_verified" for issue in result.issues) + + +def test_arrival_merge_order_is_unsupported_not_silently_corrected(): + """The control group must stay distinguishable from the treatment.""" + + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_order": "arrival"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_order") == [ + MaterializationStatus.UNSUPPORTED + ] + assert _statuses(materialization, "training.tensor_parallel_size") == [ + MaterializationStatus.ERROR + ] + assert materialization.binding.side_configs["training"]["contract"] is None + assert "arrival" in materialization.binding.side_configs["training"]["contract_error"] + + +def test_unsupported_reduction_invalidates_vllm_contract_applications(): + normalized = { + "batch": {"size": 2}, + "rollout": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_order": "arrival"}, + } + materialization = VllmRolloutMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_order") == [ + MaterializationStatus.UNSUPPORTED + ] + assert _statuses(materialization, "rollout.tensor_parallel_size") == [ + MaterializationStatus.ERROR + ] + assert materialization.binding.side_configs["rollout"]["contract"] is None + + +def test_bf16_reduction_accumulation_is_unsupported(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_acc_dtype": "bf16"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_acc_dtype") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_te_oracle_engine_is_unsupported_until_pr2_pr3(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_engine": "te_oracle"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_engine") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_vllm_cp_falls_back_to_one_in_decode_and_says_why(): + materializer = VllmRolloutMaterializer(mode=AttentionMode.DECODE) + normalized = { + "batch": {"size": 2}, + "rollout": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + + assert materializer.effective_cp_world_size({"rollout.context_parallel_size": 2}) == 1 + + materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) + contract_error = materialization.binding.side_configs["rollout"]["contract_error"] + assert "#235 PR6" in contract_error + assert MaterializationStatus.APPLIED not in {app.status for app in materialization.applications} + + +def test_decode_contract_is_refused_without_kv_cache_identity(): + with pytest.raises(AttentionContractError, match="PR6"): + VllmRolloutMaterializer(mode=AttentionMode.DECODE).build_contract(ROLLOUT_KNOBS) + + +def test_materializers_expose_distinct_implementation_fingerprints(): + megatron = MegatronAttentionMaterializer().implementation_fingerprint + vllm = VllmRolloutMaterializer().implementation_fingerprint + + assert megatron and vllm and megatron != vllm + + +def test_runtime_binding_reports_the_frozen_topology(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + binding = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS).binding + + topology = binding.topology["training"] + assert topology["tensor_parallel_size"] == 2 + assert topology["context_parallel_size"] == 2 + assert topology["world_size"] == 4 + assert topology["pipeline_parallel_size"] == 1 + assert topology["data_parallel_size"] == 1 + + +# -------------------------------------------------------------------------- +# provenance adapters +# -------------------------------------------------------------------------- + + +def test_megatron_provenance_flags_undeclared_frozen_scope_fields(): + adapter = MegatronProvenanceAdapter(SimpleNamespace(deterministic_mode=True)) + + violations = adapter.frozen_scope_violations() + + # Nothing is declared, so every assertion reads as unknown rather than as met. + assert any("expert_model_parallel_size" in text for text in violations) + assert any("fp8" in text for text in violations) + + +def test_megatron_provenance_accepts_a_conforming_dense_config(): + adapter = MegatronProvenanceAdapter( + SimpleNamespace( + deterministic_mode=True, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + sequence_parallel=False, + fp8=None, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + ) + + assert adapter.frozen_scope_violations() == ("fp8 is not declared (expected None)",) + + +def test_megatron_construction_fingerprint_tracks_fusion_changes(): + base = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=False) + fused = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=True) + + assert ( + MegatronProvenanceAdapter(base).construction_fingerprint + != MegatronProvenanceAdapter(fused).construction_fingerprint + ) + + +def test_vllm_provenance_reads_page_size_and_split_kv_policy(): + adapter = VllmProvenanceAdapter( + cache_config=SimpleNamespace(block_size=16, cache_dtype="auto"), + attention_config=SimpleNamespace(flash_attn_max_num_splits_for_cuda_graph=32), + ) + + assert adapter.kv_page_size == 16 + assert adapter.split_kv_policy == 32 + assert adapter.to_dict()["flash_attn_max_num_splits_for_cuda_graph"] == 32 + + +def test_vllm_provenance_flags_fp8_kv_cache_and_cascade_attention(): + adapter = VllmProvenanceAdapter( + model_config=SimpleNamespace(quantization=None, disable_cascade_attn=False), + cache_config=SimpleNamespace( + cache_dtype="fp8", calculate_kv_scales=False, sliding_window=None + ), + parallel_config=SimpleNamespace(pipeline_parallel_size=1, data_parallel_size=1), + ) + + violations = adapter.frozen_scope_violations() + + assert any("cache_dtype" in text for text in violations) + assert any("disable_cascade_attn" in text for text in violations) + + +# -------------------------------------------------------------------------- +# scenario config +# -------------------------------------------------------------------------- + + +SCENARIO = ( + Path(__file__).resolve().parents[1] + / "examples" + / "cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json" +) + + +def test_scenario_uses_megatron_vocabulary_only(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + training = config["baseline"]["training"] + + assert training["attention_backend"] in {"flash", "fused", "unfused", "local", "auto"} + assert training["tensor_parallel_size"] == 2 + assert training["context_parallel_size"] == 2 + assert config["baseline"]["rollout"]["batch_invariant"] is True + + +def test_scenario_knob_paths_all_exist(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + + def paths(mapping, prefix=""): + for key, value in mapping.items(): + path = f"{prefix}{key}" + if isinstance(value, dict): + yield from paths(value, f"{path}.") + else: + yield path + + declared = set(paths(config["baseline"])) + unknown = declared - set(WS2_ATTENTION_KNOBS) + assert not unknown, f"scenario declares unknown knobs: {sorted(unknown)}" + + for intervention in config["interventions"]: + assert intervention["path"] in WS2_ATTENTION_KNOBS diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py new file mode 100644 index 00000000..587c0100 --- /dev/null +++ b/tests/test_attention_preprocess.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.attention_preprocess import ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS, + H100AttentionPreprocessor, +) +from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp + + +def _has_h100_preprocess() -> bool: + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + return False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + required = ( + "rmsnorm_forward", + "rmsnorm_backward_dx", + "rmsnorm_backward_dw", + "rope_apply_sm90", + ) + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in required)) + except ImportError: + return False + + +requires_h100_preprocess = pytest.mark.skipif( + not _has_h100_preprocess(), + reason="Hopper with compiled RMSNorm and RoPE CUDA kernels is required", +) + + +def test_h100_preprocessor_uses_common_backend_ids_for_fallback(): + assert dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) == { + "qk_rmsnorm": "rlkernel.cuda.rmsnorm", + "rope": "rlkernel.cuda.rope_sm90", + } + + +def test_h100_preprocessor_fails_before_dispatch_without_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + with pytest.raises(RuntimeError, match="requires an available CUDA runtime"): + H100AttentionPreprocessor() + + +def test_h100_preprocessor_rejects_non_hopper_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (8, 0)) + with pytest.raises(RuntimeError, match="requires Hopper SM90"): + H100AttentionPreprocessor() + + +def _inputs(): + torch.manual_seed(7) + device = torch.device("cuda") + q = torch.randn(2, 4, 8, 128, device=device, dtype=torch.bfloat16) + k = torch.randn(2, 2, 8, 128, device=device, dtype=torch.bfloat16) + q_weight = torch.randn(128, device=device, dtype=torch.bfloat16) + k_weight = torch.randn(128, device=device, dtype=torch.bfloat16) + positions = torch.tensor( + [[0, 7, 2, 9, 4, 11, 6, 13], [100, 107, 102, 109, 104, 111, 106, 113]], + device=device, + dtype=torch.int64, + ) + return q, k, q_weight, k_weight, positions + + +@requires_h100_preprocess +def test_h100_preprocessor_executes_cuda_qk_norm_and_zigzag_rope(): + q, k, q_weight, k_weight, positions = _inputs() + result = H100AttentionPreprocessor()(q, k, q_weight, k_weight, positions) + + norm = NativeRMSNormOp() + rope = NativeRoPEOp() + q_ref = rope(norm(q, q_weight), positions) + k_ref = rope(norm(k, k_weight), positions) + + assert result.fallback is True + assert dict(result.backend_ids) == dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) + assert result.fallback_reason == "native_preprocess_not_supplied" + assert result.probe_id + assert result.readback_fields() == { + "preprocess_backends": dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), + "preprocess_fallback": True, + "preprocess_fallback_reason": "native_preprocess_not_supplied", + "preprocess_probe_id": result.probe_id, + "preprocess_policy_id": result.policy_id, + } + torch.testing.assert_close(result.q.float(), q_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(result.k.float(), k_ref.float(), atol=2e-2, rtol=2e-2) + + +@requires_h100_preprocess +def test_h100_preprocessor_is_bitwise_batch_invariant_for_2d_positions(): + q, k, q_weight, k_weight, positions = _inputs() + op = H100AttentionPreprocessor() + full = op(q, k, q_weight, k_weight, positions) + + for batch_index in range(q.shape[0]): + single = op( + q[batch_index : batch_index + 1], + k[batch_index : batch_index + 1], + q_weight, + k_weight, + positions[batch_index : batch_index + 1], + ) + assert torch.equal(full.q[batch_index], single.q[0]) + assert torch.equal(full.k[batch_index], single.k[0]) diff --git a/tests/test_attention_projection.py b/tests/test_attention_projection.py new file mode 100644 index 00000000..b4094351 --- /dev/null +++ b/tests/test_attention_projection.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.attention_projection import ( + O_PROJ_COLLECTIVE_CONTRACT, + QKV_COLLECTIVE_CONTRACT, + AttentionProjectionOp, + split_qkv, +) + + +def _deterministic(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return torch.mm(x.float(), weight.float()).to(torch.bfloat16) + + +def _inputs(): + torch.manual_seed(19) + return ( + torch.randn(7, 8, dtype=torch.bfloat16), + torch.randn(8, 12, dtype=torch.bfloat16), + ) + + +@pytest.mark.parametrize( + ("projection", "collective"), + [("qkv", QKV_COLLECTIVE_CONTRACT), ("o_proj", O_PROJ_COLLECTIVE_CONTRACT)], +) +def test_projection_falls_back_to_common_deterministic_path(projection, collective): + x, weight = _inputs() + result = AttentionProjectionOp(projection, deterministic=_deterministic)(x, weight) + + assert torch.equal(result.output, _deterministic(x, weight)) + assert result.plan.backend_id == "rlkernel.cuda.det_gemm" + assert result.plan.fallback is True + assert result.plan.fallback_reason == "native_projection_not_supplied" + assert result.plan.split_k is False + assert result.plan.accumulation_dtype == "torch.float32" + assert dict(result.plan.collective) == collective.to_dict() + + +def test_projection_accepts_native_only_after_bitwise_probe(): + x, weight = _inputs() + result = AttentionProjectionOp( + "qkv", + native=_deterministic, + native_backend_id="megatron.te.qkv", + deterministic=_deterministic, + )(x, weight) + + assert torch.equal(result.output, _deterministic(x, weight)) + assert result.plan.backend_id == "megatron.te.qkv" + assert result.plan.fallback is False + assert result.plan.fallback_reason is None + + +def test_projection_rejects_native_drift_and_records_reason(): + x, weight = _inputs() + + def drifting_native(a, b): + return (_deterministic(a, b).float() + 1.0).to(torch.bfloat16) + + result = AttentionProjectionOp("o_proj", native=drifting_native, deterministic=_deterministic)( + x, weight + ) + + assert result.plan.fallback is True + assert result.plan.fallback_reason == "native_projection_bitwise_probe_failed" + assert torch.equal(result.output, _deterministic(x, weight)) + + +def test_split_qkv_is_fixed_contiguous_q_k_v_order(): + projected = torch.arange(2 * 16, dtype=torch.bfloat16).reshape(2, 16) + q, k, v = split_qkv(projected, q_heads=2, kv_heads=1, head_dim=4) + + assert q.shape == (2, 8) + assert k.shape == (2, 4) + assert v.shape == (2, 4) + assert torch.equal(torch.cat((q, k, v), dim=-1), projected) + + +def test_projection_requires_bf16_and_compatible_k(): + x, weight = _inputs() + with pytest.raises(TypeError, match="BF16"): + AttentionProjectionOp("qkv", deterministic=_deterministic)(x.float(), weight) + with pytest.raises(ValueError, match="K dimensions"): + AttentionProjectionOp("qkv", deterministic=_deterministic)(x, weight[:-1]) + + +def test_o_proj_collective_contract_includes_sp_scatter_gather_and_tp_reduction(): + assert O_PROJ_COLLECTIVE_CONTRACT.sp_forward == "reduce_scatter" + assert O_PROJ_COLLECTIVE_CONTRACT.sp_backward == "all_gather" + assert O_PROJ_COLLECTIVE_CONTRACT.reduction_forward == "all_reduce" + assert O_PROJ_COLLECTIVE_CONTRACT.reduction_backward == "none" diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py new file mode 100644 index 00000000..7dff11c0 --- /dev/null +++ b/tests/test_cp_attention.py @@ -0,0 +1,712 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the WS2 deterministic CP attention reference. + +The implementation is a correctness-first prefill/chunked-prefill reference: +local KV blocks produce ``(out, lse)`` partial states and CP merges those states +with fp32 online-softmax arithmetic in logical global-block order. +""" + +import contextlib +import json +import math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + compare_cp_attention_backward, + merge_attention_partial_states, + split_kv_execution_plan_provenance, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp +from rl_engine.kernels.registry import kernel_registry + +_N_HEADS = 32 +_N_KV = 8 +_HEAD_DIM = 128 +_ATOL = 3.0e-6 +_GRAD_ATOL = 1.0e-5 + + +@contextlib.contextmanager +def _single_thread(): + prev = torch.get_num_threads() + torch.set_num_threads(1) + try: + yield + finally: + torch.set_num_threads(prev) + + +def _qkv( + batch, + sq, + skv, + *, + seed, + dtype=torch.float32, + heads=_N_HEADS, + kv_heads=_N_KV, + dim=_HEAD_DIM, +): + gen = torch.Generator().manual_seed(seed) + q = torch.randn(batch, heads, sq, dim, generator=gen, dtype=dtype) + k = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + v = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + return q, k, v + + +def _full_lse(q, k, *, causal, scale=None, key_padding_mask=None): + qf, kf = q.float(), k.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hq % hkv != 0: + raise ValueError("invalid GQA shape") + if hq != hkv: + kf = kf.repeat_interleave(hq // hkv, dim=1) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * ( + scale if scale is not None else 1.0 / math.sqrt(dim) + ) + if causal: + query_pos = torch.arange(skv - sq, skv) + key_pos = torch.arange(skv) + scores = scores.masked_fill( + (key_pos[None, :] > query_pos[:, None])[None, None, :, :], + float("-inf"), + ) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + return torch.logsumexp(scores, dim=-1) + + +def test_cp1_matches_native_attention_and_exports_lse(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 8, 8, seed=1) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + want = native.forward_fp32(q, k, v, causal=True) + want_lse = _full_lse(q, k, causal=True) + + torch.testing.assert_close(out, want, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse, want_lse, atol=_ATOL, rtol=0.0) + assert lse.dtype == torch.float32 + assert lse.shape == q.shape[:3] + + +def test_cp2_prefill_matches_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 9, 9, seed=2) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + out2, lse2 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=2) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + +def test_cp2_consumes_post_rope_qk_with_shared_global_position_metadata(): + op = DeterministicCPAttentionReferenceOp() + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 7, 7, seed=14, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([17, 103], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + + assert not torch.equal(q, pre_rope_q.float()) + assert not torch.equal(k, pre_rope_k.float()) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=1, + ) + out2, lse2 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + +def test_chunked_prefill_replay_matches_unchunked_cp2(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 10, 10, seed=3) + + with _single_thread(): + unchunked_out, unchunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + ) + chunked_out, chunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=3, + ) + + torch.testing.assert_close(chunked_out, unchunked_out, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(chunked_lse, unchunked_lse, atol=_ATOL, rtol=0.0) + + +def test_causal_mask_uses_global_positions_across_cp_boundary(): + op = DeterministicCPAttentionReferenceOp() + batch, heads, kv_heads, seq, dim = 1, 2, 1, 5, 3 + q = torch.zeros(batch, heads, seq, dim) + k = torch.zeros(batch, kv_heads, seq, dim) + v = torch.arange(seq * dim, dtype=torch.float32).reshape(1, 1, seq, dim) + out = op.forward_fp32(q, k, v, causal=True, cp_world_size=2) + + expected = torch.stack([v[0, 0, : index + 1].mean(dim=0) for index in range(seq)]) + expected = expected.reshape(1, 1, seq, dim).repeat(1, heads, 1, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + + +def test_position_offsets_apply_varlen_causal_metadata_per_batch_row(): + op = DeterministicCPAttentionReferenceOp() + q = torch.zeros(2, 2, 2, 1) + k = torch.zeros(2, 1, 4, 1) + v = torch.arange(8, dtype=torch.float32).reshape(2, 1, 4, 1) + + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=torch.tensor([0, 11]), + key_position_offsets=torch.tensor([0, 10]), + cp_world_size=2, + kv_chunk_size=1, + ) + + expected = torch.tensor([0.0, 0.5, 4.5, 5.0]).reshape(2, 1, 2, 1).repeat(1, 2, 1, 1) + expected_lse = torch.log(torch.tensor([1.0, 2.0, 2.0, 3.0])).reshape(2, 1, 2) + expected_lse = expected_lse.repeat(1, 2, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(lse, expected_lse, atol=1.0e-6, rtol=0.0) + + +def test_merge_order_uses_global_block_index_not_arrival_order(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 6, 6, seed=4) + first = op.local_partial_state( + q, + k[:, :, :3], + v[:, :, :3], + q_start=0, + k_start=0, + total_kv_len=6, + causal=True, + ) + second = op.local_partial_state( + q, + k[:, :, 3:], + v[:, :, 3:], + q_start=0, + k_start=3, + total_kv_len=6, + causal=True, + ) + + forward = merge_attention_partial_states([first, second]) + reversed_arrival = merge_attention_partial_states([second, first]) + assert torch.equal(forward.out, reversed_arrival.out) + assert torch.equal(forward.lse, reversed_arrival.lse) + + +def test_key_padding_mask_and_all_masked_rows_are_stable(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 6, 6, seed=5) + mask = torch.tensor( + [ + [True, True, True, False, False, False], + [False, False, False, False, False, False], + ], + dtype=torch.bool, + ) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=False, + key_padding_mask=mask, + cp_world_size=2, + kv_chunk_size=2, + ) + want = native.forward_fp32(q, k, v, causal=False, key_padding_mask=mask) + + torch.testing.assert_close(out[:1], want[:1], atol=_ATOL, rtol=0.0) + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert torch.isneginf(lse[1]).all() + assert torch.isfinite(out).all() + + +def test_empty_query_and_empty_kv_edges_are_stable(): + op = DeterministicCPAttentionReferenceOp() + q_empty = torch.randn(1, 2, 0, 4, requires_grad=True) + k_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + v_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + out, lse = op.forward_fp32_with_lse(q_empty, k_empty, v_empty, cp_world_size=2) + assert out.shape == (1, 2, 0, 4) + assert lse.shape == (1, 2, 0) + assert out.requires_grad + out.sum().backward() + assert torch.equal(q_empty.grad, torch.zeros_like(q_empty)) + assert torch.equal(k_empty.grad, torch.zeros_like(k_empty)) + assert torch.equal(v_empty.grad, torch.zeros_like(v_empty)) + + q = torch.randn(1, 2, 3, 4) + out, lse = op.forward_fp32_with_lse(q, k_empty, v_empty, causal=False, cp_world_size=4) + assert torch.equal(out, torch.zeros_like(out)) + assert torch.isneginf(lse).all() + + +def test_empty_kv_backward_returns_zero_grads(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 2, 3, 4, requires_grad=True) + k = torch.randn(1, 1, 0, 4, requires_grad=True) + v = torch.randn(1, 1, 0, 4, requires_grad=True) + + out = op.forward_fp32(q, k, v, causal=False, cp_world_size=4) + assert out.requires_grad + out.sum().backward() + + assert torch.equal(q.grad, torch.zeros_like(q)) + assert torch.equal(k.grad, torch.zeros_like(k)) + assert torch.equal(v.grad, torch.zeros_like(v)) + + +def test_bf16_forward_uses_fp32_merge_then_final_write(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 8, 8, seed=6, dtype=torch.bfloat16) + + out, lse = op.forward_with_lse(q, k, v, causal=True, cp_world_size=2, kv_chunk_size=2) + fp32_out, fp32_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + + assert out.dtype == torch.bfloat16 + assert lse.dtype == torch.float32 + assert torch.equal(out, fp32_out.to(torch.bfloat16)) + assert torch.equal(lse, fp32_lse) + + +def test_cp2_chunked_gradients_match_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 5, 5, seed=12, heads=4, kv_heads=2, dim=8) + q_ref = q.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + v_ref = v.detach().clone().requires_grad_(True) + q_cp = q.detach().clone().requires_grad_(True) + k_cp = k.detach().clone().requires_grad_(True) + v_cp = v.detach().clone().requires_grad_(True) + gen = torch.Generator().manual_seed(13) + dy = torch.randn(1, 4, 5, 8, generator=gen) + + with _single_thread(): + out_ref = op.forward_fp32(q_ref, k_ref, v_ref, causal=True, cp_world_size=1) + out_cp = op.forward_fp32( + q_cp, + k_cp, + v_cp, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + out_ref.backward(dy) + out_cp.backward(dy) + + torch.testing.assert_close(out_cp, out_ref, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(q_cp.grad, q_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(k_cp.grad, k_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(v_cp.grad, v_ref.grad, atol=1.0e-5, rtol=0.0) + + +def test_backward_report_cp2_prefill_matches_cp1_reference(): + q, k, v = _qkv(1, 5, 5, seed=15, heads=4, kv_heads=2, dim=8) + dout = torch.randn(1, 4, 5, 8, generator=torch.Generator().manual_seed(16)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + output_dtype=torch.float32, + ) + + assert report.reference_name == "cp1_backward_reference" + drift = report.drifts[0] + assert drift.candidate_name == "cp2_backward" + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + assert drift.out.max_abs <= _ATOL + assert drift.lse.max_abs <= _ATOL + assert len(drift.per_rank) == 2 + assert drift.per_rank[0].dq.active_count > 0 + assert drift.per_rank[1].dk.active_count > 0 + assert drift.provenance["saved_forward_state"][0] == "out" + assert drift.provenance["merge_order"] == "global_block_index" + assert drift.provenance["te_backward_oracle"] == "not_used" + assert drift.provenance["decode_backward"] == "not_supported" + json.dumps(report.to_dict()) + + +def test_backward_report_cp2_chunked_prefill_matches_cp1_reference(): + q, k, v = _qkv(1, 6, 6, seed=17, heads=4, kv_heads=2, dim=8) + dout = torch.randn(1, 4, 6, 8, generator=torch.Generator().manual_seed(18)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.float32, + ) + + drift = report.drifts[0] + assert drift.candidate_name == "cp2_chunked_backward" + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + assert drift.provenance["attention_mode"] == "chunked_prefill" + assert drift.provenance["kv_chunk_size"] == 2 + assert drift.provenance["requested_split_kv_policy"] == "fixed" + assert drift.provenance["actual_split_kv_plans"] == [ + { + "owner_cp_rank": 0, + "requested_split_kv_policy": "fixed", + "requested_split_kv_size": 2, + "actual_split_kv_policy": "fixed", + "actual_split_kv_size": 2, + "actual_split_kv_count": 2, + "actual_split_boundaries": [[0, 2], [2, 3]], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "deterministic_cp_backward_reference", + "split_kv_plan_source": "reference_execution", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + }, + { + "owner_cp_rank": 1, + "requested_split_kv_policy": "fixed", + "requested_split_kv_size": 2, + "actual_split_kv_policy": "fixed", + "actual_split_kv_size": 2, + "actual_split_kv_count": 2, + "actual_split_boundaries": [[3, 5], [5, 6]], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "deterministic_cp_backward_reference", + "split_kv_plan_source": "reference_execution", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + }, + ] + + +def test_split_kv_plan_never_crosses_cp_owner_boundaries(): + plans = split_kv_execution_plan_provenance( + 10, + cp_world_size=3, + kv_chunk_size=3, + backend="test-reference", + ) + + assert [plan["actual_split_boundaries"] for plan in plans] == [ + [[0, 3], [3, 4]], + [[4, 7]], + [[7, 10]], + ] + assert [plan["owner_cp_rank"] for plan in plans] == [0, 1, 2] + + +def test_backward_report_preserves_post_rope_position_metadata(): + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 5, 5, seed=19, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([23, 101], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + dout = torch.randn(2, 4, 5, 8, generator=torch.Generator().manual_seed(20)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.float32, + ) + + drift = report.drifts[0] + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + + +def test_qwen3_8b_local_tp2_cp2_bf16_backward_report_smoke(): + # Qwen3-8B global Hq/Hkv is 32/8. A TP=2 local shard owns 16/4 heads. + q, k, v = _qkv( + 1, + 4, + 4, + seed=21, + dtype=torch.bfloat16, + heads=16, + kv_heads=4, + dim=_HEAD_DIM, + ) + dout = torch.randn( + 1, + 16, + 4, + _HEAD_DIM, + generator=torch.Generator().manual_seed(22), + dtype=torch.bfloat16, + ) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.bfloat16, + ) + + drift = report.drifts[0] + assert drift.provenance["q_dtype"] == "bfloat16" + assert drift.provenance["output_dtype"] == "bfloat16" + assert drift.provenance["downcast_at"] == "final_write" + assert drift.dq.max_abs <= 5.0e-2 + assert drift.dk.max_abs <= 5.0e-2 + assert drift.dv.max_abs <= 5.0e-2 + + +def test_backward_report_validates_dout_shape_and_dtype(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=23, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="dout must have shape"): + op.backward_reference(q, k, v, torch.randn(1, 4, 3, 8), cp_world_size=2) + + with pytest.raises(ValueError, match="dout must be a real floating-point tensor"): + op.backward_reference( + q, + k, + v, + torch.ones(1, 4, 4, 8, dtype=torch.long), + cp_world_size=2, + ) + + with pytest.raises(ValueError, match="dout must have the same dtype"): + op.backward_reference( + q.to(torch.bfloat16), + k.to(torch.bfloat16), + v.to(torch.bfloat16), + torch.ones_like(q), + cp_world_size=2, + ) + + +def test_inputs_are_not_mutated(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 6, 6, seed=7) + mask = torch.ones(2, 6, dtype=torch.bool) + qc, kc, vc, mc = q.clone(), k.clone(), v.clone(), mask.clone() + + op.forward_fp32_with_lse(q, k, v, causal=True, key_padding_mask=mask, cp_world_size=2) + + assert torch.equal(q, qc) + assert torch.equal(k, kc) + assert torch.equal(v, vc) + assert torch.equal(mask, mc) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"cp_world_size": 0}, "cp_world_size"), + ({"cp_world_size": 2, "kv_chunk_size": 0}, "kv_chunk_size"), + ], +) +def test_invalid_parallelism_arguments_raise(kwargs, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=8) + with pytest.raises(ValueError, match=message): + op.forward_fp32_with_lse(q, k, v, causal=True, **kwargs) + + +def test_invalid_gqa_and_mask_shapes_raise(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 6, 4, _HEAD_DIM) + k = torch.randn(1, 4, 4, _HEAD_DIM) + v = torch.randn(1, 4, 4, _HEAD_DIM) + with pytest.raises(ValueError, match="not divisible"): + op.forward_fp32_with_lse(q, k, v, causal=True) + + q, k, v = _qkv(1, 4, 4, seed=9) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 3, dtype=torch.bool)) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 4)) + with pytest.raises(ValueError, match="query_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + query_position_offsets=torch.ones(2, dtype=torch.long), + ) + with pytest.raises(ValueError, match="key_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + key_position_offsets=torch.ones(1, dtype=torch.float32), + ) + + +@pytest.mark.parametrize("scale", [0.0, -1.0, float("nan"), float("inf"), True, "bad"]) +def test_invalid_scale_fails_before_attention_math(scale): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=24, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="scale must be a positive finite number"): + op.forward_fp32_with_lse(q, k, v, scale=scale) + + +def test_qkv_dtype_and_floating_contract_fails_closed(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=25, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="same dtype"): + op.forward_fp32_with_lse(q, k.to(torch.bfloat16), v) + with pytest.raises(ValueError, match="real floating-point"): + op.forward_fp32_with_lse(q.to(torch.long), k.to(torch.long), v.to(torch.long)) + + +@pytest.mark.parametrize("kwargs", [{"cp_world_size": True}, {"kv_chunk_size": True}]) +def test_boolean_parallelism_arguments_fail_closed(kwargs): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=26, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError): + op.forward_fp32_with_lse(q, k, v, **kwargs) + + +@pytest.mark.parametrize("output_dtype", [torch.long, torch.complex64, "fp32"]) +def test_nonfloating_output_dtype_fails_closed(output_dtype): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=27, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="output_dtype must be a real floating-point"): + op.forward_with_lse(q, k, v, output_dtype=output_dtype) + + +def test_partial_states_must_remain_fp32_and_colocated(): + out = torch.zeros(1, 1, 1, 1, dtype=torch.bfloat16) + lse = torch.zeros(1, 1, 1, dtype=torch.float32) + + with pytest.raises(ValueError, match="must remain FP32"): + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=1) + + +def test_overlapping_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="overlap"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=3), + AttentionPartialState(out=out, lse=lse, block_start=2, block_end=4), + ] + ) + + +def test_gapped_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="gap-free"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=2), + AttentionPartialState(out=out, lse=lse, block_start=3, block_end=4), + ] + ) + + +def test_registry_dispatches_cp_attention_reference(): + assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) + + +def test_strict_reference_is_bitwise_across_batch_cp_and_backward(): + q, k, v = _qkv(2, 4, 8, seed=41, heads=4, kv_heads=2, dim=8) + dout = torch.randn_like(q) + op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + + cp1_out, cp1_lse = op.forward_with_lse(q, k, v, cp_world_size=1, kv_chunk_size=3) + cp2_out, cp2_lse = op.forward_with_lse(q, k, v, cp_world_size=2, kv_chunk_size=3) + single_out, single_lse = op.forward_with_lse( + q[:1], k[:1], v[:1], cp_world_size=1, kv_chunk_size=3 + ) + assert torch.equal(cp1_out, cp2_out) + assert torch.equal(cp1_lse, cp2_lse) + assert torch.equal(cp1_out[:1], single_out) + assert torch.equal(cp1_lse[:1], single_lse) + + cp1 = op.backward_reference(q, k, v, dout, cp_world_size=1, kv_chunk_size=3) + cp2 = op.backward_reference(q, k, v, dout, cp_world_size=2, kv_chunk_size=3) + assert torch.equal(cp1.gradients.dq, cp2.gradients.dq) + assert torch.equal(cp1.gradients.dk, cp2.gradients.dk) + assert torch.equal(cp1.gradients.dv, cp2.gradients.dv) diff --git a/tests/test_cp_attention_transformer_engine.py b/tests/test_cp_attention_transformer_engine.py new file mode 100644 index 00000000..d98e31a1 --- /dev/null +++ b/tests/test_cp_attention_transformer_engine.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Optional Transformer Engine oracle tests for CP attention merging.""" + +from __future__ import annotations + +import importlib + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + merge_attention_partial_states, +) + + +def _te_context_parallel_module(): + try: + return importlib.import_module( + "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" + ) + except (ImportError, OSError, RuntimeError) as exc: + pytest.skip(f"Transformer Engine context-parallel attention is unavailable: {exc}") + + +def test_cp_attention_merge_matches_transformer_engine_corrections(): + te_cp = _te_context_parallel_module() + gen = torch.Generator().manual_seed(238) + out_a = torch.randn(2, 3, 5, 4, generator=gen) + out_b = torch.randn(2, 3, 5, 4, generator=gen) + lse_a = torch.randn(2, 3, 5, generator=gen) + lse_b = torch.randn(2, 3, 5, generator=gen) + + ours = merge_attention_partial_states( + [ + AttentionPartialState(out=out_b, lse=lse_b, block_start=5, block_end=9), + AttentionPartialState(out=out_a, lse=lse_a, block_start=0, block_end=5), + ] + ) + + te_lse = lse_a.clone() + te_cp.flash_attn_fwd_softmax_lse_correction(te_lse, lse_b) + te_out = te_cp.flash_attn_fwd_out_correction_init(out_a.clone(), te_lse, lse_a, seq_dim=2) + te_cp.flash_attn_fwd_out_correction(te_out, out_b, te_lse, lse_b, seq_dim=2) + + torch.testing.assert_close(ours.lse, te_lse, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(ours.out, te_out, atol=1.0e-6, rtol=0.0) From d7a863acba066321edc81c89c6193b68250c4fbe Mon Sep 17 00:00:00 2001 From: Takumi <3051000145@qq.com> Date: Wed, 19 Aug 2026 15:54:01 +0800 Subject: [PATCH 13/33] fix(attention): support batched RoPE positions --- rl_engine/kernels/attention_preprocess.py | 566 +++---- .../kernels/ops/pytorch/attention/ablation.py | 1307 +++++++++-------- 2 files changed, 948 insertions(+), 925 deletions(-) diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index e97b365f..a05a8d41 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -1,275 +1,291 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -"""Bitwise-bound QK-Norm and RoPE handoff for WS2 Attention. - -Megatron/TE and vLLM/FlashInfer remain the first-choice implementations. They -are admitted only after the same-input H100 probe is bitwise identical to the -deterministic RL-Kernel path. A failed probe (or an unavailable native backend) -selects the common RL-Kernel path for both sides and records why the fallback -was taken. The caller must pass this readback to the cross-config binder. -""" - -from __future__ import annotations - -import hashlib -import json -from dataclasses import dataclass -from types import MappingProxyType -from typing import Any, Callable, Mapping - -import torch -from torch import Tensor - -QK_RMSNORM_BACKEND_ID = "rlkernel.cuda.rmsnorm" -ROPE_BACKEND_ID = "rlkernel.cuda.rope_sm90" -NATIVE_QK_RMSNORM_BACKEND_ID = "native.qk_rmsnorm" -NATIVE_ROPE_BACKEND_ID = "native.rope" -PREPROCESS_POLICY_ID = "ws2.attention.preprocess.v2" -MANDATED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( - { - "qk_rmsnorm": QK_RMSNORM_BACKEND_ID, - "rope": ROPE_BACKEND_ID, - } -) - - -@dataclass(frozen=True) -class AttentionPreprocessResult: - """Post-QK-Norm, post-RoPE tensors plus executed backend evidence.""" - - q: Tensor - k: Tensor - backend_ids: Mapping[str, str] - fallback: bool - device_capability: tuple[int, int] - fallback_reason: str | None = None - probe_id: str = "" - policy_id: str = PREPROCESS_POLICY_ID - - def __post_init__(self) -> None: - object.__setattr__(self, "backend_ids", MappingProxyType(dict(self.backend_ids))) - - def evidence(self) -> dict[str, Any]: - return { - "backends": dict(self.backend_ids), - "fallback": self.fallback, - "device_capability": list(self.device_capability), - } - - def readback_fields(self) -> dict[str, Any]: - """Keyword fields consumed by ``AttentionRuntimeReadback``.""" - - return { - "preprocess_backends": dict(self.backend_ids), - "preprocess_fallback": self.fallback, - "preprocess_fallback_reason": self.fallback_reason, - "preprocess_probe_id": self.probe_id, - "preprocess_policy_id": self.policy_id, - } - - -class H100AttentionPreprocessor: - """Apply native QK-Norm/RoPE when the H100 probe passes. - - ``native_qk_norm`` and ``native_rope`` are framework-owned callables. They - are intentionally injected instead of importing TE/vLLM here, so the same - policy can be used by both runtimes. The deterministic callables default to - RL-Kernel's CUDA operators and are always run to establish the probe oracle. - """ - - def __init__( - self, - device: torch.device | str | int | None = None, - *, - native_qk_norm: Callable[..., Tensor] | None = None, - native_rope: Callable[..., Tensor] | None = None, - native_qk_norm_backend_id: str = NATIVE_QK_RMSNORM_BACKEND_ID, - native_rope_backend_id: str = NATIVE_ROPE_BACKEND_ID, - policy_id: str = PREPROCESS_POLICY_ID, - ) -> None: - if not torch.cuda.is_available(): - raise RuntimeError("H100AttentionPreprocessor requires an available CUDA runtime") - - current_device = torch.cuda.current_device() - self.device = torch.device("cuda", current_device) - if device is not None: - self.device = ( - torch.device("cuda", device) if isinstance(device, int) else torch.device(device) - ) - if self.device.type != "cuda": - raise RuntimeError(f"H100AttentionPreprocessor requires CUDA, got {self.device}") - if self.device.index is None: - self.device = torch.device("cuda", current_device) - - capability = torch.cuda.get_device_capability(self.device) - self.device_capability: tuple[int, int] = (int(capability[0]), int(capability[1])) - if self.device_capability[0] != 9: - raise RuntimeError( - "H100AttentionPreprocessor requires Hopper SM90; " - f"got sm_{self.device_capability[0]}{self.device_capability[1]}" - ) - - # Import only after the hardware gate so CPU tools can inspect the module. - from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp - from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op - - self.rmsnorm = RMSNormCudaOp() - self.rope = RoPESM90Op() - if not isinstance(policy_id, str) or not policy_id.strip(): - raise ValueError("policy_id must be a non-empty string") - self.native_qk_norm = native_qk_norm - self.native_rope = native_rope - self.native_qk_norm_backend_id = native_qk_norm_backend_id - self.native_rope_backend_id = native_rope_backend_id - self.policy_id = policy_id - - def __call__( - self, - q: Tensor, - k: Tensor, - q_weight: Tensor, - k_weight: Tensor, - positions: Tensor, - *, - eps: float = 1.0e-6, - theta: float = 1_000_000.0, - ) -> AttentionPreprocessResult: - return self.forward( - q, - k, - q_weight, - k_weight, - positions, - eps=eps, - theta=theta, - ) - - def forward( - self, - q: Tensor, - k: Tensor, - q_weight: Tensor, - k_weight: Tensor, - positions: Tensor, - *, - eps: float = 1.0e-6, - theta: float = 1_000_000.0, - ) -> AttentionPreprocessResult: - _validate_inputs(q, k, q_weight, k_weight, positions, self.device) - q_norm_det = self.rmsnorm(q, q_weight, eps=eps) - k_norm_det = self.rmsnorm(k, k_weight, eps=eps) - q_det = self.rope(q_norm_det, positions, theta=theta) - k_det = self.rope(k_norm_det, positions, theta=theta) - - native_qk_norm = self.native_qk_norm - native_rope = self.native_rope - if native_qk_norm is not None and native_rope is not None: - try: - q_norm_native = native_qk_norm(q, q_weight, eps=eps) - k_norm_native = native_qk_norm(k, k_weight, eps=eps) - q_native = native_rope(q_norm_native, positions, theta=theta) - k_native = native_rope(k_norm_native, positions, theta=theta) - probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) - if ( - torch.equal(q_norm_native, q_norm_det) - and torch.equal(k_norm_native, k_norm_det) - and torch.equal(q_native, q_det) - and torch.equal(k_native, k_det) - ): - return AttentionPreprocessResult( - q=q_native, - k=k_native, - backend_ids=MappingProxyType( - { - "qk_rmsnorm": self.native_qk_norm_backend_id, - "rope": self.native_rope_backend_id, - } - ), - fallback=False, - device_capability=self.device_capability, - probe_id=probe_id, - policy_id=self.policy_id, - ) - fallback_reason = "native_preprocess_bitwise_probe_failed" - except Exception as exc: # framework backend failures must fail over together - fallback_reason = f"native_preprocess_unavailable:{type(exc).__name__}" - probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) - else: - fallback_reason = "native_preprocess_not_supplied" - probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) - return AttentionPreprocessResult( - q=q_det, - k=k_det, - backend_ids=MANDATED_ATTENTION_PREPROCESS_BACKENDS, - fallback=True, - device_capability=self.device_capability, - fallback_reason=fallback_reason, - probe_id=probe_id, - policy_id=self.policy_id, - ) - - -def _probe_id( - q: Tensor, - k: Tensor, - q_weight: Tensor, - k_weight: Tensor, - positions: Tensor, - eps: float, - theta: float, -) -> str: - payload = { - "q_shape": list(q.shape), - "k_shape": list(k.shape), - "q_dtype": str(q.dtype), - "k_dtype": str(k.dtype), - "weight_dtype": str(q_weight.dtype), - "positions_shape": list(positions.shape), - "positions_sha256": hashlib.sha256(positions.detach().cpu().numpy().tobytes()).hexdigest(), - "eps": float(eps), - "theta": float(theta), - } - return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16] - - -def _validate_inputs( - q: Tensor, - k: Tensor, - q_weight: Tensor, - k_weight: Tensor, - positions: Tensor, - device: torch.device, -) -> None: - if q.dim() != 4 or k.dim() != 4: - raise ValueError("q and k must use [B, H, S, D] layout") - if q.shape[0] != k.shape[0] or q.shape[-2:] != k.shape[-2:]: - raise ValueError("q and k must have the same batch, sequence, and head dimensions") - if q.dtype is not torch.bfloat16 or k.dtype is not torch.bfloat16: - raise TypeError("the frozen H100 Attention experiment requires BF16 q and k") - if q.device != device or k.device != device: - raise ValueError(f"q and k must both be on the configured device {device}") - for name, weight in (("q_weight", q_weight), ("k_weight", k_weight)): - if weight.shape != (q.shape[-1],): - raise ValueError(f"{name} must have shape ({q.shape[-1]},)") - if weight.device != device or weight.dtype is not torch.bfloat16: - raise ValueError(f"{name} must be BF16 on {device}") - if positions.device != device: - raise ValueError(f"positions must be on {device}") - if positions.dtype not in (torch.int32, torch.int64): - raise TypeError("positions must use int32 or int64 global token indices") - expected = (q.shape[-2],) if positions.dim() == 1 else (q.shape[0], q.shape[-2]) - if positions.dim() not in (1, 2) or tuple(positions.shape) != expected: - raise ValueError(f"positions must have shape [S] or [B, S], expected {expected}") - - -__all__ = [ - "AttentionPreprocessResult", - "H100AttentionPreprocessor", - "MANDATED_ATTENTION_PREPROCESS_BACKENDS", - "QK_RMSNORM_BACKEND_ID", - "ROPE_BACKEND_ID", - "NATIVE_QK_RMSNORM_BACKEND_ID", - "NATIVE_ROPE_BACKEND_ID", - "PREPROCESS_POLICY_ID", -] +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Bitwise-bound QK-Norm and RoPE handoff for WS2 Attention. + +Megatron/TE and vLLM/FlashInfer remain the first-choice implementations. They +are admitted only after the same-input H100 probe is bitwise identical to the +deterministic RL-Kernel path. A failed probe (or an unavailable native backend) +selects the common RL-Kernel path for both sides and records why the fallback +was taken. The caller must pass this readback to the cross-config binder. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Callable, Mapping + +import torch +from torch import Tensor + +QK_RMSNORM_BACKEND_ID = "rlkernel.cuda.rmsnorm" +ROPE_BACKEND_ID = "rlkernel.cuda.rope_sm90" +NATIVE_QK_RMSNORM_BACKEND_ID = "native.qk_rmsnorm" +NATIVE_ROPE_BACKEND_ID = "native.rope" +PREPROCESS_POLICY_ID = "ws2.attention.preprocess.v2" +MANDATED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( + { + "qk_rmsnorm": QK_RMSNORM_BACKEND_ID, + "rope": ROPE_BACKEND_ID, + } +) + + +@dataclass(frozen=True) +class AttentionPreprocessResult: + """Post-QK-Norm, post-RoPE tensors plus executed backend evidence.""" + + q: Tensor + k: Tensor + backend_ids: Mapping[str, str] + fallback: bool + device_capability: tuple[int, int] + fallback_reason: str | None = None + probe_id: str = "" + policy_id: str = PREPROCESS_POLICY_ID + + def __post_init__(self) -> None: + object.__setattr__(self, "backend_ids", MappingProxyType(dict(self.backend_ids))) + + def evidence(self) -> dict[str, Any]: + return { + "backends": dict(self.backend_ids), + "fallback": self.fallback, + "device_capability": list(self.device_capability), + } + + def readback_fields(self) -> dict[str, Any]: + """Keyword fields consumed by ``AttentionRuntimeReadback``.""" + + return { + "preprocess_backends": dict(self.backend_ids), + "preprocess_fallback": self.fallback, + "preprocess_fallback_reason": self.fallback_reason, + "preprocess_probe_id": self.probe_id, + "preprocess_policy_id": self.policy_id, + } + + +class H100AttentionPreprocessor: + """Apply native QK-Norm/RoPE when the H100 probe passes. + + ``native_qk_norm`` and ``native_rope`` are framework-owned callables. They + are intentionally injected instead of importing TE/vLLM here, so the same + policy can be used by both runtimes. The deterministic callables default to + RL-Kernel's CUDA operators and are always run to establish the probe oracle. + """ + + def __init__( + self, + device: torch.device | str | int | None = None, + *, + native_qk_norm: Callable[..., Tensor] | None = None, + native_rope: Callable[..., Tensor] | None = None, + native_qk_norm_backend_id: str = NATIVE_QK_RMSNORM_BACKEND_ID, + native_rope_backend_id: str = NATIVE_ROPE_BACKEND_ID, + policy_id: str = PREPROCESS_POLICY_ID, + ) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("H100AttentionPreprocessor requires an available CUDA runtime") + + current_device = torch.cuda.current_device() + self.device = torch.device("cuda", current_device) + if device is not None: + self.device = ( + torch.device("cuda", device) if isinstance(device, int) else torch.device(device) + ) + if self.device.type != "cuda": + raise RuntimeError(f"H100AttentionPreprocessor requires CUDA, got {self.device}") + if self.device.index is None: + self.device = torch.device("cuda", current_device) + + capability = torch.cuda.get_device_capability(self.device) + self.device_capability: tuple[int, int] = (int(capability[0]), int(capability[1])) + if self.device_capability[0] != 9: + raise RuntimeError( + "H100AttentionPreprocessor requires Hopper SM90; " + f"got sm_{self.device_capability[0]}{self.device_capability[1]}" + ) + + # Import only after the hardware gate so CPU tools can inspect the module. + from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op + + self.rmsnorm = RMSNormCudaOp() + self.rope = RoPESM90Op() + if not isinstance(policy_id, str) or not policy_id.strip(): + raise ValueError("policy_id must be a non-empty string") + self.native_qk_norm = native_qk_norm + self.native_rope = native_rope + self.native_qk_norm_backend_id = native_qk_norm_backend_id + self.native_rope_backend_id = native_rope_backend_id + self.policy_id = policy_id + + def __call__( + self, + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + *, + eps: float = 1.0e-6, + theta: float = 1_000_000.0, + ) -> AttentionPreprocessResult: + return self.forward( + q, + k, + q_weight, + k_weight, + positions, + eps=eps, + theta=theta, + ) + + def forward( + self, + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + *, + eps: float = 1.0e-6, + theta: float = 1_000_000.0, + ) -> AttentionPreprocessResult: + _validate_inputs(q, k, q_weight, k_weight, positions, self.device) + q_norm_det = self.rmsnorm(q, q_weight, eps=eps) + k_norm_det = self.rmsnorm(k, k_weight, eps=eps) + q_det = _apply_deterministic_rope(self.rope, q_norm_det, positions, theta) + k_det = _apply_deterministic_rope(self.rope, k_norm_det, positions, theta) + + native_qk_norm = self.native_qk_norm + native_rope = self.native_rope + if native_qk_norm is not None and native_rope is not None: + try: + q_norm_native = native_qk_norm(q, q_weight, eps=eps) + k_norm_native = native_qk_norm(k, k_weight, eps=eps) + q_native = native_rope(q_norm_native, positions, theta=theta) + k_native = native_rope(k_norm_native, positions, theta=theta) + probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) + if ( + torch.equal(q_norm_native, q_norm_det) + and torch.equal(k_norm_native, k_norm_det) + and torch.equal(q_native, q_det) + and torch.equal(k_native, k_det) + ): + return AttentionPreprocessResult( + q=q_native, + k=k_native, + backend_ids=MappingProxyType( + { + "qk_rmsnorm": self.native_qk_norm_backend_id, + "rope": self.native_rope_backend_id, + } + ), + fallback=False, + device_capability=self.device_capability, + probe_id=probe_id, + policy_id=self.policy_id, + ) + fallback_reason = "native_preprocess_bitwise_probe_failed" + except Exception as exc: # framework backend failures must fail over together + fallback_reason = f"native_preprocess_unavailable:{type(exc).__name__}" + probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) + else: + fallback_reason = "native_preprocess_not_supplied" + probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) + return AttentionPreprocessResult( + q=q_det, + k=k_det, + backend_ids=MANDATED_ATTENTION_PREPROCESS_BACKENDS, + fallback=True, + device_capability=self.device_capability, + fallback_reason=fallback_reason, + probe_id=probe_id, + policy_id=self.policy_id, + ) + + +def _apply_deterministic_rope( + rope: Callable[..., Tensor], + x: Tensor, + positions: Tensor, + theta: float, +) -> Tensor: + """Adapt per-sample positions to the CUDA RoPE operator's 1-D contract.""" + + if positions.dim() == 1: + return rope(x, positions, theta=theta) + return torch.cat( + [rope(x[index : index + 1], positions[index], theta=theta) for index in range(x.shape[0])], + dim=0, + ) + + +def _probe_id( + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + eps: float, + theta: float, +) -> str: + payload = { + "q_shape": list(q.shape), + "k_shape": list(k.shape), + "q_dtype": str(q.dtype), + "k_dtype": str(k.dtype), + "weight_dtype": str(q_weight.dtype), + "positions_shape": list(positions.shape), + "positions_sha256": hashlib.sha256(positions.detach().cpu().numpy().tobytes()).hexdigest(), + "eps": float(eps), + "theta": float(theta), + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16] + + +def _validate_inputs( + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + device: torch.device, +) -> None: + if q.dim() != 4 or k.dim() != 4: + raise ValueError("q and k must use [B, H, S, D] layout") + if q.shape[0] != k.shape[0] or q.shape[-2:] != k.shape[-2:]: + raise ValueError("q and k must have the same batch, sequence, and head dimensions") + if q.dtype is not torch.bfloat16 or k.dtype is not torch.bfloat16: + raise TypeError("the frozen H100 Attention experiment requires BF16 q and k") + if q.device != device or k.device != device: + raise ValueError(f"q and k must both be on the configured device {device}") + for name, weight in (("q_weight", q_weight), ("k_weight", k_weight)): + if weight.shape != (q.shape[-1],): + raise ValueError(f"{name} must have shape ({q.shape[-1]},)") + if weight.device != device or weight.dtype is not torch.bfloat16: + raise ValueError(f"{name} must be BF16 on {device}") + if positions.device != device: + raise ValueError(f"positions must be on {device}") + if positions.dtype not in (torch.int32, torch.int64): + raise TypeError("positions must use int32 or int64 global token indices") + expected = (q.shape[-2],) if positions.dim() == 1 else (q.shape[0], q.shape[-2]) + if positions.dim() not in (1, 2) or tuple(positions.shape) != expected: + raise ValueError(f"positions must have shape [S] or [B, S], expected {expected}") + + +__all__ = [ + "AttentionPreprocessResult", + "H100AttentionPreprocessor", + "MANDATED_ATTENTION_PREPROCESS_BACKENDS", + "QK_RMSNORM_BACKEND_ID", + "ROPE_BACKEND_ID", + "NATIVE_QK_RMSNORM_BACKEND_ID", + "NATIVE_ROPE_BACKEND_ID", + "PREPROCESS_POLICY_ID", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index dd463906..3f8727f7 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -1,650 +1,657 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -"""Unified Attention entry point for the PR230 cross-configuration matrix. - -The matrix needs one stable callable shape even though training and rollout may -materialize different Attention backends. This adapter owns the common -contract checks and provenance only; numerical work remains in the existing -deterministic Attention implementations. -""" - -from __future__ import annotations - -import hashlib -import inspect -import json -import math -from dataclasses import dataclass, field -from types import MappingProxyType -from typing import Any, Callable, Mapping - -import torch -from torch import Tensor - -from rl_engine.kernels.attention_contract import ( - STRICT_ATTENTION_CORE_ID, - STRICT_ATTENTION_SCHEDULE_ID, - AttentionContract, - AttentionContractError, - AttentionDType, - SplitKVMode, -) - -BACKEND_ID = "rlkernel.attention.deterministic.v1" -REFERENCE_BACKEND_ID = "rlkernel.attention.reference.v1" - -_TORCH_DTYPES = { - AttentionDType.BF16: torch.bfloat16, - AttentionDType.FP16: torch.float16, - AttentionDType.FP32: torch.float32, -} - - -@dataclass(frozen=True) -class AttentionAblationConfig: - """Per-invocation settings materialized by the ablation runner.""" - - backend: str = "auto" - deterministic: bool = True - communication_backend: str = "none" - return_lse: bool = True - return_gradients: bool = False - strict_core_id: str = STRICT_ATTENTION_CORE_ID - strict_schedule: str = STRICT_ATTENTION_SCHEDULE_ID - validate: bool = True - - def __post_init__(self) -> None: - if not isinstance(self.backend, str) or not self.backend.strip(): - raise AttentionContractError("Attention backend must be a non-empty string") - if not isinstance(self.deterministic, bool): - raise AttentionContractError("deterministic must be a bool") - if ( - not isinstance(self.communication_backend, str) - or not self.communication_backend.strip() - ): - raise AttentionContractError("communication_backend must be a non-empty string") - for name in ("return_lse", "return_gradients", "validate"): - if not isinstance(getattr(self, name), bool): - raise AttentionContractError(f"{name} must be a bool") - if not isinstance(self.strict_core_id, str) or not self.strict_core_id.strip(): - raise AttentionContractError("strict_core_id must be a non-empty string") - if not isinstance(self.strict_schedule, str) or not self.strict_schedule.strip(): - raise AttentionContractError("strict_schedule must be a non-empty string") - object.__setattr__(self, "backend", self.backend.strip().lower()) - object.__setattr__(self, "communication_backend", self.communication_backend.strip()) - - -@dataclass(frozen=True) -class AttentionAblationResult: - """Standardized Attention result consumed by cross-config artifacts.""" - - out: Tensor - lse: Tensor | None - dq: Tensor | None = None - dk: Tensor | None = None - dv: Tensor | None = None - backend_id: str = BACKEND_ID - deterministic: bool = True - provenance: Mapping[str, Any] = field(default_factory=dict) - - def __post_init__(self) -> None: - if not isinstance(self.out, Tensor): - raise TypeError("Attention result out must be a torch.Tensor") - if self.lse is not None and not isinstance(self.lse, Tensor): - raise TypeError("Attention result lse must be a torch.Tensor or None") - for name in ("dq", "dk", "dv"): - value = getattr(self, name) - if value is not None and not isinstance(value, Tensor): - raise TypeError(f"Attention result {name} must be a torch.Tensor or None") - if not isinstance(self.backend_id, str) or not self.backend_id.strip(): - raise ValueError("Attention result backend_id must be non-empty") - if not isinstance(self.deterministic, bool): - raise TypeError("Attention result deterministic must be a bool") - if not isinstance(self.provenance, Mapping): - raise TypeError("Attention result provenance must be a mapping") - object.__setattr__(self, "provenance", MappingProxyType(dict(self.provenance))) - - @property - def out_lse(self) -> tuple[Tensor, Tensor | None]: - """Compatibility tuple for callers that consume ``(out, lse)``.""" - - return self.out, self.lse - - def readback(self) -> dict[str, Any]: - """Return JSON-compatible execution evidence for PR230 artifacts.""" - - return { - "backend_id": self.backend_id, - "deterministic": self.deterministic, - "out_shape": list(self.out.shape), - "out_dtype": str(self.out.dtype).replace("torch.", ""), - "lse_shape": None if self.lse is None else list(self.lse.shape), - "lse_dtype": None if self.lse is None else str(self.lse.dtype).replace("torch.", ""), - "gradients": { - "dq": self.dq is not None, - "dk": self.dk is not None, - "dv": self.dv is not None, - }, - "provenance": dict(self.provenance), - } - - -class AttentionAblationOp: - """PR230/PR314-style unified Attention wrapper. - - ``core`` and ``reference`` are injectable so the wrapper is usable by the - semantic operator session without importing CUDA at construction time. - ``core`` should expose ``forward_with_lse``; ``reference`` is the existing - pure-PyTorch CP reference with the same method. - """ - - op_class = "attention" - is_batch_invariant = True - backend_id = BACKEND_ID - - def __init__( - self, - *, - core: Any | None = None, - reference: Any | None = None, - native: Any | None = None, - cp_backend: Any | None = None, - communication_backend: str = "none", - ) -> None: - if not isinstance(communication_backend, str) or not communication_backend.strip(): - raise AttentionContractError("communication_backend must be a non-empty string") - self.core = core - self.reference = reference - self.native = native - # CP production execution is injected by the runtime adapter. Keeping - # it separate from the single-device core prevents an accidental - # fallback to the PyTorch reference when AG/RS is required. - self.cp_backend = cp_backend - self.communication_backend = communication_backend.strip() - - def __call__( - self, - q: Tensor, - k: Tensor, - v: Tensor, - *, - contract: AttentionContract, - config: AttentionAblationConfig | Mapping[str, Any] | None = None, - backend: str | Callable[..., Any] | None = None, - deterministic: bool | None = None, - return_lse: bool | None = None, - return_gradients: bool | None = None, - dout: Tensor | None = None, - communication_backend: str | None = None, - validate: bool | None = None, - **kwargs: Any, - ) -> AttentionAblationResult: - return self.apply( - q, - k, - v, - contract=contract, - config=config, - backend=backend, - deterministic=deterministic, - return_lse=return_lse, - return_gradients=return_gradients, - dout=dout, - communication_backend=communication_backend, - validate=validate, - **kwargs, - ) - - def apply( - self, - q: Tensor, - k: Tensor, - v: Tensor, - *, - contract: AttentionContract, - config: AttentionAblationConfig | Mapping[str, Any] | None = None, - backend: str | Callable[..., Any] | None = None, - deterministic: bool | None = None, - return_lse: bool | None = None, - return_gradients: bool | None = None, - dout: Tensor | None = None, - communication_backend: str | None = None, - validate: bool | None = None, - **kwargs: Any, - ) -> AttentionAblationResult: - if not isinstance(contract, AttentionContract): - raise AttentionContractError("contract must be an AttentionContract") - backend_request = ( - backend - if callable(backend) - or hasattr(backend, "forward_with_lse") - or hasattr(backend, "apply") - else None - ) - cfg = _resolve_config( - config, - backend=backend if isinstance(backend, str) else None, - deterministic=deterministic, - return_lse=return_lse, - return_gradients=return_gradients, - communication_backend=( - communication_backend - if communication_backend is not None - else self.communication_backend - ), - validate=validate, - ) - if cfg.validate: - self._validate_inputs(q, k, v, contract) - if cfg.deterministic and contract.split_kv.mode is SplitKVMode.AUTO: - raise AttentionContractError( - "deterministic Attention cannot use runtime-dependent Split-KV=auto" - ) - if cfg.deterministic and ( - cfg.strict_core_id != STRICT_ATTENTION_CORE_ID - or cfg.strict_schedule != STRICT_ATTENTION_SCHEDULE_ID - ): - raise AttentionContractError( - "strict deterministic Attention requires the canonical core and schedule" - ) - if cfg.deterministic and not cfg.return_lse: - raise AttentionContractError("strict deterministic Attention must return LSE") - if cfg.return_gradients and dout is None: - raise AttentionContractError("dout is required when return_gradients=True") - if dout is not None and dout.shape != q.shape: - raise AttentionContractError("dout must have the same shape as q") - - requested = backend_request if backend_request is not None else cfg.backend - selected, selected_id = self._select_backend( - requested, - q, - contract, - communication_backend=cfg.communication_backend, - ) - if cfg.deterministic and selected_id == "native": - raise AttentionContractError( - "deterministic=True cannot execute an unverified native Attention backend" - ) - selected_core_id = getattr(selected, "core_id", None) - selected_schedule = getattr(selected, "strict_schedule", None) - if cfg.deterministic and selected_id not in {BACKEND_ID, REFERENCE_BACKEND_ID}: - if selected_core_id != cfg.strict_core_id or selected_schedule != cfg.strict_schedule: - raise AttentionContractError( - "deterministic Attention requires the shared strict core and schedule" - ) - - call_kwargs = dict(kwargs) - call_kwargs.setdefault("causal", contract.causal) - call_kwargs.setdefault("scale", 1.0 / math.sqrt(contract.head_dim)) - call_kwargs.setdefault("cp_world_size", contract.sharding.cp_world_size) - if contract.split_kv.mode is SplitKVMode.FIXED: - call_kwargs.setdefault("kv_chunk_size", contract.split_kv.fixed_split_size) - out, lse, backend_provenance = self._invoke(selected, q, k, v, call_kwargs, contract) - if cfg.validate: - self._validate_outputs(out, lse, q, contract) - _validate_runtime_provenance( - selected, - selected_id, - backend_provenance, - cfg, - ) - - dq = dk = dv = None - if cfg.return_gradients: - dq, dk, dv = self._backward(selected, q, k, v, out, dout, call_kwargs) - - provenance = { - "schema_version": "rlkernel.attention.ablation_result.v1", - "semantic_operator": "attention", - "backend_id": selected_id, - "deterministic": cfg.deterministic, - "strict_core_id": (cfg.strict_core_id if cfg.deterministic else None), - "strict_schedule": cfg.strict_schedule if cfg.deterministic else None, - "core_id": cfg.strict_core_id if cfg.deterministic else selected_id, - "backend_deterministic": cfg.deterministic, - "native_attention_arithmetic": False if cfg.deterministic else selected_id == "native", - "communication_backend": cfg.communication_backend, - "communication_executed": bool(getattr(selected, "communication_executed", False)), - "split_kv": contract.split_kv.to_dict(), - "actual_split_kv": _actual_split_provenance( - contract, - total_kv_tokens=k.size(2), - backend=selected_id, - ), - "reduction": _reduction_provenance(contract), - "contract_fingerprint": _contract_fingerprint(contract), - "return_lse": cfg.return_lse, - "return_gradients": cfg.return_gradients, - } - provenance.update(backend_provenance) - provenance.setdefault("actual_backend", selected_id) - provenance.setdefault("communication_backend", cfg.communication_backend) - provenance.setdefault("production_ready", False) - return AttentionAblationResult( - out=out, - lse=lse if cfg.return_lse else None, - dq=dq, - dk=dk, - dv=dv, - backend_id=selected_id, - deterministic=cfg.deterministic, - provenance=provenance, - ) - - def apply_fp32(self, *args: Any, **kwargs: Any) -> AttentionAblationResult: - """Stable fingerprint entry point used by ``OperatorSession``.""" - - return self.apply(*args, **kwargs) - - def _select_backend( - self, - requested: str | Callable[..., Any], - q: Tensor, - contract: AttentionContract, - *, - communication_backend: str, - ) -> tuple[Any, str]: - if ( - callable(requested) - or hasattr(requested, "forward_with_lse") - or hasattr(requested, "apply") - ): - return requested, _callable_backend_id(requested) - normalized = str(requested).strip().lower() - if normalized in {"native", "te", "flashinfer"}: - if self.native is None: - raise AttentionContractError( - "native Attention backend was requested but no native callable was injected" - ) - return self.native, "native" - if normalized in {"reference", "pytorch_reference"}: - return self._reference_backend(), REFERENCE_BACKEND_ID - if normalized not in {"auto", "deterministic", "rlkernel"}: - raise AttentionContractError(f"unsupported Attention backend {requested!r}") - if contract.sharding.cp_world_size > 1: - if communication_backend == "self_owned_cuda_ag_rs": - if self.cp_backend is None: - raise AttentionContractError( - "CP production Attention requires an injected AG/RS backend" - ) - return self.cp_backend, _callable_backend_id(self.cp_backend) - return self._reference_backend(), REFERENCE_BACKEND_ID - if q.device.type != "cuda" or torch.version.hip is not None: - return self._reference_backend(), REFERENCE_BACKEND_ID - if self.core is None: - from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( - DeterministicAttentionOp, - ) - - self.core = DeterministicAttentionOp() - return self.core, BACKEND_ID - - def _reference_backend(self) -> Any: - if self.reference is None: - from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( - DeterministicCPAttentionReferenceOp, - ) - - self.reference = DeterministicCPAttentionReferenceOp(strict_bitwise=True) - return self.reference - - @staticmethod - def _validate_inputs(q: Tensor, k: Tensor, v: Tensor, contract: AttentionContract) -> None: - if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: - raise AttentionContractError("q, k, and v must use [B, H, S, D] layout") - expected_dtype = _TORCH_DTYPES[contract.dtype] - if ( - q.dtype is not expected_dtype - or k.dtype is not expected_dtype - or v.dtype is not expected_dtype - ): - raise AttentionContractError( - f"q, k, and v must match contract dtype {contract.dtype.value}" - ) - if q.device != k.device or q.device != v.device: - raise AttentionContractError("q, k, and v must be on the same device") - batch, q_heads, q_seq, dim = q.shape - if batch != contract.batch_size: - raise AttentionContractError( - f"q batch={batch} does not match contract batch_size={contract.batch_size}" - ) - if q_seq != contract.query_sequence_length: - raise AttentionContractError( - "q sequence length does not match AttentionContract query_sequence_length" - ) - sharding = contract.sharding - if q_heads != sharding.local_q_heads or k.shape[1] != sharding.local_kv_heads: - raise AttentionContractError("q/k head counts do not match TP sharding in contract") - if ( - k.shape[0] != batch - or v.shape[:3] != k.shape[:3] - or k.shape[-1] != dim - or v.shape[-1] != dim - ): - raise AttentionContractError("q, k, and v shapes are inconsistent") - if dim != contract.head_dim: - raise AttentionContractError("tensor head_dim does not match AttentionContract") - - @staticmethod - def _validate_outputs(out: Tensor, lse: Tensor, q: Tensor, contract: AttentionContract) -> None: - if out.shape != q.shape: - raise AttentionContractError( - f"Attention output shape {tuple(out.shape)} does not match q {tuple(q.shape)}" - ) - expected_lse = (q.shape[0], q.shape[1], q.shape[2]) - if lse.shape != expected_lse: - raise AttentionContractError( - f"attention-domain LSE shape {tuple(lse.shape)} does not match {expected_lse}" - ) - if lse.dtype is not torch.float32: - raise AttentionContractError("attention-domain LSE must remain fp32") - expected_dtype = _TORCH_DTYPES[contract.dtype] - if out.dtype is not expected_dtype: - raise AttentionContractError( - f"Attention output must be written in {contract.dtype.value}, got {out.dtype}" - ) - - @staticmethod - def _invoke( - backend: Any, - q: Tensor, - k: Tensor, - v: Tensor, - kwargs: Mapping[str, Any], - contract: AttentionContract, - ) -> tuple[Tensor, Tensor, dict[str, Any]]: - method = getattr(backend, "forward_with_lse", None) - if not callable(method): - method = getattr(backend, "apply", None) - if not callable(method): - method = backend if callable(backend) else None - if method is None: - raise AttentionContractError( - "Attention backend must expose forward_with_lse, apply, or __call__" - ) - accepted = _accepted_kwargs(method, kwargs) - if contract.sharding.cp_world_size > 1 and "cp_world_size" not in accepted: - raise AttentionContractError( - "CP>1 requires an Attention backend that explicitly accepts cp_world_size" - ) - result = method(q, k, v, **accepted) - backend_provenance: dict[str, Any] = {} - if isinstance(result, AttentionAblationResult): - out, lse = result.out, result.lse - elif isinstance(result, tuple) and len(result) == 2: - out, lse = result - else: - out = getattr(result, "out", None) - lse = getattr(result, "lse", None) - raw_provenance = getattr(result, "provenance", {}) - if isinstance(raw_provenance, Mapping): - backend_provenance = dict(raw_provenance) - if out is None or lse is None: - raise AttentionContractError( - "Attention backend must return (out, lse), AttentionAblationResult, " - "or an object with out/lse/provenance" - ) - if isinstance(result, AttentionAblationResult): - backend_provenance = dict(result.provenance) - if not isinstance(out, Tensor) or not isinstance(lse, Tensor): - raise AttentionContractError("Attention backend returned non-tensor output or LSE") - return out, lse, backend_provenance - - @staticmethod - def _backward( - backend: Any, - q: Tensor, - k: Tensor, - v: Tensor, - out: Tensor, - dout: Tensor | None, - kwargs: Mapping[str, Any], - ) -> tuple[Tensor, Tensor, Tensor]: - backward = getattr(backend, "backward_reference", None) - if callable(backward): - result = backward(q, k, v, dout, **_accepted_kwargs(backward, kwargs)) - gradients = getattr(result, "gradients", None) - if gradients is not None: - return gradients.dq, gradients.dk, gradients.dv - if not out.requires_grad: - raise AttentionContractError( - "Attention backend did not retain an autograd graph for gradients" - ) - return torch.autograd.grad( - out, - (q, k, v), - grad_outputs=dout.to(dtype=out.dtype), - allow_unused=False, - retain_graph=True, - ) - - -def _resolve_config( - config: AttentionAblationConfig | Mapping[str, Any] | None, - **overrides: Any, -) -> AttentionAblationConfig: - if config is None: - values: dict[str, Any] = {} - elif isinstance(config, AttentionAblationConfig): - values = { - name: getattr(config, name) - for name in ( - "backend", - "deterministic", - "communication_backend", - "return_lse", - "return_gradients", - "strict_core_id", - "strict_schedule", - "validate", - ) - } - elif isinstance(config, Mapping): - values = dict(config) - else: - raise AttentionContractError("config must be AttentionAblationConfig, mapping, or None") - values.update({name: value for name, value in overrides.items() if value is not None}) - return AttentionAblationConfig(**values) - - -def _accepted_kwargs(method: Callable[..., Any], kwargs: Mapping[str, Any]) -> dict[str, Any]: - try: - signature = inspect.signature(method) - except (TypeError, ValueError): - return dict(kwargs) - parameters = signature.parameters.values() - if any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters): - return dict(kwargs) - return {name: value for name, value in kwargs.items() if name in signature.parameters} - - -def _callable_backend_id(value: Any) -> str: - explicit = getattr(value, "backend_id", None) or getattr( - value, "__attention_backend_id__", None - ) - if isinstance(explicit, str) and explicit.strip(): - return explicit.strip() - return f"injected.{type(value).__module__}.{type(value).__qualname__}" - - -def _validate_runtime_provenance( - selected: Any, - selected_id: str, - runtime: Mapping[str, Any], - config: AttentionAblationConfig, -) -> None: - """Fail closed when a production strict backend did not prove its identity.""" - - if not config.deterministic: - return - if selected_id == "native": - raise AttentionContractError( - "deterministic Attention cannot execute native Attention arithmetic" - ) - - actual_core = runtime.get("strict_core_id", getattr(selected, "core_id", None)) - actual_schedule = runtime.get("strict_schedule", getattr(selected, "strict_schedule", None)) - if selected_id != REFERENCE_BACKEND_ID and ( - actual_core != config.strict_core_id or actual_schedule != config.strict_schedule - ): - raise AttentionContractError( - "deterministic Attention backend did not prove the shared strict core and schedule" - ) - - if config.communication_backend != "self_owned_cuda_ag_rs": - return - - expected = { - "strict_core_id": config.strict_core_id, - "strict_schedule": config.strict_schedule, - "actual_backend": "rlkernel.cuda.deterministic_attention", - "communication_backend": "self_owned_cuda_ag_rs", - "production_ready": True, - "native_attention_arithmetic": False, - "fallback": False, - } - mismatches = [name for name, value in expected.items() if runtime.get(name) != value] - if mismatches: - raise AttentionContractError( - "CP production Attention runtime provenance is incomplete or mismatched: " - + ", ".join(mismatches) - ) - - -def _actual_split_provenance( - contract: AttentionContract, - *, - total_kv_tokens: int, - backend: str = BACKEND_ID, -) -> dict[str, Any]: - plan = contract.split_kv.resolve( - total_kv_tokens, - backend=backend, - ) - return plan.to_dict() - - -def _reduction_provenance(contract: AttentionContract) -> dict[str, Any]: - reduction = contract.reduction - return { - "merge": reduction.merge.value, - "acc_dtype": reduction.acc_dtype.value, - "order": reduction.order.value, - "downcast_at": reduction.downcast_at.value, - "engine": reduction.engine.value, - } - - -def _contract_fingerprint(contract: AttentionContract) -> str: - payload = json.dumps(contract.to_dict(), sort_keys=True, separators=(",", ":")) - return hashlib.sha256(payload.encode("utf-8")).hexdigest() - - -__all__ = [ - "AttentionAblationConfig", - "AttentionAblationOp", - "AttentionAblationResult", - "BACKEND_ID", - "REFERENCE_BACKEND_ID", -] +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Unified Attention entry point for the PR230 cross-configuration matrix. + +The matrix needs one stable callable shape even though training and rollout may +materialize different Attention backends. This adapter owns the common +contract checks and provenance only; numerical work remains in the existing +deterministic Attention implementations. +""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import math +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Callable, Mapping, cast + +import torch +from torch import Tensor + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + AttentionContract, + AttentionContractError, + AttentionDType, + SplitKVMode, +) + +BACKEND_ID = "rlkernel.attention.deterministic.v1" +REFERENCE_BACKEND_ID = "rlkernel.attention.reference.v1" + +_TORCH_DTYPES = { + AttentionDType.BF16: torch.bfloat16, + AttentionDType.FP16: torch.float16, + AttentionDType.FP32: torch.float32, +} + + +@dataclass(frozen=True) +class AttentionAblationConfig: + """Per-invocation settings materialized by the ablation runner.""" + + backend: str = "auto" + deterministic: bool = True + communication_backend: str = "none" + return_lse: bool = True + return_gradients: bool = False + strict_core_id: str = STRICT_ATTENTION_CORE_ID + strict_schedule: str = STRICT_ATTENTION_SCHEDULE_ID + validate: bool = True + + def __post_init__(self) -> None: + if not isinstance(self.backend, str) or not self.backend.strip(): + raise AttentionContractError("Attention backend must be a non-empty string") + if not isinstance(self.deterministic, bool): + raise AttentionContractError("deterministic must be a bool") + if ( + not isinstance(self.communication_backend, str) + or not self.communication_backend.strip() + ): + raise AttentionContractError("communication_backend must be a non-empty string") + for name in ("return_lse", "return_gradients", "validate"): + if not isinstance(getattr(self, name), bool): + raise AttentionContractError(f"{name} must be a bool") + if not isinstance(self.strict_core_id, str) or not self.strict_core_id.strip(): + raise AttentionContractError("strict_core_id must be a non-empty string") + if not isinstance(self.strict_schedule, str) or not self.strict_schedule.strip(): + raise AttentionContractError("strict_schedule must be a non-empty string") + object.__setattr__(self, "backend", self.backend.strip().lower()) + object.__setattr__(self, "communication_backend", self.communication_backend.strip()) + + +@dataclass(frozen=True) +class AttentionAblationResult: + """Standardized Attention result consumed by cross-config artifacts.""" + + out: Tensor + lse: Tensor | None + dq: Tensor | None = None + dk: Tensor | None = None + dv: Tensor | None = None + backend_id: str = BACKEND_ID + deterministic: bool = True + provenance: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.out, Tensor): + raise TypeError("Attention result out must be a torch.Tensor") + if self.lse is not None and not isinstance(self.lse, Tensor): + raise TypeError("Attention result lse must be a torch.Tensor or None") + for name in ("dq", "dk", "dv"): + value = getattr(self, name) + if value is not None and not isinstance(value, Tensor): + raise TypeError(f"Attention result {name} must be a torch.Tensor or None") + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise ValueError("Attention result backend_id must be non-empty") + if not isinstance(self.deterministic, bool): + raise TypeError("Attention result deterministic must be a bool") + if not isinstance(self.provenance, Mapping): + raise TypeError("Attention result provenance must be a mapping") + object.__setattr__(self, "provenance", MappingProxyType(dict(self.provenance))) + + @property + def out_lse(self) -> tuple[Tensor, Tensor | None]: + """Compatibility tuple for callers that consume ``(out, lse)``.""" + + return self.out, self.lse + + def readback(self) -> dict[str, Any]: + """Return JSON-compatible execution evidence for PR230 artifacts.""" + + return { + "backend_id": self.backend_id, + "deterministic": self.deterministic, + "out_shape": list(self.out.shape), + "out_dtype": str(self.out.dtype).replace("torch.", ""), + "lse_shape": None if self.lse is None else list(self.lse.shape), + "lse_dtype": None if self.lse is None else str(self.lse.dtype).replace("torch.", ""), + "gradients": { + "dq": self.dq is not None, + "dk": self.dk is not None, + "dv": self.dv is not None, + }, + "provenance": dict(self.provenance), + } + + +class AttentionAblationOp: + """PR230/PR314-style unified Attention wrapper. + + ``core`` and ``reference`` are injectable so the wrapper is usable by the + semantic operator session without importing CUDA at construction time. + ``core`` should expose ``forward_with_lse``; ``reference`` is the existing + pure-PyTorch CP reference with the same method. + """ + + op_class = "attention" + is_batch_invariant = True + backend_id = BACKEND_ID + + def __init__( + self, + *, + core: Any | None = None, + reference: Any | None = None, + native: Any | None = None, + cp_backend: Any | None = None, + communication_backend: str = "none", + ) -> None: + if not isinstance(communication_backend, str) or not communication_backend.strip(): + raise AttentionContractError("communication_backend must be a non-empty string") + self.core = core + self.reference = reference + self.native = native + # CP production execution is injected by the runtime adapter. Keeping + # it separate from the single-device core prevents an accidental + # fallback to the PyTorch reference when AG/RS is required. + self.cp_backend = cp_backend + self.communication_backend = communication_backend.strip() + + def __call__( + self, + q: Tensor, + k: Tensor, + v: Tensor, + *, + contract: AttentionContract, + config: AttentionAblationConfig | Mapping[str, Any] | None = None, + backend: str | Callable[..., Any] | None = None, + deterministic: bool | None = None, + return_lse: bool | None = None, + return_gradients: bool | None = None, + dout: Tensor | None = None, + communication_backend: str | None = None, + validate: bool | None = None, + **kwargs: Any, + ) -> AttentionAblationResult: + return self.apply( + q, + k, + v, + contract=contract, + config=config, + backend=backend, + deterministic=deterministic, + return_lse=return_lse, + return_gradients=return_gradients, + dout=dout, + communication_backend=communication_backend, + validate=validate, + **kwargs, + ) + + def apply( + self, + q: Tensor, + k: Tensor, + v: Tensor, + *, + contract: AttentionContract, + config: AttentionAblationConfig | Mapping[str, Any] | None = None, + backend: str | Callable[..., Any] | None = None, + deterministic: bool | None = None, + return_lse: bool | None = None, + return_gradients: bool | None = None, + dout: Tensor | None = None, + communication_backend: str | None = None, + validate: bool | None = None, + **kwargs: Any, + ) -> AttentionAblationResult: + if not isinstance(contract, AttentionContract): + raise AttentionContractError("contract must be an AttentionContract") + backend_request = ( + backend + if callable(backend) + or hasattr(backend, "forward_with_lse") + or hasattr(backend, "apply") + else None + ) + cfg = _resolve_config( + config, + backend=backend if isinstance(backend, str) else None, + deterministic=deterministic, + return_lse=return_lse, + return_gradients=return_gradients, + communication_backend=( + communication_backend + if communication_backend is not None + else self.communication_backend + ), + validate=validate, + ) + if cfg.validate: + self._validate_inputs(q, k, v, contract) + if cfg.deterministic and contract.split_kv.mode is SplitKVMode.AUTO: + raise AttentionContractError( + "deterministic Attention cannot use runtime-dependent Split-KV=auto" + ) + if cfg.deterministic and ( + cfg.strict_core_id != STRICT_ATTENTION_CORE_ID + or cfg.strict_schedule != STRICT_ATTENTION_SCHEDULE_ID + ): + raise AttentionContractError( + "strict deterministic Attention requires the canonical core and schedule" + ) + if cfg.deterministic and not cfg.return_lse: + raise AttentionContractError("strict deterministic Attention must return LSE") + if cfg.return_gradients and dout is None: + raise AttentionContractError("dout is required when return_gradients=True") + if dout is not None and dout.shape != q.shape: + raise AttentionContractError("dout must have the same shape as q") + + requested = backend_request if backend_request is not None else cfg.backend + selected, selected_id = self._select_backend( + requested, + q, + contract, + communication_backend=cfg.communication_backend, + ) + if cfg.deterministic and selected_id == "native": + raise AttentionContractError( + "deterministic=True cannot execute an unverified native Attention backend" + ) + selected_core_id = getattr(selected, "core_id", None) + selected_schedule = getattr(selected, "strict_schedule", None) + if cfg.deterministic and selected_id not in {BACKEND_ID, REFERENCE_BACKEND_ID}: + if selected_core_id != cfg.strict_core_id or selected_schedule != cfg.strict_schedule: + raise AttentionContractError( + "deterministic Attention requires the shared strict core and schedule" + ) + + call_kwargs = dict(kwargs) + call_kwargs.setdefault("causal", contract.causal) + call_kwargs.setdefault("scale", 1.0 / math.sqrt(contract.head_dim)) + call_kwargs.setdefault("cp_world_size", contract.sharding.cp_world_size) + if contract.split_kv.mode is SplitKVMode.FIXED: + call_kwargs.setdefault("kv_chunk_size", contract.split_kv.fixed_split_size) + out, lse, backend_provenance = self._invoke(selected, q, k, v, call_kwargs, contract) + if cfg.validate: + self._validate_outputs(out, lse, q, contract) + _validate_runtime_provenance( + selected, + selected_id, + backend_provenance, + cfg, + ) + + dq = dk = dv = None + if cfg.return_gradients: + dq, dk, dv = self._backward(selected, q, k, v, out, dout, call_kwargs) + + provenance = { + "schema_version": "rlkernel.attention.ablation_result.v1", + "semantic_operator": "attention", + "backend_id": selected_id, + "deterministic": cfg.deterministic, + "strict_core_id": (cfg.strict_core_id if cfg.deterministic else None), + "strict_schedule": cfg.strict_schedule if cfg.deterministic else None, + "core_id": cfg.strict_core_id if cfg.deterministic else selected_id, + "backend_deterministic": cfg.deterministic, + "native_attention_arithmetic": False if cfg.deterministic else selected_id == "native", + "communication_backend": cfg.communication_backend, + "communication_executed": bool(getattr(selected, "communication_executed", False)), + "split_kv": contract.split_kv.to_dict(), + "actual_split_kv": _actual_split_provenance( + contract, + total_kv_tokens=k.size(2), + backend=selected_id, + ), + "reduction": _reduction_provenance(contract), + "contract_fingerprint": _contract_fingerprint(contract), + "return_lse": cfg.return_lse, + "return_gradients": cfg.return_gradients, + } + provenance.update(backend_provenance) + provenance.setdefault("actual_backend", selected_id) + provenance.setdefault("communication_backend", cfg.communication_backend) + provenance.setdefault("production_ready", False) + return AttentionAblationResult( + out=out, + lse=lse if cfg.return_lse else None, + dq=dq, + dk=dk, + dv=dv, + backend_id=selected_id, + deterministic=cfg.deterministic, + provenance=provenance, + ) + + def apply_fp32(self, *args: Any, **kwargs: Any) -> AttentionAblationResult: + """Stable fingerprint entry point used by ``OperatorSession``.""" + + return self.apply(*args, **kwargs) + + def _select_backend( + self, + requested: str | Callable[..., Any], + q: Tensor, + contract: AttentionContract, + *, + communication_backend: str, + ) -> tuple[Any, str]: + if ( + callable(requested) + or hasattr(requested, "forward_with_lse") + or hasattr(requested, "apply") + ): + return requested, _callable_backend_id(requested) + normalized = str(requested).strip().lower() + if normalized in {"native", "te", "flashinfer"}: + if self.native is None: + raise AttentionContractError( + "native Attention backend was requested but no native callable was injected" + ) + return self.native, "native" + if normalized in {"reference", "pytorch_reference"}: + return self._reference_backend(), REFERENCE_BACKEND_ID + if normalized not in {"auto", "deterministic", "rlkernel"}: + raise AttentionContractError(f"unsupported Attention backend {requested!r}") + if contract.sharding.cp_world_size > 1: + if communication_backend == "self_owned_cuda_ag_rs": + if self.cp_backend is None: + raise AttentionContractError( + "CP production Attention requires an injected AG/RS backend" + ) + return self.cp_backend, _callable_backend_id(self.cp_backend) + return self._reference_backend(), REFERENCE_BACKEND_ID + if q.device.type != "cuda" or torch.version.hip is not None: + return self._reference_backend(), REFERENCE_BACKEND_ID + if self.core is None: + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionOp, + ) + + self.core = DeterministicAttentionOp() + return self.core, BACKEND_ID + + def _reference_backend(self) -> Any: + if self.reference is None: + from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + DeterministicCPAttentionReferenceOp, + ) + + self.reference = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + return self.reference + + @staticmethod + def _validate_inputs(q: Tensor, k: Tensor, v: Tensor, contract: AttentionContract) -> None: + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise AttentionContractError("q, k, and v must use [B, H, S, D] layout") + expected_dtype = _TORCH_DTYPES[contract.dtype] + if ( + q.dtype is not expected_dtype + or k.dtype is not expected_dtype + or v.dtype is not expected_dtype + ): + raise AttentionContractError( + f"q, k, and v must match contract dtype {contract.dtype.value}" + ) + if q.device != k.device or q.device != v.device: + raise AttentionContractError("q, k, and v must be on the same device") + batch, q_heads, q_seq, dim = q.shape + if batch != contract.batch_size: + raise AttentionContractError( + f"q batch={batch} does not match contract batch_size={contract.batch_size}" + ) + if q_seq != contract.query_sequence_length: + raise AttentionContractError( + "q sequence length does not match AttentionContract query_sequence_length" + ) + sharding = contract.sharding + if q_heads != sharding.local_q_heads or k.shape[1] != sharding.local_kv_heads: + raise AttentionContractError("q/k head counts do not match TP sharding in contract") + if ( + k.shape[0] != batch + or v.shape[:3] != k.shape[:3] + or k.shape[-1] != dim + or v.shape[-1] != dim + ): + raise AttentionContractError("q, k, and v shapes are inconsistent") + if dim != contract.head_dim: + raise AttentionContractError("tensor head_dim does not match AttentionContract") + + @staticmethod + def _validate_outputs(out: Tensor, lse: Tensor, q: Tensor, contract: AttentionContract) -> None: + if out.shape != q.shape: + raise AttentionContractError( + f"Attention output shape {tuple(out.shape)} does not match q {tuple(q.shape)}" + ) + expected_lse = (q.shape[0], q.shape[1], q.shape[2]) + if lse.shape != expected_lse: + raise AttentionContractError( + f"attention-domain LSE shape {tuple(lse.shape)} does not match {expected_lse}" + ) + if lse.dtype is not torch.float32: + raise AttentionContractError("attention-domain LSE must remain fp32") + expected_dtype = _TORCH_DTYPES[contract.dtype] + if out.dtype is not expected_dtype: + raise AttentionContractError( + f"Attention output must be written in {contract.dtype.value}, got {out.dtype}" + ) + + @staticmethod + def _invoke( + backend: Any, + q: Tensor, + k: Tensor, + v: Tensor, + kwargs: Mapping[str, Any], + contract: AttentionContract, + ) -> tuple[Tensor, Tensor, dict[str, Any]]: + method = getattr(backend, "forward_with_lse", None) + if not callable(method): + method = getattr(backend, "apply", None) + if not callable(method): + method = backend if callable(backend) else None + if method is None: + raise AttentionContractError( + "Attention backend must expose forward_with_lse, apply, or __call__" + ) + accepted = _accepted_kwargs(method, kwargs) + if contract.sharding.cp_world_size > 1 and "cp_world_size" not in accepted: + raise AttentionContractError( + "CP>1 requires an Attention backend that explicitly accepts cp_world_size" + ) + result = method(q, k, v, **accepted) + backend_provenance: dict[str, Any] = {} + out: Any + lse: Any + if isinstance(result, AttentionAblationResult): + out, lse = result.out, result.lse + elif isinstance(result, tuple) and len(result) == 2: + out, lse = result + else: + out = getattr(result, "out", None) + lse = getattr(result, "lse", None) + raw_provenance = getattr(result, "provenance", {}) + if isinstance(raw_provenance, Mapping): + backend_provenance = dict(raw_provenance) + if out is None or lse is None: + raise AttentionContractError( + "Attention backend must return (out, lse), AttentionAblationResult, " + "or an object with out/lse/provenance" + ) + if isinstance(result, AttentionAblationResult): + backend_provenance = dict(result.provenance) + if not isinstance(out, Tensor) or not isinstance(lse, Tensor): + raise AttentionContractError("Attention backend returned non-tensor output or LSE") + return out, lse, backend_provenance + + @staticmethod + def _backward( + backend: Any, + q: Tensor, + k: Tensor, + v: Tensor, + out: Tensor, + dout: Tensor | None, + kwargs: Mapping[str, Any], + ) -> tuple[Tensor, Tensor, Tensor]: + backward = getattr(backend, "backward_reference", None) + if callable(backward): + result = backward(q, k, v, dout, **_accepted_kwargs(backward, kwargs)) + gradients = getattr(result, "gradients", None) + if gradients is not None: + return gradients.dq, gradients.dk, gradients.dv + if dout is None: + raise AttentionContractError("dout is required to compute Attention gradients") + if not out.requires_grad: + raise AttentionContractError( + "Attention backend did not retain an autograd graph for gradients" + ) + return cast( + tuple[Tensor, Tensor, Tensor], + torch.autograd.grad( + out, + (q, k, v), + grad_outputs=dout.to(dtype=out.dtype), + allow_unused=False, + retain_graph=True, + ), + ) + + +def _resolve_config( + config: AttentionAblationConfig | Mapping[str, Any] | None, + **overrides: Any, +) -> AttentionAblationConfig: + if config is None: + values: dict[str, Any] = {} + elif isinstance(config, AttentionAblationConfig): + values = { + name: getattr(config, name) + for name in ( + "backend", + "deterministic", + "communication_backend", + "return_lse", + "return_gradients", + "strict_core_id", + "strict_schedule", + "validate", + ) + } + elif isinstance(config, Mapping): + values = dict(config) + else: + raise AttentionContractError("config must be AttentionAblationConfig, mapping, or None") + values.update({name: value for name, value in overrides.items() if value is not None}) + return AttentionAblationConfig(**values) + + +def _accepted_kwargs(method: Callable[..., Any], kwargs: Mapping[str, Any]) -> dict[str, Any]: + try: + signature = inspect.signature(method) + except (TypeError, ValueError): + return dict(kwargs) + parameters = signature.parameters.values() + if any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters): + return dict(kwargs) + return {name: value for name, value in kwargs.items() if name in signature.parameters} + + +def _callable_backend_id(value: Any) -> str: + explicit = getattr(value, "backend_id", None) or getattr( + value, "__attention_backend_id__", None + ) + if isinstance(explicit, str) and explicit.strip(): + return explicit.strip() + return f"injected.{type(value).__module__}.{type(value).__qualname__}" + + +def _validate_runtime_provenance( + selected: Any, + selected_id: str, + runtime: Mapping[str, Any], + config: AttentionAblationConfig, +) -> None: + """Fail closed when a production strict backend did not prove its identity.""" + + if not config.deterministic: + return + if selected_id == "native": + raise AttentionContractError( + "deterministic Attention cannot execute native Attention arithmetic" + ) + + actual_core = runtime.get("strict_core_id", getattr(selected, "core_id", None)) + actual_schedule = runtime.get("strict_schedule", getattr(selected, "strict_schedule", None)) + if selected_id != REFERENCE_BACKEND_ID and ( + actual_core != config.strict_core_id or actual_schedule != config.strict_schedule + ): + raise AttentionContractError( + "deterministic Attention backend did not prove the shared strict core and schedule" + ) + + if config.communication_backend != "self_owned_cuda_ag_rs": + return + + expected = { + "strict_core_id": config.strict_core_id, + "strict_schedule": config.strict_schedule, + "actual_backend": "rlkernel.cuda.deterministic_attention", + "communication_backend": "self_owned_cuda_ag_rs", + "production_ready": True, + "native_attention_arithmetic": False, + "fallback": False, + } + mismatches = [name for name, value in expected.items() if runtime.get(name) != value] + if mismatches: + raise AttentionContractError( + "CP production Attention runtime provenance is incomplete or mismatched: " + + ", ".join(mismatches) + ) + + +def _actual_split_provenance( + contract: AttentionContract, + *, + total_kv_tokens: int, + backend: str = BACKEND_ID, +) -> dict[str, Any]: + plan = contract.split_kv.resolve( + total_kv_tokens, + backend=backend, + ) + return plan.to_dict() + + +def _reduction_provenance(contract: AttentionContract) -> dict[str, Any]: + reduction = contract.reduction + return { + "merge": reduction.merge.value, + "acc_dtype": reduction.acc_dtype.value, + "order": reduction.order.value, + "downcast_at": reduction.downcast_at.value, + "engine": reduction.engine.value, + } + + +def _contract_fingerprint(contract: AttentionContract) -> str: + payload = json.dumps(contract.to_dict(), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +__all__ = [ + "AttentionAblationConfig", + "AttentionAblationOp", + "AttentionAblationResult", + "BACKEND_ID", + "REFERENCE_BACKEND_ID", +] From 980433eed29d69fb1f0f745cf037d7514492359f Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Wed, 19 Aug 2026 15:02:08 +0000 Subject: [PATCH 14/33] feat(attention): add compact drift debug taxonomy --- .../ops/pytorch/attention/debug_taxonomy.py | 180 ++++++++++++++++++ tests/test_attention_debug_taxonomy.py | 50 +++++ 2 files changed, 230 insertions(+) create mode 100644 rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py create mode 100644 tests/test_attention_debug_taxonomy.py diff --git a/rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py b/rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py new file mode 100644 index 00000000..d9f1a6d6 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Compact root-cause taxonomy for post-training Attention drift reports.""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from rl_engine.kernels.attention_contract import AttentionContractError + +ATTENTION_DEBUG_SCHEMA_VERSION = "rlkernel.attention.debug_taxonomy.v1" + + +@dataclass(frozen=True) +class AttentionDebugAxisSpec: + """One first-line post-training Attention drift category.""" + + axis_id: str + label: str + representative_subprobe: str + subprobes: tuple[str, ...] + + def __post_init__(self) -> None: + for name in ("axis_id", "label", "representative_subprobe"): + value = getattr(self, name) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Attention debug {name} must be a non-empty string") + if not self.subprobes or any( + not isinstance(probe, str) or not probe.strip() for probe in self.subprobes + ): + raise ValueError("Attention debug subprobes must be non-empty strings") + if len(set(self.subprobes)) != len(self.subprobes): + raise ValueError( + f"Attention debug axis {self.axis_id!r} has duplicate subprobes" + ) + if self.representative_subprobe not in self.subprobes: + raise ValueError( + f"Attention debug representative {self.representative_subprobe!r} " + f"is not assigned to {self.axis_id!r}" + ) + + +ATTENTION_DEBUG_AXES = ( + AttentionDebugAxisSpec( + axis_id="position_rope", + label="Position / RoPE", + representative_subprobe="position_ids", + subprobes=( + "position_ids", + "rope_theta", + "position_offsets", + "decode_cache_position", + ), + ), + AttentionDebugAxisSpec( + axis_id="qk_preprocessing", + label="Q/K preprocessing", + representative_subprobe="qk_norm_disabled", + subprobes=( + "qk_norm_eps", + "qk_norm_disabled", + "qk_norm_weight", + "attention_scale", + "scale_placement", + ), + ), + AttentionDebugAxisSpec( + axis_id="mask_sequence_boundary", + label="Mask / sequence boundary", + representative_subprobe="causal_mask", + subprobes=("causal_mask", "key_padding_mask"), + ), + AttentionDebugAxisSpec( + axis_id="topology_head_ownership", + label="Topology / head ownership", + representative_subprobe="tp_head_ownership", + subprobes=("tp_head_ownership",), + ), + AttentionDebugAxisSpec( + axis_id="kv_cache_identity_layout", + label="KV-cache identity / layout", + representative_subprobe="kv_page_order", + subprobes=("kv_page_order", "kv_cache_content"), + ), + AttentionDebugAxisSpec( + axis_id="numerical_policy", + label="Numerical policy", + representative_subprobe="accum_dtype", + subprobes=( + "accum_dtype", + "execution_dtype", + "final_write_dtype", + "early_downcast", + ), + ), + AttentionDebugAxisSpec( + axis_id="distributed_schedule", + label="Distributed schedule", + representative_subprobe="merge_order", + subprobes=("nonstrict_cp_degree", "split_kv", "merge_order"), + ), +) + +ATTENTION_INVARIANT_CONTROLS = ( + "tp_partition_control", + "batch_composition_control", + "prefill_decode_tail_control", +) + +_ATTENTION_DEBUG_AXIS_BY_ID = MappingProxyType( + {axis.axis_id: axis for axis in ATTENTION_DEBUG_AXES} +) +_ATTENTION_DEBUG_AXIS_BY_PROBE = MappingProxyType( + {probe: axis for axis in ATTENTION_DEBUG_AXES for probe in axis.subprobes} +) +if len(_ATTENTION_DEBUG_AXIS_BY_ID) != len(ATTENTION_DEBUG_AXES): + raise RuntimeError("Attention debug axis IDs must be unique") +if len(_ATTENTION_DEBUG_AXIS_BY_PROBE) != sum( + len(axis.subprobes) for axis in ATTENTION_DEBUG_AXES +): + raise RuntimeError("Attention debug subprobes must belong to one root-cause axis") +if set(_ATTENTION_DEBUG_AXIS_BY_PROBE) & set(ATTENTION_INVARIANT_CONTROLS): + raise RuntimeError("Attention debug subprobes cannot also be invariant controls") + + +def attention_debug_probe_metadata(probe: str) -> dict[str, Any]: + """Classify one stable debug probe for report aggregation.""" + + if not isinstance(probe, str) or not probe.strip(): + raise AttentionContractError("Attention debug probe must be a non-empty string") + normalized = probe.strip() + if normalized in ATTENTION_INVARIANT_CONTROLS: + return { + "category": "invariant_control", + "root_cause_axis": None, + "root_cause_label": None, + "representative": False, + } + axis = _ATTENTION_DEBUG_AXIS_BY_PROBE.get(normalized) + if axis is None: + raise AttentionContractError(f"unknown Attention debug probe {probe!r}") + return { + "category": "root_cause_subprobe", + "root_cause_axis": axis.axis_id, + "root_cause_label": axis.label, + "representative": normalized == axis.representative_subprobe, + } + + +def attention_debug_taxonomy() -> dict[str, Any]: + """Return the compact JSON schema used by post-training drift reports.""" + + return { + "schema_version": ATTENTION_DEBUG_SCHEMA_VERSION, + "root_cause_axis_count": len(ATTENTION_DEBUG_AXES), + "subprobe_count": len(_ATTENTION_DEBUG_AXIS_BY_PROBE), + "invariant_control_count": len(ATTENTION_INVARIANT_CONTROLS), + "root_cause_axes": { + axis.axis_id: { + "label": axis.label, + "representative_subprobe": axis.representative_subprobe, + "subprobes": list(axis.subprobes), + } + for axis in ATTENTION_DEBUG_AXES + }, + "invariant_controls": list(ATTENTION_INVARIANT_CONTROLS), + } + + +__all__ = [ + "ATTENTION_DEBUG_AXES", + "ATTENTION_DEBUG_SCHEMA_VERSION", + "ATTENTION_INVARIANT_CONTROLS", + "AttentionDebugAxisSpec", + "attention_debug_probe_metadata", + "attention_debug_taxonomy", +] diff --git a/tests/test_attention_debug_taxonomy.py b/tests/test_attention_debug_taxonomy.py new file mode 100644 index 00000000..149a99fd --- /dev/null +++ b/tests/test_attention_debug_taxonomy.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from rl_engine.kernels.attention_contract import AttentionContractError +from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ( + ATTENTION_DEBUG_SCHEMA_VERSION, + attention_debug_probe_metadata, + attention_debug_taxonomy, +) + + +def test_attention_debug_taxonomy_is_compact_complete_and_unambiguous(): + taxonomy = attention_debug_taxonomy() + + assert taxonomy["schema_version"] == ATTENTION_DEBUG_SCHEMA_VERSION + assert taxonomy["root_cause_axis_count"] == 7 + assert taxonomy["subprobe_count"] == 21 + assert taxonomy["invariant_control_count"] == 3 + assert set(taxonomy["root_cause_axes"]) == { + "position_rope", + "qk_preprocessing", + "mask_sequence_boundary", + "topology_head_ownership", + "kv_cache_identity_layout", + "numerical_policy", + "distributed_schedule", + } + assert "strict_cp_degree_control" not in { + probe + for axis in taxonomy["root_cause_axes"].values() + for probe in axis["subprobes"] + } + + +def test_attention_debug_probe_metadata_separates_subprobes_and_controls(): + assert attention_debug_probe_metadata("position_ids") == { + "category": "root_cause_subprobe", + "root_cause_axis": "position_rope", + "root_cause_label": "Position / RoPE", + "representative": True, + } + control = attention_debug_probe_metadata("tp_partition_control") + assert control["category"] == "invariant_control" + assert control["root_cause_axis"] is None + + with pytest.raises(AttentionContractError, match="unknown Attention debug probe"): + attention_debug_probe_metadata("strict_cp_degree_control") From 661558291e386561ba97e6a67e73912f973f3317 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 20 Aug 2026 16:16:43 +0800 Subject: [PATCH 15/33] feat(alignment): integrate deterministic FFN and logprob wrappers --- .github/workflows/ci.yml | 12 + csrc/cuda/activation.cu | 253 ++++ .../distributed/deterministic_collective.cu | 572 +++++++++ csrc/cuda/gemm/det_gemm_kernel.cu | 121 +- csrc/ops.cpp | 79 ++ docs/design/runtime-dispatch.md | 8 + docs/operators/activation.md | 78 +- docs/operators/batch-invariant-logp.md | 36 + rl_engine/_C.pyi | 39 + rl_engine/alignment/cross_config/operators.py | 56 +- rl_engine/distributed/__init__.py | 6 + rl_engine/distributed/collectives.py | 325 +++++ rl_engine/kernels/gtest/operator_specs.py | 24 + rl_engine/kernels/logprob_contract.py | 652 ++++++++++ rl_engine/kernels/ops/cuda/__init__.py | 4 +- .../kernels/ops/cuda/activation/__init__.py | 6 + .../kernels/ops/cuda/activation/swiglu.py | 157 +++ .../kernels/ops/pytorch/activation/swiglu.py | 17 + rl_engine/kernels/ops/pytorch/ffn/__init__.py | 6 + rl_engine/kernels/ops/pytorch/ffn/ffn.py | 477 +++++++ .../ops/pytorch/loss/vocab_parallel_logp.py | 463 +++++++ .../kernels/ops/triton/activation/__init__.py | 6 + .../kernels/ops/triton/activation/swiglu.py | 251 ++++ rl_engine/kernels/registry.py | 360 +++++- setup.py | 8 +- tests/distributed/__init__.py | 2 + .../distributed/test_det_gemm_simulated_tp.py | 85 ++ .../test_deterministic_all_gather.py | 115 ++ .../test_deterministic_all_reduce.py | 116 ++ .../test_deterministic_reduce_scatter.py | 118 ++ tests/test_alignment_wrapper_interfaces.py | 159 +++ tests/test_det_gemm.py | 35 +- tests/test_logprob_contract.py | 616 +++++++++ tests/test_qwen_ffn.py | 1112 +++++++++++++++++ tests/test_swiglu.py | 525 +++++++- tests/test_vocab_parallel_logp.py | 558 +++++++++ 36 files changed, 7374 insertions(+), 83 deletions(-) create mode 100644 csrc/cuda/activation.cu create mode 100644 csrc/cuda/distributed/deterministic_collective.cu create mode 100644 rl_engine/distributed/__init__.py create mode 100644 rl_engine/distributed/collectives.py create mode 100644 rl_engine/kernels/logprob_contract.py create mode 100644 rl_engine/kernels/ops/cuda/activation/__init__.py create mode 100644 rl_engine/kernels/ops/cuda/activation/swiglu.py create mode 100644 rl_engine/kernels/ops/pytorch/ffn/__init__.py create mode 100644 rl_engine/kernels/ops/pytorch/ffn/ffn.py create mode 100644 rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py create mode 100644 rl_engine/kernels/ops/triton/activation/__init__.py create mode 100644 rl_engine/kernels/ops/triton/activation/swiglu.py create mode 100644 tests/distributed/__init__.py create mode 100644 tests/distributed/test_det_gemm_simulated_tp.py create mode 100644 tests/distributed/test_deterministic_all_gather.py create mode 100644 tests/distributed/test_deterministic_all_reduce.py create mode 100644 tests/distributed/test_deterministic_reduce_scatter.py create mode 100644 tests/test_alignment_wrapper_interfaces.py create mode 100644 tests/test_logprob_contract.py create mode 100644 tests/test_qwen_ffn.py create mode 100644 tests/test_vocab_parallel_logp.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index dc4c773e..dd1baccf 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -95,6 +95,18 @@ jobs: run: | python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu" + - name: Run WS2 Logprob Contract Tests (CPU-safe) + run: python -m pytest tests/test_logprob_contract.py -v + + - name: Run WS2 Vocab-Parallel Logprob Tests (CPU-safe) + run: python -m pytest tests/test_vocab_parallel_logp.py -v + + - name: Run WS2 Wrapper Interface and FFN Tests (CPU-safe) + run: | + python -m pytest -q \ + tests/test_alignment_wrapper_interfaces.py \ + tests/test_qwen_ffn.py + docs: runs-on: ubuntu-latest steps: diff --git a/csrc/cuda/activation.cu b/csrc/cuda/activation.cu new file mode 100644 index 00000000..8f2c6239 --- /dev/null +++ b/csrc/cuda/activation.cu @@ -0,0 +1,253 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Batch-invariant SiLU / SwiGLU CUDA kernels (WS1 elementwise activations). +// +// Semantics match NativeSiLUOp / NativeSwiGLUOp: +// silu(x) = x * sigmoid(x) (math in fp32) +// swiglu(g, u) = silu(g) * u (math in fp32) +// +// Pure elementwise / token-local: no cross-row reduction, so batch size and +// padding cannot change a row's result (Axis-A bitwise invariance). + +#include +#include +#include +#include + +namespace { + +__device__ __forceinline__ float silu_f32(float x) { + // sigmoid(x) = 1 / (1 + exp(-x)); use expf for device fp32. + const float s = 1.0f / (1.0f + expf(-x)); + return x * s; +} + +__device__ __forceinline__ float silu_grad_f32(float x) { + // d/dx [x * s] = s + x * s * (1 - s) = s * (1 + x * (1 - s)), s = sigmoid(x) + const float s = 1.0f / (1.0f + expf(-x)); + return s * (1.0f + x * (1.0f - s)); +} + +template +__global__ void silu_forward_kernel( + const scalar_t* __restrict__ x, + scalar_t* __restrict__ y, + const int64_t n) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const float xv = static_cast(x[idx]); + y[idx] = static_cast(silu_f32(xv)); +} + +template +__global__ void silu_backward_kernel( + const scalar_t* __restrict__ dy, + const scalar_t* __restrict__ x, + scalar_t* __restrict__ dx, + const int64_t n) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const float dyv = static_cast(dy[idx]); + const float xv = static_cast(x[idx]); + dx[idx] = static_cast(dyv * silu_grad_f32(xv)); +} + +template +__global__ void swiglu_forward_kernel( + const scalar_t* __restrict__ gate, + const scalar_t* __restrict__ up, + scalar_t* __restrict__ y, + const int64_t n) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const float gv = static_cast(gate[idx]); + const float uv = static_cast(up[idx]); + y[idx] = static_cast(silu_f32(gv) * uv); +} + +template +__global__ void swiglu_backward_kernel( + const scalar_t* __restrict__ dy, + const scalar_t* __restrict__ gate, + const scalar_t* __restrict__ up, + scalar_t* __restrict__ d_gate, + scalar_t* __restrict__ d_up, + const int64_t n) { + const int64_t idx = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + if (idx >= n) { + return; + } + const float dyv = static_cast(dy[idx]); + const float gv = static_cast(gate[idx]); + const float uv = static_cast(up[idx]); + const float s = silu_f32(gv); + // d_up = dy * silu(gate); d_gate = dy * up * silu'(gate) + d_up[idx] = static_cast(dyv * s); + d_gate[idx] = static_cast(dyv * uv * silu_grad_f32(gv)); +} + +static void launch_1d(int64_t n, int& threads, int64_t& blocks) { + threads = 256; + blocks = (n + threads - 1) / threads; + if (blocks == 0) { + blocks = 1; + } +} + +static void check_cuda_contig(const torch::Tensor& t, const char* name) { + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + // Supported activation dtypes only: fp16 / bf16 / fp32 (reject float64). + TORCH_CHECK( + t.scalar_type() == at::kHalf || t.scalar_type() == at::kBFloat16 || + t.scalar_type() == at::kFloat, + name, + " must be fp16, bf16, or fp32, got ", + t.scalar_type()); +} + +static void check_same_device( + const torch::Tensor& lhs, + const torch::Tensor& rhs, + const char* lhs_name, + const char* rhs_name) { + TORCH_CHECK( + lhs.device() == rhs.device(), + lhs_name, + " and ", + rhs_name, + " must be on the same CUDA device, got ", + lhs.device(), + " and ", + rhs.device()); +} + +} // namespace + +torch::Tensor silu_forward_cuda(torch::Tensor x) { + check_cuda_contig(x, "x"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(x)); + auto y = torch::empty_like(x); + const int64_t n = x.numel(); + if (n == 0) { + return y; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "silu_forward_cuda", [&] { + silu_forward_kernel<<>>( + x.data_ptr(), y.data_ptr(), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return y; +} + +torch::Tensor silu_backward_cuda(torch::Tensor dy, torch::Tensor x) { + check_cuda_contig(dy, "dy"); + check_cuda_contig(x, "x"); + check_same_device(dy, x, "dy", "x"); + TORCH_CHECK(dy.sizes() == x.sizes(), "dy and x must share shape"); + TORCH_CHECK(dy.scalar_type() == x.scalar_type(), "dy and x must share dtype"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(x)); + auto dx = torch::empty_like(x); + const int64_t n = x.numel(); + if (n == 0) { + return dx; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, x.scalar_type(), "silu_backward_cuda", [&] { + silu_backward_kernel<<>>( + dy.data_ptr(), x.data_ptr(), dx.data_ptr(), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dx; +} + +torch::Tensor swiglu_forward_cuda(torch::Tensor gate, torch::Tensor up) { + check_cuda_contig(gate, "gate"); + check_cuda_contig(up, "up"); + check_same_device(gate, up, "gate", "up"); + TORCH_CHECK(gate.sizes() == up.sizes(), "gate and up must share shape"); + TORCH_CHECK(gate.scalar_type() == up.scalar_type(), "gate and up must share dtype"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(gate)); + auto y = torch::empty_like(gate); + const int64_t n = gate.numel(); + if (n == 0) { + return y; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + gate.scalar_type(), + "swiglu_forward_cuda", + [&] { + swiglu_forward_kernel<<>>( + gate.data_ptr(), up.data_ptr(), y.data_ptr(), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return y; +} + +std::vector swiglu_backward_cuda( + torch::Tensor dy, + torch::Tensor gate, + torch::Tensor up) { + check_cuda_contig(dy, "dy"); + check_cuda_contig(gate, "gate"); + check_cuda_contig(up, "up"); + check_same_device(dy, gate, "dy", "gate"); + check_same_device(gate, up, "gate", "up"); + TORCH_CHECK(gate.sizes() == up.sizes(), "gate and up must share shape"); + TORCH_CHECK(dy.sizes() == gate.sizes(), "dy and gate must share shape"); + TORCH_CHECK(dy.scalar_type() == gate.scalar_type(), "dy and gate must share dtype"); + TORCH_CHECK(up.scalar_type() == gate.scalar_type(), "up and gate must share dtype"); + const at::cuda::OptionalCUDAGuard device_guard(device_of(gate)); + auto d_gate = torch::empty_like(gate); + auto d_up = torch::empty_like(up); + const int64_t n = gate.numel(); + if (n == 0) { + return {d_gate, d_up}; + } + int threads = 0; + int64_t blocks = 0; + launch_1d(n, threads, blocks); + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + gate.scalar_type(), + "swiglu_backward_cuda", + [&] { + swiglu_backward_kernel<<>>( + dy.data_ptr(), + gate.data_ptr(), + up.data_ptr(), + d_gate.data_ptr(), + d_up.data_ptr(), + n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {d_gate, d_up}; +} diff --git a/csrc/cuda/distributed/deterministic_collective.cu b/csrc/cuda/distributed/deterministic_collective.cu new file mode 100644 index 00000000..63a21f94 --- /dev/null +++ b/csrc/cuda/distributed/deterministic_collective.cu @@ -0,0 +1,572 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors + +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kMaxDeterministicWorldSize = 8; +constexpr int kThreads = 256; +constexpr int kMaxBlocks = 4096; + +struct PeerPointers { + const void* values[kMaxDeterministicWorldSize]; +}; + +bool is_supported_world_size(int64_t world_size) { + return world_size == 1 || world_size == 2 || world_size == 4 || world_size == 8; +} + +template +__device__ __forceinline__ T ordered_add(T lower, T upper); + +template <> +__device__ __forceinline__ float ordered_add(float lower, float upper) { + float result; + asm volatile("add.rn.f32 %0, %1, %2;" : "=f"(result) : "f"(lower), "f"(upper)); + return result; +} + +template <> +__device__ __forceinline__ half ordered_add(half lower, half upper) { + return __hadd(lower, upper); +} + +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) +template <> +__device__ __forceinline__ nv_bfloat16 ordered_add( + nv_bfloat16 lower, + nv_bfloat16 upper) { + return __hadd(lower, upper); +} +#endif + +template +__device__ __forceinline__ T fixed_tree_reduce( + const PeerPointers& peers, + int64_t index) { + static_assert( + WorldSize == 1 || WorldSize == 2 || WorldSize == 4 || WorldSize == 8, + "deterministic collectives only support TP sizes 1, 2, 4, and 8"); + const auto* rank0 = static_cast(peers.values[0]); + if constexpr (WorldSize == 1) { + return rank0[index]; + } else { + const auto* rank1 = static_cast(peers.values[1]); + const T sum01 = ordered_add(rank0[index], rank1[index]); + if constexpr (WorldSize == 2) { + return sum01; + } else { + const auto* rank2 = static_cast(peers.values[2]); + const auto* rank3 = static_cast(peers.values[3]); + const T sum23 = ordered_add(rank2[index], rank3[index]); + const T sum03 = ordered_add(sum01, sum23); + if constexpr (WorldSize == 4) { + return sum03; + } else { + const auto* rank4 = static_cast(peers.values[4]); + const auto* rank5 = static_cast(peers.values[5]); + const auto* rank6 = static_cast(peers.values[6]); + const auto* rank7 = static_cast(peers.values[7]); + const T sum45 = ordered_add(rank4[index], rank5[index]); + const T sum67 = ordered_add(rank6[index], rank7[index]); + const T sum47 = ordered_add(sum45, sum67); + return ordered_add(sum03, sum47); + } + } + } +} + +template +__global__ void deterministic_all_reduce_kernel( + PeerPointers peers, + T* output, + int64_t element_count) { + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < element_count; index += stride) { + output[index] = fixed_tree_reduce(peers, index); + } +} + +template +void launch_all_reduce( + const PeerPointers& peers, + T* output, + int64_t element_count, + int blocks, + int64_t world_size, + cudaStream_t stream) { + switch (world_size) { + case 1: + deterministic_all_reduce_kernel<<>>( + peers, output, element_count); + break; + case 2: + deterministic_all_reduce_kernel<<>>( + peers, output, element_count); + break; + case 4: + deterministic_all_reduce_kernel<<>>( + peers, output, element_count); + break; + case 8: + deterministic_all_reduce_kernel<<>>( + peers, output, element_count); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } +} + +template +__global__ void deterministic_reduce_scatter_kernel( + PeerPointers peers, + T* output, + int64_t output_element_count, + int rank) { + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + const int64_t input_offset = static_cast(rank) * output_element_count; + for (int64_t index = thread_index; index < output_element_count; index += stride) { + output[index] = fixed_tree_reduce(peers, input_offset + index); + } +} + +template +void launch_reduce_scatter( + const PeerPointers& peers, + T* output, + int64_t output_element_count, + int rank, + int blocks, + int64_t world_size, + cudaStream_t stream) { + switch (world_size) { + case 1: + deterministic_reduce_scatter_kernel<<>>( + peers, output, output_element_count, rank); + break; + case 2: + deterministic_reduce_scatter_kernel<<>>( + peers, output, output_element_count, rank); + break; + case 4: + deterministic_reduce_scatter_kernel<<>>( + peers, output, output_element_count, rank); + break; + case 8: + deterministic_reduce_scatter_kernel<<>>( + peers, output, output_element_count, rank); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } +} + +__global__ void deterministic_all_gather_kernel( + PeerPointers peers, + uint8_t* output, + int64_t input_bytes, + int64_t world_size) { + const int64_t output_bytes = input_bytes * world_size; + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < output_bytes; index += stride) { + const int peer = static_cast(index / input_bytes); + const int64_t peer_offset = index - static_cast(peer) * input_bytes; + output[index] = static_cast(peers.values[peer])[peer_offset]; + } +} + +class DeterministicCollectiveState { + public: + DeterministicCollectiveState( + torch::Tensor& staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank) + : rank_(rank), + world_size_(handles.size()), + device_index_(staging.get_device()), + capacity_bytes_(staging.numel() * staging.element_size()) { + TORCH_CHECK(staging.is_cuda(), "collective staging buffer must be CUDA"); + TORCH_CHECK(staging.is_contiguous(), "collective staging buffer must be contiguous"); + TORCH_CHECK( + staging.scalar_type() == torch::kUInt8, + "collective staging buffer must have dtype torch.uint8"); + TORCH_CHECK(capacity_bytes_ > 0, "collective staging capacity must be positive"); + TORCH_CHECK( + is_supported_world_size(world_size_), + "deterministic collectives require world size 1, 2, 4, or 8; got ", + world_size_); + TORCH_CHECK( + offsets.size() == handles.size(), + "deterministic collectives require one IPC offset per handle"); + TORCH_CHECK( + rank_ >= 0 && rank_ < world_size_, + "deterministic collective rank must be in [0, ", + world_size_, + ")"); + + for (auto& peer : peers_.values) { + peer = nullptr; + } + imported_bases_.fill(nullptr); + try { + for (int peer = 0; peer < world_size_; ++peer) { + TORCH_CHECK( + handles[peer].size() == sizeof(cudaIpcMemHandle_t), + "invalid CUDA IPC handle size for rank ", + peer); + TORCH_CHECK(offsets[peer] >= 0, "negative CUDA IPC offset for rank ", peer); + + if (peer == rank_) { + peers_.values[peer] = staging.data_ptr(); + continue; + } + + cudaIpcMemHandle_t handle{}; + auto* raw_handle = reinterpret_cast(&handle); + for (size_t byte = 0; byte < sizeof(handle); ++byte) { + TORCH_CHECK( + handles[peer][byte] >= 0 && handles[peer][byte] <= 255, + "invalid CUDA IPC handle byte for rank ", + peer); + raw_handle[byte] = static_cast(handles[peer][byte]); + } + + void* base = nullptr; + AT_CUDA_CHECK(cudaIpcOpenMemHandle( + &base, + handle, + cudaIpcMemLazyEnablePeerAccess)); + imported_bases_[peer] = base; + peers_.values[peer] = static_cast(base) + offsets[peer]; + } + } catch (...) { + close_imports(); + throw; + } + } + + ~DeterministicCollectiveState() { + int previous_device = -1; + if (cudaGetDevice(&previous_device) == cudaSuccess && previous_device != device_index_) { + if (cudaSetDevice(device_index_) != cudaSuccess) { + return; + } + } + close_imports(); + if (previous_device >= 0 && previous_device != device_index_) { + cudaSetDevice(previous_device); + } + } + + void stage(torch::Tensor& input, cudaStream_t stream) { + check_tensor(input, "input"); + const int64_t input_bytes = input.numel() * input.element_size(); + TORCH_CHECK( + input_bytes <= capacity_bytes_, + "input requires ", + input_bytes, + " bytes but staging capacity is ", + capacity_bytes_); + if (input_bytes > 0) { + AT_CUDA_CHECK(cudaMemcpyAsync( + const_cast(peers_.values[rank_]), + input.data_ptr(), + input_bytes, + cudaMemcpyDeviceToDevice, + stream)); + } + staged_bytes_ = input_bytes; + staged_scalar_type_ = input.scalar_type(); + has_staged_input_ = true; + } + + void all_reduce(torch::Tensor& output, cudaStream_t stream) const { + check_tensor(output, "output"); + TORCH_CHECK(has_staged_input_, "stage() must be called before all_reduce()"); + TORCH_CHECK( + output.scalar_type() == staged_scalar_type_, + "all-reduce output dtype must match the staged input dtype"); + TORCH_CHECK( + output.numel() * output.element_size() == staged_bytes_, + "all-reduce output size must match the staged input size"); + + const int64_t element_count = output.numel(); + if (element_count == 0) { + return; + } + const int blocks = static_cast(std::min( + kMaxBlocks, + (element_count + kThreads - 1) / kThreads)); + + switch (output.scalar_type()) { + case at::ScalarType::Float: + launch_all_reduce( + peers_, + static_cast(output.data_ptr()), + element_count, + blocks, + world_size_, + stream); + break; + case at::ScalarType::Half: + launch_all_reduce( + peers_, + static_cast(output.data_ptr()), + element_count, + blocks, + world_size_, + stream); + break; +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case at::ScalarType::BFloat16: + launch_all_reduce( + peers_, + static_cast(output.data_ptr()), + element_count, + blocks, + world_size_, + stream); + break; +#endif + default: + TORCH_CHECK( + false, + "deterministic all-reduce supports float32, float16, and bfloat16; got ", + output.scalar_type()); + } + AT_CUDA_CHECK(cudaGetLastError()); + } + + void reduce_scatter(torch::Tensor& output, cudaStream_t stream) const { + check_tensor(output, "output"); + TORCH_CHECK(has_staged_input_, "stage() must be called before reduce_scatter()"); + TORCH_CHECK( + output.scalar_type() == staged_scalar_type_, + "reduce-scatter output dtype must match the staged input dtype"); + TORCH_CHECK( + output.numel() * output.element_size() * world_size_ == staged_bytes_, + "reduce-scatter output must contain one world-size fraction of the staged input"); + + const int64_t output_element_count = output.numel(); + if (output_element_count == 0) { + return; + } + const int blocks = static_cast(std::min( + kMaxBlocks, + (output_element_count + kThreads - 1) / kThreads)); + + switch (output.scalar_type()) { + case at::ScalarType::Float: + launch_reduce_scatter( + peers_, + static_cast(output.data_ptr()), + output_element_count, + rank_, + blocks, + world_size_, + stream); + break; + case at::ScalarType::Half: + launch_reduce_scatter( + peers_, + static_cast(output.data_ptr()), + output_element_count, + rank_, + blocks, + world_size_, + stream); + break; +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case at::ScalarType::BFloat16: + launch_reduce_scatter( + peers_, + static_cast(output.data_ptr()), + output_element_count, + rank_, + blocks, + world_size_, + stream); + break; +#endif + default: + TORCH_CHECK( + false, + "deterministic reduce-scatter supports float32, float16, and bfloat16; got ", + output.scalar_type()); + } + AT_CUDA_CHECK(cudaGetLastError()); + } + + void all_gather(torch::Tensor& output, cudaStream_t stream) const { + check_tensor(output, "output"); + TORCH_CHECK(has_staged_input_, "stage() must be called before all_gather()"); + TORCH_CHECK( + output.scalar_type() == staged_scalar_type_, + "all-gather output dtype must match the staged input dtype"); + TORCH_CHECK( + output.numel() * output.element_size() == + staged_bytes_ * world_size_, + "all-gather output must contain one staged input per rank"); + + const int64_t output_bytes = output.numel() * output.element_size(); + if (output_bytes == 0) { + return; + } + const int blocks = static_cast(std::min( + kMaxBlocks, + (output_bytes + kThreads - 1) / kThreads)); + deterministic_all_gather_kernel<<>>( + peers_, + static_cast(output.data_ptr()), + staged_bytes_, + world_size_); + AT_CUDA_CHECK(cudaGetLastError()); + } + + private: + void check_tensor(const torch::Tensor& tensor, const char* name) const { + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK( + tensor.get_device() == device_index_, + name, + " must be on cuda:", + device_index_, + ", got ", + tensor.device()); + } + + void close_imports() noexcept { + for (int peer = 0; peer < world_size_; ++peer) { + if (imported_bases_[peer] != nullptr) { + cudaIpcCloseMemHandle(imported_bases_[peer]); + imported_bases_[peer] = nullptr; + } + } + } + + int64_t rank_; + int64_t world_size_; + int device_index_; + int64_t capacity_bytes_; + int64_t staged_bytes_{0}; + at::ScalarType staged_scalar_type_{at::ScalarType::Undefined}; + bool has_staged_input_{false}; + PeerPointers peers_{}; + std::array imported_bases_{}; +}; + +DeterministicCollectiveState* state_from_handle(int64_t handle) { + TORCH_CHECK(handle != 0, "deterministic collective handle is closed"); + return reinterpret_cast(handle); +} + +} // namespace + +std::tuple, int64_t> deterministic_collective_ipc_meta( + torch::Tensor& tensor) { + const c10::cuda::CUDAGuard device_guard(tensor.device()); + TORCH_CHECK(tensor.is_cuda(), "IPC tensor must be CUDA"); + TORCH_CHECK(tensor.is_contiguous(), "IPC tensor must be contiguous"); + TORCH_CHECK(tensor.numel() > 0, "cannot export an empty CUDA allocation"); + + CUdeviceptr allocation_base = 0; + size_t allocation_size = 0; + const auto pointer = reinterpret_cast(tensor.data_ptr()); + TORCH_CHECK( + cuPointerGetAttribute( + &allocation_base, + CU_POINTER_ATTRIBUTE_RANGE_START_ADDR, + pointer) == CUDA_SUCCESS, + "failed to query CUDA allocation base"); + TORCH_CHECK( + cuPointerGetAttribute( + &allocation_size, + CU_POINTER_ATTRIBUTE_RANGE_SIZE, + pointer) == CUDA_SUCCESS, + "failed to query CUDA allocation size"); + + const int64_t offset = static_cast(pointer - allocation_base); + const int64_t tensor_bytes = tensor.numel() * tensor.element_size(); + TORCH_CHECK(offset >= 0, "invalid negative CUDA allocation offset"); + TORCH_CHECK( + static_cast(offset + tensor_bytes) <= allocation_size, + "IPC tensor exceeds its CUDA allocation"); + + cudaIpcMemHandle_t handle{}; + AT_CUDA_CHECK(cudaIpcGetMemHandle( + &handle, + reinterpret_cast(allocation_base))); + const auto* raw_handle = reinterpret_cast(&handle); + std::vector bytes(sizeof(handle)); + for (size_t byte = 0; byte < sizeof(handle); ++byte) { + bytes[byte] = raw_handle[byte]; + } + return std::make_tuple(bytes, offset); +} + +int64_t deterministic_collective_create( + torch::Tensor& staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank) { + const c10::cuda::CUDAGuard device_guard(staging.device()); + auto state = std::make_unique( + staging, + handles, + offsets, + rank); + return reinterpret_cast(state.release()); +} + +void deterministic_collective_destroy(int64_t handle) { + delete state_from_handle(handle); +} + +void deterministic_collective_stage(int64_t handle, torch::Tensor& input) { + const c10::cuda::CUDAGuard device_guard(input.device()); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + state_from_handle(handle)->stage(input, stream); +} + +void deterministic_collective_all_reduce(int64_t handle, torch::Tensor& output) { + const c10::cuda::CUDAGuard device_guard(output.device()); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + state_from_handle(handle)->all_reduce(output, stream); +} + +void deterministic_collective_reduce_scatter(int64_t handle, torch::Tensor& output) { + const c10::cuda::CUDAGuard device_guard(output.device()); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + state_from_handle(handle)->reduce_scatter(output, stream); +} + +void deterministic_collective_all_gather(int64_t handle, torch::Tensor& output) { + const c10::cuda::CUDAGuard device_guard(output.device()); + auto stream = c10::cuda::getCurrentCUDAStream().stream(); + state_from_handle(handle)->all_gather(output, stream); +} diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index 4d9035fc..04281f6a 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/csrc/cuda/gemm/det_gemm_kernel.cu @@ -8,7 +8,10 @@ // fixed K order, NO split-K -> batch-invariant. // Fallback : naive FP32 scalar kernel (also the correctness ground truth). // -// Both: BF16 in / FP32 accum / no TF32 / no split-K. +// Both: BF16 in / FP32 accum / BF16 store / no TF32 / no split-K. +// K is reduced with a mid-split tree. A contiguous half-K GEMM is one child, +// so simulated TP=2 (a+b) matches TP=1. TP=8 left-fold does not. +// Leaves stay FP32 (naive: 32-wide MAC; SM90: one BK). // fwd: C = A @ B | dA = dC @ B^T | dB = A^T @ dC // Backward reuses the forward kernel on transposed operands. @@ -26,8 +29,44 @@ using nv_bf16 = __nv_bfloat16; __host__ __device__ constexpr int cdiv(int a, int b) { return (a + b - 1) / b; } +// Must match SM90 BK so an aligned-K naive tree equals the SM90 tile tree. +constexpr int K_TREE_LEAF = 32; + +__device__ __forceinline__ nv_bf16 bf16_add(nv_bf16 a, nv_bf16 b) { + return __float2bfloat16(__bfloat162float(a) + __bfloat162float(b)); +} + +// True iff [lo, hi) is a node of the mid-split tree over [0, n). +__device__ __forceinline__ bool is_mid_split_node(int lo, int hi, int n) { + int a = 0, b = n; + while (b - a > 1) { + if (a == lo && b == hi) return true; + const int m = a + (b - a) / 2; + if (hi <= m) + b = m; + else if (lo >= m) + a = m; + else + return false; + } + return a == lo && b == hi; +} + +__device__ nv_bf16 k_tree_naive(const nv_bf16* __restrict__ A, const nv_bf16* __restrict__ B, + int row, int col, int N, int K, int lo, int hi) { + if (hi - lo <= K_TREE_LEAF) { + float acc = 0.0f; + for (int k = lo; k < hi; ++k) + acc += __bfloat162float(A[row * K + k]) * __bfloat162float(B[k * N + col]); + return __float2bfloat16(acc); + } + const int mid = lo + (hi - lo) / 2; + return bf16_add(k_tree_naive(A, B, row, col, N, K, lo, mid), + k_tree_naive(A, B, row, col, N, K, mid, hi)); +} + // Naive FP32 scalar kernel (fallback + ground truth). Batch-invariant by -// construction: one thread = one output element, fixed ascending K loop. +// construction: one thread = one output element, mid-split K tree. constexpr int NAIVE_TILE = 16; __global__ void det_gemm_naive(const nv_bf16* __restrict__ A, @@ -37,10 +76,7 @@ __global__ void det_gemm_naive(const nv_bf16* __restrict__ A, const int row = blockIdx.y * NAIVE_TILE + threadIdx.y; const int col = blockIdx.x * NAIVE_TILE + threadIdx.x; if (row >= M || col >= N) return; - float acc = 0.0f; - for (int k = 0; k < K; ++k) - acc += __bfloat162float(A[row * K + k]) * __bfloat162float(B[k * N + col]); - C[row * N + col] = __float2bfloat16(acc); + C[row * N + col] = k_tree_naive(A, B, row, col, N, K, 0, K); } void launch_naive(const nv_bf16* A, const nv_bf16* B, nv_bf16* C, @@ -57,6 +93,7 @@ void launch_naive(const nv_bf16* A, const nv_bf16* B, nv_bf16* C, // passing B^T ([N,K] row-major) so the B smem tile is [BN,BK] (row=n,col=k), // matching the validated logp ldmatrix addressing. constexpr int BM = 128, BN = 64, BK = 32; +static_assert(BK == K_TREE_LEAF, "SM90 tile width must match the naive K-tree leaf"); constexpr int WARPS = 4; constexpr int WG_THREADS = WARPS * 32; // 128 constexpr int STAGES = 2; @@ -67,6 +104,7 @@ constexpr int M_TILES = WARP_M / MMA_M; // 1 constexpr int N_TILES = BN / MMA_N; // 8 constexpr int K_TILES = BK / MMA_K; // 2 constexpr int KK_GROUPS = BK / 32; // 1 +constexpr int TREE_DEPTH = 16; __device__ __forceinline__ void ldmatrix_x4(uint32_t regs[4], uint32_t addr) { asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" @@ -125,19 +163,18 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap, #pragma unroll for (int s = 0; s < STAGES; ++s) phase[s] = 0; - float acc[M_TILES][N_TILES][4]; -#pragma unroll - for (int mi = 0; mi < M_TILES; ++mi) -#pragma unroll - for (int n = 0; n < N_TILES; ++n) - acc[mi][n][0] = acc[mi][n][1] = acc[mi][n][2] = acc[mi][n][3] = 0.0f; + float tile_acc[M_TILES][N_TILES][4]; + nv_bf16 tree_v[M_TILES][N_TILES][4]; + nv_bf16 tree_stk[TREE_DEPTH][M_TILES][N_TILES][4]; + int tree_lo[TREE_DEPTH], tree_hi[TREE_DEPTH]; + int sp = 0; if (tid == 0) #pragma unroll for (int s = 0; s < STAGES - 1; ++s) if (s < kd) issue_load(s); - for (int k = 0; k < kd; ++k) { // fixed ascending K order, NO split-K + for (int k = 0; k < kd; ++k) { // fixed ascending tile order, NO split-K const int buf = k % STAGES; if (tid == 0 && k + (STAGES - 1) < kd) issue_load(k + (STAGES - 1)); det_gemm::mbar_wait(mbar[buf], phase[buf]); @@ -147,6 +184,12 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap, const uint32_t sA_buf = sA_base + buf * BM * BK * sizeof(nv_bf16); const uint32_t sB_buf = sB_base + buf * BN * BK * sizeof(nv_bf16); +#pragma unroll + for (int mi = 0; mi < M_TILES; ++mi) +#pragma unroll + for (int n = 0; n < N_TILES; ++n) + tile_acc[mi][n][0] = tile_acc[mi][n][1] = tile_acc[mi][n][2] = tile_acc[mi][n][3] = 0.0f; + uint32_t A[M_TILES][K_TILES][4]; #pragma unroll for (int mi = 0; mi < M_TILES; ++mi) { @@ -171,12 +214,43 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap, const uint32_t B1[2] = {b4[2], b4[3]}; #pragma unroll for (int mi = 0; mi < M_TILES; ++mi) { - mma_m16n8k16(A[mi][2 * kk + 0], B0, acc[mi][n]); - mma_m16n8k16(A[mi][2 * kk + 1], B1, acc[mi][n]); + mma_m16n8k16(A[mi][2 * kk + 0], B0, tile_acc[mi][n]); + mma_m16n8k16(A[mi][2 * kk + 1], B1, tile_acc[mi][n]); } } } __syncthreads(); + +#pragma unroll + for (int mi = 0; mi < M_TILES; ++mi) +#pragma unroll + for (int n = 0; n < N_TILES; ++n) +#pragma unroll + for (int i = 0; i < 4; ++i) tree_v[mi][n][i] = __float2bfloat16(tile_acc[mi][n][i]); + + int lo = k, hi = k + 1; + while (sp > 0 && tree_hi[sp - 1] == lo && is_mid_split_node(tree_lo[sp - 1], hi, kd)) { +#pragma unroll + for (int mi = 0; mi < M_TILES; ++mi) +#pragma unroll + for (int n = 0; n < N_TILES; ++n) +#pragma unroll + for (int i = 0; i < 4; ++i) + tree_v[mi][n][i] = bf16_add(tree_stk[sp - 1][mi][n][i], tree_v[mi][n][i]); + lo = tree_lo[sp - 1]; + --sp; + } + if (hi < kd) { +#pragma unroll + for (int mi = 0; mi < M_TILES; ++mi) +#pragma unroll + for (int n = 0; n < N_TILES; ++n) +#pragma unroll + for (int i = 0; i < 4; ++i) tree_stk[sp][mi][n][i] = tree_v[mi][n][i]; + tree_lo[sp] = lo; + tree_hi[sp] = hi; + ++sp; + } } #pragma unroll @@ -186,12 +260,12 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap, for (int n = 0; n < N_TILES; ++n) { const int col = col_base + n * MMA_N + (lane % 4) * 2; if (row < M && col + 1 < N) { - C[row * N + col + 0] = __float2bfloat16(acc[mi][n][0]); - C[row * N + col + 1] = __float2bfloat16(acc[mi][n][1]); + C[row * N + col + 0] = tree_v[mi][n][0]; + C[row * N + col + 1] = tree_v[mi][n][1]; } if (row + 8 < M && col + 1 < N) { - C[(row + 8) * N + col + 0] = __float2bfloat16(acc[mi][n][2]); - C[(row + 8) * N + col + 1] = __float2bfloat16(acc[mi][n][3]); + C[(row + 8) * N + col + 0] = tree_v[mi][n][2]; + C[(row + 8) * N + col + 1] = tree_v[mi][n][3]; } } } @@ -271,6 +345,15 @@ torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b) { return gemm_dispatch(a, b); } +torch::Tensor det_gemm_fwd_fp32(torch::Tensor a, torch::Tensor b) { + check_in(a, "A"); check_in(b, "B"); + a = a.contiguous(); b = b.contiguous(); + TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "det_gemm_fwd_fp32: expect 2D [M,K]@[K,N]"); + TORCH_CHECK(b.size(0) == a.size(1), "det_gemm_fwd_fp32: K mismatch"); + // Keep the FP32 running sum; only the final store is BF16. + return gemm_dispatch(a, b); +} + torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b) { check_in(dc, "dC"); check_in(b, "B"); dc = dc.contiguous(); diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 8ff8091f..9e986a14 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -80,10 +80,33 @@ torch::Tensor deterministic_logp_forward_fp32(torch::Tensor logits, torch::Tenso torch::Tensor deterministic_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices); +// Single-node TP=8 deterministic collectives. +std::tuple, int64_t> deterministic_collective_ipc_meta( + torch::Tensor& tensor); +int64_t deterministic_collective_create( + torch::Tensor& staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank); +void deterministic_collective_destroy(int64_t handle); +void deterministic_collective_stage(int64_t handle, torch::Tensor& input); +void deterministic_collective_all_reduce(int64_t handle, torch::Tensor& output); +void deterministic_collective_reduce_scatter(int64_t handle, torch::Tensor& output); +void deterministic_collective_all_gather(int64_t handle, torch::Tensor& output); + // Batch-Invariant Deterministic GEMM Declarations torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b); torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b); torch::Tensor det_gemm_db(torch::Tensor a, torch::Tensor dc); +// SiLU / SwiGLU Declarations (elementwise activation, general CUDA) +torch::Tensor silu_forward_cuda(torch::Tensor x); +torch::Tensor silu_backward_cuda(torch::Tensor dy, torch::Tensor x); +torch::Tensor swiglu_forward_cuda(torch::Tensor gate, torch::Tensor up); +std::vector swiglu_backward_cuda( + torch::Tensor dy, + torch::Tensor gate, + torch::Tensor up); + // RMSNorm Declarations & Wrappers void rmsnorm_forward_cuda( @@ -195,6 +218,26 @@ torch::Tensor rmsnorm_backward_dw( return dw; } +// SiLU / SwiGLU wrappers (WS1 elementwise activations) +torch::Tensor silu_forward(torch::Tensor x) { + return silu_forward_cuda(x); +} + +torch::Tensor silu_backward(torch::Tensor dy, torch::Tensor x) { + return silu_backward_cuda(dy, x); +} + +torch::Tensor swiglu_forward(torch::Tensor gate, torch::Tensor up) { + return swiglu_forward_cuda(gate, up); +} + +std::vector swiglu_backward( + torch::Tensor dy, + torch::Tensor gate, + torch::Tensor up) { + return swiglu_backward_cuda(dy, gate, up); +} + // Deterministic standard-softmax attention (issue #147) std::vector deterministic_attention_forward( torch::Tensor q, @@ -311,6 +354,36 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("deterministic_logp_forward_indexed_out", &deterministic_logp_forward_indexed_out, "Batch-invariant deterministic logp indexed out"); m.def("deterministic_logp_forward_indexed_fp32", &deterministic_logp_forward_indexed_fp32, "Batch-invariant deterministic logp indexed fp32"); + // Single-node TP=8 fixed-tree collectives. + m.def( + "deterministic_collective_ipc_meta", + &deterministic_collective_ipc_meta, + "Export a CUDA allocation for deterministic collectives"); + m.def( + "deterministic_collective_create", + &deterministic_collective_create, + "Create a single-node TP=8 deterministic collective state"); + m.def( + "deterministic_collective_destroy", + &deterministic_collective_destroy, + "Destroy a deterministic collective state"); + m.def( + "deterministic_collective_stage", + &deterministic_collective_stage, + "Stage one rank's input in symmetric CUDA IPC memory"); + m.def( + "deterministic_collective_all_reduce", + &deterministic_collective_all_reduce, + "Run the TP=8 deterministic fixed-tree all-reduce kernel"); + m.def( + "deterministic_collective_reduce_scatter", + &deterministic_collective_reduce_scatter, + "Run the TP=8 deterministic fixed-tree reduce-scatter kernel"); + m.def( + "deterministic_collective_all_gather", + &deterministic_collective_all_gather, + "Run the TP=8 deterministic rank-ordered all-gather kernel"); + // registry Prefix-Shared Attention m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO"); @@ -323,6 +396,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("rmsnorm_backward_dx", &rmsnorm_backward_dx, "Batch-invariant RMSNorm backward dx CUDA"); m.def("rmsnorm_backward_dw", &rmsnorm_backward_dw, "Deterministic RMSNorm backward dweight CUDA"); + // registry SiLU / SwiGLU (elementwise activation) + m.def("silu_forward", &silu_forward, "Batch-invariant SiLU forward CUDA"); + m.def("silu_backward", &silu_backward, "Batch-invariant SiLU backward CUDA"); + m.def("swiglu_forward", &swiglu_forward, "Batch-invariant SwiGLU forward CUDA"); + m.def("swiglu_backward", &swiglu_backward, "Batch-invariant SwiGLU backward CUDA"); + // Deterministic standard-softmax attention (issue #147) m.def( "deterministic_attention_forward", diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index bedf3475..23c1586a 100644 --- a/docs/design/runtime-dispatch.md +++ b/docs/design/runtime-dispatch.md @@ -11,6 +11,14 @@ logical type, and the registry selects the first available backend for the curre 4. Cache successfully constructed operator instances. 5. Skip backends that already failed in the current process. +WS2 TP-aware logprob uses the stricter `KernelRegistry.get_logprob_op(contract)` path. In +addition to platform priority, this path requires a backend capability descriptor and checks +the requested role, dtype, TP/CP layout, padded-vs-real vocab masking, inactive-token +support, vocab-domain LSE export, and deterministic TP merge semantics. Incompatible +candidates produce explicit rejection reasons and are never used as an undeclared fallback. +The contract objects and their normative reduction semantics are documented in +`rl_engine.kernels.logprob_contract`. + ## LogP Priority | Platform | Priority | diff --git a/docs/operators/activation.md b/docs/operators/activation.md index 6487d4a8..a6f2cf46 100644 --- a/docs/operators/activation.md +++ b/docs/operators/activation.md @@ -1,13 +1,15 @@ # SiLU / SwiGLU Activation -The activation operators are the element-wise core of the Qwen3/Llama gated MLP. They are -**WS1 ground-truth references** (issue #108): pure-PyTorch, fp32-accumulating definitions of -the "correct answer" that downstream fused CUDA/Triton MLP kernels are validated against. +The activation operators are the element-wise core of the Qwen3/Llama gated MLP. They +implement the WS1 dual-path contract (issue #108): pure-PyTorch fp32 ground truth, plus +CUDA and Triton candidates that validate against it. -- **SiLU** (`NativeSiLUOp`): `silu(x) = x * sigmoid(x)` — the `hidden_act="silu"` gate. -- **SwiGLU** (`NativeSwiGLUOp`): `swiglu(gate, up) = silu(gate) * up` — the gated MLP middle - stage. `gate` / `up` are the `gate_proj` / `up_proj` outputs (already at the intermediate - width); the following `down_proj` is a plain Matmul and is **not** part of this operator. +- **SiLU** (`NativeSiLUOp` / `SiLUCudaOp` / `TritonSiLUOp`): `silu(x) = x * sigmoid(x)` — + the `hidden_act="silu"` gate. +- **SwiGLU** (`NativeSwiGLUOp` / `SwiGLUCudaOp` / `TritonSwiGLUOp`): + `swiglu(gate, up) = silu(gate) * up` — the gated MLP middle stage. `gate` / `up` are the + `gate_proj` / `up_proj` outputs (already at the intermediate width); the following + `down_proj` is a plain Matmul and is **not** part of this operator. ```text hidden --gate_proj--> gate --\ @@ -29,7 +31,7 @@ y = silu(x) # [..., N] -> [..., N] h = swiglu(gate, up) # [..., I], [..., I] -> [..., I] ``` -Both ops expose the WS1 dual-path contract: +All backends expose the WS1 dual-path contract: - `forward(...)` — computes in fp32, casts back to the input dtype (Axis-B accuracy candidate / dtype-behavior path). @@ -40,7 +42,8 @@ Both ops expose the WS1 dual-path contract: | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | PyTorch fallback | `NativeSiLUOp` / `NativeSwiGLUOp` | None | fp32 ground-truth reference; CPU and any GPU. | -| CUDA / ROCm / Triton | — | — | Planned: downstream fused MLP kernels validate against this reference. | +| CUDA | `SiLUCudaOp` / `SwiGLUCudaOp` | `_C.silu_*` / `_C.swiglu_*` | General CUDA (fp16/bf16/fp32); math in fp32. | +| Triton | `TritonSiLUOp` / `TritonSwiGLUOp` | Triton JIT | Portable GPU baseline; same fp32 math contract. | ## Tensor Contract @@ -48,7 +51,7 @@ Both ops expose the WS1 dual-path contract: | --- | --- | --- | --- | | `x` (SiLU) | `[..., N]` | float (fp16/bf16/fp32) | Any shape; last dim arbitrary (Qwen3-8B `I=12288`). | | `gate` (SwiGLU) | `[..., I]` | float | `gate_proj` output. | -| `up` (SwiGLU) | `[..., I]` | float | `up_proj` output; **must share `gate`'s shape**. | +| `up` (SwiGLU) | `[..., I]` | float | `up_proj` output; **must share `gate`'s shape, dtype, and device**. | | output | same as input | `forward`: input dtype · `forward_fp32`: float32 | Same shape as input. | Element-wise and shape-agnostic: the Qwen3-8B intermediate dim `I=12288` is just one valid @@ -57,11 +60,16 @@ mutation, device/dtype follow the inputs. ## Dispatch Behavior -`kernel_registry.get_op("silu" | "swiglu")` resolves through the `OpBackend` priority map. -On `cuda` / `rocm` / `cpu` the only registered backend today is the PyTorch native op -(`PYTORCH_NATIVE_SILU` / `PYTORCH_NATIVE_SWIGLU`), so every device dispatches to the -fp32 reference. When fused kernels land, they are prepended to the priority list and the -native op becomes the fallback. +`kernel_registry.get_op("silu" | "swiglu")` resolves through the `OpBackend` priority map: + +| Platform | Priority | +| --- | --- | +| `cuda` | CUDA → Triton → PyTorch native | +| `rocm` | Triton → PyTorch native | +| `cpu` | PyTorch native | + +If the CUDA extension is not built (or symbols are missing), the registry falls back to +Triton, then to the native gold. ## Accuracy @@ -77,17 +85,30 @@ out = gate_f * torch.sigmoid(gate_f) * up.float() ``` - **Ground truth**: `forward_fp32` always accumulates in and returns fp32. -- **Dtype path**: `forward` runs the same fp32 math, then casts back to the input dtype; - it is bitwise-equal to `forward_fp32(x).to(dtype)`. +- **Dtype path**: `forward` runs the same fp32 math, then casts back to the input dtype. - **Axis A — batch invariance**: element-wise and row-independent, so a row's output is bitwise-identical regardless of batch size or padding (`torch.equal`, `atol=0`). - **Axis B — tolerance**: as `elementwise` ops, low-precision tolerance follows the - `elementwise` row of the WS1 numerical contract. + `elementwise` row of the WS1 numerical contract (`tolerance_contract.json`). + +## Ground-truth harness + +CUDA and Triton candidates are registered in `OP_SPECS` and can be checked with the +shared issue-#108 CLI: + +```bash +python scripts/check_operator.py --op silu --candidate cuda --dtype bf16 --device cuda +python scripts/check_operator.py --op swiglu --candidate triton --dtype bf16 --device cuda --check-grad +python scripts/check_operator.py --op silu --candidate pytorch --dtype fp32 --device cpu --check-grad +``` + +Gold path: `NativeSiLUOp.forward_fp32` / `NativeSwiGLUOp.forward_fp32`. ## Performance Notes -Reference operators — no fused kernel or benchmark yet. Downstream fused MLP kernels carry -their own benchmarks and are measured against this reference for correctness. +Element-wise kernels with a fixed 1-D grid (CUDA) / `BLOCK=1024` (Triton). Suitable as the +standalone WS1 activation path; fused bias+SiLU MLP kernels remain a separate future work +item and should continue to validate against this reference. ## Tests @@ -96,17 +117,22 @@ python -m pytest tests/test_swiglu.py -v ``` Covers: correctness vs an independent fp32 formula, dtype paths, Axis-A batch invariance -(slice + padding), input purity, gradient flow, the SwiGLU shape guard, and registry -dispatch. +(slice + padding), input purity, gradient flow, the SwiGLU shape guard, CUDA/Triton vs +native forward+backward, registry dispatch, and the issue-#108 `OP_SPECS` harness. ## Implementation Files -- `rl_engine/kernels/ops/pytorch/activation/swiglu.py` +- `rl_engine/kernels/ops/pytorch/activation/swiglu.py` — gold +- `rl_engine/kernels/ops/cuda/activation/swiglu.py` — CUDA wrappers +- `rl_engine/kernels/ops/triton/activation/swiglu.py` — Triton kernels +- `csrc/cuda/activation.cu` — CUDA kernels - `rl_engine/kernels/registry.py` +- `rl_engine/kernels/gtest/operator_specs.py` - `tests/test_swiglu.py` ## Known Limitations -- PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work). -- SwiGLU requires `gate` and `up` to share shape (raises `ValueError` otherwise); no - broadcasting. +- SwiGLU requires `gate` and `up` to share shape, dtype, and device; no broadcasting. +- No fused `bias + SiLU` or `chunk(y,2) + silu_and_mul` variant yet (vLLM-style + `SiluAndMul` on a packed gate/up tensor). Callers that hold a packed tensor should + split first, then call `swiglu`. diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index fbc0e9f1..f840ecfd 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -54,6 +54,39 @@ CUDA priority list when the extension exposes `_C.batch_invariant_logp_sm90` (built with `KERNEL_ALIGN_FORCE_SM90=1`) on an SM90 device. On any other build or device, dispatch is unchanged (Triton -> PyTorch). +## Tensor Parallel + +`VocabParallelLogprobOp` +(`rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py`) +**TP=1, TP=2, and TP=4 produce bit-identical results.** + +The backends above are single-shard (TP=1) references and do not yet export vocab-domain +LSE or carry vocab-shard metadata, so they are declared incompatible with strict WS2 +requests instead of being selected as a silent fallback. The contract objects are +documented in `rl_engine.kernels.logprob_contract`. + +The TP-aware implementation uses the following fixed-order construction: + +1. Split the padded vocabulary into `num_vocab_tiles` fixed tiles. +2. Each rank computes fp32 `(max, sumexp)` for the tiles it owns. Every tile + is reduced as the same contiguous `[n, tile]` shape, on any rank. +3. All tile partials are shared with `all_gather`. The collective only moves + bytes; it never does math, so it cannot round anything. +4. Every rank merges all tiles in the same fixed order, over the same + `[n, num_vocab_tiles]` shape. `LSE = M + log(sum(s_t * exp(m_t - M)))`. +5. The target logit is copied from the rank that owns it (never summed). +6. `logp = target_logit - LSE`. Inactive rows become `0.0`. + +Usage goes through the contract-aware entry point: + +```python +from rl_engine.kernels.registry import kernel_registry + +result = kernel_registry.get_logprob_op(contract) # LogprobContract from +op = result.op # rl_engine.kernels.logprob_contract +logp, lse = op(local_logits, target_ids, contract=contract, tp_group=tp_group) +``` + ## Benchmarks `benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton, and the @@ -224,3 +257,6 @@ WSL/Linux with CUDA. - `rl_engine/kernels/registry.py` - `tests/test_batch_invariant_logp.py` - `benchmarks/benchmark_batch_invariant_logp.py` +- `rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py` +- `rl_engine/kernels/logprob_contract.py` +- `tests/test_vocab_parallel_logp.py` diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index c28681fd..a5c1e816 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -2,6 +2,20 @@ # This file is a type stub for the compiled C++ extension module. import torch +def deterministic_collective_ipc_meta( + tensor: torch.Tensor, +) -> tuple[list[int], int]: ... +def deterministic_collective_create( + staging: torch.Tensor, + handles: list[list[int]], + offsets: list[int], + rank: int, +) -> int: ... +def deterministic_collective_destroy(handle: int) -> None: ... +def deterministic_collective_stage(handle: int, input: torch.Tensor) -> None: ... +def deterministic_collective_all_reduce(handle: int, output: torch.Tensor) -> None: ... +def deterministic_collective_reduce_scatter(handle: int, output: torch.Tensor) -> None: ... +def deterministic_collective_all_gather(handle: int, output: torch.Tensor) -> None: ... def fused_logp(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: ... def fused_logp_sm90(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: ... def batch_invariant_logp_sm90( @@ -156,3 +170,28 @@ def deterministic_attention_backward( scale: float, key_padding_mask: torch.Tensor | None, ) -> list[torch.Tensor]: ... +def silu_forward(x: torch.Tensor) -> torch.Tensor: ... +def silu_backward(dy: torch.Tensor, x: torch.Tensor) -> torch.Tensor: ... +def swiglu_forward(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: ... +def swiglu_backward( + dy: torch.Tensor, + gate: torch.Tensor, + up: torch.Tensor, +) -> list[torch.Tensor]: ... +def rmsnorm_forward( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> list[torch.Tensor]: ... +def rmsnorm_backward_dx( + dy: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, +) -> torch.Tensor: ... +def rmsnorm_backward_dw( + dy: torch.Tensor, + x: torch.Tensor, + rstd: torch.Tensor, + mask: torch.Tensor, +) -> torch.Tensor: ... diff --git a/rl_engine/alignment/cross_config/operators.py b/rl_engine/alignment/cross_config/operators.py index e3105a28..2b359219 100644 --- a/rl_engine/alignment/cross_config/operators.py +++ b/rl_engine/alignment/cross_config/operators.py @@ -11,6 +11,7 @@ import torch +from rl_engine.kernels.logprob_contract import LogprobContract from rl_engine.kernels.semantic_registry import ( OperatorInstanceProvenance, OperatorRequirements, @@ -201,8 +202,13 @@ def selected_logprobs_with_operator( active_mask: Optional[torch.Tensor] = None, temperature: float = 1.0, output_dtype: torch.dtype = torch.float32, + contract: Optional[LogprobContract] = None, + tp_group: Any = None, + num_vocab_tiles: Optional[int] = None, + deterministic: bool = True, + validate: bool = True, ) -> torch.Tensor: - """Apply the repository selected-logprob interface with common mask semantics.""" + """Apply legacy or contract-aware selected-logprob operators uniformly.""" if not math.isfinite(temperature) or temperature <= 0.0: raise ValueError("temperature must be finite and greater than zero") @@ -219,13 +225,49 @@ def selected_logprobs_with_operator( mask = active_mask.to(device=logits.device, dtype=torch.bool) safe_token_ids = safe_token_ids.masked_fill(~mask, 0) - scaled_logits = logits.float() / float(temperature) - if hasattr(operator, "apply_fp32") and callable(operator.apply_fp32): - selected = operator.apply_fp32(scaled_logits, safe_token_ids) - elif callable(operator): - selected = operator(scaled_logits, safe_token_ids) + if contract is not None: + if not isinstance(contract, LogprobContract): + raise TypeError("contract must be a LogprobContract") + contract_mask = torch.tensor(contract.mask.active_mask, dtype=torch.bool) + if contract_mask.numel() != token_ids.numel(): + raise ValueError("contract active_mask size must match token_ids") + if mask is not None and not torch.equal(contract_mask, mask.detach().cpu().reshape(-1)): + raise ValueError("active_mask must match contract.mask.active_mask") + kwargs: dict[str, Any] = { + "contract": contract, + "tp_group": tp_group, + "deterministic": deterministic, + "validate": validate, + } + if num_vocab_tiles is not None: + kwargs["num_vocab_tiles"] = num_vocab_tiles + if not callable(operator): + raise TypeError("contract-aware selected-logprob operator must be callable") + result = operator( + (logits / float(temperature)).reshape(-1, logits.shape[-1]), + safe_token_ids.reshape(-1), + **kwargs, + ) + if isinstance(result, tuple): + if len(result) != 2: + raise TypeError("contract-aware selected-logprob operator must return (logp, lse)") + selected, lse = result + if not isinstance(selected, torch.Tensor) or not isinstance(lse, torch.Tensor): + raise TypeError("contract-aware selected-logprob operator must return (logp, lse)") + if lse.shape != safe_token_ids.reshape(-1).shape: + raise ValueError("vocab LSE output shape must match token_ids") + else: + selected = result + if isinstance(selected, torch.Tensor): + selected = selected.reshape(token_ids.shape) else: - raise TypeError("selected-logprob operator must be callable or expose apply_fp32") + scaled_logits = logits.float() / float(temperature) + if hasattr(operator, "apply_fp32") and callable(operator.apply_fp32): + selected = operator.apply_fp32(scaled_logits, safe_token_ids) + elif callable(operator): + selected = operator(scaled_logits, safe_token_ids) + else: + raise TypeError("selected-logprob operator must be callable or expose apply_fp32") if not isinstance(selected, torch.Tensor): raise TypeError("selected-logprob operator must return a torch.Tensor") if selected.shape != token_ids.shape: diff --git a/rl_engine/distributed/__init__.py b/rl_engine/distributed/__init__.py new file mode 100644 index 00000000..37698f1a --- /dev/null +++ b/rl_engine/distributed/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from rl_engine.distributed.collectives import DeterministicCollective + +__all__ = ["DeterministicCollective"] diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py new file mode 100644 index 00000000..54b00fa5 --- /dev/null +++ b/rl_engine/distributed/collectives.py @@ -0,0 +1,325 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import socket +import threading +from types import TracebackType +from typing import Any + +import torch +import torch.distributed as dist + +_SUPPORTED_WORLD_SIZES = (1, 2, 4, 8) +_DEFAULT_MAX_SIZE_BYTES = 64 * 1024 * 1024 +_REDUCTION_DTYPES = (torch.float32, torch.float16, torch.bfloat16) + + +class DeterministicCollective: + """Correctness-first TP-invariant CUDA collectives for one eight-GPU node. + + TP sizes 1, 2, 4, and 8 use nested prefixes of the same balanced tree. + A reduction is cross-TP bitwise invariant when every rank input is the + corresponding contiguous subtree root of one canonical finest-grained + reduction, as produced by a TBIK-compatible row-parallel kernel. Every + node evaluates the lower logical subtree before the higher one. + + One instance owns a symmetric CUDA IPC staging buffer. All ranks must call + its methods in the same order with matching shapes and dtypes. Calls are + host-synchronizing by design; the first version prioritizes determinism and + lifetime safety over overlap or throughput. + """ + + def __init__( + self, + group: dist.ProcessGroup | None = None, + device: torch.device | str | int | None = None, + *, + max_size_bytes: int = _DEFAULT_MAX_SIZE_BYTES, + ) -> None: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before collectives") + if not torch.cuda.is_available(): + raise RuntimeError("deterministic collectives require CUDA") + if max_size_bytes <= 0: + raise ValueError("max_size_bytes must be positive") + + self.group = group if group is not None else dist.group.WORLD + self.rank = dist.get_rank(group=self.group) + self.world_size = dist.get_world_size(group=self.group) + if self.world_size not in _SUPPORTED_WORLD_SIZES: + raise ValueError( + "deterministic collectives require world_size in " + f"{_SUPPORTED_WORLD_SIZES}, got {self.world_size}" + ) + + if device is None: + normalized_device = torch.device("cuda", torch.cuda.current_device()) + elif isinstance(device, int): + normalized_device = torch.device("cuda", device) + else: + normalized_device = torch.device(device) + if normalized_device.type != "cuda": + raise ValueError(f"deterministic collectives require a CUDA device, got {device!r}") + if normalized_device.index is None: + normalized_device = torch.device("cuda", torch.cuda.current_device()) + if normalized_device.index != torch.cuda.current_device(): + raise ValueError( + "the collective device must be the current CUDA device; call " + f"torch.cuda.set_device({normalized_device.index}) first" + ) + + try: + from rl_engine import _C + except ImportError as exc: + raise RuntimeError( + "the RL-Kernel CUDA extension is required; rebuild with " + "`pip install --no-build-isolation -e .`" + ) from exc + required_symbols = ( + "deterministic_collective_ipc_meta", + "deterministic_collective_create", + "deterministic_collective_destroy", + "deterministic_collective_stage", + "deterministic_collective_all_reduce", + "deterministic_collective_reduce_scatter", + "deterministic_collective_all_gather", + ) + missing = [name for name in required_symbols if not hasattr(_C, name)] + if missing: + raise RuntimeError( + "the RL-Kernel CUDA extension lacks deterministic collectives: " + + ", ".join(missing) + ) + + self.device = normalized_device + self.max_size_bytes = int(max_size_bytes) + self._extension = _C + self._lock = threading.Lock() + self._handle = 0 + self._staging = torch.empty( + self.max_size_bytes, + dtype=torch.uint8, + device=self.device, + ) + + handle, offset = self._extension.deterministic_collective_ipc_meta(self._staging) + local_meta = { + "handle": handle, + "offset": int(offset), + "capacity": self.max_size_bytes, + "hostname": socket.gethostname(), + } + gathered_meta: list[dict[str, Any] | None] = [None] * self.world_size + dist.all_gather_object(gathered_meta, local_meta, group=self.group) + if any(meta is None for meta in gathered_meta): + raise RuntimeError("failed to exchange CUDA IPC metadata") + complete_meta = [meta for meta in gathered_meta if meta is not None] + hostnames = {meta["hostname"] for meta in complete_meta} + if len(hostnames) != 1: + raise ValueError("deterministic collectives require all ranks on one host") + capacities = {meta["capacity"] for meta in complete_meta} + if capacities != {self.max_size_bytes}: + raise ValueError("all ranks must use the same max_size_bytes") + + handles = [meta["handle"] for meta in complete_meta] + offsets = [meta["offset"] for meta in complete_meta] + self._handle = self._extension.deterministic_collective_create( + self._staging, + handles, + offsets, + self.rank, + ) + self._synchronize_ranks() + + def all_reduce( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Return the TBIK-compatible fixed-tree sum on every rank. + + Supported dtypes are float32, float16, and bfloat16. ``out`` may alias + ``input``; the input is staged before the output kernel starts. Cross-TP + invariance requires inputs to follow the class-level subtree contract. + """ + + self._check_open() + self._validate_reduction_input(input) + if out is None: + out = torch.empty_like(input) + self._validate_output(out, input) + + with self._lock: + self._validate_matching_signature("all_reduce", input) + self._extension.deterministic_collective_stage(self._handle, input) + self._synchronize_ranks() + self._extension.deterministic_collective_all_reduce(self._handle, out) + self._synchronize_ranks() + return out + + def all_gather( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Gather rank-ordered input bit patterns along dimension 0. + + The result is cross-TP invariant when every TP configuration partitions + the same global dimension 0 into equal contiguous rank-ordered shards. + """ + + self._check_open() + self._validate_gather_input(input) + output_shape = (input.size(0) * self.world_size, *input.shape[1:]) + if out is None: + out = torch.empty(output_shape, dtype=input.dtype, device=input.device) + self._validate_sharded_output(out, input, output_shape) + + with self._lock: + self._validate_matching_signature("all_gather", input) + self._extension.deterministic_collective_stage(self._handle, input) + self._synchronize_ranks() + self._extension.deterministic_collective_all_gather(self._handle, out) + self._synchronize_ranks() + return out + + def reduce_scatter( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """TBIK-compatible fixed-tree sum, then a rank-ordered dimension-0 scatter. + + Concatenating all rank outputs is cross-TP invariant when inputs follow + the class-level subtree contract and the global dimension-0 size is fixed. + """ + + self._check_open() + self._validate_reduction_input(input) + if input.dim() == 0: + raise ValueError("reduce_scatter input must have at least one dimension") + if input.size(0) % self.world_size != 0: + raise ValueError( + "reduce_scatter input.size(0) must be divisible by " + f"world_size={self.world_size}; got {input.size(0)}" + ) + output_shape = (input.size(0) // self.world_size, *input.shape[1:]) + if out is None: + out = torch.empty(output_shape, dtype=input.dtype, device=input.device) + self._validate_sharded_output(out, input, output_shape) + + with self._lock: + self._validate_matching_signature("reduce_scatter", input) + self._extension.deterministic_collective_stage(self._handle, input) + self._synchronize_ranks() + self._extension.deterministic_collective_reduce_scatter(self._handle, out) + self._synchronize_ranks() + return out + + def close(self) -> None: + """Release imported CUDA IPC mappings after the last collective call.""" + + handle = getattr(self, "_handle", 0) + if not handle: + return + torch.cuda.synchronize(self.device) + self._handle = 0 + self._extension.deterministic_collective_destroy(handle) + + def __enter__(self) -> DeterministicCollective: + self._check_open() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.close() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + def _check_open(self) -> None: + if not getattr(self, "_handle", 0): + raise RuntimeError("deterministic collective is closed") + + def _validate_reduction_input(self, input: torch.Tensor) -> None: + if not input.is_cuda or input.device != self.device: + raise ValueError(f"input must be on {self.device}, got {input.device}") + if not input.is_contiguous(): + raise ValueError("input must be contiguous") + if input.dtype not in _REDUCTION_DTYPES: + raise TypeError( + "deterministic reductions support float32, float16, and bfloat16; " + f"got {input.dtype}" + ) + input_bytes = input.numel() * input.element_size() + if input_bytes > self.max_size_bytes: + raise ValueError( + f"input requires {input_bytes} bytes but max_size_bytes={self.max_size_bytes}" + ) + + def _validate_gather_input(self, input: torch.Tensor) -> None: + if not input.is_cuda or input.device != self.device: + raise ValueError(f"input must be on {self.device}, got {input.device}") + if not input.is_contiguous(): + raise ValueError("input must be contiguous") + if input.dim() == 0: + raise ValueError("all_gather input must have at least one dimension") + input_bytes = input.numel() * input.element_size() + if input_bytes > self.max_size_bytes: + raise ValueError( + f"input requires {input_bytes} bytes but max_size_bytes={self.max_size_bytes}" + ) + + def _validate_output(self, output: torch.Tensor, input: torch.Tensor) -> None: + if output.device != input.device: + raise ValueError("out must be on the same device as input") + if output.dtype != input.dtype: + raise TypeError("out must have the same dtype as input") + if output.shape != input.shape: + raise ValueError("out must have the same shape as input") + if not output.is_contiguous(): + raise ValueError("out must be contiguous") + + def _validate_sharded_output( + self, + output: torch.Tensor, + input: torch.Tensor, + output_shape: tuple[int, ...], + ) -> None: + if output.device != input.device: + raise ValueError("out must be on the same device as input") + if output.dtype != input.dtype: + raise TypeError("out must have the same dtype as input") + if output.shape != output_shape: + raise ValueError(f"out must have shape {output_shape}, got {tuple(output.shape)}") + if not output.is_contiguous(): + raise ValueError("out must be contiguous") + + def _validate_matching_signature(self, op_name: str, input: torch.Tensor) -> None: + signature = (op_name, tuple(input.shape), str(input.dtype), input.numel()) + signatures: list[tuple[Any, ...] | None] = [None] * self.world_size + dist.all_gather_object(signatures, signature, group=self.group) + if any(peer_signature != signature for peer_signature in signatures): + raise ValueError( + f"all ranks must call {op_name} with matching shapes and dtypes; got {signatures}" + ) + + def _synchronize_ranks(self) -> None: + torch.cuda.synchronize(self.device) + backend = dist.get_backend(self.group) + if backend == dist.Backend.NCCL or str(backend).lower() == "nccl": + dist.barrier(group=self.group, device_ids=[self.device.index]) + else: + dist.barrier(group=self.group) diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index c3d27848..b22c609a 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -112,6 +112,30 @@ def _load_object(path: str) -> Any: }, grad_input_names=("x",), ), + "silu": OperatorSpec( + name="silu", + op_class="elementwise", + gold_path="rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp", + "triton": "rl_engine.kernels.ops.triton.activation.swiglu.TritonSiLUOp", + "cuda": "rl_engine.kernels.ops.cuda.activation.swiglu.SiLUCudaOp", + }, + grad_input_names=("x",), + ), + "swiglu": OperatorSpec( + name="swiglu", + op_class="elementwise", + gold_path="rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp", + "triton": "rl_engine.kernels.ops.triton.activation.swiglu.TritonSwiGLUOp", + "cuda": "rl_engine.kernels.ops.cuda.activation.swiglu.SwiGLUCudaOp", + }, + grad_input_names=("gate", "up"), + ), "batch_invariant_logp": OperatorSpec( name="batch_invariant_logp", op_class="logprob", diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py new file mode 100644 index 00000000..b0bcd7ab --- /dev/null +++ b/rl_engine/kernels/logprob_contract.py @@ -0,0 +1,652 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for TP-aware selected-token log-probability. + +The objects in this module describe a vocab-parallel logprob invocation: + +``selected_logp[t] = logits[t, target[t]] - logsumexp_vocab(logits[t, :])`` + +Under vocab-parallel tensor parallelism the vocabulary-wide ``logsumexp`` +requires cross-rank reduction. This module only *describes* that invocation +(shard ownership, merge semantics, mask/ignore-index metadata); it does not +shard tensors, launch collectives, or implement the ``(max, sumexp)`` merge. +Keeping description and materialization separate lets dispatch reject an +incompatible backend before any numerically different path is launched. + +Context parallelism is a declared non-merge axis: CP partitions tokens, never +the vocabulary, so the logprob reduction spans TP vocab shards only. CP rank +metadata is carried for provenance and must never widen the merge. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + +# Policy keywords accepted by KernelRegistry.get_logprob_op; a backend id must +# never shadow one of these, or it becomes unselectable by id. +RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) +# Backend tiers; determinism is a separate axis (DeterminismScope). +IMPLEMENTATION_KINDS = frozenset({"production", "reference"}) + + +class LogprobContractError(ValueError): + """Raised when logprob metadata does not describe a valid invocation.""" + + +class LogprobRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class LogprobDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class LogprobMerge(str, Enum): + """Merge primitive for per-shard ``(local_max, local_sumexp)`` partials.""" + + MAX_SUMEXP = "max_sumexp" + + +class MergeAxis(str, Enum): + """The only reduction axis of this contract; CP is a non-merge axis.""" + + TP_VOCAB = "tp_vocab" + + +class ReductionOrder(str, Enum): + GLOBAL_VOCAB_SHARD_INDEX = "global_vocab_shard_index" + + +class ReductionTransport(str, Enum): + """Collectives move partial states only; they never reduce numerically.""" + + ALL_GATHER = "all_gather" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +class DeterminismScope(str, Enum): + """Strength of the reduction's determinism guarantee. + + ``fixed_topology``: bitwise-reproducible for one fixed TP degree; results + at different TP degrees are compared against the #108 tolerance table. + + ``cross_tp_bitwise``: additionally bitwise-equal across TP degrees. This + requires the entire reduction to follow a global tile-level structure that + is independent of TP partitioning: a fixed tile decomposition of the + vocabulary plus a fixed merge order and rescaling tree over those tiles, + identical at every TP degree, so the TP degree only selects which rank + computes which tiles and never changes the floating-point grouping. + Fixed shard-order merging alone is not sufficient, because shard + boundaries would still group the combines differently across degrees. + """ + + FIXED_TOPOLOGY = "fixed_topology" + CROSS_TP_BITWISE = "cross_tp_bitwise" + + +class MaskMode(str, Enum): + """How a backend consumes inactive-token information. + + The contract permits inactive targets that do not hold ``ignore_index``, + so an ``ignore_index``-only backend cannot serve a contract with inactive + tokens. + """ + + EXPLICIT_ACTIVE_MASK = "explicit_active_mask" + IGNORE_INDEX = "ignore_index" + + +class TPPlacement(str, Enum): + REPLICATED = "replicated" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise LogprobContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LogprobContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise LogprobContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _plain_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise LogprobContractError(f"{field} must be an integer; got {value!r}") + return value + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical vocab-parallel TP ownership for one logprob invocation. + + ``vocab_shard_bounds`` lists every TP rank's half-open ``[start, end)`` + vocab range, indexed by rank; the full table is required on every rank and + must form a contiguous ``[0, padded_vocab_size)`` partition. + ``padded_vocab_size`` is the shard-covered (weight) vocabulary, + ``real_vocab_size`` the tokenizer vocabulary; padding columns occupy + ``[real_vocab_size, padded_vocab_size)``. + """ + + tp_rank: int + tp_world_size: int + vocab_shard_bounds: tuple[tuple[int, int], ...] + real_vocab_size: int + padded_vocab_size: int + cp_rank: int = 0 + cp_world_size: int = 1 + + def __post_init__(self) -> None: + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + if tp_rank >= tp_world_size: + raise LogprobContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + if cp_rank >= cp_world_size: + raise LogprobContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" + ) + + real_vocab_size = _positive_int(self.real_vocab_size, "real_vocab_size") + padded_vocab_size = _positive_int(self.padded_vocab_size, "padded_vocab_size") + if padded_vocab_size < real_vocab_size: + raise LogprobContractError( + f"padded_vocab_size={padded_vocab_size} must not be smaller than " + f"real_vocab_size={real_vocab_size}" + ) + + try: + bounds = tuple((pair[0], pair[1]) for pair in self.vocab_shard_bounds) + except (TypeError, IndexError) as exc: + raise LogprobContractError( + "vocab_shard_bounds must be an iterable of (start, end) integer pairs" + ) from exc + if len(bounds) != tp_world_size: + raise LogprobContractError( + "vocab_shard_bounds must declare exactly one (start, end) pair per TP rank; " + f"got {len(bounds)} pairs for tp_world_size={tp_world_size}" + ) + expected_start = 0 + for rank, (start, end) in enumerate(bounds): + start = _plain_int(start, f"vocab_shard_bounds[{rank}][0]") + end = _plain_int(end, f"vocab_shard_bounds[{rank}][1]") + if end <= start: + raise LogprobContractError( + f"vocab_shard_bounds[{rank}] must satisfy end > start; got [{start}, {end})" + ) + if start != expected_start: + raise LogprobContractError( + "vocab_shard_bounds must form a contiguous [0, padded_vocab_size) " + f"partition in TP-rank order; rank {rank} starts at {start}, " + f"expected {expected_start}" + ) + expected_start = end + if expected_start != padded_vocab_size: + raise LogprobContractError( + "vocab_shard_bounds must cover padded_vocab_size exactly; " + f"covered {expected_start}, declared {padded_vocab_size}" + ) + object.__setattr__(self, "vocab_shard_bounds", bounds) + + @property + def local_vocab_start(self) -> int: + return self.vocab_shard_bounds[self.tp_rank][0] + + @property + def local_vocab_end(self) -> int: + return self.vocab_shard_bounds[self.tp_rank][1] + + @property + def local_vocab_size(self) -> int: + start, end = self.vocab_shard_bounds[self.tp_rank] + return end - start + + def owner_rank(self, token_id: int) -> int: + """Return the unique TP rank owning ``token_id``; error outside real vocab.""" + + token_id = _plain_int(token_id, "token_id") + if token_id < 0 or token_id >= self.real_vocab_size: + raise LogprobContractError( + f"token_id={token_id} is outside the real vocabulary " + f"[0, {self.real_vocab_size}); mask it as inactive instead" + ) + for rank, (start, end) in enumerate(self.vocab_shard_bounds): + if start <= token_id < end: + return rank + raise LogprobContractError( + f"token_id={token_id} is not covered by any declared vocab shard" + ) + + +@dataclass(frozen=True) +class MaskSpec: + """Active-token mask and ignore index for one logprob invocation. + + Inactive tokens are excluded from drift aggregates and from the + single-owner target gather; their targets may legally hold ``ignore_index``. + """ + + num_tokens: int + active_mask: tuple[bool, ...] + ignore_index: int = -100 + _active_token_count: int = field(init=False, repr=False, compare=False) + _active_mask_sha256: str = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + num_tokens = _positive_int(self.num_tokens, "num_tokens") + _plain_int(self.ignore_index, "ignore_index") + try: + active_mask = tuple(self.active_mask) + except TypeError as exc: + raise LogprobContractError("active_mask must be an iterable of booleans") from exc + for index, value in enumerate(active_mask): + if not isinstance(value, bool): + raise LogprobContractError(f"active_mask[{index}] must be a bool; got {value!r}") + if len(active_mask) != num_tokens: + raise LogprobContractError( + "active_mask must contain exactly one entry per token; " + f"got {len(active_mask)} entries for num_tokens={num_tokens}" + ) + object.__setattr__(self, "active_mask", active_mask) + object.__setattr__(self, "_active_token_count", sum(active_mask)) + object.__setattr__( + self, "_active_mask_sha256", hashlib.sha256(bytes(active_mask)).hexdigest() + ) + + @property + def active_token_count(self) -> int: + return self._active_token_count + + @property + def active_mask_sha256(self) -> str: + """Compact mask identity for provenance and cross-rank agreement.""" + return self._active_mask_sha256 + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic TP-vocab ``(max, sumexp)`` merge semantics. + + Every rank first masks local columns whose global id lies in + ``[real_vocab_size, padded_vocab_size)`` to ``-inf`` (padding never + contributes to the logsumexp), then computes ``m_l = max(local_logits)`` + and ``s_l = sum(exp(local_logits - m_l))`` in fp32. Partials travel by + all-gather -- collectives are transport only, never a numerical + reduction -- and every rank merges in fixed global vocab-shard index + order:: + + M = max_l(m_l) + S = sum_l(s_l * exp(m_l - M)) + LSE = M + log(S) + selected_logp = target_logit - LSE + + The selected target logit comes from a masked single-owner gather; + downcast happens only at the final write. The identity partial for a + padding-only shard, or a row whose local columns are all ``-inf`` after + masking, is ``(m_l, s_l) = (-inf, 0)``: a partial with ``s_l = 0`` + contributes nothing to the merge regardless of its ``m_l``, and + implementations must use this identity directly rather than evaluate + ``exp(-inf - (-inf))``, which would poison the merge with NaN. Averaging + per-rank logsumexp values, or letting a collective reduce numerically, is + never conformant at either determinism scope. + """ + + merge: LogprobMerge = LogprobMerge.MAX_SUMEXP + merge_axis: MergeAxis = MergeAxis.TP_VOCAB + acc_dtype: LogprobDType = LogprobDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_VOCAB_SHARD_INDEX + transport: ReductionTransport = ReductionTransport.ALL_GATHER + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + determinism_scope: DeterminismScope = DeterminismScope.CROSS_TP_BITWISE + + def __post_init__(self) -> None: + object.__setattr__( + self, + "determinism_scope", + _enum_value(DeterminismScope, self.determinism_scope, "determinism_scope"), + ) + object.__setattr__(self, "merge", _enum_value(LogprobMerge, self.merge, "merge")) + object.__setattr__( + self, "merge_axis", _enum_value(MergeAxis, self.merge_axis, "merge_axis") + ) + object.__setattr__( + self, "acc_dtype", _enum_value(LogprobDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "transport", _enum_value(ReductionTransport, self.transport, "transport") + ) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not LogprobDType.FP32: + raise LogprobContractError( + f"TP logprob accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class LogprobOutputSpec: + """Output surface every conforming backend must produce: fp32 selected + logprob and fp32 vocab-domain LSE, replicated across the TP group.""" + + selected_logp_dtype: LogprobDType = LogprobDType.FP32 + lse_dtype: LogprobDType = LogprobDType.FP32 + tp_placement: TPPlacement = TPPlacement.REPLICATED + + def __post_init__(self) -> None: + object.__setattr__( + self, + "selected_logp_dtype", + _enum_value(LogprobDType, self.selected_logp_dtype, "selected_logp_dtype"), + ) + object.__setattr__( + self, "lse_dtype", _enum_value(LogprobDType, self.lse_dtype, "lse_dtype") + ) + object.__setattr__( + self, "tp_placement", _enum_value(TPPlacement, self.tp_placement, "tp_placement") + ) + if self.selected_logp_dtype is not LogprobDType.FP32: + raise LogprobContractError( + f"selected logprob output must be fp32; got {self.selected_logp_dtype.value}" + ) + if self.lse_dtype is not LogprobDType.FP32: + raise LogprobContractError(f"vocab LSE output must be fp32; got {self.lse_dtype.value}") + + +@dataclass(frozen=True) +class LogprobContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: LogprobRole + dtype: LogprobDType + mask: MaskSpec + sharding: ShardingSpec + reduction: ReductionSpec + output: LogprobOutputSpec = field(default_factory=LogprobOutputSpec) + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(LogprobRole, self.role, "role")) + object.__setattr__(self, "dtype", _enum_value(LogprobDType, self.dtype, "dtype")) + if not isinstance(self.mask, MaskSpec): + raise LogprobContractError("mask must be a MaskSpec") + if not isinstance(self.sharding, ShardingSpec): + raise LogprobContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise LogprobContractError("reduction must be a ReductionSpec") + if not isinstance(self.output, LogprobOutputSpec): + raise LogprobContractError("output must be a LogprobOutputSpec") + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise LogprobContractError( + "export_lse must be True for the WS2 vocab-domain LSE drift contract" + ) + if 0 <= self.mask.ignore_index < self.sharding.real_vocab_size: + raise LogprobContractError( + f"ignore_index={self.mask.ignore_index} must not collide with the real " + f"vocabulary [0, {self.sharding.real_vocab_size})" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "vocab_shard_bounds": [list(pair) for pair in self.sharding.vocab_shard_bounds], + "real_vocab_size": self.sharding.real_vocab_size, + "padded_vocab_size": self.sharding.padded_vocab_size, + "local_vocab_start": self.sharding.local_vocab_start, + "local_vocab_end": self.sharding.local_vocab_end, + } + reduction = { + "merge": self.reduction.merge.value, + "merge_axis": self.reduction.merge_axis.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "transport": self.reduction.transport.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + "determinism_scope": self.reduction.determinism_scope.value, + "cp_is_merge_axis": False, + } + # The digest stands in for the raw per-token mask, which would + # dominate the provenance size. + mask = { + "num_tokens": self.mask.num_tokens, + "active_token_count": self.mask.active_token_count, + "active_mask_sha256": self.mask.active_mask_sha256, + "ignore_index": self.mask.ignore_index, + } + output = { + "selected_logp_dtype": self.output.selected_logp_dtype.value, + "lse_dtype": self.output.lse_dtype.value, + "tp_placement": self.output.tp_placement.value, + } + return { + "semantic_operator": "selected_token_logprob", + "role": self.role.value, + "dtype": self.dtype.value, + "export_lse": self.export_lse, + "lse_domain": "vocab", + "mask": mask, + "sharding": sharding, + "reduction": reduction, + "output": output, + } + + def cross_rank_fingerprint(self) -> str: + """Rank-independent identity for preflight agreement across ranks. + + Excludes ``tp_rank``/``cp_rank`` (and their derived local bounds) so + every rank of one logical invocation computes the same value. + All-gathering this fingerprint together with the resolved backend id + and aborting on mismatch is the documented preflight for distributed + dispatch; ``requested_backend="auto"`` is not distributed-safe + without it. + """ + + payload = self.to_dict() + payload["sharding"] = { + key: value + for key, value in payload["sharding"].items() + if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} + } + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +@dataclass(frozen=True) +class LogprobBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[LogprobRole] + dtypes: frozenset[LogprobDType] + tp_world_sizes: tuple[int, ...] | None = None + cp_world_sizes: tuple[int, ...] | None = None + supports_vocab_padding: bool = False + mask_modes: frozenset[MaskMode] = frozenset() + exports_vocab_lse: bool = False + determinism_scopes: frozenset[DeterminismScope] = frozenset() + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise LogprobContractError("backend_id must be a non-empty string") + if self.backend_id.strip().lower() in RESERVED_DISPATCH_POLICIES: + raise LogprobContractError( + f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" + ) + object.__setattr__(self, "backend_id", self.backend_id.strip()) + try: + roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) + dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) + except TypeError as exc: + raise LogprobContractError("roles and dtypes must be iterables of enum values") from exc + if not roles or not dtypes: + raise LogprobContractError("backend roles and dtypes must not be empty") + tp_world_sizes = self._validated_world_sizes(self.tp_world_sizes, "tp_world_sizes") + cp_world_sizes = self._validated_world_sizes(self.cp_world_sizes, "cp_world_sizes") + try: + mask_modes = frozenset( + _enum_value(MaskMode, value, "mask_modes") for value in self.mask_modes + ) + determinism_scopes = frozenset( + _enum_value(DeterminismScope, value, "determinism_scopes") + for value in self.determinism_scopes + ) + except TypeError as exc: + raise LogprobContractError( + "mask_modes and determinism_scopes must be iterables of enum values" + ) from exc + for flag_name in ("supports_vocab_padding", "exports_vocab_lse"): + if not isinstance(getattr(self, flag_name), bool): + raise LogprobContractError(f"{flag_name} must be a bool") + if self.implementation_kind not in IMPLEMENTATION_KINDS: + raise LogprobContractError( + f"implementation_kind must be one of: {', '.join(sorted(IMPLEMENTATION_KINDS))}" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "mask_modes", mask_modes) + object.__setattr__(self, "determinism_scopes", determinism_scopes) + + @staticmethod + def _validated_world_sizes( + values: tuple[int, ...] | None, field: str + ) -> tuple[int, ...] | None: + if values is None: + return None + try: + sizes = tuple(values) + except TypeError as exc: + raise LogprobContractError(f"{field} must be an iterable of integers") from exc + if not sizes: + raise LogprobContractError(f"{field} must not be empty; use None for unrestricted") + for value in sizes: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LogprobContractError(f"{field} must contain positive values; got {value!r}") + if len(set(sizes)) != len(sizes): + raise LogprobContractError(f"{field} must not contain duplicates") + return sizes + + def incompatibilities(self, contract: LogprobContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if self.cp_world_sizes is not None and cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if ( + contract.sharding.padded_vocab_size != contract.sharding.real_vocab_size + and not self.supports_vocab_padding + ): + reasons.append("padded-vs-real vocab masking is unsupported") + if ( + contract.mask.active_token_count != contract.mask.num_tokens + and MaskMode.EXPLICIT_ACTIVE_MASK not in self.mask_modes + ): + # Inactive targets need not hold ignore_index (see MaskMode). + reasons.append("explicit active-token masking is unsupported") + if contract.export_lse and not self.exports_vocab_lse: + reasons.append("vocab-domain LSE export is unsupported") + if contract.reduction.determinism_scope not in self.determinism_scopes: + reasons.append( + f"determinism_scope={contract.reduction.determinism_scope.value} is unsupported" + ) + return tuple(reasons) + + def supports(self, contract: LogprobContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes) if self.cp_world_sizes else None, + "supports_vocab_padding": self.supports_vocab_padding, + "mask_modes": sorted(mode.value for mode in self.mask_modes), + "exports_vocab_lse": self.exports_vocab_lse, + "determinism_scopes": sorted(scope.value for scope in self.determinism_scopes), + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class LogprobDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: LogprobBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "IMPLEMENTATION_KINDS", + "RESERVED_DISPATCH_POLICIES", + "DeterminismScope", + "DowncastPoint", + "LogprobBackendCapability", + "LogprobContract", + "LogprobContractError", + "LogprobDType", + "LogprobDispatchResult", + "LogprobMerge", + "LogprobOutputSpec", + "LogprobRole", + "MaskMode", + "MaskSpec", + "MergeAxis", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "ReductionTransport", + "ShardingSpec", + "TPPlacement", +] diff --git a/rl_engine/kernels/ops/cuda/__init__.py b/rl_engine/kernels/ops/cuda/__init__.py index 5f1ae8f8..94a16716 100644 --- a/rl_engine/kernels/ops/cuda/__init__.py +++ b/rl_engine/kernels/ops/cuda/__init__.py @@ -1,2 +1,2 @@ -# append matmul to the existing imports -from . import attention, loss, matmul, norm # noqa: F401 +# append matmul / activation to the existing imports +from . import activation, attention, loss, matmul, norm # noqa: F401 diff --git a/rl_engine/kernels/ops/cuda/activation/__init__.py b/rl_engine/kernels/ops/cuda/activation/__init__.py new file mode 100644 index 00000000..214a0007 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/activation/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .swiglu import SiLUCudaOp, SwiGLUCudaOp + +__all__ = ["SiLUCudaOp", "SwiGLUCudaOp"] diff --git a/rl_engine/kernels/ops/cuda/activation/swiglu.py b/rl_engine/kernels/ops/cuda/activation/swiglu.py new file mode 100644 index 00000000..bba3d65e --- /dev/null +++ b/rl_engine/kernels/ops/cuda/activation/swiglu.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""CUDA SiLU / SwiGLU ops matching NativeSiLUOp / NativeSwiGLUOp (WS1 ground truth). + +Math is performed in fp32 inside the CUDA kernels and rounded back to the input +dtype on store — the same dual-path contract as the PyTorch references: + + silu(x) = x * sigmoid(x) + swiglu(g, u) = silu(g) * u + +Element-wise and row-independent, so Axis-A batch invariance holds bitwise. +""" + +from __future__ import annotations + +import torch +from torch import Tensor + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.utils.logger import logger + +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) + + +def _validate_dtype(x: Tensor, name: str) -> None: + if x.dtype not in _SUPPORTED_DTYPES: + raise TypeError(f"{name} must have dtype fp16, bf16, or fp32, got {x.dtype}.") + + +def _require_cuda_activation() -> None: + if not _EXT_AVAILABLE or _C is None: + raise RuntimeError("CUDA activation kernels require the compiled rl_engine._C extension.") + if not hasattr(_C, "silu_forward") or not hasattr(_C, "swiglu_forward"): + raise RuntimeError( + "CUDA activation symbols (silu_forward / swiglu_forward) are not compiled " + "into _C. Rebuild the extension with csrc/cuda/activation.cu." + ) + + +class _SiLUCudaFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor) -> Tensor: + x_c = x.contiguous() + y = _C.silu_forward(x_c) + ctx.save_for_backward(x_c) + return y + + @staticmethod + def backward(ctx, grad_out: Tensor): + (x,) = ctx.saved_tensors + dx = None + if ctx.needs_input_grad[0]: + dx = _C.silu_backward(grad_out.contiguous(), x) + return dx + + +class _SwiGLUCudaFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, gate: Tensor, up: Tensor) -> Tensor: + gate_c = gate.contiguous() + up_c = up.contiguous() + y = _C.swiglu_forward(gate_c, up_c) + ctx.save_for_backward(gate_c, up_c) + return y + + @staticmethod + def backward(ctx, grad_out: Tensor): + gate, up = ctx.saved_tensors + d_gate = d_up = None + if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]: + grads = _C.swiglu_backward(grad_out.contiguous(), gate, up) + if ctx.needs_input_grad[0]: + d_gate = grads[0] + if ctx.needs_input_grad[1]: + d_up = grads[1] + return d_gate, d_up + + +class SiLUCudaOp: + """CUDA SiLU: ``out = x * sigmoid(x)``, math in fp32.""" + + op_class = "elementwise" + + def __init__(self) -> None: + _require_cuda_activation() + logger.info("Successfully linked to precompiled _C.silu_forward kernel.") + + def __call__(self, x: Tensor) -> Tensor: + return self.forward(x) + + def forward(self, x: Tensor) -> Tensor: + if x.device.type != "cuda": + raise RuntimeError(f"SiLUCudaOp requires a CUDA tensor, got device '{x.device}'.") + _validate_dtype(x, "x") + return _SiLUCudaFunction.apply(x) + + def forward_fp32(self, x: Tensor) -> Tensor: + """Ground-truth path: force fp32 input so the kernel stores fp32 output.""" + if x.device.type != "cuda": + raise RuntimeError(f"SiLUCudaOp requires a CUDA tensor, got device '{x.device}'.") + _validate_dtype(x, "x") + return _SiLUCudaFunction.apply(x.float()) + + +class SwiGLUCudaOp: + """CUDA SwiGLU: ``out = silu(gate) * up``, math in fp32.""" + + op_class = "elementwise" + + def __init__(self) -> None: + _require_cuda_activation() + logger.info("Successfully linked to precompiled _C.swiglu_forward kernel.") + + def __call__(self, gate: Tensor, up: Tensor) -> Tensor: + return self.forward(gate, up) + + def forward(self, gate: Tensor, up: Tensor) -> Tensor: + if gate.device.type != "cuda" or up.device.type != "cuda": + raise RuntimeError( + f"SwiGLUCudaOp requires CUDA tensors, got gate='{gate.device}', up='{up.device}'." + ) + if gate.device != up.device: + raise RuntimeError( + f"gate and up must be on the same CUDA device, got " + f"'{gate.device}' and '{up.device}'." + ) + if gate.shape != up.shape: + raise ValueError( + f"gate and up must share shape, got tuple(gate.shape)=" + f"{tuple(gate.shape)} vs tuple(up.shape)={tuple(up.shape)}" + ) + _validate_dtype(gate, "gate") + _validate_dtype(up, "up") + if gate.dtype != up.dtype: + raise TypeError(f"gate and up must share dtype, got {gate.dtype} and {up.dtype}.") + return _SwiGLUCudaFunction.apply(gate, up) + + def forward_fp32(self, gate: Tensor, up: Tensor) -> Tensor: + if gate.device.type != "cuda" or up.device.type != "cuda": + raise RuntimeError( + f"SwiGLUCudaOp requires CUDA tensors, got gate='{gate.device}', up='{up.device}'." + ) + if gate.device != up.device: + raise RuntimeError( + f"gate and up must be on the same CUDA device, got " + f"'{gate.device}' and '{up.device}'." + ) + if gate.shape != up.shape: + raise ValueError( + f"gate and up must share shape, got tuple(gate.shape)=" + f"{tuple(gate.shape)} vs tuple(up.shape)={tuple(up.shape)}" + ) + _validate_dtype(gate, "gate") + _validate_dtype(up, "up") + if gate.dtype != up.dtype: + raise TypeError(f"gate and up must share dtype, got {gate.dtype} and {up.dtype}.") + return _SwiGLUCudaFunction.apply(gate.float(), up.float()) diff --git a/rl_engine/kernels/ops/pytorch/activation/swiglu.py b/rl_engine/kernels/ops/pytorch/activation/swiglu.py index 70cea940..c3d8fde3 100644 --- a/rl_engine/kernels/ops/pytorch/activation/swiglu.py +++ b/rl_engine/kernels/ops/pytorch/activation/swiglu.py @@ -6,6 +6,13 @@ import torch import torch.nn as nn +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) + + +def _validate_dtype(x: torch.Tensor, name: str) -> None: + if x.dtype not in _SUPPORTED_DTYPES: + raise TypeError(f"{name} must have dtype fp16, bf16, or fp32, got {x.dtype}.") + class NativeSiLUOp(nn.Module): """ @@ -36,6 +43,7 @@ def forward_fp32(self, x: torch.Tensor) -> torch.Tensor: # ------------------------------------------------------------------ # @staticmethod def _silu(x: torch.Tensor, *, output_dtype: torch.dtype) -> torch.Tensor: + _validate_dtype(x, "x") x_f = x.float() out = x_f * torch.sigmoid(x_f) return out.to(output_dtype) @@ -81,6 +89,15 @@ def _swiglu( f"gate and up must share shape, got tuple(gate.shape)=" f"{tuple(gate.shape)} vs tuple(up.shape)={tuple(up.shape)}" ) + if gate.device != up.device: + raise RuntimeError( + f"gate and up must be on the same device, got '{gate.device}' and '{up.device}'." + ) + _validate_dtype(gate, "gate") + _validate_dtype(up, "up") + if gate.dtype != up.dtype: + raise TypeError(f"gate and up must share dtype, got {gate.dtype} and {up.dtype}.") + gate_f = gate.float() out = gate_f * torch.sigmoid(gate_f) * up.float() return out.to(output_dtype) diff --git a/rl_engine/kernels/ops/pytorch/ffn/__init__.py b/rl_engine/kernels/ops/pytorch/ffn/__init__.py new file mode 100644 index 00000000..19a54d23 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/ffn/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .ffn import BACKEND_ID, Qwen3FFNOp, qwen3_ffn + +__all__ = ["BACKEND_ID", "Qwen3FFNOp", "qwen3_ffn"] diff --git a/rl_engine/kernels/ops/pytorch/ffn/ffn.py b/rl_engine/kernels/ops/pytorch/ffn/ffn.py new file mode 100644 index 00000000..c2507a44 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/ffn/ffn.py @@ -0,0 +1,477 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Bias-free gated FFN assembled from deterministic CUDA kernels.""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import Tensor + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + +QWEN3_8B_HIDDEN_SIZE = 4096 +QWEN3_8B_INTERMEDIATE_SIZE = 12288 +BACKEND_ID = "rlkernel.ffn.qwen3.deterministic.v1" + +_DET_GEMM_SYMBOLS = ( + "det_gemm_fwd", + "det_gemm_db", +) +_SWIGLU_SYMBOLS = ( + "swiglu_forward", + "swiglu_backward", +) +_REQUIRED_SYMBOLS = _DET_GEMM_SYMBOLS + _SWIGLU_SYMBOLS +_COLLECTIVE_MIN_CAPACITY_BYTES = 64 * 1024 * 1024 +_COLLECTIVES: dict[tuple[int, int, int, int], Any] = {} + + +def _require_ffn_kernels(*, disable_split_k: bool) -> None: + required = _REQUIRED_SYMBOLS if disable_split_k else _SWIGLU_SYMBOLS + missing = [name for name in required if not hasattr(_C, name)] + if not _EXT_AVAILABLE or _C is None or missing: + suffix = f" Missing symbols: {', '.join(missing)}." if missing else "" + needed = ( + "compiled deterministic GEMM and SwiGLU CUDA kernels" + if disable_split_k + else "compiled SwiGLU CUDA kernels" + ) + raise RuntimeError(f"qwen3_ffn requires the {needed}.{suffix}") + + +def _gemm_fwd(a: Tensor, b: Tensor, *, disable_split_k: bool) -> Tensor: + if disable_split_k: + return _C.det_gemm_fwd(a, b) + # cuBLASLt / CUTLASS: may use split-K. Detach so Autograd.Function owns backward. + with torch.no_grad(): + return torch.matmul(a, b) + + +def _gemm_db(a: Tensor, grad_output: Tensor, *, disable_split_k: bool) -> Tensor: + if disable_split_k: + return _C.det_gemm_db(a, grad_output) + with torch.no_grad(): + return torch.matmul(a.t().contiguous(), grad_output) + + +def _require_parallel_group(group: Any, name: str): + if group is None: + return None + + import torch.distributed as dist + + if not dist.is_available(): + raise RuntimeError(f"{name}-parallel FFN requires torch.distributed.") + if not dist.is_initialized(): + raise RuntimeError(f"{name}-parallel FFN requires an initialized process group.") + if dist.get_world_size(group=group) <= 1: + raise ValueError(f"{name}_group must contain at least two ranks.") + return dist + + +def _validate_ffn_inputs( + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, +) -> None: + tensors = { + "rmsnorm_output": rmsnorm_output, + "gate_weight": gate_weight, + "up_weight": up_weight, + "down_weight": down_weight, + } + for name, tensor in tensors.items(): + if not isinstance(tensor, Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor)!r}.") + + if rmsnorm_output.dim() < 1: + raise ValueError("rmsnorm_output must have at least one dimension.") + if rmsnorm_output.numel() == 0: + raise ValueError("rmsnorm_output must contain at least one token.") + for name, weight in ( + ("gate_weight", gate_weight), + ("up_weight", up_weight), + ("down_weight", down_weight), + ): + if weight.dim() != 2: + raise ValueError(f"{name} must be 2-D, got shape {tuple(weight.shape)}.") + + hidden_size = rmsnorm_output.size(-1) + intermediate_size = gate_weight.size(0) + expected_shapes = { + "gate_weight": (intermediate_size, hidden_size), + "up_weight": (intermediate_size, hidden_size), + "down_weight": (hidden_size, intermediate_size), + } + for name, expected in expected_shapes.items(): + actual = tuple(tensors[name].shape) + if actual != expected: + raise ValueError(f"{name} must have shape {expected}, got {actual}.") + + for name, tensor in tensors.items(): + if tensor.dtype != torch.bfloat16: + raise TypeError(f"{name} must have dtype bfloat16, got {tensor.dtype}.") + if not tensor.is_cuda: + raise RuntimeError(f"{name} must be on a CUDA device, got '{tensor.device}'.") + if tensor.device != rmsnorm_output.device: + raise RuntimeError( + f"all FFN inputs must be on {rmsnorm_output.device}, " + f"got {name} on {tensor.device}." + ) + + +def _collective_for_group(group: Any, *, min_size_bytes: int): + if group is None: + return None + + import torch.distributed as dist + + from rl_engine.distributed import DeterministicCollective + + rank = dist.get_rank(group=group) + world_size = dist.get_world_size(group=group) + device_index = torch.cuda.current_device() + key = (id(group), rank, world_size, device_index) + cached = _COLLECTIVES.get(key) + if cached is not None and cached.max_size_bytes >= min_size_bytes: + return cached + if cached is not None: + cached.close() + + collective = DeterministicCollective( + group=group, + max_size_bytes=max(_COLLECTIVE_MIN_CAPACITY_BYTES, min_size_bytes), + ) + _COLLECTIVES[key] = collective + return collective + + +def _all_gather_tokens(tensor: Tensor, collective: Any) -> Tensor: + return collective.all_gather(tensor.contiguous()) + + +def _reduce_scatter_tokens(tensor: Tensor, collective: Any) -> Tensor: + world_size = collective.world_size + if tensor.size(0) % world_size != 0: + raise ValueError( + "the gathered token count must be divisible by the tensor-parallel " + f"world size, got {tensor.size(0)} and {world_size}." + ) + return collective.reduce_scatter(tensor.contiguous()) + + +def _all_reduce_inplace(tensor: Tensor, collective: Any) -> Tensor: + return collective.all_reduce(tensor, out=tensor) + + +class _DeterministicFFNFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + tp_group: Any, + cp_group: Any, + sequence_parallel: bool, + disable_split_k: bool, + ) -> Tensor: + tp_dist = _require_parallel_group(tp_group, "tensor") + _require_parallel_group(cp_group, "context") + if sequence_parallel and tp_dist is None: + raise ValueError("sequence_parallel requires a tensor-parallel group.") + + input_shape = rmsnorm_output.shape + rmsnorm_output_2d = rmsnorm_output.reshape(-1, input_shape[-1]).contiguous() + tp_world = tp_dist.get_world_size(group=tp_group) if tp_dist is not None else 1 + gemm_tokens = rmsnorm_output_2d.size(0) * (tp_world if sequence_parallel else 1) + element_size = rmsnorm_output_2d.element_size() + min_size_bytes = max( + gemm_tokens * rmsnorm_output_2d.size(1) * element_size, + gemm_tokens * gate_weight.size(0) * element_size, + gate_weight.numel() * element_size, + up_weight.numel() * element_size, + down_weight.numel() * element_size, + ) + # Create TP before CP so every rank follows the same group order. + tp_collective = _collective_for_group(tp_group, min_size_bytes=min_size_bytes) + cp_collective = _collective_for_group(cp_group, min_size_bytes=min_size_bytes) + + if sequence_parallel: + rmsnorm_output_2d = _all_gather_tokens(rmsnorm_output_2d, tp_collective) + + # The model stores projection weights as [out, in]; GEMM consumes [K, N]. + gate = _gemm_fwd( + rmsnorm_output_2d, + gate_weight.t().contiguous(), + disable_split_k=disable_split_k, + ) + up = _gemm_fwd( + rmsnorm_output_2d, + up_weight.t().contiguous(), + disable_split_k=disable_split_k, + ) + activated = _C.swiglu_forward(gate, up) + output = _gemm_fwd(activated, down_weight.t().contiguous(), disable_split_k=disable_split_k) + + if sequence_parallel: + output = _reduce_scatter_tokens(output, tp_collective) + elif tp_collective is not None: + output = _all_reduce_inplace(output, tp_collective) + + ctx.save_for_backward( + rmsnorm_output_2d, + gate, + up, + activated, + gate_weight, + up_weight, + down_weight, + ) + ctx.input_shape = input_shape + ctx.tp_collective = tp_collective + ctx.cp_collective = cp_collective + ctx.sequence_parallel = sequence_parallel + ctx.disable_split_k = disable_split_k + return output.reshape(*input_shape[:-1], output.size(-1)) + + @staticmethod + def backward(ctx, grad_output: Tensor): + ( + rmsnorm_output, + gate, + up, + activated, + gate_weight, + up_weight, + down_weight, + ) = ctx.saved_tensors + tp_collective = ctx.tp_collective + cp_collective = ctx.cp_collective + disable_split_k = ctx.disable_split_k + grad_output = grad_output.reshape(-1, grad_output.size(-1)).contiguous() + if ctx.sequence_parallel: + grad_output = _all_gather_tokens(grad_output, tp_collective) + + # Down weight gradients must see every CP token so gemm_db's K-tree + # matches CP=1. Local dW + AllReduce is a different parenthesization + # whenever T is not a complete mid-split tree of 32-wide leaves. + if cp_collective is not None: + activated_full = _all_gather_tokens(activated, cp_collective) + grad_output_full = _all_gather_tokens(grad_output, cp_collective) + grad_down_weight = ( + _gemm_db(activated_full, grad_output_full, disable_split_k=disable_split_k) + .t() + .contiguous() + ) + else: + grad_down_weight = ( + _gemm_db(activated, grad_output, disable_split_k=disable_split_k).t().contiguous() + ) + + # Down input-gradient shards concatenate across TP; no TP reduction. + grad_activated = _gemm_fwd(grad_output, down_weight, disable_split_k=disable_split_k) + grad_gate, grad_up = _C.swiglu_backward(grad_activated, gate, up) + + if cp_collective is not None: + rmsnorm_full = _all_gather_tokens(rmsnorm_output, cp_collective) + grad_gate_full = _all_gather_tokens(grad_gate, cp_collective) + grad_up_full = _all_gather_tokens(grad_up, cp_collective) + grad_gate_weight = ( + _gemm_db(rmsnorm_full, grad_gate_full, disable_split_k=disable_split_k) + .t() + .contiguous() + ) + grad_up_weight = ( + _gemm_db(rmsnorm_full, grad_up_full, disable_split_k=disable_split_k) + .t() + .contiguous() + ) + else: + grad_gate_weight = ( + _gemm_db(rmsnorm_output, grad_gate, disable_split_k=disable_split_k) + .t() + .contiguous() + ) + grad_up_weight = ( + _gemm_db(rmsnorm_output, grad_up, disable_split_k=disable_split_k).t().contiguous() + ) + + # Gate/Up input gradients reduce across TP, then add locally. + grad_rmsnorm_from_gate = _gemm_fwd(grad_gate, gate_weight, disable_split_k=disable_split_k) + if ctx.sequence_parallel: + grad_rmsnorm_from_gate = _reduce_scatter_tokens( + grad_rmsnorm_from_gate, + tp_collective, + ) + elif tp_collective is not None: + grad_rmsnorm_from_gate = _all_reduce_inplace( + grad_rmsnorm_from_gate, + tp_collective, + ) + + grad_rmsnorm_from_up = _gemm_fwd(grad_up, up_weight, disable_split_k=disable_split_k) + if ctx.sequence_parallel: + grad_rmsnorm_from_up = _reduce_scatter_tokens( + grad_rmsnorm_from_up, + tp_collective, + ) + elif tp_collective is not None: + grad_rmsnorm_from_up = _all_reduce_inplace( + grad_rmsnorm_from_up, + tp_collective, + ) + + grad_rmsnorm_output = grad_rmsnorm_from_gate.add_(grad_rmsnorm_from_up) + return ( + grad_rmsnorm_output.reshape(ctx.input_shape), + grad_gate_weight, + grad_up_weight, + grad_down_weight, + None, + None, + None, + None, + ) + + +def qwen3_ffn( + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + tp_group: Any = None, + cp_group: Any = None, + sequence_parallel: bool = False, + deterministic: bool | None = None, + disable_split_k: bool | None = None, +) -> Tensor: + """Apply a bias-free SiLU-gated FFN with deterministic backward kernels. + + Args: + rmsnorm_output: RMSNorm output, shape ``[..., H]``. + gate_weight: Gate projection weight in ``[out, in]`` layout, shape + ``[I_local, H]``. + up_weight: Up projection weight in ``[out, in]`` layout, shape + ``[I_local, H]``. + down_weight: Down projection weight in ``[out, in]`` layout, shape + ``[H, I_local]``. + tp_group: Optional tensor-parallel process group. Gate and Up are + column-parallel; Down is row-parallel. Reductions use the + deterministic fixed-tree collectives rather than NCCL. + cp_group: Optional context-parallel process group. Each rank owns + different token rows and the same local weight shards. Weight + gradients AllGather tokens along CP and run the full-token + ``det_gemm_db`` so they match CP=1 bitwise. + sequence_parallel: Whether ``rmsnorm_output`` and the returned output + are sharded on the flattened token dimension across ``tp_group``. + Token gather/scatter use the deterministic AllGather and + ReduceScatter. + deterministic: Select the RL-Kernel fixed-reduction GEMM when True + (default), or the production ``torch.matmul`` GEMM when False. + disable_split_k: Compatibility alias for ``deterministic``. New code + should use ``deterministic`` because Split-K is only one possible + implementation detail of the production GEMM. + + Returns: + FFN output with shape ``[..., H]``. + """ + if not isinstance(sequence_parallel, bool): + raise TypeError("sequence_parallel must be a bool.") + deterministic = _resolve_deterministic_mode(deterministic, disable_split_k) + _validate_ffn_inputs(rmsnorm_output, gate_weight, up_weight, down_weight) + _require_ffn_kernels(disable_split_k=deterministic) + return _DeterministicFFNFunction.apply( + rmsnorm_output, + gate_weight, + up_weight, + down_weight, + tp_group, + cp_group, + sequence_parallel, + deterministic, + ) + + +def _resolve_deterministic_mode( + deterministic: bool | None, + disable_split_k: bool | None, +) -> bool: + if deterministic is not None and not isinstance(deterministic, bool): + raise TypeError("deterministic must be a bool or None.") + if disable_split_k is not None and not isinstance(disable_split_k, bool): + raise TypeError("disable_split_k must be a bool or None.") + if ( + deterministic is not None + and disable_split_k is not None + and deterministic != disable_split_k + ): + raise ValueError("deterministic and disable_split_k select conflicting FFN backends.") + if deterministic is not None: + return deterministic + if disable_split_k is not None: + return disable_split_k + return True + + +class Qwen3FFNOp: + """Instantiable Qwen3 FFN wrapper for semantic operator dispatch.""" + + op_class = "ffn" + is_batch_invariant = True + backend_id = BACKEND_ID + + def __call__( + self, + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + tp_group: Any = None, + cp_group: Any = None, + sequence_parallel: bool = False, + deterministic: bool | None = None, + disable_split_k: bool | None = None, + ) -> Tensor: + return self.apply( + rmsnorm_output, + gate_weight, + up_weight, + down_weight, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + deterministic=deterministic, + disable_split_k=disable_split_k, + ) + + def apply( + self, + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + tp_group: Any = None, + cp_group: Any = None, + sequence_parallel: bool = False, + deterministic: bool | None = None, + disable_split_k: bool | None = None, + ) -> Tensor: + return qwen3_ffn( + rmsnorm_output, + gate_weight, + up_weight, + down_weight, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + deterministic=deterministic, + disable_split_k=disable_split_k, + ) diff --git a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py new file mode 100644 index 00000000..5c3f56ae --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic vocab-parallel TP selected-token logprob reference (issue #241 PR3). + +Implements the WS2 contract in ``rl_engine.kernels.logprob_contract`` with a +TP-independent vocab tile decomposition: the padded vocabulary is split into +``num_vocab_tiles`` fixed tiles, every tile's fp32 ``(max, sumexp)`` partial is +computed from a contiguous ``[n, tile]`` tensor, all tile partials travel by +all-gather (transport only), and every rank merges them in global tile-index +order over a fixed ``[n, num_vocab_tiles]`` shape. The TP degree only decides +which rank computes which tiles and never changes any floating-point grouping, +so outputs and gradients are bitwise-identical across TP degrees +(``DeterminismScope.CROSS_TP_BITWISE``) as long as ``num_vocab_tiles`` is held +fixed. A fixed per-shard merge order alone cannot provide this property: +shard boundaries would regroup the combines differently at each degree. + +Consequences of the tile structure: + +- ``num_vocab_tiles`` is part of the numerical identity. It must be pinned + across ranks (enforced by the preflight) and across the TP degrees being + compared; it is never derived from the shard layout. +- Every shard boundary must be tile-aligned; misalignment fails loudly. +- At TP=1 the result matches the WS1 ``NativeBatchInvariantLogpOp`` only + within the #108 logprob tolerance, not bitwise — the WS1 op reduces the + whole ``[n, V]`` row at once, which groups the sums differently. + +Preconditions: logits over the real vocabulary must be finite. A row whose +real-vocab logits are all ``-inf`` has no finite logsumexp; with +``validate=True`` such a row fails loudly if it is active. + +The selected logprob is zero-filled at inactive rows (``MaskSpec.active_mask`` +is the sole authority; with validation enabled an active row can never legally +hold ``ignore_index``). The vocab-domain LSE is returned for every row and is +differentiable everywhere, including inactive rows. + +``deterministic=False`` trades the guarantee for speed: each rank reduces its +whole shard in one pass and the per-shard partials are merged in shard order, +so the floating-point grouping changes with the TP degree and nothing is +promised about reproducibility. ``num_vocab_tiles`` is ignored (no tile +alignment is required), and a contract declaring +``determinism_scope=cross_tp_bitwise`` is rejected loudly — the fast path +cannot honor it. Batch invariance is unaffected: rows never mix either way. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobContract, + LogprobContractError, + LogprobDType, +) + +BACKEND_ID = "pytorch-vocab-parallel-logp-ws2" +DEFAULT_NUM_VOCAB_TILES = 64 + +_TORCH_TO_CONTRACT_DTYPE = { + torch.bfloat16: LogprobDType.BF16, + torch.float16: LogprobDType.FP16, + torch.float32: LogprobDType.FP32, +} + + +def _require_distributed_initialized(): + import torch.distributed as dist + + if not dist.is_available(): + raise LogprobContractError("vocab-parallel logprob requires torch.distributed.") + if not dist.is_initialized(): + raise LogprobContractError( + "vocab-parallel logprob requires an initialized process group when " + "the contract declares tp_world_size > 1." + ) + return dist + + +def _tile_size(contract: LogprobContract, num_vocab_tiles: int) -> int: + if isinstance(num_vocab_tiles, bool) or not isinstance(num_vocab_tiles, int): + raise LogprobContractError( + f"num_vocab_tiles must be a positive integer; got {num_vocab_tiles!r}" + ) + if num_vocab_tiles <= 0: + raise LogprobContractError( + f"num_vocab_tiles must be a positive integer; got {num_vocab_tiles}" + ) + padded = contract.sharding.padded_vocab_size + if padded % num_vocab_tiles != 0: + raise LogprobContractError( + f"num_vocab_tiles={num_vocab_tiles} must divide " f"padded_vocab_size={padded} exactly" + ) + tile = padded // num_vocab_tiles + for rank, (start, end) in enumerate(contract.sharding.vocab_shard_bounds): + if start % tile != 0 or end % tile != 0: + raise LogprobContractError( + f"vocab_shard_bounds[{rank}]=[{start}, {end}) is not aligned to the " + f"vocab tile size {tile} (num_vocab_tiles={num_vocab_tiles}); " + "cross-TP bitwise determinism requires tile-aligned shard bounds" + ) + return tile + + +def _validate_invocation( + local_logits: torch.Tensor, + target_ids: torch.Tensor, + contract: LogprobContract, + tp_group: Any, +) -> None: + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + if local_logits.dim() != 2: + raise LogprobContractError( + f"local_logits must be 2-D [num_tokens, local_vocab]; got {local_logits.dim()}-D" + ) + if target_ids.dim() != 1 or target_ids.shape[0] != local_logits.shape[0]: + raise LogprobContractError( + f"target_ids must be 1-D with one entry per token; got shape " + f"{tuple(target_ids.shape)} for {local_logits.shape[0]} tokens" + ) + sharding = contract.sharding + if local_logits.shape[1] != sharding.local_vocab_size: + raise LogprobContractError( + f"local_logits has {local_logits.shape[1]} vocab columns but the contract " + f"declares local shard [{sharding.local_vocab_start}, " + f"{sharding.local_vocab_end}) of size {sharding.local_vocab_size}" + ) + if local_logits.shape[0] != contract.mask.num_tokens: + raise LogprobContractError( + f"local_logits has {local_logits.shape[0]} tokens but MaskSpec declares " + f"num_tokens={contract.mask.num_tokens}" + ) + declared = _TORCH_TO_CONTRACT_DTYPE.get(local_logits.dtype) + if declared is not contract.dtype: + raise LogprobContractError( + f"local_logits dtype {local_logits.dtype} does not match the contract " + f"dtype {contract.dtype.value}" + ) + if sharding.tp_world_size > 1: + dist = _require_distributed_initialized() + group_rank = dist.get_rank(group=tp_group) + group_world = dist.get_world_size(group=tp_group) + if group_world != sharding.tp_world_size: + raise LogprobContractError( + f"tp_group world size {group_world} does not match the contract " + f"tp_world_size={sharding.tp_world_size}; pass the TP subgroup, " + "not the global group" + ) + if group_rank != sharding.tp_rank: + raise LogprobContractError( + f"tp_group rank {group_rank} does not match the contract " + f"tp_rank={sharding.tp_rank}" + ) + + +def _validate_active_targets( + target_1d: torch.Tensor, active_mask: torch.Tensor, real_vocab_size: int +) -> None: + bad = active_mask & ((target_1d < 0) | (target_1d >= real_vocab_size)) + if bool(bad.any().item()): + bad_values = target_1d[bad] + raise LogprobContractError( + "active target_ids must lie in the real vocabulary " + f"[0, {real_vocab_size}); got values in " + f"[{int(bad_values.min().item())}, {int(bad_values.max().item())}] " + "on active rows" + ) + + +def _preflight_cross_rank_agreement( + contract: LogprobContract, tp_group: Any, num_vocab_tiles: int, deterministic: bool +) -> None: + """All-gather (fingerprint, backend id, tile count, mode) and abort on mismatch. + + The tile count travels as ``None`` when ``deterministic=False``: the fast + path never uses it, so ranks must not fail preflight over an irrelevant + value — but they must never disagree on the mode itself, or they would + issue different collectives. + """ + + dist = _require_distributed_initialized() + payload = ( + contract.cross_rank_fingerprint(), + BACKEND_ID, + int(num_vocab_tiles) if deterministic else None, + bool(deterministic), + ) + world = dist.get_world_size(group=tp_group) + gathered: list[Any] = [None] * world + dist.all_gather_object(gathered, payload, group=tp_group) + mismatched = [(rank, other) for rank, other in enumerate(gathered) if other != payload] + if mismatched: + rank, other = mismatched[0] + raise LogprobContractError( + "cross-rank preflight failed: rank " + f"{contract.sharding.tp_rank} has {payload} but rank {rank} has {other}; " + "all TP ranks must agree on the contract fingerprint, backend id, " + "num_vocab_tiles, and deterministic mode before any collective" + ) + + +def _local_tile_stats(z_masked: torch.Tensor, tile: int) -> tuple[torch.Tensor, torch.Tensor]: + """fp32 per-tile ``(max, sumexp)`` partials for this rank's shard. + + Each tile is reduced as a contiguous ``[n, tile]`` tensor so the reduction + shape and layout are identical no matter which rank computes the tile or + what the local shard size is. An all-``-inf`` (padding-only) tile yields + the identity partial ``(-inf, 0)`` without evaluating ``exp(-inf - (-inf))``. + """ + + n, local_vocab = z_masked.shape + m_parts: list[torch.Tensor] = [] + s_parts: list[torch.Tensor] = [] + for tile_index in range(local_vocab // tile): + block = z_masked[:, tile_index * tile : (tile_index + 1) * tile].contiguous() + m_t = block.max(dim=-1).values + finite = m_t > float("-inf") + m_safe = torch.where(finite, m_t, torch.zeros_like(m_t)) + s_t = (block - m_safe.unsqueeze(-1)).exp().sum(dim=-1) + s_t = torch.where(finite, s_t, torch.zeros_like(s_t)) + m_parts.append(m_t) + s_parts.append(s_t) + return torch.stack(m_parts, dim=1), torch.stack(s_parts, dim=1) + + +def _gather_tile_stats( + local_m: torch.Tensor, + local_s: torch.Tensor, + contract: LogprobContract, + tp_group: Any, + tile_counts: list[int], +) -> tuple[torch.Tensor, torch.Tensor]: + """Assemble every rank's partials in global shard order. + + ``tile_counts`` holds each rank's partial count: one per tile on the + deterministic path, exactly one per shard on the fast path. + """ + + sharding = contract.sharding + if sharding.tp_world_size == 1: + return local_m.contiguous(), local_s.contiguous() + + dist = _require_distributed_initialized() + n = local_m.shape[0] + max_tiles = max(tile_counts) + packed = local_m.new_zeros((n, max_tiles, 2)) + packed[:, : local_m.shape[1], 0] = local_m + packed[:, : local_s.shape[1], 1] = local_s + packed = packed.contiguous() + gathered = [torch.empty_like(packed) for _ in range(sharding.tp_world_size)] + dist.all_gather(gathered, packed, group=tp_group) + + m_parts = [gathered[rank][:, : tile_counts[rank], 0] for rank in range(len(tile_counts))] + s_parts = [gathered[rank][:, : tile_counts[rank], 1] for rank in range(len(tile_counts))] + return torch.cat(m_parts, dim=1).contiguous(), torch.cat(s_parts, dim=1).contiguous() + + +def _gather_target_logit( + z_masked: torch.Tensor, + safe_target: torch.Tensor, + contract: LogprobContract, + tp_group: Any, +) -> torch.Tensor: + """Exact selected-target logit via a select-by-owner copy.""" + + sharding = contract.sharding + n = z_masked.shape[0] + start = sharding.local_vocab_start + local_vocab = sharding.local_vocab_size + local_idx = (safe_target - start).clamp(0, max(local_vocab - 1, 0)) + owns = (safe_target >= start) & (safe_target < sharding.local_vocab_end) + rows = torch.arange(n, device=z_masked.device) + local_contrib = torch.where( + owns, z_masked[rows, local_idx], torch.zeros_like(safe_target, dtype=z_masked.dtype) + ).contiguous() + + if sharding.tp_world_size == 1: + stacked = local_contrib.unsqueeze(0) + else: + dist = _require_distributed_initialized() + gathered = [torch.empty_like(local_contrib) for _ in range(sharding.tp_world_size)] + dist.all_gather(gathered, local_contrib, group=tp_group) + stacked = torch.stack(gathered, dim=0) + + starts = torch.tensor( + [bound_start for bound_start, _ in sharding.vocab_shard_bounds], + device=safe_target.device, + dtype=torch.long, + ) + owner = torch.bucketize(safe_target, starts, right=True) - 1 + return stacked[owner, rows] + + +def _merge_tile_partials(m_all: torch.Tensor, s_all: torch.Tensor) -> torch.Tensor: + """Fixed-order (max, sumexp) merge over [n, num_vocab_tiles].""" + + M = m_all.max(dim=1).values + finite = M > float("-inf") + M_safe = torch.where(finite, M, torch.zeros_like(M)) + terms = s_all * (m_all - M_safe.unsqueeze(1)).exp() + S = terms.sum(dim=1) + return M + S.log() + + +class _VocabParallelLogprobFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, local_logits, target_1d, active_mask, contract, tp_group, tile): + z_masked = local_logits.float() + sharding = contract.sharding + global_ids = torch.arange( + sharding.local_vocab_start, sharding.local_vocab_end, device=z_masked.device + ) + padding_cols = global_ids >= sharding.real_vocab_size + if bool(padding_cols.any()): + z_masked = z_masked.masked_fill(padding_cols.unsqueeze(0), float("-inf")) + + safe_target = torch.where(active_mask, target_1d, torch.zeros_like(target_1d)) + + if tile is not None: + local_m, local_s = _local_tile_stats(z_masked, tile) + tile_counts = [(end - start) // tile for start, end in sharding.vocab_shard_bounds] + else: + # Fast path: one (max, sumexp) partial over the whole shard. The + # grouping now depends on the shard layout, so no cross-TP claim. + local_m, local_s = _local_tile_stats(z_masked, z_masked.shape[1]) + tile_counts = [1] * sharding.tp_world_size + m_all, s_all = _gather_tile_stats(local_m, local_s, contract, tp_group, tile_counts) + target_logit = _gather_target_logit(z_masked, safe_target, contract, tp_group) + lse = _merge_tile_partials(m_all, s_all) + + selected_logp = torch.where(active_mask, target_logit - lse, torch.zeros_like(lse)) + + ctx.save_for_backward(z_masked, lse, safe_target, active_mask, padding_cols) + ctx.local_vocab_start = sharding.local_vocab_start + ctx.local_vocab_size = sharding.local_vocab_size + ctx.input_dtype = local_logits.dtype + ctx.set_materialize_grads(False) + return selected_logp, lse + + @staticmethod + def backward(ctx, grad_logp, grad_lse): + if not ctx.needs_input_grad[0] or (grad_logp is None and grad_lse is None): + return None, None, None, None, None, None + + z_masked, lse, safe_target, active_mask, padding_cols = ctx.saved_tensors + n, local_vocab = z_masked.shape + finite_row = torch.isfinite(lse) + lse_safe = torch.where(finite_row, lse, torch.zeros_like(lse)) + p = (z_masked - lse_safe.unsqueeze(1)).exp() + p = torch.where(finite_row.unsqueeze(1), p, torch.zeros_like(p)) + + grad = torch.zeros_like(z_masked) + if grad_logp is not None: + local_idx = (safe_target - ctx.local_vocab_start).clamp(0, max(local_vocab - 1, 0)) + owns = (safe_target >= ctx.local_vocab_start) & ( + safe_target < ctx.local_vocab_start + local_vocab + ) + onehot = torch.zeros_like(z_masked) + hit = owns & active_mask + rows = torch.arange(n, device=z_masked.device)[hit] + onehot[rows, local_idx[hit]] = 1.0 + g_logp = torch.where(active_mask, grad_logp, torch.zeros_like(grad_logp)) + grad = grad + g_logp.unsqueeze(1) * (onehot - p) + if grad_lse is not None: + grad = grad + grad_lse.unsqueeze(1) * p + if bool(padding_cols.any()): + grad = grad.masked_fill(padding_cols.unsqueeze(0), 0.0) + return grad.to(ctx.input_dtype), None, None, None, None, None + + +class VocabParallelLogprobOp: + """Vocab-parallel selected-token logprob (WS2 reference). + + Deterministic by default (cross-TP bitwise, tile-ordered merge); + ``deterministic=False`` selects the faster whole-shard reduction with no + reproducibility guarantee. + """ + + op_class = "logprob" + is_batch_invariant = True + backend_id = BACKEND_ID + + def __init__(self) -> None: + pass + + def __call__( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any = None, + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, + validate: bool = True, + deterministic: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + deterministic=deterministic, + ) + + def apply( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any = None, + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, + validate: bool = True, + deterministic: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + if not isinstance(deterministic, bool): + raise LogprobContractError(f"deterministic must be a bool; got {deterministic!r}") + if deterministic: + tile = _tile_size(contract, num_vocab_tiles) + elif contract.reduction.determinism_scope is DeterminismScope.CROSS_TP_BITWISE: + raise LogprobContractError( + "deterministic=False cannot honor determinism_scope=cross_tp_bitwise: " + "the fast path reduces each shard in one piece, so its floating-point " + "grouping changes with the TP degree; keep deterministic=True or relax " + "the contract to determinism_scope=fixed_topology" + ) + else: + tile = None + _validate_invocation(local_logits, target_ids, contract, tp_group) + + target_1d = target_ids.reshape(-1).to(device=local_logits.device, dtype=torch.long) + active_mask = torch.tensor( + contract.mask.active_mask, dtype=torch.bool, device=local_logits.device + ) + if validate: + _validate_active_targets(target_1d, active_mask, contract.sharding.real_vocab_size) + if contract.sharding.tp_world_size > 1: + _preflight_cross_rank_agreement(contract, tp_group, num_vocab_tiles, deterministic) + + selected_logp, lse = _VocabParallelLogprobFunction.apply( + local_logits, target_1d, active_mask, contract, tp_group, tile + ) + + if validate and bool((~torch.isfinite(lse) & active_mask).any().item()): + raise LogprobContractError( + "non-finite logsumexp on an active row: logits over the real " + "vocabulary must be finite for every active token" + ) + return selected_logp, lse + + +__all__ = [ + "BACKEND_ID", + "DEFAULT_NUM_VOCAB_TILES", + "VocabParallelLogprobOp", +] diff --git a/rl_engine/kernels/ops/triton/activation/__init__.py b/rl_engine/kernels/ops/triton/activation/__init__.py new file mode 100644 index 00000000..f6950928 --- /dev/null +++ b/rl_engine/kernels/ops/triton/activation/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .swiglu import TritonSiLUOp, TritonSwiGLUOp + +__all__ = ["TritonSiLUOp", "TritonSwiGLUOp"] diff --git a/rl_engine/kernels/ops/triton/activation/swiglu.py b/rl_engine/kernels/ops/triton/activation/swiglu.py new file mode 100644 index 00000000..6fb66313 --- /dev/null +++ b/rl_engine/kernels/ops/triton/activation/swiglu.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Triton SiLU / SwiGLU ops matching NativeSiLUOp / NativeSwiGLUOp (WS1 ground truth). + +Math is performed in fp32 inside the Triton kernels and rounded back to the input +dtype on store — the same dual-path contract as the PyTorch references: + + silu(x) = x * sigmoid(x) + swiglu(g, u) = silu(g) * u + +Element-wise and row-independent, so Axis-A batch invariance holds bitwise. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from torch import Tensor + +_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32) + + +def _validate_dtype(x: Tensor, name: str) -> None: + if x.dtype not in _SUPPORTED_DTYPES: + raise TypeError(f"{name} must have dtype fp16, bf16, or fp32, got {x.dtype}.") + + +@triton.jit +def _silu_fwd_kernel(x_ptr, y_ptr, n_elements, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) + s = 1.0 / (1.0 + tl.exp(-x)) + y = x * s + tl.store(y_ptr + offs, y.to(y_ptr.dtype.element_ty), mask=mask) + + +@triton.jit +def _silu_bwd_kernel(dy_ptr, x_ptr, dx_ptr, n_elements, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + dy = tl.load(dy_ptr + offs, mask=mask, other=0.0).to(tl.float32) + x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) + s = 1.0 / (1.0 + tl.exp(-x)) + # silu'(x) = s * (1 + x * (1 - s)) + dx = dy * s * (1.0 + x * (1.0 - s)) + tl.store(dx_ptr + offs, dx.to(dx_ptr.dtype.element_ty), mask=mask) + + +@triton.jit +def _swiglu_fwd_kernel(gate_ptr, up_ptr, y_ptr, n_elements, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + g = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32) + u = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32) + s = 1.0 / (1.0 + tl.exp(-g)) + y = (g * s) * u + tl.store(y_ptr + offs, y.to(y_ptr.dtype.element_ty), mask=mask) + + +@triton.jit +def _swiglu_bwd_kernel( + dy_ptr, gate_ptr, up_ptr, d_gate_ptr, d_up_ptr, n_elements, BLOCK: tl.constexpr +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + dy = tl.load(dy_ptr + offs, mask=mask, other=0.0).to(tl.float32) + g = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32) + u = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32) + s = 1.0 / (1.0 + tl.exp(-g)) + silu_g = g * s + d_up = dy * silu_g + d_gate = dy * u * s * (1.0 + g * (1.0 - s)) + tl.store(d_up_ptr + offs, d_up.to(d_up_ptr.dtype.element_ty), mask=mask) + tl.store(d_gate_ptr + offs, d_gate.to(d_gate_ptr.dtype.element_ty), mask=mask) + + +_BLOCK = 1024 + + +def _launch_silu_fwd(x: Tensor) -> Tensor: + x_c = x.contiguous() + y = torch.empty_like(x_c) + n = x_c.numel() + if n == 0: + return y + grid = (triton.cdiv(n, _BLOCK),) + _silu_fwd_kernel[grid](x_c, y, n, BLOCK=_BLOCK) + return y + + +def _launch_silu_bwd(dy: Tensor, x: Tensor) -> Tensor: + dy_c = dy.contiguous() + x_c = x.contiguous() + dx = torch.empty_like(x_c) + n = x_c.numel() + if n == 0: + return dx + grid = (triton.cdiv(n, _BLOCK),) + _silu_bwd_kernel[grid](dy_c, x_c, dx, n, BLOCK=_BLOCK) + return dx + + +def _launch_swiglu_fwd(gate: Tensor, up: Tensor) -> Tensor: + gate_c = gate.contiguous() + up_c = up.contiguous() + y = torch.empty_like(gate_c) + n = gate_c.numel() + if n == 0: + return y + grid = (triton.cdiv(n, _BLOCK),) + _swiglu_fwd_kernel[grid](gate_c, up_c, y, n, BLOCK=_BLOCK) + return y + + +def _launch_swiglu_bwd(dy: Tensor, gate: Tensor, up: Tensor) -> tuple[Tensor, Tensor]: + dy_c = dy.contiguous() + gate_c = gate.contiguous() + up_c = up.contiguous() + d_gate = torch.empty_like(gate_c) + d_up = torch.empty_like(up_c) + n = gate_c.numel() + if n == 0: + return d_gate, d_up + grid = (triton.cdiv(n, _BLOCK),) + _swiglu_bwd_kernel[grid](dy_c, gate_c, up_c, d_gate, d_up, n, BLOCK=_BLOCK) + return d_gate, d_up + + +class _SiLUTritonFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor) -> Tensor: + x_c = x.contiguous() + y = _launch_silu_fwd(x_c) + ctx.save_for_backward(x_c) + return y + + @staticmethod + def backward(ctx, grad_out: Tensor): + (x,) = ctx.saved_tensors + dx = None + if ctx.needs_input_grad[0]: + dx = _launch_silu_bwd(grad_out, x) + return dx + + +class _SwiGLUTritonFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, gate: Tensor, up: Tensor) -> Tensor: + gate_c = gate.contiguous() + up_c = up.contiguous() + y = _launch_swiglu_fwd(gate_c, up_c) + ctx.save_for_backward(gate_c, up_c) + return y + + @staticmethod + def backward(ctx, grad_out: Tensor): + gate, up = ctx.saved_tensors + d_gate = d_up = None + if ctx.needs_input_grad[0] or ctx.needs_input_grad[1]: + dg, du = _launch_swiglu_bwd(grad_out, gate, up) + if ctx.needs_input_grad[0]: + d_gate = dg + if ctx.needs_input_grad[1]: + d_up = du + return d_gate, d_up + + +class TritonSiLUOp: + """Triton SiLU: ``out = x * sigmoid(x)``, math in fp32.""" + + op_class = "elementwise" + + def __call__(self, x: Tensor) -> Tensor: + return self.forward(x) + + def forward(self, x: Tensor) -> Tensor: + if x.device.type not in ("cuda", "hip", "xpu"): + raise RuntimeError(f"TritonSiLUOp requires a GPU tensor, got device '{x.device}'.") + _validate_dtype(x, "x") + return _SiLUTritonFunction.apply(x) + + def forward_fp32(self, x: Tensor) -> Tensor: + if x.device.type not in ("cuda", "hip", "xpu"): + raise RuntimeError(f"TritonSiLUOp requires a GPU tensor, got device '{x.device}'.") + _validate_dtype(x, "x") + return _SiLUTritonFunction.apply(x.float()) + + +class TritonSwiGLUOp: + """Triton SwiGLU: ``out = silu(gate) * up``, math in fp32.""" + + op_class = "elementwise" + + def __call__(self, gate: Tensor, up: Tensor) -> Tensor: + return self.forward(gate, up) + + def forward(self, gate: Tensor, up: Tensor) -> Tensor: + if gate.device.type not in ("cuda", "hip", "xpu") or up.device.type not in ( + "cuda", + "hip", + "xpu", + ): + raise RuntimeError( + f"TritonSwiGLUOp requires GPU tensors, got gate='{gate.device}', up='{up.device}'." + ) + if gate.device != up.device: + raise RuntimeError( + f"gate and up must be on the same GPU device, got " + f"'{gate.device}' and '{up.device}'." + ) + if gate.shape != up.shape: + raise ValueError( + f"gate and up must share shape, got tuple(gate.shape)=" + f"{tuple(gate.shape)} vs tuple(up.shape)={tuple(up.shape)}" + ) + _validate_dtype(gate, "gate") + _validate_dtype(up, "up") + if gate.dtype != up.dtype: + raise TypeError(f"gate and up must share dtype, got {gate.dtype} and {up.dtype}.") + return _SwiGLUTritonFunction.apply(gate, up) + + def forward_fp32(self, gate: Tensor, up: Tensor) -> Tensor: + if gate.device.type not in ("cuda", "hip", "xpu") or up.device.type not in ( + "cuda", + "hip", + "xpu", + ): + raise RuntimeError( + f"TritonSwiGLUOp requires GPU tensors, got gate='{gate.device}', up='{up.device}'." + ) + if gate.device != up.device: + raise RuntimeError( + f"gate and up must be on the same GPU device, got " + f"'{gate.device}' and '{up.device}'." + ) + if gate.shape != up.shape: + raise ValueError( + f"gate and up must share shape, got tuple(gate.shape)=" + f"{tuple(gate.shape)} vs tuple(up.shape)={tuple(up.shape)}" + ) + _validate_dtype(gate, "gate") + _validate_dtype(up, "up") + if gate.dtype != up.dtype: + raise TypeError(f"gate and up must share dtype, got {gate.dtype} and {up.dtype}.") + return _SwiGLUTritonFunction.apply(gate.float(), up.float()) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index e9a7d075..747e6227 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -8,6 +8,8 @@ from enum import Enum, EnumMeta from typing import Any, Dict, Optional, Set, Type +import torch + from rl_engine.kernels.attention_contract import ( AttentionBackendCapability, AttentionContract, @@ -17,6 +19,17 @@ AttentionMode, AttentionRole, ) +from rl_engine.kernels.logprob_contract import ( + IMPLEMENTATION_KINDS, + DeterminismScope, + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDispatchResult, + LogprobDType, + LogprobRole, + MaskMode, +) from rl_engine.kernels.semantic_registry import ( OperatorBackendDescriptor, OperatorFallbackPolicy, @@ -91,6 +104,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_BATCH_INVARIANT_LOGP_SM90 = ( "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90Op" ) + # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) + PYTORCH_VOCAB_PARALLEL_LOGP = ( + "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" + ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" @@ -105,6 +122,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_ROPE_SM90 = "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op" PYTORCH_NATIVE_SILU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp" PYTORCH_NATIVE_SWIGLU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp" + CUDA_SILU = "rl_engine.kernels.ops.cuda.activation.swiglu.SiLUCudaOp" + CUDA_SWIGLU = "rl_engine.kernels.ops.cuda.activation.swiglu.SwiGLUCudaOp" + TRITON_SILU = "rl_engine.kernels.ops.triton.activation.swiglu.TritonSiLUOp" + TRITON_SWIGLU = "rl_engine.kernels.ops.triton.activation.swiglu.TritonSwiGLUOp" # WS1 pure-PyTorch ground-truth attention reference (hand-written fp32 softmax). # Distinct from PYTORCH_ATTN above, which is the production SDPA fallback. @@ -174,6 +195,27 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, version_or_build_fingerprint="runtime-native-unresolved-v1", ), + OperatorBackendDescriptor( + semantic_op="selected_logprob", + backend_id="pytorch-vocab-parallel-logp-ws2", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu", "cuda", "rocm"}), + supported_dtypes=frozenset({"float32", "bfloat16", "float16"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "algorithm": "fixed_global_vocab_tiles", + "deterministic": True, + "cross_tp_bitwise": True, + "exports_vocab_lse": True, + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.DISTRIBUTED_CONTEXT, + implementation_class_or_factory=( + "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp." "VocabParallelLogprobOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="VocabParallelLogprobOp-fixed-tiles-v1", + ), OperatorBackendDescriptor( semantic_op="attention", backend_id="rlkernel.attention.deterministic.v1", @@ -213,6 +255,40 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, version_or_build_fingerprint="runtime-native-attention-unresolved-v1", ), + OperatorBackendDescriptor( + semantic_op="ffn", + backend_id="rlkernel.ffn.qwen3.deterministic.v1", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cuda"}), + supported_dtypes=frozenset({"bfloat16"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "algorithm": "qwen3_swiglu_fixed_reduction_gemm", + "batch_invariant": True, + "deterministic": True, + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.DISTRIBUTED_CONTEXT, + implementation_class_or_factory=("rl_engine.kernels.ops.pytorch.ffn.ffn.Qwen3FFNOp"), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="Qwen3FFNOp-fixed-reduction-v1", + ), + OperatorBackendDescriptor( + semantic_op="ffn", + backend_id="native", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"*"}), + supported_dtypes=frozenset({"*"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "selection": "runtime_native", + "strict_observable": False, + }, + lifecycle=OperatorLifecycle.DISTRIBUTED_CONTEXT, + implementation_class_or_factory=None, + fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, + version_or_build_fingerprint="runtime-native-ffn-unresolved-v1", + ), ) @@ -319,6 +395,47 @@ def __init__(self): ), } + # Truthful descriptors for the existing WS1 batch-invariant logp + # implementations: single-shard (TP=1), ignore-index masking only, no + # vocab-shard metadata, no vocab-domain LSE export. + common_logprob_roles = frozenset({LogprobRole.TRAIN, LogprobRole.INFER}) + common_logprob_dtypes = frozenset({LogprobDType.BF16, LogprobDType.FP16, LogprobDType.FP32}) + base_logprob_capabilities = { + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP: LogprobBackendCapability( + backend_id="pytorch-batch-invariant-logp-ws1", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="reference", + ), + OpBackend.TRITON_BATCH_INVARIANT_LOGP: LogprobBackendCapability( + backend_id="triton-batch-invariant-logp-ws1", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", + ), + OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90: LogprobBackendCapability( + backend_id="cuda-batch-invariant-logp-sm90-ws1", + roles=common_logprob_roles, + dtypes=frozenset({LogprobDType.BF16, LogprobDType.FP32}), + tp_world_sizes=(1,), + supports_vocab_padding=False, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), + exports_vocab_lse=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", + ), + } + self._priority_map = { "cuda": { "logp": [ @@ -347,7 +464,11 @@ def __init__(self): OpBackend.CUDA_DETERMINISTIC_LOGP, OpBackend.PYTORCH_NATIVE, ], - "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], + "attn": [ + OpBackend.FLASH_ATTN, + OpBackend.TRITON_GENERIC, + OpBackend.PYTORCH_ATTN, + ], "attention": [ OpBackend.CUDA_DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION, @@ -374,8 +495,17 @@ def __init__(self): "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], - "silu": [OpBackend.PYTORCH_NATIVE_SILU], - "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], + "silu": [ + OpBackend.CUDA_SILU, + OpBackend.TRITON_SILU, + OpBackend.PYTORCH_NATIVE_SILU, + ], + "swiglu": [ + OpBackend.CUDA_SWIGLU, + OpBackend.TRITON_SWIGLU, + OpBackend.PYTORCH_NATIVE_SWIGLU, + ], + # Default dispatch logic for new operators "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rope": [ OpBackend.CUDA_ROPE_SM90, @@ -405,7 +535,10 @@ def __init__(self): "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], - "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], + "linear_logp": [ + OpBackend.TRITON_LINEAR_LOGP, + OpBackend.PYTORCH_LINEAR_LOGP, + ], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "pack": [OpBackend.PYTORCH_PACK], "det_gemm": [OpBackend.TRITON_DET_GEMM], @@ -417,8 +550,8 @@ def __init__(self): "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], - "silu": [OpBackend.PYTORCH_NATIVE_SILU], - "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], + "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], + "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -450,6 +583,47 @@ def __init__(self): self._adjust_priority_for_hardware() self._adjust_priority_from_env() + # WS2 dispatch owns its candidate list, seeded from the legacy + # batch_invariant_logp priority but decoupled afterwards: neither + # path's registrations may affect the other. + self._logprob_candidates: Dict[str, list] = { + platform: list(ops.get("batch_invariant_logp", [])) + for platform, ops in self._priority_map.items() + } + # Capabilities are scoped per platform: the same backend enum may + # truthfully declare different support on cuda vs rocm vs cpu. + self._logprob_capabilities: Dict[str, Dict[OpBackend, LogprobBackendCapability]] = { + platform: { + backend: base_logprob_capabilities[backend] + for backend in candidates + if backend in base_logprob_capabilities + } + for platform, candidates in self._logprob_candidates.items() + } + + # deterministic vocab-parallel TP logprob reference. + ws2_tp_logprob_capability = LogprobBackendCapability( + backend_id="pytorch-vocab-parallel-logp-ws2", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=None, + cp_world_sizes=None, + supports_vocab_padding=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), + exports_vocab_lse=True, + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", + ) + for ws2_platform in self._priority_map: + self.register_logprob_backend( + OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP, + ws2_tp_logprob_capability, + platform=ws2_platform, + prepend=True, + ) + def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() if rocm_attn_backend in {"flash_attn", "flash-attn", "flash_attention"}: @@ -529,24 +703,158 @@ def get_op(self, op_type: str) -> Any: ) for backend in candidates: - if backend.name in self._instance_cache: - return self._instance_cache[backend.name] - if backend.name in self._failed_backends: - continue + op_instance = self._get_or_create_backend(backend) + if op_instance is not None: + return op_instance + + raise RuntimeError(f"No functional backend found for {op_type} on {platform}") - op_class = self._load_backend(backend) - if op_class: - try: - op_instance = op_class() - self._instance_cache[backend.name] = op_instance - return op_instance - except Exception as exc: - logger.error(f"Failed to instantiate {backend.name}: {exc}") - self._failed_backends.add(backend.name) + def register_logprob_backend( + self, + backend: OpBackend, + capability: LogprobBackendCapability, + *, + platform: Optional[str] = None, + prepend: bool = False, + ) -> None: + """Register (or replace) a backend for WS2 contract-aware logprob dispatch. + + This is the supported seam for making a new backend selectable by + ``get_logprob_op`` (e.g. the deterministic vocab-parallel TP reference + from issue #241 PR 3) without touching the legacy ``get_op`` priority + lists. Registering the same backend again replaces its capability + without duplicating the candidate entry. + """ + + if not isinstance(backend, OpBackend): + raise LogprobContractError("backend must be an OpBackend") + if not isinstance(capability, LogprobBackendCapability): + raise LogprobContractError("capability must be a LogprobBackendCapability") + resolved_platform = platform if platform is not None else self._platform() + if resolved_platform not in self._priority_map: + raise LogprobContractError( + f"unsupported platform {resolved_platform!r}; expected one of " + f"{sorted(self._priority_map)}" + ) + candidates = self._logprob_candidates.setdefault(resolved_platform, []) + self._logprob_capabilities.setdefault(resolved_platform, {})[backend] = capability + if backend not in candidates: + if prepend: + candidates.insert(0, backend) else: - self._failed_backends.add(backend.name) + candidates.append(backend) - raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + def get_logprob_op( + self, + contract: LogprobContract, + *, + requested_backend: str = "auto", + ) -> LogprobDispatchResult: + """Resolve only a backend that explicitly supports the WS2 logprob contract. + + This entry point is intentionally separate from legacy ``get_op`` so + existing callers retain their current behavior while WS2 callers cannot + silently fall back to a backend with different distributed semantics. + + ``requested_backend`` is either a case-insensitive policy keyword + (``auto`` | ``production`` | ``reference`` | ``deterministic``) or an + exact, case-sensitive stable backend id. Strictness comes from the + contract's capability checks, not from this policy string, so the + default is ``auto``. + """ + + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + if not isinstance(requested_backend, str) or not requested_backend.strip(): + raise LogprobContractError("requested_backend must be a non-empty string") + requested_backend = requested_backend.strip() + if requested_backend.lower() == "deterministic": + raise LogprobContractError( + 'requested_backend="deterministic" is not a dispatch policy; request ' + "determinism through ReductionSpec.determinism_scope and match it against " + "backend determinism_scopes instead" + ) + + platform = self._platform() + candidates = self._logprob_candidates.get(platform, []) + rejected: list[str] = [] + # provenance["fallback"] reports only capability/load rejections of + # otherwise-eligible candidates; skips caused purely by the caller's + # own requested_backend policy filter are not fallbacks. + capability_rejections = 0 + + platform_capabilities = self._logprob_capabilities.get(platform, {}) + for backend in candidates: + capability = platform_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no LogprobBackendCapability declared") + capability_rejections += 1 + continue + policy_mismatch = self._logprob_policy_mismatch(requested_backend, capability) + if policy_mismatch is not None: + # Excluded by the caller's own policy: never a fallback, even + # if the candidate would also have failed capability checks. + rejected.append(f"{backend.name}: {policy_mismatch}") + continue + capability_incompat = list(capability.incompatibilities(contract)) + if capability_incompat: + rejected.append(f"{backend.name}: " + "; ".join(capability_incompat)) + capability_rejections += 1 + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + capability_rejections += 1 + continue + + provenance = { + "requested_backend": requested_backend, + "actual_backend": capability.backend_id, + "backend_enum": backend.name, + "platform": platform, + "fallback": capability_rejections > 0, + "prior_rejections": list(rejected), + "contract": contract.to_dict(), + "capability": capability.to_dict(), + } + return LogprobDispatchResult( + op=op, + capability=capability, + provenance=provenance, + ) + + details = " | ".join(rejected) if rejected else "no candidates registered" + requested = contract.to_dict() + raise RuntimeError( + "No logprob backend supports the requested WS2 contract on " + f"{platform}: role={requested['role']}, dtype={requested['dtype']}, " + f"TP={contract.sharding.tp_world_size}, CP={contract.sharding.cp_world_size}, " + f"padded_vocab={contract.sharding.padded_vocab_size}, " + f"real_vocab={contract.sharding.real_vocab_size}. Rejections: {details}" + ) + + @staticmethod + def _logprob_policy_mismatch( + requested_backend: str, + capability: LogprobBackendCapability, + ) -> str | None: + policy = requested_backend.lower() + if policy == "auto": + return None + if policy in IMPLEMENTATION_KINDS: + if capability.implementation_kind == policy: + return None + return ( + f"implementation_kind={capability.implementation_kind} does not satisfy " + f"requested_backend={policy}" + ) + if capability.backend_id == requested_backend: + return None + return ( + f"backend_id={capability.backend_id} does not match " + f"requested_backend={requested_backend}" + ) def get_attention_op( self, @@ -637,6 +945,16 @@ def _platform() -> str: return "cuda" return "cpu" + def _platform_for_device(self, device: torch.device | str | None) -> str: + if device is None: + return self._platform() + resolved = torch.device(device) + if resolved.type == "cuda": + return "rocm" if torch.version.hip is not None else "cuda" + if resolved.type in self._priority_map: + return resolved.type + return "cpu" + def _get_or_create_backend(self, backend: OpBackend) -> Optional[Any]: if backend.name in self._instance_cache: return self._instance_cache[backend.name] diff --git a/setup.py b/setup.py index a5679a85..3d36a0ad 100644 --- a/setup.py +++ b/setup.py @@ -81,7 +81,9 @@ def get_extensions(): "csrc/cuda/attention/prefix_shared_attention.cu", "csrc/cuda/gemm/det_gemm_kernel.cu", "csrc/cuda/rmsnorm.cu", + "csrc/cuda/activation.cu", "csrc/cuda/attention/deterministic_attention.cu", + "csrc/cuda/distributed/deterministic_collective.cu", ] cc_major, cc_minor = torch.cuda.get_device_capability() @@ -146,6 +148,9 @@ def get_extensions(): cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] extra_link_args = list(torch_rpath) + if os.name != "nt": + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") sm90_srcs = [ "csrc/cuda/fused_logp_sm90.cu", @@ -160,7 +165,8 @@ def get_extensions(): cuda_sources.extend(present_sm90) nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - extra_link_args.append("-lcuda") + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") # det_gemm SM90 (mma.sync + TMA) path — independent of the fused_logp # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in diff --git a/tests/distributed/__init__.py b/tests/distributed/__init__.py new file mode 100644 index 00000000..86cf4c9d --- /dev/null +++ b/tests/distributed/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/tests/distributed/test_det_gemm_simulated_tp.py b/tests/distributed/test_det_gemm_simulated_tp.py new file mode 100644 index 00000000..c4824e85 --- /dev/null +++ b/tests/distributed/test_det_gemm_simulated_tp.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Simulate TP by splitting GEMM K, without a real process group. + +TP=2 is one BF16 add of two shards (`a+b == b+a`), so it can match TP=1. +TP=8 left-fold is a different parenthesization from the kernel K-tree, so it +must not match. Real NCCL / custom AllReduce is out of scope here. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.cuda.matmul import deterministic_gemm + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="det_gemm requires CUDA SM80+", +) + +DEV = "cuda" + + +def _rand(*shape): + return torch.randn(*shape, device=DEV, dtype=torch.bfloat16) + + +def _k_shards(a: torch.Tensor, b: torch.Tensor, tp: int) -> list[torch.Tensor]: + k = a.shape[1] + assert k % tp == 0 + width = k // tp + return [ + deterministic_gemm( + a[:, i * width : (i + 1) * width].contiguous(), + b[i * width : (i + 1) * width].contiguous(), + ) + for i in range(tp) + ] + + +def _left_fold(parts: list[torch.Tensor]) -> torch.Tensor: + acc = parts[0] + for part in parts[1:]: + acc = acc + part + return acc + + +def test_simulated_tp2_matches_full(): + # Two shards: AllReduce is a+b, and BF16 add is commutative. + torch.manual_seed(8) + m, k, n = 16, 256, 64 + a, b = _rand(m, k), _rand(k, n) + left, right = _k_shards(a, b, 2) + full = deterministic_gemm(a, b) + assert torch.equal(full, left + right) + assert torch.equal(left + right, right + left) + + +def test_simulated_tp8_left_fold_does_not_match_full(): + # Eight shards left-folded: ((((s0+s1)+s2)+...)+s7) is not the kernel tree + # ((s0+s1)+(s2+s3))+((s4+s5)+(s6+s7)), so this must diverge. + torch.manual_seed(8) + m, k, n = 16, 256, 64 + a, b = _rand(m, k), _rand(k, n) + full = deterministic_gemm(a, b) + folded = _left_fold(_k_shards(a, b, 8)) + n_mismatch = int((full != folded).sum().item()) + assert n_mismatch > 0, "TP=8 left-fold unexpectedly matched TP=1" + + +def test_simulated_tp2_is_batch_invariant(): + torch.manual_seed(9) + k, n = 256, 64 + b = _rand(k, n) + row = _rand(1, k) + big = _rand(32, k) + big[0] = row[0] + + def tp2(x): + left, right = _k_shards(x, b, 2) + return left + right + + assert torch.equal(tp2(row)[0], tp2(big)[0]) + assert torch.equal(tp2(row)[0], deterministic_gemm(big, b)[0]) diff --git a/tests/distributed/test_deterministic_all_gather.py b/tests/distributed/test_deterministic_all_gather.py new file mode 100644 index 00000000..d6058dd3 --- /dev/null +++ b/tests/distributed/test_deterministic_all_gather.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import os +import socket +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.distributed import DeterministicCollective + +_MAX_WORLD_SIZE = 8 +_TP_SIZES = (1, 2, 4, 8) +_EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + +pytestmark = [ + pytest.mark.skipif( + _EXTERNAL_WORLD_SIZE != 1, + reason="this cross-TP test owns its worker processes; run pytest directly", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _MAX_WORLD_SIZE, + reason="requires eight visible CUDA GPUs", + ), +] + + +def _all_gather_tensors( + value: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, +) -> list[torch.Tensor]: + gathered = [torch.empty_like(value) for _ in range(world_size)] + dist.all_gather(gathered, value, group=group) + return gathered + + +def _make_global_input(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + if dtype == torch.int64: + return torch.arange( + _MAX_WORLD_SIZE * 13 * 7, + device=device, + dtype=dtype, + ).reshape(_MAX_WORLD_SIZE * 13, 7) + generator = torch.Generator().manual_seed(20260817) + return torch.randn( + _MAX_WORLD_SIZE * 13, + 7, + dtype=torch.float32, + generator=generator, + ).to(device=device, dtype=dtype) + + +def _worker(rank: int, port: int) -> None: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=_MAX_WORLD_SIZE, + timeout=timedelta(minutes=5), + ) + try: + groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} + for tp_size, group in groups.items(): + if rank < tp_size: + with DeterministicCollective( + group=group, + device=device, + max_size_bytes=1024 * 1024, + ) as collective: + group_rank = dist.get_rank(group=group) + for dtype in (torch.float32, torch.bfloat16, torch.int64): + expected = _make_global_input(device, dtype) + input = expected.chunk(tp_size, dim=0)[group_rank].contiguous() + + output = collective.all_gather(input) + assert torch.equal(output, expected) + + peer_outputs = _all_gather_tensors(output, group, tp_size) + assert all(torch.equal(peer_output, output) for peer_output in peer_outputs) + + baseline = output.clone() + for _ in range(3): + repeated = collective.all_gather(input) + assert torch.equal(repeated, baseline) + + provided = torch.empty_like(expected) + returned = collective.all_gather(input, out=provided) + assert returned is provided + assert torch.equal(provided, expected) + dist.barrier() + finally: + dist.destroy_process_group() + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_deterministic_all_gather_cross_tp_cuda() -> None: + mp.spawn( + _worker, + args=(_find_free_port(),), + nprocs=_MAX_WORLD_SIZE, + join=True, + ) diff --git a/tests/distributed/test_deterministic_all_reduce.py b/tests/distributed/test_deterministic_all_reduce.py new file mode 100644 index 00000000..eb406881 --- /dev/null +++ b/tests/distributed/test_deterministic_all_reduce.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import os +import socket +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.distributed import DeterministicCollective + +_MAX_WORLD_SIZE = 8 +_TP_SIZES = (1, 2, 4, 8) +_EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + +pytestmark = [ + pytest.mark.skipif( + _EXTERNAL_WORLD_SIZE != 1, + reason="this cross-TP test owns its worker processes; run pytest directly", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _MAX_WORLD_SIZE, + reason="requires eight visible CUDA GPUs", + ), +] + + +def _fixed_tree_reference(values: list[torch.Tensor]) -> torch.Tensor: + level = values + while len(level) > 1: + level = [level[index] + level[index + 1] for index in range(0, len(level), 2)] + return level[0] + + +def _all_gather_tensors( + value: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, +) -> list[torch.Tensor]: + gathered = [torch.empty_like(value) for _ in range(world_size)] + dist.all_gather(gathered, value, group=group) + return gathered + + +def _worker(rank: int, port: int) -> None: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=_MAX_WORLD_SIZE, + timeout=timedelta(minutes=5), + ) + try: + groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} + for tp_size, group in groups.items(): + if rank < tp_size: + with DeterministicCollective( + group=group, + device=device, + max_size_bytes=1024 * 1024, + ) as collective: + group_rank = dist.get_rank(group=group) + leaves_per_rank = _MAX_WORLD_SIZE // tp_size + start = group_rank * leaves_per_rank + for dtype in (torch.float32, torch.float16, torch.bfloat16): + generator = torch.Generator().manual_seed(20260815) + leaves_tensor = torch.randn( + _MAX_WORLD_SIZE, + 257, + dtype=torch.float32, + generator=generator, + ).to(device=device, dtype=dtype) + leaves = list(leaves_tensor.unbind()) + input = _fixed_tree_reference(leaves[start : start + leaves_per_rank]) + expected = _fixed_tree_reference(leaves) + + output = collective.all_reduce(input) + assert torch.equal(output, expected) + + peer_outputs = _all_gather_tensors(output, group, tp_size) + assert all(torch.equal(peer_output, output) for peer_output in peer_outputs) + + baseline = output.clone() + for _ in range(3): + repeated = collective.all_reduce(input) + assert torch.equal(repeated, baseline) + + inplace = input.clone() + returned = collective.all_reduce(inplace, out=inplace) + assert returned is inplace + assert torch.equal(inplace, expected) + dist.barrier() + finally: + dist.destroy_process_group() + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_deterministic_all_reduce_cross_tp_cuda() -> None: + mp.spawn( + _worker, + args=(_find_free_port(),), + nprocs=_MAX_WORLD_SIZE, + join=True, + ) diff --git a/tests/distributed/test_deterministic_reduce_scatter.py b/tests/distributed/test_deterministic_reduce_scatter.py new file mode 100644 index 00000000..3125f6c3 --- /dev/null +++ b/tests/distributed/test_deterministic_reduce_scatter.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import os +import socket +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.distributed import DeterministicCollective + +_MAX_WORLD_SIZE = 8 +_TP_SIZES = (1, 2, 4, 8) +_EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + +pytestmark = [ + pytest.mark.skipif( + _EXTERNAL_WORLD_SIZE != 1, + reason="this cross-TP test owns its worker processes; run pytest directly", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _MAX_WORLD_SIZE, + reason="requires eight visible CUDA GPUs", + ), +] + + +def _fixed_tree_reference(values: list[torch.Tensor]) -> torch.Tensor: + level = values + while len(level) > 1: + level = [level[index] + level[index + 1] for index in range(0, len(level), 2)] + return level[0] + + +def _all_gather_tensors( + value: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, +) -> list[torch.Tensor]: + gathered = [torch.empty_like(value) for _ in range(world_size)] + dist.all_gather(gathered, value, group=group) + return gathered + + +def _worker(rank: int, port: int) -> None: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=_MAX_WORLD_SIZE, + timeout=timedelta(minutes=5), + ) + try: + groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} + for tp_size, group in groups.items(): + if rank < tp_size: + with DeterministicCollective( + group=group, + device=device, + max_size_bytes=1024 * 1024, + ) as collective: + group_rank = dist.get_rank(group=group) + leaves_per_rank = _MAX_WORLD_SIZE // tp_size + start = group_rank * leaves_per_rank + for dtype in (torch.float32, torch.float16, torch.bfloat16): + generator = torch.Generator().manual_seed(20260816) + leaves_tensor = torch.randn( + _MAX_WORLD_SIZE, + _MAX_WORLD_SIZE * 17, + 19, + dtype=torch.float32, + generator=generator, + ).to(device=device, dtype=dtype) + leaves = list(leaves_tensor.unbind()) + input = _fixed_tree_reference(leaves[start : start + leaves_per_rank]) + reduced = _fixed_tree_reference(leaves) + expected = reduced.chunk(tp_size, dim=0)[group_rank] + + output = collective.reduce_scatter(input) + assert torch.equal(output, expected) + + peer_outputs = _all_gather_tensors(output, group, tp_size) + assert torch.equal(torch.cat(peer_outputs, dim=0), reduced) + + baseline = output.clone() + for _ in range(3): + repeated = collective.reduce_scatter(input) + assert torch.equal(repeated, baseline) + + provided = torch.empty_like(expected) + returned = collective.reduce_scatter(input, out=provided) + assert returned is provided + assert torch.equal(provided, expected) + dist.barrier() + finally: + dist.destroy_process_group() + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_deterministic_reduce_scatter_cross_tp_cuda() -> None: + mp.spawn( + _worker, + args=(_find_free_port(),), + nprocs=_MAX_WORLD_SIZE, + join=True, + ) diff --git a/tests/test_alignment_wrapper_interfaces.py b/tests/test_alignment_wrapper_interfaces.py new file mode 100644 index 00000000..c9fbad70 --- /dev/null +++ b/tests/test_alignment_wrapper_interfaces.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Common invocation-surface tests for Attention, FFN, and logprob wrappers.""" + +from __future__ import annotations + +import inspect + +import pytest +import torch + +from rl_engine.alignment.cross_config.operators import selected_logprobs_with_operator +from rl_engine.kernels.logprob_contract import ( + LogprobContract, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.pytorch.attention import AttentionAblationOp +from rl_engine.kernels.ops.pytorch.ffn import Qwen3FFNOp +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import VocabParallelLogprobOp +from rl_engine.kernels.registry import KernelRegistry +from rl_engine.kernels.semantic_registry import OperatorRequirements + + +def _logprob_contract( + *, active: tuple[bool, ...], determinism_scope: str = "cross_tp_bitwise" +) -> LogprobContract: + return LogprobContract( + role="train", + dtype="fp32", + mask=MaskSpec(num_tokens=len(active), active_mask=active), + sharding=ShardingSpec( + tp_rank=0, + tp_world_size=1, + vocab_shard_bounds=((0, 8),), + real_vocab_size=7, + padded_vocab_size=8, + ), + reduction=ReductionSpec(determinism_scope=determinism_scope), + ) + + +@pytest.mark.parametrize( + "wrapper", + [AttentionAblationOp, Qwen3FFNOp, VocabParallelLogprobOp], +) +def test_alignment_wrappers_share_the_deterministic_switch(wrapper): + parameters = inspect.signature(wrapper.__call__).parameters + + assert "deterministic" in parameters + assert parameters["deterministic"].kind is inspect.Parameter.KEYWORD_ONLY + assert parameters["deterministic"].default in {None, True} + + +@pytest.mark.parametrize( + ("wrapper", "backend_id"), + [ + (AttentionAblationOp, "rlkernel.attention.deterministic.v1"), + (Qwen3FFNOp, "rlkernel.ffn.qwen3.deterministic.v1"), + (VocabParallelLogprobOp, "pytorch-vocab-parallel-logp-ws2"), + ], +) +def test_alignment_wrapper_backend_ids_are_stable(wrapper, backend_id): + assert wrapper.backend_id == backend_id + + +@pytest.mark.parametrize("deterministic", [True, False]) +def test_contract_aware_logprob_bridge_accepts_logp_lse_result(deterministic): + active = (True, False, True, True) + contract = _logprob_contract( + active=active, + determinism_scope="cross_tp_bitwise" if deterministic else "fixed_topology", + ) + logits = torch.tensor( + [ + [[1.0, 2.0, 0.0, -1.0, 3.0, 0.5, -0.5, 100.0], [0.0] * 8], + [[2.0, 0.0, 1.0, -2.0, 0.5, 3.0, -1.0, 100.0], [0.0] * 8], + ] + ) + targets = torch.tensor([[4, -100], [5, 0]]) + mask = torch.tensor(active).reshape_as(targets) + + actual = selected_logprobs_with_operator( + VocabParallelLogprobOp(), + logits, + targets, + active_mask=mask, + contract=contract, + num_vocab_tiles=4, + deterministic=deterministic, + ) + + real_logits = logits[..., :7] + safe_targets = targets.masked_fill(~mask, 0) + expected = torch.gather( + torch.log_softmax(real_logits, dim=-1), + -1, + safe_targets.unsqueeze(-1), + ).squeeze(-1) + expected = expected.masked_fill(~mask, 0.0) + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-6) + + +def test_contract_aware_logprob_bridge_rejects_a_different_active_mask(): + contract = _logprob_contract(active=(True, False)) + + with pytest.raises(ValueError, match="active_mask must match"): + selected_logprobs_with_operator( + VocabParallelLogprobOp(), + torch.zeros((2, 8)), + torch.tensor([0, 0]), + active_mask=torch.tensor([True, True]), + contract=contract, + num_vocab_tiles=4, + ) + + +def test_semantic_catalog_exposes_exact_rlkernel_and_native_axes(): + catalog = KernelRegistry().semantic + + for semantic_op, backend_id in ( + ("attention", "rlkernel.attention.deterministic.v1"), + ("ffn", "rlkernel.ffn.qwen3.deterministic.v1"), + ("selected_logprob", "pytorch-vocab-parallel-logp-ws2"), + ): + rlkernel = catalog.backend_descriptor(semantic_op, backend_id) + native = catalog.backend_descriptor(semantic_op, "native") + assert rlkernel is not None + assert rlkernel.determinism_or_alignment_properties["deterministic"] is True + assert rlkernel.is_strictly_observable + assert native is not None + assert native.fallback_policy.value == "runtime_managed" + + +def test_semantic_ffn_and_logprob_wrappers_are_instantiable(): + catalog = KernelRegistry().semantic + session = catalog.session() + + ffn = session.instantiate( + session.resolve( + semantic_op="ffn", + requested_backend="rlkernel.ffn.qwen3.deterministic.v1", + target="training", + requirements=OperatorRequirements(device="cuda", dtype="bfloat16"), + ) + ) + logprob = session.instantiate( + session.resolve( + semantic_op="selected_logprob", + requested_backend="pytorch-vocab-parallel-logp-ws2", + target="training", + requirements=OperatorRequirements(device="cpu", dtype="float32"), + ) + ) + + assert isinstance(ffn, Qwen3FFNOp) + assert isinstance(logprob, VocabParallelLogprobOp) diff --git a/tests/test_det_gemm.py b/tests/test_det_gemm.py index 21181bb8..24f7e546 100644 --- a/tests/test_det_gemm.py +++ b/tests/test_det_gemm.py @@ -38,6 +38,24 @@ def _rand(*shape): return torch.randn(*shape, device=DEV, dtype=torch.bfloat16) +# Matches det_gemm_kernel.cu: mid-split K-tree, FP32 leaf width 32, BF16 internal adds. +_K_TREE_LEAF = 32 + + +def _k_tree_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + a = a.detach().contiguous() + b = b.detach().contiguous() + k = a.shape[1] + + def rec(lo: int, hi: int) -> torch.Tensor: + if hi - lo <= _K_TREE_LEAF: + return (a[:, lo:hi].float() @ b[lo:hi, :].float()).to(dtype=torch.bfloat16) + mid = lo + (hi - lo) // 2 + return rec(lo, mid) + rec(mid, hi) + + return rec(0, k) + + @pytest.mark.parametrize("name,gemm", _BACKENDS) def test_forward_batch_invariance(name, gemm): # A row's output must not change when other rows join the batch. @@ -81,7 +99,10 @@ def test_forward_correctness(name, gemm): M, K, N = 128, 2048, 2048 a, b = _rand(M, K), _rand(K, N) out = gemm(a, b).float() - ref = a.float() @ b.float() + if name == "cuda": + ref = _k_tree_gemm(a, b).float() + else: + ref = a.float() @ b.float() contract = load_contract() thresholds = contract["accuracy"]["default"]["reduction"]["bfloat16"] torch.testing.assert_close(out, ref, atol=thresholds["atol"], rtol=thresholds["rtol"]) @@ -111,6 +132,18 @@ def test_backward_correctness(name, gemm): b = _rand(K, N).requires_grad_(True) g = _rand(M, N) gemm(a, b).backward(g) + if name == "cuda": + da = _k_tree_gemm(g, b.t().contiguous()) + db = _k_tree_gemm(a.detach().t().contiguous(), g) + contract = load_contract() + thresholds = contract["accuracy"]["default"]["reduction"]["bfloat16"] + torch.testing.assert_close( + a.grad.float(), da.float(), atol=thresholds["atol"], rtol=thresholds["rtol"] + ) + torch.testing.assert_close( + b.grad.float(), db.float(), atol=thresholds["atol"], rtol=thresholds["rtol"] + ) + return af = a.detach().float().requires_grad_(True) bf = b.detach().float().requires_grad_(True) (af @ bf).backward(g.float()) diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py new file mode 100644 index 00000000..a961f312 --- /dev/null +++ b/tests/test_logprob_contract.py @@ -0,0 +1,616 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 TP-aware logprob contract and contract-aware dispatch tests (issue #241).""" + +from __future__ import annotations + +import json +from dataclasses import replace + +import pytest + +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDType, + LogprobOutputSpec, + LogprobRole, + MaskMode, + MaskSpec, + ReductionSpec, + ShardingSpec, + TPPlacement, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + +QWEN3_REAL_VOCAB = 151936 +QWEN3_PADDED_VOCAB = 152064 + + +def _even_bounds(padded_vocab: int, tp_world_size: int) -> tuple[tuple[int, int], ...]: + shard = padded_vocab // tp_world_size + return tuple( + (rank * shard, padded_vocab if rank == tp_world_size - 1 else (rank + 1) * shard) + for rank in range(tp_world_size) + ) + + +def _sharding( + *, + tp_rank: int = 0, + tp_world_size: int = 2, + cp_rank: int = 0, + cp_world_size: int = 2, + real_vocab_size: int = QWEN3_REAL_VOCAB, + padded_vocab_size: int = QWEN3_PADDED_VOCAB, + vocab_shard_bounds: tuple[tuple[int, int], ...] | None = None, +) -> ShardingSpec: + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=( + vocab_shard_bounds + if vocab_shard_bounds is not None + else _even_bounds(padded_vocab_size, tp_world_size) + ), + real_vocab_size=real_vocab_size, + padded_vocab_size=padded_vocab_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + ) + + +def _mask( + *, + num_tokens: int = 8, + active_mask: tuple[bool, ...] | None = None, + ignore_index: int = -100, +) -> MaskSpec: + return MaskSpec( + num_tokens=num_tokens, + active_mask=( + active_mask + if active_mask is not None + else (False, False, True, True, True, True, True, False) + ), + ignore_index=ignore_index, + ) + + +def _contract( + *, + role: str = "train", + dtype: str = "bf16", + mask: MaskSpec | None = None, + sharding: ShardingSpec | None = None, + reduction: ReductionSpec | None = None, +) -> LogprobContract: + return LogprobContract( + role=role, + dtype=dtype, + mask=mask if mask is not None else _mask(), + sharding=sharding if sharding is not None else _sharding(), + reduction=reduction if reduction is not None else ReductionSpec(), + ) + + +def _declared_tp_backend() -> LogprobBackendCapability: + return LogprobBackendCapability( + backend_id="test-deterministic-tp-logprob", + roles=frozenset({LogprobRole.TRAIN, LogprobRole.INFER}), + dtypes=frozenset({LogprobDType.BF16}), + tp_world_sizes=(1, 2, 4), + cp_world_sizes=None, + supports_vocab_padding=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), + exports_vocab_lse=True, + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", + ) + + +def test_qwen3_tp2_bf16_contract_is_representable_and_serializable(): + contract = _contract() + + assert contract.sharding.tp_world_size == 2 + assert contract.sharding.cp_world_size == 2 + assert contract.sharding.local_vocab_start == 0 + assert contract.sharding.local_vocab_end == QWEN3_PADDED_VOCAB // 2 + assert contract.sharding.local_vocab_size == QWEN3_PADDED_VOCAB // 2 + assert contract.mask.active_token_count == 5 + assert contract.reduction.acc_dtype is LogprobDType.FP32 + assert contract.to_dict()["reduction"] == { + "merge": "max_sumexp", + "merge_axis": "tp_vocab", + "acc_dtype": "fp32", + "order": "global_vocab_shard_index", + "transport": "all_gather", + "downcast_at": "final_write", + "engine": "in_op_reference", + "determinism_scope": "cross_tp_bitwise", + "cp_is_merge_axis": False, + } + json.dumps(contract.to_dict()) + + +@pytest.mark.parametrize("tp_world_size", [1, 2, 4]) +def test_pr4_sweep_tp_degrees_are_representable(tp_world_size): + sharding = _sharding(tp_world_size=tp_world_size, cp_world_size=1) + + assert len(sharding.vocab_shard_bounds) == tp_world_size + assert sharding.vocab_shard_bounds[-1][1] == QWEN3_PADDED_VOCAB + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == tp_world_size - 1 + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("tp_rank", 2, "tp_rank=2"), + ("cp_rank", 2, "cp_rank=2"), + ("real_vocab_size", 0, "positive integer"), + ("padded_vocab_size", QWEN3_REAL_VOCAB - 1, "must not be smaller"), + ], +) +def test_invalid_rank_and_vocab_metadata_fail_loudly(field, value, message): + values = { + "tp_rank": 0, + "tp_world_size": 2, + "cp_rank": 0, + "cp_world_size": 2, + "real_vocab_size": QWEN3_REAL_VOCAB, + "padded_vocab_size": QWEN3_PADDED_VOCAB, + "vocab_shard_bounds": _even_bounds(QWEN3_PADDED_VOCAB, 2), + } + values[field] = value + + with pytest.raises(LogprobContractError, match=message): + ShardingSpec(**values) + + +@pytest.mark.parametrize( + ("bounds", "message"), + [ + ((), "one \\(start, end\\) pair per TP rank"), + (((0, 76032),), "one \\(start, end\\) pair per TP rank"), + (((0, 76032), (76032, 76032)), "end > start"), + (((0, 76000), (76032, 152064)), "contiguous"), + (((0, 76064), (76032, 152064)), "contiguous"), + (((0, 76032), (76032, 152000)), "cover padded_vocab_size exactly"), + ], +) +def test_incomplete_or_overlapping_vocab_shard_bounds_fail_loudly(bounds, message): + with pytest.raises(LogprobContractError, match=message): + _sharding(vocab_shard_bounds=bounds) + + +def test_owner_rank_is_unique_and_rejects_out_of_real_vocab_targets(): + sharding = _sharding() + + assert sharding.owner_rank(0) == 0 + assert sharding.owner_rank(QWEN3_PADDED_VOCAB // 2 - 1) == 0 + assert sharding.owner_rank(QWEN3_PADDED_VOCAB // 2) == 1 + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == 1 + + with pytest.raises(LogprobContractError, match="outside the real vocabulary"): + sharding.owner_rank(-1) + with pytest.raises(LogprobContractError, match="outside the real vocabulary"): + sharding.owner_rank(QWEN3_REAL_VOCAB) + + +def test_active_token_mask_metadata_is_validated(): + with pytest.raises(LogprobContractError, match="one entry per token"): + _mask(num_tokens=4) + + with pytest.raises(LogprobContractError, match="must be a bool"): + MaskSpec(num_tokens=2, active_mask=(True, 1)) + + all_inactive = _mask(num_tokens=3, active_mask=(False, False, False)) + assert all_inactive.active_token_count == 0 + + +def test_reduction_requires_fp32_accumulation_and_known_semantics(): + with pytest.raises(LogprobContractError, match="must be fp32"): + ReductionSpec(acc_dtype="bf16") + + with pytest.raises(LogprobContractError, match="merge must be one of"): + ReductionSpec(merge="lse_average") + + with pytest.raises(LogprobContractError, match="transport must be one of"): + ReductionSpec(transport="all_reduce") + + +def test_contract_component_types_and_lse_export_are_enforced(): + with pytest.raises(LogprobContractError, match="mask must be a MaskSpec"): + LogprobContract( + role="train", + dtype="bf16", + mask=None, + sharding=_sharding(), + reduction=ReductionSpec(), + ) + + with pytest.raises(LogprobContractError, match="export_lse must be True"): + replace(_contract(), export_lse=False) + + +def test_ignore_index_must_not_collide_with_the_real_vocabulary(): + with pytest.raises(LogprobContractError, match="must not collide"): + _contract(mask=_mask(ignore_index=5)) + + padding_column = QWEN3_REAL_VOCAB + 1 + contract = _contract(mask=_mask(ignore_index=padding_column)) + assert contract.mask.ignore_index == padding_column + + +def _restrict_to_ws1_candidates(registry: KernelRegistry) -> None: + """Drop the #241 PR3 vocab-parallel reference so only WS1 backends remain.""" + + platform = registry._platform() + registry._logprob_candidates[platform] = [ + backend + for backend in registry._logprob_candidates[platform] + if backend is not OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP + ] + + +def test_current_ws1_backend_rejects_strict_tp_contract_without_fallback(): + registry = KernelRegistry() + _restrict_to_ws1_candidates(registry) + + with pytest.raises(RuntimeError) as exc_info: + registry.get_logprob_op(_contract()) + + message = str(exc_info.value) + assert "TP=2 is unsupported" in message + assert "vocab-domain LSE export is unsupported" in message + assert "determinism_scope=cross_tp_bitwise is unsupported" in message + assert "padded-vs-real vocab masking is unsupported" in message + + +def test_current_ws1_backend_rejects_padded_vocab_even_at_tp1(): + registry = KernelRegistry() + _restrict_to_ws1_candidates(registry) + contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + + with pytest.raises(RuntimeError) as exc_info: + registry.get_logprob_op(contract) + + message = str(exc_info.value) + assert "TP=1 is unsupported" not in message + assert "padded-vs-real vocab masking is unsupported" in message + + +def test_ws1_rejections_recorded_when_vocab_parallel_reference_resolves(): + """The WS1 backends still reject strict contracts; they are skipped with + recorded reasons while dispatch resolves the #241 PR3 reference.""" + + registry = KernelRegistry() + platform = registry._platform() + # Order the WS1 backends ahead of the reference so their rejections are + # exercised on the way to a successful resolution. + candidates = registry._logprob_candidates[platform] + candidates.remove(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) + candidates.append(OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP) + + result = registry.get_logprob_op(_contract()) + assert result.capability.backend_id == "pytorch-vocab-parallel-logp-ws2" + assert result.provenance["fallback"] is True + rejections = " | ".join(result.provenance["prior_rejections"]) + assert "vocab-domain LSE export is unsupported" in rejections + + +def test_undeclared_backend_capability_is_never_selected(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [OpBackend.PYTORCH_NATIVE] + + with pytest.raises(RuntimeError, match="no LogprobBackendCapability declared"): + registry.get_logprob_op(_contract()) + + +def test_declared_compatible_backend_resolves_and_records_provenance(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="reference") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert result.provenance["requested_backend"] == "reference" + assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" + assert result.provenance["fallback"] is False + assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 + assert result.provenance["contract"]["sharding"]["real_vocab_size"] == QWEN3_REAL_VOCAB + assert result.provenance["contract"]["reduction"]["cp_is_merge_axis"] is False + json.dumps(result.provenance) + + +def test_requested_stable_backend_id_is_enforced(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"): + registry.get_logprob_op(_contract(), requested_backend="another-backend") + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" + + +def test_cp_is_a_non_merge_axis_and_cp_agnostic_backends_accept_any_cp_degree(): + capability = _declared_tp_backend() + cp2_contract = _contract(sharding=_sharding(cp_world_size=2, cp_rank=1)) + + assert capability.incompatibilities(cp2_contract) == () + + cp_restricted = replace(capability, cp_world_sizes=(1,)) + assert cp_restricted.incompatibilities(cp2_contract) == ("CP=2 is unsupported",) + + +def test_inactive_tokens_require_explicit_active_mask_support(): + capability = replace(_declared_tp_backend(), mask_modes=frozenset({MaskMode.IGNORE_INDEX})) + contract = _contract() + + assert "explicit active-token masking is unsupported" in ( + capability.incompatibilities(contract) + ) + + fully_active = _contract(mask=_mask(num_tokens=3, active_mask=(True, True, True))) + assert capability.incompatibilities(fully_active) == () + + +def test_backend_id_must_not_shadow_a_reserved_policy_keyword(): + with pytest.raises(LogprobContractError, match="reserved dispatch policy keyword"): + replace(_declared_tp_backend(), backend_id="Deterministic") + + +def test_default_auto_policy_resolves_any_compatible_implementation_kind(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), implementation_kind="reference"), + platform=platform, + ) + + result = registry.get_logprob_op(_contract()) + + assert result.provenance["requested_backend"] == "auto" + assert result.capability.implementation_kind == "reference" + + +def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="REFERENCE") + assert result.capability.backend_id == "test-deterministic-tp-logprob" + + with pytest.raises(RuntimeError, match="does not match requested_backend"): + registry.get_logprob_op(_contract(), requested_backend="Test-Deterministic-TP-Logprob") + + +def test_policy_only_skips_are_not_reported_as_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="other-compatible-backend"), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + + assert result.provenance["fallback"] is False + assert len(result.provenance["prior_rejections"]) == 1 + + +def test_capability_rejections_are_reported_as_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,)), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract()) + + assert result.provenance["fallback"] is True + assert "TP=2 is unsupported" in result.provenance["prior_rejections"][0] + + +def test_ws2_candidate_list_is_decoupled_from_the_legacy_priority_map(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform].insert(0, OpBackend.PYTORCH_NATIVE) + + legacy = registry._priority_map[platform]["batch_invariant_logp"] + assert OpBackend.PYTORCH_NATIVE not in legacy + + legacy.insert(0, OpBackend.PYTORCH_GEMM) + assert OpBackend.PYTORCH_GEMM not in registry._logprob_candidates[platform] + + +def test_register_logprob_backend_is_the_public_registration_seam(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + capability = _declared_tp_backend() + + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, capability, platform=platform + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(capability, backend_id="replacement-backend"), + platform=platform, + ) + + assert registry._logprob_candidates[platform] == [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + result = registry.get_logprob_op(_contract()) + assert result.capability.backend_id == "replacement-backend" + + with pytest.raises(LogprobContractError, match="capability must be"): + registry.register_logprob_backend(OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, None) + + +def test_backend_id_whitespace_is_normalized_for_dispatch(): + capability = replace(_declared_tp_backend(), backend_id=" padded-id ") + assert capability.backend_id == "padded-id" + + +def test_capabilities_are_scoped_per_platform(): + registry = KernelRegistry() + platform = registry._platform() + other = "rocm" if platform != "rocm" else "cpu" + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="other-platform-backend"), + platform=other, + ) + + result = registry.get_logprob_op(_contract()) + + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert ( + registry._logprob_capabilities[other][OpBackend.PYTORCH_BATCH_INVARIANT_LOGP].backend_id + == "other-platform-backend" + ) + + +def test_register_logprob_backend_rejects_unknown_platform(): + registry = KernelRegistry() + + with pytest.raises(LogprobContractError, match="unsupported platform"): + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + _declared_tp_backend(), + platform="cuda-typo", + ) + + +def test_non_iterable_roles_and_dtypes_raise_contract_errors(): + with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): + replace(_declared_tp_backend(), roles=None) + + with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): + replace(_declared_tp_backend(), dtypes=42) + + +def test_requested_deterministic_policy_is_a_loud_error(): + registry = KernelRegistry() + + with pytest.raises(LogprobContractError, match="determinism_scope"): + registry.get_logprob_op(_contract(), requested_backend="deterministic") + + +def test_determinism_scope_is_part_of_the_typed_contract(): + fixed_only = replace( + _declared_tp_backend(), + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + ) + + assert "determinism_scope=cross_tp_bitwise is unsupported" in ( + fixed_only.incompatibilities(_contract()) + ) + + relaxed = _contract(reduction=ReductionSpec(determinism_scope="fixed_topology")) + assert fixed_only.incompatibilities(relaxed) == () + + +def test_policy_filtered_candidates_never_count_toward_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,)), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + + assert result.provenance["fallback"] is False + assert len(result.provenance["prior_rejections"]) == 1 + + +def test_output_spec_is_pinned_to_fp32_replicated(): + with pytest.raises(LogprobContractError, match="must be fp32"): + LogprobOutputSpec(selected_logp_dtype="bf16") + with pytest.raises(LogprobContractError, match="must be fp32"): + LogprobOutputSpec(lse_dtype="bf16") + + assert LogprobOutputSpec().tp_placement is TPPlacement.REPLICATED + assert _contract().to_dict()["output"] == { + "selected_logp_dtype": "fp32", + "lse_dtype": "fp32", + "tp_placement": "replicated", + } + + +def test_cross_rank_fingerprint_is_rank_independent_and_content_sensitive(): + rank0 = _contract(sharding=_sharding(tp_rank=0)) + rank1 = _contract(sharding=_sharding(tp_rank=1, cp_rank=1)) + + assert rank0.cross_rank_fingerprint() == rank1.cross_rank_fingerprint() + + different_mask = _contract( + mask=_mask(active_mask=(True, True, True, True, True, True, True, False)) + ) + assert rank0.cross_rank_fingerprint() != different_mask.cross_rank_fingerprint() + + +def test_provenance_records_the_active_mask_digest(): + provenance_mask = _contract().to_dict()["mask"] + + assert provenance_mask["active_mask_sha256"] == _mask().active_mask_sha256 + assert len(provenance_mask["active_mask_sha256"]) == 64 + + same_count_different_mask = _mask( + active_mask=(True, True, True, True, True, False, False, False) + ) + assert same_count_different_mask.active_token_count == _mask().active_token_count + assert same_count_different_mask.active_mask_sha256 != _mask().active_mask_sha256 + + +def test_padding_only_shard_is_constructible_for_the_identity_partial(): + sharding = _sharding( + vocab_shard_bounds=((0, QWEN3_REAL_VOCAB), (QWEN3_REAL_VOCAB, QWEN3_PADDED_VOCAB)), + ) + + assert sharding.local_vocab_start == 0 + assert sharding.vocab_shard_bounds[1] == (QWEN3_REAL_VOCAB, QWEN3_PADDED_VOCAB) + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == 0 diff --git a/tests/test_qwen_ffn.py b/tests/test_qwen_ffn.py new file mode 100644 index 00000000..9b2b1d21 --- /dev/null +++ b/tests/test_qwen_ffn.py @@ -0,0 +1,1112 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the deterministic Qwen3 dense FFN. + +Covers single-GPU correctness, token boundaries, Qwen3-8B shapes, TP/CP/SP +bitwise alignment, and DeterministicCollective cache lifetime. +""" + +from __future__ import annotations + +import queue +import tempfile +import traceback +from datetime import timedelta +from pathlib import Path + +import pytest +import torch +import torch.multiprocessing as mp +import torch.nn.functional as F + +import rl_engine.kernels.ops.pytorch.ffn.ffn as ffn_module +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.pytorch.ffn.ffn import ( + QWEN3_8B_HIDDEN_SIZE, + QWEN3_8B_INTERMEDIATE_SIZE, + qwen3_ffn, +) + +_REQUIRED_SYMBOLS = ( + "det_gemm_fwd", + "det_gemm_db", + "swiglu_forward", + "swiglu_backward", +) +_HIDDEN = 64 +_INTERMEDIATE = 512 +_TOPOLOGY_TOKENS = 256 +_CP_TOKEN_COUNTS = (8, 32, 64, 96, 128, 256) +_TOKEN_BOUNDARY_COUNTS = (8, 31, 32, 33, 64, 96, 128) +_WORLD2_CONFIGS = ( + ("tp2_sp", 2, 1, True, _TOPOLOGY_TOKENS), + ("cp2", 1, 2, False, _TOPOLOGY_TOKENS), + *((f"cp2_T{token_count}", 1, 2, False, token_count) for token_count in _CP_TOKEN_COUNTS), +) +_WORLD4_CONFIGS = ( + ("tp4", 4, 1, False, _TOPOLOGY_TOKENS), + ("cp4", 1, 4, False, _TOPOLOGY_TOKENS), + ("tp2_cp2", 2, 2, False, _TOPOLOGY_TOKENS), + ("tp2_cp2_sp", 2, 2, True, _TOPOLOGY_TOKENS), + *((f"cp4_T{token_count}", 1, 4, False, token_count) for token_count in _CP_TOKEN_COUNTS), +) +_WORLD8_WORLD_GROUP_CONFIGS = ( + ("tp8", 8, 1, False, _TOPOLOGY_TOKENS), + ("tp8_sp", 8, 1, True, _TOPOLOGY_TOKENS), + ("cp8", 1, 8, False, _TOPOLOGY_TOKENS), + *((f"cp8_T{token_count}", 1, 8, False, token_count) for token_count in _CP_TOKEN_COUNTS), +) +_WORLD8_TP2_CP4_CONFIGS = (("tp2_cp4", 2, 4, False, _TOPOLOGY_TOKENS),) +_WORLD8_TP4_CP2_CONFIGS = ( + ("tp4_cp2", 4, 2, False, _TOPOLOGY_TOKENS), + ("tp4_cp2_sp", 4, 2, True, _TOPOLOGY_TOKENS), +) + + +def _has_sm90_ffn_devices(count: int) -> bool: + return ( + _EXT_AVAILABLE + and torch.distributed.is_available() + and torch.distributed.is_nccl_available() + and torch.cuda.device_count() >= count + and all(torch.cuda.get_device_capability(index)[0] == 9 for index in range(count)) + and all(hasattr(_C, name) for name in _REQUIRED_SYMBOLS) + ) + + +def _has_sm90_ffn() -> bool: + return ( + torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] == 9 + and _EXT_AVAILABLE + and all(hasattr(_C, name) for name in _REQUIRED_SYMBOLS) + ) + + +requires_cuda_ffn = pytest.mark.skipif( + not _has_sm90_ffn(), + reason="FFN optimized-path validation requires SM90 and the GEMM/SwiGLU extension", +) + + +class _TorchKernelStub: + def __init__(self) -> None: + self.calls: list[str] = [] + + def det_gemm_fwd(self, a, b): + self.calls.append("det_gemm_fwd") + return a @ b + + def det_gemm_db(self, a, grad_output): + self.calls.append("det_gemm_db") + return a.t().contiguous() @ grad_output + + def swiglu_forward(self, gate, up): + self.calls.append("swiglu_forward") + return gate * torch.sigmoid(gate) * up + + def swiglu_backward(self, grad_output, gate, up): + self.calls.append("swiglu_backward") + sigmoid = torch.sigmoid(gate) + grad_gate = grad_output * up * sigmoid * (1.0 + gate * (1.0 - sigmoid)) + grad_up = grad_output * gate * sigmoid + return grad_gate, grad_up + + +def _reference(hidden_states, gate_weight, up_weight, down_weight): + gate = hidden_states @ gate_weight.t() + up = hidden_states @ up_weight.t() + activated = F.silu(gate) * up + return (activated @ down_weight.t()), gate, up, activated + + +def _randn(shape, *, seed, device="cpu", dtype=torch.float32): + generator = torch.Generator(device="cpu").manual_seed(seed) + value = torch.randn(*shape, generator=generator, dtype=torch.float32) * 0.02 + return value.to(device=device, dtype=dtype) + + +def _close_ffn_collectives() -> None: + for collective in list(ffn_module._COLLECTIVES.values()): + collective.close() + ffn_module._COLLECTIVES.clear() + + +def _shard_ranges( + rank: int, + *, + tp_size: int, + cp_size: int, + sequence_parallel: bool, + token_count: int, + intermediate_size: int, +) -> tuple[int, int, int, int]: + tp_rank = rank % tp_size + cp_rank = rank // tp_size + cp_tokens = token_count // cp_size + local_tokens = cp_tokens // tp_size if sequence_parallel else cp_tokens + token_start = cp_rank * cp_tokens + if sequence_parallel: + token_start += tp_rank * local_tokens + token_end = token_start + local_tokens + local_i = intermediate_size // tp_size + feat_start = tp_rank * local_i + feat_end = feat_start + local_i + return token_start, token_end, feat_start, feat_end + + +def _spawn_nccl_workers(worker, world_size: int, worker_args=(), *, timeout: int = 180) -> None: + if not _has_sm90_ffn_devices(world_size): + pytest.skip(f"requires {world_size} SM90 GPUs, NCCL, and the GEMM/SwiGLU extension") + + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmpdir: + init_method = (Path(tmpdir) / "nccl_init").as_uri() + result_queue = ctx.Queue() + processes = [ + ctx.Process( + target=worker, + args=(rank, world_size, init_method, result_queue, *worker_args), + ) + for rank in range(world_size) + ] + for process in processes: + process.start() + results = [] + try: + for _ in processes: + results.append(result_queue.get(timeout=timeout)) + except queue.Empty: + for process in processes: + if process.is_alive(): + process.terminate() + pytest.fail(f"timed out waiting for {world_size} FFN workers") + finally: + for process in processes: + process.join(timeout=10) + if process.is_alive(): + process.terminate() + + for result in sorted(results, key=lambda item: item["rank"]): + assert result["ok"], result.get("traceback") or result.get("failures") + for process in processes: + assert process.exitcode == 0 + + +def _distributed_ffn_backward_nccl_worker( + rank, + world_size, + init_method, + result_queue, + cp_size, + sequence_parallel, +): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + + tp_size = world_size // cp_size + tp_groups = [ + dist.new_group(list(range(cp_rank * tp_size, (cp_rank + 1) * tp_size))) + for cp_rank in range(cp_size) + ] + cp_groups = [ + dist.new_group([cp_rank * tp_size + tp_rank for cp_rank in range(cp_size)]) + for tp_rank in range(tp_size) + ] + tp_rank = rank % tp_size + cp_rank = rank // tp_size + tp_group = tp_groups[cp_rank] + cp_group = cp_groups[tp_rank] if cp_size > 1 else None + + token_count, hidden_size, intermediate_size = 8, 64, 128 + token_start, token_end, feature_start, feature_end = _shard_ranges( + rank, + tp_size=tp_size, + cp_size=cp_size, + sequence_parallel=sequence_parallel, + token_count=token_count, + intermediate_size=intermediate_size, + ) + local_tokens = token_end - token_start + + device = torch.device("cuda", rank) + rmsnorm_output = _randn( + (token_count, hidden_size), seed=40, device=device, dtype=torch.bfloat16 + ) + gate_weight = _randn( + (intermediate_size, hidden_size), + seed=41, + device=device, + dtype=torch.bfloat16, + ) + up_weight = _randn( + (intermediate_size, hidden_size), + seed=42, + device=device, + dtype=torch.bfloat16, + ) + down_weight = _randn( + (hidden_size, intermediate_size), + seed=43, + device=device, + dtype=torch.bfloat16, + ) + grad_output = _randn( + (token_count, hidden_size), seed=44, device=device, dtype=torch.bfloat16 + ) + + reference_inputs = [ + value.detach().float().requires_grad_(True) + for value in (rmsnorm_output, gate_weight, up_weight, down_weight) + ] + reference_output, _, _, _ = _reference(*reference_inputs) + reference_output.backward(grad_output.float()) + + local_grad_output = grad_output[token_start:token_end].contiguous() + actual_inputs = [ + value.detach().clone().requires_grad_(True) + for value in ( + rmsnorm_output[token_start:token_end].contiguous(), + gate_weight[feature_start:feature_end].contiguous(), + up_weight[feature_start:feature_end].contiguous(), + down_weight[:, feature_start:feature_end].contiguous(), + ) + ] + actual_output = qwen3_ffn( + *actual_inputs, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + actual_output.backward(local_grad_output) + + expected_grads = ( + reference_inputs[0].grad[token_start:token_end], + reference_inputs[1].grad[feature_start:feature_end], + reference_inputs[2].grad[feature_start:feature_end], + reference_inputs[3].grad[:, feature_start:feature_end], + ) + torch.testing.assert_close( + actual_output.float(), + reference_output[token_start:token_end].detach(), + atol=5e-2, + rtol=2e-2, + ) + for actual, expected in zip(actual_inputs, expected_grads, strict=True): + torch.testing.assert_close( + actual.grad.float(), + expected, + atol=5e-2, + rtol=2e-2, + ) + + slice_size = max(1, local_tokens // 2) + slice_start = (local_tokens - slice_size) // 2 + slice_end = slice_start + slice_size + slice_inputs = [ + value.detach().clone().requires_grad_(True) + for value in ( + actual_inputs[0][slice_start:slice_end].contiguous(), + actual_inputs[1], + actual_inputs[2], + actual_inputs[3], + ) + ] + slice_output = qwen3_ffn( + *slice_inputs, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + slice_output.backward(local_grad_output[slice_start:slice_end]) + + assert torch.equal( + slice_output, + actual_output[slice_start:slice_end], + ), "FFN output changed with the local token batch size" + assert torch.equal( + slice_inputs[0].grad, + actual_inputs[0].grad[slice_start:slice_end], + ), "FFN input gradient changed with the local token batch size" + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _tp1_vs_tpn_train_infer_worker(rank, world_size, init_method, result_queue, expect_match): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + device = torch.device("cuda", rank) + token_count, hidden_size, intermediate_size = 16, 64, 256 + hidden = _randn((token_count, hidden_size), seed=50, device=device, dtype=torch.bfloat16) + gate_weight = _randn( + (intermediate_size, hidden_size), + seed=51, + device=device, + dtype=torch.bfloat16, + ) + up_weight = _randn( + (intermediate_size, hidden_size), + seed=52, + device=device, + dtype=torch.bfloat16, + ) + down_weight = _randn( + (hidden_size, intermediate_size), + seed=53, + device=device, + dtype=torch.bfloat16, + ) + grad_output = _randn( + (token_count, hidden_size), seed=54, device=device, dtype=torch.bfloat16 + ) + + with torch.no_grad(): + infer_tp1 = qwen3_ffn(hidden, gate_weight, up_weight, down_weight) + + tp1_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + train_tp1 = qwen3_ffn(*tp1_inputs) + train_tp1.backward(grad_output) + assert torch.equal(infer_tp1, train_tp1.detach()), "TP=1 train/infer forward mismatch" + + local_i = intermediate_size // world_size + feat_start = rank * local_i + feat_end = feat_start + local_i + shard = ( + hidden, + gate_weight[feat_start:feat_end].contiguous(), + up_weight[feat_start:feat_end].contiguous(), + down_weight[:, feat_start:feat_end].contiguous(), + ) + with torch.no_grad(): + infer_tpn = qwen3_ffn(*shard, tp_group=dist.group.WORLD) + + tpn_inputs = [value.detach().clone().requires_grad_(True) for value in shard] + train_tpn = qwen3_ffn(*tpn_inputs, tp_group=dist.group.WORLD) + train_tpn.backward(grad_output) + assert torch.equal( + infer_tpn, train_tpn.detach() + ), f"TP={world_size} train/infer forward mismatch" + + infer_match = torch.equal(infer_tp1, infer_tpn) + train_match = torch.equal(train_tp1.detach(), train_tpn.detach()) + hidden_match = torch.equal(tp1_inputs[0].grad, tpn_inputs[0].grad) + if expect_match: + assert infer_match, f"TP=1 vs TP={world_size} infer forward mismatch" + assert train_match, f"TP=1 vs TP={world_size} train forward mismatch" + assert hidden_match, f"TP=1 vs TP={world_size} hidden grad mismatch" + else: + assert not infer_match, f"TP=1 vs TP={world_size} infer forward unexpectedly matched" + assert not train_match, f"TP=1 vs TP={world_size} train forward unexpectedly matched" + assert not hidden_match, f"TP=1 vs TP={world_size} hidden grad unexpectedly matched" + + assert torch.equal( + tp1_inputs[1].grad[feat_start:feat_end], tpn_inputs[1].grad + ), f"TP=1 vs TP={world_size} gate weight grad mismatch" + assert torch.equal( + tp1_inputs[2].grad[feat_start:feat_end], tpn_inputs[2].grad + ), f"TP=1 vs TP={world_size} up weight grad mismatch" + assert torch.equal( + tp1_inputs[3].grad[:, feat_start:feat_end], tpn_inputs[3].grad + ), f"TP=1 vs TP={world_size} down weight grad mismatch" + + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _make_topology_inputs(token_count, device): + hidden = _randn((token_count, _HIDDEN), seed=60, device=device, dtype=torch.bfloat16) + gate = _randn((_INTERMEDIATE, _HIDDEN), seed=61, device=device, dtype=torch.bfloat16) + up = _randn((_INTERMEDIATE, _HIDDEN), seed=62, device=device, dtype=torch.bfloat16) + down = _randn((_HIDDEN, _INTERMEDIATE), seed=63, device=device, dtype=torch.bfloat16) + grad = _randn((token_count, _HIDDEN), seed=64, device=device, dtype=torch.bfloat16) + return hidden, gate, up, down, grad + + +def _canonical(hidden, gate, up, down, grad): + with torch.no_grad(): + infer = qwen3_ffn(hidden, gate, up, down) + inputs = [value.detach().clone().requires_grad_(True) for value in (hidden, gate, up, down)] + train = qwen3_ffn(*inputs) + train.backward(grad) + return infer, train, inputs + + +def _mesh_groups(dist, tp_size, cp_size): + world_size = dist.get_world_size() + if tp_size == world_size and cp_size == 1: + return [dist.group.WORLD], [] + if cp_size == world_size and tp_size == 1: + return [], [dist.group.WORLD] + tp_groups = [] + if tp_size > 1: + for cp_rank in range(cp_size): + ranks = list(range(cp_rank * tp_size, (cp_rank + 1) * tp_size)) + tp_groups.append(dist.new_group(ranks)) + cp_groups = [] + if cp_size > 1: + for tp_rank in range(tp_size): + ranks = [cp_rank * tp_size + tp_rank for cp_rank in range(cp_size)] + cp_groups.append(dist.new_group(ranks)) + return tp_groups, cp_groups + + +def _run_topology_config( + rank, + dist, + meshes, + *, + name, + tp_size, + cp_size, + sequence_parallel, + hidden, + gate, + up, + down, + grad, + infer_ref, + train_ref, + ref_inputs, +): + key = (tp_size, cp_size) + if key not in meshes: + meshes[key] = _mesh_groups(dist, tp_size, cp_size) + tp_groups, cp_groups = meshes[key] + tp_rank = rank % tp_size + cp_rank = rank // tp_size + tp_group = tp_groups[cp_rank] if tp_size > 1 else None + cp_group = cp_groups[tp_rank] if cp_size > 1 else None + token_start, token_end, feat_start, feat_end = _shard_ranges( + rank, + tp_size=tp_size, + cp_size=cp_size, + sequence_parallel=sequence_parallel, + token_count=hidden.size(0), + intermediate_size=_INTERMEDIATE, + ) + shard = ( + hidden[token_start:token_end].contiguous(), + gate[feat_start:feat_end].contiguous(), + up[feat_start:feat_end].contiguous(), + down[:, feat_start:feat_end].contiguous(), + ) + with torch.no_grad(): + infer = qwen3_ffn( + *shard, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + inputs = [value.detach().clone().requires_grad_(True) for value in shard] + train = qwen3_ffn( + *inputs, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + train.backward(grad[token_start:token_end].contiguous()) + + assert torch.equal(infer, train.detach()), f"{name}: train/infer forward mismatch" + assert torch.equal( + infer, infer_ref[token_start:token_end] + ), f"{name}: infer forward mismatch vs TP=1/CP=1" + assert torch.equal( + train.detach(), train_ref.detach()[token_start:token_end] + ), f"{name}: train forward mismatch vs TP=1/CP=1" + assert torch.equal( + inputs[0].grad, ref_inputs[0].grad[token_start:token_end] + ), f"{name}: hidden grad mismatch vs TP=1/CP=1" + + weight_checks = ( + (1, ref_inputs[1].grad[feat_start:feat_end], "gate"), + (2, ref_inputs[2].grad[feat_start:feat_end], "up"), + (3, ref_inputs[3].grad[:, feat_start:feat_end], "down"), + ) + for index, expected, label in weight_checks: + assert torch.equal( + inputs[index].grad, expected + ), f"{name}: {label} weight grad mismatch vs TP=1/CP=1" + + +def _topology_worker(rank, world_size, init_method, result_queue, configs): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(minutes=5), + device_id=torch.device("cuda", rank), + ) + device = torch.device("cuda", rank) + meshes = {} + canonical = {} + for name, tp_size, cp_size, sequence_parallel, token_count in configs: + if token_count not in canonical: + tensors = _make_topology_inputs(token_count, device) + canonical[token_count] = (*tensors, *_canonical(*tensors)) + hidden, gate, up, down, grad, infer_ref, train_ref, ref_inputs = canonical[token_count] + _run_topology_config( + rank, + dist, + meshes, + name=name, + tp_size=tp_size, + cp_size=cp_size, + sequence_parallel=sequence_parallel, + hidden=hidden, + gate=gate, + up=up, + down=down, + grad=grad, + infer_ref=infer_ref, + train_ref=train_ref, + ref_inputs=ref_inputs, + ) + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _ffn_tensors(token_count, device, seed, *, hidden=_HIDDEN, intermediate=_INTERMEDIATE): + rmsnorm = _randn((token_count, hidden), seed=seed, device=device, dtype=torch.bfloat16) + gate = _randn((intermediate, hidden), seed=seed + 1, device=device, dtype=torch.bfloat16) + up = _randn((intermediate, hidden), seed=seed + 2, device=device, dtype=torch.bfloat16) + down = _randn((hidden, intermediate), seed=seed + 3, device=device, dtype=torch.bfloat16) + return rmsnorm, gate, up, down + + +def _cache_worker(rank, world_size, init_method, result_queue): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + device = torch.device("cuda", rank) + ffn_module._COLLECTIVE_MIN_CAPACITY_BYTES = 64 + _close_ffn_collectives() + + small = _ffn_tensors(8, device, seed=100, intermediate=128) + first = qwen3_ffn(*small, tp_group=dist.group.WORLD) + assert len(ffn_module._COLLECTIVES) == 1 + ((cache_key, first_collective),) = ffn_module._COLLECTIVES.items() + first_handle = first_collective._handle + first_capacity = first_collective.max_size_bytes + assert first_handle != 0 + + repeated = qwen3_ffn(*small, tp_group=dist.group.WORLD) + assert torch.equal(first, repeated) + assert len(ffn_module._COLLECTIVES) == 1 + assert ffn_module._COLLECTIVES[cache_key] is first_collective + assert first_collective._handle == first_handle + + large = _ffn_tensors(256, device, seed=110, intermediate=128) + grown = qwen3_ffn(*large, tp_group=dist.group.WORLD) + assert grown.shape[0] == 256 + assert len(ffn_module._COLLECTIVES) == 1 + grown_collective = next(iter(ffn_module._COLLECTIVES.values())) + assert grown_collective is not first_collective + assert first_collective._handle == 0 + assert grown_collective.max_size_bytes > first_capacity + assert grown_collective._handle != 0 + + _close_ffn_collectives() + assert ffn_module._COLLECTIVES == {} + recreated = qwen3_ffn(*small, tp_group=dist.group.WORLD) + assert torch.equal(recreated, first) + assert len(ffn_module._COLLECTIVES) == 1 + recreated_collective = next(iter(ffn_module._COLLECTIVES.values())) + assert recreated_collective is not grown_collective + assert recreated_collective._handle != 0 + + rebuilt_group = dist.new_group(ranks=[0, 1]) + rebuilt = qwen3_ffn(*small, tp_group=rebuilt_group) + assert torch.equal(rebuilt, first) + assert len(ffn_module._COLLECTIVES) == 2 + + _close_ffn_collectives() + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _uneven_sp_worker(rank, world_size, init_method, result_queue): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + device = torch.device("cuda", rank) + hidden, gate, up, down = _ffn_tensors(2, device, seed=130, intermediate=128) + local_i = 128 // world_size + feat_start = rank * local_i + feat_end = feat_start + local_i + local_hidden = hidden[:2] if rank == 0 else hidden[:1] + try: + qwen3_ffn( + local_hidden, + gate[feat_start:feat_end].contiguous(), + up[feat_start:feat_end].contiguous(), + down[:, feat_start:feat_end].contiguous(), + tp_group=dist.group.WORLD, + sequence_parallel=True, + ) + result_queue.put( + { + "ok": False, + "rank": rank, + "failures": "uneven SP tokens should have raised", + } + ) + except ValueError as exc: + message = str(exc) + if "matching shapes" not in message and "world_size" not in message: + raise + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _qwen3_8b_weights(device): + hidden = _randn((8, QWEN3_8B_HIDDEN_SIZE), seed=90, device=device, dtype=torch.bfloat16) + gate = _randn( + (QWEN3_8B_INTERMEDIATE_SIZE, QWEN3_8B_HIDDEN_SIZE), + seed=91, + device=device, + dtype=torch.bfloat16, + ) + up = _randn( + (QWEN3_8B_INTERMEDIATE_SIZE, QWEN3_8B_HIDDEN_SIZE), + seed=92, + device=device, + dtype=torch.bfloat16, + ) + down = _randn( + (QWEN3_8B_HIDDEN_SIZE, QWEN3_8B_INTERMEDIATE_SIZE), + seed=93, + device=device, + dtype=torch.bfloat16, + ) + grad = _randn((8, QWEN3_8B_HIDDEN_SIZE), seed=94, device=device, dtype=torch.bfloat16) + return hidden, gate, up, down, grad + + +def _qwen3_8b_tp2_worker(rank, world_size, init_method, result_queue): + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + device = torch.device("cuda", rank) + hidden, gate, up, down, grad = _qwen3_8b_weights(device) + with torch.no_grad(): + infer_tp1 = qwen3_ffn(hidden, gate, up, down) + tp1_inputs = [ + value.detach().clone().requires_grad_(True) for value in (hidden, gate, up, down) + ] + train_tp1 = qwen3_ffn(*tp1_inputs) + train_tp1.backward(grad) + + local_i = QWEN3_8B_INTERMEDIATE_SIZE // world_size + feat_start = rank * local_i + feat_end = feat_start + local_i + shard = ( + hidden, + gate[feat_start:feat_end].contiguous(), + up[feat_start:feat_end].contiguous(), + down[:, feat_start:feat_end].contiguous(), + ) + with torch.no_grad(): + infer_tp2 = qwen3_ffn(*shard, tp_group=dist.group.WORLD) + tp2_inputs = [value.detach().clone().requires_grad_(True) for value in shard] + train_tp2 = qwen3_ffn(*tp2_inputs, tp_group=dist.group.WORLD) + train_tp2.backward(grad) + + assert torch.equal(infer_tp1, infer_tp2) + assert torch.equal(train_tp1.detach(), train_tp2.detach()) + assert torch.equal(tp1_inputs[0].grad, tp2_inputs[0].grad) + assert torch.equal(tp1_inputs[1].grad[feat_start:feat_end], tp2_inputs[1].grad) + assert torch.equal(tp1_inputs[2].grad[feat_start:feat_end], tp2_inputs[2].grad) + assert torch.equal(tp1_inputs[3].grad[:, feat_start:feat_end], tp2_inputs[3].grad) + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def test_qwen_ffn_qwen3_8b_dimensions_are_pinned(): + assert QWEN3_8B_HIDDEN_SIZE == 4096 + assert QWEN3_8B_INTERMEDIATE_SIZE == 12288 + + +def test_qwen_ffn_backward_matches_autograd_reference(monkeypatch): + stub = _TorchKernelStub() + monkeypatch.setattr(ffn_module, "_C", stub) + monkeypatch.setattr(ffn_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) + + hidden = _randn((2, 3, 8), seed=0) + gate_weight = _randn((12, 8), seed=1) + up_weight = _randn((12, 8), seed=2) + down_weight = _randn((8, 12), seed=3) + grad_output = _randn(hidden.shape, seed=4) + + ref_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + expected, _, _, _ = _reference(*ref_inputs) + expected.backward(grad_output) + + actual_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + actual = qwen3_ffn(*actual_inputs) + actual.backward(grad_output) + + torch.testing.assert_close(actual, expected.detach()) + for actual_input, reference in zip(actual_inputs, ref_inputs, strict=True): + torch.testing.assert_close(actual_input.grad, reference.grad) + + assert stub.calls.count("det_gemm_fwd") == 6 + assert stub.calls.count("det_gemm_db") == 3 + assert stub.calls.count("swiglu_forward") == 1 + assert stub.calls.count("swiglu_backward") == 1 + + +def test_qwen_ffn_disable_split_k_false_uses_torch_matmul(monkeypatch): + stub = _TorchKernelStub() + monkeypatch.setattr(ffn_module, "_C", stub) + monkeypatch.setattr(ffn_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) + + hidden = _randn((2, 3, 8), seed=0) + gate_weight = _randn((12, 8), seed=1) + up_weight = _randn((12, 8), seed=2) + down_weight = _randn((8, 12), seed=3) + grad_output = _randn(hidden.shape, seed=4) + + ref_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + expected, _, _, _ = _reference(*ref_inputs) + expected.backward(grad_output) + + actual_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + actual = qwen3_ffn(*actual_inputs, disable_split_k=False) + actual.backward(grad_output) + + torch.testing.assert_close(actual, expected.detach()) + for actual_input, reference in zip(actual_inputs, ref_inputs, strict=True): + torch.testing.assert_close(actual_input.grad, reference.grad) + + assert stub.calls.count("det_gemm_fwd") == 0 + assert stub.calls.count("det_gemm_db") == 0 + assert stub.calls.count("swiglu_forward") == 1 + assert stub.calls.count("swiglu_backward") == 1 + + +def test_qwen_ffn_deterministic_false_uses_production_gemm(monkeypatch): + modes = [] + monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) + monkeypatch.setattr(ffn_module, "_require_ffn_kernels", lambda **kwargs: modes.append(kwargs)) + monkeypatch.setattr( + ffn_module._DeterministicFFNFunction, + "apply", + lambda *args: args[-1], + ) + + tensors = [torch.empty(1)] * 4 + assert qwen3_ffn(*tensors, deterministic=False) is False + assert modes == [{"disable_split_k": False}] + + +def test_qwen_ffn_rejects_conflicting_backend_switches(): + tensors = [torch.empty(1)] * 4 + + with pytest.raises(ValueError, match="conflicting FFN backends"): + qwen3_ffn(*tensors, deterministic=True, disable_split_k=False) + + +def test_qwen_ffn_rejects_non_bool_deterministic(): + tensors = [torch.empty(1)] * 4 + + with pytest.raises(TypeError, match="deterministic must be a bool or None"): + qwen3_ffn(*tensors, deterministic=1) # type: ignore[arg-type] + + +def test_qwen_ffn_rejects_non_bool_disable_split_k(): + hidden = torch.empty((2, 8), dtype=torch.bfloat16) + gate_weight = torch.empty((12, 8), dtype=torch.bfloat16) + up_weight = torch.empty((12, 8), dtype=torch.bfloat16) + down_weight = torch.empty((8, 12), dtype=torch.bfloat16) + with pytest.raises(TypeError, match="disable_split_k must be a bool"): + qwen3_ffn( + hidden, + gate_weight, + up_weight, + down_weight, + disable_split_k=1, # type: ignore[arg-type] + ) + + +def test_qwen_ffn_rejects_non_huggingface_weight_layout(): + hidden = torch.empty((2, 8), dtype=torch.bfloat16) + gate_weight = torch.empty((8, 12), dtype=torch.bfloat16) + up_weight = torch.empty((12, 8), dtype=torch.bfloat16) + down_weight = torch.empty((8, 12), dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="gate_weight must have shape"): + qwen3_ffn(hidden, gate_weight, up_weight, down_weight) + + +@requires_cuda_ffn +@pytest.mark.parametrize("disable_split_k", [True, False]) +def test_qwen_ffn_cuda_forward_backward_matches_fp32_reference(disable_split_k): + hidden = _randn((2, 3, 64), seed=10, device="cuda", dtype=torch.bfloat16) + gate_weight = _randn((128, 64), seed=11, device="cuda", dtype=torch.bfloat16) + up_weight = _randn((128, 64), seed=12, device="cuda", dtype=torch.bfloat16) + down_weight = _randn((64, 128), seed=13, device="cuda", dtype=torch.bfloat16) + grad_output = _randn(hidden.shape, seed=14, device="cuda", dtype=torch.bfloat16) + + ref_inputs = [ + value.detach().cpu().float().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + expected, _, _, _ = _reference(*ref_inputs) + expected.backward(grad_output.cpu().float()) + + actual_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, gate_weight, up_weight, down_weight) + ] + actual = qwen3_ffn(*actual_inputs, disable_split_k=disable_split_k) + actual.backward(grad_output) + + torch.testing.assert_close( + actual.cpu().float(), + expected.detach(), + atol=5e-2, + rtol=2e-2, + ) + for actual_input, reference in zip(actual_inputs, ref_inputs, strict=True): + torch.testing.assert_close( + actual_input.grad.cpu().float(), + reference.grad, + atol=5e-2, + rtol=2e-2, + ) + + +@requires_cuda_ffn +def test_qwen_ffn_cuda_forward_and_input_gradient_are_batch_invariant(): + gate_weight = _randn((128, 64), seed=20, device="cuda", dtype=torch.bfloat16) + up_weight = _randn((128, 64), seed=21, device="cuda", dtype=torch.bfloat16) + down_weight = _randn((64, 128), seed=22, device="cuda", dtype=torch.bfloat16) + hidden = _randn((6, 64), seed=23, device="cuda", dtype=torch.bfloat16) + grad_output = _randn(hidden.shape, seed=24, device="cuda", dtype=torch.bfloat16) + + full_hidden = hidden.detach().clone().requires_grad_(True) + full_output = qwen3_ffn(full_hidden, gate_weight, up_weight, down_weight) + full_output.backward(grad_output) + + slice_hidden = hidden[2:4].detach().clone().requires_grad_(True) + slice_output = qwen3_ffn(slice_hidden, gate_weight, up_weight, down_weight) + slice_output.backward(grad_output[2:4]) + + assert torch.equal(slice_output, full_output[2:4]) + assert torch.equal(slice_hidden.grad, full_hidden.grad[2:4]) + + +@requires_cuda_ffn +def test_qwen_ffn_cuda_train_and_infer_forward_are_bitwise_identical(): + hidden = _randn((16, 64), seed=30, device="cuda", dtype=torch.bfloat16) + gate_weight = _randn((128, 64), seed=31, device="cuda", dtype=torch.bfloat16) + up_weight = _randn((128, 64), seed=32, device="cuda", dtype=torch.bfloat16) + down_weight = _randn((64, 128), seed=33, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + infer = qwen3_ffn(hidden, gate_weight, up_weight, down_weight) + train_hidden = hidden.detach().clone().requires_grad_(True) + train = qwen3_ffn(train_hidden, gate_weight, up_weight, down_weight) + assert torch.equal(infer, train.detach()) + + +@requires_cuda_ffn +@pytest.mark.parametrize("token_count", _TOKEN_BOUNDARY_COUNTS) +def test_qwen_ffn_output_and_hidden_grad_are_batch_invariant(token_count): + device = torch.device("cuda", 0) + gate = _randn((256, _HIDDEN), seed=70, device=device, dtype=torch.bfloat16) + up = _randn((256, _HIDDEN), seed=71, device=device, dtype=torch.bfloat16) + down = _randn((_HIDDEN, 256), seed=72, device=device, dtype=torch.bfloat16) + hidden = _randn((token_count, _HIDDEN), seed=73, device=device, dtype=torch.bfloat16) + grad = _randn((token_count, _HIDDEN), seed=74, device=device, dtype=torch.bfloat16) + + full_hidden = hidden.detach().clone().requires_grad_(True) + full_output = qwen3_ffn(full_hidden, gate, up, down) + full_output.backward(grad) + slice_end = min(8, token_count) + slice_hidden = hidden[:slice_end].detach().clone().requires_grad_(True) + slice_output = qwen3_ffn(slice_hidden, gate, up, down) + slice_output.backward(grad[:slice_end]) + + assert torch.equal(slice_output, full_output[:slice_end]) + assert torch.equal(slice_hidden.grad, full_hidden.grad[:slice_end]) + + +@requires_cuda_ffn +def test_qwen_ffn_qwen3_8b_shapes_run_and_are_batch_invariant(): + device = torch.device("cuda", 0) + hidden, gate, up, down, grad = _qwen3_8b_weights(device) + + with torch.no_grad(): + infer = qwen3_ffn(hidden, gate, up, down) + full_hidden = hidden.detach().clone().requires_grad_(True) + full_gate = gate.detach().clone().requires_grad_(True) + full_up = up.detach().clone().requires_grad_(True) + full_down = down.detach().clone().requires_grad_(True) + train = qwen3_ffn(full_hidden, full_gate, full_up, full_down) + train.backward(grad) + assert torch.equal(infer, train.detach()) + + slice_hidden = hidden[:4].detach().clone().requires_grad_(True) + slice_out = qwen3_ffn(slice_hidden, full_gate, full_up, full_down) + slice_out.backward(grad[:4]) + assert torch.equal(slice_out, train.detach()[:4]) + assert torch.equal(slice_hidden.grad, full_hidden.grad[:4]) + + +@requires_cuda_ffn +def test_qwen_ffn_sequence_parallel_requires_tensor_parallel_group(): + device = torch.device("cuda", 0) + hidden, gate, up, down = _ffn_tensors(8, device, seed=120, intermediate=128) + with pytest.raises(ValueError, match="sequence_parallel requires a tensor-parallel group"): + qwen3_ffn(hidden, gate, up, down, sequence_parallel=True) + + +def test_qwen_ffn_tp_correctness_and_batch_invariance(): + _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 2, (1, False), timeout=90) + + +def test_qwen_ffn_tp_sp_correctness_and_batch_invariance(): + _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 2, (1, True), timeout=90) + + +def test_qwen_ffn_tp_cp_correctness_and_batch_invariance(): + _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 4, (2, False), timeout=90) + + +def test_qwen_ffn_tp_cp_sp_correctness_and_batch_invariance(): + _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 4, (2, True), timeout=90) + + +def test_qwen_ffn_tp1_vs_tp2_train_infer_bitwise_identical(): + _spawn_nccl_workers(_tp1_vs_tpn_train_infer_worker, 2, (True,), timeout=120) + + +def test_qwen_ffn_tp1_vs_tp8_train_infer_bitwise_identical(): + _spawn_nccl_workers(_tp1_vs_tpn_train_infer_worker, 8, (True,), timeout=120) + + +def test_qwen_ffn_world2_tp_sp_and_cp_match_tp1_cp1_bitwise(): + _spawn_nccl_workers(_topology_worker, 2, (_WORLD2_CONFIGS,), timeout=120) + + +def test_qwen_ffn_world4_tp_cp_sp_match_tp1_cp1_bitwise(): + _spawn_nccl_workers(_topology_worker, 4, (_WORLD4_CONFIGS,), timeout=120) + + +def test_qwen_ffn_world8_tp8_and_cp8_match_tp1_cp1_bitwise(): + _spawn_nccl_workers(_topology_worker, 8, (_WORLD8_WORLD_GROUP_CONFIGS,), timeout=120) + + +def test_qwen_ffn_world8_tp2_cp4_match_tp1_cp1_bitwise(): + _spawn_nccl_workers(_topology_worker, 8, (_WORLD8_TP2_CP4_CONFIGS,), timeout=120) + + +def test_qwen_ffn_world8_tp4_cp2_match_tp1_cp1_bitwise(): + _spawn_nccl_workers(_topology_worker, 8, (_WORLD8_TP4_CP2_CONFIGS,), timeout=120) + + +def test_qwen_ffn_collective_cache_reuses_grows_closes_and_rebuilds_group(): + _spawn_nccl_workers(_cache_worker, 2, timeout=120) + + +def test_qwen_ffn_sequence_parallel_rejects_uneven_tokens(): + _spawn_nccl_workers(_uneven_sp_worker, 2, timeout=90) + + +def test_qwen_ffn_qwen3_8b_shapes_tp2_matches_tp1_bitwise(): + _spawn_nccl_workers(_qwen3_8b_tp2_worker, 2, timeout=180) diff --git a/tests/test_swiglu.py b/tests/test_swiglu.py index bd9ff32c..22b26336 100644 --- a/tests/test_swiglu.py +++ b/tests/test_swiglu.py @@ -1,20 +1,92 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +"""SiLU / SwiGLU tests: native gold + CUDA / Triton candidates vs ground truth. + +Covers: +- Native correctness (fp32 formula, dtype path, shape guard) +- Axis A batch invariance (slice + padding, forward + backward) +- CUDA / Triton forward+backward vs NativeSiLUOp / NativeSwiGLUOp (issue #108 harness) +- Registry dispatch + OP_SPECS candidate paths +""" + +from __future__ import annotations + +import argparse + import pytest import torch +from rl_engine.kernels.gtest.op_checks import run_operator_suite +from rl_engine.kernels.gtest.operator_specs import ( + make_candidate, + make_operator_case, + operator_names, +) from rl_engine.kernels.ops.pytorch.activation.swiglu import NativeSiLUOp, NativeSwiGLUOp from rl_engine.kernels.registry import kernel_registry +try: + from rl_engine.kernels.ops.triton.activation.swiglu import TritonSiLUOp, TritonSwiGLUOp + + _HAS_TRITON_ACTIVATION = True +except ImportError: # pragma: no cover - triton may be missing in CPU-only builds. + _HAS_TRITON_ACTIVATION = False + TritonSiLUOp = None # type: ignore[misc, assignment] + TritonSwiGLUOp = None # type: ignore[misc, assignment] + +try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + from rl_engine.kernels.ops.cuda.activation.swiglu import SiLUCudaOp, SwiGLUCudaOp + + _HAS_CUDA_ACTIVATION = ( + _EXT_AVAILABLE and hasattr(_C, "silu_forward") and hasattr(_C, "swiglu_forward") + ) +except ImportError: # pragma: no cover - extension may be missing in CPU-only builds. + _HAS_CUDA_ACTIVATION = False + SiLUCudaOp = None # type: ignore[misc, assignment] + SwiGLUCudaOp = None # type: ignore[misc, assignment] + # Qwen3-8B SwiGLU intermediate dim (gate/up_proj output width). _INTERMEDIATE = 12288 # Shared helper -def _rand(shape, *, seed, dtype=torch.float32): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=dtype) +def _rand(shape, *, seed, dtype=torch.float32, device="cpu"): + gen = torch.Generator(device="cpu").manual_seed(seed) + t = torch.randn(*shape, generator=gen, dtype=torch.float32) + return t.to(device=device, dtype=dtype) + + +def _dtype_tolerance(dtype: torch.dtype) -> tuple[float, float]: + # Matches elementwise row of tolerance_contract.json (issue #108). + if dtype is torch.float32: + return 1e-5, 1e-5 + if dtype is torch.float16: + return 1e-3, 1e-3 + if dtype is torch.bfloat16: + return 2e-2, 1.6e-2 + raise ValueError(f"unsupported dtype: {dtype}") + + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +requires_cuda_activation = pytest.mark.skipif( + not (torch.cuda.is_available() and _HAS_CUDA_ACTIVATION), + reason="CUDA SiLU/SwiGLU extension is not available", +) +requires_triton_activation = pytest.mark.skipif( + not _HAS_TRITON_ACTIVATION, + reason="Triton SiLU/SwiGLU is not available", +) +requires_nvidia_cuda = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is not None, + reason="NVIDIA CUDA is required", +) + + +# --------------------------------------------------------------------------- +# Native gold (PyTorch reference) +# --------------------------------------------------------------------------- @pytest.mark.parametrize("dtype", (torch.float32, torch.bfloat16, torch.float16)) @@ -50,6 +122,17 @@ def test_native_swiglu_rejects_mismatched_shape(): NativeSwiGLUOp().forward(gate, up) +def test_native_activation_rejects_invalid_dtypes(): + with pytest.raises(TypeError, match="fp16, bf16, or fp32"): + NativeSiLUOp().forward(torch.ones(8, dtype=torch.int32)) + + with pytest.raises(TypeError, match="share dtype"): + NativeSwiGLUOp().forward( + torch.ones(8, dtype=torch.float16), + torch.ones(8, dtype=torch.bfloat16), + ) + + # Axis A -- batch invariance, bitwise (the WS1 "aligned" property). # A row's output must not depend on how many rows share the batch. def test_silu_batch_invariance_slice(): @@ -165,6 +248,436 @@ def test_swiglu_backward_batch_invariance_slice(): assert torch.equal(up_slice.grad, grad_up_full_sliced) -def test_registry_dispatches_native_activation_ops(): - assert isinstance(kernel_registry.get_op("silu"), NativeSiLUOp) - assert isinstance(kernel_registry.get_op("swiglu"), NativeSwiGLUOp) +def test_registry_dispatches_native_activation_ops_on_cpu(): + assert isinstance(kernel_registry.get_op("silu", device="cpu"), NativeSiLUOp) + assert isinstance(kernel_registry.get_op("swiglu", device="cpu"), NativeSwiGLUOp) + + +# --------------------------------------------------------------------------- +# CUDA / Triton candidates vs native gold (RMSNorm-style) +# --------------------------------------------------------------------------- + + +def _silu_impls(): + return [ + pytest.param("triton", marks=requires_triton_activation, id="triton"), + pytest.param("cuda", marks=requires_cuda_activation, id="cuda"), + ] + + +@requires_nvidia_cuda +def test_cuda_activation_symbols_are_built_on_cuda_host(): + if not _HAS_CUDA_ACTIVATION: + pytest.skip( + "CUDA is available but SiLU/SwiGLU symbols are missing from rl_engine._C; " + "rebuild the extension from the current source tree to exercise the CUDA path" + ) + + +def _make_silu_op(impl: str): + if impl == "cuda": + return SiLUCudaOp() + if impl == "triton": + return TritonSiLUOp() + raise ValueError(impl) + + +def _make_swiglu_op(impl: str): + if impl == "cuda": + return SwiGLUCudaOp() + if impl == "triton": + return TritonSwiGLUOp() + raise ValueError(impl) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + "shape", + [ + (1, 64), + (8, 256), + (4, 32, 512), + (2, 8, _INTERMEDIATE), # Qwen3-8B intermediate width + ], +) +def test_cuda_triton_silu_matches_native_forward_and_backward(impl, dtype, shape): + if impl == "cuda" and not _HAS_CUDA_ACTIVATION: + pytest.skip("CUDA SiLU extension is not available") + + native = NativeSiLUOp() + cand = _make_silu_op(impl) + + x_cpu = _rand(shape, seed=0, dtype=torch.float32) + dy_cpu = _rand(shape, seed=1, dtype=torch.float32) + + x_ref = x_cpu.to(dtype).float().detach().requires_grad_(True) + dy_ref = dy_cpu.to(dtype).float() + y_ref = native.forward_fp32(x_ref) + y_ref.backward(dy_ref) + + x_gpu = x_cpu.to(device="cuda", dtype=dtype).detach().requires_grad_(True) + dy_gpu = dy_cpu.to(device="cuda", dtype=dtype) + y_gpu = cand.forward(x_gpu) + y_gpu.backward(dy_gpu) + + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(y_gpu.detach().cpu().float(), y_ref.detach(), atol=atol, rtol=rtol) + torch.testing.assert_close(x_gpu.grad.detach().cpu().float(), x_ref.grad, atol=atol, rtol=rtol) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +@pytest.mark.parametrize( + "shape", + [ + (1, 64), + (8, 256), + (4, 32, 512), + (2, 8, _INTERMEDIATE), + ], +) +def test_cuda_triton_swiglu_matches_native_forward_and_backward(impl, dtype, shape): + if impl == "cuda" and not _HAS_CUDA_ACTIVATION: + pytest.skip("CUDA SwiGLU extension is not available") + + native = NativeSwiGLUOp() + cand = _make_swiglu_op(impl) + + gate_cpu = _rand(shape, seed=2, dtype=torch.float32) + up_cpu = _rand(shape, seed=3, dtype=torch.float32) + dy_cpu = _rand(shape, seed=4, dtype=torch.float32) + + gate_ref = gate_cpu.to(dtype).float().detach().requires_grad_(True) + up_ref = up_cpu.to(dtype).float().detach().requires_grad_(True) + dy_ref = dy_cpu.to(dtype).float() + y_ref = native.forward_fp32(gate_ref, up_ref) + y_ref.backward(dy_ref) + + gate_gpu = gate_cpu.to(device="cuda", dtype=dtype).detach().requires_grad_(True) + up_gpu = up_cpu.to(device="cuda", dtype=dtype).detach().requires_grad_(True) + dy_gpu = dy_cpu.to(device="cuda", dtype=dtype) + y_gpu = cand.forward(gate_gpu, up_gpu) + y_gpu.backward(dy_gpu) + + atol, rtol = _dtype_tolerance(dtype) + torch.testing.assert_close(y_gpu.detach().cpu().float(), y_ref.detach(), atol=atol, rtol=rtol) + torch.testing.assert_close( + gate_gpu.grad.detach().cpu().float(), gate_ref.grad, atol=atol, rtol=rtol + ) + torch.testing.assert_close( + up_gpu.grad.detach().cpu().float(), up_ref.grad, atol=atol, rtol=rtol + ) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_silu_batch_invariance_bitwise(impl): + if impl == "cuda" and not _HAS_CUDA_ACTIVATION: + pytest.skip("CUDA SiLU extension is not available") + + op = _make_silu_op(impl) + x = _rand((8, 32, 256), seed=5, dtype=torch.bfloat16, device="cuda") + full = op.forward(x) + assert torch.equal(op.forward(x[:1]), full[:1]) + assert torch.equal(op.forward(x[3:5]), full[3:5]) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_swiglu_batch_invariance_bitwise(impl): + if impl == "cuda" and not _HAS_CUDA_ACTIVATION: + pytest.skip("CUDA SwiGLU extension is not available") + + op = _make_swiglu_op(impl) + gate = _rand((8, 32, 256), seed=6, dtype=torch.bfloat16, device="cuda") + up = _rand((8, 32, 256), seed=7, dtype=torch.bfloat16, device="cuda") + full = op.forward(gate, up) + assert torch.equal(op.forward(gate[:1], up[:1]), full[:1]) + assert torch.equal(op.forward(gate[3:5], up[3:5]), full[3:5]) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_silu_deterministic_repeat(impl): + if impl == "cuda" and not _HAS_CUDA_ACTIVATION: + pytest.skip("CUDA SiLU extension is not available") + + op = _make_silu_op(impl) + x = _rand((32, 1024), seed=8, dtype=torch.bfloat16, device="cuda") + dy = _rand((32, 1024), seed=9, dtype=torch.bfloat16, device="cuda") + + def _run(): + x_r = x.detach().clone().requires_grad_(True) + y = op.forward(x_r) + y.backward(dy) + return y.detach(), x_r.grad.detach() + + y0, dx0 = _run() + torch.cuda.synchronize() + for _ in range(5): + y, dx = _run() + torch.cuda.synchronize() + assert torch.equal(y0, y) + assert torch.equal(dx0, dx) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_swiglu_rejects_mismatched_shape(impl): + if impl == "cuda" and not _HAS_CUDA_ACTIVATION: + pytest.skip("CUDA SwiGLU extension is not available") + + op = _make_swiglu_op(impl) + gate = torch.randn(2, 3, device="cuda") + up = torch.randn(2, 4, device="cuda") + with pytest.raises(ValueError, match="share shape"): + op.forward(gate, up) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_silu_backward_batch_invariance_bitwise(impl): + op = _make_silu_op(impl) + x = _rand((6, 4, 64), seed=10, dtype=torch.bfloat16, device="cuda") + dy = _rand(x.shape, seed=11, dtype=torch.bfloat16, device="cuda") + + x_full = x.detach().clone().requires_grad_(True) + op.forward(x_full).backward(dy) + full_grad = x_full.grad[2:4].clone() + + x_slice = x[2:4].detach().clone().requires_grad_(True) + op.forward(x_slice).backward(dy[2:4]) + assert torch.equal(x_slice.grad, full_grad) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_swiglu_backward_batch_invariance_bitwise(impl): + op = _make_swiglu_op(impl) + gate = _rand((6, 4, 64), seed=12, dtype=torch.bfloat16, device="cuda") + up = _rand((6, 4, 64), seed=13, dtype=torch.bfloat16, device="cuda") + dy = _rand(gate.shape, seed=14, dtype=torch.bfloat16, device="cuda") + + gate_full = gate.detach().clone().requires_grad_(True) + up_full = up.detach().clone().requires_grad_(True) + op.forward(gate_full, up_full).backward(dy) + full_d_gate = gate_full.grad[2:4].clone() + full_d_up = up_full.grad[2:4].clone() + + gate_slice = gate[2:4].detach().clone().requires_grad_(True) + up_slice = up[2:4].detach().clone().requires_grad_(True) + op.forward(gate_slice, up_slice).backward(dy[2:4]) + assert torch.equal(gate_slice.grad, full_d_gate) + assert torch.equal(up_slice.grad, full_d_up) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_silu_padding_invariance_forward_and_backward(impl): + op = _make_silu_op(impl) + x = _rand((4, 64), seed=15, dtype=torch.bfloat16, device="cuda") + dy = _rand(x.shape, seed=16, dtype=torch.bfloat16, device="cuda") + x_padded = torch.cat( + [x, _rand((3, 64), seed=17, dtype=torch.bfloat16, device="cuda")], dim=0 + ).requires_grad_(True) + dy_padded = torch.cat([dy, _rand((3, 64), seed=18, dtype=torch.bfloat16, device="cuda")], dim=0) + + y_padded = op.forward(x_padded) + y_padded.backward(dy_padded) + + x_real = x.detach().clone().requires_grad_(True) + y_real = op.forward(x_real) + y_real.backward(dy) + assert torch.equal(y_padded[:4], y_real) + assert torch.equal(x_padded.grad[:4], x_real.grad) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_swiglu_padding_invariance_forward_and_backward(impl): + op = _make_swiglu_op(impl) + gate = _rand((4, 64), seed=19, dtype=torch.bfloat16, device="cuda") + up = _rand((4, 64), seed=20, dtype=torch.bfloat16, device="cuda") + dy = _rand(gate.shape, seed=21, dtype=torch.bfloat16, device="cuda") + gate_padded = torch.cat( + [gate, _rand((3, 64), seed=22, dtype=torch.bfloat16, device="cuda")], dim=0 + ).requires_grad_(True) + up_padded = torch.cat( + [up, _rand((3, 64), seed=23, dtype=torch.bfloat16, device="cuda")], dim=0 + ).requires_grad_(True) + dy_padded = torch.cat([dy, _rand((3, 64), seed=24, dtype=torch.bfloat16, device="cuda")], dim=0) + + y_padded = op.forward(gate_padded, up_padded) + y_padded.backward(dy_padded) + + gate_real = gate.detach().clone().requires_grad_(True) + up_real = up.detach().clone().requires_grad_(True) + y_real = op.forward(gate_real, up_real) + y_real.backward(dy) + assert torch.equal(y_padded[:4], y_real) + assert torch.equal(gate_padded.grad[:4], gate_real.grad) + assert torch.equal(up_padded.grad[:4], up_real.grad) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_swiglu_deterministic_repeat(impl): + op = _make_swiglu_op(impl) + gate = _rand((16, 256), seed=25, dtype=torch.bfloat16, device="cuda") + up = _rand((16, 256), seed=26, dtype=torch.bfloat16, device="cuda") + dy = _rand(gate.shape, seed=27, dtype=torch.bfloat16, device="cuda") + + def _run(): + gate_r = gate.detach().clone().requires_grad_(True) + up_r = up.detach().clone().requires_grad_(True) + y = op.forward(gate_r, up_r) + y.backward(dy) + return y.detach(), gate_r.grad.detach(), up_r.grad.detach() + + expected = _run() + for _ in range(5): + actual = _run() + torch.cuda.synchronize() + assert all(torch.equal(lhs, rhs) for lhs, rhs in zip(expected, actual, strict=True)) + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_activation_handles_noncontiguous_and_empty_inputs(impl): + silu = _make_silu_op(impl) + swiglu = _make_swiglu_op(impl) + + x = torch.randn(5, 3, device="cuda", dtype=torch.bfloat16).T.requires_grad_(True) + assert not x.is_contiguous() + silu.forward(x).sum().backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + + gate = torch.randn(5, 3, device="cuda", dtype=torch.bfloat16).T.requires_grad_(True) + up = torch.randn(5, 3, device="cuda", dtype=torch.bfloat16).T.requires_grad_(True) + assert not gate.is_contiguous() and not up.is_contiguous() + swiglu.forward(gate, up).sum().backward() + assert gate.grad is not None and up.grad is not None + + empty = torch.empty((0, 64), device="cuda", dtype=torch.bfloat16, requires_grad=True) + silu_empty = silu.forward(empty) + silu_empty.sum().backward() + assert silu_empty.shape == empty.shape and empty.grad.shape == empty.shape + + empty_gate = empty.detach().clone().requires_grad_(True) + empty_up = empty.detach().clone().requires_grad_(True) + swiglu_empty = swiglu.forward(empty_gate, empty_up) + swiglu_empty.sum().backward() + assert swiglu_empty.shape == empty_gate.shape + assert empty_gate.grad.shape == empty_gate.shape and empty_up.grad.shape == empty_up.shape + + +@requires_cuda +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_activation_rejects_invalid_dtypes(impl): + silu = _make_silu_op(impl) + swiglu = _make_swiglu_op(impl) + with pytest.raises(TypeError, match="fp16, bf16, or fp32"): + silu.forward(torch.ones(8, device="cuda", dtype=torch.int32)) + with pytest.raises(TypeError, match="share dtype"): + swiglu.forward( + torch.ones(8, device="cuda", dtype=torch.float16), + torch.ones(8, device="cuda", dtype=torch.bfloat16), + ) + + +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires at least two CUDA devices") +@pytest.mark.parametrize("impl", _silu_impls()) +def test_cuda_triton_swiglu_rejects_cross_device_inputs(impl): + op = _make_swiglu_op(impl) + gate = torch.ones(8, device="cuda:0") + up = torch.ones(8, device="cuda:1") + with pytest.raises(RuntimeError, match=r"same .*device"): + op.forward(gate, up) + + +# --------------------------------------------------------------------------- +# Issue #108 ground-truth harness (OP_SPECS + check_operator path) +# --------------------------------------------------------------------------- + + +def _spec_args(op: str, **overrides) -> argparse.Namespace: + values = dict( + op=op, + candidate="pytorch", + arch_key=None, + batch=2, + seq=4, + vocab=17, + seed=123, + input_mode="random", + constant_value=0.5, + token_value=3, + normalized_dim=8, + k_dim=8, + n_dim=8, + theta=1.0e6, + eps=1.0e-6, + ) + values.update(overrides) + return argparse.Namespace(**values) + + +def test_silu_swiglu_registered_in_op_specs(): + assert "silu" in operator_names() + assert "swiglu" in operator_names() + + +def test_silu_pytorch_candidate_suite_passes_issue_108_helper(): + args = _spec_args("silu", candidate="pytorch") + report = run_operator_suite( + "silu", + candidates=[make_candidate(args)], + cases=[make_operator_case(args, torch.float32, torch.device("cpu"))], + check_grad=True, + ) + assert report.passed + + +def test_swiglu_pytorch_candidate_suite_passes_issue_108_helper(): + args = _spec_args("swiglu", candidate="pytorch") + report = run_operator_suite( + "swiglu", + candidates=[make_candidate(args)], + cases=[make_operator_case(args, torch.float32, torch.device("cpu"))], + check_grad=True, + ) + assert report.passed + + +@requires_cuda +@pytest.mark.parametrize( + "candidate", + [ + pytest.param("triton", marks=requires_triton_activation, id="triton"), + pytest.param("cuda", marks=requires_cuda_activation, id="cuda"), + ], +) +@pytest.mark.parametrize("op_name", ["silu", "swiglu"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_silu_swiglu_cuda_triton_issue_108_harness(candidate, op_name, dtype): + if candidate == "cuda" and not _HAS_CUDA_ACTIVATION: + pytest.skip("CUDA activation extension is not available") + if candidate == "triton" and not _HAS_TRITON_ACTIVATION: + pytest.skip("Triton SiLU/SwiGLU is not available") + + args = _spec_args(op_name, candidate=candidate, batch=2, seq=8) + device = torch.device("cuda") + report = run_operator_suite( + op_name, + candidates=[make_candidate(args)], + cases=[make_operator_case(args, dtype, device)], + check_grad=True, + ) + assert report.passed, ( + f"{op_name}/{candidate}/{dtype} failed against gold: " + f"{report.candidates[0].cases[0].outputs}" + ) diff --git a/tests/test_vocab_parallel_logp.py b/tests/test_vocab_parallel_logp.py new file mode 100644 index 00000000..1dee8bff --- /dev/null +++ b/tests/test_vocab_parallel_logp.py @@ -0,0 +1,558 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic vocab-parallel TP logprob reference tests (issue #241 PR3). + +Bit-level determinism assertions compare raw bit patterns via +``tensor.view(torch.int32)`` rather than ``torch.equal``: value equality +treats ``-0.0 == 0.0`` as equal and ``NaN != NaN`` as different, neither of +which is what a bitwise claim means. +""" + +from __future__ import annotations + +import queue +import tempfile +import traceback +from pathlib import Path + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.kernels.gtest.tolerance import load_contract +from rl_engine.kernels.logprob_contract import ( + DeterminismScope, + LogprobContract, + LogprobContractError, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp import NativeBatchInvariantLogpOp +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + BACKEND_ID, + VocabParallelLogprobOp, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + +REAL_VOCAB = 27 +PADDED_VOCAB = 32 +NUM_TILES = 8 +NUM_TOKENS = 6 +ACTIVE = (True, True, True, True, True, False) + + +def _even_bounds(padded: int, world: int) -> tuple[tuple[int, int], ...]: + shard = padded // world + return tuple( + (rank * shard, padded if rank == world - 1 else (rank + 1) * shard) for rank in range(world) + ) + + +def _contract( + *, + tp_rank: int = 0, + tp_world_size: int = 1, + bounds: tuple[tuple[int, int], ...] | None = None, + real_vocab: int = REAL_VOCAB, + padded_vocab: int = PADDED_VOCAB, + num_tokens: int = NUM_TOKENS, + active: tuple[bool, ...] = ACTIVE, + dtype: str = "fp32", + determinism_scope: str = "cross_tp_bitwise", +) -> LogprobContract: + return LogprobContract( + role="train", + dtype=dtype, + mask=MaskSpec(num_tokens=num_tokens, active_mask=active), + sharding=ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=( + bounds if bounds is not None else _even_bounds(padded_vocab, tp_world_size) + ), + real_vocab_size=real_vocab, + padded_vocab_size=padded_vocab, + ), + reduction=ReductionSpec(determinism_scope=determinism_scope), + ) + + +def _inputs(dtype=torch.float32, seed: int = 2026): + torch.manual_seed(seed) + logits = torch.randn(NUM_TOKENS, PADDED_VOCAB, dtype=torch.float32).to(dtype) + targets = torch.tensor([1, 5, REAL_VOCAB - 1, 0, 13, -100]) + return logits, targets + + +def _bits(tensor: torch.Tensor) -> torch.Tensor: + view_dtype = {torch.float32: torch.int32, torch.bfloat16: torch.int16}[tensor.dtype] + return tensor.contiguous().view(view_dtype) + + +def _bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and bool((_bits(a) == _bits(b)).all()) + + +def _case_shard_size_mismatch(): + logits, targets = _inputs() + return logits, targets, _contract(tp_rank=0, tp_world_size=2), NUM_TILES, "vocab columns" + + +def _case_mask_length_mismatch(): + logits, targets = _inputs() + contract = _contract(num_tokens=NUM_TOKENS + 1, active=ACTIVE + (True,)) + return logits, targets, contract, NUM_TILES, "num_tokens" + + +def _case_dtype_mismatch(): + logits, targets = _inputs() + return logits, targets, _contract(dtype="bf16"), NUM_TILES, "dtype" + + +def _case_tile_misaligned_bounds(): + # Tile size is 32/8 = 4; a boundary at 6 is misaligned. + logits, targets = _inputs() + contract = _contract(tp_world_size=2, bounds=((0, 6), (6, 32))) + return logits[:, :6], targets, contract, NUM_TILES, "tile" + + +def _case_bad_num_vocab_tiles(): + logits, targets = _inputs() + return logits, targets, _contract(), 7, "num_vocab_tiles" + + +def _case_active_target_out_of_real_vocab(): + logits, targets = _inputs() + bad_targets = targets.clone() + bad_targets[0] = REAL_VOCAB # padding column, active row + return logits, bad_targets, _contract(), NUM_TILES, "real vocabulary" + + +def _case_all_inf_active_row(): + logits, targets = _inputs() + poisoned = logits.clone() + poisoned[0, :] = float("-inf") + return poisoned, targets, _contract(), NUM_TILES, "non-finite" + + +@pytest.mark.parametrize( + "case", + [ + _case_shard_size_mismatch, + _case_mask_length_mismatch, + _case_dtype_mismatch, + _case_tile_misaligned_bounds, + _case_bad_num_vocab_tiles, + _case_active_target_out_of_real_vocab, + _case_all_inf_active_row, + ], + ids=lambda fn: fn.__name__.removeprefix("_case_"), +) +def test_invalid_invocations_fail_loudly(case): + logits, targets, contract, num_tiles, match = case() + with pytest.raises(LogprobContractError, match=match): + VocabParallelLogprobOp()(logits, targets, contract=contract, num_vocab_tiles=num_tiles) + + +class TestSingleRank: + def test_repeated_runs_are_bitwise_identical(self): + contract = _contract() + logits, targets = _inputs() + op = VocabParallelLogprobOp() + logp_a, lse_a = op(logits, targets, contract=contract, num_vocab_tiles=NUM_TILES) + logp_b, lse_b = op(logits, targets, contract=contract, num_vocab_tiles=NUM_TILES) + assert _bitwise_equal(logp_a, logp_b) + assert _bitwise_equal(lse_a, lse_b) + + def test_batch_invariance_same_row_any_context(self): + contract_full = _contract() + logits, targets = _inputs() + op = VocabParallelLogprobOp() + logp_full, lse_full = op(logits, targets, contract=contract_full, num_vocab_tiles=NUM_TILES) + + contract_single = _contract(num_tokens=1, active=(True,)) + logp_one, lse_one = op( + logits[2:3], targets[2:3], contract=contract_single, num_vocab_tiles=NUM_TILES + ) + assert _bitwise_equal(logp_full[2:3], logp_one) + assert _bitwise_equal(lse_full[2:3], lse_one) + + def test_matches_ws1_batch_invariant_logp_within_contract_tolerance(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + contract = _contract(padded_vocab=REAL_VOCAB + 5) + # Use a real==padded contract so the WS1 op sees identical logits. + contract = _contract(real_vocab=PADDED_VOCAB, padded_vocab=PADDED_VOCAB) + logits, targets = _inputs() + logp, _ = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + ws1 = NativeBatchInvariantLogpOp().apply(logits, targets) + active = torch.tensor(ACTIVE) + assert torch.allclose( + logp[active], ws1[active], atol=tolerance["atol"], rtol=tolerance["rtol"] + ) + + def test_padding_columns_are_excluded_and_finite(self): + contract = _contract() + logits, targets = _inputs() + boosted = logits.clone() + boosted[:, REAL_VOCAB:] = 1e4 # huge padding logits must not leak into LSE + logp, lse = VocabParallelLogprobOp()( + boosted, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + ref_lse = torch.logsumexp(boosted[:, :REAL_VOCAB].float(), dim=-1) + assert torch.isfinite(logp).all() and torch.isfinite(lse).all() + assert torch.allclose(lse, ref_lse, atol=1e-5) + + def test_inactive_rows_zero_filled_lse_still_exported(self): + contract = _contract() + logits, targets = _inputs() + logp, lse = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + assert logp[-1].item() == 0.0 + assert torch.isfinite(lse[-1]) + + +class TestNonDeterministicPath: + """deterministic=False: the fast whole-shard reduction with no guarantee.""" + + def test_rejects_a_cross_tp_bitwise_contract(self): + logits, targets = _inputs() + with pytest.raises(LogprobContractError, match="cross_tp_bitwise"): + VocabParallelLogprobOp()(logits, targets, contract=_contract(), deterministic=False) + + def test_rejects_a_non_bool_flag(self): + logits, targets = _inputs() + with pytest.raises(LogprobContractError, match="deterministic must be a bool"): + VocabParallelLogprobOp()(logits, targets, contract=_contract(), deterministic=1) + + def test_matches_the_deterministic_path_within_tolerance(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + logits, targets = _inputs() + relaxed = _contract(determinism_scope="fixed_topology") + logp_fast, lse_fast = VocabParallelLogprobOp()( + logits, targets, contract=relaxed, deterministic=False + ) + logp_det, lse_det = VocabParallelLogprobOp()( + logits, targets, contract=_contract(), num_vocab_tiles=NUM_TILES + ) + assert torch.allclose(logp_fast, logp_det, atol=tolerance["atol"], rtol=tolerance["rtol"]) + assert torch.allclose(lse_fast, lse_det, atol=tolerance["atol"], rtol=tolerance["rtol"]) + + def test_ignores_tile_constraints(self): + # 7 does not divide the padded vocab; the deterministic path rejects it, + # the fast path never looks at it. + logits, targets = _inputs() + relaxed = _contract(determinism_scope="fixed_topology") + logp, lse = VocabParallelLogprobOp()( + logits, targets, contract=relaxed, num_vocab_tiles=7, deterministic=False + ) + ref_lse = torch.logsumexp(logits[:, :REAL_VOCAB].float(), dim=-1) + assert torch.allclose(lse, ref_lse, atol=1e-5) + assert torch.isfinite(logp).all() + + def test_grads_match_autograd_oracle(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + relaxed = _contract(determinism_scope="fixed_topology") + logits, targets = _inputs() + x = logits.clone().requires_grad_(True) + logp, lse = VocabParallelLogprobOp()(x, targets, contract=relaxed, deterministic=False) + (logp.sum() + 0.5 * lse.sum()).backward() + + y = logits.clone().requires_grad_(True) + ref_lse = torch.logsumexp(y[:, :REAL_VOCAB].float(), dim=-1) + safe = targets.clamp(0, REAL_VOCAB - 1) + ref_logp = y[torch.arange(NUM_TOKENS), safe].float() - ref_lse + ref_logp = torch.where(torch.tensor(ACTIVE), ref_logp, torch.zeros_like(ref_logp)) + (ref_logp.sum() + 0.5 * ref_lse.sum()).backward() + + assert torch.allclose(x.grad, y.grad, atol=tolerance["atol"], rtol=tolerance["rtol"]) + assert bool((x.grad[:, REAL_VOCAB:] == 0).all()) + + +class TestBackward: + def test_grads_match_autograd_oracle(self): + tolerance = load_contract()["accuracy"]["default"]["logprob"]["float32"] + contract = _contract() + logits, targets = _inputs() + x = logits.clone().requires_grad_(True) + logp, lse = VocabParallelLogprobOp()( + x, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + (logp.sum() + 0.5 * lse.sum()).backward() + + y = logits.clone().requires_grad_(True) + ref_lse = torch.logsumexp(y[:, :REAL_VOCAB].float(), dim=-1) + safe = targets.clamp(0, REAL_VOCAB - 1) + ref_logp = y[torch.arange(NUM_TOKENS), safe].float() - ref_lse + ref_logp = torch.where(torch.tensor(ACTIVE), ref_logp, torch.zeros_like(ref_logp)) + (ref_logp.sum() + 0.5 * ref_lse.sum()).backward() + + assert torch.allclose(x.grad, y.grad, atol=tolerance["atol"], rtol=tolerance["rtol"]) + assert bool((x.grad[:, REAL_VOCAB:] == 0).all()) + + # No grad requested -> outputs detached from autograd entirely. + logp_ng, lse_ng = VocabParallelLogprobOp()( + logits, targets, contract=contract, num_vocab_tiles=NUM_TILES + ) + assert not logp_ng.requires_grad and not lse_ng.requires_grad + + def test_inactive_rows_grad_asymmetry(self): + """The logp term is zeroed on inactive rows; the lse term still flows — + lse is a row property exported (and differentiable) for every row.""" + + contract = _contract() + logits, targets = _inputs() + + x = logits.clone().requires_grad_(True) + _, lse = VocabParallelLogprobOp()(x, targets, contract=contract, num_vocab_tiles=NUM_TILES) + lse.sum().backward() + assert bool((x.grad[-1, :REAL_VOCAB].abs() > 0).any()) + + z = logits.clone().requires_grad_(True) + logp, _ = VocabParallelLogprobOp()(z, targets, contract=contract, num_vocab_tiles=NUM_TILES) + logp.sum().backward() + assert bool((z.grad[-1] == 0).all()) + + +def test_dispatch_resolves_reference_and_leaves_legacy_untouched(): + registry = KernelRegistry() + contract = _contract() + + result = registry.get_logprob_op(contract) + assert result.capability.backend_id == BACKEND_ID + assert result.provenance["fallback"] is False + assert isinstance(result.op, VocabParallelLogprobOp) + assert ( + result.provenance["contract"]["reduction"]["determinism_scope"] + == DeterminismScope.CROSS_TP_BITWISE.value + ) + + by_id = registry.get_logprob_op(contract, requested_backend=BACKEND_ID) + assert by_id.capability.backend_id == BACKEND_ID + by_kind = registry.get_logprob_op(contract, requested_backend="reference") + assert by_kind.capability.backend_id == BACKEND_ID + + for ops in registry._priority_map.values(): + for candidates in ops.values(): + assert OpBackend.PYTORCH_VOCAB_PARALLEL_LOGP not in candidates + + +# --------------------------------------------------------------------------- +# Multi-rank gloo tests (spawn pattern from tests/test_linear_logp.py) +# --------------------------------------------------------------------------- + + +def _gloo_available() -> bool: + return torch.distributed.is_available() and torch.distributed.is_gloo_available() + + +requires_gloo = pytest.mark.skipif( + not _gloo_available(), reason="requires torch.distributed with the gloo backend" +) + +_WORLD_SIZE = 4 +_UNEVEN_BOUNDS = ((0, 4), (4, 16), (16, 24), (24, 32)) # tile-aligned (tile=4) + + +def _tp_worker(rank, world_size, init_method, result_queue, scenario): + import torch.distributed as dist + + torch.set_num_threads(1) + try: + dist.init_process_group( + backend="gloo", init_method=init_method, rank=rank, world_size=world_size + ) + dtype = torch.bfloat16 if scenario == "bf16" else torch.float32 + dtype_name = "bf16" if scenario == "bf16" else "fp32" + bounds = _UNEVEN_BOUNDS if scenario == "uneven" else _even_bounds(PADDED_VOCAB, world_size) + logits, targets = _inputs(dtype=dtype) + start, end = bounds[rank] + + op = VocabParallelLogprobOp() + tiles = 16 if scenario == "preflight" and rank == 0 else NUM_TILES + contract_tp = _contract( + tp_rank=rank, tp_world_size=world_size, bounds=bounds, dtype=dtype_name + ) + + if scenario == "preflight": + try: + op( + logits[:, start:end].contiguous().clone(), + targets, + contract=contract_tp, + tp_group=dist.group.WORLD, + num_vocab_tiles=tiles, + ) + result_queue.put({"ok": False, "rank": rank, "traceback": "no error raised"}) + except LogprobContractError: + result_queue.put({"ok": True, "rank": rank}) + return + + if scenario == "mode_mismatch": + # Same relaxed contract everywhere; only rank 0 runs deterministic. + relaxed = _contract( + tp_rank=rank, + tp_world_size=world_size, + bounds=bounds, + determinism_scope="fixed_topology", + ) + try: + op( + logits[:, start:end].contiguous().clone(), + targets, + contract=relaxed, + tp_group=dist.group.WORLD, + num_vocab_tiles=NUM_TILES, + deterministic=(rank == 0), + ) + result_queue.put({"ok": False, "rank": rank, "traceback": "no error raised"}) + except LogprobContractError: + result_queue.put({"ok": True, "rank": rank}) + return + + if scenario == "nondeterministic": + relaxed = _contract( + tp_rank=rank, + tp_world_size=world_size, + bounds=bounds, + determinism_scope="fixed_topology", + ) + shard = logits[:, start:end].contiguous().clone().requires_grad_(True) + logp_tp, lse_tp = op( + shard, targets, contract=relaxed, tp_group=dist.group.WORLD, deterministic=False + ) + (logp_tp.sum() + 0.5 * lse_tp.sum()).backward() + + full = logits.clone().requires_grad_(True) + relaxed_tp1 = _contract(determinism_scope="fixed_topology") + logp_one, lse_one = op(full, targets, contract=relaxed_tp1, deterministic=False) + (logp_one.sum() + 0.5 * lse_one.sum()).backward() + + result_queue.put( + { + "ok": True, + "rank": rank, + "logp_close": torch.allclose(logp_tp, logp_one, atol=1e-6), + "lse_close": torch.allclose(lse_tp, lse_one, atol=1e-6), + "grad_close": torch.allclose(shard.grad, full.grad[:, start:end], atol=1e-6), + "logp": logp_tp.detach().float(), + "lse": lse_tp.detach().float(), + } + ) + return + + shard = logits[:, start:end].contiguous().clone().requires_grad_(True) + logp_tp, lse_tp = op( + shard, + targets, + contract=contract_tp, + tp_group=dist.group.WORLD, + num_vocab_tiles=NUM_TILES, + ) + (logp_tp.sum() + 0.5 * lse_tp.sum()).backward() + + # In-process TP=1 run of the same op on the full logits: the cross-TP + # bitwise claim is TP=n output == TP=1 output, bit for bit. + full = logits.clone().requires_grad_(True) + contract_tp1 = _contract(dtype=dtype_name) + logp_one, lse_one = op(full, targets, contract=contract_tp1, num_vocab_tiles=NUM_TILES) + (logp_one.sum() + 0.5 * lse_one.sum()).backward() + + result_queue.put( + { + "ok": True, + "rank": rank, + "logp_bits_match": _bitwise_equal(logp_tp, logp_one), + "lse_bits_match": _bitwise_equal(lse_tp, lse_one), + "grad_bits_match": _bitwise_equal(shard.grad, full.grad[:, start:end]), + "logp": logp_tp.detach().float(), + "lse": lse_tp.detach().float(), + } + ) + except Exception: + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _run_gloo_scenario(scenario): + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmpdir: + init_method = (Path(tmpdir) / "gloo_init").as_uri() + result_queue = ctx.Queue() + processes = [ + ctx.Process( + target=_tp_worker, + args=(rank, _WORLD_SIZE, init_method, result_queue, scenario), + ) + for rank in range(_WORLD_SIZE) + ] + results = [] + try: + for process in processes: + process.start() + for _ in range(_WORLD_SIZE): + try: + results.append(result_queue.get(timeout=60)) + except queue.Empty: + for process in processes: + if process.is_alive(): + process.terminate() + pytest.fail("timed out waiting for vocab-parallel gloo workers") + finally: + for process in processes: + process.join(timeout=10) + if process.is_alive(): + process.terminate() + results.sort(key=lambda item: item["rank"]) + for result in results: + assert result["ok"], result.get("traceback") + for process in processes: + assert process.exitcode == 0 + return results + + +@requires_gloo +@pytest.mark.parametrize("scenario", ["even", "uneven", "bf16"]) +def test_tp4_bitwise_identical_to_tp1(scenario): + results = _run_gloo_scenario(scenario) + for result in results: + assert result["logp_bits_match"], f"rank {result['rank']} logp bits differ from TP=1" + assert result["lse_bits_match"], f"rank {result['rank']} lse bits differ from TP=1" + assert result["grad_bits_match"], f"rank {result['rank']} grad bits differ from TP=1" + # Outputs are replicated: every rank must hold identical bits. + for other in results[1:]: + assert _bitwise_equal(results[0]["logp"], other["logp"]) + assert _bitwise_equal(results[0]["lse"], other["lse"]) + + +@requires_gloo +def test_preflight_rejects_mismatched_num_vocab_tiles(): + _run_gloo_scenario("preflight") + + +@requires_gloo +def test_preflight_rejects_mixed_deterministic_modes(): + _run_gloo_scenario("mode_mismatch") + + +@requires_gloo +def test_tp4_nondeterministic_matches_tp1_within_tolerance(): + """No bitwise claim: the fast path's grouping changes with the TP degree. + The values must still agree within fp32 tolerance, and the outputs stay + replicated bitwise across ranks — every rank merges the same partials.""" + + results = _run_gloo_scenario("nondeterministic") + for result in results: + assert result["logp_close"], f"rank {result['rank']} logp differs from TP=1 beyond atol" + assert result["lse_close"], f"rank {result['rank']} lse differs from TP=1 beyond atol" + assert result["grad_close"], f"rank {result['rank']} grads differ from TP=1 beyond atol" + for other in results[1:]: + assert _bitwise_equal(results[0]["logp"], other["logp"]) + assert _bitwise_equal(results[0]["lse"], other["lse"]) From a450ef560fc98e3dbec44c0ad6b9bb98fc84eb5a Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 20 Aug 2026 16:39:56 +0800 Subject: [PATCH 16/33] fix(registry): preserve explicit device dispatch --- rl_engine/kernels/registry.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 4f1dcf86..8a6bb74c 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -718,6 +718,16 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + def _platform_for_device(self, device: torch.device | str | None) -> str: + if device is None: + return self._platform() + resolved = torch.device(device) + if resolved.type == "cuda": + return "rocm" if torch.version.hip is not None else "cuda" + if resolved.type in self._priority_map: + return resolved.type + return "cpu" + def register_logprob_backend( self, backend: OpBackend, From 39be7a59ba6f073cd13a50ef1cf1f5559b2f6a82 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 20 Aug 2026 19:18:36 +0800 Subject: [PATCH 17/33] feat(alignment): add compact debug matrix --- docs/design/ws2-attention-debug-matrix.md | 82 +++++++ docs/design/ws2-module-debug-matrix.md | 43 ++++ ...config_qwen3_8b_megatron_tp2_cp2_vllm.json | 11 +- .../alignment/cross_config/debug_matrix.py | 194 ++++++++++++++++ .../ops/pytorch/attention/debug_matrix.py | 215 ++++++++++++++++++ tests/test_attention_cross_config_binding.py | 2 + tests/test_attention_debug_matrix.py | 82 +++++++ tests/test_module_debug_matrix.py | 61 +++++ 8 files changed, 689 insertions(+), 1 deletion(-) create mode 100644 docs/design/ws2-attention-debug-matrix.md create mode 100644 docs/design/ws2-module-debug-matrix.md create mode 100644 rl_engine/alignment/cross_config/debug_matrix.py create mode 100644 rl_engine/kernels/ops/pytorch/attention/debug_matrix.py create mode 100644 tests/test_attention_debug_matrix.py create mode 100644 tests/test_module_debug_matrix.py diff --git a/docs/design/ws2-attention-debug-matrix.md b/docs/design/ws2-attention-debug-matrix.md new file mode 100644 index 00000000..32dacb79 --- /dev/null +++ b/docs/design/ws2-attention-debug-matrix.md @@ -0,0 +1,82 @@ +# Attention Debug Matrix + +The post-training Attention tool replays one frozen rollout, keeps the sample +identity fixed, changes one factor at a time, and reports both mismatch metrics +and a small set of invariance controls. +This is intentionally separate from the runtime knob catalog in +`examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json`. + +## Matrix shape + +`rl_engine.kernels.ops.pytorch.attention.debug_matrix` is the source of truth for +the compact matrix manifest: + +| Rows | Role | Meaning | +| --- | --- | --- | +| `A0` | baseline | Strict deterministic replay with the same train/rollout identity. | +| `A1`-`A3`, `A5`-`A7` | root-cause probes | One representative probe for each numerical or stateful Attention category. | +| `A4` | comparability gate | Head/sequence ownership must match before a numerical comparison is meaningful. | +| `C0`-`C2` | invariant controls | Cases that must remain exactly zero; a nonzero result invalidates the run. | + +The seven `A` rows are: + +| Row | Category | Representative probe | +| --- | --- | --- | +| `A1` | Position / RoPE | `position_ids` | +| `A2` | Q/K preprocessing | `qk_norm_disabled` | +| `A3` | Mask / sequence boundary | `causal_mask` | +| `A4` | Topology / head ownership | `tp_head_ownership`; reject as `comparable=false`, do not report a drift scalar. | +| `A5` | KV-cache identity / layout | `kv_page_order` | +| `A6` | Numerical policy | `accum_dtype` | +| `A7` | Distributed schedule | `merge_order` | + +The representative row is a fast first-line diagnosis. The existing taxonomy +keeps the secondary probes for a second pass, so users do not need to run all +21 probes for every post-training incident. + +## Replay contract + +Every row reuses the same: + +- checkpoint and model version; +- token IDs, selected-token IDs, masks, and positions; +- KV-cache and packing metadata; +- pre-update model state; +- train/rollout sample ordering. + +The matrix is one-at-a-time. It does not create a Cartesian product of Attention +knobs. Each diagnostic row records its own phase-local baseline, which allows +rollout-only and train-only debugging to be compared independently. A changed topology is a +gate failure (`comparable=false`), not a meaningful numerical drift sample. A changed +Split-KV plan or merge order is useful as a diagnostic injection (`A7`), but is not an +accepted production comparison until the actual plan again matches on both sides. + +## Metrics + +The replay report should include selected-token mismatch metrics: +`train_rollout_logprob_abs_diff`, forward `mismatch_kl`, and `mismatch_k3_kl`. +The Attention artifact additionally records `out`, `lse`, `dQ`, `dK`, and `dV` +maximum absolute drift. These are diagnostics; pass/fail still uses the fixed +repository numerical contract and never a user-supplied tolerance. + +`C0` (`tp_partition_control`), `C1` (`batch_composition_control`), and `C2` +(`prefill_decode_tail_control`) must be bitwise zero. They catch accidental +batch dependence, invalid preserved head ownership, and cache-position mistakes +before a root-cause row is trusted. + +## Runtime knobs versus debug probes + +The cross-configuration planner continues to own executable lifecycle knobs such +as TP/CP topology, CP communication, fusion boundary, and reduction policy. +Those knobs are materialized and read back by the runtime adapters. The compact +`A0`-`A7` matrix is a triage layer over that execution evidence; it does not +pretend that unsupported values such as `reduction_order=arrival` or BF16 CP +accumulation are production implementations. + +The portable manifest is available from: + +```python +from rl_engine.kernels.ops.pytorch.attention.debug_matrix import attention_debug_matrix + +manifest = attention_debug_matrix() +``` diff --git a/docs/design/ws2-module-debug-matrix.md b/docs/design/ws2-module-debug-matrix.md new file mode 100644 index 00000000..046fadb6 --- /dev/null +++ b/docs/design/ws2-module-debug-matrix.md @@ -0,0 +1,43 @@ +# Module Debug Matrix + +The cross-configuration debug surface uses a fixed replay and changes exactly +one factor at a time. It compares two edges of the same frozen rollout: + +1. training score versus rollout prefill; +2. rollout prefill versus rollout decode. + +The source of truth is +`rl_engine.alignment.cross_config.debug_matrix.module_debug_matrix`. The +manifest is reporting metadata, not a second collection of runtime flags. + +## Comparison gates + +Before reporting numerical drift, the run must prove that the selected tokens, +active mask, model state, and logical ownership map are unchanged. A failed +gate records `comparable=false`; it does not produce a drift number. This keeps +an input or sharding error from being misdiagnosed as kernel rounding. + +## First-line axes + +| Module | Baseline | Gate rows | Diagnostic rows | Invariant control | +| --- | --- | --- | --- | --- | +| Attention | `A0` | `A4` head/sequence ownership | `A1` position/RoPE, `A2` Q/K preprocessing, `A3` mask boundary, `A5` KV state, `A6` precision/rounding, `A7` block plan/merge | `C0`-`C2` | +| FFN / GEMM | `F0` | `F1` TP weight ownership | `F2` SwiGLU rounding, `F3` K-reduction/Split-K, `F4` token collective | `FC0` batch-invariant row replay | +| Selected-token logp | `L0` | `L1` vocabulary ownership/padding domain, `L2` selected token/active mask | `L3` vocabulary LSE tile/merge | `LC0` batch-invariant row replay | + +The logp wrapper receives local logits and therefore owns the vocabulary +log-sum-exp, selected-token gather, mask, padding, and TP merge. It does not +own the language-model-head GEMM. The FFN/GEMM wrapper owns the K reduction and +SwiGLU rounding; both must be diagnosed before attributing a logp difference to +the vocabulary merge. + +## Numerical program + +For each diagnostic row, record the same inputs and constants, intermediate +precision and final write, reduction grouping and logical order, and every +state handoff. Scheduler choices that do not change the numerical program or +logical ownership are intentionally not exposed as mismatch axes. + +Use the compact row first. The operator-specific secondary probes remain +available for a focused follow-up only after the representative row reproduces +the mismatch. diff --git a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json index 11f9cef1..ed668820 100644 --- a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json +++ b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json @@ -24,7 +24,16 @@ "prefill only; strict PR4 acceptance covers CP=2 prefill/chunked prefill.", "A decode request that becomes CP=1 is a blocking fallback and is tested", "separately by the PR6 logical KV replay harness." - ] + ], + "debug_matrix": { + "schema_version": "rlkernel.debug_matrix.v1", + "method": "fixed_replay_one_at_a_time", + "module_manifest": "rl_engine.alignment.cross_config.debug_matrix", + "modules": ["attention", "ffn", "logp"], + "comparison_edges": ["train_vs_rollout_prefill", "rollout_prefill_vs_decode"], + "replay_identity": "fixed rollout tokens, masks, positions, cache metadata, and pre-update state", + "cartesian_product": false + } }, "baseline": { "batch": { diff --git a/rl_engine/alignment/cross_config/debug_matrix.py b/rl_engine/alignment/cross_config/debug_matrix.py new file mode 100644 index 00000000..8f7b7041 --- /dev/null +++ b/rl_engine/alignment/cross_config/debug_matrix.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Compact, module-level mismatch axes for post-training drift triage. + +The manifest is deliberately a reporting contract. It names the first +diagnostic probe for each semantic operator without turning the probes into +runtime configuration or expanding a Cartesian product of settings. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from rl_engine.kernels.ops.pytorch.attention.debug_matrix import ( + ATTENTION_DEBUG_MATRIX, + ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION, +) +from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ATTENTION_DEBUG_AXES + + +DEBUG_MATRIX_SCHEMA_VERSION = "rlkernel.debug_matrix.v1" + + +@dataclass(frozen=True) +class ModuleDebugAxis: + """One first-line axis exposed by a semantic operator.""" + + module: str + axis_id: str + label: str + representative_probe: str + kind: str = "diagnostic" + + def __post_init__(self) -> None: + for field in ("module", "axis_id", "label", "representative_probe"): + value = getattr(self, field) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"module debug {field} must be a non-empty string") + if self.kind not in {"diagnostic", "gate"}: + raise ValueError(f"unknown module debug axis kind {self.kind!r}") + + def to_dict(self, *, row_id: str) -> dict[str, str]: + return { + "row": row_id, + "id": self.axis_id, + "label": self.label, + "representative_probe": self.representative_probe, + "kind": self.kind, + } + + +def _attention_axes() -> tuple[ModuleDebugAxis, ...]: + return tuple( + ModuleDebugAxis( + module="attention", + axis_id=axis.axis_id, + label=axis.label, + representative_probe=axis.representative_subprobe, + kind="gate" if axis.axis_id == "topology_head_ownership" else "diagnostic", + ) + for axis in ATTENTION_DEBUG_AXES + ) + + +MODULE_DEBUG_AXES: MappingProxyType[str, tuple[ModuleDebugAxis, ...]] = ( + MappingProxyType( + { + "attention": _attention_axes(), + "ffn": ( + ModuleDebugAxis( + "ffn", + "weight_shard_ownership", + "Weight shard / TP ownership", + "tp_weight_ownership", + "gate", + ), + ModuleDebugAxis( + "ffn", + "swiglu_rounding", + "SwiGLU intermediate rounding", + "swiglu_one_round", + ), + ModuleDebugAxis( + "ffn", + "gemm_reduction", + "GEMM K-reduction / Split-K policy", + "k_reduction_split_k", + ), + ModuleDebugAxis( + "ffn", + "token_collective", + "Token gather / reduce-scatter", + "sequence_parallel", + ), + ), + "logp": ( + ModuleDebugAxis( + "logp", + "vocab_shard_ownership", + "Vocabulary shard / TP ownership", + "vocab_shard_bounds", + "gate", + ), + ModuleDebugAxis( + "logp", + "selected_token_identity", + "Selected-token and active-mask identity", + "selected_token_active_mask", + "gate", + ), + ModuleDebugAxis( + "logp", + "vocab_lse_reduction", + "Vocabulary LSE tile / merge policy", + "vocab_tile_merge", + ), + ), + } + ) +) + + +_MODULE_DEBUG_AXIS_BY_ID = MappingProxyType( + { + (axis.module, axis.axis_id): axis + for axes in MODULE_DEBUG_AXES.values() + for axis in axes + } +) + + +def _module_rows(module: str) -> dict[str, Any]: + axes = MODULE_DEBUG_AXES[module] + baseline = {"attention": "A0", "ffn": "F0", "logp": "L0"}[module] + controls = { + "attention": ["C0", "C1", "C2"], + "ffn": ["FC0"], + "logp": ["LC0"], + }[module] + rows = [baseline] + [ + f"{module[0].upper()}{index}" for index in range(1, len(axes) + 1) + ] + return { + "baseline_row": baseline, + "rows": rows, + "invariant_controls": controls, + "axes": [ + axis.to_dict(row_id=rows[index]) for index, axis in enumerate(axes, start=1) + ], + } + + +def module_debug_matrix() -> dict[str, Any]: + """Return the portable matrix manifest shared by all three operators.""" + + return { + "schema_version": DEBUG_MATRIX_SCHEMA_VERSION, + "method": "fixed_replay_one_at_a_time", + "cartesian_product": False, + "comparison_edges": [ + "train_vs_rollout_prefill", + "rollout_prefill_vs_decode", + ], + "replay_identity": ( + "same checkpoint, token IDs, selected-token IDs, masks, positions, " + "cache metadata, and pre-update model state" + ), + "modules": {module: _module_rows(module) for module in MODULE_DEBUG_AXES}, + "attention_compatibility": { + "schema_version": ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION, + "rows": [row.to_dict() for row in ATTENTION_DEBUG_MATRIX], + }, + } + + +def module_debug_axis(module: str, axis_id: str) -> ModuleDebugAxis: + """Look up one stable axis by module and identifier.""" + + try: + return _MODULE_DEBUG_AXIS_BY_ID[(module.strip(), axis_id.strip())] + except (AttributeError, KeyError) as exc: + raise ValueError(f"unknown module debug axis {module!r}/{axis_id!r}") from exc + + +__all__ = [ + "DEBUG_MATRIX_SCHEMA_VERSION", + "MODULE_DEBUG_AXES", + "ModuleDebugAxis", + "module_debug_axis", + "module_debug_matrix", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py b/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py new file mode 100644 index 00000000..306e1db9 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay matrix for post-training Attention drift triage. + +The matrix deliberately separates a fixed replay baseline, one-at-a-time root- +cause probes, and invariant controls. It is a reporting/debug contract, not a +second runtime knob catalog and it never implies a Cartesian product. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from rl_engine.kernels.attention_contract import AttentionContractError +from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ( + ATTENTION_DEBUG_AXES, + ATTENTION_INVARIANT_CONTROLS, +) + +ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION = "rlkernel.attention.debug_matrix.v1" + + +@dataclass(frozen=True) +class AttentionDebugMatrixRow: + """One independent case in the replay matrix.""" + + row_id: str + label: str + category: str + probe: str | None = None + root_cause_axis: str | None = None + expected: str = "diagnostic" + + def __post_init__(self) -> None: + if not self.row_id.strip() or not self.label.strip(): + raise ValueError("Attention debug matrix rows need an id and label") + if self.category not in { + "baseline", + "root_cause", + "comparability_gate", + "invariant_control", + }: + raise ValueError( + f"unknown Attention debug matrix category {self.category!r}" + ) + if self.category == "baseline": + if self.probe is not None or self.root_cause_axis is not None: + raise ValueError("baseline rows cannot name a probe or root-cause axis") + if self.expected != "baseline": + raise ValueError("baseline rows must use expected='baseline'") + elif not self.probe or not self.probe.strip(): + raise ValueError("non-baseline rows need a probe") + if self.category in {"root_cause", "comparability_gate"}: + if not self.root_cause_axis: + raise ValueError("root-cause and gate rows need root_cause_axis") + if self.category == "root_cause": + if self.expected != "diagnostic": + raise ValueError("root-cause rows must use expected='diagnostic'") + if self.category == "comparability_gate" and self.expected != "rejected": + raise ValueError("comparability gates must use expected='rejected'") + if self.category == "invariant_control": + if self.expected != "exact_zero": + raise ValueError("invariant controls must use expected='exact_zero'") + + def to_dict(self) -> dict[str, Any]: + return { + "id": self.row_id, + "label": self.label, + "category": self.category, + "probe": self.probe, + "root_cause_axis": self.root_cause_axis, + "expected": self.expected, + } + + +def _build_rows() -> tuple[AttentionDebugMatrixRow, ...]: + rows: list[AttentionDebugMatrixRow] = [ + AttentionDebugMatrixRow( + row_id="A0", + label="Strict replay baseline", + category="baseline", + expected="baseline", + ) + ] + for index, axis in enumerate(ATTENTION_DEBUG_AXES, start=1): + is_gate = axis.axis_id == "topology_head_ownership" + rows.append( + AttentionDebugMatrixRow( + row_id=f"A{index}", + label=axis.label, + category="comparability_gate" if is_gate else "root_cause", + probe=axis.representative_subprobe, + root_cause_axis=axis.axis_id, + expected="rejected" if is_gate else "diagnostic", + ) + ) + for index, probe in enumerate(ATTENTION_INVARIANT_CONTROLS): + rows.append( + AttentionDebugMatrixRow( + row_id=f"C{index}", + label="Invariant control", + category="invariant_control", + probe=probe, + expected="exact_zero", + ) + ) + return tuple(rows) + + +ATTENTION_DEBUG_MATRIX = _build_rows() +_ROWS_BY_ID = MappingProxyType({row.row_id: row for row in ATTENTION_DEBUG_MATRIX}) + + +def validate_attention_debug_matrix() -> None: + """Validate matrix coverage and the no-Cartesian-product invariant.""" + + if len(_ROWS_BY_ID) != len(ATTENTION_DEBUG_MATRIX): + raise RuntimeError("Attention debug matrix row IDs must be unique") + baselines = [row for row in ATTENTION_DEBUG_MATRIX if row.category == "baseline"] + if len(baselines) != 1 or baselines[0].row_id != "A0": + raise RuntimeError( + "Attention debug matrix must contain exactly one A0 baseline" + ) + + axis_ids = {axis.axis_id for axis in ATTENTION_DEBUG_AXES} + diagnostic_rows = [ + row + for row in ATTENTION_DEBUG_MATRIX + if row.category in {"root_cause", "comparability_gate"} + ] + if {row.root_cause_axis for row in diagnostic_rows} != axis_ids: + raise RuntimeError("every root-cause axis must have one representative row") + if len(diagnostic_rows) != len(axis_ids): + raise RuntimeError("root-cause representatives must be one-at-a-time") + gates = [row for row in diagnostic_rows if row.category == "comparability_gate"] + if [row.root_cause_axis for row in gates] != ["topology_head_ownership"]: + raise RuntimeError("only topology/head ownership is a comparability gate") + + controls = [ + row for row in ATTENTION_DEBUG_MATRIX if row.category == "invariant_control" + ] + if {row.probe for row in controls} != set(ATTENTION_INVARIANT_CONTROLS): + raise RuntimeError("invariant-control coverage is incomplete") + if any(row.expected != "exact_zero" for row in controls): + raise RuntimeError("invariant controls must be exact-zero checks") + + axis_by_id = {axis.axis_id: axis for axis in ATTENTION_DEBUG_AXES} + for row in diagnostic_rows: + axis = axis_by_id[row.root_cause_axis] + if row.probe != axis.representative_subprobe: + raise RuntimeError( + f"matrix row {row.row_id} is not the representative probe for {axis.axis_id}" + ) + + +validate_attention_debug_matrix() + + +def attention_debug_matrix() -> dict[str, Any]: + """Return the portable matrix manifest used by reports and tooling.""" + + return { + "schema_version": ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION, + "method": "fixed_replay_one_at_a_time", + "baseline_row": "A0", + "cartesian_product": False, + "replay_identity": ( + "same checkpoint, token IDs, selected-token IDs, masks, positions, " + "cache metadata, and pre-update model state" + ), + "row_baseline": "each diagnostic row is compared with its own phase-local A0 baseline", + "metrics": [ + "train_rollout_logprob_abs_diff", + "mismatch_kl", + "mismatch_k3_kl", + "out_max_abs", + "lse_max_abs", + "dq_max_abs", + "dk_max_abs", + "dv_max_abs", + ], + "rows": [row.to_dict() for row in ATTENTION_DEBUG_MATRIX], + "topology_gate": [ + "checkpoint/model/token identity", + "TP head ownership", + "CP sequence ownership", + "actual Split-KV plan", + ], + } + + +def attention_debug_matrix_row(row_id: str) -> AttentionDebugMatrixRow: + """Look up a stable matrix row by its report ID.""" + + if not isinstance(row_id, str) or not row_id.strip(): + raise AttentionContractError("Attention debug matrix row ID must be non-empty") + try: + return _ROWS_BY_ID[row_id.strip()] + except KeyError as exc: + raise AttentionContractError( + f"unknown Attention debug matrix row {row_id!r}" + ) from exc + + +__all__ = [ + "ATTENTION_DEBUG_MATRIX", + "ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION", + "AttentionDebugMatrixRow", + "attention_debug_matrix", + "attention_debug_matrix_row", + "validate_attention_debug_matrix", +] diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index c1c8d35d..429f654f 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -1274,6 +1274,8 @@ def test_scenario_uses_megatron_vocabulary_only(): assert training["tensor_parallel_size"] == 2 assert training["context_parallel_size"] == 2 assert config["baseline"]["rollout"]["batch_invariant"] is True + assert config["scenario"]["debug_matrix"]["modules"] == ["attention", "ffn", "logp"] + assert config["scenario"]["debug_matrix"]["cartesian_product"] is False def test_scenario_knob_paths_all_exist(): diff --git a/tests/test_attention_debug_matrix.py b/tests/test_attention_debug_matrix.py new file mode 100644 index 00000000..86710882 --- /dev/null +++ b/tests/test_attention_debug_matrix.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from rl_engine.kernels.attention_contract import AttentionContractError +from rl_engine.kernels.ops.pytorch.attention.debug_matrix import ( + ATTENTION_DEBUG_MATRIX, + ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION, + attention_debug_matrix, + attention_debug_matrix_row, +) + + +def test_attention_debug_matrix_is_fixed_replay_oat_with_controls(): + manifest = attention_debug_matrix() + + assert manifest["schema_version"] == ATTENTION_DEBUG_MATRIX_SCHEMA_VERSION + assert manifest["method"] == "fixed_replay_one_at_a_time" + assert manifest["baseline_row"] == "A0" + assert manifest["cartesian_product"] is False + assert [row["id"] for row in manifest["rows"]] == [ + "A0", + "A1", + "A2", + "A3", + "A4", + "A5", + "A6", + "A7", + "C0", + "C1", + "C2", + ] + assert sum(row["category"] == "root_cause" for row in manifest["rows"]) == 6 + assert sum(row["category"] == "comparability_gate" for row in manifest["rows"]) == 1 + assert sum(row["category"] == "invariant_control" for row in manifest["rows"]) == 3 + + +def test_attention_debug_matrix_rows_have_stable_semantics(): + baseline = attention_debug_matrix_row("A0") + assert baseline.category == "baseline" + assert baseline.expected == "baseline" + + root = attention_debug_matrix_row("A1") + assert root.category == "root_cause" + assert root.probe == "position_ids" + assert root.root_cause_axis == "position_rope" + assert root.expected == "diagnostic" + + topology = attention_debug_matrix_row("A4") + assert topology.category == "comparability_gate" + assert topology.expected == "rejected" + + control = attention_debug_matrix_row("C0") + assert control.category == "invariant_control" + assert control.probe == "tp_partition_control" + assert control.expected == "exact_zero" + + with pytest.raises( + AttentionContractError, match="unknown Attention debug matrix row" + ): + attention_debug_matrix_row("A8") + + +def test_matrix_manifest_matches_taxonomy_representatives(): + from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ( + ATTENTION_DEBUG_AXES, + ) + + root_rows = [ + row + for row in ATTENTION_DEBUG_MATRIX + if row.category in {"root_cause", "comparability_gate"} + ] + assert [row.root_cause_axis for row in root_rows] == [ + axis.axis_id for axis in ATTENTION_DEBUG_AXES + ] + assert [row.probe for row in root_rows] == [ + axis.representative_subprobe for axis in ATTENTION_DEBUG_AXES + ] diff --git a/tests/test_module_debug_matrix.py b/tests/test_module_debug_matrix.py new file mode 100644 index 00000000..3830b783 --- /dev/null +++ b/tests/test_module_debug_matrix.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +from rl_engine.alignment.cross_config.debug_matrix import ( + DEBUG_MATRIX_SCHEMA_VERSION, + MODULE_DEBUG_AXES, + module_debug_axis, + module_debug_matrix, +) + + +def test_module_matrix_is_fixed_replay_and_compact(): + manifest = module_debug_matrix() + + assert manifest["schema_version"] == DEBUG_MATRIX_SCHEMA_VERSION + assert manifest["method"] == "fixed_replay_one_at_a_time" + assert manifest["cartesian_product"] is False + assert manifest["comparison_edges"] == [ + "train_vs_rollout_prefill", + "rollout_prefill_vs_decode", + ] + assert set(manifest["modules"]) == {"attention", "ffn", "logp"} + assert manifest["modules"]["attention"]["rows"] == [ + "A0", + "A1", + "A2", + "A3", + "A4", + "A5", + "A6", + "A7", + ] + assert manifest["modules"]["ffn"]["rows"] == ["F0", "F1", "F2", "F3", "F4"] + assert manifest["modules"]["logp"]["rows"] == ["L0", "L1", "L2", "L3"] + assert [row["row"] for row in manifest["modules"]["ffn"]["axes"]] == [ + "F1", + "F2", + "F3", + "F4", + ] + + +def test_module_axes_mark_identity_and_topology_as_gates(): + assert module_debug_axis("attention", "topology_head_ownership").kind == "gate" + assert module_debug_axis("ffn", "weight_shard_ownership").kind == "gate" + assert module_debug_axis("logp", "vocab_shard_ownership").kind == "gate" + assert module_debug_axis("logp", "selected_token_identity").kind == "gate" + assert module_debug_axis("ffn", "gemm_reduction").kind == "diagnostic" + + with pytest.raises(ValueError, match="unknown module debug axis"): + module_debug_axis("ffn", "not_an_axis") + + +def test_module_axes_have_unique_ids_and_probes(): + for module, axes in MODULE_DEBUG_AXES.items(): + assert all(axis.module == module for axis in axes) + assert len({axis.axis_id for axis in axes}) == len(axes) + assert len({axis.representative_probe for axis in axes}) == len(axes) From 2cf8d5dab4a92e0a9067f0448ee5e4142e4dc0fa Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 20 Aug 2026 22:08:25 +0800 Subject: [PATCH 18/33] feat(attention): reuse qualified TE QK normalization --- .../cross_config/attention_binding.py | 73 ++-- rl_engine/kernels/attention_preprocess.py | 354 +++++++++++++----- rl_engine/kernels/attention_projection.py | 12 +- tests/test_attention_cross_config_binding.py | 15 +- tests/test_attention_preprocess.py | 80 +++- tests/test_attention_projection.py | 13 + 6 files changed, 419 insertions(+), 128 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 13f87000..4249c314 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -57,13 +57,16 @@ validate_split_kv_plan_set_alignment, ) from rl_engine.kernels.attention_preprocess import ( + ALLOWED_ATTENTION_PREPROCESS_BACKENDS, MANDATED_ATTENTION_PREPROCESS_BACKENDS, PREPROCESS_POLICY_ID, ) from rl_engine.kernels.attention_projection import ( + CUDA_DETERMINISTIC_PROJECTION_BACKEND_ID, O_PROJ_COLLECTIVE_CONTRACT, PROJECTION_POLICY_ID, QKV_COLLECTIVE_CONTRACT, + ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, ) __all__ = [ @@ -976,21 +979,20 @@ def bind_attention_runtime_readbacks( message="QK-Norm/RoPE policy IDs differ between runtimes", ) ) - if rollout.preprocess_fallback and training.preprocess_fallback: - for name in MANDATED_ATTENTION_PREPROCESS_BACKENDS: - rollout_backend = rollout.preprocess_backends.get(name) - training_backend = training.preprocess_backends.get(name) - if rollout_backend != training_backend: - missing_scope_evidence.append( - BindingIssue( - code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, - tier=BindingTier.SEMANTIC, - field=f"preprocess.{name}", - rollout=rollout_backend, - training=training_backend, - message="fallback sides must use the same deterministic preprocess backend", - ) + for name in MANDATED_ATTENTION_PREPROCESS_BACKENDS: + rollout_backend = rollout.preprocess_backends.get(name) + training_backend = training.preprocess_backends.get(name) + if rollout_backend != training_backend: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"preprocess.{name}", + rollout=rollout_backend, + training=training_backend, + message="training and rollout must execute the same preprocess backend", ) + ) for projection in ("qkv", "o_proj"): rollout_plan = rollout.projection_plans.get(projection, {}) training_plan = training.projection_plans.get(projection, {}) @@ -1009,22 +1011,18 @@ def bind_attention_runtime_readbacks( ), ) ) - if rollout_fallback and training_fallback: - for field_name in ("backend_id", "policy_id", "split_k", "reduction_order"): - if rollout_plan.get(field_name) != training_plan.get(field_name): - missing_scope_evidence.append( - BindingIssue( - code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, - tier=BindingTier.SEMANTIC, - field=f"projection.{projection}.{field_name}", - rollout=rollout_plan.get(field_name), - training=training_plan.get(field_name), - message=( - "deterministic projection fallback evidence differs " - "between sides" - ), - ) + for field_name in ("backend_id", "policy_id", "split_k", "reduction_order"): + if rollout_plan.get(field_name) != training_plan.get(field_name): + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"projection.{projection}.{field_name}", + rollout=rollout_plan.get(field_name), + training=training_plan.get(field_name), + message="projection execution evidence differs between sides", ) + ) return bind_attention_contracts( rollout_contract=rollout.contract, training_contract=training.contract, @@ -1256,7 +1254,7 @@ def _attention_preprocess_issues( ), ) ) - elif actual != mandated and not actual.startswith("native."): + elif actual not in ALLOWED_ATTENTION_PREPROCESS_BACKENDS[name]: issues.append( BindingIssue( code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, @@ -1266,13 +1264,13 @@ def _attention_preprocess_issues( training=actual if side == "training" else None, message=( f"{side} executed {actual!r}; " - f"the H100 experiment requires {mandated!r} or a verified native backend" + f"the experiment requires {mandated!r} or a verified platform backend" ), ) ) if readback.preprocess_fallback and any( - readback.preprocess_backends.get(name) != mandated - for name, mandated in MANDATED_ATTENTION_PREPROCESS_BACKENDS.items() + str(readback.preprocess_backends.get(name, "")).startswith("transformer_engine.") + for name in MANDATED_ATTENTION_PREPROCESS_BACKENDS ): issues.append( BindingIssue( @@ -1280,8 +1278,7 @@ def _attention_preprocess_issues( tier=BindingTier.SEMANTIC, field=f"{side}.preprocess.fallback", message=( - f"{side} reported a fallback but did not use the common deterministic " - "QK-Norm/RoPE backends" + f"{side} reported fallback while still claiming a vendor preprocess backend" ), ) ) @@ -1361,7 +1358,11 @@ def _attention_projection_issues( ) ) if plan.get("fallback"): - if backend_id != "rlkernel.cuda.det_gemm" or not plan.get("fallback_reason"): + deterministic_backends = { + CUDA_DETERMINISTIC_PROJECTION_BACKEND_ID, + ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, + } + if backend_id not in deterministic_backends or not plan.get("fallback_reason"): issues.append( BindingIssue( code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index a05a8d41..6d1dda0d 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -1,14 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Bitwise-bound QK-Norm and RoPE handoff for WS2 Attention. - -Megatron/TE and vLLM/FlashInfer remain the first-choice implementations. They -are admitted only after the same-input H100 probe is bitwise identical to the -deterministic RL-Kernel path. A failed probe (or an unavailable native backend) -selects the common RL-Kernel path for both sides and records why the fallback -was taken. The caller must pass this readback to the cross-config binder. -""" +"""Bitwise-bound QK-Norm and RoPE handoff for WS2 Attention. + +Transformer Engine RMSNorm is the first-choice QK-Norm implementation on CUDA +and ROCm. It is admitted only after a same-input bitwise probe against the +platform RL-Kernel path. RoPE remains on the shared RL-Kernel implementation so +training and paged-cache rollout expose the same post-RoPE boundary. +""" from __future__ import annotations @@ -16,22 +15,123 @@ import json from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Callable, Mapping - -import torch -from torch import Tensor - -QK_RMSNORM_BACKEND_ID = "rlkernel.cuda.rmsnorm" -ROPE_BACKEND_ID = "rlkernel.cuda.rope_sm90" -NATIVE_QK_RMSNORM_BACKEND_ID = "native.qk_rmsnorm" -NATIVE_ROPE_BACKEND_ID = "native.rope" -PREPROCESS_POLICY_ID = "ws2.attention.preprocess.v2" -MANDATED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( +from typing import Any, Callable, Mapping + +import torch +from torch import Tensor +from torch.autograd import Function + +QK_RMSNORM_BACKEND_ID = "rlkernel.cuda.rmsnorm" +ROPE_BACKEND_ID = "rlkernel.cuda.rope_sm90" +ROCM_QK_RMSNORM_BACKEND_ID = "rlkernel.rocm.triton_rmsnorm" +ROCM_ROPE_BACKEND_ID = "rlkernel.rocm.deterministic_rope" +TE_CUDA_QK_RMSNORM_BACKEND_ID = "transformer_engine.cuda.rmsnorm" +TE_ROCM_QK_RMSNORM_BACKEND_ID = "transformer_engine.rocm.rmsnorm" +NATIVE_QK_RMSNORM_BACKEND_ID = TE_CUDA_QK_RMSNORM_BACKEND_ID +NATIVE_ROPE_BACKEND_ID = "native.rope" +PREPROCESS_POLICY_ID = "ws2.attention.preprocess.v3" +MANDATED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( { "qk_rmsnorm": QK_RMSNORM_BACKEND_ID, "rope": ROPE_BACKEND_ID, } -) +) +ROCM_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( + { + "qk_rmsnorm": ROCM_QK_RMSNORM_BACKEND_ID, + "rope": ROCM_ROPE_BACKEND_ID, + } +) +ALLOWED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, frozenset[str]] = MappingProxyType( + { + "qk_rmsnorm": frozenset( + { + QK_RMSNORM_BACKEND_ID, + ROCM_QK_RMSNORM_BACKEND_ID, + TE_CUDA_QK_RMSNORM_BACKEND_ID, + TE_ROCM_QK_RMSNORM_BACKEND_ID, + } + ), + "rope": frozenset({ROPE_BACKEND_ID, ROCM_ROPE_BACKEND_ID, NATIVE_ROPE_BACKEND_ID}), + } +) + + +class TransformerEngineRMSNormUnavailable(RuntimeError): + """Raised when the exact TE RMSNorm functional contract is unavailable.""" + + +def _load_transformer_engine_rmsnorm(): + try: + import transformer_engine + import transformer_engine.pytorch # noqa: F401 - loads the platform extension + from transformer_engine.pytorch.constants import TE_DType + from transformer_engine.pytorch.cpp_extensions import rmsnorm_bwd, rmsnorm_fwd + except (ImportError, OSError, RuntimeError) as exc: + raise TransformerEngineRMSNormUnavailable( + "Transformer Engine RMSNorm forward/backward is unavailable" + ) from exc + return rmsnorm_fwd, rmsnorm_bwd, TE_DType, getattr(transformer_engine, "__version__", "unknown") + + +class _TransformerEngineRMSNormFunction(Function): + @staticmethod + def forward(ctx, x: Tensor, weight: Tensor, eps: float) -> Tensor: + rmsnorm_fwd, rmsnorm_bwd, te_dtype, _version = _load_transformer_engine_rmsnorm() + original_shape = x.shape + hidden = original_shape[-1] + x_2d = x.contiguous().view(-1, hidden) + weight_1d = weight.contiguous().view(hidden) + try: + out, _unused, rsigma = rmsnorm_fwd( + x_2d, + weight_1d, + float(eps), + None, + None, + te_dtype[x.dtype], + 0, + False, + ) + except (KeyError, TypeError, ValueError, RuntimeError) as exc: + raise TransformerEngineRMSNormUnavailable( + "installed Transformer Engine has an incompatible RMSNorm API" + ) from exc + ctx.save_for_backward(x_2d, rsigma, weight_1d) + ctx.rmsnorm_bwd = rmsnorm_bwd + ctx.original_shape = original_shape + return out.view(original_shape) + + @staticmethod + def backward(ctx, grad_output: Tensor): + x_2d, rsigma, weight = ctx.saved_tensors + dy = grad_output.contiguous().view_as(x_2d) + try: + dx, dw = ctx.rmsnorm_bwd(dy, x_2d, rsigma, weight, 0, False) + except (TypeError, ValueError, RuntimeError) as exc: + raise TransformerEngineRMSNormUnavailable( + "installed Transformer Engine has an incompatible RMSNorm backward API" + ) from exc + return dx.view(ctx.original_shape), dw.view_as(weight), None + + +class TransformerEngineRMSNormOp: + """External-weight TE RMSNorm used by both training and rollout adapters.""" + + def __init__(self) -> None: + _fwd, _bwd, _dtype, version = _load_transformer_engine_rmsnorm() + platform = "rocm" if torch.version.hip is not None else "cuda" + self.backend_id = f"transformer_engine.{platform}.rmsnorm" + self.package_version = version + + def __call__(self, x: Tensor, weight: Tensor, *, eps: float = 1.0e-6) -> Tensor: + if x.dtype not in (torch.float16, torch.bfloat16) or weight.dtype != x.dtype: + raise TypeError("Transformer Engine RMSNorm requires matching FP16/BF16 input and weight") + if not x.is_cuda or not weight.is_cuda or x.device != weight.device: + raise ValueError("Transformer Engine RMSNorm requires input and weight on one GPU") + if weight.shape != (x.shape[-1],): + raise ValueError("RMSNorm weight must match the input hidden dimension") + return _TransformerEngineRMSNormFunction.apply(x, weight, float(eps)) @dataclass(frozen=True) @@ -69,8 +169,8 @@ def readback_fields(self) -> dict[str, Any]: } -class H100AttentionPreprocessor: - """Apply native QK-Norm/RoPE when the H100 probe passes. +class H100AttentionPreprocessor: + """Reuse TE QK-Norm with the RL-Kernel SM90 RoPE boundary. ``native_qk_norm`` and ``native_rope`` are framework-owned callables. They are intentionally injected instead of importing TE/vLLM here, so the same @@ -84,9 +184,11 @@ def __init__( *, native_qk_norm: Callable[..., Tensor] | None = None, native_rope: Callable[..., Tensor] | None = None, - native_qk_norm_backend_id: str = NATIVE_QK_RMSNORM_BACKEND_ID, - native_rope_backend_id: str = NATIVE_ROPE_BACKEND_ID, - policy_id: str = PREPROCESS_POLICY_ID, + native_qk_norm_backend_id: str = NATIVE_QK_RMSNORM_BACKEND_ID, + native_rope_backend_id: str = NATIVE_ROPE_BACKEND_ID, + reuse_transformer_engine_qk_norm: bool = True, + require_transformer_engine_qk_norm: bool = False, + policy_id: str = PREPROCESS_POLICY_ID, ) -> None: if not torch.cuda.is_available(): raise RuntimeError("H100AttentionPreprocessor requires an available CUDA runtime") @@ -114,12 +216,20 @@ def __init__( from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op - self.rmsnorm = RMSNormCudaOp() - self.rope = RoPESM90Op() - if not isinstance(policy_id, str) or not policy_id.strip(): - raise ValueError("policy_id must be a non-empty string") - self.native_qk_norm = native_qk_norm - self.native_rope = native_rope + self.rmsnorm = RMSNormCudaOp() + self.rope = RoPESM90Op() + self.deterministic_backend_ids = MANDATED_ATTENTION_PREPROCESS_BACKENDS + if not isinstance(policy_id, str) or not policy_id.strip(): + raise ValueError("policy_id must be a non-empty string") + if native_qk_norm is None and reuse_transformer_engine_qk_norm: + try: + native_qk_norm = TransformerEngineRMSNormOp() + native_qk_norm_backend_id = native_qk_norm.backend_id + except TransformerEngineRMSNormUnavailable: + if require_transformer_engine_qk_norm: + raise + self.native_qk_norm = native_qk_norm + self.native_rope = native_rope self.native_qk_norm_backend_id = native_qk_norm_backend_id self.native_rope_backend_id = native_rope_backend_id self.policy_id = policy_id @@ -159,55 +269,120 @@ def forward( _validate_inputs(q, k, q_weight, k_weight, positions, self.device) q_norm_det = self.rmsnorm(q, q_weight, eps=eps) k_norm_det = self.rmsnorm(k, k_weight, eps=eps) - q_det = _apply_deterministic_rope(self.rope, q_norm_det, positions, theta) - k_det = _apply_deterministic_rope(self.rope, k_norm_det, positions, theta) - - native_qk_norm = self.native_qk_norm - native_rope = self.native_rope - if native_qk_norm is not None and native_rope is not None: - try: - q_norm_native = native_qk_norm(q, q_weight, eps=eps) - k_norm_native = native_qk_norm(k, k_weight, eps=eps) - q_native = native_rope(q_norm_native, positions, theta=theta) - k_native = native_rope(k_norm_native, positions, theta=theta) - probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) - if ( - torch.equal(q_norm_native, q_norm_det) - and torch.equal(k_norm_native, k_norm_det) - and torch.equal(q_native, q_det) - and torch.equal(k_native, k_det) - ): - return AttentionPreprocessResult( - q=q_native, - k=k_native, - backend_ids=MappingProxyType( - { - "qk_rmsnorm": self.native_qk_norm_backend_id, - "rope": self.native_rope_backend_id, - } - ), + q_det = _apply_deterministic_rope(self.rope, q_norm_det, positions, theta) + k_det = _apply_deterministic_rope(self.rope, k_norm_det, positions, theta) + probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) + + native_qk_norm = self.native_qk_norm + native_rope = self.native_rope + if native_qk_norm is not None: + try: + q_norm_native = native_qk_norm(q, q_weight, eps=eps) + k_norm_native = native_qk_norm(k, k_weight, eps=eps) + norm_matches = torch.equal(q_norm_native, q_norm_det) and torch.equal( + k_norm_native, k_norm_det + ) + if norm_matches and native_rope is None: + return AttentionPreprocessResult( + q=_apply_deterministic_rope( + self.rope, q_norm_native, positions, theta + ), + k=_apply_deterministic_rope( + self.rope, k_norm_native, positions, theta + ), + backend_ids=MappingProxyType( + { + "qk_rmsnorm": self.native_qk_norm_backend_id, + "rope": self.deterministic_backend_ids["rope"], + } + ), fallback=False, device_capability=self.device_capability, probe_id=probe_id, - policy_id=self.policy_id, - ) - fallback_reason = "native_preprocess_bitwise_probe_failed" - except Exception as exc: # framework backend failures must fail over together - fallback_reason = f"native_preprocess_unavailable:{type(exc).__name__}" - probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) - else: - fallback_reason = "native_preprocess_not_supplied" - probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) - return AttentionPreprocessResult( - q=q_det, - k=k_det, - backend_ids=MANDATED_ATTENTION_PREPROCESS_BACKENDS, + policy_id=self.policy_id, + ) + if norm_matches and native_rope is not None: + q_native = native_rope(q_norm_native, positions, theta=theta) + k_native = native_rope(k_norm_native, positions, theta=theta) + if torch.equal(q_native, q_det) and torch.equal(k_native, k_det): + return AttentionPreprocessResult( + q=q_native, + k=k_native, + backend_ids=MappingProxyType( + { + "qk_rmsnorm": self.native_qk_norm_backend_id, + "rope": self.native_rope_backend_id, + } + ), + fallback=False, + device_capability=self.device_capability, + probe_id=probe_id, + policy_id=self.policy_id, + ) + fallback_reason = "native_preprocess_bitwise_probe_failed" + except Exception as exc: # framework backend failures use the common path + fallback_reason = f"native_preprocess_unavailable:{type(exc).__name__}" + else: + fallback_reason = "native_preprocess_not_supplied" + return AttentionPreprocessResult( + q=q_det, + k=k_det, + backend_ids=self.deterministic_backend_ids, fallback=True, device_capability=self.device_capability, fallback_reason=fallback_reason, probe_id=probe_id, - policy_id=self.policy_id, - ) + policy_id=self.policy_id, + ) + + +class RocmAttentionPreprocessor(H100AttentionPreprocessor): + """Reuse TE QK-Norm with the RL-Kernel deterministic ROCm RoPE boundary.""" + + def __init__( + self, + device: torch.device | str | int | None = None, + *, + native_qk_norm: Callable[..., Tensor] | None = None, + reuse_transformer_engine_qk_norm: bool = True, + require_transformer_engine_qk_norm: bool = True, + policy_id: str = PREPROCESS_POLICY_ID, + ) -> None: + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("RocmAttentionPreprocessor requires an available ROCm runtime") + current_device = torch.cuda.current_device() + self.device = torch.device("cuda", current_device) + if device is not None: + self.device = ( + torch.device("cuda", device) if isinstance(device, int) else torch.device(device) + ) + if self.device.type != "cuda": + raise RuntimeError(f"RocmAttentionPreprocessor requires ROCm, got {self.device}") + if self.device.index is None: + self.device = torch.device("cuda", current_device) + + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RocmDeterministicRoPEOp + from rl_engine.kernels.ops.triton.rmsnorm_triton import RMSNormTritonOp + + self.rmsnorm = RMSNormTritonOp() + self.rope = RocmDeterministicRoPEOp() + self.deterministic_backend_ids = ROCM_ATTENTION_PREPROCESS_BACKENDS + self.device_capability = (0, 0) + if native_qk_norm is None and reuse_transformer_engine_qk_norm: + try: + native_qk_norm = TransformerEngineRMSNormOp() + except TransformerEngineRMSNormUnavailable: + if require_transformer_engine_qk_norm: + raise + self.native_qk_norm = native_qk_norm + self.native_rope = None + self.native_qk_norm_backend_id = ( + getattr(native_qk_norm, "backend_id", TE_ROCM_QK_RMSNORM_BACKEND_ID) + if native_qk_norm is not None + else TE_ROCM_QK_RMSNORM_BACKEND_ID + ) + self.native_rope_backend_id = ROCM_ROPE_BACKEND_ID + self.policy_id = policy_id def _apply_deterministic_rope( @@ -261,8 +436,8 @@ def _validate_inputs( raise ValueError("q and k must use [B, H, S, D] layout") if q.shape[0] != k.shape[0] or q.shape[-2:] != k.shape[-2:]: raise ValueError("q and k must have the same batch, sequence, and head dimensions") - if q.dtype is not torch.bfloat16 or k.dtype is not torch.bfloat16: - raise TypeError("the frozen H100 Attention experiment requires BF16 q and k") + if q.dtype is not torch.bfloat16 or k.dtype is not torch.bfloat16: + raise TypeError("the frozen Attention experiment requires BF16 q and k") if q.device != device or k.device != device: raise ValueError(f"q and k must both be on the configured device {device}") for name, weight in (("q_weight", q_weight), ("k_weight", k_weight)): @@ -279,13 +454,22 @@ def _validate_inputs( raise ValueError(f"positions must have shape [S] or [B, S], expected {expected}") -__all__ = [ - "AttentionPreprocessResult", - "H100AttentionPreprocessor", - "MANDATED_ATTENTION_PREPROCESS_BACKENDS", - "QK_RMSNORM_BACKEND_ID", - "ROPE_BACKEND_ID", - "NATIVE_QK_RMSNORM_BACKEND_ID", - "NATIVE_ROPE_BACKEND_ID", - "PREPROCESS_POLICY_ID", -] +__all__ = [ + "ALLOWED_ATTENTION_PREPROCESS_BACKENDS", + "AttentionPreprocessResult", + "H100AttentionPreprocessor", + "MANDATED_ATTENTION_PREPROCESS_BACKENDS", + "ROCM_ATTENTION_PREPROCESS_BACKENDS", + "RocmAttentionPreprocessor", + "QK_RMSNORM_BACKEND_ID", + "ROCM_QK_RMSNORM_BACKEND_ID", + "ROCM_ROPE_BACKEND_ID", + "ROPE_BACKEND_ID", + "NATIVE_QK_RMSNORM_BACKEND_ID", + "NATIVE_ROPE_BACKEND_ID", + "PREPROCESS_POLICY_ID", + "TE_CUDA_QK_RMSNORM_BACKEND_ID", + "TE_ROCM_QK_RMSNORM_BACKEND_ID", + "TransformerEngineRMSNormOp", + "TransformerEngineRMSNormUnavailable", +] diff --git a/rl_engine/kernels/attention_projection.py b/rl_engine/kernels/attention_projection.py index a7a91bfe..1e23ceda 100644 --- a/rl_engine/kernels/attention_projection.py +++ b/rl_engine/kernels/attention_projection.py @@ -27,6 +27,8 @@ QKV_PROJECTION = "qkv" O_PROJ_PROJECTION = "o_proj" PROJECTION_POLICY_ID = "ws2.attention.projection.v1" +CUDA_DETERMINISTIC_PROJECTION_BACKEND_ID = "rlkernel.cuda.det_gemm" +ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID = "rlkernel.rocm.triton_det_gemm" @dataclass(frozen=True) @@ -143,6 +145,7 @@ def __init__( native: ProjectionCallable | None = None, native_backend_id: str | None = None, deterministic: ProjectionCallable | None = None, + deterministic_backend_id: str | None = None, policy_id: str = PROJECTION_POLICY_ID, ) -> None: if projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: @@ -151,6 +154,11 @@ def __init__( self.native = native self.native_backend_id = native_backend_id or f"native.{projection}" self.deterministic = deterministic or DetGemmOp() + self.deterministic_backend_id = deterministic_backend_id or ( + ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID + if torch.version.hip is not None + else CUDA_DETERMINISTIC_PROJECTION_BACKEND_ID + ) self.policy_id = policy_id self.collective = ( QKV_COLLECTIVE_CONTRACT if projection == QKV_PROJECTION else O_PROJ_COLLECTIVE_CONTRACT @@ -191,7 +199,7 @@ def __call__(self, x: Tensor, weight: Tensor) -> ProjectionResult: deterministic_out, ProjectionPlan( projection=self.projection, - backend_id="rlkernel.cuda.det_gemm", + backend_id=self.deterministic_backend_id, fallback=True, fallback_reason=reason, probe_id=probe_id, @@ -242,6 +250,7 @@ def _probe_id(x: Tensor, weight: Tensor) -> str: __all__ = [ "AttentionProjectionOp", + "CUDA_DETERMINISTIC_PROJECTION_BACKEND_ID", "O_PROJ_COLLECTIVE_CONTRACT", "O_PROJ_PROJECTION", "PROJECTION_POLICY_ID", @@ -250,5 +259,6 @@ def _probe_id(x: Tensor, weight: Tensor) -> str: "ProjectionResult", "QKV_COLLECTIVE_CONTRACT", "QKV_PROJECTION", + "ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID", "split_qkv", ] diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 429f654f..1f3baa92 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -957,17 +957,23 @@ def test_strict_runtime_readback_accepts_common_deterministic_preprocess_fallbac assert result.passed -def test_strict_runtime_readback_accepts_distinct_verified_native_preprocess_backends(): +def test_strict_runtime_readback_rejects_distinct_native_preprocess_backends(): rollout = replace( _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), - preprocess_backends={"qk_rmsnorm": "native.vllm.rmsnorm", "rope": "native.vllm.rope"}, + preprocess_backends={ + "qk_rmsnorm": "transformer_engine.rocm.rmsnorm", + "rope": "rlkernel.rocm.deterministic_rope", + }, preprocess_probe_id="rollout-probe", ) training = replace( _readback( MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" ), - preprocess_backends={"qk_rmsnorm": "native.te.rmsnorm", "rope": "native.te.rope"}, + preprocess_backends={ + "qk_rmsnorm": "transformer_engine.cuda.rmsnorm", + "rope": "rlkernel.cuda.rope_sm90", + }, preprocess_probe_id="training-probe", ) @@ -980,7 +986,8 @@ def test_strict_runtime_readback_accepts_distinct_verified_native_preprocess_bac training_backend_id="megatron.te", ) - assert result.passed + assert not result.passed + assert result.issues_by_code(BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH) @pytest.mark.parametrize( diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py index 587c0100..c1264b5a 100644 --- a/tests/test_attention_preprocess.py +++ b/tests/test_attention_preprocess.py @@ -9,6 +9,7 @@ from rl_engine.kernels.attention_preprocess import ( MANDATED_ATTENTION_PREPROCESS_BACKENDS, H100AttentionPreprocessor, + TE_ROCM_QK_RMSNORM_BACKEND_ID, ) from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp @@ -44,6 +45,79 @@ def test_h100_preprocessor_uses_common_backend_ids_for_fallback(): } +def test_preprocessor_can_reuse_qk_norm_without_fusing_rope(): + def rmsnorm(x, weight, *, eps): + del eps + return (x.float() * weight.float()).to(torch.bfloat16) + + def rope(x, positions, *, theta): + del positions, theta + return x + + op = object.__new__(H100AttentionPreprocessor) + op.device = torch.device("cpu") + op.device_capability = (0, 0) + op.rmsnorm = rmsnorm + op.rope = rope + op.deterministic_backend_ids = { + "qk_rmsnorm": "rlkernel.rocm.triton_rmsnorm", + "rope": "rlkernel.rocm.deterministic_rope", + } + op.native_qk_norm = rmsnorm + op.native_rope = None + op.native_qk_norm_backend_id = TE_ROCM_QK_RMSNORM_BACKEND_ID + op.native_rope_backend_id = "unused" + op.policy_id = "test" + q = torch.randn(2, 4, 3, 8, dtype=torch.bfloat16) + k = torch.randn(2, 2, 3, 8, dtype=torch.bfloat16) + weight = torch.randn(8, dtype=torch.bfloat16) + positions = torch.arange(3, dtype=torch.int64) + + result = op(q, k, weight, weight, positions) + + assert result.fallback is False + assert result.backend_ids == { + "qk_rmsnorm": TE_ROCM_QK_RMSNORM_BACKEND_ID, + "rope": "rlkernel.rocm.deterministic_rope", + } + + +def test_preprocessor_falls_back_atomically_when_vendor_qk_norm_drifts(): + def deterministic(x, weight, *, eps): + del weight, eps + return x + + def drifting(x, weight, *, eps): + del weight, eps + return (x.float() + 1).to(torch.bfloat16) + + def rope(x, positions, *, theta): + del positions, theta + return x + + op = object.__new__(H100AttentionPreprocessor) + op.device = torch.device("cpu") + op.device_capability = (0, 0) + op.rmsnorm = deterministic + op.rope = rope + op.deterministic_backend_ids = dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) + op.native_qk_norm = drifting + op.native_rope = None + op.native_qk_norm_backend_id = "transformer_engine.cuda.rmsnorm" + op.native_rope_backend_id = "unused" + op.policy_id = "test" + q = torch.randn(1, 2, 3, 8, dtype=torch.bfloat16) + k = torch.randn(1, 1, 3, 8, dtype=torch.bfloat16) + weight = torch.randn(8, dtype=torch.bfloat16) + positions = torch.arange(3, dtype=torch.int64) + + result = op(q, k, weight, weight, positions) + + assert result.fallback is True + assert result.backend_ids == MANDATED_ATTENTION_PREPROCESS_BACKENDS + assert result.fallback_reason == "native_preprocess_bitwise_probe_failed" + + def test_h100_preprocessor_fails_before_dispatch_without_cuda(monkeypatch): monkeypatch.setattr(torch.cuda, "is_available", lambda: False) with pytest.raises(RuntimeError, match="requires an available CUDA runtime"): @@ -76,7 +150,9 @@ def _inputs(): @requires_h100_preprocess def test_h100_preprocessor_executes_cuda_qk_norm_and_zigzag_rope(): q, k, q_weight, k_weight, positions = _inputs() - result = H100AttentionPreprocessor()(q, k, q_weight, k_weight, positions) + result = H100AttentionPreprocessor(reuse_transformer_engine_qk_norm=False)( + q, k, q_weight, k_weight, positions + ) norm = NativeRMSNormOp() rope = NativeRoPEOp() @@ -101,7 +177,7 @@ def test_h100_preprocessor_executes_cuda_qk_norm_and_zigzag_rope(): @requires_h100_preprocess def test_h100_preprocessor_is_bitwise_batch_invariant_for_2d_positions(): q, k, q_weight, k_weight, positions = _inputs() - op = H100AttentionPreprocessor() + op = H100AttentionPreprocessor(reuse_transformer_engine_qk_norm=False) full = op(q, k, q_weight, k_weight, positions) for batch_index in range(q.shape[0]): diff --git a/tests/test_attention_projection.py b/tests/test_attention_projection.py index b4094351..40717181 100644 --- a/tests/test_attention_projection.py +++ b/tests/test_attention_projection.py @@ -9,6 +9,7 @@ O_PROJ_COLLECTIVE_CONTRACT, QKV_COLLECTIVE_CONTRACT, AttentionProjectionOp, + ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, split_qkv, ) @@ -57,6 +58,18 @@ def test_projection_accepts_native_only_after_bitwise_probe(): assert result.plan.fallback_reason is None +def test_projection_records_rocm_deterministic_backend(): + x, weight = _inputs() + result = AttentionProjectionOp( + "qkv", + deterministic=_deterministic, + deterministic_backend_id=ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, + )(x, weight) + + assert result.plan.backend_id == ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID + assert result.plan.fallback is True + + def test_projection_rejects_native_drift_and_records_reason(): x, weight = _inputs() From 2666f1843566a9c4be5dd6c6726f4ac37fd09840 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 20 Aug 2026 22:12:49 +0800 Subject: [PATCH 19/33] fix(attention): fail closed for required TE QK norm --- rl_engine/kernels/attention_preprocess.py | 22 +++++++++++++++++----- tests/test_attention_preprocess.py | 2 ++ 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index 6d1dda0d..69352d1a 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -11,9 +11,10 @@ from __future__ import annotations -import hashlib -import json -from dataclasses import dataclass +import hashlib +import json +from dataclasses import dataclass +from functools import lru_cache from types import MappingProxyType from typing import Any, Callable, Mapping @@ -61,6 +62,7 @@ class TransformerEngineRMSNormUnavailable(RuntimeError): """Raised when the exact TE RMSNorm functional contract is unavailable.""" +@lru_cache(maxsize=1) def _load_transformer_engine_rmsnorm(): try: import transformer_engine @@ -126,7 +128,9 @@ def __init__(self) -> None: def __call__(self, x: Tensor, weight: Tensor, *, eps: float = 1.0e-6) -> Tensor: if x.dtype not in (torch.float16, torch.bfloat16) or weight.dtype != x.dtype: - raise TypeError("Transformer Engine RMSNorm requires matching FP16/BF16 input and weight") + raise TypeError( + "Transformer Engine RMSNorm requires matching FP16/BF16 input and weight" + ) if not x.is_cuda or not weight.is_cuda or x.device != weight.device: raise ValueError("Transformer Engine RMSNorm requires input and weight on one GPU") if weight.shape != (x.shape[-1],): @@ -187,7 +191,7 @@ def __init__( native_qk_norm_backend_id: str = NATIVE_QK_RMSNORM_BACKEND_ID, native_rope_backend_id: str = NATIVE_ROPE_BACKEND_ID, reuse_transformer_engine_qk_norm: bool = True, - require_transformer_engine_qk_norm: bool = False, + require_transformer_engine_qk_norm: bool = True, policy_id: str = PREPROCESS_POLICY_ID, ) -> None: if not torch.cuda.is_available(): @@ -230,6 +234,9 @@ def __init__( raise self.native_qk_norm = native_qk_norm self.native_rope = native_rope + self.require_native_qk_norm = bool( + require_transformer_engine_qk_norm and reuse_transformer_engine_qk_norm + ) self.native_qk_norm_backend_id = native_qk_norm_backend_id self.native_rope_backend_id = native_rope_backend_id self.policy_id = policy_id @@ -324,6 +331,8 @@ def forward( fallback_reason = f"native_preprocess_unavailable:{type(exc).__name__}" else: fallback_reason = "native_preprocess_not_supplied" + if self.require_native_qk_norm: + raise TransformerEngineRMSNormUnavailable(fallback_reason) return AttentionPreprocessResult( q=q_det, k=k_det, @@ -376,6 +385,9 @@ def __init__( raise self.native_qk_norm = native_qk_norm self.native_rope = None + self.require_native_qk_norm = bool( + require_transformer_engine_qk_norm and reuse_transformer_engine_qk_norm + ) self.native_qk_norm_backend_id = ( getattr(native_qk_norm, "backend_id", TE_ROCM_QK_RMSNORM_BACKEND_ID) if native_qk_norm is not None diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py index c1264b5a..610a86d3 100644 --- a/tests/test_attention_preprocess.py +++ b/tests/test_attention_preprocess.py @@ -65,6 +65,7 @@ def rope(x, positions, *, theta): } op.native_qk_norm = rmsnorm op.native_rope = None + op.require_native_qk_norm = True op.native_qk_norm_backend_id = TE_ROCM_QK_RMSNORM_BACKEND_ID op.native_rope_backend_id = "unused" op.policy_id = "test" @@ -103,6 +104,7 @@ def rope(x, positions, *, theta): op.deterministic_backend_ids = dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) op.native_qk_norm = drifting op.native_rope = None + op.require_native_qk_norm = False op.native_qk_norm_backend_id = "transformer_engine.cuda.rmsnorm" op.native_rope_backend_id = "unused" op.policy_id = "test" From bff396eaa7a8e1a438676d2609e2029c04d58f34 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 20 Aug 2026 23:11:29 +0800 Subject: [PATCH 20/33] style(attention): satisfy PR230 pre-commit hooks --- .../alignment/cross_config/debug_matrix.py | 119 +++-- rl_engine/kernels/attention_preprocess.py | 414 +++++++++--------- .../ops/pytorch/attention/debug_matrix.py | 16 +- .../ops/pytorch/attention/debug_taxonomy.py | 8 +- tests/test_attention_debug_matrix.py | 8 +- tests/test_attention_debug_taxonomy.py | 4 +- tests/test_attention_preprocess.py | 2 +- tests/test_attention_projection.py | 2 +- 8 files changed, 270 insertions(+), 303 deletions(-) diff --git a/rl_engine/alignment/cross_config/debug_matrix.py b/rl_engine/alignment/cross_config/debug_matrix.py index 8f7b7041..32a49470 100644 --- a/rl_engine/alignment/cross_config/debug_matrix.py +++ b/rl_engine/alignment/cross_config/debug_matrix.py @@ -20,7 +20,6 @@ ) from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ATTENTION_DEBUG_AXES - DEBUG_MATRIX_SCHEMA_VERSION = "rlkernel.debug_matrix.v1" @@ -65,70 +64,64 @@ def _attention_axes() -> tuple[ModuleDebugAxis, ...]: ) -MODULE_DEBUG_AXES: MappingProxyType[str, tuple[ModuleDebugAxis, ...]] = ( - MappingProxyType( - { - "attention": _attention_axes(), - "ffn": ( - ModuleDebugAxis( - "ffn", - "weight_shard_ownership", - "Weight shard / TP ownership", - "tp_weight_ownership", - "gate", - ), - ModuleDebugAxis( - "ffn", - "swiglu_rounding", - "SwiGLU intermediate rounding", - "swiglu_one_round", - ), - ModuleDebugAxis( - "ffn", - "gemm_reduction", - "GEMM K-reduction / Split-K policy", - "k_reduction_split_k", - ), - ModuleDebugAxis( - "ffn", - "token_collective", - "Token gather / reduce-scatter", - "sequence_parallel", - ), +MODULE_DEBUG_AXES: MappingProxyType[str, tuple[ModuleDebugAxis, ...]] = MappingProxyType( + { + "attention": _attention_axes(), + "ffn": ( + ModuleDebugAxis( + "ffn", + "weight_shard_ownership", + "Weight shard / TP ownership", + "tp_weight_ownership", + "gate", ), - "logp": ( - ModuleDebugAxis( - "logp", - "vocab_shard_ownership", - "Vocabulary shard / TP ownership", - "vocab_shard_bounds", - "gate", - ), - ModuleDebugAxis( - "logp", - "selected_token_identity", - "Selected-token and active-mask identity", - "selected_token_active_mask", - "gate", - ), - ModuleDebugAxis( - "logp", - "vocab_lse_reduction", - "Vocabulary LSE tile / merge policy", - "vocab_tile_merge", - ), + ModuleDebugAxis( + "ffn", + "swiglu_rounding", + "SwiGLU intermediate rounding", + "swiglu_one_round", ), - } - ) + ModuleDebugAxis( + "ffn", + "gemm_reduction", + "GEMM K-reduction / Split-K policy", + "k_reduction_split_k", + ), + ModuleDebugAxis( + "ffn", + "token_collective", + "Token gather / reduce-scatter", + "sequence_parallel", + ), + ), + "logp": ( + ModuleDebugAxis( + "logp", + "vocab_shard_ownership", + "Vocabulary shard / TP ownership", + "vocab_shard_bounds", + "gate", + ), + ModuleDebugAxis( + "logp", + "selected_token_identity", + "Selected-token and active-mask identity", + "selected_token_active_mask", + "gate", + ), + ModuleDebugAxis( + "logp", + "vocab_lse_reduction", + "Vocabulary LSE tile / merge policy", + "vocab_tile_merge", + ), + ), + } ) _MODULE_DEBUG_AXIS_BY_ID = MappingProxyType( - { - (axis.module, axis.axis_id): axis - for axes in MODULE_DEBUG_AXES.values() - for axis in axes - } + {(axis.module, axis.axis_id): axis for axes in MODULE_DEBUG_AXES.values() for axis in axes} ) @@ -140,16 +133,12 @@ def _module_rows(module: str) -> dict[str, Any]: "ffn": ["FC0"], "logp": ["LC0"], }[module] - rows = [baseline] + [ - f"{module[0].upper()}{index}" for index in range(1, len(axes) + 1) - ] + rows = [baseline] + [f"{module[0].upper()}{index}" for index in range(1, len(axes) + 1)] return { "baseline_row": baseline, "rows": rows, "invariant_controls": controls, - "axes": [ - axis.to_dict(row_id=rows[index]) for index, axis in enumerate(axes, start=1) - ], + "axes": [axis.to_dict(row_id=rows[index]) for index, axis in enumerate(axes, start=1)], } diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index 69352d1a..6e3f463f 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -1,6 +1,6 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + """Bitwise-bound QK-Norm and RoPE handoff for WS2 Attention. Transformer Engine RMSNorm is the first-choice QK-Norm implementation on CUDA @@ -8,14 +8,14 @@ platform RL-Kernel path. RoPE remains on the shared RL-Kernel implementation so training and paged-cache rollout expose the same post-RoPE boundary. """ - -from __future__ import annotations - + +from __future__ import annotations + import hashlib import json from dataclasses import dataclass from functools import lru_cache -from types import MappingProxyType +from types import MappingProxyType from typing import Any, Callable, Mapping import torch @@ -32,10 +32,10 @@ NATIVE_ROPE_BACKEND_ID = "native.rope" PREPROCESS_POLICY_ID = "ws2.attention.preprocess.v3" MANDATED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( - { - "qk_rmsnorm": QK_RMSNORM_BACKEND_ID, - "rope": ROPE_BACKEND_ID, - } + { + "qk_rmsnorm": QK_RMSNORM_BACKEND_ID, + "rope": ROPE_BACKEND_ID, + } ) ROCM_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( { @@ -136,90 +136,90 @@ def __call__(self, x: Tensor, weight: Tensor, *, eps: float = 1.0e-6) -> Tensor: if weight.shape != (x.shape[-1],): raise ValueError("RMSNorm weight must match the input hidden dimension") return _TransformerEngineRMSNormFunction.apply(x, weight, float(eps)) - - -@dataclass(frozen=True) -class AttentionPreprocessResult: - """Post-QK-Norm, post-RoPE tensors plus executed backend evidence.""" - - q: Tensor - k: Tensor - backend_ids: Mapping[str, str] - fallback: bool - device_capability: tuple[int, int] - fallback_reason: str | None = None - probe_id: str = "" - policy_id: str = PREPROCESS_POLICY_ID - - def __post_init__(self) -> None: - object.__setattr__(self, "backend_ids", MappingProxyType(dict(self.backend_ids))) - - def evidence(self) -> dict[str, Any]: - return { - "backends": dict(self.backend_ids), - "fallback": self.fallback, - "device_capability": list(self.device_capability), - } - - def readback_fields(self) -> dict[str, Any]: - """Keyword fields consumed by ``AttentionRuntimeReadback``.""" - - return { - "preprocess_backends": dict(self.backend_ids), - "preprocess_fallback": self.fallback, - "preprocess_fallback_reason": self.fallback_reason, - "preprocess_probe_id": self.probe_id, - "preprocess_policy_id": self.policy_id, - } - - + + +@dataclass(frozen=True) +class AttentionPreprocessResult: + """Post-QK-Norm, post-RoPE tensors plus executed backend evidence.""" + + q: Tensor + k: Tensor + backend_ids: Mapping[str, str] + fallback: bool + device_capability: tuple[int, int] + fallback_reason: str | None = None + probe_id: str = "" + policy_id: str = PREPROCESS_POLICY_ID + + def __post_init__(self) -> None: + object.__setattr__(self, "backend_ids", MappingProxyType(dict(self.backend_ids))) + + def evidence(self) -> dict[str, Any]: + return { + "backends": dict(self.backend_ids), + "fallback": self.fallback, + "device_capability": list(self.device_capability), + } + + def readback_fields(self) -> dict[str, Any]: + """Keyword fields consumed by ``AttentionRuntimeReadback``.""" + + return { + "preprocess_backends": dict(self.backend_ids), + "preprocess_fallback": self.fallback, + "preprocess_fallback_reason": self.fallback_reason, + "preprocess_probe_id": self.probe_id, + "preprocess_policy_id": self.policy_id, + } + + class H100AttentionPreprocessor: """Reuse TE QK-Norm with the RL-Kernel SM90 RoPE boundary. - - ``native_qk_norm`` and ``native_rope`` are framework-owned callables. They - are intentionally injected instead of importing TE/vLLM here, so the same - policy can be used by both runtimes. The deterministic callables default to - RL-Kernel's CUDA operators and are always run to establish the probe oracle. - """ - - def __init__( - self, - device: torch.device | str | int | None = None, - *, - native_qk_norm: Callable[..., Tensor] | None = None, - native_rope: Callable[..., Tensor] | None = None, + + ``native_qk_norm`` and ``native_rope`` are framework-owned callables. They + are intentionally injected instead of importing TE/vLLM here, so the same + policy can be used by both runtimes. The deterministic callables default to + RL-Kernel's CUDA operators and are always run to establish the probe oracle. + """ + + def __init__( + self, + device: torch.device | str | int | None = None, + *, + native_qk_norm: Callable[..., Tensor] | None = None, + native_rope: Callable[..., Tensor] | None = None, native_qk_norm_backend_id: str = NATIVE_QK_RMSNORM_BACKEND_ID, native_rope_backend_id: str = NATIVE_ROPE_BACKEND_ID, reuse_transformer_engine_qk_norm: bool = True, require_transformer_engine_qk_norm: bool = True, policy_id: str = PREPROCESS_POLICY_ID, - ) -> None: - if not torch.cuda.is_available(): - raise RuntimeError("H100AttentionPreprocessor requires an available CUDA runtime") - - current_device = torch.cuda.current_device() - self.device = torch.device("cuda", current_device) - if device is not None: - self.device = ( - torch.device("cuda", device) if isinstance(device, int) else torch.device(device) - ) - if self.device.type != "cuda": - raise RuntimeError(f"H100AttentionPreprocessor requires CUDA, got {self.device}") - if self.device.index is None: - self.device = torch.device("cuda", current_device) - - capability = torch.cuda.get_device_capability(self.device) - self.device_capability: tuple[int, int] = (int(capability[0]), int(capability[1])) - if self.device_capability[0] != 9: - raise RuntimeError( - "H100AttentionPreprocessor requires Hopper SM90; " - f"got sm_{self.device_capability[0]}{self.device_capability[1]}" - ) - - # Import only after the hardware gate so CPU tools can inspect the module. - from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp - from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op - + ) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("H100AttentionPreprocessor requires an available CUDA runtime") + + current_device = torch.cuda.current_device() + self.device = torch.device("cuda", current_device) + if device is not None: + self.device = ( + torch.device("cuda", device) if isinstance(device, int) else torch.device(device) + ) + if self.device.type != "cuda": + raise RuntimeError(f"H100AttentionPreprocessor requires CUDA, got {self.device}") + if self.device.index is None: + self.device = torch.device("cuda", current_device) + + capability = torch.cuda.get_device_capability(self.device) + self.device_capability: tuple[int, int] = (int(capability[0]), int(capability[1])) + if self.device_capability[0] != 9: + raise RuntimeError( + "H100AttentionPreprocessor requires Hopper SM90; " + f"got sm_{self.device_capability[0]}{self.device_capability[1]}" + ) + + # Import only after the hardware gate so CPU tools can inspect the module. + from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op + self.rmsnorm = RMSNormCudaOp() self.rope = RoPESM90Op() self.deterministic_backend_ids = MANDATED_ATTENTION_PREPROCESS_BACKENDS @@ -237,45 +237,45 @@ def __init__( self.require_native_qk_norm = bool( require_transformer_engine_qk_norm and reuse_transformer_engine_qk_norm ) - self.native_qk_norm_backend_id = native_qk_norm_backend_id - self.native_rope_backend_id = native_rope_backend_id - self.policy_id = policy_id - - def __call__( - self, - q: Tensor, - k: Tensor, - q_weight: Tensor, - k_weight: Tensor, - positions: Tensor, - *, - eps: float = 1.0e-6, - theta: float = 1_000_000.0, - ) -> AttentionPreprocessResult: - return self.forward( - q, - k, - q_weight, - k_weight, - positions, - eps=eps, - theta=theta, - ) - - def forward( - self, - q: Tensor, - k: Tensor, - q_weight: Tensor, - k_weight: Tensor, - positions: Tensor, - *, - eps: float = 1.0e-6, - theta: float = 1_000_000.0, - ) -> AttentionPreprocessResult: - _validate_inputs(q, k, q_weight, k_weight, positions, self.device) - q_norm_det = self.rmsnorm(q, q_weight, eps=eps) - k_norm_det = self.rmsnorm(k, k_weight, eps=eps) + self.native_qk_norm_backend_id = native_qk_norm_backend_id + self.native_rope_backend_id = native_rope_backend_id + self.policy_id = policy_id + + def __call__( + self, + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + *, + eps: float = 1.0e-6, + theta: float = 1_000_000.0, + ) -> AttentionPreprocessResult: + return self.forward( + q, + k, + q_weight, + k_weight, + positions, + eps=eps, + theta=theta, + ) + + def forward( + self, + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + *, + eps: float = 1.0e-6, + theta: float = 1_000_000.0, + ) -> AttentionPreprocessResult: + _validate_inputs(q, k, q_weight, k_weight, positions, self.device) + q_norm_det = self.rmsnorm(q, q_weight, eps=eps) + k_norm_det = self.rmsnorm(k, k_weight, eps=eps) q_det = _apply_deterministic_rope(self.rope, q_norm_det, positions, theta) k_det = _apply_deterministic_rope(self.rope, k_norm_det, positions, theta) probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) @@ -291,21 +291,17 @@ def forward( ) if norm_matches and native_rope is None: return AttentionPreprocessResult( - q=_apply_deterministic_rope( - self.rope, q_norm_native, positions, theta - ), - k=_apply_deterministic_rope( - self.rope, k_norm_native, positions, theta - ), + q=_apply_deterministic_rope(self.rope, q_norm_native, positions, theta), + k=_apply_deterministic_rope(self.rope, k_norm_native, positions, theta), backend_ids=MappingProxyType( { "qk_rmsnorm": self.native_qk_norm_backend_id, "rope": self.deterministic_backend_ids["rope"], } ), - fallback=False, - device_capability=self.device_capability, - probe_id=probe_id, + fallback=False, + device_capability=self.device_capability, + probe_id=probe_id, policy_id=self.policy_id, ) if norm_matches and native_rope is not None: @@ -337,10 +333,10 @@ def forward( q=q_det, k=k_det, backend_ids=self.deterministic_backend_ids, - fallback=True, - device_capability=self.device_capability, - fallback_reason=fallback_reason, - probe_id=probe_id, + fallback=True, + device_capability=self.device_capability, + fallback_reason=fallback_reason, + probe_id=probe_id, policy_id=self.policy_id, ) @@ -395,77 +391,77 @@ def __init__( ) self.native_rope_backend_id = ROCM_ROPE_BACKEND_ID self.policy_id = policy_id - - -def _apply_deterministic_rope( - rope: Callable[..., Tensor], - x: Tensor, - positions: Tensor, - theta: float, -) -> Tensor: - """Adapt per-sample positions to the CUDA RoPE operator's 1-D contract.""" - - if positions.dim() == 1: - return rope(x, positions, theta=theta) - return torch.cat( - [rope(x[index : index + 1], positions[index], theta=theta) for index in range(x.shape[0])], - dim=0, - ) - - -def _probe_id( - q: Tensor, - k: Tensor, - q_weight: Tensor, - k_weight: Tensor, - positions: Tensor, - eps: float, - theta: float, -) -> str: - payload = { - "q_shape": list(q.shape), - "k_shape": list(k.shape), - "q_dtype": str(q.dtype), - "k_dtype": str(k.dtype), - "weight_dtype": str(q_weight.dtype), - "positions_shape": list(positions.shape), - "positions_sha256": hashlib.sha256(positions.detach().cpu().numpy().tobytes()).hexdigest(), - "eps": float(eps), - "theta": float(theta), - } - return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16] - - -def _validate_inputs( - q: Tensor, - k: Tensor, - q_weight: Tensor, - k_weight: Tensor, - positions: Tensor, - device: torch.device, -) -> None: - if q.dim() != 4 or k.dim() != 4: - raise ValueError("q and k must use [B, H, S, D] layout") - if q.shape[0] != k.shape[0] or q.shape[-2:] != k.shape[-2:]: - raise ValueError("q and k must have the same batch, sequence, and head dimensions") + + +def _apply_deterministic_rope( + rope: Callable[..., Tensor], + x: Tensor, + positions: Tensor, + theta: float, +) -> Tensor: + """Adapt per-sample positions to the CUDA RoPE operator's 1-D contract.""" + + if positions.dim() == 1: + return rope(x, positions, theta=theta) + return torch.cat( + [rope(x[index : index + 1], positions[index], theta=theta) for index in range(x.shape[0])], + dim=0, + ) + + +def _probe_id( + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + eps: float, + theta: float, +) -> str: + payload = { + "q_shape": list(q.shape), + "k_shape": list(k.shape), + "q_dtype": str(q.dtype), + "k_dtype": str(k.dtype), + "weight_dtype": str(q_weight.dtype), + "positions_shape": list(positions.shape), + "positions_sha256": hashlib.sha256(positions.detach().cpu().numpy().tobytes()).hexdigest(), + "eps": float(eps), + "theta": float(theta), + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16] + + +def _validate_inputs( + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + device: torch.device, +) -> None: + if q.dim() != 4 or k.dim() != 4: + raise ValueError("q and k must use [B, H, S, D] layout") + if q.shape[0] != k.shape[0] or q.shape[-2:] != k.shape[-2:]: + raise ValueError("q and k must have the same batch, sequence, and head dimensions") if q.dtype is not torch.bfloat16 or k.dtype is not torch.bfloat16: raise TypeError("the frozen Attention experiment requires BF16 q and k") - if q.device != device or k.device != device: - raise ValueError(f"q and k must both be on the configured device {device}") - for name, weight in (("q_weight", q_weight), ("k_weight", k_weight)): - if weight.shape != (q.shape[-1],): - raise ValueError(f"{name} must have shape ({q.shape[-1]},)") - if weight.device != device or weight.dtype is not torch.bfloat16: - raise ValueError(f"{name} must be BF16 on {device}") - if positions.device != device: - raise ValueError(f"positions must be on {device}") - if positions.dtype not in (torch.int32, torch.int64): - raise TypeError("positions must use int32 or int64 global token indices") - expected = (q.shape[-2],) if positions.dim() == 1 else (q.shape[0], q.shape[-2]) - if positions.dim() not in (1, 2) or tuple(positions.shape) != expected: - raise ValueError(f"positions must have shape [S] or [B, S], expected {expected}") - - + if q.device != device or k.device != device: + raise ValueError(f"q and k must both be on the configured device {device}") + for name, weight in (("q_weight", q_weight), ("k_weight", k_weight)): + if weight.shape != (q.shape[-1],): + raise ValueError(f"{name} must have shape ({q.shape[-1]},)") + if weight.device != device or weight.dtype is not torch.bfloat16: + raise ValueError(f"{name} must be BF16 on {device}") + if positions.device != device: + raise ValueError(f"positions must be on {device}") + if positions.dtype not in (torch.int32, torch.int64): + raise TypeError("positions must use int32 or int64 global token indices") + expected = (q.shape[-2],) if positions.dim() == 1 else (q.shape[0], q.shape[-2]) + if positions.dim() not in (1, 2) or tuple(positions.shape) != expected: + raise ValueError(f"positions must have shape [S] or [B, S], expected {expected}") + + __all__ = [ "ALLOWED_ATTENTION_PREPROCESS_BACKENDS", "AttentionPreprocessResult", diff --git a/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py b/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py index 306e1db9..cbceab9c 100644 --- a/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py +++ b/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py @@ -43,9 +43,7 @@ def __post_init__(self) -> None: "comparability_gate", "invariant_control", }: - raise ValueError( - f"unknown Attention debug matrix category {self.category!r}" - ) + raise ValueError(f"unknown Attention debug matrix category {self.category!r}") if self.category == "baseline": if self.probe is not None or self.root_cause_axis is not None: raise ValueError("baseline rows cannot name a probe or root-cause axis") @@ -121,9 +119,7 @@ def validate_attention_debug_matrix() -> None: raise RuntimeError("Attention debug matrix row IDs must be unique") baselines = [row for row in ATTENTION_DEBUG_MATRIX if row.category == "baseline"] if len(baselines) != 1 or baselines[0].row_id != "A0": - raise RuntimeError( - "Attention debug matrix must contain exactly one A0 baseline" - ) + raise RuntimeError("Attention debug matrix must contain exactly one A0 baseline") axis_ids = {axis.axis_id for axis in ATTENTION_DEBUG_AXES} diagnostic_rows = [ @@ -139,9 +135,7 @@ def validate_attention_debug_matrix() -> None: if [row.root_cause_axis for row in gates] != ["topology_head_ownership"]: raise RuntimeError("only topology/head ownership is a comparability gate") - controls = [ - row for row in ATTENTION_DEBUG_MATRIX if row.category == "invariant_control" - ] + controls = [row for row in ATTENTION_DEBUG_MATRIX if row.category == "invariant_control"] if {row.probe for row in controls} != set(ATTENTION_INVARIANT_CONTROLS): raise RuntimeError("invariant-control coverage is incomplete") if any(row.expected != "exact_zero" for row in controls): @@ -200,9 +194,7 @@ def attention_debug_matrix_row(row_id: str) -> AttentionDebugMatrixRow: try: return _ROWS_BY_ID[row_id.strip()] except KeyError as exc: - raise AttentionContractError( - f"unknown Attention debug matrix row {row_id!r}" - ) from exc + raise AttentionContractError(f"unknown Attention debug matrix row {row_id!r}") from exc __all__ = [ diff --git a/rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py b/rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py index d9f1a6d6..02bc021c 100644 --- a/rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py +++ b/rl_engine/kernels/ops/pytorch/attention/debug_taxonomy.py @@ -33,9 +33,7 @@ def __post_init__(self) -> None: ): raise ValueError("Attention debug subprobes must be non-empty strings") if len(set(self.subprobes)) != len(self.subprobes): - raise ValueError( - f"Attention debug axis {self.axis_id!r} has duplicate subprobes" - ) + raise ValueError(f"Attention debug axis {self.axis_id!r} has duplicate subprobes") if self.representative_subprobe not in self.subprobes: raise ValueError( f"Attention debug representative {self.representative_subprobe!r} " @@ -118,9 +116,7 @@ def __post_init__(self) -> None: ) if len(_ATTENTION_DEBUG_AXIS_BY_ID) != len(ATTENTION_DEBUG_AXES): raise RuntimeError("Attention debug axis IDs must be unique") -if len(_ATTENTION_DEBUG_AXIS_BY_PROBE) != sum( - len(axis.subprobes) for axis in ATTENTION_DEBUG_AXES -): +if len(_ATTENTION_DEBUG_AXIS_BY_PROBE) != sum(len(axis.subprobes) for axis in ATTENTION_DEBUG_AXES): raise RuntimeError("Attention debug subprobes must belong to one root-cause axis") if set(_ATTENTION_DEBUG_AXIS_BY_PROBE) & set(ATTENTION_INVARIANT_CONTROLS): raise RuntimeError("Attention debug subprobes cannot also be invariant controls") diff --git a/tests/test_attention_debug_matrix.py b/tests/test_attention_debug_matrix.py index 86710882..eb23e513 100644 --- a/tests/test_attention_debug_matrix.py +++ b/tests/test_attention_debug_matrix.py @@ -58,16 +58,12 @@ def test_attention_debug_matrix_rows_have_stable_semantics(): assert control.probe == "tp_partition_control" assert control.expected == "exact_zero" - with pytest.raises( - AttentionContractError, match="unknown Attention debug matrix row" - ): + with pytest.raises(AttentionContractError, match="unknown Attention debug matrix row"): attention_debug_matrix_row("A8") def test_matrix_manifest_matches_taxonomy_representatives(): - from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ( - ATTENTION_DEBUG_AXES, - ) + from rl_engine.kernels.ops.pytorch.attention.debug_taxonomy import ATTENTION_DEBUG_AXES root_rows = [ row diff --git a/tests/test_attention_debug_taxonomy.py b/tests/test_attention_debug_taxonomy.py index 149a99fd..1d783132 100644 --- a/tests/test_attention_debug_taxonomy.py +++ b/tests/test_attention_debug_taxonomy.py @@ -29,9 +29,7 @@ def test_attention_debug_taxonomy_is_compact_complete_and_unambiguous(): "distributed_schedule", } assert "strict_cp_degree_control" not in { - probe - for axis in taxonomy["root_cause_axes"].values() - for probe in axis["subprobes"] + probe for axis in taxonomy["root_cause_axes"].values() for probe in axis["subprobes"] } diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py index 610a86d3..9b800dca 100644 --- a/tests/test_attention_preprocess.py +++ b/tests/test_attention_preprocess.py @@ -8,8 +8,8 @@ from rl_engine.kernels.attention_preprocess import ( MANDATED_ATTENTION_PREPROCESS_BACKENDS, - H100AttentionPreprocessor, TE_ROCM_QK_RMSNORM_BACKEND_ID, + H100AttentionPreprocessor, ) from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp diff --git a/tests/test_attention_projection.py b/tests/test_attention_projection.py index 40717181..05d02598 100644 --- a/tests/test_attention_projection.py +++ b/tests/test_attention_projection.py @@ -8,8 +8,8 @@ from rl_engine.kernels.attention_projection import ( O_PROJ_COLLECTIVE_CONTRACT, QKV_COLLECTIVE_CONTRACT, - AttentionProjectionOp, ROCM_DETERMINISTIC_PROJECTION_BACKEND_ID, + AttentionProjectionOp, split_qkv, ) From a14ef3607b74475205cb7f86fbeed819e4d28aad Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 20 Aug 2026 23:41:47 +0800 Subject: [PATCH 21/33] feat(attention): bind qualified vendor cores by platform --- .../cross_config/attention_binding.py | 64 ++++-- rl_engine/kernels/attention_preprocess.py | 10 +- .../kernels/ops/pytorch/attention/ablation.py | 55 ++--- tests/test_attention_ablation.py | 209 +++++++++++++++--- tests/test_attention_cross_config_binding.py | 86 ++++++- 5 files changed, 336 insertions(+), 88 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 4249c314..f6a40ff5 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -44,8 +44,6 @@ from typing import Any, Optional from rl_engine.kernels.attention_contract import ( - STRICT_ATTENTION_CORE_ID, - STRICT_ATTENTION_SCHEDULE_ID, AttentionContract, AttentionContractError, AttentionDType, @@ -159,6 +157,8 @@ class AttentionRuntimeReadback: actual_backend: str | None = None communication_backend: str | None = None production_ready: bool = False + attention_fallback: bool = False + reference_only: bool = False def __post_init__(self) -> None: if not isinstance(self.contract, AttentionContract): @@ -216,6 +216,10 @@ def __post_init__(self) -> None: raise ValueError(f"{name} must be a non-empty string when provided") if not isinstance(self.production_ready, bool): raise TypeError("production_ready must be a bool") + if not isinstance(self.attention_fallback, bool): + raise TypeError("attention_fallback must be a bool") + if not isinstance(self.reference_only, bool): + raise TypeError("reference_only must be a bool") normalized_projection_plans: dict[str, Mapping[str, Any]] = {} for name, plan in self.projection_plans.items(): if not isinstance(name, str) or not isinstance(plan, Mapping): @@ -268,6 +272,8 @@ def to_dict(self) -> dict[str, Any]: "actual_backend": self.actual_backend, "communication_backend": self.communication_backend, "production_ready": self.production_ready, + "fallback": self.attention_fallback, + "reference_only": self.reference_only, }, "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), } @@ -963,7 +969,7 @@ def bind_attention_runtime_readbacks( rollout=rollout.preprocess_fallback, training=training.preprocess_fallback, message=( - "both runtimes must either pass the native H100 bitwise probe or " + "both runtimes must either pass the platform vendor bitwise probe or " "use the same deterministic preprocess fallback" ), ) @@ -1050,6 +1056,8 @@ def bind_attention_runtime_readbacks( "runtime.actual_backend": rollout.actual_backend, "runtime.communication_backend": rollout.communication_backend, "runtime.production_ready": rollout.production_ready, + "runtime.fallback": rollout.attention_fallback, + "runtime.reference_only": rollout.reference_only, **{ f"projection.{projection}": dict(plan) for projection, plan in rollout.projection_plans.items() @@ -1072,6 +1080,8 @@ def bind_attention_runtime_readbacks( "runtime.actual_backend": training.actual_backend, "runtime.communication_backend": training.communication_backend, "runtime.production_ready": training.production_ready, + "runtime.fallback": training.attention_fallback, + "runtime.reference_only": training.reference_only, **{ f"projection.{projection}": dict(plan) for projection, plan in training.projection_plans.items() @@ -1087,7 +1097,7 @@ def _strict_attention_core_issues( if not readback.strict_mode: return [] issues = [] - if readback.actual_backend != "rlkernel.cuda.deterministic_attention": + if not readback.actual_backend: issues.append( BindingIssue( code=BindingErrorCode.ATTENTION_BACKEND_MISSING, @@ -1095,10 +1105,14 @@ def _strict_attention_core_issues( field=f"{side}.runtime.actual_backend", rollout=readback.actual_backend if side == "rollout" else None, training=readback.actual_backend if side == "training" else None, - message=(f"{side} strict Attention did not execute the CUDA deterministic core"), + message=(f"{side} strict Attention did not report its executed production core"), ) ) - if readback.communication_backend != "self_owned_cuda_ag_rs": + supported_communication = {"self_owned_cuda_ag_rs", "cuda_ag_rs", "rccl_ag_rs"} + if ( + readback.contract.sharding.cp_world_size > 1 + and readback.communication_backend not in supported_communication + ): issues.append( BindingIssue( code=BindingErrorCode.ATTENTION_BACKEND_MISSING, @@ -1107,7 +1121,7 @@ def _strict_attention_core_issues( rollout=readback.communication_backend if side == "rollout" else None, training=readback.communication_backend if side == "training" else None, message=( - f"{side} strict CP Attention did not execute the self-owned CUDA AG/RS path" + f"{side} strict CP Attention did not execute a supported self-owned AG/RS path" ), ) ) @@ -1122,37 +1136,31 @@ def _strict_attention_core_issues( message=(f"{side} evidence is reference-only and cannot close the production gate"), ) ) - if readback.strict_core_id != STRICT_ATTENTION_CORE_ID: + if not readback.strict_core_id: issues.append( BindingIssue( code=BindingErrorCode.ATTENTION_CORE_MISSING, tier=BindingTier.SEMANTIC, field=f"{side}.strict.core_id", - message=( - f"{side} strict Attention did not execute the shared core " - f"{STRICT_ATTENTION_CORE_ID!r}" - ), + message=(f"{side} strict Attention did not report its exact shared core identity"), ) ) - if readback.strict_schedule != STRICT_ATTENTION_SCHEDULE_ID: + if not readback.strict_schedule: issues.append( BindingIssue( code=BindingErrorCode.ATTENTION_CORE_SCHEDULE, tier=BindingTier.SEMANTIC, field=f"{side}.strict.schedule", - message=( - f"{side} strict Attention did not execute the canonical schedule " - f"{STRICT_ATTENTION_SCHEDULE_ID!r}" - ), + message=(f"{side} strict Attention did not report its exact reduction schedule"), ) ) - if readback.native_attention_arithmetic: + if readback.attention_fallback or readback.reference_only: issues.append( BindingIssue( - code=BindingErrorCode.ATTENTION_NATIVE_ARITHMETIC, + code=BindingErrorCode.ATTENTION_BACKEND_MISSING, tier=BindingTier.SEMANTIC, - field=f"{side}.strict.native_attention_arithmetic", - message=f"{side} strict Attention entered native TE/FlashInfer arithmetic", + field=f"{side}.runtime.production_path", + message=f"{side} strict Attention executed a fallback or reference-only path", ) ) if ( @@ -1210,6 +1218,20 @@ def _strict_attention_core_pair_issues( message="training and rollout executed different strict Attention schedules", ) ] + if ( + rollout.strict_mode + and rollout.native_attention_arithmetic != training.native_attention_arithmetic + ): + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_NATIVE_ARITHMETIC, + tier=BindingTier.SEMANTIC, + field="strict.native_attention_arithmetic", + rollout=rollout.native_attention_arithmetic, + training=training.native_attention_arithmetic, + message="training and rollout disagree on vendor Attention arithmetic", + ) + ] if rollout.strict_mode and rollout.actual_backend != training.actual_backend: return [ BindingIssue( diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index 6e3f463f..9833d022 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -140,7 +140,11 @@ def __call__(self, x: Tensor, weight: Tensor, *, eps: float = 1.0e-6) -> Tensor: @dataclass(frozen=True) class AttentionPreprocessResult: - """Post-QK-Norm, post-RoPE tensors plus executed backend evidence.""" + """Post-QK-Norm, post-RoPE tensors plus executed backend evidence. + + ``probe_id`` identifies the probe configuration, not tensor contents, and + must never be used as an admission-result cache key. + """ q: Tensor k: Tensor @@ -278,7 +282,7 @@ def forward( k_norm_det = self.rmsnorm(k, k_weight, eps=eps) q_det = _apply_deterministic_rope(self.rope, q_norm_det, positions, theta) k_det = _apply_deterministic_rope(self.rope, k_norm_det, positions, theta) - probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) + probe_id = _probe_configuration_id(q, k, q_weight, k_weight, positions, eps, theta) native_qk_norm = self.native_qk_norm native_rope = self.native_rope @@ -409,7 +413,7 @@ def _apply_deterministic_rope( ) -def _probe_id( +def _probe_configuration_id( q: Tensor, k: Tensor, q_weight: Tensor, diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index 3f8727f7..12deabe1 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -5,8 +5,8 @@ The matrix needs one stable callable shape even though training and rollout may materialize different Attention backends. This adapter owns the common -contract checks and provenance only; numerical work remains in the existing -deterministic Attention implementations. +contract checks and provenance only; numerical work remains in the qualified +CUDA/ROCm production cores or the explicit RL-Kernel reference core. """ from __future__ import annotations @@ -33,6 +33,7 @@ BACKEND_ID = "rlkernel.attention.deterministic.v1" REFERENCE_BACKEND_ID = "rlkernel.attention.reference.v1" +_STRICT_AG_RS_BACKENDS = frozenset({"self_owned_cuda_ag_rs", "cuda_ag_rs", "rccl_ag_rs"}) _TORCH_DTYPES = { AttentionDType.BF16: torch.bfloat16, @@ -157,7 +158,7 @@ def __init__( self.core = core self.reference = reference self.native = native - # CP production execution is injected by the runtime adapter. Keeping + # CP production execution is injected by the runtime adapter. Keeping # it separate from the single-device core prevents an accidental # fallback to the PyTorch reference when AG/RS is required. self.cp_backend = cp_backend @@ -241,13 +242,6 @@ def apply( raise AttentionContractError( "deterministic Attention cannot use runtime-dependent Split-KV=auto" ) - if cfg.deterministic and ( - cfg.strict_core_id != STRICT_ATTENTION_CORE_ID - or cfg.strict_schedule != STRICT_ATTENTION_SCHEDULE_ID - ): - raise AttentionContractError( - "strict deterministic Attention requires the canonical core and schedule" - ) if cfg.deterministic and not cfg.return_lse: raise AttentionContractError("strict deterministic Attention must return LSE") if cfg.return_gradients and dout is None: @@ -363,7 +357,7 @@ def _select_backend( if normalized not in {"auto", "deterministic", "rlkernel"}: raise AttentionContractError(f"unsupported Attention backend {requested!r}") if contract.sharding.cp_world_size > 1: - if communication_backend == "self_owned_cuda_ag_rs": + if communication_backend in _STRICT_AG_RS_BACKENDS: if self.cp_backend is None: raise AttentionContractError( "CP production Attention requires an injected AG/RS backend" @@ -464,7 +458,7 @@ def _invoke( "Attention backend must expose forward_with_lse, apply, or __call__" ) accepted = _accepted_kwargs(method, kwargs) - if contract.sharding.cp_world_size > 1 and "cp_world_size" not in accepted: + if contract.sharding.cp_world_size > 1 and not _declares_keyword(method, "cp_world_size"): raise AttentionContractError( "CP>1 requires an Attention backend that explicitly accepts cp_world_size" ) @@ -566,6 +560,19 @@ def _accepted_kwargs(method: Callable[..., Any], kwargs: Mapping[str, Any]) -> d return {name: value for name, value in kwargs.items() if name in signature.parameters} +def _declares_keyword(method: Callable[..., Any], name: str) -> bool: + """Return whether ``method`` explicitly declares a keyword-capable parameter.""" + + try: + parameter = inspect.signature(method).parameters.get(name) + except (TypeError, ValueError): + return False + return parameter is not None and parameter.kind in { + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + } + + def _callable_backend_id(value: Any) -> str: explicit = getattr(value, "backend_id", None) or getattr( value, "__attention_backend_id__", None @@ -590,31 +597,27 @@ def _validate_runtime_provenance( "deterministic Attention cannot execute native Attention arithmetic" ) - actual_core = runtime.get("strict_core_id", getattr(selected, "core_id", None)) - actual_schedule = runtime.get("strict_schedule", getattr(selected, "strict_schedule", None)) - if selected_id != REFERENCE_BACKEND_ID and ( - actual_core != config.strict_core_id or actual_schedule != config.strict_schedule - ): - raise AttentionContractError( - "deterministic Attention backend did not prove the shared strict core and schedule" - ) - - if config.communication_backend != "self_owned_cuda_ag_rs": + external_production_backend = selected_id not in {BACKEND_ID, REFERENCE_BACKEND_ID} + if not external_production_backend: return expected = { "strict_core_id": config.strict_core_id, "strict_schedule": config.strict_schedule, - "actual_backend": "rlkernel.cuda.deterministic_attention", - "communication_backend": "self_owned_cuda_ag_rs", "production_ready": True, - "native_attention_arithmetic": False, "fallback": False, + "reference_only": False, } + if config.communication_backend in _STRICT_AG_RS_BACKENDS: + expected["communication_backend"] = config.communication_backend mismatches = [name for name, value in expected.items() if runtime.get(name) != value] + if not isinstance(runtime.get("actual_backend"), str) or not runtime["actual_backend"].strip(): + mismatches.append("actual_backend") + if not isinstance(runtime.get("native_attention_arithmetic"), bool): + mismatches.append("native_attention_arithmetic") if mismatches: raise AttentionContractError( - "CP production Attention runtime provenance is incomplete or mismatched: " + "production Attention runtime provenance is incomplete or mismatched: " + ", ".join(mismatches) ) diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py index 70dbd2ca..ae36237c 100644 --- a/tests/test_attention_ablation.py +++ b/tests/test_attention_ablation.py @@ -22,6 +22,7 @@ from rl_engine.kernels.ops.pytorch.attention.ablation import ( BACKEND_ID, REFERENCE_BACKEND_ID, + AttentionAblationConfig, AttentionAblationOp, ) @@ -68,6 +69,38 @@ def _qkv(): ) +def _cp2_contract() -> AttentionContract: + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=2, + head_dim=4, + causal=True, + causal_offsets=(0,), + sharding=ShardingSpec( + tp_rank=0, + tp_world_size=1, + cp_rank=0, + cp_world_size=2, + global_q_heads=2, + global_kv_heads=1, + local_q_head_start=0, + local_q_heads=2, + local_kv_head_start=0, + local_kv_heads=1, + global_sequence_length=4, + local_sequence_length=2, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 2), + ), + reduction=ReductionSpec(), + split_kv=SplitKVSpec.disabled(), + ) + + def test_attention_wrapper_has_unified_result_and_provenance(): q, k, v = _qkv() result = AttentionAblationOp()(q, k, v, contract=_contract()) @@ -106,44 +139,106 @@ def forward_with_lse(self, q, k, v, *, causal, scale): assert torch.equal(result.out, q) -def test_cp_production_configuration_fails_closed_without_ag_rs_backend(): +def test_wrapper_owned_deterministic_core_does_not_require_external_provenance(): q, k, v = _qkv() - cp_sharding = ShardingSpec( - tp_rank=0, - tp_world_size=1, - cp_rank=0, - cp_world_size=2, - global_q_heads=2, - global_kv_heads=1, - local_q_head_start=0, - local_q_heads=2, - local_kv_head_start=0, - local_kv_heads=1, - global_sequence_length=4, - local_sequence_length=2, - global_block_indices=(0,), - global_block_token_starts=(0,), - local_block_offsets=(0, 2), - ) - contract = AttentionContract( - role=AttentionRole.TRAIN, - mode=AttentionMode.PREFILL, - dtype=AttentionDType.BF16, - batch_size=1, - query_sequence_length=2, - head_dim=4, - causal=True, - causal_offsets=(0,), - sharding=cp_sharding, - reduction=ReductionSpec(), - split_kv=SplitKVSpec.disabled(), + + class WrapperOwnedCore: + backend_id = BACKEND_ID + + def __call__(self, q, k, v, *, causal, scale): + del k, v, causal, scale + return q.clone(), torch.zeros(q.shape[:3], dtype=torch.float32) + + result = AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + backend=WrapperOwnedCore(), ) + + assert result.backend_id == BACKEND_ID + assert result.deterministic is True + + +def test_cp_production_configuration_fails_closed_without_ag_rs_backend(): + q, k, v = _qkv() with pytest.raises(AttentionContractError, match="injected AG/RS backend"): AttentionAblationOp(communication_backend="self_owned_cuda_ag_rs")( - q[:, :, :2], k[:, :, :2], v[:, :, :2], contract=contract + q[:, :, :2], k[:, :, :2], v[:, :, :2], contract=_cp2_contract() + ) + + +def test_cp_backend_must_explicitly_declare_cp_world_size(): + q, k, v = _qkv() + + class KwargsOnlyBackend: + backend_id = "test.kwargs_only" + core_id = STRICT_ATTENTION_CORE_ID + strict_schedule = STRICT_ATTENTION_SCHEDULE_ID + + def __call__(self, q, k, v, **kwargs): + del k, v, kwargs + return q, torch.zeros(q.shape[:3], dtype=torch.float32) + + with pytest.raises(AttentionContractError, match="explicitly accepts cp_world_size"): + AttentionAblationOp( + cp_backend=KwargsOnlyBackend(), + communication_backend="cuda_ag_rs", + )( + q[:, :, :2], + k[:, :, :2], + v[:, :, :2], + contract=_cp2_contract(), ) +@pytest.mark.parametrize("communication_backend", ["cuda_ag_rs", "rccl_ag_rs"]) +def test_cp_wrapper_accepts_exact_platform_vendor_core(communication_backend): + q, k, v = _qkv() + + class VendorCPBackend: + backend_id = "vendor.strict_attention" + core_id = "vendor.strict_core.v1" + strict_schedule = "vendor.strict_schedule.v1" + + def __call__(self, q, k, v, *, causal, scale, cp_world_size): + del k, v, causal, scale + assert cp_world_size == 2 + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:3], dtype=torch.float32), + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": self.backend_id, + "communication_backend": communication_backend, + "production_ready": True, + "native_attention_arithmetic": True, + "fallback": False, + "reference_only": False, + }, + ) + + result = AttentionAblationOp( + cp_backend=VendorCPBackend(), + communication_backend=communication_backend, + )( + q[:, :, :2], + k[:, :, :2], + v[:, :, :2], + contract=_cp2_contract(), + config=AttentionAblationConfig( + strict_core_id=VendorCPBackend.core_id, + strict_schedule=VendorCPBackend.strict_schedule, + ), + ) + + assert result.provenance["actual_backend"] == VendorCPBackend.backend_id + assert result.provenance["communication_backend"] == communication_backend + assert result.provenance["native_attention_arithmetic"] is True + + def test_cp_production_wrapper_preserves_runtime_backend_provenance(): q, k, v = _qkv() @@ -165,6 +260,7 @@ def __call__(self, q, k, v, *, causal, scale): "production_ready": True, "native_attention_arithmetic": False, "fallback": False, + "reference_only": False, }, ) @@ -183,6 +279,57 @@ def __call__(self, q, k, v, *, causal, scale): assert result.provenance["production_ready"] is True +@pytest.mark.parametrize( + ("missing_field", "replacement"), + [ + ("production_ready", None), + ("fallback", True), + ("reference_only", True), + ], +) +def test_vendor_production_core_fails_closed_without_exact_provenance(missing_field, replacement): + q, k, v = _qkv() + + class VendorCore: + backend_id = "vendor.strict_attention" + core_id = "vendor.strict_core.v1" + strict_schedule = "vendor.strict_schedule.v1" + + def __call__(self, q, k, v, *, causal, scale): + del k, v, causal, scale + provenance = { + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": self.backend_id, + "production_ready": True, + "native_attention_arithmetic": True, + "fallback": False, + "reference_only": False, + } + if replacement is None: + del provenance[missing_field] + else: + provenance[missing_field] = replacement + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:3], dtype=torch.float32), + provenance=provenance, + ) + + with pytest.raises(AttentionContractError, match="runtime provenance"): + AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + backend=VendorCore(), + config=AttentionAblationConfig( + strict_core_id=VendorCore.core_id, + strict_schedule=VendorCore.strict_schedule, + ), + ) + + def test_deterministic_attention_rejects_runtime_split_kv_auto(): q, k, v = _qkv() contract = _contract(split_kv=SplitKVSpec.auto(strict_consistency=False)) diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 1f3baa92..3e12ad34 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -819,13 +819,9 @@ def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): @pytest.mark.parametrize( ("changes", "expected_code"), [ - ( - {"native_attention_arithmetic": True}, - BindingErrorCode.ATTENTION_NATIVE_ARITHMETIC, - ), ( {"strict_core_id": "different.core"}, - BindingErrorCode.ATTENTION_CORE_MISSING, + BindingErrorCode.ATTENTION_CORE_MISMATCH, ), ( {"strict_schedule": "different_schedule"}, @@ -861,6 +857,81 @@ def test_strict_runtime_readback_rejects_non_shared_attention_arithmetic(changes assert result.issues_by_code(expected_code) +@pytest.mark.parametrize( + ("actual_backend", "core_id", "schedule", "communication_backend"), + [ + ( + "flash_attention_4.cute", + "rlkernel.attention.flash_attention4.num_splits1.v1", + "single_batch_flash_attention4_num_splits1", + "cuda_ag_rs", + ), + ( + "aiter.rocm.ck_dense_mha", + "rlkernel.attention.aiter_ck.num_splits1.v1", + "single_batch_aiter_ck_num_splits1", + "rccl_ag_rs", + ), + ], +) +def test_strict_runtime_readback_accepts_same_qualified_vendor_core( + actual_backend, core_id, schedule, communication_backend +): + rollout = replace( + _strict_readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + strict_core_id=core_id, + strict_schedule=schedule, + native_attention_arithmetic=True, + actual_backend=actual_backend, + communication_backend=communication_backend, + ) + training = replace( + _strict_readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ), + strict_core_id=core_id, + strict_schedule=schedule, + native_attention_arithmetic=True, + actual_backend=actual_backend, + communication_backend=communication_backend, + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id=actual_backend, + training_backend_id=actual_backend, + ) + + assert result.passed + + +def test_strict_runtime_readback_rejects_vendor_arithmetic_mismatch(): + rollout = _strict_readback( + VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback" + ) + training = replace( + _strict_readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ), + native_attention_arithmetic=True, + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id=rollout.actual_backend, + training_backend_id=training.actual_backend, + ) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.ATTENTION_NATIVE_ARITHMETIC) + + def test_strict_runtime_readback_accepts_shared_no_split_k_core(): rollout = _strict_readback( VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback" @@ -895,10 +966,11 @@ def test_strict_runtime_readback_accepts_shared_no_split_k_core(): ("field", "value", "error_code"), [ ( - "actual_backend", - "rlkernel.attention.cp_reference", + "reference_only", + True, BindingErrorCode.ATTENTION_BACKEND_MISSING, ), + ("attention_fallback", True, BindingErrorCode.ATTENTION_BACKEND_MISSING), ( "communication_backend", "p2p_nccl_reference", From d0d572af0f4e222a56e18a3f59dbafbbeb7379d7 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 20 Aug 2026 23:51:30 +0800 Subject: [PATCH 22/33] fix(attention): fail closed on missing ROCm integration --- rl_engine/kernels/attention_preprocess.py | 14 ++++++++++---- .../kernels/ops/pytorch/attention/debug_matrix.py | 5 ++++- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index 9833d022..746990de 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -224,8 +224,8 @@ def __init__( from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op - self.rmsnorm = RMSNormCudaOp() - self.rope = RoPESM90Op() + self.rmsnorm: Callable[..., Tensor] = RMSNormCudaOp() + self.rope: Callable[..., Tensor] = RoPESM90Op() self.deterministic_backend_ids = MANDATED_ATTENTION_PREPROCESS_BACKENDS if not isinstance(policy_id, str) or not policy_id.strip(): raise ValueError("policy_id must be a non-empty string") @@ -370,11 +370,17 @@ def __init__( if self.device.index is None: self.device = torch.device("cuda", current_device) - from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RocmDeterministicRoPEOp + from rl_engine.kernels.ops.cuda.rotary_embedding import rope as rope_module from rl_engine.kernels.ops.triton.rmsnorm_triton import RMSNormTritonOp + rocm_rope_type = getattr(rope_module, "RocmDeterministicRoPEOp", None) + if rocm_rope_type is None: + raise RuntimeError( + "ROCm Attention preprocessing requires RocmDeterministicRoPEOp " + "from the ROCm Attention integration" + ) self.rmsnorm = RMSNormTritonOp() - self.rope = RocmDeterministicRoPEOp() + self.rope = rocm_rope_type() self.deterministic_backend_ids = ROCM_ATTENTION_PREPROCESS_BACKENDS self.device_capability = (0, 0) if native_qk_norm is None and reuse_transformer_engine_qk_norm: diff --git a/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py b/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py index cbceab9c..ee76e5f8 100644 --- a/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py +++ b/rl_engine/kernels/ops/pytorch/attention/debug_matrix.py @@ -143,7 +143,10 @@ def validate_attention_debug_matrix() -> None: axis_by_id = {axis.axis_id: axis for axis in ATTENTION_DEBUG_AXES} for row in diagnostic_rows: - axis = axis_by_id[row.root_cause_axis] + axis_id = row.root_cause_axis + if axis_id is None: + raise RuntimeError(f"diagnostic row {row.row_id} must name a root-cause axis") + axis = axis_by_id[axis_id] if row.probe != axis.representative_subprobe: raise RuntimeError( f"matrix row {row.row_id} is not the representative probe for {axis.axis_id}" From ee1e83ff3cb8f03d8b299c33be0d69466b1ae511 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 20 Aug 2026 23:56:52 +0800 Subject: [PATCH 23/33] fix(attention): close strict vendor admission gaps --- .../kernels/ops/pytorch/attention/ablation.py | 15 +++- tests/test_attention_ablation.py | 78 +++++++++++++++++-- 2 files changed, 84 insertions(+), 9 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index 12deabe1..bea6fb52 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -238,9 +238,17 @@ def apply( ) if cfg.validate: self._validate_inputs(q, k, v, contract) - if cfg.deterministic and contract.split_kv.mode is SplitKVMode.AUTO: + if cfg.deterministic and contract.split_kv.mode is not SplitKVMode.DISABLED: raise AttentionContractError( - "deterministic Attention cannot use runtime-dependent Split-KV=auto" + "strict deterministic Attention requires Split-KV to be disabled" + ) + if ( + cfg.deterministic + and contract.sharding.cp_world_size > 1 + and cfg.communication_backend not in _STRICT_AG_RS_BACKENDS + ): + raise AttentionContractError( + "strict CP Attention requires an explicit CUDA AG/RS or ROCm RCCL AG/RS backend" ) if cfg.deterministic and not cfg.return_lse: raise AttentionContractError("strict deterministic Attention must return LSE") @@ -604,6 +612,7 @@ def _validate_runtime_provenance( expected = { "strict_core_id": config.strict_core_id, "strict_schedule": config.strict_schedule, + "actual_backend": selected_id, "production_ready": True, "fallback": False, "reference_only": False, @@ -611,8 +620,6 @@ def _validate_runtime_provenance( if config.communication_backend in _STRICT_AG_RS_BACKENDS: expected["communication_backend"] = config.communication_backend mismatches = [name for name, value in expected.items() if runtime.get(name) != value] - if not isinstance(runtime.get("actual_backend"), str) or not runtime["actual_backend"].strip(): - mismatches.append("actual_backend") if not isinstance(runtime.get("native_attention_arithmetic"), bool): mismatches.append("native_attention_arithmetic") if mismatches: diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py index ae36237c..e07201da 100644 --- a/tests/test_attention_ablation.py +++ b/tests/test_attention_ablation.py @@ -255,7 +255,7 @@ def __call__(self, q, k, v, *, causal, scale): provenance={ "strict_core_id": STRICT_ATTENTION_CORE_ID, "strict_schedule": STRICT_ATTENTION_SCHEDULE_ID, - "actual_backend": "rlkernel.cuda.deterministic_attention", + "actual_backend": self.backend_id, "communication_backend": "self_owned_cuda_ag_rs", "production_ready": True, "native_attention_arithmetic": False, @@ -274,7 +274,7 @@ def __call__(self, q, k, v, *, causal, scale): contract=_contract(), backend=StrictCPBackend(), ) - assert result.provenance["actual_backend"] == "rlkernel.cuda.deterministic_attention" + assert result.provenance["actual_backend"] == StrictCPBackend.backend_id assert result.provenance["communication_backend"] == "self_owned_cuda_ag_rs" assert result.provenance["production_ready"] is True @@ -330,13 +330,81 @@ def __call__(self, q, k, v, *, causal, scale): ) -def test_deterministic_attention_rejects_runtime_split_kv_auto(): +@pytest.mark.parametrize( + "split_kv", + [SplitKVSpec.auto(strict_consistency=False), SplitKVSpec.fixed(2)], +) +def test_deterministic_attention_requires_split_kv_disabled(split_kv): q, k, v = _qkv() - contract = _contract(split_kv=SplitKVSpec.auto(strict_consistency=False)) - with pytest.raises(AttentionContractError, match="Split-KV=auto"): + contract = _contract(split_kv=split_kv) + with pytest.raises(AttentionContractError, match="Split-KV to be disabled"): AttentionAblationOp()(q, k, v, contract=contract) +def test_explicit_cp_callable_cannot_bypass_ag_rs_requirement(): + q, k, v = _qkv() + + class VendorCPBackend: + backend_id = "vendor.strict_attention" + core_id = "vendor.strict_core.v1" + strict_schedule = "vendor.strict_schedule.v1" + + def __call__(self, q, k, v, *, causal, scale, cp_world_size): + del k, v, causal, scale, cp_world_size + return q, torch.zeros(q.shape[:3], dtype=torch.float32) + + with pytest.raises(AttentionContractError, match="requires an explicit CUDA AG/RS"): + AttentionAblationOp()( + q[:, :, :2], + k[:, :, :2], + v[:, :, :2], + contract=_cp2_contract(), + backend=VendorCPBackend(), + config=AttentionAblationConfig( + strict_core_id=VendorCPBackend.core_id, + strict_schedule=VendorCPBackend.strict_schedule, + ), + ) + + +def test_vendor_core_actual_backend_must_match_selected_backend(): + q, k, v = _qkv() + + class VendorCore: + backend_id = "vendor.strict_attention" + core_id = "vendor.strict_core.v1" + strict_schedule = "vendor.strict_schedule.v1" + + def __call__(self, q, k, v, *, causal, scale): + del k, v, causal, scale + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:3], dtype=torch.float32), + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": "vendor.other_attention", + "production_ready": True, + "native_attention_arithmetic": True, + "fallback": False, + "reference_only": False, + }, + ) + + with pytest.raises(AttentionContractError, match="actual_backend"): + AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + backend=VendorCore(), + config=AttentionAblationConfig( + strict_core_id=VendorCore.core_id, + strict_schedule=VendorCore.strict_schedule, + ), + ) + + def test_deterministic_native_backend_requires_explicit_native_callable(): q, k, v = _qkv() with pytest.raises(AttentionContractError, match="native Attention backend"): From 00ee500d22ce35978a70e9fee9663f5db95f8e03 Mon Sep 17 00:00:00 2001 From: Takumi <3051000145@qq.com> Date: Fri, 21 Aug 2026 16:14:32 +0800 Subject: [PATCH 24/33] feat(alignment): add offline cross-config drift report --- docs/design/ws2-module-debug-matrix.md | 10 + docs/usage/cross-config-drift-report.md | 40 + pyproject.toml | 97 +- rl_engine/alignment/cross_config/__main__.py | 331 ++-- .../alignment/cross_config/drift_report.py | 1413 +++++++++++++++++ .../alignment/cross_config/drift_viewer.py | 316 ++++ setup.py | 618 +++---- tests/test_cross_config_drift_report.py | 344 ++++ 8 files changed, 2680 insertions(+), 489 deletions(-) create mode 100644 docs/usage/cross-config-drift-report.md create mode 100644 rl_engine/alignment/cross_config/drift_report.py create mode 100644 rl_engine/alignment/cross_config/drift_viewer.py create mode 100644 tests/test_cross_config_drift_report.py diff --git a/docs/design/ws2-module-debug-matrix.md b/docs/design/ws2-module-debug-matrix.md index 046fadb6..89302c9c 100644 --- a/docs/design/ws2-module-debug-matrix.md +++ b/docs/design/ws2-module-debug-matrix.md @@ -41,3 +41,13 @@ logical ownership are intentionally not exposed as mismatch axes. Use the compact row first. The operator-specific secondary probes remain available for a focused follow-up only after the representative row reproduces the mismatch. + +## Report artifacts + +The matrix reporter is owned by RL-Kernel because it consumes the sealed +cross-configuration attempt contract. It reads `actual.json`, +`comparison.json`, and `token_diffs.pt` only after `COMPLETE` validation, so +the report displays actual selected operators and the materialized topology, +not merely requested flags. See +[Cross-Configuration Drift Report](../usage/cross-config-drift-report.md) for +the `.rlk-drift` bundle, desktop viewer, static image, and trace entry points. diff --git a/docs/usage/cross-config-drift-report.md b/docs/usage/cross-config-drift-report.md new file mode 100644 index 00000000..3902bc5c --- /dev/null +++ b/docs/usage/cross-config-drift-report.md @@ -0,0 +1,40 @@ +# Cross-Configuration Drift Report + +The cross-configuration runner writes a sealed, append-only attempt directory. +The report tool reads that directory after `COMPLETE` has been published and +renders the recorded comparison rather than re-running a model. It shows actual +operator provenance, requested TP/CP and dtype axes, the fixed comparison +threshold, and the worst active selected-token delta. + +Generate the offline desktop bundle from one completed attempt: + +```bash +python -m rl_engine.alignment.cross_config report \ + runs//cases//attempt-0001 \ + --output /tmp/qwen3-tp2-cp2.rlk-drift +``` + +The `.rlk-drift` file contains sanitized report JSON, a Chrome Trace Event JSON +trace, and a PNG preview. It excludes checkpoints, prompts, and raw score +tensors. Install the optional local viewer and open it without a browser: + +```bash +python -m pip install "rl-engine[drift-viewer]" +rlk-drift-view /tmp/qwen3-tp2-cp2.rlk-drift +``` + +The viewer has an expandable track tree, horizontal zoom/scroll, selectable +events, and an event-details panel. The horizontal scale is explicitly marked +as a sample ordinal unless the source artifacts contain real timestamps. + +Other output suffixes expose the same sealed evidence in a form suitable for a +specific workflow: + +| Suffix | Artifact | Intended use | +| --- | --- | --- | +| `.rlk-drift` | Offline desktop bundle | Interactive post-training triage | +| `.png`, `.jpg` | Static summary | Pull request or issue attachment | +| `.json` | Chrome Trace Event | Perfetto or trace-viewer inspection | + +An unsealed or malformed attempt is rejected. A failed identity gate stays +`not comparable`; the tool does not turn it into a numerical drift value. diff --git a/pyproject.toml b/pyproject.toml index 7ffe391c..3aa5ded1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,51 +1,52 @@ -[build-system] -requires = ["setuptools>=64", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "RL-Kernel" -version = "0.1.0" -description = "High-performance RL training engine focused on kernel fusion and memory efficiency." -readme = "README.md" -requires-python = ">=3.10" -license = {text = "Apache-2.0"} -authors = [ - {name = "RL-Kernel Contributors"} -] -dependencies = [ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", -] - -[project.optional-dependencies] -cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] +[build-system] +requires = ["setuptools>=64", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "RL-Kernel" +version = "0.1.0" +description = "High-performance RL training engine focused on kernel fusion and memory efficiency." +readme = "README.md" +requires-python = ">=3.10" +license = {text = "Apache-2.0"} +authors = [ + {name = "RL-Kernel Contributors"} +] +dependencies = [ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", +] + +[project.optional-dependencies] +cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] rocm = ["aiter"] vllm = ["vllm>=0.6.0"] +drift-viewer = ["Pillow>=10", "PySide6>=6.6"] dev = ["pytest", "black", "isort", "ruff", "mypy", "pre-commit"] - -[tool.setuptools.packages.find] -where = ["."] -include = ["rl_engine*"] - -[tool.ruff] -line-length = 100 - -[tool.ruff.lint] -select = ["E", "F", "B"] -ignore = [] - -[tool.ruff.lint.per-file-ignores] -"__init__.py" = ["F401"] - -[tool.mypy] -ignore_missing_imports = true -follow_imports = "silent" - -[tool.pytest.ini_options] -markers = [ - "smoke_operator: temporary smoke-only operator plumbing tests", - "unit: CPU-safe unit tests", -] + +[tool.setuptools.packages.find] +where = ["."] +include = ["rl_engine*"] + +[tool.ruff] +line-length = 100 + +[tool.ruff.lint] +select = ["E", "F", "B"] +ignore = [] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["F401"] + +[tool.mypy] +ignore_missing_imports = true +follow_imports = "silent" + +[tool.pytest.ini_options] +markers = [ + "smoke_operator: temporary smoke-only operator plumbing tests", + "unit: CPU-safe unit tests", +] diff --git a/rl_engine/alignment/cross_config/__main__.py b/rl_engine/alignment/cross_config/__main__.py index cd047ca3..058173a4 100644 --- a/rl_engine/alignment/cross_config/__main__.py +++ b/rl_engine/alignment/cross_config/__main__.py @@ -1,138 +1,199 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -"""Command-line interface for cross-configuration experiments.""" - -from __future__ import annotations - -import argparse -import json -import sys -from pathlib import Path -from typing import Any, Optional, Sequence - -from rl_engine.alignment.cross_config.artifacts import ArtifactStore -from rl_engine.alignment.cross_config.config import ExperimentConfig, load_config -from rl_engine.alignment.cross_config.execution_plan import build_execution_plan - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description=__doc__) - commands = parser.add_subparsers(dest="command", required=True) - +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Command-line interface for cross-configuration experiments.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any, Optional, Sequence + +from rl_engine.alignment.cross_config.artifacts import ArtifactStore +from rl_engine.alignment.cross_config.config import ExperimentConfig, load_config +from rl_engine.alignment.cross_config.execution_plan import build_execution_plan + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + plan = commands.add_parser("plan", help="validate and persist a plan without execution") - _add_common_arguments(plan) - + _add_common_arguments(plan) + run = commands.add_parser("run", help="execute a plan with an explicit runtime adapter") - _add_common_arguments(run) - run.add_argument( - "--runtime", - required=True, - choices=("cpu-smoke",), - help="Runtime adapter; only the temporary CPU smoke adapter ships in V1", - ) - run.add_argument( - "--allow-smoke-operators", - action="store_true", - help="Explicitly authorize temporary smoke-only operator backends", - ) - run.add_argument( - "--timeout-seconds", - type=float, - default=30.0, - help="Per paired-scoring attempt deadline", - ) - run.add_argument( - "--no-resume", - action="store_true", - help="Create new attempts even when matching COMPLETE artifacts exist", - ) - return parser - - -def _add_common_arguments(parser: argparse.ArgumentParser) -> None: - parser.add_argument("config", type=Path, help="Versioned experiment JSON") - parser.add_argument( - "--output-root", - type=Path, - default=Path("runs"), - help="Append-only artifact root (default: runs)", - ) - - -def main(argv: Optional[Sequence[str]] = None) -> int: - args = build_parser().parse_args(argv) - try: - config = load_config(args.config) - if args.command == "plan": - summary = record_plan(config, args.output_root) - else: - summary = _run(config, args) - except Exception as exc: - summary = { - "schema_version": "cross_config.cli_summary.v1", - "status": "error", - "error_type": f"{type(exc).__module__}.{type(exc).__qualname__}", - "error": str(exc), - } - print(json.dumps(summary, sort_keys=True)) + _add_common_arguments(run) + run.add_argument( + "--runtime", + required=True, + choices=("cpu-smoke",), + help="Runtime adapter; only the temporary CPU smoke adapter ships in V1", + ) + run.add_argument( + "--allow-smoke-operators", + action="store_true", + help="Explicitly authorize temporary smoke-only operator backends", + ) + run.add_argument( + "--timeout-seconds", + type=float, + default=30.0, + help="Per paired-scoring attempt deadline", + ) + run.add_argument( + "--no-resume", + action="store_true", + help="Create new attempts even when matching COMPLETE artifacts exist", + ) + + report = commands.add_parser( + "report", + help="render a sealed attempt as an offline drift report", + ) + report.add_argument( + "attempt_dir", type=Path, help="Completed cross-config attempt directory" + ) + report.add_argument( + "--output", + "-o", + type=Path, + required=True, + help="Output .rlk-drift bundle, .png/.jpg image, or .json trace", + ) + report.add_argument("--title", default=None, help="Optional report title") + report.add_argument( + "--no-preview", action="store_true", help="Omit PNG from a .rlk-drift bundle" + ) + return parser + + +def _add_common_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("config", type=Path, help="Versioned experiment JSON") + parser.add_argument( + "--output-root", + type=Path, + default=Path("runs"), + help="Append-only artifact root (default: runs)", + ) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = build_parser().parse_args(argv) + try: + if args.command == "report": + summary = _report(args) + else: + config = load_config(args.config) + if args.command == "plan": + summary = record_plan(config, args.output_root) + else: + summary = _run(config, args) + except Exception as exc: + summary = { + "schema_version": "cross_config.cli_summary.v1", + "status": "error", + "error_type": f"{type(exc).__module__}.{type(exc).__qualname__}", + "error": str(exc), + } + print(json.dumps(summary, sort_keys=True)) print(f"cross-configuration error: {type(exc).__name__}: {exc}", file=sys.stderr) - return 2 - - print(json.dumps(summary, sort_keys=True)) - if args.command == "plan": - print( - f"planned {summary['planned_case_count']} cases; no runtime was created", - file=sys.stderr, - ) - return 0 - print( - f"CPU smoke: {summary['status']} ({len(summary['cases'])} cases)", - file=sys.stderr, - ) - for case in summary["cases"]: - print( - f" {case['case_id']}: {case['status']}; actual backends " - f"rollout={case['rollout_backend']}, training={case['training_backend']}; " - f"worst sample/token={case['worst_token_index']}; " - f"mismatches={case['mismatch_count']}; resumed={case['resumed']}", - file=sys.stderr, - ) - return 0 if summary["status"] == "pass" else 1 - - -def record_plan(config: ExperimentConfig, output_root: Path) -> dict[str, Any]: - plan = build_execution_plan(config) - store = ArtifactStore(output_root) - experiment_dir = store.initialize_experiment( - config.definition.experiment_id, - experiment=plan.experiment, - plan=plan.rows(), - ) - return { - "schema_version": "cross_config.cli_summary.v1", - "status": "planned", - "experiment_id": config.definition.experiment_id, - "scenario_id": config.definition.scenario_id, - "planned_case_count": len(plan.entries), - "planning_issues": [issue.to_dict() for issue in plan.issues], - "artifact_dir": str(experiment_dir), - } - - -def _run(config: ExperimentConfig, args: argparse.Namespace) -> dict[str, Any]: - if args.runtime != "cpu-smoke": # pragma: no cover - argparse owns choices - raise ValueError(f"unsupported runtime {args.runtime!r}") - from rl_engine.alignment.testing.cpu_cross_config import run_cpu_experiment - - return run_cpu_experiment( - config, - output_root=args.output_root, - allow_smoke_operators=args.allow_smoke_operators, - timeout_seconds=args.timeout_seconds, - resume=not args.no_resume, - ) - - -if __name__ == "__main__": # pragma: no cover - raise SystemExit(main()) + return 2 + + print(json.dumps(summary, sort_keys=True)) + if args.command == "plan": + print( + f"planned {summary['planned_case_count']} cases; no runtime was created", + file=sys.stderr, + ) + return 0 + if args.command == "report": + print( + f"drift report: {summary['status']} ({summary['output']})", + file=sys.stderr, + ) + return 0 + print( + f"CPU smoke: {summary['status']} ({len(summary['cases'])} cases)", + file=sys.stderr, + ) + for case in summary["cases"]: + print( + f" {case['case_id']}: {case['status']}; actual backends " + f"rollout={case['rollout_backend']}, training={case['training_backend']}; " + f"worst sample/token={case['worst_token_index']}; " + f"mismatches={case['mismatch_count']}; resumed={case['resumed']}", + file=sys.stderr, + ) + return 0 if summary["status"] == "pass" else 1 + + +def record_plan(config: ExperimentConfig, output_root: Path) -> dict[str, Any]: + plan = build_execution_plan(config) + store = ArtifactStore(output_root) + experiment_dir = store.initialize_experiment( + config.definition.experiment_id, + experiment=plan.experiment, + plan=plan.rows(), + ) + return { + "schema_version": "cross_config.cli_summary.v1", + "status": "planned", + "experiment_id": config.definition.experiment_id, + "scenario_id": config.definition.scenario_id, + "planned_case_count": len(plan.entries), + "planning_issues": [issue.to_dict() for issue in plan.issues], + "artifact_dir": str(experiment_dir), + } + + +def _run(config: ExperimentConfig, args: argparse.Namespace) -> dict[str, Any]: + if args.runtime != "cpu-smoke": # pragma: no cover - argparse owns choices + raise ValueError(f"unsupported runtime {args.runtime!r}") + from rl_engine.alignment.testing.cpu_cross_config import run_cpu_experiment + + return run_cpu_experiment( + config, + output_root=args.output_root, + allow_smoke_operators=args.allow_smoke_operators, + timeout_seconds=args.timeout_seconds, + resume=not args.no_resume, + ) + + +def _report(args: argparse.Namespace) -> dict[str, Any]: + from rl_engine.alignment.cross_config.drift_report import ( + build_cross_config_attempt_report, + write_drift_bundle, + write_drift_report_image, + write_drift_trace, + ) + + report = build_cross_config_attempt_report(args.attempt_dir, title=args.title) + suffix = args.output.suffix.lower() + if suffix == ".rlk-drift": + output = write_drift_bundle( + report, args.output, include_preview=not args.no_preview + ) + artifact_type = "bundle" + elif suffix in {".png", ".jpg", ".jpeg"}: + output = write_drift_report_image(report, args.output) + artifact_type = "image" + elif suffix == ".json": + output = write_drift_trace(report, args.output) + artifact_type = "trace" + else: + raise ValueError("report output must use .rlk-drift, .png/.jpg, or .json") + return { + "schema_version": "cross_config.drift_report_summary.v1", + "status": report["status"], + "artifact_type": artifact_type, + "output": str(output), + "case_id": report["axes"].get("case_id"), + "attempt_id": report["axes"].get("attempt_id"), + } + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/rl_engine/alignment/cross_config/drift_report.py b/rl_engine/alignment/cross_config/drift_report.py new file mode 100644 index 00000000..e8771239 --- /dev/null +++ b/rl_engine/alignment/cross_config/drift_report.py @@ -0,0 +1,1413 @@ +"""Build self-contained profiler-style reports from cross-config artifacts. + +The reporter is deliberately downstream of execution. It validates a sealed +attempt directory, then renders the recorded comparison, token deltas, actual +operator provenance, and materialized topology without changing either the +training or rollout path. +""" + +from __future__ import annotations + +import html +import io +import json +import zipfile +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch + +from rl_engine.alignment.cross_config.artifacts import ArtifactStore + + +REPORT_SCHEMA_VERSION = 2 +_STATUS_LABELS = {"pass": "PASS", "warning": "WARN", "failure": "FAIL", "info": "INFO"} +_STATUS_COLORS = { + "pass": "#39d98a", + "warning": "#f4b942", + "failure": "#ff5c77", + "info": "#70a7ff", +} +_REPORT_IMAGE_COLORS = { + "background": "#eef0f2", + "panel": "#ffffff", + "panel_alt": "#f7f8f9", + "track": "#c8cdd2", + "line": "#cbd1d7", + "grid": "#e5e8eb", + "text": "#20252b", + "muted": "#66707a", + "green": "#43a64b", + "yellow": "#d79700", + "red": "#d64242", + "blue": "#3677c8", + "purple": "#7659bb", + "pass": "#43a64b", + "warning": "#d79700", + "failure": "#d64242", + "info": "#3677c8", +} + + +def build_cross_config_attempt_report( + attempt_dir: str | Path, + *, + title: str | None = None, + validate: bool = True, +) -> dict[str, Any]: + """Render one immutable cross-configuration attempt into a drift report. + + ``attempt_dir`` must be a directory emitted by :class:`PairedRunner`. The + report never derives a numerical claim from an incomplete attempt. A + caller can set ``validate=False`` only when it has already validated the + same sealed directory in the current process. + """ + + directory = Path(attempt_dir) + store = ArtifactStore(directory.parent) + if validate: + store.validate_completed_attempt(directory) + + comparison = _read_attempt_json(directory / "comparison.json") + actual = _read_attempt_json(directory / "actual.json") + materialized = _read_attempt_json(directory / "materialized.json") + identity = _read_attempt_json(directory / "identity.json") + token_payload = store.load_tensor_bundle(directory / "token_diffs.pt") + tensors = token_payload["tensors"] + active_mask = tensors["active_mask"].to(dtype=torch.bool) + absolute_diff = tensors["absolute_diff"].to(dtype=torch.float32) + active_count = int(active_mask.sum().item()) + worst_token = _worst_token(absolute_diff, active_mask) + + materialized_case = _plain_mapping(materialized.get("materialized_case")) + case = _plain_mapping(materialized_case.get("case")) + requested = _plain_mapping(case.get("requested")) + rollout_requested = _plain_mapping(requested.get("rollout")) + training_requested = _plain_mapping(requested.get("training")) + rollout_actual = _plain_mapping(actual.get("rollout")) + training_actual = _plain_mapping(actual.get("training")) + diagnostics = _plain_mapping(comparison.get("diagnostics")) + + failures: list[dict[str, Any]] = [] + if not bool(comparison.get("comparable")): + failures.append( + { + "code": "not_comparable", + "message": "The sealed attempt failed an identity or artifact comparison gate.", + } + ) + for field in ("identity_errors", "artifact_errors"): + for value in comparison.get(field) or (): + failures.append({"code": field, "message": str(value)}) + if not bool(comparison.get("passed")) and not failures: + failures.append( + { + "code": str(comparison.get("status", "comparison_failed")), + "message": "The selected-token comparison exceeded its fixed contract.", + } + ) + + metrics = { + "active_token_count": active_count, + "mismatch_count": int(comparison.get("mismatch_count", 0)), + "max_abs_dlogp": 0.0 if worst_token is None else worst_token["abs_dlogp"], + "fixed_threshold": comparison.get("fixed_threshold"), + "comparison_status": comparison.get("status"), + "comparable": bool(comparison.get("comparable")), + "passed": bool(comparison.get("passed")), + } + axes = { + "case_id": comparison.get("case_id"), + "attempt_id": comparison.get("attempt_id"), + "rollout_tp": rollout_requested.get("tensor_parallel_size"), + "rollout_cp": rollout_requested.get("context_parallel_size"), + "rollout_dtype": rollout_requested.get("dtype"), + "training_dtype": training_requested.get("compute_dtype"), + "training_sharding": training_requested.get("sharding"), + "contract_fingerprint": comparison.get("contract_fingerprint"), + } + rollout_backend = _first_present( + rollout_actual.get("backend_id"), rollout_actual.get("actual_backend") + ) + training_backend = _first_present( + training_actual.get("backend_id"), training_actual.get("actual_backend") + ) + provenance = { + "operator_source": actual.get("operator_source"), + "execution_fingerprint": actual.get("execution_fingerprint"), + "environment_fingerprint": actual.get("environment_fingerprint"), + "actual_backend": _backend_label(rollout_backend, training_backend), + "rollout_backend": rollout_backend, + "training_backend": training_backend, + "rollout": rollout_actual, + "training": training_actual, + } + samples = _samples_from_identity(identity, active_mask) + replay_manifest = { + "mode": "cross_config_fixed_replay", + "samples": samples, + "batch_invariance_cases": (), + "validation": {"warnings": (), "failures": failures}, + "runtime_provenance": provenance, + } + result_cube = { + "mode": "cross_config_fixed_replay", + "axes": axes, + "metrics": metrics, + "worst_token": worst_token or {}, + "metadata_validation": replay_manifest["validation"], + "runtime_provenance": provenance, + "diagnostics": diagnostics, + } + return build_drift_report( + replay_manifest=replay_manifest, + result_cube=result_cube, + runtime_provenance=provenance, + title=title + or f"Cross-config drift: {comparison.get('case_id', directory.parent.name)} / {directory.name}", + ) + + +def build_drift_report( + *, + replay_manifest: Mapping[str, Any] | None = None, + result_cube: Mapping[str, Any] | None = None, + runtime_provenance: Mapping[str, Any] | None = None, + title: str | None = None, +) -> dict[str, Any]: + """Normalize existing audit artifacts into a visual diagnostic report.""" + + manifest = _plain_mapping(replay_manifest) + cube = _plain_mapping(result_cube) + metrics = _plain_mapping(cube.get("metrics")) + axes = _plain_mapping(cube.get("axes")) + validation = _plain_mapping( + cube.get("metadata_validation") or manifest.get("validation") + ) + provenance = _plain_mapping( + runtime_provenance + or cube.get("runtime_provenance") + or manifest.get("runtime_provenance") + ) + samples = [ + _plain_mapping(item) + for item in manifest.get("samples", []) + if isinstance(item, Mapping) + ] + warnings = [ + _plain_mapping(item) + for item in validation.get("warnings", []) + if isinstance(item, Mapping) + ] + failures = [ + _plain_mapping(item) + for item in validation.get("failures", []) + if isinstance(item, Mapping) + ] + + max_abs_dlogp = _number(metrics.get("max_abs_dlogp")) + warning_count = _number(metrics.get("warning_count"), default=0.0) or 0.0 + metadata_warning_count = ( + _number(metrics.get("metadata_warning_count"), default=float(len(warnings))) + or 0.0 + ) + metadata_failure_count = ( + _number(metrics.get("metadata_failure_count"), default=float(len(failures))) + or 0.0 + ) + runtime_fallback = bool( + metrics.get("runtime_fallback") or provenance.get("fallback") + ) + strict_failure = bool( + metrics.get("runtime_strict_failure") or provenance.get("strict_failure") + ) + comparison_passed = metrics.get("passed") is True + + if strict_failure or metadata_failure_count > 0: + status = "failure" + elif warning_count > 0 or metadata_warning_count > 0 or runtime_fallback: + status = "warning" + elif comparison_passed: + status = "pass" + elif (max_abs_dlogp or 0.0) > 0.0: + status = "warning" + else: + status = "pass" + + span = max(1.0, float(len(samples))) + has_timestamps = any( + _number(sample.get("start_ts")) is not None for sample in samples + ) + timeline_mode = "timestamp" if has_timestamps else "ordinal_diagnostic" + events: list[dict[str, Any]] = [] + + train_status = ( + "failure" + if status == "failure" + else "warning" if status == "warning" else "pass" + ) + events.append( + { + "id": "train-audit", + "kind": "bar", + "lane": "Training audit", + "start": 0.0, + "end": span, + "label": "training-side audit", + "status": train_status, + "details": { + "mode": manifest.get("mode", cube.get("mode", "unknown")), + "rank": cube.get("rank"), + }, + } + ) + for position, sample in enumerate(samples): + sample_label = sample.get("sample_index") + if sample_label is None: + sample_label = sample.get("rollout_id") + if sample_label is None: + sample_label = position + start = _number(sample.get("start_ts"), default=float(position)) + end = _number(sample.get("end_ts"), default=float(start or position) + 0.82) + if end is None or start is None or end <= start: + start, end = float(position), float(position) + 0.82 + events.append( + { + "id": f"rollout-{position}", + "kind": "bar", + "lane": "Rollout samples", + "start": start, + "end": end, + "label": f"sample {sample_label}", + "status": "info", + "details": sample, + } + ) + + actual_backend = _first_present( + provenance.get("actual_backend"), + provenance.get("backend_id"), + ( + cube.get("axes", {}).get("logp_backend") + if isinstance(cube.get("axes"), Mapping) + else None + ), + provenance.get("requested_backend"), + "unknown backend", + ) + operator_status = ( + "failure" if strict_failure else "warning" if runtime_fallback else "pass" + ) + events.append( + { + "id": "operator-backend", + "kind": "bar", + "lane": "Operator / backend", + "start": 0.12, + "end": max(0.94, span - 0.12), + "label": str(actual_backend), + "status": operator_status, + "details": provenance, + } + ) + + worst_token = _plain_mapping(cube.get("worst_token")) + if worst_token: + marker_position = ( + _number(worst_token.get("sample_position"), default=max(0.0, span - 0.5)) + or 0.0 + ) + events.append( + { + "id": "worst-drift", + "kind": "marker", + "lane": "Drift markers", + "start": marker_position + 0.41, + "end": marker_position + 0.41, + "label": f"|dlogp| {_format_number(worst_token.get('abs_dlogp'))}", + "status": ( + "failure" + if status == "failure" + else "warning" if status == "warning" else "pass" + ), + "details": worst_token, + } + ) + if warnings or failures or warning_count: + marker_status = "failure" if failures else "warning" + events.append( + { + "id": "validation-marker", + "kind": "marker", + "lane": "Drift markers", + "start": max(0.2, span - 0.22), + "end": max(0.2, span - 0.22), + "label": f"{len(failures)} failures / {len(warnings)} warnings", + "status": marker_status, + "details": { + "warnings": warnings, + "failures": failures, + "dlogp_warning_count": warning_count, + }, + } + ) + + return { + "schema_version": REPORT_SCHEMA_VERSION, + "title": title or "RL-Kernel cross-config drift report", + "status": status, + "status_label": _STATUS_LABELS[status], + "timeline_mode": timeline_mode, + "timeline_note": ( + "Real artifact timestamps are shown." + if timeline_mode == "timestamp" + else "No artifact timestamps were available; positions are stable sample ordinals, not elapsed time." + ), + "lanes": [ + "Training audit", + "Rollout samples", + "Operator / backend", + "Drift markers", + ], + "events": events, + "axes": axes, + "metrics": metrics, + "worst_token": worst_token, + "validation": validation, + "runtime_provenance": provenance, + "sample_count": len(samples), + "replay_case_count": len(manifest.get("batch_invariance_cases", []) or []), + "manifest_fingerprint": manifest.get("fingerprint"), + "cube_fingerprint": cube.get("fingerprint"), + } + + +def build_drift_trace(report: Mapping[str, Any]) -> dict[str, Any]: + """Convert a consistency report into Chrome Trace Event JSON. + + The resulting file can be opened directly in Perfetto and expanded by + process/thread track. It deliberately uses the report's timestamp mode; + ordinal reports remain diagnostic sample positions, never fabricated time. + """ + + normalized = _plain_mapping(report) + events = [_plain_mapping(event) for event in normalized.get("events", [])] + status = str(normalized.get("status", "info")) + timeline_mode = str(normalized.get("timeline_mode", "ordinal_diagnostic")) + lanes = [ + ("Audit", None), + ("Training audit", "Training audit"), + ("Rollout samples", "Rollout samples"), + ("Execution", None), + ("Operator / backend", "Operator / backend"), + ("Token comparison", "Token comparison"), + ("Validation", None), + ("Drift markers", "Drift markers"), + ] + lane_ids = {name: 100 + index for index, (name, _) in enumerate(lanes)} + process_id = 1 + trace_events: list[dict[str, Any]] = [ + { + "name": "process_name", + "ph": "M", + "pid": process_id, + "args": {"name": "RL-Kernel cross-config drift"}, + }, + ] + for lane, event_lane in lanes: + if event_lane is None: + continue + trace_events.append( + { + "name": "thread_name", + "ph": "M", + "pid": process_id, + "tid": lane_ids[lane], + "args": {"name": lane}, + } + ) + + def trace_time(value: Any) -> float: + number = _number(value, default=0.0) or 0.0 + # Chrome Trace timestamps are microseconds. In ordinal mode the same + # scale is retained only to make adjacent sample positions visible. + return number * 1_000_000.0 + + event_colors = { + "pass": "good", + "warning": "terrible", + "failure": "bad", + "info": "thread_state_running", + } + for event in events: + event_lane = str(event.get("lane", "Drift markers")) + tid = lane_ids.get(event_lane, lane_ids["Drift markers"]) + start = _number(event.get("start"), default=0.0) or 0.0 + end = _number(event.get("end"), default=start) or start + details = ( + event.get("details") if isinstance(event.get("details"), Mapping) else {} + ) + args = { + "event_id": str(event.get("id", "")), + "status": str(event.get("status", status)), + "timeline_mode": timeline_mode, + "timeline_note": normalized.get("timeline_note", ""), + "details": details, + } + color = event_colors.get( + str(event.get("status", "info")), "thread_state_running" + ) + if event.get("kind") == "marker": + trace_events.append( + { + "name": str(event.get("label", "marker")), + "cat": "consistency.drift", + "ph": "I", + "s": "t", + "pid": process_id, + "tid": tid, + "ts": trace_time(start), + "cname": color, + "args": args, + } + ) + continue + trace_events.append( + { + "name": str(event.get("label", event.get("id", "event"))), + "cat": "consistency.audit", + "ph": "X", + "pid": process_id, + "tid": tid, + "ts": trace_time(start), + "dur": max(1.0, trace_time(end) - trace_time(start)), + "cname": color, + "args": args, + } + ) + + metrics = ( + normalized.get("metrics") + if isinstance(normalized.get("metrics"), Mapping) + else {} + ) + max_abs_dlogp = _number(metrics.get("max_abs_dlogp")) + if max_abs_dlogp is not None: + trace_events.append( + { + "name": "max |dlogp|", + "cat": "consistency.metric", + "ph": "C", + "pid": process_id, + "tid": lane_ids["Token comparison"], + "ts": 0.0, + "args": {"max_abs_dlogp": max_abs_dlogp}, + } + ) + + return { + "traceEvents": trace_events, + "displayTimeUnit": "ms", + "metadata": { + "report_title": normalized.get( + "title", "RL-Kernel cross-config drift report" + ), + "status": status, + "timeline_mode": timeline_mode, + "timeline_note": normalized.get("timeline_note", ""), + "schema_version": normalized.get("schema_version", REPORT_SCHEMA_VERSION), + }, + } + + +def write_drift_trace(report: Mapping[str, Any], path: str | Path) -> Path: + """Write a Chrome Trace Event JSON file for Perfetto or trace viewers.""" + + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text( + json.dumps(build_drift_trace(report), ensure_ascii=True, indent=2), + encoding="utf-8", + ) + return output + + +def write_drift_bundle( + report: Mapping[str, Any], + path: str | Path, + *, + include_preview: bool = True, +) -> Path: + """Write a self-contained ``.rlk-drift`` desktop-viewer bundle. + + The bundle contains sanitized report JSON and a portable trace. The PNG + preview is included for PR/issue sharing but is not required by the viewer. + Raw train dumps are intentionally not copied into the bundle. + """ + + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + normalized = _plain_mapping(report) + trace = build_drift_trace(normalized) + manifest = { + "format": "rl_kernel.cross_config_drift", + "bundle_version": 1, + "report_schema_version": normalized.get( + "schema_version", REPORT_SCHEMA_VERSION + ), + "files": [ + "manifest.json", + "report.json", + "trace.json", + *(["preview.png"] if include_preview else []), + ], + "preview_included": bool(include_preview), + "title": normalized.get("title", "RL-Kernel cross-config drift report"), + "status": normalized.get("status", "info"), + } + with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as archive: + archive.writestr( + "manifest.json", json.dumps(manifest, ensure_ascii=True, indent=2) + ) + archive.writestr( + "report.json", json.dumps(normalized, ensure_ascii=True, indent=2) + ) + archive.writestr("trace.json", json.dumps(trace, ensure_ascii=True, indent=2)) + if include_preview: + preview = render_drift_report_image(normalized) + buffer = io.BytesIO() + preview.save(buffer, format="PNG", optimize=True) + archive.writestr("preview.png", buffer.getvalue()) + return output + + +def load_drift_bundle(path: str | Path) -> dict[str, Any]: + """Load a ``.rlk-drift`` bundle without importing the optional GUI.""" + + with zipfile.ZipFile(path, "r") as archive: + names = set(archive.namelist()) + if "report.json" not in names or "trace.json" not in names: + raise ValueError( + "invalid RL-Kernel drift bundle: report.json and trace.json are required" + ) + report = json.loads(archive.read("report.json")) + trace = json.loads(archive.read("trace.json")) + manifest = ( + json.loads(archive.read("manifest.json")) + if "manifest.json" in names + else {} + ) + return {"manifest": manifest, "report": report, "trace": trace} + + +def render_drift_report(report: Mapping[str, Any]) -> str: + """Render a report as a self-contained HTML document with an SVG timeline.""" + + normalized = _plain_mapping(report) + events = [_plain_mapping(event) for event in normalized.get("events", [])] + lanes = [str(lane) for lane in normalized.get("lanes", [])] + width = 1180 + left = 190 + right = 28 + top = 46 + row_height = 54 + timeline_width = width - left - right + timeline_span = max( + 1.0, + max( + (_number(event.get("end"), default=1.0) or 1.0 for event in events), + default=1.0, + ), + ) + svg_height = top + row_height * len(lanes) + 42 + event_map = {str(event.get("id")): event for event in events} + + def x(value: float) -> float: + return ( + left + max(0.0, min(timeline_span, value)) / timeline_span * timeline_width + ) + + svg_parts = [ + f'', + f'', + ] + grid_steps = min(12, max(2, int(timeline_span) + 1)) + for index in range(grid_steps + 1): + value = timeline_span * index / grid_steps + xpos = x(value) + svg_parts.append( + f'' + ) + svg_parts.append( + f'{html.escape(_format_number(value))}' + ) + for lane_index, lane in enumerate(lanes): + ypos = top + lane_index * row_height + svg_parts.append( + f'{html.escape(lane)}' + ) + svg_parts.append( + f'' + ) + + for event in events: + lane_index = ( + lanes.index(str(event.get("lane"))) + if str(event.get("lane")) in lanes + else 0 + ) + ypos = top + lane_index * row_height + status = str(event.get("status", "info")) + color = _STATUS_COLORS.get(status, _STATUS_COLORS["info"]) + event_id = html.escape(str(event.get("id")), quote=True) + label = html.escape(_truncate(str(event.get("label", "event")), 28)) + title = html.escape( + f"{event.get('label', 'event')} [{_STATUS_LABELS.get(status, status.upper())}]", + quote=True, + ) + if event.get("kind") == "marker": + xpos = x(_number(event.get("start"), default=0.0) or 0.0) + points = f"{xpos:.2f},{ypos + 7} {xpos + 9:.2f},{ypos + 16} {xpos:.2f},{ypos + 25} {xpos - 9:.2f},{ypos + 16}" + svg_parts.append( + f'{title}' + ) + svg_parts.append( + f'{label}' + ) + else: + start = _number(event.get("start"), default=0.0) or 0.0 + end = _number(event.get("end"), default=start + 0.5) or start + 0.5 + xpos = x(start) + event_width = max(8.0, x(end) - xpos) + svg_parts.append( + f'{title}' + ) + if event_width > 60: + svg_parts.append( + f'{label}' + ) + else: + svg_parts.append( + f'{label}' + ) + svg_parts.append("") + svg = "".join(svg_parts) + + metrics = ( + normalized.get("metrics") + if isinstance(normalized.get("metrics"), Mapping) + else {} + ) + metric_cards = [ + ("Max |dlogp|", _format_number(metrics.get("max_abs_dlogp"))), + ("Active tokens", _format_number(metrics.get("active_token_count"))), + ("Warnings", _format_number(metrics.get("warning_count"), default="0")), + ("Replay cases", str(normalized.get("replay_case_count", 0))), + ] + cards_html = "".join( + f'
{html.escape(name)}
{html.escape(value)}
' + for name, value in metric_cards + ) + axes_html = _render_key_value_table( + normalized.get("axes"), empty="No normalized axes were recorded." + ) + provenance_html = _render_key_value_table( + normalized.get("runtime_provenance"), + empty="No runtime provenance was recorded.", + ) + validation_html = _render_validation(normalized.get("validation")) + event_json = json.dumps( + event_map, ensure_ascii=True, separators=(",", ":") + ).replace("<", "\\u003c") + status = str(normalized.get("status", "info")) + status_color = _STATUS_COLORS.get(status, _STATUS_COLORS["info"]) + title = html.escape( + str(normalized.get("title", "RL-Kernel cross-config drift report")) + ) + timeline_note = html.escape(str(normalized.get("timeline_note", ""))) + + return f""" + +{title} +
+

{title}

{timeline_note}
{html.escape(_STATUS_LABELS.get(status, status.upper()))}
+
{cards_html}
+

Operator drift timeline {html.escape(str(normalized.get("timeline_mode", "diagnostic")))}

{svg}
+

Selected event

Select a bar or marker in the timeline.
+

Normalized axes

{axes_html}

Runtime provenance

{provenance_html}
+

Validation

{validation_html}
+
""" + + +def write_drift_report(report: Mapping[str, Any], path: str | Path) -> Path: + """Write a self-contained HTML report and return its resolved path.""" + + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(render_drift_report(report), encoding="utf-8") + return output + + +def render_drift_report_image(report: Mapping[str, Any], *, width: int = 2400) -> Any: + """Render a static profiler-style report image. + + The image is intentionally self-contained and suitable for attaching to a + PR, issue, or debug artifact. It uses a light, high-density profiler + layout: a compact tool chrome, a hierarchical track tree, a fine-grained + ruler/grid, thin event bars, and tabular details below the timeline. + """ + + from PIL import Image, ImageDraw, ImageFont + + normalized = _plain_mapping(report) + metrics = ( + normalized.get("metrics") + if isinstance(normalized.get("metrics"), Mapping) + else {} + ) + events = [_plain_mapping(event) for event in normalized.get("events", [])] + status = str(normalized.get("status", "info")) + status_color = _REPORT_IMAGE_COLORS.get(status, _REPORT_IMAGE_COLORS["info"]) + bg = _REPORT_IMAGE_COLORS["background"] + panel_bg = _REPORT_IMAGE_COLORS["panel"] + panel_alt = _REPORT_IMAGE_COLORS["panel_alt"] + line = _REPORT_IMAGE_COLORS["line"] + text_color = _REPORT_IMAGE_COLORS["text"] + muted = _REPORT_IMAGE_COLORS["muted"] + + height = 1680 + image = Image.new("RGB", (width, height), bg) + draw = ImageDraw.Draw(image) + regular = _load_report_font(22) + small = _load_report_font(18) + tiny = _load_report_font(15) + label_font = _load_report_font(18, bold=True) + section_font = _load_report_font(20, bold=True) + title_font = _load_report_font(30, bold=True) + metric_font = _load_report_font(25, bold=True) + mono = _load_report_font(16, mono=True) + + def rect( + box: tuple[int, int, int, int], + fill: str, + radius: int = 2, + outline: str | None = None, + ) -> None: + draw.rectangle(box, fill=fill, outline=outline, width=1 if outline else 1) + + def write( + x: int, + y: int, + value: Any, + font: Any = regular, + fill: str = text_color, + anchor: str | None = None, + ) -> None: + draw.text((x, y), str(value), font=font, fill=fill, anchor=anchor) + + def fit(value: Any, limit: int) -> str: + return _truncate(str(value), limit) + + def key_value_rows(values: Any, limit: int = 33) -> list[tuple[str, str]]: + if not isinstance(values, Mapping): + return [] + return [ + (fit(key, limit), fit(_format_value(values[key]), 47)) + for key in sorted(values, key=str) + ] + + # Nsight-like tool chrome: compact title, run metadata, and controls. + draw.rectangle((0, 0, width, 42), fill="#252a30") + write(30, 12, "RL-KERNEL CROSS-CONFIG ANALYSIS", tiny, "#e9edf1") + write( + width - 30, + 12, + "static report | schema v" + + str(normalized.get("schema_version", REPORT_SCHEMA_VERSION)), + tiny, + "#b7c0c8", + anchor="ra", + ) + write( + 34, + 60, + normalized.get("title", "RL-Kernel cross-config drift report"), + title_font, + text_color, + ) + write(34, 101, "run / audit replay", tiny, muted) + write(180, 101, normalized.get("timeline_note", ""), tiny, muted) + badge_text = _STATUS_LABELS.get(status, status.upper()) + draw.rectangle( + (width - 188, 61, width - 34, 101), fill=panel_bg, outline=status_color, width=2 + ) + write(width - 111, 81, badge_text, label_font, status_color, anchor="mm") + + # Compact metric strip. These are values, not dashboard cards. + strip_y = 127 + draw.rectangle( + (34, strip_y, width - 34, strip_y + 64), fill=panel_bg, outline=line, width=1 + ) + strip_items = [ + ("MAX |DLOGP|", _format_number(metrics.get("max_abs_dlogp")), status_color), + ( + "ACTIVE TOKENS", + _format_number(metrics.get("active_token_count")), + _REPORT_IMAGE_COLORS["blue"], + ), + ( + "WARNINGS", + _format_number(metrics.get("warning_count"), default="0"), + _REPORT_IMAGE_COLORS["yellow"], + ), + ( + "REPLAY CASES", + normalized.get("replay_case_count", 0), + _REPORT_IMAGE_COLORS["purple"], + ), + ] + strip_width = (width - 68) // len(strip_items) + for index, (name, value, accent) in enumerate(strip_items): + x0 = 34 + index * strip_width + if index: + draw.line((x0, strip_y + 10, x0, strip_y + 54), fill=line, width=1) + draw.rectangle((x0, strip_y, x0 + 4, strip_y + 64), fill=accent) + write(x0 + 18, strip_y + 12, name, tiny, muted) + write(x0 + 18, strip_y + 34, value, metric_font, text_color) + + # Timeline panel: strict left track tree + dense aligned lanes. + timeline_y = 216 + timeline_h = 620 + rect((34, timeline_y, width - 34, timeline_y + timeline_h), panel_bg, 2, line) + write(52, timeline_y + 14, "CAPTURED EXECUTION", section_font, text_color) + mode = str(normalized.get("timeline_mode", "diagnostic")) + mode_color = ( + _REPORT_IMAGE_COLORS["yellow"] + if mode != "timestamp" + else _REPORT_IMAGE_COLORS["blue"] + ) + write(width - 52, timeline_y + 18, mode, tiny, mode_color, anchor="ra") + + lanes = [ + ("Audit", None), + ("Training audit", "Training audit"), + ("Rollout samples", "Rollout samples"), + ("Execution", None), + ("Operator / backend", "Operator / backend"), + ("Token comparison", "Token comparison"), + ("Validation", None), + ("Drift markers", "Drift markers"), + ] + left = 370 + right = width - 54 + ruler_y = timeline_y + 59 + chart_top = timeline_y + 91 + row_height = 62 + chart_bottom = chart_top + row_height * len(lanes) + max_end = max( + (_number(event.get("end"), default=1.0) or 1.0 for event in events), default=1.0 + ) + span = max(1.0, max_end) + + def x_for(value: float) -> int: + return int(left + max(0.0, min(span, value)) / span * (right - left)) + + # Group header, ruler, minor grid, and hierarchical track labels. + draw.rectangle((34, timeline_y + 42, width - 34, timeline_y + 80), fill="#f1f3f5") + write(52, timeline_y + 54, "TRACKS", tiny, muted) + write(left + 8, timeline_y + 54, "TIME / SAMPLE ORDINAL", tiny, muted) + draw.line((left, ruler_y + 16, right, ruler_y + 16), fill=line, width=1) + for tick in range(0, 33): + value = span * tick / 32 + xpos = x_for(value) + major = tick % 4 == 0 + draw.line( + (xpos, ruler_y + (7 if major else 13), xpos, chart_bottom), + fill=line if major else _REPORT_IMAGE_COLORS["grid"], + width=1, + ) + if major: + write(xpos, ruler_y - 3, _format_number(value), tiny, muted, anchor="ma") + + for lane_index, (lane, event_lane) in enumerate(lanes): + y0 = chart_top + lane_index * row_height + draw.rectangle( + (34, y0, right, y0 + row_height), + fill=panel_alt if lane_index % 2 else panel_bg, + ) + draw.line((34, y0 + row_height, right, y0 + row_height), fill=line, width=1) + if event_lane is None: + draw.rectangle((34, y0, left, y0 + row_height), fill="#e8ebee") + write(52, y0 + 20, lane.upper(), label_font, text_color) + else: + write(52, y0 + 22, "|--", tiny, muted) + write(92, y0 + 21, lane, tiny, text_color) + + event_colors = { + "pass": _REPORT_IMAGE_COLORS["green"], + "warning": _REPORT_IMAGE_COLORS["yellow"], + "failure": _REPORT_IMAGE_COLORS["red"], + "info": _REPORT_IMAGE_COLORS["blue"], + } + for event in events: + event_lane = str(event.get("lane", "Drift markers")) + lane_index = next( + (idx for idx, (_, name) in enumerate(lanes) if name == event_lane), + len(lanes) - 1, + ) + y0 = chart_top + lane_index * row_height + color = event_colors.get(str(event.get("status", "info")), event_colors["info"]) + start = _number(event.get("start"), default=0.0) or 0.0 + end = _number(event.get("end"), default=start + 0.5) or start + 0.5 + x0 = x_for(start) + x1 = max(x0 + 10, x_for(end)) + if event.get("kind") == "marker": + mid = x_for(start) + draw.line((mid, y0 + 6, mid, y0 + row_height - 6), fill=color, width=3) + draw.polygon( + [(mid - 7, y0 + 7), (mid + 7, y0 + 7), (mid, y0 + 17)], fill=color + ) + marker_label = fit(event.get("label", "marker"), 34) + label_x = mid + 14 if mid <= right - 250 else max(left + 10, mid - 240) + write(label_x, y0 + 22, marker_label, tiny, color) + else: + bar_y = y0 + 22 + draw.rectangle((x0, bar_y, x1, bar_y + 18), fill=color) + if x1 - x0 >= 130: + write( + x0 + 10, + bar_y + 2, + fit(event.get("label", "event"), 42), + tiny, + "#ffffff", + ) + else: + write( + min(x1 + 10, right - 260), + bar_y + 2, + fit(event.get("label", "event"), 34), + tiny, + text_color, + ) + + # Keep the comparison track meaningful even when the dump only contains a + # scalar worst-token summary instead of per-token samples. + worst = ( + normalized.get("worst_token") + if isinstance(normalized.get("worst_token"), Mapping) + else {} + ) + comparison_y = chart_top + 5 * row_height + 31 + draw.line( + (left + 18, comparison_y, right - 18, comparison_y), + fill=_REPORT_IMAGE_COLORS["track"], + width=3, + ) + write(left + 18, comparison_y - 22, "train vs rollout", tiny, muted) + if worst: + sample_position = _number(worst.get("sample_position"), default=0.0) or 0.0 + comparison_x = x_for(sample_position + 0.5) + draw.line( + (comparison_x, comparison_y - 15, comparison_x, comparison_y + 15), + fill=status_color, + width=3, + ) + write( + comparison_x + 10, + comparison_y - 10, + f"delta={_format_number(worst.get('abs_dlogp'))}", + tiny, + status_color, + ) + + # A compact legend makes the static image readable without hover state. + legend_y = timeline_y + timeline_h - 30 + for x0, label, color in ( + (52, "PASS", _REPORT_IMAGE_COLORS["green"]), + (130, "WARN", _REPORT_IMAGE_COLORS["yellow"]), + (214, "FAIL", _REPORT_IMAGE_COLORS["red"]), + (298, "INFO", _REPORT_IMAGE_COLORS["blue"]), + ): + draw.rectangle((x0, legend_y + 3, x0 + 14, legend_y + 15), fill=color) + write(x0 + 22, legend_y, label, tiny, muted) + write( + right, + legend_y, + "positions are sample ordinals when timestamps are absent", + tiny, + muted, + anchor="ra", + ) + + # Diagnostic summary panels. + summary_y = 850 + summary_h = 220 + half_gap = 18 + half_width = (width - 68 - half_gap) // 2 + rect((34, summary_y, 34 + half_width, summary_y + summary_h), panel_bg, 2, line) + rect( + (34 + half_width + half_gap, summary_y, width - 34, summary_y + summary_h), + panel_bg, + 2, + line, + ) + write(54, summary_y + 18, "DRIFT SUMMARY", section_font) + write(54, summary_y + 57, "observed maximum", tiny, muted) + max_value = _number(metrics.get("max_abs_dlogp"), default=0.0) or 0.0 + bar_x = 54 + bar_y = summary_y + 91 + bar_width = half_width - 110 + rect( + (bar_x, bar_y, bar_x + bar_width, bar_y + 24), _REPORT_IMAGE_COLORS["track"], 4 + ) + fill_width = ( + int(min(1.0, max_value / max(max_value, 1.0)) * bar_width) if max_value else 0 + ) + if fill_width: + rect((bar_x, bar_y, bar_x + max(8, fill_width), bar_y + 24), status_color, 4) + write(bar_x, bar_y + 30, "0", tiny, muted) + write( + bar_x + bar_width, + bar_y + 30, + _format_number(max(max_value, 1.0)), + tiny, + muted, + anchor="ra", + ) + worst = ( + normalized.get("worst_token") + if isinstance(normalized.get("worst_token"), Mapping) + else {} + ) + worst_text = ( + f"worst token: sample={worst.get('sample_position', '-')} token={worst.get('token_position', '-')} " + f"|dlogp|={_format_number(worst.get('abs_dlogp'))}" + ) + write(54, summary_y + 157, fit(worst_text, 92), mono, text_color) + right_x = 34 + half_width + half_gap + 20 + write(right_x, summary_y + 18, "SELECTED EVENT", section_font) + warning_items = ( + normalized.get("validation") + if isinstance(normalized.get("validation"), Mapping) + else {} + ) + warnings = warning_items.get("warnings") or [] + failures = warning_items.get("failures") or [] + issue_text = "No anomaly recorded." + if failures or warnings: + first = (failures or warnings)[0] + first = first if isinstance(first, Mapping) else {"message": first} + issue_text = ( + f"{str(first.get('code', 'validation'))}: {str(first.get('message', ''))}" + ) + elif worst: + issue_text = ( + f"worst token: sample={worst.get('sample_position', '-')} " + f"token={worst.get('token_position', '-')} " + f"|dlogp|={_format_number(worst.get('abs_dlogp'))}" + ) + write(right_x, summary_y + 57, fit(issue_text, 95), small, status_color) + write(right_x, summary_y + 101, "worst token / validation marker", tiny, muted) + write( + right_x, + summary_y + 140, + f"status={badge_text} warnings={len(warnings)} failures={len(failures)}", + mono, + text_color, + ) + + # Bottom tables: axis capture and runtime provenance are the actionable part + # of the image for a post-training user. + table_y = 1095 + table_h = 400 + rect((34, table_y, 34 + half_width, table_y + table_h), panel_bg, 2, line) + rect( + (34 + half_width + half_gap, table_y, width - 34, table_y + table_h), + panel_bg, + 2, + line, + ) + write(54, table_y + 18, "CAPTURED EXECUTION AXES", section_font) + write( + 34 + half_width + half_gap + 20, + table_y + 18, + "RUNTIME PROVENANCE", + section_font, + ) + axes_rows = key_value_rows(normalized.get("axes")) + provenance_rows = key_value_rows(normalized.get("runtime_provenance")) + + def table( + rows: list[tuple[str, str]], x0: int, y0: int, w: int, max_rows: int = 8 + ) -> None: + row_y = y0 + for index, (key, value) in enumerate(rows[:max_rows]): + if index % 2 == 0: + draw.rectangle((x0, row_y - 3, x0 + w, row_y + 35), fill=panel_alt) + write(x0 + 14, row_y + 8, key, tiny, muted) + write(x0 + 270, row_y + 8, value, mono, text_color) + draw.line((x0, row_y + 38, x0 + w, row_y + 38), fill=line, width=1) + row_y += 39 + if not rows: + write(x0 + 14, row_y + 8, "not recorded", small, muted) + + table(axes_rows, 54, table_y + 66, half_width - 40) + table( + provenance_rows, 34 + half_width + half_gap + 20, table_y + 66, half_width - 40 + ) + validation_label = ( + "VALIDATION: PASS" if status == "pass" else f"VALIDATION: {badge_text}" + ) + if failures or warnings: + validation_label = f"VALIDATION: {badge_text} | {fit(issue_text, 86)}" + write(70, height - 58, validation_label, tiny, status_color) + footer = "static diagnostic image | schema v" + str( + normalized.get("schema_version", REPORT_SCHEMA_VERSION) + ) + write(width - 54, height - 28, footer, tiny, muted, anchor="ra") + return image + + +def write_drift_report_image(report: Mapping[str, Any], path: str | Path) -> Path: + """Write a static PNG/JPEG consistency drift report image.""" + + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + image = render_drift_report_image(report) + suffix = output.suffix.lower() + image_format = "JPEG" if suffix in {".jpg", ".jpeg"} else "PNG" + if image_format == "JPEG": + image.save(output, format=image_format, quality=95, optimize=True) + else: + image.save(output, format=image_format, optimize=True) + return output + + +def _read_attempt_json(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"invalid cross-config report input: {path}") from exc + if not isinstance(value, Mapping): + raise ValueError(f"cross-config report input must be a JSON object: {path}") + return _plain_mapping(value) + + +def _worst_token( + absolute_diff: torch.Tensor, + active_mask: torch.Tensor, +) -> dict[str, Any] | None: + if absolute_diff.shape != active_mask.shape: + raise ValueError( + "token_diffs absolute_diff and active_mask must have matching shapes" + ) + if not bool(active_mask.any()): + return None + values = absolute_diff.masked_fill(~active_mask, -1.0) + flat_index = int(values.reshape(-1).argmax().item()) + coordinates = list(torch.unravel_index(torch.tensor(flat_index), values.shape)) + sample_position = int(coordinates[0].item()) if coordinates else 0 + token_position = int(coordinates[-1].item()) if coordinates else flat_index + return { + "sample_position": sample_position, + "token_position": token_position, + "abs_dlogp": float(values.reshape(-1)[flat_index].item()), + } + + +def _samples_from_identity( + identity: Mapping[str, Any], + active_mask: torch.Tensor, +) -> list[dict[str, Any]]: + logical_identity = _plain_mapping(identity.get("identity")) + token_ids = logical_identity.get("token_ids") + samples: list[dict[str, Any]] = [] + for index in range(active_mask.shape[0] if active_mask.ndim else 0): + row = ( + token_ids[index] + if isinstance(token_ids, list) and index < len(token_ids) + else None + ) + samples.append( + { + "sample_position": index, + "sample_index": index, + "token_count": ( + len(row) if isinstance(row, list) else int(active_mask.shape[-1]) + ), + "active_token_count": int(active_mask[index].sum().item()), + } + ) + return samples + + +def _render_key_value_table(values: Any, *, empty: str) -> str: + if not isinstance(values, Mapping) or not values: + return f'
{html.escape(empty)}
' + rows = [] + for key in sorted(values, key=str): + rows.append( + f"{html.escape(str(key))}{html.escape(_format_value(values[key]))}" + ) + return "" + "".join(rows) + "
" + + +def _render_validation(validation: Any) -> str: + if not isinstance(validation, Mapping): + return '
No validation record was captured.
' + warnings = validation.get("warnings") or [] + failures = validation.get("failures") or [] + if not warnings and not failures: + return ( + '
PASS: no metadata validation warnings or failures.
' + ) + rows = [] + for kind, items in (("failure", failures), ("warning", warnings)): + for item in items: + item = item if isinstance(item, Mapping) else {"message": item} + rows.append( + f'{kind.upper()}{html.escape(str(item.get("code", "")))}{html.escape(str(item.get("message", "")))}' + ) + return ( + "" + + "".join(rows) + + "
StatusCodeMessage
" + ) + + +def _plain_mapping(value: Any) -> dict[str, Any]: + return ( + {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, Mapping) + else {} + ) + + +def _json_safe(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, Mapping): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + return str(value) + + +def _first_present(*values: Any) -> Any: + for value in values: + if value not in (None, ""): + return value + return None + + +def _backend_label(rollout: Any, training: Any) -> str | None: + """Show one actual backend only when both execution sides agree.""" + + if rollout in (None, "") and training in (None, ""): + return None + if rollout in (None, ""): + return f"training={training}" + if training in (None, ""): + return f"rollout={rollout}" + if rollout == training: + return str(rollout) + return f"rollout={rollout}; training={training}" + + +def _number(value: Any, default: float | None = None) -> float | None: + if value is None or isinstance(value, bool): + return default + try: + return float(value) + except (TypeError, ValueError): + return default + + +def _format_number(value: Any, default: str = "-") -> str: + number = _number(value) + if number is None: + return default + if abs(number) >= 1000 or (abs(number) < 0.001 and number != 0): + return f"{number:.3e}" + return f"{number:.6f}".rstrip("0").rstrip(".") + + +def _format_value(value: Any) -> str: + if isinstance(value, (dict, list, tuple)): + return json.dumps(value, ensure_ascii=True, sort_keys=True) + return str(value) + + +def _truncate(value: str, limit: int) -> str: + return value if len(value) <= limit else value[: max(1, limit - 3)] + "..." + + +def _load_report_font(size: int, *, bold: bool = False, mono: bool = False) -> Any: + """Load a platform font with deterministic fallbacks for report images.""" + + from PIL import ImageFont + + candidates = [] + if mono: + candidates.extend( + [ + "C:/Windows/Fonts/consola.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf", + ] + ) + elif bold: + candidates.extend( + [ + "C:/Windows/Fonts/segoeuib.ttf", + "C:/Windows/Fonts/arialbd.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", + ] + ) + else: + candidates.extend( + [ + "C:/Windows/Fonts/segoeui.ttf", + "C:/Windows/Fonts/arial.ttf", + "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", + ] + ) + for candidate in candidates: + try: + return ImageFont.truetype(candidate, size=size) + except OSError: + continue + return ImageFont.load_default() diff --git a/rl_engine/alignment/cross_config/drift_viewer.py b/rl_engine/alignment/cross_config/drift_viewer.py new file mode 100644 index 00000000..75bf7580 --- /dev/null +++ b/rl_engine/alignment/cross_config/drift_viewer.py @@ -0,0 +1,316 @@ +"""Optional offline Qt viewer for ``.rlk-drift`` cross-config bundles.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +from rl_engine.alignment.cross_config.drift_report import load_drift_bundle + + +_LANES = [ + ("Audit", None), + ("Training audit", "Training audit"), + ("Rollout samples", "Rollout samples"), + ("Execution", None), + ("Operator / backend", "Operator / backend"), + ("Token comparison", "Token comparison"), + ("Validation", None), + ("Drift markers", "Drift markers"), +] +_COLORS = { + "pass": "#54a64f", + "warning": "#e0a000", + "failure": "#d53f3f", + "info": "#72a6c8", +} + + +def _load_qt() -> tuple[Any, Any, Any]: + try: + from PySide6 import QtCore, QtGui, QtWidgets + except ImportError as exc: # pragma: no cover - depends on workstation extras + raise RuntimeError( + "The offline viewer requires the optional GUI dependency. " + "Install it with: pip install 'rl-engine[drift-viewer]'" + ) from exc + return QtCore, QtGui, QtWidgets + + +def _format_event(event: dict[str, Any]) -> str: + details = event.get("details") if isinstance(event.get("details"), dict) else {} + payload = { + "id": event.get("id"), + "lane": event.get("lane"), + "label": event.get("label"), + "status": event.get("status"), + "start": event.get("start"), + "end": event.get("end"), + "details": details, + } + return json.dumps(payload, ensure_ascii=True, indent=2, sort_keys=True) + + +def _run_qt(bundle_path: Path) -> int: # pragma: no cover - GUI path + QtCore, QtGui, QtWidgets = _load_qt() + bundle = load_drift_bundle(bundle_path) + report = bundle["report"] + events = [event for event in report.get("events", []) if isinstance(event, dict)] + title = str(report.get("title", "RL-Kernel cross-config drift report")) + status = str(report.get("status", "info")) + status_color = _COLORS.get(status, _COLORS["info"]) + + class EventItem(QtWidgets.QGraphicsRectItem): + def __init__(self, rect: Any, event: dict[str, Any], color: str) -> None: + super().__init__(rect) + self.event = event + self.setBrush(QtGui.QColor(color)) + self.setPen(QtGui.QPen(QtGui.QColor(color))) + self.setFlag( + QtWidgets.QGraphicsItem.GraphicsItemFlag.ItemIsSelectable, True + ) + self.setToolTip(str(event.get("label", event.get("id", "event")))) + + class TimelineView(QtWidgets.QGraphicsView): + def __init__(self, scene: Any) -> None: + super().__init__(scene) + self.setRenderHint(QtGui.QPainter.RenderHint.Antialiasing, False) + self.setDragMode(QtWidgets.QGraphicsView.DragMode.ScrollHandDrag) + self.setTransformationAnchor( + QtWidgets.QGraphicsView.ViewportAnchor.AnchorUnderMouse + ) + self.setResizeAnchor( + QtWidgets.QGraphicsView.ViewportAnchor.AnchorViewCenter + ) + self.setHorizontalScrollBarPolicy( + QtCore.Qt.ScrollBarPolicy.ScrollBarAlwaysOn + ) + self.setVerticalScrollBarPolicy( + QtCore.Qt.ScrollBarPolicy.ScrollBarAlwaysOff + ) + self.setStyleSheet("QGraphicsView { border: 0; background: #ffffff; }") + + def wheelEvent(self, event: Any) -> None: + factor = 1.18 if event.angleDelta().y() > 0 else 1.0 / 1.18 + self.scale(factor, 1.0) + event.accept() + + class Window(QtWidgets.QMainWindow): + def __init__(self) -> None: + super().__init__() + self.setWindowTitle(f"RL-Kernel Cross-Config Drift - {title}") + self.resize(1500, 900) + self._scene_width = 2200.0 + self._row_height = 52.0 + self._left_label_width = 8.0 + self._span = max( + 1.0, + max( + (float(event.get("end", 1.0) or 1.0) for event in events), + default=1.0, + ), + ) + self._scene = QtWidgets.QGraphicsScene(self) + self._scene.setBackgroundBrush(QtGui.QColor("#ffffff")) + self._view = TimelineView(self._scene) + self._tree = QtWidgets.QTreeWidget() + self._tree.setHeaderLabel("Tracks") + self._tree.setMinimumWidth(245) + self._tree.setStyleSheet( + "QTreeWidget { background: #f4f5f6; border: 0; " + "font-family: 'Segoe UI'; font-size: 10pt; }" + ) + self._details = QtWidgets.QPlainTextEdit() + self._details.setReadOnly(True) + self._details.setPlaceholderText( + "Select an event to inspect its audit details." + ) + self._details.setStyleSheet( + "QPlainTextEdit { background: #ffffff; border-top: 1px solid #c9ced3; " + "font-family: Consolas; font-size: 10pt; padding: 8px; }" + ) + self._build_tree() + self._build_scene() + self._scene.selectionChanged.connect(self._show_selected) + self._tree.itemSelectionChanged.connect(self._select_lane) + self._build_layout() + + def _build_tree(self) -> None: + groups: dict[str, Any] = {} + for label, event_lane in _LANES: + if event_lane is None: + item = QtWidgets.QTreeWidgetItem([label]) + item.setFlags(item.flags() & ~QtCore.Qt.ItemFlag.ItemIsSelectable) + self._tree.addTopLevelItem(item) + groups[label] = item + continue + parent = groups.get( + "Audit" + if event_lane in {"Training audit", "Rollout samples"} + else "Execution" + ) + if event_lane == "Drift markers": + parent = groups.get("Validation") + item = QtWidgets.QTreeWidgetItem([label]) + item.setData(0, QtCore.Qt.ItemDataRole.UserRole, event_lane) + if parent is not None: + parent.addChild(item) + else: + self._tree.addTopLevelItem(item) + self._tree.expandAll() + + def _x_for(self, value: float) -> float: + return ( + 120.0 + + max(0.0, min(self._span, value)) / self._span * self._scene_width + ) + + def _build_scene(self) -> None: + font = QtGui.QFont("Segoe UI", 9) + muted = QtGui.QColor("#69737c") + line = QtGui.QColor("#d6dade") + grid = QtGui.QColor("#e8eaec") + for tick in range(0, 17): + x = self._x_for(self._span * tick / 16.0) + self._scene.addLine( + x, + 0, + x, + len(_LANES) * self._row_height, + QtGui.QPen(line if tick % 4 == 0 else grid), + ) + label = self._scene.addText(f"{self._span * tick / 16.0:g}", font) + label.setDefaultTextColor(muted) + label.setPos(x + 3, -24) + + lane_index = { + name: index + for index, (_, name) in enumerate(_LANES) + if name is not None + } + for index, (label, event_lane) in enumerate(_LANES): + y = index * self._row_height + self._scene.addLine( + 0, + y + self._row_height, + self._scene_width + 140, + y + self._row_height, + QtGui.QPen(line), + ) + if event_lane is None: + block = self._scene.addRect( + 0, + y, + self._scene_width + 140, + self._row_height, + QtGui.QPen(), + QtGui.QBrush(QtGui.QColor("#eef0f2")), + ) + block.setZValue(-2) + else: + text = self._scene.addText(label, font) + text.setDefaultTextColor(QtGui.QColor("#252a2f")) + text.setPos(12, y + 17) + + for event in events: + event_lane = str(event.get("lane", "Drift markers")) + if event_lane not in lane_index: + continue + y = lane_index[event_lane] * self._row_height + 16 + start = float(event.get("start", 0.0) or 0.0) + end = float(event.get("end", start) or start) + x0 = self._x_for(start) + x1 = max(x0 + 10.0, self._x_for(end)) + color = _COLORS.get(str(event.get("status", "info")), _COLORS["info"]) + if event.get("kind") == "marker": + item = EventItem(QtCore.QRectF(x0 - 5, y + 5, 10, 10), event, color) + else: + item = EventItem(QtCore.QRectF(x0, y, x1 - x0, 18), event, color) + self._scene.addItem(item) + if x1 - x0 > 130 and event.get("kind") != "marker": + text = self._scene.addText(str(event.get("label", "event")), font) + text.setDefaultTextColor(QtGui.QColor("#ffffff")) + text.setPos(x0 + 8, y - 1) + text.setZValue(1) + self._scene.setSceneRect( + 0, -28, self._scene_width + 140, len(_LANES) * self._row_height + 40 + ) + + def _build_layout(self) -> None: + toolbar = QtWidgets.QToolBar() + toolbar.setMovable(False) + toolbar.setStyleSheet( + "QToolBar { background: #ffffff; border-bottom: 1px solid #c9ced3; }" + ) + status_label = QtWidgets.QLabel(f" {title} STATUS: {status.upper()} ") + status_label.setStyleSheet( + f"color: {status_color}; font-weight: 600; padding: 5px;" + ) + toolbar.addWidget(status_label) + fit_button = QtWidgets.QAction("Fit", self) + fit_button.triggered.connect( + lambda: self._view.fitInView( + self._scene.sceneRect(), QtCore.Qt.AspectRatioMode.KeepAspectRatio + ) + ) + toolbar.addAction(fit_button) + toolbar.addWidget(QtWidgets.QLabel(f" {report.get('timeline_note', '')}")) + self.addToolBar(toolbar) + + right = QtWidgets.QSplitter(QtCore.Qt.Orientation.Vertical) + right.addWidget(self._view) + right.addWidget(self._details) + right.setStretchFactor(0, 4) + right.setStretchFactor(1, 1) + root = QtWidgets.QSplitter(QtCore.Qt.Orientation.Horizontal) + root.addWidget(self._tree) + root.addWidget(right) + root.setStretchFactor(1, 1) + self.setCentralWidget(root) + + def _select_lane(self) -> None: + items = self._tree.selectedItems() + if not items: + return + lane = items[0].data(0, QtCore.Qt.ItemDataRole.UserRole) + if not lane: + return + for item in self._scene.items(): + if isinstance(item, EventItem): + item.setOpacity(1.0 if item.event.get("lane") == lane else 0.28) + + def _show_selected(self) -> None: + selected = self._scene.selectedItems() + self._details.setPlainText( + _format_event(selected[0].event) if selected else "" + ) + + app = QtWidgets.QApplication.instance() or QtWidgets.QApplication(sys.argv) + window = Window() + window.show() + window._view.fitInView( + window._scene.sceneRect(), QtCore.Qt.AspectRatioMode.KeepAspectRatio + ) + return app.exec() + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + description="Open an RL-Kernel drift bundle in the offline desktop viewer." + ) + parser.add_argument("bundle", type=Path, help=".rlk-drift bundle") + args = parser.parse_args(argv) + if args.bundle.suffix.lower() != ".rlk-drift": + parser.error("the offline viewer expects a .rlk-drift bundle") + try: + return _run_qt(args.bundle) + except RuntimeError as exc: + parser.error(str(exc)) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/setup.py b/setup.py index b31a7edb..79f882d9 100644 --- a/setup.py +++ b/setup.py @@ -1,307 +1,313 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -import importlib.util -import os -import warnings -from pathlib import Path - -from setuptools import find_packages, setup - - -def _load_envs_module(): - envs_path = Path(__file__).with_name("envs.py") - spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"failed to load environment helpers from {envs_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -envs = _load_envs_module() - - -def _load_torch_extension_tools(): - try: - import torch - except ModuleNotFoundError as exc: - if exc.name != "torch": - raise - return None, None, None - - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - - # CUDAExtension is also the supported extension entry point for ROCm - # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when - # torch.version.hip is set. - return torch, BuildExtension, CUDAExtension - - -def _native_extension_required() -> bool: - """Whether the caller explicitly requested a native extension build.""" - return ( - envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) - or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) - or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) - or envs.env_flag("FORCE_CUDA") - ) - - -def _cuda_define_from_env(name: str, macro: str) -> list[str]: - value = os.environ.get(name) - if value is None: - return [] - parsed = int(value) - if parsed <= 0: - raise ValueError(f"{name} must be positive, got {value!r}") - return [f"-D{macro}={parsed}"] - - -_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( - "-Xfatbin", - "-compress-all", - "-gencode", - "--generate-code", - "--expt-", - "-lineinfo", - "-allow-unsupported-compiler", - "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", -) -_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { - "-Xfatbin", - "-gencode", - "--generate-code", -} - - -def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: - """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" - filtered_flags = [] - skip_next = False - for flag in flags: - if skip_next: - skip_next = False - continue - if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: - skip_next = True - continue - if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): - continue - filtered_flags.append(flag) - return filtered_flags - - -def get_extensions(): - torch, _, CUDAExtension = _load_torch_extension_tools() - if torch is None: - message = ( - "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " - "CUDA/ROCm PyTorch build first, then run " - "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." - ) - if _native_extension_required(): - raise RuntimeError(message) - warnings.warn( - f"{message} Continuing with the pure-Python fallback because no native extension " - "was explicitly requested.", - RuntimeWarning, - stacklevel=2, - ) - return [] - - extensions = [] - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] - if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": - torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = getattr(torch.version, "hip", None) is not None - - # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, - # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also - # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add - # --offload-arch. Do not require a visible GPU when a ROCm target was - # explicitly selected. - no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() - if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: - raise RuntimeError( - "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " - "Set one or more ';'-separated targets, for example " - "PYTORCH_ROCM_ARCH='gfx942;gfx950'." - ) - - if is_rocm or torch.cuda.is_available(): - cuda_sources = [ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cu", - "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/gemm/det_gemm_kernel.cu", - "csrc/cuda/rmsnorm.cu", - "csrc/cuda/activation.cu", - "csrc/cuda/attention/deterministic_attention.cu", - "csrc/cuda/distributed/deterministic_collective.cu", - ] - if not is_rocm: - # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). - # The ROCm dispatcher falls back to PyTorch SDPA for this operator. - cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") - - nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") - if not is_rocm: - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - ) - ) - if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): - nvcc_flags.append("-lineinfo") - if ( - not is_rocm - and os.name == "nt" - and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) - ): - nvcc_flags.append("-allow-unsupported-compiler") - nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") - - cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] - extra_link_args = list(torch_rpath) - if os.name != "nt": - # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). - extra_link_args.append("-lcuda") - - if not is_rocm: - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - # Single-card batch-invariant embedding/lm-head. - "csrc/cuda/embedding_lm_head_sm90.cu", - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - if enable_sm90 and present_sm90: - tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant - cuda_sources.extend(present_sm90) +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import importlib.util +import os +import warnings +from pathlib import Path + +from setuptools import find_packages, setup + + +def _load_envs_module(): + envs_path = Path(__file__).with_name("envs.py") + spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load environment helpers from {envs_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +envs = _load_envs_module() + + +def _load_torch_extension_tools(): + try: + import torch + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + return None, None, None + + from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + # CUDAExtension is also the supported extension entry point for ROCm + # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when + # torch.version.hip is set. + return torch, BuildExtension, CUDAExtension + + +def _native_extension_required() -> bool: + """Whether the caller explicitly requested a native extension build.""" + return ( + envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) + or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or envs.env_flag("FORCE_CUDA") + ) + + +def _cuda_define_from_env(name: str, macro: str) -> list[str]: + value = os.environ.get(name) + if value is None: + return [] + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive, got {value!r}") + return [f"-D{macro}={parsed}"] + + +_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( + "-Xfatbin", + "-compress-all", + "-gencode", + "--generate-code", + "--expt-", + "-lineinfo", + "-allow-unsupported-compiler", + "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", +) +_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { + "-Xfatbin", + "-gencode", + "--generate-code", +} + + +def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: + """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" + filtered_flags = [] + skip_next = False + for flag in flags: + if skip_next: + skip_next = False + continue + if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: + skip_next = True + continue + if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): + continue + filtered_flags.append(flag) + return filtered_flags + + +def get_extensions(): + torch, _, CUDAExtension = _load_torch_extension_tools() + if torch is None: + message = ( + "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " + "CUDA/ROCm PyTorch build first, then run " + "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." + ) + if _native_extension_required(): + raise RuntimeError(message) + warnings.warn( + f"{message} Continuing with the pure-Python fallback because no native extension " + "was explicitly requested.", + RuntimeWarning, + stacklevel=2, + ) + return [] + + extensions = [] + torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") + torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] + if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": + torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") + is_rocm = getattr(torch.version, "hip", None) is not None + + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, + # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also + # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add + # --offload-arch. Do not require a visible GPU when a ROCm target was + # explicitly selected. + no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() + if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: + raise RuntimeError( + "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " + "Set one or more ';'-separated targets, for example " + "PYTORCH_ROCM_ARCH='gfx942;gfx950'." + ) + + if is_rocm or torch.cuda.is_available(): + cuda_sources = [ + "csrc/ops.cpp", + "csrc/fused_logp_kernel.cu", + "csrc/deterministic_logp_kernel.cu", + "csrc/cuda/gemm/det_gemm_kernel.cu", + "csrc/cuda/rmsnorm.cu", + "csrc/cuda/activation.cu", + "csrc/cuda/attention/deterministic_attention.cu", + "csrc/cuda/distributed/deterministic_collective.cu", + ] + if not is_rocm: + # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). + # The ROCm dispatcher falls back to PyTorch SDPA for this operator. + cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") + + nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] + if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + nvcc_flags.append("--use_fast_math") + if not is_rocm: + cc_major, cc_minor = torch.cuda.get_device_capability() + enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + if not enable_sm90: + # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. + nvcc_flags.append( + f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" + ) + nvcc_flags.append("--expt-relaxed-constexpr") + nvcc_flags.append("--expt-extended-lambda") + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + ) + ) + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + nvcc_flags.append("-lineinfo") + if ( + not is_rocm + and os.name == "nt" + and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) + ): + nvcc_flags.append("-allow-unsupported-compiler") + nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") + + cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] + extra_link_args = list(torch_rpath) + if os.name != "nt": + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") + + if not is_rocm: + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + # Single-card batch-invariant embedding/lm-head. + "csrc/cuda/embedding_lm_head_sm90.cu", + ] + enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + if enable_sm90 and present_sm90: + tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant + cuda_sources.extend(present_sm90) nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") - - if is_rocm: - nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) - - extensions.append( - CUDAExtension( - name="rl_engine._C", - sources=cuda_sources, - include_dirs=[], - extra_compile_args={ - "cxx": cxx_flags, - "nvcc": nvcc_flags, - }, - extra_link_args=extra_link_args, - ) - ) - - if _native_extension_required() and not extensions: - raise RuntimeError( - "rl_engine._C was requested but no CUDA/ROCm build environment is available. " - "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " - "PYTORCH_ROCM_ARCH to the target architecture." - ) - - return extensions - - -def get_cmdclass(): - _, BuildExtension, _ = _load_torch_extension_tools() - if BuildExtension is None: - return {} - return {"build_ext": BuildExtension} - - -setup( - name="rl-engine", - version="0.1.0", - packages=find_packages(include=["rl_engine", "rl_engine.*"]), - install_requires=[ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", - ], - ext_modules=get_extensions(), - cmdclass=get_cmdclass(), - extras_require={ - "cuda": ["flashinfer"], - "rocm": ["aiter"], - "vllm": ["vllm>=0.6.0"], - }, - python_requires=">=3.10", - include_package_data=True, - zip_safe=False, -) + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp + # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in + # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. + enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" + if enable_det_gemm_sm90: + tma_arch = f"{cc_major}{cc_minor}a" + arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" + if arch_flag not in nvcc_flags: + nvcc_flags.append(arch_flag) + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + if is_rocm: + nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) + + extensions.append( + CUDAExtension( + name="rl_engine._C", + sources=cuda_sources, + include_dirs=[], + extra_compile_args={ + "cxx": cxx_flags, + "nvcc": nvcc_flags, + }, + extra_link_args=extra_link_args, + ) + ) + + if _native_extension_required() and not extensions: + raise RuntimeError( + "rl_engine._C was requested but no CUDA/ROCm build environment is available. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " + "PYTORCH_ROCM_ARCH to the target architecture." + ) + + return extensions + + +def get_cmdclass(): + _, BuildExtension, _ = _load_torch_extension_tools() + if BuildExtension is None: + return {} + return {"build_ext": BuildExtension} + + +setup( + name="rl-engine", + version="0.1.0", + packages=find_packages(include=["rl_engine", "rl_engine.*"]), + install_requires=[ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", + ], + ext_modules=get_extensions(), + cmdclass=get_cmdclass(), + extras_require={ + "cuda": ["flashinfer"], + "rocm": ["aiter"], + "vllm": ["vllm>=0.6.0"], + "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], + }, + entry_points={ + "console_scripts": [ + "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", + ], + }, + python_requires=">=3.10", + include_package_data=True, + zip_safe=False, +) diff --git a/tests/test_cross_config_drift_report.py b/tests/test_cross_config_drift_report.py new file mode 100644 index 00000000..bb0fe3d7 --- /dev/null +++ b/tests/test_cross_config_drift_report.py @@ -0,0 +1,344 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest +import torch + +import rl_engine.alignment.cross_config.__main__ as cli_main +from rl_engine.alignment.cross_config.drift_report import ( + build_cross_config_attempt_report, + build_drift_report, + build_drift_trace, + load_drift_bundle, + render_drift_report, + render_drift_report_image, + write_drift_bundle, + write_drift_report, + write_drift_report_image, + write_drift_trace, +) +from rl_engine.alignment.cross_config.artifacts import ArtifactStore + + +def _artifacts(*, with_timestamp: bool = False): + sample = { + "sample_position": 0, + "sample_index": 17, + "rollout_id": 23, + "batch_layout_fingerprint": "layout-1", + } + if with_timestamp: + sample.update(start_ts=10.0, end_ts=10.5) + manifest = { + "mode": "audit", + "samples": [sample], + "batch_invariance_cases": [{"case": "same_sample_alone"}], + "validation": {"warnings": [], "failures": []}, + "runtime_provenance": {"operator": "linear_logp"}, + } + cube = { + "mode": "audit", + "rank": 0, + "axes": {"dtype": "bf16", "cp": 1, "logp_backend": "rlk.linear_logp.fast"}, + "metrics": { + "active_token_count": 2, + "max_abs_dlogp": 0.125, + "warning_count": 1, + "metadata_warning_count": 0, + "metadata_failure_count": 0, + }, + "worst_token": {"abs_dlogp": 0.125, "sample_position": 0, "token_position": 4}, + "metadata_validation": {"warnings": [], "failures": []}, + "runtime_provenance": { + "actual_backend": "rl_engine.linear_logp", + "fallback": False, + }, + } + return manifest, cube + + +def _completed_attempt(root: Path) -> Path: + store = ArtifactStore(root) + attempt = store.create_attempt("drift-report", "case-1") + envelope = {"case_id": "case-1", "attempt_id": attempt.name} + requested = { + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "dtype": "bfloat16", + }, + "training": {"compute_dtype": "bfloat16", "sharding": "tp2-cp2"}, + } + store.write_json( + attempt, + "requested", + { + **envelope, + "schema_version": "cross_config.requested.v1", + "case": {"requested": requested}, + }, + ) + store.write_json( + attempt, + "materialized", + { + **envelope, + "schema_version": "cross_config.materialized_envelope.v1", + "materialized_case": {"case": {"requested": requested}}, + }, + ) + store.write_json( + attempt, + "identity", + { + **envelope, + "schema_version": "cross_config.identity_envelope.v1", + "identity": {"token_ids": [[11, 12], [21, 22]]}, + }, + ) + store.write_json( + attempt, + "actual", + { + **envelope, + "schema_version": "cross_config.actual.v1", + "operator_source": "exact_resolution_and_instance", + "execution_fingerprint": "execution-sha", + "environment_fingerprint": "environment-sha", + "rollout": {"backend_id": "rlkernel.attention.deterministic.v1"}, + "training": {"backend_id": "rlkernel.ffn.qwen3.deterministic.v1"}, + }, + ) + store.write_json( + attempt, + "comparison", + { + **envelope, + "schema_version": "cross_config.alignment_result.v1", + "status": "pass", + "comparable": True, + "passed": True, + "mismatch_count": 0, + "fixed_threshold": 0.0, + "contract_fingerprint": "contract-sha", + "diagnostics": {}, + }, + ) + scores = { + "selected_logprobs": torch.zeros((2, 2)), + "active_mask": torch.tensor([[False, True], [False, True]]), + } + + def write_tensor_bundle(name: str, tensors: dict[str, torch.Tensor]) -> None: + torch.save( + {"schema_version": 1, "tensors": tensors, "metadata": envelope}, + attempt / f"{name}.pt", + ) + + for name in ("score_rollout", "score_training"): + write_tensor_bundle(name, scores) + write_tensor_bundle( + "token_diffs", + { + "rollout_logprobs": torch.zeros((2, 2)), + "training_logprobs": torch.zeros((2, 2)), + "active_mask": scores["active_mask"], + "absolute_diff": torch.zeros((2, 2)), + "mismatch_mask": torch.zeros((2, 2), dtype=torch.bool), + }, + ) + store.complete_attempt( + attempt, + summary={ + **envelope, + "schema_version": "cross_config.complete.v1", + "status": "pass", + }, + ) + return attempt + + +@pytest.mark.unit +def test_report_uses_ordinal_timeline_without_fabricating_timestamps(): + manifest, cube = _artifacts() + + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + + assert report["timeline_mode"] == "ordinal_diagnostic" + assert "not elapsed time" in report["timeline_note"] + assert {event["lane"] for event in report["events"]} == { + "Training audit", + "Rollout samples", + "Operator / backend", + "Drift markers", + } + assert report["status"] == "warning" + + +@pytest.mark.unit +def test_report_prefers_actual_backend_and_timestamp_mode(): + manifest, cube = _artifacts(with_timestamp=True) + cube["runtime_provenance"] = { + "requested_backend": "registry", + "actual_backend": "native.linear_logp", + "fallback": True, + } + + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + operator = next( + event for event in report["events"] if event["id"] == "operator-backend" + ) + + assert report["timeline_mode"] == "timestamp" + assert operator["label"] == "native.linear_logp" + assert operator["status"] == "warning" + + +@pytest.mark.unit +def test_rendered_report_is_self_contained_and_escapes_details(tmp_path: Path): + manifest, cube = _artifacts() + manifest["validation"]["warnings"] = [ + {"code": "bad<&", "message": "value "} + ] + cube["metadata_validation"] = manifest["validation"] + + report = build_drift_report( + replay_manifest=manifest, + result_cube=cube, + title="", + ) + html = render_drift_report(report) + output = write_drift_report(report, tmp_path / "drift.html") + + assert output.exists() + assert "" not in html + assert "value " not in html + assert "ordinal_diagnostic" in html + assert "Operator / backend" in html + assert "http://" not in html + assert "https://" not in html + assert "detail.innerHTML" not in html + assert "detail.replaceChildren" in html + + +@pytest.mark.unit +def test_static_report_image_is_shareable_png(tmp_path: Path): + pytest.importorskip("PIL") + manifest, cube = _artifacts() + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + + image = render_drift_report_image(report) + assert image.size == (2400, 1680) + + output = write_drift_report_image(report, tmp_path / "drift.png") + assert output.read_bytes()[:8] == b"\x89PNG\r\n\x1a\n" + + +@pytest.mark.unit +def test_consistency_trace_is_expandable_chrome_trace_json(tmp_path: Path): + manifest, cube = _artifacts() + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + + trace = build_drift_trace(report) + assert trace["metadata"]["timeline_mode"] == "ordinal_diagnostic" + assert any( + event.get("ph") == "M" and event.get("name") == "thread_name" + for event in trace["traceEvents"] + ) + assert any( + event.get("ph") == "X" and event.get("cat") == "consistency.audit" + for event in trace["traceEvents"] + ) + assert any( + event.get("ph") == "I" and event.get("cat") == "consistency.drift" + for event in trace["traceEvents"] + ) + + output = write_drift_trace(report, tmp_path / "drift.json") + assert output.read_text(encoding="utf-8").startswith('{\n "traceEvents"') + + +@pytest.mark.unit +def test_consistency_bundle_contains_report_trace_and_preview(tmp_path: Path): + manifest, cube = _artifacts() + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + + output = write_drift_bundle(report, tmp_path / "drift.rlk-drift") + assert output.exists() + + import zipfile + + with zipfile.ZipFile(output) as archive: + assert set(archive.namelist()) == { + "manifest.json", + "report.json", + "trace.json", + "preview.png", + } + + loaded = load_drift_bundle(output) + assert loaded["manifest"]["format"] == "rl_kernel.cross_config_drift" + assert loaded["manifest"]["preview_included"] is True + assert loaded["report"]["status"] == "warning" + assert loaded["trace"]["metadata"]["timeline_mode"] == "ordinal_diagnostic" + + +@pytest.mark.unit +def test_consistency_bundle_can_omit_preview(tmp_path: Path): + manifest, cube = _artifacts() + report = build_drift_report(replay_manifest=manifest, result_cube=cube) + + output = write_drift_bundle( + report, + tmp_path / "drift-no-preview.rlk-drift", + include_preview=False, + ) + + import zipfile + + with zipfile.ZipFile(output) as archive: + assert "preview.png" not in archive.namelist() + bundle_manifest = json.loads(archive.read("manifest.json")) + assert bundle_manifest["preview_included"] is False + assert bundle_manifest["files"] == ["manifest.json", "report.json", "trace.json"] + + +@pytest.mark.unit +def test_attempt_report_reads_only_sealed_cross_config_artifacts(tmp_path: Path): + attempt_dir = _completed_attempt(tmp_path) + + report = build_cross_config_attempt_report(attempt_dir) + + assert report["status"] == "pass" + assert report["axes"]["rollout_tp"] == 2 + assert report["axes"]["rollout_cp"] == 2 + assert report["metrics"]["active_token_count"] == 2 + assert ( + report["runtime_provenance"]["operator_source"] + == "exact_resolution_and_instance" + ) + + +@pytest.mark.unit +def test_attempt_report_rejects_unsealed_artifact_directory(tmp_path: Path): + store = ArtifactStore(tmp_path) + attempt = store.create_attempt("unfinished", "case-1") + + with pytest.raises(Exception, match="COMPLETE"): + build_cross_config_attempt_report(attempt) + + +@pytest.mark.unit +def test_cli_writes_offline_bundle_from_sealed_attempt(tmp_path: Path, capsys): + attempt = _completed_attempt(tmp_path) + output = tmp_path / "report.rlk-drift" + + assert ( + cli_main.main(["report", str(attempt), "--output", str(output), "--no-preview"]) + == 0 + ) + + assert output.is_file() + assert "drift report: pass" in capsys.readouterr().err From fb5196b9ae5ca8331cb8f245d4a7701afb0f5db6 Mon Sep 17 00:00:00 2001 From: Takumi <3051000145@qq.com> Date: Fri, 21 Aug 2026 16:18:30 +0800 Subject: [PATCH 25/33] fix(alignment): finalize offline drift report entrypoints From 3617b28aebbb39b6af44e7ccadcebbee0f6ad993 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Fri, 21 Aug 2026 17:56:57 +0800 Subject: [PATCH 26/33] feat(ws2): add Vime logprob provider for CP metadata --- rl_engine/integrations/__init__.py | 5 + rl_engine/integrations/vime/__init__.py | 9 + rl_engine/integrations/vime/logp.py | 229 ++++++++++++++++++ .../ops/pytorch/loss/vocab_parallel_logp.py | 138 ++++++++++- tests/test_vime_logprob_provider.py | 93 +++++++ tests/test_vocab_parallel_logp.py | 23 ++ 6 files changed, 486 insertions(+), 11 deletions(-) create mode 100644 rl_engine/integrations/__init__.py create mode 100644 rl_engine/integrations/vime/__init__.py create mode 100644 rl_engine/integrations/vime/logp.py create mode 100644 tests/test_vime_logprob_provider.py diff --git a/rl_engine/integrations/__init__.py b/rl_engine/integrations/__init__.py new file mode 100644 index 00000000..29bcc68f --- /dev/null +++ b/rl_engine/integrations/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Optional framework adapters owned by RL-Kernel.""" + diff --git a/rl_engine/integrations/vime/__init__.py b/rl_engine/integrations/vime/__init__.py new file mode 100644 index 00000000..44470314 --- /dev/null +++ b/rl_engine/integrations/vime/__init__.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Vime adapter entry points without a Vime runtime dependency.""" + +from .logp import ProviderResult, SelectedLogprobProviderUnavailable, provider + +__all__ = ["ProviderResult", "SelectedLogprobProviderUnavailable", "provider"] + diff --git a/rl_engine/integrations/vime/logp.py b/rl_engine/integrations/vime/logp.py new file mode 100644 index 00000000..6a47e725 --- /dev/null +++ b/rl_engine/integrations/vime/logp.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime selected-logprob provider for Vime's Megatron backend. + +The adapter intentionally accepts and returns structural objects: RL-Kernel +never imports Vime. Vime remains responsible for constructing locally owned +CP token rows and response masks; this provider owns only the TP-vocabulary +reduction. CP rank and layout are recorded and validated as row ownership +metadata, never passed to the numerical merge. +""" + +from __future__ import annotations + +import os +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import torch + +from rl_engine.kernels.logprob_contract import ( + LogprobContract, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp import ( + BACKEND_ID, + DEFAULT_NUM_VOCAB_TILES, +) +from rl_engine.kernels.registry import kernel_registry + + +class SelectedLogprobProviderUnavailable(RuntimeError): + """Request Vime's native provider fallback in ``auto`` mode. + + Vime recognizes the marker instead of importing this class, which keeps + the dependency direction from Vime to RL-Kernel at runtime only. + """ + + selected_logprob_provider_unavailable = True + + +@dataclass(frozen=True) +class ProviderResult: + """Structural result understood by the Vime provider boundary.""" + + selected_logprobs: torch.Tensor + entropy: torch.Tensor | None + backend_id: str + contract_id: str + provenance: Mapping[str, Any] + + +def _as_positive_int(value: Any, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise SelectedLogprobProviderUnavailable(f"{name} must be a positive integer; got {value!r}") + return value + + +def _metadata(request: Any) -> Mapping[str, Any]: + value = getattr(request, "metadata", None) + if not isinstance(value, Mapping): + raise SelectedLogprobProviderUnavailable("request.metadata must provide vocab-parallel metadata") + return value + + +def _request_tensor(request: Any, name: str) -> torch.Tensor: + value = getattr(request, name, None) + if not isinstance(value, torch.Tensor): + raise SelectedLogprobProviderUnavailable(f"request.{name} must be a torch.Tensor") + return value + + +def _tp_coordinates(tp_group: Any) -> tuple[int, int]: + if tp_group is not None and hasattr(tp_group, "rank") and hasattr(tp_group, "size"): + return int(tp_group.rank()), int(tp_group.size()) + + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return dist.get_rank(group=tp_group), dist.get_world_size(group=tp_group) + return 0, 1 + + +def _tile_count(metadata: Mapping[str, Any], padded_vocab_size: int) -> int: + configured = metadata.get("num_vocab_tiles", os.getenv("RL_KERNEL_LOGPROB_NUM_VOCAB_TILES")) + if configured is None or configured == "": + configured = DEFAULT_NUM_VOCAB_TILES + try: + tiles = int(configured) + except (TypeError, ValueError) as exc: + raise SelectedLogprobProviderUnavailable( + f"num_vocab_tiles must be an integer; got {configured!r}" + ) from exc + if tiles <= 0 or padded_vocab_size % tiles: + raise SelectedLogprobProviderUnavailable( + f"num_vocab_tiles={tiles} must divide padded_vocab_size={padded_vocab_size}" + ) + return tiles + + +def _contract_for_request(request: Any) -> tuple[LogprobContract, int]: + logits = _request_tensor(request, "logits") + targets = _request_tensor(request, "target_ids") + metadata = _metadata(request) + if logits.ndim != 2 or targets.shape != (logits.shape[0],): + raise SelectedLogprobProviderUnavailable("request must contain local [T, V] logits and aligned [T] targets") + if logits.dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise SelectedLogprobProviderUnavailable(f"unsupported logit dtype {logits.dtype}") + if targets.device != logits.device: + raise SelectedLogprobProviderUnavailable("target_ids must share the local logits device") + + cp = getattr(request, "context_parallel", None) + cp_world_size = _as_positive_int(getattr(cp, "world_size", None), "context_parallel.world_size") + cp_rank = getattr(cp, "rank", None) + if isinstance(cp_rank, bool) or not isinstance(cp_rank, int) or not 0 <= cp_rank < cp_world_size: + raise SelectedLogprobProviderUnavailable( + f"context_parallel.rank={cp_rank!r} is invalid for CP={cp_world_size}" + ) + if getattr(cp, "layout", None) not in ({"single"} if cp_world_size == 1 else {"zigzag", "allgather"}): + raise SelectedLogprobProviderUnavailable("context_parallel layout does not describe local CP token ownership") + + tp_rank, tp_world_size = _tp_coordinates(getattr(request, "tensor_parallel_group", None)) + declared_tp_rank = metadata.get("tp_rank") + declared_tp_world_size = metadata.get("tp_world_size") + if declared_tp_rank is not None and declared_tp_rank != tp_rank: + raise SelectedLogprobProviderUnavailable( + f"metadata tp_rank={declared_tp_rank} disagrees with TP group rank={tp_rank}" + ) + if declared_tp_world_size is not None and declared_tp_world_size != tp_world_size: + raise SelectedLogprobProviderUnavailable( + f"metadata tp_world_size={declared_tp_world_size} disagrees with TP group size={tp_world_size}" + ) + + real_vocab_size = _as_positive_int(metadata.get("real_vocab_size"), "real_vocab_size") + padded_vocab_size = _as_positive_int(metadata.get("padded_vocab_size"), "padded_vocab_size") + if logits.shape[1] * tp_world_size != padded_vocab_size: + raise SelectedLogprobProviderUnavailable( + "local vocab width and TP group do not cover padded_vocab_size exactly: " + f"{logits.shape[1]} * {tp_world_size} != {padded_vocab_size}" + ) + if real_vocab_size > padded_vocab_size: + raise SelectedLogprobProviderUnavailable("real_vocab_size must not exceed padded_vocab_size") + + bounds = tuple( + (rank * logits.shape[1], (rank + 1) * logits.shape[1]) for rank in range(tp_world_size) + ) + active_mask = (True,) * logits.shape[0] + contract = LogprobContract( + role=LogprobRole.TRAIN, + dtype={ + torch.bfloat16: LogprobDType.BF16, + torch.float16: LogprobDType.FP16, + torch.float32: LogprobDType.FP32, + }[logits.dtype], + mask=MaskSpec(num_tokens=logits.shape[0], active_mask=active_mask), + sharding=ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=bounds, + real_vocab_size=real_vocab_size, + padded_vocab_size=padded_vocab_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + ), + reduction=ReductionSpec(), + ) + return contract, _tile_count(metadata, padded_vocab_size) + + +def provider(request: Any) -> ProviderResult: + """Compute Vime selected logprobs on the explicit WS2 TP/CP contract. + + Top-p replay is deliberately unavailable until it has a separately + validated fixed-order mask contract. In Vime ``auto`` mode this signals + native execution; in ``strict`` mode it fails instead of changing sampled + distribution semantics. + """ + + if getattr(request, "log_prob_keep_mask", None) is not None: + raise SelectedLogprobProviderUnavailable( + "RL-Kernel WS2 logprob does not yet materialize Vime top-p replay masks" + ) + + contract, num_vocab_tiles = _contract_for_request(request) + dispatch = kernel_registry.get_logprob_op(contract, requested_backend=BACKEND_ID) + if dispatch.provenance["actual_backend"] != BACKEND_ID or dispatch.provenance["fallback"]: + raise RuntimeError("explicit WS2 backend dispatch changed during materialization") + if getattr(request, "with_entropy", False): + selected_logp, _lse, entropy = dispatch.op.apply_with_entropy( + request.logits, + request.target_ids, + contract=contract, + tp_group=getattr(request, "tensor_parallel_group", None), + num_vocab_tiles=num_vocab_tiles, + with_entropy_grad=bool(getattr(request, "with_entropy_grad", False)), + ) + else: + selected_logp, _lse = dispatch.op( + request.logits, + request.target_ids, + contract=contract, + tp_group=getattr(request, "tensor_parallel_group", None), + num_vocab_tiles=num_vocab_tiles, + ) + entropy = None + provenance = dict(dispatch.provenance) + provenance["cp_row_ownership"] = { + "cp_rank": contract.sharding.cp_rank, + "cp_world_size": contract.sharding.cp_world_size, + "layout": getattr(request.context_parallel, "layout"), + "local_token_rows": int(request.logits.shape[0]), + "cp_is_merge_axis": False, + } + provenance["num_vocab_tiles"] = num_vocab_tiles + return ProviderResult( + selected_logprobs=selected_logp.unsqueeze(-1), + entropy=entropy, + backend_id=dispatch.capability.backend_id, + contract_id=contract.cross_rank_fingerprint(), + provenance=provenance, + ) + + +__all__ = ["ProviderResult", "SelectedLogprobProviderUnavailable", "provider"] diff --git a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py index 5c3f56ae..d5e8884f 100644 --- a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py @@ -294,6 +294,28 @@ def _gather_target_logit( return stacked[owner, rows] +def _gather_entropy_partials( + local_entropy: torch.Tensor, + contract: LogprobContract, + tp_group: Any, +) -> torch.Tensor: + """Merge per-shard entropy contributions in TP-rank order. + + Entropy is not part of the WS2 selected-logprob contract, but the Vime + adapter needs it for the existing training loss surface. The collective + transports independent per-shard contributions; every TP rank performs + the same explicit rank-ordered sum afterwards. + """ + + if contract.sharding.tp_world_size == 1: + return local_entropy + + dist = _require_distributed_initialized() + gathered = [torch.empty_like(local_entropy) for _ in range(contract.sharding.tp_world_size)] + dist.all_gather(gathered, local_entropy.contiguous(), group=tp_group) + return torch.stack(gathered, dim=0).sum(dim=0) + + def _merge_tile_partials(m_all: torch.Tensor, s_all: torch.Tensor) -> torch.Tensor: """Fixed-order (max, sumexp) merge over [n, num_vocab_tiles].""" @@ -307,7 +329,17 @@ def _merge_tile_partials(m_all: torch.Tensor, s_all: torch.Tensor) -> torch.Tens class _VocabParallelLogprobFunction(torch.autograd.Function): @staticmethod - def forward(ctx, local_logits, target_1d, active_mask, contract, tp_group, tile): + def forward( + ctx, + local_logits, + target_1d, + active_mask, + contract, + tp_group, + tile, + with_entropy, + with_entropy_grad, + ): z_masked = local_logits.float() sharding = contract.sharding global_ids = torch.arange( @@ -333,19 +365,40 @@ def forward(ctx, local_logits, target_1d, active_mask, contract, tp_group, tile) selected_logp = torch.where(active_mask, target_logit - lse, torch.zeros_like(lse)) - ctx.save_for_backward(z_masked, lse, safe_target, active_mask, padding_cols) + if with_entropy: + finite_row = torch.isfinite(lse) + lse_safe = torch.where(finite_row, lse, torch.zeros_like(lse)) + probabilities = (z_masked - lse_safe.unsqueeze(1)).exp() + probabilities = torch.where(finite_row.unsqueeze(1), probabilities, torch.zeros_like(probabilities)) + finite_logits = torch.isfinite(z_masked) + log_gap = torch.where( + finite_logits, + lse_safe.unsqueeze(1) - z_masked, + torch.zeros_like(z_masked), + ) + local_entropy = (probabilities * log_gap).sum(dim=1) + entropy = _gather_entropy_partials(local_entropy, contract, tp_group) + else: + entropy = local_logits.new_empty((0,), dtype=torch.float32) + + ctx.save_for_backward(z_masked, lse, safe_target, active_mask, padding_cols, entropy) ctx.local_vocab_start = sharding.local_vocab_start ctx.local_vocab_size = sharding.local_vocab_size ctx.input_dtype = local_logits.dtype + ctx.with_entropy_grad = bool(with_entropy and with_entropy_grad) ctx.set_materialize_grads(False) - return selected_logp, lse + if with_entropy and not ctx.with_entropy_grad: + ctx.mark_non_differentiable(entropy) + return selected_logp, lse, entropy @staticmethod - def backward(ctx, grad_logp, grad_lse): - if not ctx.needs_input_grad[0] or (grad_logp is None and grad_lse is None): - return None, None, None, None, None, None + def backward(ctx, grad_logp, grad_lse, grad_entropy): + if not ctx.needs_input_grad[0] or ( + grad_logp is None and grad_lse is None and grad_entropy is None + ): + return None, None, None, None, None, None, None, None - z_masked, lse, safe_target, active_mask, padding_cols = ctx.saved_tensors + z_masked, lse, safe_target, active_mask, padding_cols, entropy = ctx.saved_tensors n, local_vocab = z_masked.shape finite_row = torch.isfinite(lse) lse_safe = torch.where(finite_row, lse, torch.zeros_like(lse)) @@ -366,9 +419,13 @@ def backward(ctx, grad_logp, grad_lse): grad = grad + g_logp.unsqueeze(1) * (onehot - p) if grad_lse is not None: grad = grad + grad_lse.unsqueeze(1) * p + if ctx.with_entropy_grad and grad_entropy is not None: + entropy_input = lse_safe.unsqueeze(1) - z_masked - entropy.unsqueeze(1) + entropy_input = torch.where(torch.isfinite(z_masked), entropy_input, torch.zeros_like(entropy_input)) + grad = grad + grad_entropy.unsqueeze(1) * p * entropy_input if bool(padding_cols.any()): grad = grad.masked_fill(padding_cols.unsqueeze(0), 0.0) - return grad.to(ctx.input_dtype), None, None, None, None, None + return grad.to(ctx.input_dtype), None, None, None, None, None, None, None class VocabParallelLogprobOp: @@ -418,6 +475,58 @@ def apply( validate: bool = True, deterministic: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: + selected_logp, lse, _ = self._apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + with_entropy=False, + with_entropy_grad=False, + ) + return selected_logp, lse + + def apply_with_entropy( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any = None, + num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, + validate: bool = True, + with_entropy_grad: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return selected logprob, vocabulary LSE, and full-vocabulary entropy. + + The method is intentionally separate from :meth:`apply` so the WS2 + selected-logprob surface remains unchanged for existing callers. + """ + + return self._apply( + local_logits, + target_ids, + contract=contract, + tp_group=tp_group, + num_vocab_tiles=num_vocab_tiles, + validate=validate, + with_entropy=True, + with_entropy_grad=with_entropy_grad, + ) + + def _apply( + self, + local_logits: torch.Tensor, + target_ids: torch.Tensor, + *, + contract: LogprobContract, + tp_group: Any, + num_vocab_tiles: int, + validate: bool, + with_entropy: bool, + with_entropy_grad: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if not isinstance(contract, LogprobContract): raise LogprobContractError("contract must be a LogprobContract") if not isinstance(deterministic, bool): @@ -444,8 +553,15 @@ def apply( if contract.sharding.tp_world_size > 1: _preflight_cross_rank_agreement(contract, tp_group, num_vocab_tiles, deterministic) - selected_logp, lse = _VocabParallelLogprobFunction.apply( - local_logits, target_1d, active_mask, contract, tp_group, tile + selected_logp, lse, entropy = _VocabParallelLogprobFunction.apply( + local_logits, + target_1d, + active_mask, + contract, + tp_group, + tile, + with_entropy, + with_entropy_grad, ) if validate and bool((~torch.isfinite(lse) & active_mask).any().item()): @@ -453,7 +569,7 @@ def apply( "non-finite logsumexp on an active row: logits over the real " "vocabulary must be finite for every active token" ) - return selected_logp, lse + return selected_logp, lse, entropy __all__ = [ diff --git a/tests/test_vime_logprob_provider.py b/tests/test_vime_logprob_provider.py new file mode 100644 index 00000000..021f4e9f --- /dev/null +++ b/tests/test_vime_logprob_provider.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU coverage for the optional Vime WS2 selected-logprob adapter.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.integrations.vime.logp import SelectedLogprobProviderUnavailable, provider + + +def _request(*, cp_rank: int = 0, with_entropy: bool = False, keep_mask=None): + logits = torch.tensor( + [[0.25, -0.5, 1.0, 0.1, -0.3, 0.6, -0.7, 0.4] for _ in range(3)], + dtype=torch.float32, + requires_grad=True, + ) + return SimpleNamespace( + logits=logits, + target_ids=torch.tensor([2, 5, 0]), + tensor_parallel_group=None, + context_parallel=SimpleNamespace( + world_size=2, + rank=cp_rank, + layout="zigzag", + ), + with_entropy=with_entropy, + with_entropy_grad=with_entropy, + log_prob_keep_mask=keep_mask, + metadata={ + "real_vocab_size": 7, + "padded_vocab_size": 8, + "tp_rank": 0, + "tp_world_size": 1, + "num_vocab_tiles": 4, + }, + ) + + +def test_provider_runs_locally_with_cp2_row_metadata(): + request = _request(cp_rank=1) + + result = provider(request) + reference = torch.log_softmax(request.logits[:, :7], dim=-1)[ + torch.arange(request.logits.size(0)), request.target_ids + ] + + assert result.selected_logprobs.shape == (3, 1) + torch.testing.assert_close(result.selected_logprobs.squeeze(-1), reference) + assert result.backend_id == "pytorch-vocab-parallel-logp-ws2" + assert result.provenance["cp_row_ownership"] == { + "cp_rank": 1, + "cp_world_size": 2, + "layout": "zigzag", + "local_token_rows": 3, + "cp_is_merge_axis": False, + } + + +def test_provider_entropy_preserves_vime_semantics_and_autograd(): + request = _request(with_entropy=True) + + result = provider(request) + reference_logits = request.logits.detach().clone().requires_grad_(True) + log_probs = torch.log_softmax(reference_logits[:, :7], dim=-1) + reference_logp = log_probs[torch.arange(reference_logits.size(0)), request.target_ids] + reference_entropy = -(log_probs.exp() * log_probs).sum(dim=-1) + + torch.testing.assert_close(result.selected_logprobs.squeeze(-1), reference_logp) + torch.testing.assert_close(result.entropy, reference_entropy) + (result.selected_logprobs.sum() + result.entropy.sum()).backward() + (reference_logp.sum() + reference_entropy.sum()).backward() + torch.testing.assert_close(request.logits.grad[:, :7], reference_logits.grad[:, :7]) + assert bool((request.logits.grad[:, 7] == 0).all()) + + +def test_provider_rejects_top_p_replay_without_changing_its_semantics(): + request = _request(keep_mask=torch.ones((3, 8), dtype=torch.bool)) + + with pytest.raises(SelectedLogprobProviderUnavailable, match="top-p replay"): + provider(request) + + +def test_provider_rejects_local_vocab_metadata_that_cannot_describe_tp_ownership(): + request = _request() + request.metadata["padded_vocab_size"] = 16 + + with pytest.raises(SelectedLogprobProviderUnavailable, match="cover padded_vocab_size"): + provider(request) diff --git a/tests/test_vocab_parallel_logp.py b/tests/test_vocab_parallel_logp.py index 1dee8bff..eefc1d6a 100644 --- a/tests/test_vocab_parallel_logp.py +++ b/tests/test_vocab_parallel_logp.py @@ -317,6 +317,29 @@ def test_inactive_rows_grad_asymmetry(self): logp.sum().backward() assert bool((z.grad[-1] == 0).all()) + def test_entropy_matches_full_vocab_oracle_and_backpropagates(self): + contract = _contract() + logits, targets = _inputs() + x = logits.clone().requires_grad_(True) + logp, _lse, entropy = VocabParallelLogprobOp().apply_with_entropy( + x, + targets, + contract=contract, + num_vocab_tiles=NUM_TILES, + ) + + y = logits.clone().requires_grad_(True) + ref_log_probs = torch.log_softmax(y[:, :REAL_VOCAB].float(), dim=-1) + safe = targets.clamp(0, REAL_VOCAB - 1) + ref_logp = ref_log_probs[torch.arange(NUM_TOKENS), safe] + ref_logp = torch.where(torch.tensor(ACTIVE), ref_logp, torch.zeros_like(ref_logp)) + ref_entropy = -(ref_log_probs.exp() * ref_log_probs).sum(dim=-1) + + torch.testing.assert_close(entropy, ref_entropy) + (logp.sum() + entropy.sum()).backward() + (ref_logp.sum() + ref_entropy.sum()).backward() + torch.testing.assert_close(x.grad, y.grad, atol=2e-5, rtol=2e-5) + def test_dispatch_resolves_reference_and_leaves_legacy_untouched(): registry = KernelRegistry() From a9acf48d557f759cad852c8c046cd164f7cec59f Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Fri, 21 Aug 2026 18:04:33 +0800 Subject: [PATCH 27/33] fix(ws2): keep TP entropy merge order explicit --- docs/operators/batch-invariant-logp.md | 23 +++++++++++++++++++ .../ops/pytorch/loss/vocab_parallel_logp.py | 5 +++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index f840ecfd..bb05c038 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -87,6 +87,29 @@ op = result.op # rl_engine.kernels.logprob_ logp, lse = op(local_logits, target_ids, contract=contract, tp_group=tp_group) ``` +### Vime CP=2 runtime provider + +The optional Vime adapter is owned by RL-Kernel and can be selected without +patching Megatron or vLLM: + +```text +--selected-logprob-provider rl_engine.integrations.vime.logp.provider +--selected-logprob-provider-mode strict +``` + +Vime passes the local `[T, V_local]` logits, shifted targets, TP subgroup, +and CP row-ownership metadata. The provider builds the same `LogprobContract` +used by the distributed report, dispatches the explicit +`pytorch-vocab-parallel-logp-ws2` backend, and returns selected logp as `[T, 1]`. +When entropy is requested, it uses the same fixed TP-rank order and returns +full-vocabulary entropy for the existing loss surface. CP rank/layout are +recorded in provenance and never participate in the vocabulary LSE merge. + +The provider fails closed for undeclared real/padded vocabulary sizes, TP/CP +metadata mismatches, unsupported top-p replay masks, and backend fallback. +`auto` mode may then use Vime's native path; `strict` mode reports the +configuration error. This adapter does not import Vime. + ## Benchmarks `benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton, and the diff --git a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py index d5e8884f..4fc432a1 100644 --- a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py @@ -313,7 +313,10 @@ def _gather_entropy_partials( dist = _require_distributed_initialized() gathered = [torch.empty_like(local_entropy) for _ in range(contract.sharding.tp_world_size)] dist.all_gather(gathered, local_entropy.contiguous(), group=tp_group) - return torch.stack(gathered, dim=0).sum(dim=0) + merged = gathered[0].clone() + for partial in gathered[1:]: + merged = merged + partial + return merged def _merge_tile_partials(m_all: torch.Tensor, s_all: torch.Tensor) -> torch.Tensor: From d718bb939a89bc2001a1aacc27064c78ed629795 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Fri, 21 Aug 2026 18:52:07 +0800 Subject: [PATCH 28/33] feat(vime): add Qwen3 TP2 CP2 validation example --- examples/vime_qwen3_8b_tp2_cp2/README.md | 69 +++++ .../qwen3_8b_tp2_cp2.json | 43 +++ examples/vime_qwen3_8b_tp2_cp2/run.py | 261 ++++++++++++++++++ rl_engine/integrations/vime/logp.py | 17 ++ tests/test_vime_qwen3_example.py | 126 +++++++++ 5 files changed, 516 insertions(+) create mode 100644 examples/vime_qwen3_8b_tp2_cp2/README.md create mode 100644 examples/vime_qwen3_8b_tp2_cp2/qwen3_8b_tp2_cp2.json create mode 100644 examples/vime_qwen3_8b_tp2_cp2/run.py create mode 100644 tests/test_vime_qwen3_example.py diff --git a/examples/vime_qwen3_8b_tp2_cp2/README.md b/examples/vime_qwen3_8b_tp2_cp2/README.md new file mode 100644 index 00000000..30d3734a --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/README.md @@ -0,0 +1,69 @@ +# Vime Qwen3-8B TP=2 CP=2 validation + +This example is the recommended reproducible entry point for the Vime-side +selected-logprob integration. It keeps framework glue in Vime and keeps the +numerical provider, contract, provenance, and report in RL-Kernel. + +The example is deliberately strict: + +- Megatron training uses `TP=2`, `CP=2`, `PP=1`, and four actor ranks. +- vLLM rollout uses processed logprobs and `top_p=1.0`. +- Vime must load `rl_engine.integrations.vime.logp.provider` in `strict` mode. +- A native fallback or a missing provider marker is not reported as a pass. +- Attention and FFN are not declared consistent from configuration alone. They + require executed Megatron and vLLM readbacks, so the report marks them + `unclaimed` until those artifacts are supplied. The readback must use + `rlkernel.operator_runtime_evidence.v1` and report exact-zero comparison + metrics for both sides. + +The Vime companion must be installed or checked out separately. This example +does not modify `vllm-project/vime`. + +## Dry run + +```bash +python examples/vime_qwen3_8b_tp2_cp2/run.py \ + --vime-root /path/to/RL-Align/vime \ + --rl-kernel-root /path/to/RL-Kernel \ + --output reports/qwen3_8b_tp2_cp2.validation.json +``` + +## Execute + +The Vime script expects model/checkpoint/data paths through environment +variables. Override them before adding `--run`: + +```bash +export MODEL_ROOT=/models/Qwen3-8B +export TORCH_DIST_ROOT=/models/Qwen3-8B_torch_dist +export PROMPT_DATA=/data/dapo-math-17k.jsonl +export RL_KERNEL_ROOT=/path/to/RL-Kernel + +python examples/vime_qwen3_8b_tp2_cp2/run.py \ + --vime-root /path/to/RL-Align/vime \ + --rl-kernel-root "$RL_KERNEL_ROOT" \ + --output reports/qwen3_8b_tp2_cp2.validation.json \ + --run +``` + +When the Megatron/vLLM launch also emits the operator readback artifact, pass +it explicitly: + +```bash +python examples/vime_qwen3_8b_tp2_cp2/run.py \ + --vime-root /path/to/RL-Align/vime \ + --rl-kernel-root "$RL_KERNEL_ROOT" \ + --runtime-evidence reports/qwen3_8b_tp2_cp2.runtime-evidence.json \ + --output reports/qwen3_8b_tp2_cp2.validation.json \ + --run +``` + +The evidence file is intentionally post-execution. It must include training +and rollout identities for `attention` and `ffn`, plus `passed: true` and +exact-zero `out`, backward, and (for attention) `LSE` comparison metrics. A +configured backend without this readback remains `unclaimed`. + +The runner writes a JSON report and a sibling combined log. The report records +the exact command, both repository revisions, provider backend identity, strict +fallback status, and the claim boundary. It does not fabricate numerical drift +when the GPU run was not executed. diff --git a/examples/vime_qwen3_8b_tp2_cp2/qwen3_8b_tp2_cp2.json b/examples/vime_qwen3_8b_tp2_cp2/qwen3_8b_tp2_cp2.json new file mode 100644 index 00000000..93332cbb --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/qwen3_8b_tp2_cp2.json @@ -0,0 +1,43 @@ +{ + "schema_version": "rlkernel.vime_qwen3_8b_tp2_cp2.v1", + "model": "Qwen/Qwen3-8B", + "training": { + "framework": "megatron", + "tensor_model_parallel_size": 2, + "context_parallel_size": 2, + "pipeline_model_parallel_size": 1, + "world_size": 4, + "dtype": "bf16" + }, + "rollout": { + "framework": "vllm", + "top_p": 1.0, + "logprobs_mode": "processed_logprobs" + }, + "selected_logprob_provider": { + "path": "rl_engine.integrations.vime.logp.provider", + "mode": "strict", + "backend_id": "pytorch-vocab-parallel-logp-ws2", + "real_vocab_size": 151936, + "padded_vocab_size": 152064, + "num_vocab_tiles": 64 + }, + "operator_evidence": { + "logp": { + "training": "rl-kernel-provider", + "rollout": "vllm-native-processed-logprobs", + "required_runtime_marker": "Selected-logprob provider active" + }, + "attention": { + "training": "runtime-readback-required", + "rollout": "runtime-readback-required", + "status": "not_claimed_without_Megatron_and_vLLM_readback" + }, + "ffn": { + "training": "runtime-readback-required", + "rollout": "runtime-readback-required", + "status": "not_claimed_without_Megatron_and_vLLM_readback" + } + }, + "vime_script": "scripts/run-qwen3-8B-rlkernel-tp2-cp2.sh" +} diff --git a/examples/vime_qwen3_8b_tp2_cp2/run.py b/examples/vime_qwen3_8b_tp2_cp2/run.py new file mode 100644 index 00000000..4823dd58 --- /dev/null +++ b/examples/vime_qwen3_8b_tp2_cp2/run.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Run and archive the Vime Qwen3-8B TP=2/CP=2 validation entry point. + +This is an integration example, not a synthetic pass generator. A dry run +only records the exact launch contract. ``--run`` executes Vime and records +whether the strict RL-Kernel provider was actually observed in the log. The +report deliberately leaves attention/FFN unclaimed until both framework +readbacks are supplied. +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Mapping + + +DEFAULT_CONFIG = Path(__file__).with_name("qwen3_8b_tp2_cp2.json") +PROVIDER_MARKER = "Selected-logprob provider active" +FALLBACK_MARKERS = ("using native path", "fallback=True", "fallback=true") +RUNTIME_EVIDENCE_SCHEMA = "rlkernel.operator_runtime_evidence.v1" +_OPERATOR_METRICS = { + "attention": ("out_max_abs", "lse_max_abs", "dq_max_abs", "dk_max_abs", "dv_max_abs"), + "ffn": ("out_max_abs", "dx_max_abs", "dw_max_abs"), +} + + +def load_config(path: Path) -> dict[str, Any]: + with path.open(encoding="utf-8") as handle: + value = json.load(handle) + if not isinstance(value, dict): + raise ValueError("example config must contain a JSON object") + return value + + +def validate_config(config: Mapping[str, Any]) -> None: + training = config.get("training") + rollout = config.get("rollout") + provider = config.get("selected_logprob_provider") + if not isinstance(training, Mapping) or not isinstance(rollout, Mapping) or not isinstance(provider, Mapping): + raise ValueError("training, rollout, and selected_logprob_provider sections are required") + expected = { + "tensor_model_parallel_size": 2, + "context_parallel_size": 2, + "pipeline_model_parallel_size": 1, + "world_size": 4, + } + for name, value in expected.items(): + if training.get(name) != value: + raise ValueError(f"training.{name} must be {value!r}") + if rollout.get("top_p") != 1.0: + raise ValueError("rollout.top_p must remain 1.0 for the strict provider contract") + if provider.get("mode") != "strict": + raise ValueError("selected_logprob_provider.mode must be strict") + if provider.get("path") != "rl_engine.integrations.vime.logp.provider": + raise ValueError("example must use the RL-Kernel Vime provider") + if provider.get("backend_id") != "pytorch-vocab-parallel-logp-ws2": + raise ValueError("example must pin the WS2 vocab-parallel backend") + + +def load_runtime_evidence(path: Path | None) -> dict[str, Any] | None: + """Load post-execution readback without treating configuration as evidence.""" + + if path is None: + return None + with path.open(encoding="utf-8") as handle: + value = json.load(handle) + if not isinstance(value, dict) or value.get("schema_version") != RUNTIME_EVIDENCE_SCHEMA: + raise ValueError(f"runtime evidence must use schema {RUNTIME_EVIDENCE_SCHEMA!r}") + return value + + +def _operator_evidence_status(evidence: Mapping[str, Any] | None, operator: str) -> str: + if evidence is None: + return "unclaimed" + operators = evidence.get("operators") + item = operators.get(operator) if isinstance(operators, Mapping) else None + if not isinstance(item, Mapping): + return "unclaimed" + training = item.get("training") + rollout = item.get("rollout") + comparison = item.get("comparison") + if not isinstance(training, Mapping) or not isinstance(rollout, Mapping): + return "unclaimed" + if not isinstance(comparison, Mapping) or comparison.get("passed") is not True: + return "failed" + required_identity = ("implementation_id", "backend_id", "contract_id") + if any(not training.get(name) or not rollout.get(name) for name in required_identity): + return "failed" + if training["implementation_id"] != rollout["implementation_id"]: + return "failed" + for metric in _OPERATOR_METRICS[operator]: + value = comparison.get(metric) + if not isinstance(value, (int, float)) or isinstance(value, bool) or value != 0.0: + return "failed" + return "passed" + + +def validate_runtime_evidence(evidence: Mapping[str, Any] | None) -> None: + """Reject malformed evidence before it can affect a report.""" + + if evidence is None: + return + for operator in _OPERATOR_METRICS: + status = _operator_evidence_status(evidence, operator) + if status == "failed": + raise ValueError(f"runtime evidence for {operator} is incomplete or non-zero") + + +def _revision(path: Path) -> str | None: + try: + result = subprocess.run( + ["git", "-C", str(path), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return None + return result.stdout.strip() or None + + +def build_environment(vime_root: Path, rl_kernel_root: Path) -> dict[str, str]: + env = dict(os.environ) + existing = [str(vime_root), str(rl_kernel_root), "/root/Megatron-LM"] + if env.get("PYTHONPATH"): + existing.append(env["PYTHONPATH"]) + env["PYTHONPATH"] = os.pathsep.join(existing) + env["RL_KERNEL_ROOT"] = str(rl_kernel_root) + env["TP_SIZE"] = "2" + env["CP_SIZE"] = "2" + env["ROLLOUT_TOP_P"] = "1.0" + return env + + +def build_command(config: Mapping[str, Any], vime_root: Path) -> list[str]: + script = vime_root / str(config.get("vime_script", "")) + if not script.is_file(): + raise FileNotFoundError(f"Vime entry script does not exist: {script}") + return ["bash", str(script)] + + +def build_report( + config: Mapping[str, Any], + *, + vime_root: Path, + rl_kernel_root: Path, + command: list[str], + status: str, + returncode: int | None, + log_text: str, + log_path: Path | None, + runtime_evidence: Mapping[str, Any] | None = None, + runtime_evidence_path: Path | None = None, +) -> dict[str, Any]: + provider_active = PROVIDER_MARKER in log_text + fallback_observed = any(marker in log_text for marker in FALLBACK_MARKERS) + strict_provider_passed = status == "passed" and provider_active and not fallback_observed + effective_status = "passed" if strict_provider_passed else ("failed" if status == "passed" else status) + attention_status = _operator_evidence_status(runtime_evidence, "attention") + ffn_status = _operator_evidence_status(runtime_evidence, "ffn") + return { + "schema_version": "rlkernel.vime_validation_report.v1", + "created_at": datetime.now(timezone.utc).isoformat(), + "status": effective_status, + "claim_boundary": { + "qwen3_8b_tp2_cp2_vime_training": strict_provider_passed, + "attention_train_infer_consistency": attention_status, + "ffn_train_infer_consistency": ffn_status, + "reason": ( + "attention and FFN require executed Megatron/vLLM runtime readbacks; " + "the evidence contract accepts only exact-zero comparison metrics" + ), + }, + "config": dict(config), + "topology": config["training"], + "provider": { + "configured_path": config["selected_logprob_provider"]["path"], + "configured_mode": config["selected_logprob_provider"]["mode"], + "backend_id": config["selected_logprob_provider"]["backend_id"], + "active_observed": provider_active, + "fallback_observed": fallback_observed, + }, + "command": command, + "returncode": returncode, + "artifacts": { + "log": None if log_path is None else str(log_path), + "runtime_evidence": None if runtime_evidence_path is None else str(runtime_evidence_path), + }, + "runtime_evidence": None if runtime_evidence is None else dict(runtime_evidence), + "revisions": { + "vime": _revision(vime_root), + "rl_kernel": _revision(rl_kernel_root), + }, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG) + parser.add_argument("--vime-root", type=Path, default=Path(os.environ.get("VIME_ROOT", "."))) + parser.add_argument("--rl-kernel-root", type=Path, default=Path(os.environ.get("RL_KERNEL_ROOT", "."))) + parser.add_argument("--output", type=Path, default=Path("qwen3_8b_tp2_cp2.validation.json")) + parser.add_argument( + "--runtime-evidence", + type=Path, + default=None, + help="post-execution Megatron/vLLM operator readback JSON (strict exact-zero contract)", + ) + parser.add_argument("--run", action="store_true", help="execute the Vime script") + args = parser.parse_args(argv) + + config = load_config(args.config) + validate_config(config) + runtime_evidence = load_runtime_evidence(args.runtime_evidence) + validate_runtime_evidence(runtime_evidence) + vime_root = args.vime_root.resolve() + rl_kernel_root = args.rl_kernel_root.resolve() + command = build_command(config, vime_root) + + status = "not_run" + returncode: int | None = None + log_text = "" + log_path: Path | None = None + if args.run: + args.output.parent.mkdir(parents=True, exist_ok=True) + log_path = args.output.with_suffix(".log") + env = build_environment(vime_root, rl_kernel_root) + with log_path.open("w", encoding="utf-8") as log_handle: + process = subprocess.run(command, cwd=vime_root, env=env, stdout=log_handle, stderr=subprocess.STDOUT) + returncode = process.returncode + log_text = log_path.read_text(encoding="utf-8", errors="replace") + status = "passed" if returncode == 0 else "failed" + + report = build_report( + config, + vime_root=vime_root, + rl_kernel_root=rl_kernel_root, + command=command, + status=status, + returncode=returncode, + log_text=log_text, + log_path=log_path, + runtime_evidence=runtime_evidence, + runtime_evidence_path=args.runtime_evidence, + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 if report["status"] in {"passed", "not_run"} else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/rl_engine/integrations/vime/logp.py b/rl_engine/integrations/vime/logp.py index 6a47e725..db436d69 100644 --- a/rl_engine/integrations/vime/logp.py +++ b/rl_engine/integrations/vime/logp.py @@ -209,6 +209,23 @@ def provider(request: Any) -> ProviderResult: ) entropy = None provenance = dict(dispatch.provenance) + provenance["request"] = { + "logits_shape": list(request.logits.shape), + "logits_dtype": str(request.logits.dtype).replace("torch.", ""), + "target_shape": list(request.target_ids.shape), + "target_dtype": str(request.target_ids.dtype).replace("torch.", ""), + "real_vocab_size": contract.sharding.real_vocab_size, + "padded_vocab_size": contract.sharding.padded_vocab_size, + "tp_rank": contract.sharding.tp_rank, + "tp_world_size": contract.sharding.tp_world_size, + "cp_rank": contract.sharding.cp_rank, + "cp_world_size": contract.sharding.cp_world_size, + } + provenance["execution"] = { + "role": "vime_training_selected_logprob", + "strict_backend": True, + "top_p_replay": False, + } provenance["cp_row_ownership"] = { "cp_rank": contract.sharding.cp_rank, "cp_world_size": contract.sharding.cp_world_size, diff --git a/tests/test_vime_qwen3_example.py b/tests/test_vime_qwen3_example.py new file mode 100644 index 00000000..007d4ba1 --- /dev/null +++ b/tests/test_vime_qwen3_example.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from examples.vime_qwen3_8b_tp2_cp2.run import ( + build_report, + load_config, + validate_config, + validate_runtime_evidence, +) + + +ROOT = Path(__file__).parents[1] +CONFIG = ROOT / "examples" / "vime_qwen3_8b_tp2_cp2" / "qwen3_8b_tp2_cp2.json" + + +def test_qwen3_example_config_is_strict_and_explicit(): + config = load_config(CONFIG) + validate_config(config) + assert config["training"]["tensor_model_parallel_size"] == 2 + assert config["training"]["context_parallel_size"] == 2 + assert config["selected_logprob_provider"]["mode"] == "strict" + + +def test_qwen3_example_report_does_not_claim_unread_back_attention_or_ffn(tmp_path): + config = load_config(CONFIG) + report = build_report( + config, + vime_root=tmp_path / "vime", + rl_kernel_root=tmp_path / "rl-kernel", + command=["bash", "run.sh"], + status="passed", + returncode=0, + log_text="Selected-logprob provider active: backend_id=pytorch-vocab-parallel-logp-ws2", + log_path=tmp_path / "run.log", + ) + assert report["status"] == "passed" + assert report["claim_boundary"]["qwen3_8b_tp2_cp2_vime_training"] is True + assert report["claim_boundary"]["attention_train_infer_consistency"] == "unclaimed" + assert report["claim_boundary"]["ffn_train_infer_consistency"] == "unclaimed" + assert report["provider"]["fallback_observed"] is False + + +def test_qwen3_example_fails_closed_when_provider_marker_is_missing(tmp_path): + config = load_config(CONFIG) + report = build_report( + config, + vime_root=tmp_path / "vime", + rl_kernel_root=tmp_path / "rl-kernel", + command=["bash", "run.sh"], + status="passed", + returncode=0, + log_text="training completed without provider provenance", + log_path=None, + ) + assert report["status"] == "failed" + assert report["claim_boundary"]["qwen3_8b_tp2_cp2_vime_training"] is False + + +def _runtime_evidence(): + return { + "schema_version": "rlkernel.operator_runtime_evidence.v1", + "operators": { + "attention": { + "training": {"implementation_id": "rlk.attn", "backend_id": "rlk", "contract_id": "a"}, + "rollout": {"implementation_id": "rlk.attn", "backend_id": "rlk", "contract_id": "a"}, + "comparison": { + "passed": True, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "dq_max_abs": 0.0, + "dk_max_abs": 0.0, + "dv_max_abs": 0.0, + }, + }, + "ffn": { + "training": {"implementation_id": "rlk.ffn", "backend_id": "rlk", "contract_id": "f"}, + "rollout": {"implementation_id": "rlk.ffn", "backend_id": "rlk", "contract_id": "f"}, + "comparison": { + "passed": True, + "out_max_abs": 0.0, + "dx_max_abs": 0.0, + "dw_max_abs": 0.0, + }, + }, + }, + } + + +def test_qwen3_example_accepts_only_exact_zero_runtime_evidence(tmp_path): + evidence = _runtime_evidence() + validate_runtime_evidence(evidence) + config = load_config(CONFIG) + report = build_report( + config, + vime_root=tmp_path / "vime", + rl_kernel_root=tmp_path / "rl-kernel", + command=["bash", "run.sh"], + status="passed", + returncode=0, + log_text="Selected-logprob provider active: backend_id=pytorch-vocab-parallel-logp-ws2", + log_path=None, + runtime_evidence=evidence, + ) + assert report["claim_boundary"]["attention_train_infer_consistency"] == "passed" + assert report["claim_boundary"]["ffn_train_infer_consistency"] == "passed" + + +def test_qwen3_example_rejects_nonzero_runtime_evidence(): + evidence = _runtime_evidence() + evidence["operators"]["attention"]["comparison"]["out_max_abs"] = 1e-6 + with pytest.raises(ValueError, match="attention"): + validate_runtime_evidence(evidence) + + +@pytest.mark.parametrize("bad_path", ["", "other.provider"]) +def test_qwen3_example_rejects_non_rlkernel_provider(bad_path): + config = load_config(CONFIG) + config["selected_logprob_provider"]["path"] = bad_path + with pytest.raises(ValueError, match="RL-Kernel Vime provider"): + validate_config(config) From 3c67928510944bd4bc84e8222e098d1fee79c810 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Fri, 21 Aug 2026 19:22:08 +0800 Subject: [PATCH 29/33] docs(example): document 8-GPU Qwen3 validation topology --- examples/vime_qwen3_8b_tp2_cp2/README.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/examples/vime_qwen3_8b_tp2_cp2/README.md b/examples/vime_qwen3_8b_tp2_cp2/README.md index 30d3734a..2039f118 100644 --- a/examples/vime_qwen3_8b_tp2_cp2/README.md +++ b/examples/vime_qwen3_8b_tp2_cp2/README.md @@ -19,6 +19,13 @@ The example is deliberately strict: The Vime companion must be installed or checked out separately. This example does not modify `vllm-project/vime`. +The executable entry point is the Vime script +`scripts/run-qwen3-8B-rlkernel-tp2-cp2.sh`. Its default topology is an +8-GPU H100 node with four Megatron actor GPUs and four vLLM rollout GPUs. The +script refuses to start on a different GPU count or GPU class. Set +`COLOCATE=1` only when intentionally testing the colocated path; that mode is +not the default 8-GPU train/infer split. + ## Dry run ```bash @@ -38,6 +45,7 @@ export MODEL_ROOT=/models/Qwen3-8B export TORCH_DIST_ROOT=/models/Qwen3-8B_torch_dist export PROMPT_DATA=/data/dapo-math-17k.jsonl export RL_KERNEL_ROOT=/path/to/RL-Kernel +export MEGATRON_ROOT=/path/to/Megatron-LM python examples/vime_qwen3_8b_tp2_cp2/run.py \ --vime-root /path/to/RL-Align/vime \ @@ -46,6 +54,12 @@ python examples/vime_qwen3_8b_tp2_cp2/run.py \ --run ``` +For a real 8xH100 run, the model, Megatron torch-dist checkpoint, prompt data, +and Megatron checkout must already exist on the host. The first run can omit +`VIME_CKPT`; the script will initialize from `TORCH_DIST_ROOT` and save the +Vime checkpoint there. Use `NUM_ROLLOUT=1` for the integration smoke test and +increase it only after the provider marker is observed. + When the Megatron/vLLM launch also emits the operator readback artifact, pass it explicitly: From 11617393c56a70aa16a630e2e642c2e6bf0db712 Mon Sep 17 00:00:00 2001 From: Codex H100 Validation Date: Fri, 21 Aug 2026 15:36:43 +0000 Subject: [PATCH 30/33] fix(integration): preserve test tolerance and logprob controls --- rl_engine/kernels/gtest/tolerance.py | 1 + rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py | 4 ++++ 2 files changed, 5 insertions(+) diff --git a/rl_engine/kernels/gtest/tolerance.py b/rl_engine/kernels/gtest/tolerance.py index 2660fb0f..d4537dc1 100644 --- a/rl_engine/kernels/gtest/tolerance.py +++ b/rl_engine/kernels/gtest/tolerance.py @@ -14,6 +14,7 @@ from __future__ import annotations +import hashlib import json import math from dataclasses import asdict, dataclass diff --git a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py index 4fc432a1..92e2e25c 100644 --- a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py @@ -485,6 +485,7 @@ def apply( tp_group=tp_group, num_vocab_tiles=num_vocab_tiles, validate=validate, + deterministic=deterministic, with_entropy=False, with_entropy_grad=False, ) @@ -500,6 +501,7 @@ def apply_with_entropy( num_vocab_tiles: int = DEFAULT_NUM_VOCAB_TILES, validate: bool = True, with_entropy_grad: bool = True, + deterministic: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Return selected logprob, vocabulary LSE, and full-vocabulary entropy. @@ -514,6 +516,7 @@ def apply_with_entropy( tp_group=tp_group, num_vocab_tiles=num_vocab_tiles, validate=validate, + deterministic=deterministic, with_entropy=True, with_entropy_grad=with_entropy_grad, ) @@ -527,6 +530,7 @@ def _apply( tp_group: Any, num_vocab_tiles: int, validate: bool, + deterministic: bool, with_entropy: bool, with_entropy_grad: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: From 784c05447dde8f119166884f2649a7effc54f6f5 Mon Sep 17 00:00:00 2001 From: Codex H100 Validation Date: Fri, 21 Aug 2026 15:46:37 +0000 Subject: [PATCH 31/33] feat(integrations): route Megatron and vLLM through ablation plan --- rl_engine/integrations/ablation.py | 130 ++++++++++++++++++ rl_engine/integrations/megatron.py | 39 ++++++ rl_engine/integrations/runtime.py | 103 ++++++++++++++ rl_engine/integrations/vllm.py | 39 ++++++ tests/test_framework_operator_integrations.py | 69 ++++++++++ 5 files changed, 380 insertions(+) create mode 100644 rl_engine/integrations/ablation.py create mode 100644 rl_engine/integrations/megatron.py create mode 100644 rl_engine/integrations/runtime.py create mode 100644 rl_engine/integrations/vllm.py create mode 100644 tests/test_framework_operator_integrations.py diff --git a/rl_engine/integrations/ablation.py b/rl_engine/integrations/ablation.py new file mode 100644 index 00000000..a2ccdb82 --- /dev/null +++ b/rl_engine/integrations/ablation.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Executable production/RL-Kernel routes for the module debug matrix.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from types import MappingProxyType +from typing import Any, Mapping + +from rl_engine.alignment.cross_config.debug_matrix import module_debug_matrix + + +class Implementation(str, Enum): + PRODUCTION = "production" + RL_KERNEL = "rl_kernel" + + +@dataclass(frozen=True) +class OperatorAblationCase: + module: str + case_id: str + training: Implementation + rollout: Implementation + purpose: str + diagnostic_axes: tuple[str, ...] + + def implementation_for(self, target: str) -> Implementation: + normalized = target.strip().lower() + if normalized == "training": + return self.training + if normalized == "rollout": + return self.rollout + raise ValueError("target must be 'training' or 'rollout'") + + def to_dict(self) -> dict[str, Any]: + return { + "module": self.module, + "case_id": self.case_id, + "training_implementation": self.training.value, + "rollout_implementation": self.rollout.value, + "purpose": self.purpose, + "diagnostic_axes": list(self.diagnostic_axes), + } + + +_CASE_DEFINITIONS = ( + ("P/P", Implementation.PRODUCTION, Implementation.PRODUCTION, "native baseline"), + ("R/R", Implementation.RL_KERNEL, Implementation.RL_KERNEL, "RL-Kernel control"), + ("P/R", Implementation.PRODUCTION, Implementation.RL_KERNEL, "rollout-only mismatch"), + ("R/P", Implementation.RL_KERNEL, Implementation.PRODUCTION, "training-only mismatch"), +) + + +def operator_ablation_cases(module: str) -> tuple[OperatorAblationCase, ...]: + normalized = module.strip().lower() + matrix = module_debug_matrix() + try: + axes = tuple(str(axis["id"]) for axis in matrix["modules"][normalized]["axes"]) + except KeyError as exc: + raise ValueError(f"unknown ablation module {module!r}") from exc + return tuple( + OperatorAblationCase(normalized, case_id, training, rollout, purpose, axes) + for case_id, training, rollout, purpose in _CASE_DEFINITIONS + ) + + +def operator_ablation_case(module: str, case_id: str) -> OperatorAblationCase: + normalized = case_id.strip().upper() + for case in operator_ablation_cases(module): + if case.case_id == normalized: + return case + raise ValueError(f"unknown ablation case {case_id!r}") + + +@dataclass(frozen=True) +class IntegrationPlan: + """One independently selectable P/R case for Attention, FFN, and Logp.""" + + cases: Mapping[str, OperatorAblationCase] + + def __post_init__(self) -> None: + normalized = dict(self.cases) + if set(normalized) != {"attention", "ffn", "logp"}: + raise ValueError("integration plan must define attention, ffn, and logp") + for module, case in normalized.items(): + if not isinstance(case, OperatorAblationCase) or case.module != module: + raise ValueError(f"invalid integration case for {module!r}") + object.__setattr__(self, "cases", MappingProxyType(normalized)) + + @classmethod + def from_case_ids( + cls, + *, + attention: str = "P/P", + ffn: str = "P/P", + logp: str = "P/P", + ) -> "IntegrationPlan": + return cls( + { + "attention": operator_ablation_case("attention", attention), + "ffn": operator_ablation_case("ffn", ffn), + "logp": operator_ablation_case("logp", logp), + } + ) + + def implementation_for(self, module: str, target: str) -> Implementation: + try: + case = self.cases[module.strip().lower()] + except KeyError as exc: + raise ValueError(f"unknown integration module {module!r}") from exc + return case.implementation_for(target) + + def to_dict(self) -> dict[str, Any]: + matrix = module_debug_matrix() + return { + "schema_version": matrix["schema_version"], + "cases": {module: case.to_dict() for module, case in self.cases.items()}, + } + + +__all__ = [ + "Implementation", + "IntegrationPlan", + "OperatorAblationCase", + "operator_ablation_case", + "operator_ablation_cases", +] diff --git a/rl_engine/integrations/megatron.py b/rl_engine/integrations/megatron.py new file mode 100644 index 00000000..104bf623 --- /dev/null +++ b/rl_engine/integrations/megatron.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Megatron-side operator boundary owned by RL-Kernel.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +from rl_engine.integrations.ablation import IntegrationPlan +from rl_engine.integrations.runtime import FrameworkOperatorIntegration + + +class MegatronIntegration(FrameworkOperatorIntegration): + def __init__( + self, + plan: IntegrationPlan, + *, + rl_kernel_operators: Mapping[str, Callable[..., Any]], + ) -> None: + super().__init__( + framework="megatron", + target="training", + plan=plan, + rl_kernel_operators=rl_kernel_operators, + ) + + def attention(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("attention", native, *args, **kwargs) + + def ffn(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("ffn", native, *args, **kwargs) + + def logp(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("logp", native, *args, **kwargs) + + +__all__ = ["MegatronIntegration"] diff --git a/rl_engine/integrations/runtime.py b/rl_engine/integrations/runtime.py new file mode 100644 index 00000000..58948cb1 --- /dev/null +++ b/rl_engine/integrations/runtime.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fail-closed operator routing shared by framework integrations.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from threading import Lock +from typing import Any + +from rl_engine.integrations.ablation import Implementation, IntegrationPlan + + +@dataclass(frozen=True) +class OperatorReadback: + framework: str + target: str + module: str + case_id: str + implementation: str + backend_id: str + call_count: int + + def to_dict(self) -> dict[str, Any]: + return self.__dict__.copy() + + +class FrameworkOperatorIntegration: + """Route framework calls without importing or modifying framework packages.""" + + def __init__( + self, + *, + framework: str, + target: str, + plan: IntegrationPlan, + rl_kernel_operators: Mapping[str, Callable[..., Any]], + ) -> None: + self.framework = framework + self.target = target + self.plan = plan + self._rl_kernel_operators = dict(rl_kernel_operators) + self._counts: Counter[str] = Counter() + self._readbacks: dict[str, OperatorReadback] = {} + self._lock = Lock() + + def execute( + self, + module: str, + native: Callable[..., Any], + *args: Any, + **kwargs: Any, + ) -> Any: + normalized = module.strip().lower() + if not callable(native): + raise TypeError("native operator must be callable") + implementation = self.plan.implementation_for(normalized, self.target) + if implementation is Implementation.PRODUCTION: + selected = native + else: + selected = self._rl_kernel_operators.get(normalized) + if selected is None: + raise RuntimeError( + f"{self.framework} {normalized} selected RL-Kernel but no operator was installed" + ) + result = selected(*args, **kwargs) + backend_id = getattr(selected, "backend_id", None) + if not isinstance(backend_id, str) or not backend_id.strip(): + backend_id = ( + f"{self.framework}.production.{normalized}" + if implementation is Implementation.PRODUCTION + else f"rlkernel.{normalized}.unidentified" + ) + with self._lock: + self._counts[normalized] += 1 + case = self.plan.cases[normalized] + self._readbacks[normalized] = OperatorReadback( + framework=self.framework, + target=self.target, + module=normalized, + case_id=case.case_id, + implementation=implementation.value, + backend_id=backend_id, + call_count=self._counts[normalized], + ) + return result + + def readback(self) -> dict[str, Any]: + with self._lock: + return { + "framework": self.framework, + "target": self.target, + "plan": self.plan.to_dict(), + "operators": { + module: readback.to_dict() for module, readback in self._readbacks.items() + }, + } + + +__all__ = ["FrameworkOperatorIntegration", "OperatorReadback"] diff --git a/rl_engine/integrations/vllm.py b/rl_engine/integrations/vllm.py new file mode 100644 index 00000000..1d760850 --- /dev/null +++ b/rl_engine/integrations/vllm.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""vLLM-side operator boundary owned by RL-Kernel.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +from rl_engine.integrations.ablation import IntegrationPlan +from rl_engine.integrations.runtime import FrameworkOperatorIntegration + + +class VllmIntegration(FrameworkOperatorIntegration): + def __init__( + self, + plan: IntegrationPlan, + *, + rl_kernel_operators: Mapping[str, Callable[..., Any]], + ) -> None: + super().__init__( + framework="vllm", + target="rollout", + plan=plan, + rl_kernel_operators=rl_kernel_operators, + ) + + def attention(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("attention", native, *args, **kwargs) + + def ffn(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("ffn", native, *args, **kwargs) + + def logp(self, native: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: + return self.execute("logp", native, *args, **kwargs) + + +__all__ = ["VllmIntegration"] diff --git a/tests/test_framework_operator_integrations.py b/tests/test_framework_operator_integrations.py new file mode 100644 index 00000000..55f63aa8 --- /dev/null +++ b/tests/test_framework_operator_integrations.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import pytest + +from rl_engine.integrations import IntegrationPlan, MegatronIntegration, VllmIntegration + + +class FakeOperator: + def __init__(self, value: str, backend_id: str): + self.value = value + self.backend_id = backend_id + + def __call__(self, payload: str) -> str: + return f"{self.value}:{payload}" + + +def _operators(): + return { + "attention": FakeOperator("rlk-attn", "rlkernel.attention.deterministic.v1"), + "ffn": FakeOperator("rlk-ffn", "rlkernel.ffn.qwen3.deterministic.v1"), + "logp": FakeOperator("rlk-logp", "pytorch-vocab-parallel-logp-ws2"), + } + + +def test_plan_uses_module_matrix_cases_without_cartesian_expansion(): + plan = IntegrationPlan.from_case_ids(attention="P/R", ffn="R/P", logp="R/R") + + assert plan.implementation_for("attention", "training").value == "production" + assert plan.implementation_for("attention", "rollout").value == "rl_kernel" + assert plan.implementation_for("ffn", "training").value == "rl_kernel" + assert plan.implementation_for("ffn", "rollout").value == "production" + assert plan.to_dict()["schema_version"] == "rlkernel.debug_matrix.v1" + + +def test_megatron_and_vllm_route_the_same_plan_on_opposite_sides(): + plan = IntegrationPlan.from_case_ids(attention="P/R", ffn="R/P", logp="R/R") + megatron = MegatronIntegration(plan, rl_kernel_operators=_operators()) + vllm = VllmIntegration(plan, rl_kernel_operators=_operators()) + native = FakeOperator("native", "production.backend") + + assert megatron.attention(native, "x") == "native:x" + assert vllm.attention(native, "x") == "rlk-attn:x" + assert megatron.ffn(native, "x") == "rlk-ffn:x" + assert vllm.ffn(native, "x") == "native:x" + assert megatron.logp(native, "x") == "rlk-logp:x" + assert vllm.logp(native, "x") == "rlk-logp:x" + + assert megatron.readback()["operators"]["attention"]["backend_id"] == "production.backend" + assert vllm.readback()["operators"]["attention"]["backend_id"] == ( + "rlkernel.attention.deterministic.v1" + ) + + +def test_rl_kernel_selection_fails_closed_when_operator_is_missing(): + plan = IntegrationPlan.from_case_ids(attention="R/R") + integration = MegatronIntegration(plan, rl_kernel_operators={}) + + with pytest.raises(RuntimeError, match="no operator was installed"): + integration.attention(lambda value: value, "x") + + +def test_integration_modules_do_not_import_optional_frameworks(): + import sys + + assert "megatron" not in sys.modules + assert "vllm" not in sys.modules From ca49fd596b18f53a241c6cc962c48d6dfbadd6e6 Mon Sep 17 00:00:00 2001 From: Codex H100 Validation Date: Fri, 21 Aug 2026 15:51:42 +0000 Subject: [PATCH 32/33] feat(integrations): export framework integration API --- rl_engine/integrations/__init__.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/rl_engine/integrations/__init__.py b/rl_engine/integrations/__init__.py index 29bcc68f..90d45a5d 100644 --- a/rl_engine/integrations/__init__.py +++ b/rl_engine/integrations/__init__.py @@ -1,5 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Optional framework adapters owned by RL-Kernel.""" +"""Framework integration entry points owned by RL-Kernel.""" + +from rl_engine.integrations.ablation import ( + Implementation, + IntegrationPlan, + OperatorAblationCase, + operator_ablation_case, + operator_ablation_cases, +) +from rl_engine.integrations.megatron import MegatronIntegration +from rl_engine.integrations.vllm import VllmIntegration + +__all__ = [ + "Implementation", + "IntegrationPlan", + "MegatronIntegration", + "OperatorAblationCase", + "VllmIntegration", + "operator_ablation_case", + "operator_ablation_cases", +] From 27178679971811e605df9d79590c177b916ce580 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sat, 22 Aug 2026 00:19:33 +0800 Subject: [PATCH 33/33] fix(integration): satisfy test branch CI contracts --- examples/vime_qwen3_8b_tp2_cp2/run.py | 24 +- rl_engine/alignment/cross_config/__main__.py | 376 +++++++++--------- .../alignment/cross_config/drift_report.py | 204 +++------- .../alignment/cross_config/drift_viewer.py | 58 +-- rl_engine/integrations/__init__.py | 1 - rl_engine/integrations/runtime.py | 9 +- rl_engine/integrations/vime/__init__.py | 1 - rl_engine/integrations/vime/logp.py | 33 +- rl_engine/kernels/gtest/tolerance.py | 6 +- .../ops/pytorch/loss/vocab_parallel_logp.py | 8 +- tests/test_cross_config_drift_report.py | 20 +- tests/test_vime_qwen3_example.py | 26 +- 12 files changed, 335 insertions(+), 431 deletions(-) diff --git a/examples/vime_qwen3_8b_tp2_cp2/run.py b/examples/vime_qwen3_8b_tp2_cp2/run.py index 4823dd58..72e3744b 100644 --- a/examples/vime_qwen3_8b_tp2_cp2/run.py +++ b/examples/vime_qwen3_8b_tp2_cp2/run.py @@ -16,12 +16,10 @@ import json import os import subprocess -import sys from datetime import datetime, timezone from pathlib import Path from typing import Any, Mapping - DEFAULT_CONFIG = Path(__file__).with_name("qwen3_8b_tp2_cp2.json") PROVIDER_MARKER = "Selected-logprob provider active" FALLBACK_MARKERS = ("using native path", "fallback=True", "fallback=true") @@ -44,7 +42,11 @@ def validate_config(config: Mapping[str, Any]) -> None: training = config.get("training") rollout = config.get("rollout") provider = config.get("selected_logprob_provider") - if not isinstance(training, Mapping) or not isinstance(rollout, Mapping) or not isinstance(provider, Mapping): + if ( + not isinstance(training, Mapping) + or not isinstance(rollout, Mapping) + or not isinstance(provider, Mapping) + ): raise ValueError("training, rollout, and selected_logprob_provider sections are required") expected = { "tensor_model_parallel_size": 2, @@ -163,7 +165,9 @@ def build_report( provider_active = PROVIDER_MARKER in log_text fallback_observed = any(marker in log_text for marker in FALLBACK_MARKERS) strict_provider_passed = status == "passed" and provider_active and not fallback_observed - effective_status = "passed" if strict_provider_passed else ("failed" if status == "passed" else status) + effective_status = ( + "passed" if strict_provider_passed else ("failed" if status == "passed" else status) + ) attention_status = _operator_evidence_status(runtime_evidence, "attention") ffn_status = _operator_evidence_status(runtime_evidence, "ffn") return { @@ -192,7 +196,9 @@ def build_report( "returncode": returncode, "artifacts": { "log": None if log_path is None else str(log_path), - "runtime_evidence": None if runtime_evidence_path is None else str(runtime_evidence_path), + "runtime_evidence": ( + None if runtime_evidence_path is None else str(runtime_evidence_path) + ), }, "runtime_evidence": None if runtime_evidence is None else dict(runtime_evidence), "revisions": { @@ -206,7 +212,9 @@ def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG) parser.add_argument("--vime-root", type=Path, default=Path(os.environ.get("VIME_ROOT", "."))) - parser.add_argument("--rl-kernel-root", type=Path, default=Path(os.environ.get("RL_KERNEL_ROOT", "."))) + parser.add_argument( + "--rl-kernel-root", type=Path, default=Path(os.environ.get("RL_KERNEL_ROOT", ".")) + ) parser.add_argument("--output", type=Path, default=Path("qwen3_8b_tp2_cp2.validation.json")) parser.add_argument( "--runtime-evidence", @@ -234,7 +242,9 @@ def main(argv: list[str] | None = None) -> int: log_path = args.output.with_suffix(".log") env = build_environment(vime_root, rl_kernel_root) with log_path.open("w", encoding="utf-8") as log_handle: - process = subprocess.run(command, cwd=vime_root, env=env, stdout=log_handle, stderr=subprocess.STDOUT) + process = subprocess.run( + command, cwd=vime_root, env=env, stdout=log_handle, stderr=subprocess.STDOUT + ) returncode = process.returncode log_text = log_path.read_text(encoding="utf-8", errors="replace") status = "passed" if returncode == 0 else "failed" diff --git a/rl_engine/alignment/cross_config/__main__.py b/rl_engine/alignment/cross_config/__main__.py index 058173a4..611dddd2 100644 --- a/rl_engine/alignment/cross_config/__main__.py +++ b/rl_engine/alignment/cross_config/__main__.py @@ -1,199 +1,195 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -"""Command-line interface for cross-configuration experiments.""" - -from __future__ import annotations - -import argparse -import json -import sys -from pathlib import Path -from typing import Any, Optional, Sequence - -from rl_engine.alignment.cross_config.artifacts import ArtifactStore -from rl_engine.alignment.cross_config.config import ExperimentConfig, load_config -from rl_engine.alignment.cross_config.execution_plan import build_execution_plan - - -def build_parser() -> argparse.ArgumentParser: - parser = argparse.ArgumentParser(description=__doc__) - commands = parser.add_subparsers(dest="command", required=True) - +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Command-line interface for cross-configuration experiments.""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any, Optional, Sequence + +from rl_engine.alignment.cross_config.artifacts import ArtifactStore +from rl_engine.alignment.cross_config.config import ExperimentConfig, load_config +from rl_engine.alignment.cross_config.execution_plan import build_execution_plan + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + plan = commands.add_parser("plan", help="validate and persist a plan without execution") - _add_common_arguments(plan) - + _add_common_arguments(plan) + run = commands.add_parser("run", help="execute a plan with an explicit runtime adapter") - _add_common_arguments(run) - run.add_argument( - "--runtime", - required=True, - choices=("cpu-smoke",), - help="Runtime adapter; only the temporary CPU smoke adapter ships in V1", - ) - run.add_argument( - "--allow-smoke-operators", - action="store_true", - help="Explicitly authorize temporary smoke-only operator backends", - ) - run.add_argument( - "--timeout-seconds", - type=float, - default=30.0, - help="Per paired-scoring attempt deadline", - ) - run.add_argument( - "--no-resume", - action="store_true", - help="Create new attempts even when matching COMPLETE artifacts exist", - ) - - report = commands.add_parser( - "report", - help="render a sealed attempt as an offline drift report", - ) - report.add_argument( - "attempt_dir", type=Path, help="Completed cross-config attempt directory" - ) - report.add_argument( - "--output", - "-o", - type=Path, - required=True, + _add_common_arguments(run) + run.add_argument( + "--runtime", + required=True, + choices=("cpu-smoke",), + help="Runtime adapter; only the temporary CPU smoke adapter ships in V1", + ) + run.add_argument( + "--allow-smoke-operators", + action="store_true", + help="Explicitly authorize temporary smoke-only operator backends", + ) + run.add_argument( + "--timeout-seconds", + type=float, + default=30.0, + help="Per paired-scoring attempt deadline", + ) + run.add_argument( + "--no-resume", + action="store_true", + help="Create new attempts even when matching COMPLETE artifacts exist", + ) + + report = commands.add_parser( + "report", + help="render a sealed attempt as an offline drift report", + ) + report.add_argument("attempt_dir", type=Path, help="Completed cross-config attempt directory") + report.add_argument( + "--output", + "-o", + type=Path, + required=True, help="Output .rlk-drift bundle, .png/.jpg image, or .json trace", - ) - report.add_argument("--title", default=None, help="Optional report title") - report.add_argument( - "--no-preview", action="store_true", help="Omit PNG from a .rlk-drift bundle" - ) - return parser - - -def _add_common_arguments(parser: argparse.ArgumentParser) -> None: - parser.add_argument("config", type=Path, help="Versioned experiment JSON") - parser.add_argument( - "--output-root", - type=Path, - default=Path("runs"), - help="Append-only artifact root (default: runs)", - ) - - -def main(argv: Optional[Sequence[str]] = None) -> int: - args = build_parser().parse_args(argv) - try: - if args.command == "report": - summary = _report(args) - else: - config = load_config(args.config) - if args.command == "plan": - summary = record_plan(config, args.output_root) - else: - summary = _run(config, args) - except Exception as exc: - summary = { - "schema_version": "cross_config.cli_summary.v1", - "status": "error", - "error_type": f"{type(exc).__module__}.{type(exc).__qualname__}", - "error": str(exc), - } - print(json.dumps(summary, sort_keys=True)) + ) + report.add_argument("--title", default=None, help="Optional report title") + report.add_argument( + "--no-preview", action="store_true", help="Omit PNG from a .rlk-drift bundle" + ) + return parser + + +def _add_common_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("config", type=Path, help="Versioned experiment JSON") + parser.add_argument( + "--output-root", + type=Path, + default=Path("runs"), + help="Append-only artifact root (default: runs)", + ) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = build_parser().parse_args(argv) + try: + if args.command == "report": + summary = _report(args) + else: + config = load_config(args.config) + if args.command == "plan": + summary = record_plan(config, args.output_root) + else: + summary = _run(config, args) + except Exception as exc: + summary = { + "schema_version": "cross_config.cli_summary.v1", + "status": "error", + "error_type": f"{type(exc).__module__}.{type(exc).__qualname__}", + "error": str(exc), + } + print(json.dumps(summary, sort_keys=True)) print(f"cross-configuration error: {type(exc).__name__}: {exc}", file=sys.stderr) - return 2 - - print(json.dumps(summary, sort_keys=True)) - if args.command == "plan": - print( - f"planned {summary['planned_case_count']} cases; no runtime was created", - file=sys.stderr, - ) - return 0 - if args.command == "report": - print( - f"drift report: {summary['status']} ({summary['output']})", - file=sys.stderr, - ) - return 0 - print( - f"CPU smoke: {summary['status']} ({len(summary['cases'])} cases)", - file=sys.stderr, - ) - for case in summary["cases"]: - print( - f" {case['case_id']}: {case['status']}; actual backends " - f"rollout={case['rollout_backend']}, training={case['training_backend']}; " - f"worst sample/token={case['worst_token_index']}; " - f"mismatches={case['mismatch_count']}; resumed={case['resumed']}", - file=sys.stderr, - ) - return 0 if summary["status"] == "pass" else 1 - - -def record_plan(config: ExperimentConfig, output_root: Path) -> dict[str, Any]: - plan = build_execution_plan(config) - store = ArtifactStore(output_root) - experiment_dir = store.initialize_experiment( - config.definition.experiment_id, - experiment=plan.experiment, - plan=plan.rows(), - ) - return { - "schema_version": "cross_config.cli_summary.v1", - "status": "planned", - "experiment_id": config.definition.experiment_id, - "scenario_id": config.definition.scenario_id, - "planned_case_count": len(plan.entries), - "planning_issues": [issue.to_dict() for issue in plan.issues], - "artifact_dir": str(experiment_dir), - } - - -def _run(config: ExperimentConfig, args: argparse.Namespace) -> dict[str, Any]: - if args.runtime != "cpu-smoke": # pragma: no cover - argparse owns choices - raise ValueError(f"unsupported runtime {args.runtime!r}") - from rl_engine.alignment.testing.cpu_cross_config import run_cpu_experiment - - return run_cpu_experiment( - config, - output_root=args.output_root, - allow_smoke_operators=args.allow_smoke_operators, - timeout_seconds=args.timeout_seconds, - resume=not args.no_resume, - ) - - -def _report(args: argparse.Namespace) -> dict[str, Any]: - from rl_engine.alignment.cross_config.drift_report import ( + return 2 + + print(json.dumps(summary, sort_keys=True)) + if args.command == "plan": + print( + f"planned {summary['planned_case_count']} cases; no runtime was created", + file=sys.stderr, + ) + return 0 + if args.command == "report": + print( + f"drift report: {summary['status']} ({summary['output']})", + file=sys.stderr, + ) + return 0 + print( + f"CPU smoke: {summary['status']} ({len(summary['cases'])} cases)", + file=sys.stderr, + ) + for case in summary["cases"]: + print( + f" {case['case_id']}: {case['status']}; actual backends " + f"rollout={case['rollout_backend']}, training={case['training_backend']}; " + f"worst sample/token={case['worst_token_index']}; " + f"mismatches={case['mismatch_count']}; resumed={case['resumed']}", + file=sys.stderr, + ) + return 0 if summary["status"] == "pass" else 1 + + +def record_plan(config: ExperimentConfig, output_root: Path) -> dict[str, Any]: + plan = build_execution_plan(config) + store = ArtifactStore(output_root) + experiment_dir = store.initialize_experiment( + config.definition.experiment_id, + experiment=plan.experiment, + plan=plan.rows(), + ) + return { + "schema_version": "cross_config.cli_summary.v1", + "status": "planned", + "experiment_id": config.definition.experiment_id, + "scenario_id": config.definition.scenario_id, + "planned_case_count": len(plan.entries), + "planning_issues": [issue.to_dict() for issue in plan.issues], + "artifact_dir": str(experiment_dir), + } + + +def _run(config: ExperimentConfig, args: argparse.Namespace) -> dict[str, Any]: + if args.runtime != "cpu-smoke": # pragma: no cover - argparse owns choices + raise ValueError(f"unsupported runtime {args.runtime!r}") + from rl_engine.alignment.testing.cpu_cross_config import run_cpu_experiment + + return run_cpu_experiment( + config, + output_root=args.output_root, + allow_smoke_operators=args.allow_smoke_operators, + timeout_seconds=args.timeout_seconds, + resume=not args.no_resume, + ) + + +def _report(args: argparse.Namespace) -> dict[str, Any]: + from rl_engine.alignment.cross_config.drift_report import ( build_cross_config_attempt_report, write_drift_bundle, write_drift_report_image, - write_drift_trace, - ) - - report = build_cross_config_attempt_report(args.attempt_dir, title=args.title) - suffix = args.output.suffix.lower() - if suffix == ".rlk-drift": - output = write_drift_bundle( - report, args.output, include_preview=not args.no_preview - ) - artifact_type = "bundle" - elif suffix in {".png", ".jpg", ".jpeg"}: - output = write_drift_report_image(report, args.output) - artifact_type = "image" - elif suffix == ".json": - output = write_drift_trace(report, args.output) - artifact_type = "trace" + write_drift_trace, + ) + + report = build_cross_config_attempt_report(args.attempt_dir, title=args.title) + suffix = args.output.suffix.lower() + if suffix == ".rlk-drift": + output = write_drift_bundle(report, args.output, include_preview=not args.no_preview) + artifact_type = "bundle" + elif suffix in {".png", ".jpg", ".jpeg"}: + output = write_drift_report_image(report, args.output) + artifact_type = "image" + elif suffix == ".json": + output = write_drift_trace(report, args.output) + artifact_type = "trace" else: raise ValueError("report output must use .rlk-drift, .png/.jpg, or .json") - return { - "schema_version": "cross_config.drift_report_summary.v1", - "status": report["status"], - "artifact_type": artifact_type, - "output": str(output), - "case_id": report["axes"].get("case_id"), - "attempt_id": report["axes"].get("attempt_id"), - } - - -if __name__ == "__main__": # pragma: no cover - raise SystemExit(main()) + return { + "schema_version": "cross_config.drift_report_summary.v1", + "status": report["status"], + "artifact_type": artifact_type, + "output": str(output), + "case_id": report["axes"].get("case_id"), + "attempt_id": report["axes"].get("attempt_id"), + } + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/rl_engine/alignment/cross_config/drift_report.py b/rl_engine/alignment/cross_config/drift_report.py index e8771239..7f11227a 100644 --- a/rl_engine/alignment/cross_config/drift_report.py +++ b/rl_engine/alignment/cross_config/drift_report.py @@ -1,3 +1,4 @@ +# flake8: noqa: E501 """Build self-contained profiler-style reports from cross-config artifacts. The reporter is deliberately downstream of execution. It validates a sealed @@ -20,7 +21,6 @@ from rl_engine.alignment.cross_config.artifacts import ArtifactStore - REPORT_SCHEMA_VERSION = 2 _STATUS_LABELS = {"pass": "PASS", "warning": "WARN", "failure": "FAIL", "info": "INFO"} _STATUS_COLORS = { @@ -182,46 +182,30 @@ def build_drift_report( cube = _plain_mapping(result_cube) metrics = _plain_mapping(cube.get("metrics")) axes = _plain_mapping(cube.get("axes")) - validation = _plain_mapping( - cube.get("metadata_validation") or manifest.get("validation") - ) + validation = _plain_mapping(cube.get("metadata_validation") or manifest.get("validation")) provenance = _plain_mapping( - runtime_provenance - or cube.get("runtime_provenance") - or manifest.get("runtime_provenance") + runtime_provenance or cube.get("runtime_provenance") or manifest.get("runtime_provenance") ) samples = [ - _plain_mapping(item) - for item in manifest.get("samples", []) - if isinstance(item, Mapping) + _plain_mapping(item) for item in manifest.get("samples", []) if isinstance(item, Mapping) ] warnings = [ - _plain_mapping(item) - for item in validation.get("warnings", []) - if isinstance(item, Mapping) + _plain_mapping(item) for item in validation.get("warnings", []) if isinstance(item, Mapping) ] failures = [ - _plain_mapping(item) - for item in validation.get("failures", []) - if isinstance(item, Mapping) + _plain_mapping(item) for item in validation.get("failures", []) if isinstance(item, Mapping) ] max_abs_dlogp = _number(metrics.get("max_abs_dlogp")) warning_count = _number(metrics.get("warning_count"), default=0.0) or 0.0 metadata_warning_count = ( - _number(metrics.get("metadata_warning_count"), default=float(len(warnings))) - or 0.0 + _number(metrics.get("metadata_warning_count"), default=float(len(warnings))) or 0.0 ) metadata_failure_count = ( - _number(metrics.get("metadata_failure_count"), default=float(len(failures))) - or 0.0 - ) - runtime_fallback = bool( - metrics.get("runtime_fallback") or provenance.get("fallback") - ) - strict_failure = bool( - metrics.get("runtime_strict_failure") or provenance.get("strict_failure") + _number(metrics.get("metadata_failure_count"), default=float(len(failures))) or 0.0 ) + runtime_fallback = bool(metrics.get("runtime_fallback") or provenance.get("fallback")) + strict_failure = bool(metrics.get("runtime_strict_failure") or provenance.get("strict_failure")) comparison_passed = metrics.get("passed") is True if strict_failure or metadata_failure_count > 0: @@ -236,16 +220,12 @@ def build_drift_report( status = "pass" span = max(1.0, float(len(samples))) - has_timestamps = any( - _number(sample.get("start_ts")) is not None for sample in samples - ) + has_timestamps = any(_number(sample.get("start_ts")) is not None for sample in samples) timeline_mode = "timestamp" if has_timestamps else "ordinal_diagnostic" events: list[dict[str, Any]] = [] train_status = ( - "failure" - if status == "failure" - else "warning" if status == "warning" else "pass" + "failure" if status == "failure" else "warning" if status == "warning" else "pass" ) events.append( { @@ -296,9 +276,7 @@ def build_drift_report( provenance.get("requested_backend"), "unknown backend", ) - operator_status = ( - "failure" if strict_failure else "warning" if runtime_fallback else "pass" - ) + operator_status = "failure" if strict_failure else "warning" if runtime_fallback else "pass" events.append( { "id": "operator-backend", @@ -315,8 +293,7 @@ def build_drift_report( worst_token = _plain_mapping(cube.get("worst_token")) if worst_token: marker_position = ( - _number(worst_token.get("sample_position"), default=max(0.0, span - 0.5)) - or 0.0 + _number(worst_token.get("sample_position"), default=max(0.0, span - 0.5)) or 0.0 ) events.append( { @@ -445,9 +422,7 @@ def trace_time(value: Any) -> float: tid = lane_ids.get(event_lane, lane_ids["Drift markers"]) start = _number(event.get("start"), default=0.0) or 0.0 end = _number(event.get("end"), default=start) or start - details = ( - event.get("details") if isinstance(event.get("details"), Mapping) else {} - ) + details = event.get("details") if isinstance(event.get("details"), Mapping) else {} args = { "event_id": str(event.get("id", "")), "status": str(event.get("status", status)), @@ -455,9 +430,7 @@ def trace_time(value: Any) -> float: "timeline_note": normalized.get("timeline_note", ""), "details": details, } - color = event_colors.get( - str(event.get("status", "info")), "thread_state_running" - ) + color = event_colors.get(str(event.get("status", "info")), "thread_state_running") if event.get("kind") == "marker": trace_events.append( { @@ -487,11 +460,7 @@ def trace_time(value: Any) -> float: } ) - metrics = ( - normalized.get("metrics") - if isinstance(normalized.get("metrics"), Mapping) - else {} - ) + metrics = _plain_mapping(normalized.get("metrics")) max_abs_dlogp = _number(metrics.get("max_abs_dlogp")) if max_abs_dlogp is not None: trace_events.append( @@ -510,9 +479,7 @@ def trace_time(value: Any) -> float: "traceEvents": trace_events, "displayTimeUnit": "ms", "metadata": { - "report_title": normalized.get( - "title", "RL-Kernel cross-config drift report" - ), + "report_title": normalized.get("title", "RL-Kernel cross-config drift report"), "status": status, "timeline_mode": timeline_mode, "timeline_note": normalized.get("timeline_note", ""), @@ -553,9 +520,7 @@ def write_drift_bundle( manifest = { "format": "rl_kernel.cross_config_drift", "bundle_version": 1, - "report_schema_version": normalized.get( - "schema_version", REPORT_SCHEMA_VERSION - ), + "report_schema_version": normalized.get("schema_version", REPORT_SCHEMA_VERSION), "files": [ "manifest.json", "report.json", @@ -567,12 +532,8 @@ def write_drift_bundle( "status": normalized.get("status", "info"), } with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as archive: - archive.writestr( - "manifest.json", json.dumps(manifest, ensure_ascii=True, indent=2) - ) - archive.writestr( - "report.json", json.dumps(normalized, ensure_ascii=True, indent=2) - ) + archive.writestr("manifest.json", json.dumps(manifest, ensure_ascii=True, indent=2)) + archive.writestr("report.json", json.dumps(normalized, ensure_ascii=True, indent=2)) archive.writestr("trace.json", json.dumps(trace, ensure_ascii=True, indent=2)) if include_preview: preview = render_drift_report_image(normalized) @@ -593,11 +554,7 @@ def load_drift_bundle(path: str | Path) -> dict[str, Any]: ) report = json.loads(archive.read("report.json")) trace = json.loads(archive.read("trace.json")) - manifest = ( - json.loads(archive.read("manifest.json")) - if "manifest.json" in names - else {} - ) + manifest = json.loads(archive.read("manifest.json")) if "manifest.json" in names else {} return {"manifest": manifest, "report": report, "trace": trace} @@ -624,9 +581,7 @@ def render_drift_report(report: Mapping[str, Any]) -> str: event_map = {str(event.get("id")): event for event in events} def x(value: float) -> float: - return ( - left + max(0.0, min(timeline_span, value)) / timeline_span * timeline_width - ) + return left + max(0.0, min(timeline_span, value)) / timeline_span * timeline_width svg_parts = [ f'', @@ -652,11 +607,7 @@ def x(value: float) -> float: ) for event in events: - lane_index = ( - lanes.index(str(event.get("lane"))) - if str(event.get("lane")) in lanes - else 0 - ) + lane_index = lanes.index(str(event.get("lane"))) if str(event.get("lane")) in lanes else 0 ypos = top + lane_index * row_height status = str(event.get("status", "info")) color = _STATUS_COLORS.get(status, _STATUS_COLORS["info"]) @@ -694,11 +645,7 @@ def x(value: float) -> float: svg_parts.append("") svg = "".join(svg_parts) - metrics = ( - normalized.get("metrics") - if isinstance(normalized.get("metrics"), Mapping) - else {} - ) + metrics = _plain_mapping(normalized.get("metrics")) metric_cards = [ ("Max |dlogp|", _format_number(metrics.get("max_abs_dlogp"))), ("Active tokens", _format_number(metrics.get("active_token_count"))), @@ -717,14 +664,12 @@ def x(value: float) -> float: empty="No runtime provenance was recorded.", ) validation_html = _render_validation(normalized.get("validation")) - event_json = json.dumps( - event_map, ensure_ascii=True, separators=(",", ":") - ).replace("<", "\\u003c") + event_json = json.dumps(event_map, ensure_ascii=True, separators=(",", ":")).replace( + "<", "\\u003c" + ) status = str(normalized.get("status", "info")) status_color = _STATUS_COLORS.get(status, _STATUS_COLORS["info"]) - title = html.escape( - str(normalized.get("title", "RL-Kernel cross-config drift report")) - ) + title = html.escape(str(normalized.get("title", "RL-Kernel cross-config drift report"))) timeline_note = html.escape(str(normalized.get("timeline_note", ""))) return f""" @@ -791,14 +736,10 @@ def render_drift_report_image(report: Mapping[str, Any], *, width: int = 2400) - ruler/grid, thin event bars, and tabular details below the timeline. """ - from PIL import Image, ImageDraw, ImageFont + from PIL import Image, ImageDraw normalized = _plain_mapping(report) - metrics = ( - normalized.get("metrics") - if isinstance(normalized.get("metrics"), Mapping) - else {} - ) + metrics = _plain_mapping(normalized.get("metrics")) events = [_plain_mapping(event) for event in normalized.get("events", [])] status = str(normalized.get("status", "info")) status_color = _REPORT_IMAGE_COLORS.get(status, _REPORT_IMAGE_COLORS["info"]) @@ -856,8 +797,7 @@ def key_value_rows(values: Any, limit: int = 33) -> list[tuple[str, str]]: write( width - 30, 12, - "static report | schema v" - + str(normalized.get("schema_version", REPORT_SCHEMA_VERSION)), + "static report | schema v" + str(normalized.get("schema_version", REPORT_SCHEMA_VERSION)), tiny, "#b7c0c8", anchor="ra", @@ -872,16 +812,12 @@ def key_value_rows(values: Any, limit: int = 33) -> list[tuple[str, str]]: write(34, 101, "run / audit replay", tiny, muted) write(180, 101, normalized.get("timeline_note", ""), tiny, muted) badge_text = _STATUS_LABELS.get(status, status.upper()) - draw.rectangle( - (width - 188, 61, width - 34, 101), fill=panel_bg, outline=status_color, width=2 - ) + draw.rectangle((width - 188, 61, width - 34, 101), fill=panel_bg, outline=status_color, width=2) write(width - 111, 81, badge_text, label_font, status_color, anchor="mm") # Compact metric strip. These are values, not dashboard cards. strip_y = 127 - draw.rectangle( - (34, strip_y, width - 34, strip_y + 64), fill=panel_bg, outline=line, width=1 - ) + draw.rectangle((34, strip_y, width - 34, strip_y + 64), fill=panel_bg, outline=line, width=1) strip_items = [ ("MAX |DLOGP|", _format_number(metrics.get("max_abs_dlogp")), status_color), ( @@ -916,9 +852,7 @@ def key_value_rows(values: Any, limit: int = 33) -> list[tuple[str, str]]: write(52, timeline_y + 14, "CAPTURED EXECUTION", section_font, text_color) mode = str(normalized.get("timeline_mode", "diagnostic")) mode_color = ( - _REPORT_IMAGE_COLORS["yellow"] - if mode != "timestamp" - else _REPORT_IMAGE_COLORS["blue"] + _REPORT_IMAGE_COLORS["yellow"] if mode != "timestamp" else _REPORT_IMAGE_COLORS["blue"] ) write(width - 52, timeline_y + 18, mode, tiny, mode_color, anchor="ra") @@ -938,9 +872,7 @@ def key_value_rows(values: Any, limit: int = 33) -> list[tuple[str, str]]: chart_top = timeline_y + 91 row_height = 62 chart_bottom = chart_top + row_height * len(lanes) - max_end = max( - (_number(event.get("end"), default=1.0) or 1.0 for event in events), default=1.0 - ) + max_end = max((_number(event.get("end"), default=1.0) or 1.0 for event in events), default=1.0) span = max(1.0, max_end) def x_for(value: float) -> int: @@ -998,9 +930,7 @@ def x_for(value: float) -> int: if event.get("kind") == "marker": mid = x_for(start) draw.line((mid, y0 + 6, mid, y0 + row_height - 6), fill=color, width=3) - draw.polygon( - [(mid - 7, y0 + 7), (mid + 7, y0 + 7), (mid, y0 + 17)], fill=color - ) + draw.polygon([(mid - 7, y0 + 7), (mid + 7, y0 + 7), (mid, y0 + 17)], fill=color) marker_label = fit(event.get("label", "marker"), 34) label_x = mid + 14 if mid <= right - 250 else max(left + 10, mid - 240) write(label_x, y0 + 22, marker_label, tiny, color) @@ -1026,11 +956,7 @@ def x_for(value: float) -> int: # Keep the comparison track meaningful even when the dump only contains a # scalar worst-token summary instead of per-token samples. - worst = ( - normalized.get("worst_token") - if isinstance(normalized.get("worst_token"), Mapping) - else {} - ) + worst = _plain_mapping(normalized.get("worst_token")) comparison_y = chart_top + 5 * row_height + 31 draw.line( (left + 18, comparison_y, right - 18, comparison_y), @@ -1091,12 +1017,8 @@ def x_for(value: float) -> int: bar_x = 54 bar_y = summary_y + 91 bar_width = half_width - 110 - rect( - (bar_x, bar_y, bar_x + bar_width, bar_y + 24), _REPORT_IMAGE_COLORS["track"], 4 - ) - fill_width = ( - int(min(1.0, max_value / max(max_value, 1.0)) * bar_width) if max_value else 0 - ) + rect((bar_x, bar_y, bar_x + bar_width, bar_y + 24), _REPORT_IMAGE_COLORS["track"], 4) + fill_width = int(min(1.0, max_value / max(max_value, 1.0)) * bar_width) if max_value else 0 if fill_width: rect((bar_x, bar_y, bar_x + max(8, fill_width), bar_y + 24), status_color, 4) write(bar_x, bar_y + 30, "0", tiny, muted) @@ -1108,11 +1030,7 @@ def x_for(value: float) -> int: muted, anchor="ra", ) - worst = ( - normalized.get("worst_token") - if isinstance(normalized.get("worst_token"), Mapping) - else {} - ) + worst = _plain_mapping(normalized.get("worst_token")) worst_text = ( f"worst token: sample={worst.get('sample_position', '-')} token={worst.get('token_position', '-')} " f"|dlogp|={_format_number(worst.get('abs_dlogp'))}" @@ -1120,20 +1038,14 @@ def x_for(value: float) -> int: write(54, summary_y + 157, fit(worst_text, 92), mono, text_color) right_x = 34 + half_width + half_gap + 20 write(right_x, summary_y + 18, "SELECTED EVENT", section_font) - warning_items = ( - normalized.get("validation") - if isinstance(normalized.get("validation"), Mapping) - else {} - ) + warning_items = _plain_mapping(normalized.get("validation")) warnings = warning_items.get("warnings") or [] failures = warning_items.get("failures") or [] issue_text = "No anomaly recorded." if failures or warnings: first = (failures or warnings)[0] first = first if isinstance(first, Mapping) else {"message": first} - issue_text = ( - f"{str(first.get('code', 'validation'))}: {str(first.get('message', ''))}" - ) + issue_text = f"{str(first.get('code', 'validation'))}: {str(first.get('message', ''))}" elif worst: issue_text = ( f"worst token: sample={worst.get('sample_position', '-')} " @@ -1171,9 +1083,7 @@ def x_for(value: float) -> int: axes_rows = key_value_rows(normalized.get("axes")) provenance_rows = key_value_rows(normalized.get("runtime_provenance")) - def table( - rows: list[tuple[str, str]], x0: int, y0: int, w: int, max_rows: int = 8 - ) -> None: + def table(rows: list[tuple[str, str]], x0: int, y0: int, w: int, max_rows: int = 8) -> None: row_y = y0 for index, (key, value) in enumerate(rows[:max_rows]): if index % 2 == 0: @@ -1186,12 +1096,8 @@ def table( write(x0 + 14, row_y + 8, "not recorded", small, muted) table(axes_rows, 54, table_y + 66, half_width - 40) - table( - provenance_rows, 34 + half_width + half_gap + 20, table_y + 66, half_width - 40 - ) - validation_label = ( - "VALIDATION: PASS" if status == "pass" else f"VALIDATION: {badge_text}" - ) + table(provenance_rows, 34 + half_width + half_gap + 20, table_y + 66, half_width - 40) + validation_label = "VALIDATION: PASS" if status == "pass" else f"VALIDATION: {badge_text}" if failures or warnings: validation_label = f"VALIDATION: {badge_text} | {fit(issue_text, 86)}" write(70, height - 58, validation_label, tiny, status_color) @@ -1232,9 +1138,7 @@ def _worst_token( active_mask: torch.Tensor, ) -> dict[str, Any] | None: if absolute_diff.shape != active_mask.shape: - raise ValueError( - "token_diffs absolute_diff and active_mask must have matching shapes" - ) + raise ValueError("token_diffs absolute_diff and active_mask must have matching shapes") if not bool(active_mask.any()): return None values = absolute_diff.masked_fill(~active_mask, -1.0) @@ -1257,18 +1161,12 @@ def _samples_from_identity( token_ids = logical_identity.get("token_ids") samples: list[dict[str, Any]] = [] for index in range(active_mask.shape[0] if active_mask.ndim else 0): - row = ( - token_ids[index] - if isinstance(token_ids, list) and index < len(token_ids) - else None - ) + row = token_ids[index] if isinstance(token_ids, list) and index < len(token_ids) else None samples.append( { "sample_position": index, "sample_index": index, - "token_count": ( - len(row) if isinstance(row, list) else int(active_mask.shape[-1]) - ), + "token_count": (len(row) if isinstance(row, list) else int(active_mask.shape[-1])), "active_token_count": int(active_mask[index].sum().item()), } ) @@ -1292,9 +1190,7 @@ def _render_validation(validation: Any) -> str: warnings = validation.get("warnings") or [] failures = validation.get("failures") or [] if not warnings and not failures: - return ( - '
PASS: no metadata validation warnings or failures.
' - ) + return '
PASS: no metadata validation warnings or failures.
' rows = [] for kind, items in (("failure", failures), ("warning", warnings)): for item in items: diff --git a/rl_engine/alignment/cross_config/drift_viewer.py b/rl_engine/alignment/cross_config/drift_viewer.py index 75bf7580..80d8bd65 100644 --- a/rl_engine/alignment/cross_config/drift_viewer.py +++ b/rl_engine/alignment/cross_config/drift_viewer.py @@ -10,7 +10,6 @@ from rl_engine.alignment.cross_config.drift_report import load_drift_bundle - _LANES = [ ("Audit", None), ("Training audit", "Training audit"), @@ -63,34 +62,24 @@ def _run_qt(bundle_path: Path) -> int: # pragma: no cover - GUI path status = str(report.get("status", "info")) status_color = _COLORS.get(status, _COLORS["info"]) - class EventItem(QtWidgets.QGraphicsRectItem): + class EventItem(QtWidgets.QGraphicsRectItem): # type: ignore[name-defined] def __init__(self, rect: Any, event: dict[str, Any], color: str) -> None: super().__init__(rect) self.event = event self.setBrush(QtGui.QColor(color)) self.setPen(QtGui.QPen(QtGui.QColor(color))) - self.setFlag( - QtWidgets.QGraphicsItem.GraphicsItemFlag.ItemIsSelectable, True - ) + self.setFlag(QtWidgets.QGraphicsItem.GraphicsItemFlag.ItemIsSelectable, True) self.setToolTip(str(event.get("label", event.get("id", "event")))) - class TimelineView(QtWidgets.QGraphicsView): + class TimelineView(QtWidgets.QGraphicsView): # type: ignore[name-defined] def __init__(self, scene: Any) -> None: super().__init__(scene) self.setRenderHint(QtGui.QPainter.RenderHint.Antialiasing, False) self.setDragMode(QtWidgets.QGraphicsView.DragMode.ScrollHandDrag) - self.setTransformationAnchor( - QtWidgets.QGraphicsView.ViewportAnchor.AnchorUnderMouse - ) - self.setResizeAnchor( - QtWidgets.QGraphicsView.ViewportAnchor.AnchorViewCenter - ) - self.setHorizontalScrollBarPolicy( - QtCore.Qt.ScrollBarPolicy.ScrollBarAlwaysOn - ) - self.setVerticalScrollBarPolicy( - QtCore.Qt.ScrollBarPolicy.ScrollBarAlwaysOff - ) + self.setTransformationAnchor(QtWidgets.QGraphicsView.ViewportAnchor.AnchorUnderMouse) + self.setResizeAnchor(QtWidgets.QGraphicsView.ViewportAnchor.AnchorViewCenter) + self.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarPolicy.ScrollBarAlwaysOn) + self.setVerticalScrollBarPolicy(QtCore.Qt.ScrollBarPolicy.ScrollBarAlwaysOff) self.setStyleSheet("QGraphicsView { border: 0; background: #ffffff; }") def wheelEvent(self, event: Any) -> None: @@ -98,7 +87,7 @@ def wheelEvent(self, event: Any) -> None: self.scale(factor, 1.0) event.accept() - class Window(QtWidgets.QMainWindow): + class Window(QtWidgets.QMainWindow): # type: ignore[name-defined] def __init__(self) -> None: super().__init__() self.setWindowTitle(f"RL-Kernel Cross-Config Drift - {title}") @@ -125,9 +114,7 @@ def __init__(self) -> None: ) self._details = QtWidgets.QPlainTextEdit() self._details.setReadOnly(True) - self._details.setPlaceholderText( - "Select an event to inspect its audit details." - ) + self._details.setPlaceholderText("Select an event to inspect its audit details.") self._details.setStyleSheet( "QPlainTextEdit { background: #ffffff; border-top: 1px solid #c9ced3; " "font-family: Consolas; font-size: 10pt; padding: 8px; }" @@ -148,9 +135,7 @@ def _build_tree(self) -> None: groups[label] = item continue parent = groups.get( - "Audit" - if event_lane in {"Training audit", "Rollout samples"} - else "Execution" + "Audit" if event_lane in {"Training audit", "Rollout samples"} else "Execution" ) if event_lane == "Drift markers": parent = groups.get("Validation") @@ -163,10 +148,7 @@ def _build_tree(self) -> None: self._tree.expandAll() def _x_for(self, value: float) -> float: - return ( - 120.0 - + max(0.0, min(self._span, value)) / self._span * self._scene_width - ) + return 120.0 + max(0.0, min(self._span, value)) / self._span * self._scene_width def _build_scene(self) -> None: font = QtGui.QFont("Segoe UI", 9) @@ -186,11 +168,7 @@ def _build_scene(self) -> None: label.setDefaultTextColor(muted) label.setPos(x + 3, -24) - lane_index = { - name: index - for index, (_, name) in enumerate(_LANES) - if name is not None - } + lane_index = {name: index for index, (_, name) in enumerate(_LANES) if name is not None} for index, (label, event_lane) in enumerate(_LANES): y = index * self._row_height self._scene.addLine( @@ -246,9 +224,7 @@ def _build_layout(self) -> None: "QToolBar { background: #ffffff; border-bottom: 1px solid #c9ced3; }" ) status_label = QtWidgets.QLabel(f" {title} STATUS: {status.upper()} ") - status_label.setStyleSheet( - f"color: {status_color}; font-weight: 600; padding: 5px;" - ) + status_label.setStyleSheet(f"color: {status_color}; font-weight: 600; padding: 5px;") toolbar.addWidget(status_label) fit_button = QtWidgets.QAction("Fit", self) fit_button.triggered.connect( @@ -284,16 +260,12 @@ def _select_lane(self) -> None: def _show_selected(self) -> None: selected = self._scene.selectedItems() - self._details.setPlainText( - _format_event(selected[0].event) if selected else "" - ) + self._details.setPlainText(_format_event(selected[0].event) if selected else "") app = QtWidgets.QApplication.instance() or QtWidgets.QApplication(sys.argv) window = Window() window.show() - window._view.fitInView( - window._scene.sceneRect(), QtCore.Qt.AspectRatioMode.KeepAspectRatio - ) + window._view.fitInView(window._scene.sceneRect(), QtCore.Qt.AspectRatioMode.KeepAspectRatio) return app.exec() diff --git a/rl_engine/integrations/__init__.py b/rl_engine/integrations/__init__.py index 90d45a5d..24eac10c 100644 --- a/rl_engine/integrations/__init__.py +++ b/rl_engine/integrations/__init__.py @@ -22,4 +22,3 @@ "operator_ablation_case", "operator_ablation_cases", ] - diff --git a/rl_engine/integrations/runtime.py b/rl_engine/integrations/runtime.py index 58948cb1..7ede1224 100644 --- a/rl_engine/integrations/runtime.py +++ b/rl_engine/integrations/runtime.py @@ -58,14 +58,17 @@ def execute( if not callable(native): raise TypeError("native operator must be callable") implementation = self.plan.implementation_for(normalized, self.target) + selected: Callable[..., Any] if implementation is Implementation.PRODUCTION: selected = native else: - selected = self._rl_kernel_operators.get(normalized) - if selected is None: + rl_kernel_operator = self._rl_kernel_operators.get(normalized) + if rl_kernel_operator is None: raise RuntimeError( - f"{self.framework} {normalized} selected RL-Kernel but no operator was installed" + f"{self.framework} {normalized} selected RL-Kernel " + "but no operator was installed" ) + selected = rl_kernel_operator result = selected(*args, **kwargs) backend_id = getattr(selected, "backend_id", None) if not isinstance(backend_id, str) or not backend_id.strip(): diff --git a/rl_engine/integrations/vime/__init__.py b/rl_engine/integrations/vime/__init__.py index 44470314..3c35c88f 100644 --- a/rl_engine/integrations/vime/__init__.py +++ b/rl_engine/integrations/vime/__init__.py @@ -6,4 +6,3 @@ from .logp import ProviderResult, SelectedLogprobProviderUnavailable, provider __all__ = ["ProviderResult", "SelectedLogprobProviderUnavailable", "provider"] - diff --git a/rl_engine/integrations/vime/logp.py b/rl_engine/integrations/vime/logp.py index db436d69..ef1b6d19 100644 --- a/rl_engine/integrations/vime/logp.py +++ b/rl_engine/integrations/vime/logp.py @@ -57,14 +57,18 @@ class ProviderResult: def _as_positive_int(value: Any, name: str) -> int: if isinstance(value, bool) or not isinstance(value, int) or value <= 0: - raise SelectedLogprobProviderUnavailable(f"{name} must be a positive integer; got {value!r}") + raise SelectedLogprobProviderUnavailable( + f"{name} must be a positive integer; got {value!r}" + ) return value def _metadata(request: Any) -> Mapping[str, Any]: value = getattr(request, "metadata", None) if not isinstance(value, Mapping): - raise SelectedLogprobProviderUnavailable("request.metadata must provide vocab-parallel metadata") + raise SelectedLogprobProviderUnavailable( + "request.metadata must provide vocab-parallel metadata" + ) return value @@ -108,7 +112,9 @@ def _contract_for_request(request: Any) -> tuple[LogprobContract, int]: targets = _request_tensor(request, "target_ids") metadata = _metadata(request) if logits.ndim != 2 or targets.shape != (logits.shape[0],): - raise SelectedLogprobProviderUnavailable("request must contain local [T, V] logits and aligned [T] targets") + raise SelectedLogprobProviderUnavailable( + "request must contain local [T, V] logits and aligned [T] targets" + ) if logits.dtype not in (torch.bfloat16, torch.float16, torch.float32): raise SelectedLogprobProviderUnavailable(f"unsupported logit dtype {logits.dtype}") if targets.device != logits.device: @@ -117,12 +123,20 @@ def _contract_for_request(request: Any) -> tuple[LogprobContract, int]: cp = getattr(request, "context_parallel", None) cp_world_size = _as_positive_int(getattr(cp, "world_size", None), "context_parallel.world_size") cp_rank = getattr(cp, "rank", None) - if isinstance(cp_rank, bool) or not isinstance(cp_rank, int) or not 0 <= cp_rank < cp_world_size: + if ( + isinstance(cp_rank, bool) + or not isinstance(cp_rank, int) + or not 0 <= cp_rank < cp_world_size + ): raise SelectedLogprobProviderUnavailable( f"context_parallel.rank={cp_rank!r} is invalid for CP={cp_world_size}" ) - if getattr(cp, "layout", None) not in ({"single"} if cp_world_size == 1 else {"zigzag", "allgather"}): - raise SelectedLogprobProviderUnavailable("context_parallel layout does not describe local CP token ownership") + if getattr(cp, "layout", None) not in ( + {"single"} if cp_world_size == 1 else {"zigzag", "allgather"} + ): + raise SelectedLogprobProviderUnavailable( + "context_parallel layout does not describe local CP token ownership" + ) tp_rank, tp_world_size = _tp_coordinates(getattr(request, "tensor_parallel_group", None)) declared_tp_rank = metadata.get("tp_rank") @@ -133,7 +147,8 @@ def _contract_for_request(request: Any) -> tuple[LogprobContract, int]: ) if declared_tp_world_size is not None and declared_tp_world_size != tp_world_size: raise SelectedLogprobProviderUnavailable( - f"metadata tp_world_size={declared_tp_world_size} disagrees with TP group size={tp_world_size}" + f"metadata tp_world_size={declared_tp_world_size} disagrees with " + f"TP group size={tp_world_size}" ) real_vocab_size = _as_positive_int(metadata.get("real_vocab_size"), "real_vocab_size") @@ -144,7 +159,9 @@ def _contract_for_request(request: Any) -> tuple[LogprobContract, int]: f"{logits.shape[1]} * {tp_world_size} != {padded_vocab_size}" ) if real_vocab_size > padded_vocab_size: - raise SelectedLogprobProviderUnavailable("real_vocab_size must not exceed padded_vocab_size") + raise SelectedLogprobProviderUnavailable( + "real_vocab_size must not exceed padded_vocab_size" + ) bounds = tuple( (rank * logits.shape[1], (rank + 1) * logits.shape[1]) for rank in range(tp_world_size) diff --git a/rl_engine/kernels/gtest/tolerance.py b/rl_engine/kernels/gtest/tolerance.py index d4537dc1..4fb5bcbf 100644 --- a/rl_engine/kernels/gtest/tolerance.py +++ b/rl_engine/kernels/gtest/tolerance.py @@ -1014,10 +1014,14 @@ def resolve_logprob_threshold(dtype: Any) -> float: raise ValueError(f"invalid WS1 logprob threshold for dtype {dtype_name!r}") return threshold + def tolerance_contract_fingerprint() -> str: - canonical = json.dumps(load_contract(), ensure_ascii=True, separators=(",", ":"), sort_keys=True).encode("utf-8") + canonical = json.dumps( + load_contract(), ensure_ascii=True, separators=(",", ":"), sort_keys=True + ).encode("utf-8") return hashlib.sha256(canonical).hexdigest() + # Private compatibility alias for callers outside this package that have not # migrated to the public normalizer yet. _dtype_name = normalize_dtype_name diff --git a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py index 92e2e25c..3745382a 100644 --- a/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/vocab_parallel_logp.py @@ -372,7 +372,9 @@ def forward( finite_row = torch.isfinite(lse) lse_safe = torch.where(finite_row, lse, torch.zeros_like(lse)) probabilities = (z_masked - lse_safe.unsqueeze(1)).exp() - probabilities = torch.where(finite_row.unsqueeze(1), probabilities, torch.zeros_like(probabilities)) + probabilities = torch.where( + finite_row.unsqueeze(1), probabilities, torch.zeros_like(probabilities) + ) finite_logits = torch.isfinite(z_masked) log_gap = torch.where( finite_logits, @@ -424,7 +426,9 @@ def backward(ctx, grad_logp, grad_lse, grad_entropy): grad = grad + grad_lse.unsqueeze(1) * p if ctx.with_entropy_grad and grad_entropy is not None: entropy_input = lse_safe.unsqueeze(1) - z_masked - entropy.unsqueeze(1) - entropy_input = torch.where(torch.isfinite(z_masked), entropy_input, torch.zeros_like(entropy_input)) + entropy_input = torch.where( + torch.isfinite(z_masked), entropy_input, torch.zeros_like(entropy_input) + ) grad = grad + grad_entropy.unsqueeze(1) * p * entropy_input if bool(padding_cols.any()): grad = grad.masked_fill(padding_cols.unsqueeze(0), 0.0) diff --git a/tests/test_cross_config_drift_report.py b/tests/test_cross_config_drift_report.py index bb0fe3d7..9996591b 100644 --- a/tests/test_cross_config_drift_report.py +++ b/tests/test_cross_config_drift_report.py @@ -7,6 +7,7 @@ import torch import rl_engine.alignment.cross_config.__main__ as cli_main +from rl_engine.alignment.cross_config.artifacts import ArtifactStore from rl_engine.alignment.cross_config.drift_report import ( build_cross_config_attempt_report, build_drift_report, @@ -19,7 +20,6 @@ write_drift_report_image, write_drift_trace, ) -from rl_engine.alignment.cross_config.artifacts import ArtifactStore def _artifacts(*, with_timestamp: bool = False): @@ -187,9 +187,7 @@ def test_report_prefers_actual_backend_and_timestamp_mode(): } report = build_drift_report(replay_manifest=manifest, result_cube=cube) - operator = next( - event for event in report["events"] if event["id"] == "operator-backend" - ) + operator = next(event for event in report["events"] if event["id"] == "operator-backend") assert report["timeline_mode"] == "timestamp" assert operator["label"] == "native.linear_logp" @@ -199,9 +197,7 @@ def test_report_prefers_actual_backend_and_timestamp_mode(): @pytest.mark.unit def test_rendered_report_is_self_contained_and_escapes_details(tmp_path: Path): manifest, cube = _artifacts() - manifest["validation"]["warnings"] = [ - {"code": "bad<&", "message": "value "} - ] + manifest["validation"]["warnings"] = [{"code": "bad<&", "message": "value "}] cube["metadata_validation"] = manifest["validation"] report = build_drift_report( @@ -315,10 +311,7 @@ def test_attempt_report_reads_only_sealed_cross_config_artifacts(tmp_path: Path) assert report["axes"]["rollout_tp"] == 2 assert report["axes"]["rollout_cp"] == 2 assert report["metrics"]["active_token_count"] == 2 - assert ( - report["runtime_provenance"]["operator_source"] - == "exact_resolution_and_instance" - ) + assert report["runtime_provenance"]["operator_source"] == "exact_resolution_and_instance" @pytest.mark.unit @@ -335,10 +328,7 @@ def test_cli_writes_offline_bundle_from_sealed_attempt(tmp_path: Path, capsys): attempt = _completed_attempt(tmp_path) output = tmp_path / "report.rlk-drift" - assert ( - cli_main.main(["report", str(attempt), "--output", str(output), "--no-preview"]) - == 0 - ) + assert cli_main.main(["report", str(attempt), "--output", str(output), "--no-preview"]) == 0 assert output.is_file() assert "drift report: pass" in capsys.readouterr().err diff --git a/tests/test_vime_qwen3_example.py b/tests/test_vime_qwen3_example.py index 007d4ba1..91d9e2a2 100644 --- a/tests/test_vime_qwen3_example.py +++ b/tests/test_vime_qwen3_example.py @@ -2,7 +2,6 @@ from __future__ import annotations -import json from pathlib import Path import pytest @@ -14,7 +13,6 @@ validate_runtime_evidence, ) - ROOT = Path(__file__).parents[1] CONFIG = ROOT / "examples" / "vime_qwen3_8b_tp2_cp2" / "qwen3_8b_tp2_cp2.json" @@ -67,8 +65,16 @@ def _runtime_evidence(): "schema_version": "rlkernel.operator_runtime_evidence.v1", "operators": { "attention": { - "training": {"implementation_id": "rlk.attn", "backend_id": "rlk", "contract_id": "a"}, - "rollout": {"implementation_id": "rlk.attn", "backend_id": "rlk", "contract_id": "a"}, + "training": { + "implementation_id": "rlk.attn", + "backend_id": "rlk", + "contract_id": "a", + }, + "rollout": { + "implementation_id": "rlk.attn", + "backend_id": "rlk", + "contract_id": "a", + }, "comparison": { "passed": True, "out_max_abs": 0.0, @@ -79,8 +85,16 @@ def _runtime_evidence(): }, }, "ffn": { - "training": {"implementation_id": "rlk.ffn", "backend_id": "rlk", "contract_id": "f"}, - "rollout": {"implementation_id": "rlk.ffn", "backend_id": "rlk", "contract_id": "f"}, + "training": { + "implementation_id": "rlk.ffn", + "backend_id": "rlk", + "contract_id": "f", + }, + "rollout": { + "implementation_id": "rlk.ffn", + "backend_id": "rlk", + "contract_id": "f", + }, "comparison": { "passed": True, "out_max_abs": 0.0,