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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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/41] 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 cb359fe6c7f9d9f5900a9515d0fcf2fd811404d2 Mon Sep 17 00:00:00 2001 From: "yangbosong.ljx" Date: Sun, 19 Jul 2026 14:49:37 +0000 Subject: [PATCH 09/41] feat: add CP-aware attention contract --- docs/design/runtime-dispatch.md | 6 + docs/design/ws2-cp-attention-contract.md | 167 +++++++ docs/operators/attention.md | 11 + rl_engine/kernels/attention_contract.py | 593 +++++++++++++++++++++++ rl_engine/kernels/registry.py | 156 ++++++ tests/test_attention_contract.py | 286 +++++++++++ 6 files changed, 1219 insertions(+) create mode 100644 docs/design/ws2-cp-attention-contract.md create mode 100644 rl_engine/kernels/attention_contract.py create mode 100644 tests/test_attention_contract.py diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index bedf3475..41946a97 100644 --- a/docs/design/runtime-dispatch.md +++ b/docs/design/runtime-dispatch.md @@ -11,6 +11,12 @@ 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 Attention uses the stricter `KernelRegistry.get_attention_op(contract)` path. In addition to +platform priority, this path requires a backend capability descriptor and checks the requested +role, mode, dtype, TP/CP layout, LSE export, deterministic merge, packed varlen, and KV-cache +semantics. Incompatible candidates produce explicit rejection reasons and are never used as an +undeclared fallback. See [WS2 CP-aware Attention contract](ws2-cp-attention-contract.md). + ## LogP Priority | Platform | Priority | diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md new file mode 100644 index 00000000..dfd053ff --- /dev/null +++ b/docs/design/ws2-cp-attention-contract.md @@ -0,0 +1,167 @@ +# 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; +- `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. + +## Qwen3-8B TP=4 CP=4 Example + +```python +from rl_engine.kernels.attention_contract import ( + AttentionContract, + ReductionSpec, + ShardingSpec, +) + +sharding = ShardingSpec( + tp_rank=0, + tp_world_size=4, + cp_rank=0, + cp_world_size=4, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=0, + local_q_heads=8, + local_kv_head_start=0, + local_kv_heads=2, + global_sequence_length=4096, + local_sequence_length=1024, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 1024), +) + +contract = AttentionContract( + role="infer", + mode="prefill", + dtype="bf16", + batch_size=1, + query_sequence_length=1024, + head_dim=128, + causal=True, + causal_offsets=(0,), + sharding=sharding, + reduction=ReductionSpec(), +) +``` + +The TP fields preserve the global Qwen3 GQA mapping: each rank owns 8 of 32 query heads and 2 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, 7) +global_block_token_starts=(0, 3584) +local_block_offsets=(0, 512, 1024) +``` + +This metadata is sufficient for a later implementation to restore logical global order without +using ring arrival order. + +## 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; +- global token positions for every logical cached token; +- a prefix-cache key when prefix caching is enabled. + +Missing decode cache identity is an error at contract construction time. + +## 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. +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=4/CP=4 construction, GQA ownership errors, non-contiguous CP blocks, +packed varlen metadata, decode cache identity, undeclared backend rejection, no incompatible +fallback, and JSON-compatible provenance. diff --git a/docs/operators/attention.md b/docs/operators/attention.md index e3cb5f9b..8bec7d22 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -83,6 +83,17 @@ path. When fused attention kernels land, they are prepended to the priority list op becomes the fallback. The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. +### WS2 CP-aware dispatch + +WS2 distributed callers use a separate contract-aware entry point, +`kernel_registry.get_attention_op(contract)`. It validates explicit TP/CP ownership, fixed +`(out, lse)` merge semantics, causal or packed-sequence offsets, and decode KV-cache identity +before selecting a backend. Legacy `get_op("attention")` behavior remains unchanged. + +Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic +CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as +a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md). + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py new file mode 100644 index 00000000..b2269d04 --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,593 @@ +# 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 +from enum import Enum +from typing import Any, Iterable, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +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" + + +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 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, ...] + prefix_cache_enabled: bool = False + prefix_cache_key: str | None = None + + 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") + 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 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)}" + ) + + 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" + ) + for row_index, row in enumerate(block_table): + 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 not isinstance(self.prefix_cache_enabled, bool): + raise AttentionContractError("prefix_cache_enabled must be a bool") + 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" + ) + + 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 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 + kv_cache: KVCacheSpec | 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") + _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.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 = len(self.sharding.packed_sequence_offsets) - 1 + 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.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), + "prefix_cache_enabled": self.kv_cache.prefix_cache_enabled, + "prefix_cache_key": self.kv_cache.prefix_cache_key, + } + 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, + "kv_cache": kv_cache, + } + + +@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 + 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 field in ( + "exports_attention_lse", + "deterministic_cp_merge", + "supports_packed_varlen", + "supports_kv_cache", + ): + if not isinstance(getattr(self, field), bool): + raise AttentionContractError(f"{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") + 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, + "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", + "ShardingSpec", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 041ed3e1..cab339fd 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -6,6 +6,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.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -135,6 +144,39 @@ def __init__(self): self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() + # These descriptors report what existing WS1 implementations actually + # support. Neither implementation exports attention-domain LSE yet, so + # a strict WS2 request is rejected until the deterministic CP reference + # backend lands instead of silently selecting an incompatible fallback. + 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", + ), + } + self._priority_map = { "cuda": { "logp": [ @@ -314,6 +356,120 @@ 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. + + 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. + """ + + 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() + op_type = "kv_cache_attention" if contract.mode is AttentionMode.DECODE else "attention" + candidates = self._priority_map.get(platform, {}).get(op_type, []) + 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 + + 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(), + } + return AttentionDispatchResult( + op=op, + capability=capability, + provenance=provenance, + ) + + 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, + ) -> str | None: + 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}" + ) + + def _platform(self) -> 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) -> Any | None: + 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]: """Dynamic loading technique: Import modules only when needed and check environment dependencies. diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py new file mode 100644 index 00000000..907b776d --- /dev/null +++ b/tests/test_attention_contract.py @@ -0,0 +1,286 @@ +# 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, + ShardingSpec, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + + +def _sharding( + *, + tp_rank: int = 0, + tp_world_size: int = 4, + cp_rank: int = 0, + cp_world_size: int = 4, + global_sequence_length: int = 4096, + local_sequence_length: int = 1024, + global_block_indices: tuple[int, ...] = (0,), + global_block_token_starts: tuple[int, ...] = (0,), + local_block_offsets: tuple[int, ...] = (0, 1024), + 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, +) -> AttentionContract: + resolved_sharding = sharding or _sharding() + return AttentionContract( + role=role, + mode=mode, + dtype="bf16", + batch_size=batch_size, + query_sequence_length=(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, + ) + + +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=(4,), + cp_world_sizes=(1, 2, 4), + exports_attention_lse=True, + deterministic_cp_merge=True, + supports_packed_varlen=True, + supports_kv_cache=True, + implementation_kind="deterministic", + ) + + +def test_qwen3_tp4_cp4_contract_is_representable_and_serializable(): + contract = _contract() + + assert contract.sharding.local_q_heads == 8 + assert contract.sharding.local_kv_heads == 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()) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("tp_rank", 4, "tp_rank=4"), + ("cp_rank", 4, "cp_rank=4"), + ("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": 4, + "cp_rank": 0, + "cp_world_size": 4, + "global_q_heads": 32, + "global_kv_heads": 8, + "local_q_head_start": 0, + "local_q_heads": 8, + "local_kv_head_start": 0, + "local_kv_heads": 2, + "global_sequence_length": 4096, + "local_sequence_length": 1024, + "global_block_indices": (0,), + "global_block_token_starts": (0,), + "local_block_offsets": (0, 1024), + } + 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, 256, 1024)) + assert sharding.packed_sequence_offsets == (0, 256, 1024) + + +def test_non_contiguous_cp_blocks_have_explicit_global_and_local_offsets(): + sharding = _sharding( + global_block_indices=(0, 7), + global_block_token_starts=(0, 3584), + local_block_offsets=(0, 512, 1024), + ) + + assert sharding.global_block_indices == (0, 7) + assert sharding.global_block_token_starts == (0, 3584) + assert sharding.local_block_offsets == (0, 512, 1024) + + with pytest.raises(AttentionContractError, match="non-overlapping and ordered"): + _sharding( + global_block_indices=(0, 1), + global_block_token_starts=(0, 256), + local_block_offsets=(0, 512, 1024), + ) + + +def test_reduction_requires_fp32_accumulation(): + with pytest.raises(AttentionContractError, match="must be fp32"): + ReductionSpec(acc_dtype="bf16") + + +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_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)), + 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)), + prefix_cache_enabled=True, + ) + + +def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): + registry = KernelRegistry() + + with pytest.raises(RuntimeError) as exc_info: + registry.get_attention_op(_contract()) + + message = str(exc_info.value) + assert "CP=4 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]["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() + 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"]["cp_world_size"] == 4 + json.dumps(result.provenance) + + +def test_requested_stable_backend_id_is_enforced(): + registry = KernelRegistry() + 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, 1024)), + causal_offsets=(0, 0), + ) + + assert capability.incompatibilities(contract) == ("packed varlen layout is unsupported",) From 33119ffebf6e7912fce59358b9f6ef5d9d65346d Mon Sep 17 00:00:00 2001 From: "yangbosong.ljx" Date: Sun, 19 Jul 2026 15:12:26 +0000 Subject: [PATCH 10/41] fix: tighten attention contract validation --- docs/design/ws2-cp-attention-contract.md | 11 +++- rl_engine/kernels/attention_contract.py | 56 ++++++++++++++++++- tests/test_attention_contract.py | 69 ++++++++++++++++++++++++ 3 files changed, 133 insertions(+), 3 deletions(-) diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md index dfd053ff..0bc1be02 100644 --- a/docs/design/ws2-cp-attention-contract.md +++ b/docs/design/ws2-cp-attention-contract.md @@ -35,6 +35,10 @@ Construction performs validation immediately. A structurally valid contract mean 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. + ## Qwen3-8B TP=4 CP=4 Example ```python @@ -122,10 +126,15 @@ 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 when prefix caching is enabled. -Missing decode cache identity is an error at contract construction time. +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. Missing or inconsistent decode cache identity is an error at contract +construction time. ## Contract-Aware Dispatch diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index b2269d04..2607b737 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -276,12 +276,14 @@ class KVCacheSpec: 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 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" ) @@ -289,6 +291,10 @@ def __post_init__(self) -> None: 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" @@ -298,6 +304,20 @@ def __post_init__(self) -> None: "global_token_positions must describe every logical cached token; " f"expected {sum(kv_seq_lens)}, got {len(global_token_positions)}" ) + token_offset = 0 + 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" + ) + token_offset += sequence_length try: block_table = tuple(tuple(row) for row in self.block_table) @@ -309,13 +329,37 @@ def __post_init__(self) -> None: raise AttentionContractError( "block_table must contain one non-empty row per kv_seq_lens entry" ) - for row_index, row in enumerate(block_table): + for row_index, (row, sequence_length) in enumerate( + zip(block_table, kv_seq_lens, strict=True) + ): + 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" + ) + active_blocks.append(block) + + expected_blocks = (sequence_length + page_size - 1) // page_size + if len(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(active_blocks)}" + ) + if len(set(active_blocks)) != len(active_blocks): + raise AttentionContractError( + f"block_table row {row_index} contains duplicate active page ids" + ) if not isinstance(self.prefix_cache_enabled, bool): raise AttentionContractError("prefix_cache_enabled must be a bool") @@ -362,6 +406,13 @@ def __post_init__(self) -> None: raise AttentionContractError("sharding must be a ShardingSpec") if not isinstance(self.reduction, ReductionSpec): raise AttentionContractError("reduction must be a ReductionSpec") + 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: @@ -372,7 +423,7 @@ def __post_init__(self) -> None: 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 = len(self.sharding.packed_sequence_offsets) - 1 + expected_causal_offsets = batch_size offset_owner = "packed sequence" else: expected_causal_offsets = batch_size @@ -441,6 +492,7 @@ def to_dict(self) -> dict[str, Any]: "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, } diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index 907b776d..a07274be 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -209,6 +209,7 @@ def test_decode_requires_complete_kv_cache_identity(): 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", ) @@ -223,10 +224,67 @@ def test_prefix_cache_key_is_required_only_when_prefix_cache_is_enabled(): 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, + ) + + +@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", + ) + + assert cache.block_table == ((3,), (3,)) + + def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): registry = KernelRegistry() @@ -281,6 +339,17 @@ def test_packed_layout_requires_declared_backend_support(): contract = _contract( sharding=_sharding(packed_sequence_offsets=(0, 512, 1024)), causal_offsets=(0, 0), + batch_size=2, ) assert capability.incompatibilities(contract) == ("packed varlen layout is unsupported",) + + +def test_packed_sequence_count_must_match_logical_batch_size(): + sharding = _sharding(packed_sequence_offsets=(0, 512, 1024)) + + 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 From b7ba64b58977cd59bec17b8a4c3f37cddbbb2ce0 Mon Sep 17 00:00:00 2001 From: "yangbosong.ljx" Date: Sun, 19 Jul 2026 15:28:36 +0000 Subject: [PATCH 11/41] fix: constrain shared KV prefix pages --- docs/design/ws2-cp-attention-contract.md | 8 ++- rl_engine/kernels/attention_contract.py | 61 ++++++++++++++++++ tests/test_attention_contract.py | 82 ++++++++++++++++++++++++ 3 files changed, 149 insertions(+), 2 deletions(-) diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md index 0bc1be02..becd90a7 100644 --- a/docs/design/ws2-cp-attention-contract.md +++ b/docs/design/ws2-cp-attention-contract.md @@ -128,12 +128,16 @@ Decode additionally requires `KVCacheSpec` with: - a block/page table; - the physical page size; - global token positions for every logical cached token; -- a prefix-cache key when prefix caching is enabled. +- 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. Missing or inconsistent decode cache identity is an error at contract +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. ## Contract-Aware Dispatch diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 2607b737..cfb4e8f1 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -279,6 +279,7 @@ class KVCacheSpec: 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") @@ -305,6 +306,7 @@ def __post_init__(self) -> None: 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 @@ -317,6 +319,7 @@ def __post_init__(self) -> None: "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 try: @@ -329,6 +332,7 @@ def __post_init__(self) -> None: 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) ): @@ -360,9 +364,13 @@ def __post_init__(self) -> None: raise AttentionContractError( f"block_table row {row_index} contains duplicate active page ids" ) + active_block_rows.append(tuple(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" @@ -371,6 +379,58 @@ def __post_init__(self) -> 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(active_blocks) < shared_prefix_page_count for active_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, (active_blocks, positions) in enumerate( + zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 + ): + if active_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, active_blocks in enumerate(active_block_rows): + for page_id in active_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) @@ -495,6 +555,7 @@ def to_dict(self) -> dict[str, Any]: "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, } return { "semantic_operator": "standard_softmax_attention", diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index a07274be..9d0959c8 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -280,9 +280,91 @@ def test_prefix_pages_may_be_shared_across_sequences(): 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(): From 0f1fd760402126a4aaebe2620f95fb6c8012b751 Mon Sep 17 00:00:00 2001 From: "yangbosong.ljx" Date: Sun, 19 Jul 2026 15:59:08 +0000 Subject: [PATCH 12/41] fix: satisfy attention contract type checks --- .github/workflows/ci.yml | 3 +++ rl_engine/kernels/attention_contract.py | 24 +++++++++++------------- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..26f0575c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -76,6 +76,9 @@ jobs: run: | python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu" + - name: Run WS2 Attention Contract Tests (CPU-safe) + run: python -m pytest tests/test_attention_contract.py -v + docs: runs-on: ubuntu-latest steps: diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index cfb4e8f1..2c68e7dd 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -336,7 +336,7 @@ def __post_init__(self) -> None: for row_index, (row, sequence_length) in enumerate( zip(block_table, kv_seq_lens, strict=True) ): - active_blocks: list[int] = [] + 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: @@ -352,19 +352,19 @@ def __post_init__(self) -> None: "block_table -1 padding must be trailing; " f"row {row_index} contains an active block after padding" ) - active_blocks.append(block) + row_active_blocks.append(block) expected_blocks = (sequence_length + page_size - 1) // page_size - if len(active_blocks) != expected_blocks: + 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(active_blocks)}" + f"row {row_index} expected {expected_blocks}, got {len(row_active_blocks)}" ) - if len(set(active_blocks)) != len(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(active_blocks)) + 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") @@ -386,9 +386,7 @@ def __post_init__(self) -> None: shared_prefix_pages: tuple[int, ...] = () if shared_prefix_page_count > 0: - if any( - len(active_blocks) < shared_prefix_page_count for active_blocks in active_block_rows - ): + 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" ) @@ -400,10 +398,10 @@ def __post_init__(self) -> None: 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, (active_blocks, positions) in enumerate( + for sequence_index, (row_blocks, positions) in enumerate( zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 ): - if active_blocks[:shared_prefix_page_count] != shared_prefix_pages: + 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" @@ -416,8 +414,8 @@ def __post_init__(self) -> None: exclusive_page_owners: dict[int, int] = {} shared_prefix_page_ids = set(shared_prefix_pages) - for sequence_index, active_blocks in enumerate(active_block_rows): - for page_id in active_blocks[shared_prefix_page_count:]: + 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" From 6d826df393ba6d4972d1706cd5ffd704a44e4456 Mon Sep 17 00:00:00 2001 From: "yangbosong.ljx" Date: Mon, 20 Jul 2026 09:31:13 +0000 Subject: [PATCH 13/41] fix: validate attention position metadata --- docs/design/ws2-cp-attention-contract.md | 8 +++++ rl_engine/kernels/attention_contract.py | 21 ++++++++++++- tests/test_attention_contract.py | 38 +++++++++++++++++++++++- 3 files changed, 65 insertions(+), 2 deletions(-) diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md index becd90a7..f9bb6ec1 100644 --- a/docs/design/ws2-cp-attention-contract.md +++ b/docs/design/ws2-cp-attention-contract.md @@ -39,6 +39,10 @@ materialize it. 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=4 CP=4 Example ```python @@ -140,6 +144,10 @@ for copy-on-write before divergent decode. When prefix caching is disabled, no a 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: diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 2c68e7dd..0f14fdea 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -322,6 +322,17 @@ def __post_init__(self) -> None: 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: @@ -458,12 +469,20 @@ def __post_init__(self) -> None: 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") - _positive_int(self.query_sequence_length, "query_sequence_length") + 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 ( + 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: diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index 9d0959c8..2a14d856 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -67,6 +67,7 @@ def _contract( kv_cache: KVCacheSpec | None = None, causal_offsets: tuple[int, ...] = (0,), batch_size: int = 1, + query_sequence_length: int | None = None, ) -> AttentionContract: resolved_sharding = sharding or _sharding() return AttentionContract( @@ -74,7 +75,11 @@ def _contract( mode=mode, dtype="bf16", batch_size=batch_size, - query_sequence_length=(1 if mode == "decode" else resolved_sharding.local_sequence_length), + 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, @@ -200,6 +205,26 @@ def test_causal_attention_requires_explicit_offset(): 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=2048) + + 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") @@ -240,6 +265,17 @@ def test_cache_positions_must_match_kv_sequence_count(): ) +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"): From ed05a09f66f12fa85e2cb7f5dbf68410178ecb41 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Mon, 20 Jul 2026 22:59:48 +0800 Subject: [PATCH 14/41] feat(attention): add deterministic CP reference Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/operators/attention.md | 22 + rl_engine/kernels/gtest/operator_inputs.py | 22 + .../ops/pytorch/attention/cp_attention.py | 548 ++++++++++++++++++ rl_engine/kernels/registry.py | 9 + tests/test_cp_attention.py | 385 ++++++++++++ tests/test_operator_inputs.py | 1 + 6 files changed, 987 insertions(+) create mode 100644 rl_engine/kernels/ops/pytorch/attention/cp_attention.py create mode 100644 tests/test_cp_attention.py diff --git a/docs/operators/attention.md b/docs/operators/attention.md index e3cb5f9b..97dfed29 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -83,6 +83,15 @@ path. When fused attention kernels land, they are prepended to the priority list op becomes the fallback. The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. +`kernel_registry.get_op("cp_attention")` resolves to +`DeterministicCPAttentionReferenceOp`, the WS2 correctness-first context-parallel +reference. It emulates CP prefill and chunked-prefill by splitting logical query +and KV sequence blocks, computing per-block `(out, lse)` partial states, and +merging them in fp32 by global KV block index. This path is not a production +fused backend; it defines the CP/LSE merge behavior that downstream fused paths +must match. Optional per-batch `query_position_offsets` / `key_position_offsets` +cover varlen causal-mask metadata while keeping the dense tensor layout. + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -135,6 +144,7 @@ memory. ```bash python -m pytest tests/test_attention.py -v +python -m pytest tests/test_cp_attention.py -v ``` Covers: `forward_fp32` vs an independent fp32 reference (bitwise), strict-fp32 under hostile @@ -144,15 +154,27 @@ invariance (slice + chunked, bitwise; padding is near-equality only, see below), gradient flow, registry dispatch, and a GPU-only LARGE Qwen3-8B real-shape smoke test. +`tests/test_cp_attention.py` covers the WS2 CP reference: CP=1 vs standard +attention, CP=2 prefill vs CP=1, chunked-prefill replay, global-position causal +masking across CP boundaries, order-independent LSE merge by global block index, +padding/all-masked stability, BF16 final-write behavior, input purity, argument +validation, and registry dispatch. +`make_operator_inputs("cp_attention", ...)` also emits a CP=2 chunked-prefill +synthetic case for local harnesses. + ## Implementation Files - `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` +- `rl_engine/kernels/ops/pytorch/attention/cp_attention.py` - `rl_engine/kernels/registry.py` - `tests/test_attention.py` +- `tests/test_cp_attention.py` ## Known Limitations - PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work). +- `cp_attention` is a PyTorch reference for CP prefill/chunked-prefill semantics, + not a distributed runtime or fused kernel. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - The naive path materializes the full `[B, Hq, Sq, Skv]` scores tensor — no query-chunking, so the LARGE load point is memory-heavy and GPU-only. diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index f124cafb..8d26ac04 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -28,6 +28,7 @@ def make_operator_inputs( "rms_norm": _make_rms_norm_inputs, "matmul": _make_matmul_inputs, "attention": _make_attention_inputs, + "cp_attention": _make_cp_attention_inputs, "logp": _make_logp_inputs, "linear_logp": _make_linear_logp_inputs, "rope": _make_rope_inputs, @@ -50,6 +51,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "rms_norm": f"{batch}x{seq}x{_normalized_dim(args)}", "matmul": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", + "cp_attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}xcp2", "logp": f"{batch}x{seq}x{vocab}", "linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", @@ -107,6 +109,26 @@ def _make_attention_inputs( } +def _make_cp_attention_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + return { + "q": _floating_tensor( + (batch, DEFAULT_N_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0 + ), + "k": _floating_tensor( + (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 1 + ), + "v": _floating_tensor( + (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 2 + ), + "causal": True, + "cp_world_size": 2, + "kv_chunk_size": max(1, seq // 2), + } + + def _make_logp_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: 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..c658b134 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -0,0 +1,548 @@ +# 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.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.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") + + +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 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. + """ + + 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. + """ + + 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(q.dtype if output_dtype is None else 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 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. + """ + + _validate_qkv(q, k, v) + 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) + 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") + + 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 _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 not overlap") + 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) % k.size(1) != 0: + raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") + + +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 + + +CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp + +__all__ = [ + "AttentionPartialState", + "CPAttentionReferenceOp", + "DeterministicCPAttentionReferenceOp", + "merge_attention_partial_states", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 041ed3e1..d463bb5b 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -72,6 +72,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_KV_CACHE_ATTN = ( "rl_engine.kernels.ops.pytorch.attention.kv_cache.NativeKVCacheAttnOp" ) + # WS2 correctness-first context-parallel attention reference. It emulates + # CP prefill/chunked-prefill with fp32 attention-domain LSE merges. + 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 @@ -165,6 +171,7 @@ def __init__(self): ], "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_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], @@ -188,6 +195,7 @@ def __init__(self): OpBackend.TRITON_GENERIC, ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], @@ -206,6 +214,7 @@ def __init__(self): "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py new file mode 100644 index 00000000..dd0f53f7 --- /dev/null +++ b/tests/test_cp_attention.py @@ -0,0 +1,385 @@ +# 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 math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + merge_attention_partial_states, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.registry import kernel_registry + +_N_HEADS = 32 +_N_KV = 8 +_HEAD_DIM = 128 +_ATOL = 3.0e-6 + + +@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_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_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), + ) + + +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_registry_dispatches_cp_attention_reference(): + assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index bb1a2220..9e3eac34 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -36,6 +36,7 @@ def _args(**overrides): "rms_norm", "matmul", "attention", + "cp_attention", "logp", "linear_logp", "rope", From 8a4f9eb190ed19053d996583a034e5aaeb1ab502 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Mon, 27 Jul 2026 22:00:02 +0800 Subject: [PATCH 15/41] fix: align attention contract target with issue scope --- docs/design/ws2-cp-attention-contract.md | 26 ++++----- tests/test_attention_contract.py | 67 +++++++++++++----------- 2 files changed, 48 insertions(+), 45 deletions(-) diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md index f9bb6ec1..3e259fa8 100644 --- a/docs/design/ws2-cp-attention-contract.md +++ b/docs/design/ws2-cp-attention-contract.md @@ -43,7 +43,7 @@ For full `prefill`, `query_sequence_length` equals the local sequence length des `ShardingSpec`. Chunked prefill and decode may use shorter query lengths than their available KV context. -## Qwen3-8B TP=4 CP=4 Example +## Qwen3-8B TP=2 CP=2 Example ```python from rl_engine.kernels.attention_contract import ( @@ -54,20 +54,20 @@ from rl_engine.kernels.attention_contract import ( sharding = ShardingSpec( tp_rank=0, - tp_world_size=4, + tp_world_size=2, cp_rank=0, - cp_world_size=4, + cp_world_size=2, global_q_heads=32, global_kv_heads=8, local_q_head_start=0, - local_q_heads=8, + local_q_heads=16, local_kv_head_start=0, - local_kv_heads=2, + local_kv_heads=4, global_sequence_length=4096, - local_sequence_length=1024, + local_sequence_length=2048, global_block_indices=(0,), global_block_token_starts=(0,), - local_block_offsets=(0, 1024), + local_block_offsets=(0, 2048), ) contract = AttentionContract( @@ -75,7 +75,7 @@ contract = AttentionContract( mode="prefill", dtype="bf16", batch_size=1, - query_sequence_length=1024, + query_sequence_length=2048, head_dim=128, causal=True, causal_offsets=(0,), @@ -84,14 +84,14 @@ contract = AttentionContract( ) ``` -The TP fields preserve the global Qwen3 GQA mapping: each rank owns 8 of 32 query heads and 2 of +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, 7) -global_block_token_starts=(0, 3584) -local_block_offsets=(0, 512, 1024) +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 @@ -183,6 +183,6 @@ Contract and dispatch behavior are covered by: python -m pytest tests/test_attention_contract.py -q ``` -The tests include Qwen3 TP=4/CP=4 construction, GQA ownership errors, non-contiguous CP blocks, +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, and JSON-compatible provenance. diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index 2a14d856..feab4c3a 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -27,14 +27,14 @@ def _sharding( *, tp_rank: int = 0, - tp_world_size: int = 4, + tp_world_size: int = 2, cp_rank: int = 0, - cp_world_size: int = 4, + cp_world_size: int = 2, global_sequence_length: int = 4096, - local_sequence_length: int = 1024, + local_sequence_length: int = 2048, global_block_indices: tuple[int, ...] = (0,), global_block_token_starts: tuple[int, ...] = (0,), - local_block_offsets: tuple[int, ...] = (0, 1024), + local_block_offsets: tuple[int, ...] = (0, 2048), packed_sequence_offsets: tuple[int, ...] | None = None, ) -> ShardingSpec: local_q_heads = 32 // tp_world_size @@ -97,8 +97,8 @@ def _declared_cp_backend() -> AttentionBackendCapability: {AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL, AttentionMode.DECODE} ), dtypes=frozenset({AttentionDType.BF16}), - tp_world_sizes=(4,), - cp_world_sizes=(1, 2, 4), + tp_world_sizes=(2,), + cp_world_sizes=(1, 2), exports_attention_lse=True, deterministic_cp_merge=True, supports_packed_varlen=True, @@ -107,11 +107,13 @@ def _declared_cp_backend() -> AttentionBackendCapability: ) -def test_qwen3_tp4_cp4_contract_is_representable_and_serializable(): +def test_qwen3_tp2_cp2_contract_is_representable_and_serializable(): contract = _contract() - assert contract.sharding.local_q_heads == 8 - assert contract.sharding.local_kv_heads == 2 + 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", @@ -126,8 +128,8 @@ def test_qwen3_tp4_cp4_contract_is_representable_and_serializable(): @pytest.mark.parametrize( ("field", "value", "message"), [ - ("tp_rank", 4, "tp_rank=4"), - ("cp_rank", 4, "cp_rank=4"), + ("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"), ], @@ -135,20 +137,20 @@ def test_qwen3_tp4_cp4_contract_is_representable_and_serializable(): def test_invalid_rank_and_cp_order_metadata_fail_loudly(field, value, message): values = { "tp_rank": 0, - "tp_world_size": 4, + "tp_world_size": 2, "cp_rank": 0, - "cp_world_size": 4, + "cp_world_size": 2, "global_q_heads": 32, "global_kv_heads": 8, "local_q_head_start": 0, - "local_q_heads": 8, + "local_q_heads": 16, "local_kv_head_start": 0, - "local_kv_heads": 2, + "local_kv_heads": 4, "global_sequence_length": 4096, - "local_sequence_length": 1024, + "local_sequence_length": 2048, "global_block_indices": (0,), "global_block_token_starts": (0,), - "local_block_offsets": (0, 1024), + "local_block_offsets": (0, 2048), } values[field] = value @@ -171,26 +173,26 @@ def test_sequence_range_and_packed_offsets_are_validated(): with pytest.raises(AttentionContractError, match="final packed_sequence_offsets"): _sharding(packed_sequence_offsets=(0, 512)) - sharding = _sharding(packed_sequence_offsets=(0, 256, 1024)) - assert sharding.packed_sequence_offsets == (0, 256, 1024) + 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, 7), - global_block_token_starts=(0, 3584), - local_block_offsets=(0, 512, 1024), + global_block_indices=(0, 3), + global_block_token_starts=(0, 3072), + local_block_offsets=(0, 1024, 2048), ) - assert sharding.global_block_indices == (0, 7) - assert sharding.global_block_token_starts == (0, 3584) - assert sharding.local_block_offsets == (0, 512, 1024) + 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, 256), - local_block_offsets=(0, 512, 1024), + global_block_token_starts=(0, 512), + local_block_offsets=(0, 1024, 2048), ) @@ -207,7 +209,7 @@ def test_causal_attention_requires_explicit_offset(): 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=2048) + _contract(mode="prefill", query_sequence_length=1024) chunked = _contract(mode="chunked_prefill", query_sequence_length=512) decode = _contract( @@ -410,7 +412,7 @@ def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): registry.get_attention_op(_contract()) message = str(exc_info.value) - assert "CP=4 is unsupported" in message + 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 @@ -435,7 +437,8 @@ def test_declared_compatible_backend_resolves_and_records_provenance(): 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"]["cp_world_size"] == 4 + assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 + assert result.provenance["contract"]["sharding"]["cp_world_size"] == 2 json.dumps(result.provenance) @@ -455,7 +458,7 @@ def test_requested_stable_backend_id_is_enforced(): 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, 1024)), + sharding=_sharding(packed_sequence_offsets=(0, 512, 2048)), causal_offsets=(0, 0), batch_size=2, ) @@ -464,7 +467,7 @@ def test_packed_layout_requires_declared_backend_support(): def test_packed_sequence_count_must_match_logical_batch_size(): - sharding = _sharding(packed_sequence_offsets=(0, 512, 1024)) + 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) From d39f0a50a19559d97811684f01e671720f64a448 Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Sun, 2 Aug 2026 23:39:25 +0800 Subject: [PATCH 16/41] feat(attention): add rope contract metadata Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/design/ws2-cp-attention-contract.md | 41 ++++++- rl_engine/kernels/attention_contract.py | 142 +++++++++++++++++++++++ tests/test_attention_contract.py | 75 ++++++++++++ 3 files changed, 257 insertions(+), 1 deletion(-) diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md index 3e259fa8..95320101 100644 --- a/docs/design/ws2-cp-attention-contract.md +++ b/docs/design/ws2-cp-attention-contract.md @@ -29,6 +29,7 @@ belong to later work in #235. - `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 @@ -49,6 +50,7 @@ context. from rl_engine.kernels.attention_contract import ( AttentionContract, ReductionSpec, + RoPESpec, ShardingSpec, ) @@ -81,6 +83,18 @@ contract = AttentionContract( 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", + ), ) ``` @@ -97,6 +111,29 @@ 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: @@ -160,6 +197,8 @@ 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 @@ -185,4 +224,4 @@ 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, and JSON-compatible provenance. +fallback, RoPE metadata validation, and JSON-compatible provenance. diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 0f14fdea..7deeddf0 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -55,6 +55,22 @@ class ReductionEngine(str, Enum): IN_OP_REFERENCE = "in_op_reference" +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) @@ -447,6 +463,65 @@ def __post_init__(self) -> None: 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 field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, field) + if values is None: + continue + normalized = _integer_tuple(values, field) + if not normalized or any(value < 0 for value in normalized): + raise AttentionContractError(f"{field} must contain non-negative positions") + object.__setattr__(self, 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.""" @@ -462,6 +537,7 @@ class AttentionContract: sharding: ShardingSpec reduction: ReductionSpec kv_cache: KVCacheSpec | None = None + rope: RoPESpec | None = None export_lse: bool = True def __post_init__(self) -> None: @@ -520,6 +596,27 @@ def __post_init__(self) -> 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 field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, field) + if offsets is not None and len(offsets) != batch_size: + raise AttentionContractError( + f"{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( @@ -574,6 +671,32 @@ def to_dict(self) -> dict[str, Any]: "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, @@ -591,6 +714,7 @@ def to_dict(self) -> dict[str, Any]: "sharding": sharding, "reduction": reduction, "kv_cache": kv_cache, + "rope": rope, } @@ -608,6 +732,8 @@ class AttentionBackendCapability: 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 implementation_kind: str = "production" def __post_init__(self) -> None: @@ -635,6 +761,8 @@ def __post_init__(self) -> None: "deterministic_cp_merge", "supports_packed_varlen", "supports_kv_cache", + "supports_rope_metadata", + "supports_fused_rope_attention", ): if not isinstance(getattr(self, field), bool): raise AttentionContractError(f"{field} must be a bool") @@ -675,6 +803,14 @@ def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: 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") return tuple(reasons) def supports(self, contract: AttentionContract) -> bool: @@ -692,6 +828,8 @@ def to_dict(self) -> dict[str, Any]: "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, "implementation_kind": self.implementation_kind, } @@ -719,5 +857,9 @@ class AttentionDispatchResult: "ReductionEngine", "ReductionOrder", "ReductionSpec", + "RoPECastPoint", + "RoPEFusionBoundary", + "RoPESpec", + "RoPEState", "ShardingSpec", ] diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index feab4c3a..b986fc15 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -19,6 +19,8 @@ AttentionRole, KVCacheSpec, ReductionSpec, + RoPEFusionBoundary, + RoPESpec, ShardingSpec, ) from rl_engine.kernels.registry import KernelRegistry, OpBackend @@ -68,6 +70,7 @@ def _contract( 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( @@ -86,6 +89,7 @@ def _contract( sharding=resolved_sharding, reduction=ReductionSpec(), kv_cache=kv_cache, + rope=rope, ) @@ -125,6 +129,55 @@ def test_qwen3_tp2_cp2_contract_is_representable_and_serializable(): 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"), [ @@ -466,6 +519,28 @@ def test_packed_layout_requires_declared_backend_support(): 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)) From 0480ce81fe2cf8a543f6aa6613a6ec2cf8550ede Mon Sep 17 00:00:00 2001 From: inaniloquentee <3051000145@qq.com> Date: Sun, 2 Aug 2026 23:40:00 +0800 Subject: [PATCH 17/41] docs(attention): clarify rope boundary for cp reference Signed-off-by: inaniloquentee <3051000145@qq.com> --- docs/operators/attention.md | 13 ++++++- .../ops/pytorch/attention/cp_attention.py | 8 ++++ tests/test_cp_attention.py | 38 +++++++++++++++++++ 3 files changed, 57 insertions(+), 2 deletions(-) diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 9aac608e..988fad8c 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -94,6 +94,12 @@ fused backend; it defines the CP/LSE merge behavior that downstream fused paths must match. Optional per-batch `query_position_offsets` / `key_position_offsets` cover varlen causal-mask metadata while keeping the dense tensor layout. +For Qwen3 WS2, `cp_attention` consumes post-QK-Norm, post-RoPE Q/K. It does not +call `NativeRoPEOp` internally and does not hide RoPE inside the CP merge. The +position offsets passed to CP attention must describe the same absolute token +positions used when RoPE was applied, so PR3 validates the post-RoPE Q/K boundary +while PR7 can later validate production fused `RoPE+Attention` kernels. + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -158,8 +164,9 @@ gradient flow, registry dispatch, and a GPU-only LARGE Qwen3-8B real-shape smoke test. `tests/test_cp_attention.py` covers the WS2 CP reference: CP=1 vs standard -attention, CP=2 prefill vs CP=1, chunked-prefill replay, global-position causal -masking across CP boundaries, order-independent LSE merge by global block index, +attention, CP=2 prefill vs CP=1, post-RoPE Q/K input semantics with shared +global position metadata, chunked-prefill replay, global-position causal masking +across CP boundaries, order-independent LSE merge by global block index, padding/all-masked stability, BF16 final-write behavior, input purity, argument validation, and registry dispatch. `make_operator_inputs("cp_attention", ...)` also emits a CP=2 chunked-prefill @@ -246,6 +253,8 @@ for measured peak memory at representative shapes. - Full materialization of scores/P limits practical sequence length. - `cp_attention` is a PyTorch reference for CP prefill/chunked-prefill semantics, not a distributed runtime or fused kernel. +- `cp_attention` consumes post-RoPE Q/K for Qwen3 WS2; RoPE execution and fused + `RoPE+Attention` backend alignment are outside PR3. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op). - No FP8, no multi-GPU / sequence-parallel. diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 734d97b9..302ad73f 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -84,6 +84,12 @@ def merge_attention_partial_states( 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 @@ -267,6 +273,8 @@ def local_partial_state( ``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) diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index dd0f53f7..03c861dd 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -19,6 +19,7 @@ merge_attention_partial_states, ) 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 @@ -106,6 +107,43 @@ def test_cp2_prefill_matches_cp1_reference(): 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) From f2c6acb5b59f3051eea646bedc831a6b99fd5628 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 4 Aug 2026 22:20:11 +0800 Subject: [PATCH 18/41] feat(alignment): bind rollout and training attention contracts (#235 PR4) Wire the CP attention path into the cross-configuration planner/runtime for the Qwen3-8B TP=2 CP=2 BF16 target. The PR4 criterion "rollout and training descriptors bind to the same semantic attention contract" cannot hold literally: training runs full-sequence prefill over a CP-sharded sequence while rollout runs vLLM paged-KV chunked prefill, so the two AttentionContract instances always differ. Binding is therefore split into three tiers -- identity must match bit for bit, reduction semantics must match each other and the WS2 mandate, and materialization differences are recorded and measured rather than rejected. reduction.engine stays in the recorded tier so a Transformer Engine merge oracle on one side does not fail the binding; reduction.order and acc_dtype stay in the semantic tier because that is the WS2 claim. Also adds the first two framework-shaped RuntimeMaterializer implementations. Before this the only one was CpuSmokeMaterializer over a synthetic CPU model, and every named scenario was planning-only. Neither adapter imports megatron or vllm, so the binding rules run on CPU in CI. Determinism is probed on both sides and compared, because the two frameworks mean different things by it: Megatron asserts NCCL_ALGO and leaves TF32 and BF16 reduced-precision reduction unmanaged, while vLLM hard-sets ten NCCL variables and disables both. Mismatches in NCCL_ALGO, NCCL_PROTO and CUBLAS_WORKSPACE_CONFIG are blocking; the rest are recorded. Fixes a latent break on the way: the planner normalizes dtype knobs to torch spellings (bfloat16) while AttentionDType uses short ones (bf16), so passing a normalized knob into the enum raised. Stacked on #236 (attention contract) and #238 (deterministic CP reference), on top of #230 (cross-configuration framework). Part of #235 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Q3Ar3z9fHEBFQQHddSEMaw --- .../ws2-attention-cross-config-integration.md | 135 ++++ ...config_qwen3_8b_megatron_tp2_cp2_vllm.json | 87 +++ .../cross_config/adapters/__init__.py | 33 + .../cross_config/adapters/_common.py | 263 ++++++++ .../alignment/cross_config/adapters/knobs.py | 160 +++++ .../cross_config/adapters/megatron.py | 381 +++++++++++ .../alignment/cross_config/adapters/vllm.py | 442 +++++++++++++ .../cross_config/attention_binding.py | 528 +++++++++++++++ .../alignment/cross_config/determinism.py | 305 +++++++++ tests/test_attention_cross_config_binding.py | 612 ++++++++++++++++++ 10 files changed, 2946 insertions(+) create mode 100644 docs/design/ws2-attention-cross-config-integration.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 tests/test_attention_cross_config_binding.py 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..06966ebe --- /dev/null +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -0,0 +1,135 @@ +# 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 version, weight version, tokenizer, token ids, active mask, position ids, padding side, pre-update state, Q/KV heads, head dim, RoPE theta/scaling/rotary dim, QK-Norm, cached global token positions, KV sequence lengths | equal bit for bit | `comparable=False`; no drift number from the pair means anything | +| `SEMANTIC` | `reduction.merge`, `reduction.acc_dtype`, `reduction.order`, `reduction.downcast_at`, `export_lse`, cross-side determinism mode | 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, CP/TP world sizes, local sequence length | free to differ | none; 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. +* **`reduction.order` and `reduction.acc_dtype` are `SEMANTIC`.** This is the entire + WS2 claim: merge order and accumulation precision are decided by the contract, not + by whichever backend happens to be selected. + +`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. + +## 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` (adds `kv_page_size` from + `CacheConfig.block_size` and `split_kv_policy` from + `AttentionConfig.flash_attn_max_num_splits_for_cuda_graph`) and + `VllmRolloutMaterializer`. + +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 | +| `rollout.context_parallel_size>1` with `mode=decode` | `FALLBACK` | vLLM CP covers prefill only; recorded with the reason | + +## 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 attention; +* decode-mode materialization, which needs the validated `KVCacheSpec` from #235 PR6 + and is refused with that reference rather than stubbed; +* Transformer Engine calls of any kind (PR4's TE plan is policy and provenance only); +* distributed drift benchmarks and report artifacts (#235 PR5); +* fused production backend alignment (#235 PR7) and backward (#235 PR8). 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..e74d3e6c --- /dev/null +++ b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json @@ -0,0 +1,87 @@ +{ + "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; a decode-mode contract runs at CP=1." + ] + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 1, + "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/rl_engine/alignment/cross_config/adapters/__init__.py b/rl_engine/alignment/cross_config/adapters/__init__.py new file mode 100644 index 00000000..c2db2134 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/__init__.py @@ -0,0 +1,33 @@ +# 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, 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", + "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..33b4d91c --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/_common.py @@ -0,0 +1,263 @@ +# 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.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, +) + +__all__ = [ + "QWEN3_8B", + "Qwen3ModelSpec", + "application", + "attention_dtype", + "build_reduction_spec", + "build_sharding_spec", + "causal_offsets_for", + "flatten", + "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 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..c1cf0b84 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/knobs.py @@ -0,0 +1,160 @@ +# 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( + # vLLM: AttentionConfig.flash_attn_max_num_splits_for_cuda_graph + "attention.split_kv_policy", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + 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..8b96b159 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -0,0 +1,381 @@ +# 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. It does not +launch ``torchrun``, initialize process groups, or execute attention. Binding a +constructed Megatron model to this contract is the next step and needs the 2-node +x 2-GPU environment that #239 fixes. +""" + +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, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + 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, + ): + 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 + + @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), + 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 + applications.append( + application( + descriptor, + requested, + requested, + requested, + MaterializationStatus.APPLIED, + "bound to the training-side attention contract", + frozen_scope_violations=list(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, + "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, + ), + ) 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..a8de61c9 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -0,0 +1,442 @@ +# 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. +""" + +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, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + 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]: + """The split-KV knob #235 PR5/PR7 needs; #236 has no field for it yet.""" + + 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, + "split_kv_policy": 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, + ): + 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 + + @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), + 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 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( + application( + descriptor, + requested, + requested, + requested, + MaterializationStatus.APPLIED, + "bound to the rollout-side attention contract", + frozen_scope_violations=list(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, + "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, + ), + ) 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..fcd4c203 --- /dev/null +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -0,0 +1,528 @@ +# 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 typing import Any, Optional + +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionDType, + AttentionMerge, + AttentionRole, + DowncastPoint, + ReductionOrder, +) + +__all__ = [ + "ATTENTION_LSE_DOMAIN", + "AttentionBindingError", + "AttentionBindingResult", + "BindingErrorCode", + "BindingIssue", + "BindingTier", + "IDENTITY_FIELDS", + "NULLABLE_IDENTITY_FIELDS", + "RECORDED_FIELDS", + "SEMANTIC_REDUCTION_FIELDS", + "WS2_ATTENTION_REDUCTION_MANDATE", + "bind_attention_contracts", + "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" + + +#: 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", + # 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", +) + + +#: 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", + "kv_cache.page_size", + "kv_cache.prefix_cache_enabled", + "kv_cache.block_table_shape", + "sharding.cp_world_size", + "sharding.tp_world_size", + "sharding.local_sequence_length", +) + + +@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.v1" + + 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) -> 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, + "sharding.cp_world_size": contract.sharding.cp_world_size, + "sharding.tp_world_size": contract.sharding.tp_world_size, + "sharding.local_sequence_length": contract.sharding.local_sequence_length, + } + 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), + ], + } + ) + return view + + +#: 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, +) -> 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``. + """ + + 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) + + # ---- 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}" + ), + ) + ) + + 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["backend_id"] = rollout_backend_id + training_recorded = _recorded_view(training_contract) + 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, + "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, + "lse_domain": ATTENTION_LSE_DOMAIN, + "rollout_backend": rollout_backend_id, + "training_backend": training_backend_id, + } + ), + recorded_differences=recorded_differences, + provenance=provenance, + ) + + +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/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py new file mode 100644 index 00000000..058dfa44 --- /dev/null +++ b/tests/test_attention_cross_config_binding.py @@ -0,0 +1,612 @@ +# 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 pathlib import Path +from types import SimpleNamespace + +import pytest + +from rl_engine.alignment.cross_config.adapters import ( + QWEN3_8B, + WS2_ATTENTION_KNOBS, + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) +from rl_engine.alignment.cross_config.attention_binding import ( + ATTENTION_LSE_DOMAIN, + AttentionBindingError, + BindingErrorCode, + BindingTier, + bind_attention_contracts, + 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 AttentionContractError, AttentionMode + +pytestmark = pytest.mark.unit + + +TRAINING_KNOBS = { + "batch.size": 2, + "training.tensor_parallel_size": 2, + "training.context_parallel_size": 2, + "training.compute_dtype": "bf16", +} + +ROLLOUT_KNOBS = { + "batch.size": 2, + "rollout.tensor_parallel_size": 2, + "rollout.context_parallel_size": 1, + "rollout.dtype": "bf16", +} + + +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", + "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 _bind(rollout_identity=None, training_identity=None, **kwargs): + rollout, training = _contracts() + return bind_attention_contracts( + rollout_contract=kwargs.pop("rollout_contract", rollout), + training_contract=kwargs.pop("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", + **kwargs, + ) + + +# -------------------------------------------------------------------------- +# tier 1: identity +# -------------------------------------------------------------------------- + + +def test_matching_identity_binds_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 == () + # Training runs CP=2 full prefill, rollout runs CP=1 chunked prefill. Those + # differences are recorded, not rejected. + assert "mode" in result.recorded_differences + assert "sharding.cp_world_size" 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" + + +# -------------------------------------------------------------------------- +# 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 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 materialization.binding.side_configs["training"]["contract"] is None + assert "arrival" in materialization.binding.side_configs["training"]["contract_error"] + + +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 + + +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 + + +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 From 4ad305b0829b21da0bbb610b325148c337ab7e3b Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 4 Aug 2026 22:23:29 +0800 Subject: [PATCH 19/41] fix(alignment): bind dtype, batch size and split-KV policy (#235 PR4) Three fields could differ between the two sides without the binding noticing. dtype was in no tier at all, so a BF16 rollout could bind to an FP16 training pass and produce a drift number attributable to nothing. It joins the semantic tier, with allow_dtype_difference for the #235 PR5 sweep that deliberately scores BF16 against an FP32 reference. batch_size was likewise unchecked. Batch invariance is a claim about results not changing with batch makeup, so two sides scoring different batches are not comparable and it belongs to identity. split_kv_policy has no field in the #236 contract, so it only reached side_configs and never took part in binding. Callers now pass it through rollout_recorded_extra / training_recorded_extra so the difference is at least visible in provenance; it can move into the contract once #236 grows the field. Part of #235 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Q3Ar3z9fHEBFQQHddSEMaw --- .../cross_config/attention_binding.py | 55 ++++++++++++++++- tests/test_attention_cross_config_binding.py | 59 +++++++++++++++++++ 2 files changed, 111 insertions(+), 3 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index fcd4c203..48e382b8 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -61,6 +61,7 @@ "IDENTITY_FIELDS", "NULLABLE_IDENTITY_FIELDS", "RECORDED_FIELDS", + "SEMANTIC_CONTRACT_FIELDS", "SEMANTIC_REDUCTION_FIELDS", "WS2_ATTENTION_REDUCTION_MANDATE", "bind_attention_contracts", @@ -126,6 +127,9 @@ class BindingErrorCode(str, Enum): "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", @@ -142,6 +146,14 @@ class BindingErrorCode(str, Enum): ) +#: 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",) + + #: 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 @@ -172,6 +184,11 @@ class BindingErrorCode(str, Enum): "sharding.cp_world_size", "sharding.tp_world_size", "sharding.local_sequence_length", + # Supplied by the caller, not by the contract: #236 has no split-KV field yet, so + # the value comes from vLLM's flash_attn_max_num_splits_for_cuda_graph via the + # adapter. Recorded so split-KV differences are at least visible in provenance + # until #236 grows the field and it can move into the contract proper. + "split_kv_policy", ) @@ -264,7 +281,10 @@ def _reduction_view(contract: AttentionContract) -> dict[str, Any]: } -def _recorded_view(contract: AttentionContract) -> dict[str, Any]: +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] = { @@ -297,6 +317,8 @@ def _recorded_view(contract: AttentionContract) -> dict[str, Any]: ], } ) + if extra: + view.update(extra) return view @@ -324,6 +346,9 @@ def bind_attention_contracts( training_backend_id: str, determinism_issues: Sequence[BindingIssue] = (), require_full_identity: bool = True, + allow_dtype_difference: bool = False, + 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. @@ -335,6 +360,13 @@ def bind_attention_contracts( ``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. + + ``rollout_recorded_extra`` / ``training_recorded_extra`` carry materialization + facts that #236 does not yet model -- today that is ``split_kv_policy``. They are + merged into the recorded tier, never into identity or semantics. """ if rollout_contract.role is not AttentionRole.INFER: @@ -419,6 +451,22 @@ def bind_attention_contracts( ) ) + 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" + ), + ) + ) + for side, contract in (("rollout", rollout_contract), ("training", training_contract)): if not contract.export_lse: issues.append( @@ -441,9 +489,9 @@ def bind_attention_contracts( issues.extend(determinism_issues) # ---- tier 3: recorded differences -------------------------------------- - rollout_recorded = _recorded_view(rollout_contract) + rollout_recorded = _recorded_view(rollout_contract, rollout_recorded_extra) rollout_recorded["backend_id"] = rollout_backend_id - training_recorded = _recorded_view(training_contract) + training_recorded = _recorded_view(training_contract, training_recorded_extra) training_recorded["backend_id"] = training_backend_id recorded_differences: dict[str, dict[str, Any]] = {} @@ -464,6 +512,7 @@ def bind_attention_contracts( provenance = { "lse_domain": ATTENTION_LSE_DOMAIN, + "dtype": training_contract.dtype.value, "rollout": { "contract": rollout_contract.to_dict(), "backend_id": rollout_backend_id, diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 058dfa44..815daf12 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -71,6 +71,7 @@ def _identity(**overrides): "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", } @@ -224,6 +225,64 @@ def test_lse_domain_is_recorded_as_attention_domain(): 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_split_kv_policy_difference_is_recorded(): + """#236 has no split-KV field, so the adapter supplies it to the recorded tier.""" + + result = _bind( + rollout_recorded_extra={"split_kv_policy": 8}, + training_recorded_extra={"split_kv_policy": None}, + ) + + assert result.passed + assert result.recorded_differences["split_kv_policy"] == { + "rollout": 8, + "training": None, + } + + +def test_matching_split_kv_policy_is_not_reported_as_a_difference(): + result = _bind( + rollout_recorded_extra={"split_kv_policy": 32}, + training_recorded_extra={"split_kv_policy": 32}, + ) + + assert "split_kv_policy" not in result.recorded_differences + + # -------------------------------------------------------------------------- # role and input validation # -------------------------------------------------------------------------- From deab4ed47e44ec827ab5cd2e7856579e9cddaf8c Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 12 Aug 2026 22:24:21 +0800 Subject: [PATCH 20/41] fix(attention): complete Split-KV contract validation --- docs/operators/attention.md | 108 ++- rl_engine/kernels/attention_contract.py | 683 +++++++++++++++++- rl_engine/kernels/gtest/operator_inputs.py | 79 +- rl_engine/kernels/gtest/operator_specs.py | 133 ++++ .../kernels/ops/cuda/attention/__init__.py | 40 + rl_engine/kernels/registry.py | 179 ++++- tests/test_attention_contract.py | 136 +++- 7 files changed, 1310 insertions(+), 48 deletions(-) diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 8bec7d22..bc7f4a38 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -48,7 +48,7 @@ The op exposes the WS1 dual-path contract: | Backend | Wrapper | Native symbol | Status | | --- | --- | --- | --- | | PyTorch fallback | `NativeAttentionOp` | None | fp32 ground-truth reference; CPU and any GPU. | -| CUDA / ROCm / Triton | — | — | Planned: downstream fused attention kernels validate against this reference. | +| CUDA deterministic | `DeterministicAttentionOp` | `_C.deterministic_attention_forward/backward` | Batch-invariant CUDA implementation (issue #147). | ## Tensor Contract @@ -76,12 +76,14 @@ the inputs' device. ## Dispatch Behavior `kernel_registry.get_op("attention")` resolves through the `OpBackend` priority map. On -`cuda` / `rocm` / `cpu` the only registered backend today is the PyTorch native op -(`PYTORCH_NATIVE_ATTENTION`), so every device dispatches to this op. Calling it (`__call__` -> -`forward(...)`) computes in the input dtype; `forward_fp32(...)` is the explicit fp32 golden -path. When fused attention kernels land, they are prepended to the priority list and the native -op becomes the fallback. The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is -a separate dispatch chain and is unaffected. +`cuda` the priority is: + +1. `CUDA_DETERMINISTIC_ATTENTION` — `DeterministicAttentionOp` (batch-invariant, fixed-order). +2. `PYTORCH_NATIVE_ATTENTION` — `NativeAttentionOp` (fallback). + +Calling it (`__call__` -> `forward(...)`) computes in the input dtype; `forward_fp32(...)` is +the explicit fp32 golden path (NativeAttentionOp only). The production `"attn"` op_type +(SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. ### WS2 CP-aware dispatch @@ -94,6 +96,21 @@ Existing WS1 implementations do not yet export attention-domain LSE or implement CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md). +Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow +`disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block +boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason. +Runtime-selected `auto` plans are diagnostic only unless both training and rollout export and +validate the same actual plan. + +The rank-aware drift benchmark can emit a CPU smoke artifact or a torchrun-friendly GPU report: + +```bash +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 \ + --cp-world-sizes 2 --kv-chunk-sizes none,1 --include-backward \ + --output artifacts/ws2-cp-attention-drift.json +``` + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -146,6 +163,7 @@ memory. ```bash python -m pytest tests/test_attention.py -v +python -m pytest tests/test_cp_attention.py -v ``` Covers: `forward_fp32` vs an independent fp32 reference (bitwise), strict-fp32 under hostile @@ -157,14 +175,80 @@ GPU-only LARGE Qwen3-8B real-shape smoke test. ## Implementation Files -- `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` +- `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` — ground-truth reference +- `rl_engine/kernels/ops/cuda/attention/deterministic_attn.py` — CUDA deterministic op +- `csrc/cuda/attention/deterministic_attention.cu` — CUDA kernels - `rl_engine/kernels/registry.py` - `tests/test_attention.py` +- `tests/test_deterministic_attention_cuda.py` + +## Fixed Reduction Order (CUDA Deterministic Backend) + +The CUDA `DeterministicAttentionOp` pins reduction order for batch-invariance: + +**QK kernel**: D-dimension FP32 accumulation, `d = 0 .. D-1`. Each `scores[b,hq,q,k]` has +exactly one writer thread. + +**Softmax + LSE kernel**: Each `(b, hq, q)` row processed by one CTA with fixed 256 threads. +Max and sum-exp use a fixed shared-memory tree reduction (power-of-two stride). No split by +batch size, sequence length, or SM count. + +**PV kernel**: K-dimension FP32 accumulation, `k = 0 .. Skv-1`. Each `out[b,hq,q,d]` has +exactly one writer thread. + +**Backward dK/dV**: Per `(b, hkv, k, d)` output element, a single thread accumulates over +query heads in group order then query positions: +```text +for local = 0 .. g-1: # g = Hq / Hkv + hq = hkv * g + local + for q = 0 .. Sq-1: + acc += ... +``` +No cross-CTA atomics. No launch-order dependent accumulation. + +## Prefill / Decode / KV-cache Shared Contract + +All inference modes use **the same standard attention kernels** (only Sq/Skv differ): + +- **Prefill**: `Sq == Skv`. Causal mask `key_index <= Skv - Sq + query_index`. +- **Chunked-prefill**: each chunk uses `Sq = chunk_size`, `Skv = past + chunk_size`. + Same causal offset formula produces identical results to full prefill at matching positions. +- **Decode**: `Sq = 1` (or few), `Skv = full_context`. Same kernel, same offset. +- **KV-cache**: caller does `k_full = cat([k_cache, k_new], dim=2)` then calls this op. + No separate KV-cache softmax implementation allowed. + +Hooks: +- `forward(q, k, v, ...)` — main path (registry, #108 harness). Differentiable. +- `forward_with_lse(q, k, v, ...)` — returns `(out, lse)` for LSE verification, debugging, + and future KV-cache / training integration. +- `backward_reference(q, k, v, dout, ...)` — runs the deterministic training backward + validation path and returns `dq`, `dk`, `dv`, `out`, `lse`, and provenance. +- `compare_cp_attention_backward(q, k, v, dout, ...)` — compares CP=1 backward against + CP/chunked-prefill backward and emits whole-tensor plus per-logical-rank drift stats. + +## Tolerance + +| Scenario | Comparison | Tolerance | +| --- | --- | --- | +| Same physical shape, varying batch position/size/chunk | bitwise | `batch_invariance` (atol=0, rtol=0) | +| Chunked-prefill on/off at same position | bitwise | `batch_invariance` | +| Prefill tail vs decode slice | bitwise | `batch_invariance` | +| CUDA vs `forward_fp32` output/grad | tolerance | `accuracy.default.attention` | +| Valid-only vs padded (reduction width differs) | near-equal | accuracy tolerance (NOT bitwise) | + +## Memory Tradeoff (First Version) + +The first version materializes full FP32 `scores [B, Hq, Sq, Skv]` and `P [B, Hq, Sq, Skv]`. +Memory cost: `4 * B * Hq * Sq * Skv` bytes per tensor. For Qwen3-8B at B=8, Sq=Skv=4096, +Hq=32: each tensor is ~17 GB. This is acceptable for correctness verification and moderate +sequence lengths but OOM-prone for long sequences. See `benchmarks/benchmark_deterministic_attention.py` +for measured peak memory at representative shapes. ## Known Limitations -- PyTorch fallback only; no fused CUDA/Triton backend yet (downstream work). +- First version: `D=128` only (Qwen3-8B alignment). +- Supported dtypes: BF16, FP16. +- Full materialization of scores/P limits practical sequence length. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). -- The naive path materializes the full `[B, Hq, Sq, Skv]` scores tensor — no query-chunking, - so the LARGE load point is memory-heavy and GPU-only. -- Covers softmax attention only; QK-Norm and RoPE are applied before the call. +- CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op). +- No FP8, no multi-GPU / sequence-parallel. diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 7deeddf0..bb0104ad 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -11,7 +11,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum from typing import Any, Iterable, TypeVar @@ -55,6 +55,12 @@ 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" @@ -284,6 +290,645 @@ def __post_init__(self) -> None: ) +@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.""" @@ -536,6 +1181,7 @@ class AttentionContract: 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 @@ -551,6 +1197,8 @@ def __post_init__(self) -> None: 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 @@ -713,6 +1361,7 @@ def to_dict(self) -> dict[str, Any]: "lse_domain": "attention", "sharding": sharding, "reduction": reduction, + "split_kv": self.split_kv.to_dict(), "kv_cache": kv_cache, "rope": rope, } @@ -734,6 +1383,10 @@ class AttentionBackendCapability: 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: @@ -763,6 +1416,10 @@ def __post_init__(self) -> None: "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, field), bool): raise AttentionContractError(f"{field} must be a bool") @@ -811,6 +1468,17 @@ def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: 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: @@ -830,6 +1498,10 @@ def to_dict(self) -> dict[str, Any]: "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, } @@ -862,4 +1534,13 @@ class AttentionDispatchResult: "RoPESpec", "RoPEState", "ShardingSpec", + "build_split_kv_runtime_plan_set", + "SplitKVExecutionPlan", + "SplitKVMode", + "SplitKVRuntimeCoordinate", + "SplitKVRuntimePlanEntry", + "SplitKVRuntimePlanSet", + "SplitKVSpec", + "validate_split_kv_alignment", + "validate_split_kv_plan_set_alignment", ] diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index f124cafb..78a9cfed 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -27,9 +27,12 @@ def make_operator_inputs( builders = { "rms_norm": _make_rms_norm_inputs, "matmul": _make_matmul_inputs, + "det_gemm": _make_det_gemm_inputs, "attention": _make_attention_inputs, + "cp_attention": _make_attention_inputs, "logp": _make_logp_inputs, "linear_logp": _make_linear_logp_inputs, + "batch_invariant_logp": _make_batch_invariant_logp_inputs, "rope": _make_rope_inputs, "silu": _make_silu_inputs, "swiglu": _make_swiglu_inputs, @@ -49,14 +52,17 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: names = { "rms_norm": f"{batch}x{seq}x{_normalized_dim(args)}", "matmul": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", + "det_gemm": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", + "cp_attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}xcp2", "logp": f"{batch}x{seq}x{vocab}", "linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", + "batch_invariant_logp": f"{batch}x{seq}x{vocab}", "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", "silu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", "swiglu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", - "embedding": f"{batch}x{seq}x{vocab}x{DEFAULT_HIDDEN}", - "lm_head": f"{batch}x{seq}x{vocab}", + "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", + "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", } try: @@ -89,23 +95,51 @@ def _make_matmul_inputs( } -def _make_attention_inputs( +def _make_det_gemm_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: batch, seq = _batch_seq(args) + k_dim = _matmul_k(args) + n_dim = _matmul_n(args) + m_dim = batch * seq return { - "q": _floating_tensor( - (batch, DEFAULT_N_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0 - ), - "k": _floating_tensor( - (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 1 - ), - "v": _floating_tensor( - (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 2 - ), - "causal": True, + "a": _floating_tensor((m_dim, k_dim), args, dtype, device, offset=0), + "b": _floating_tensor((k_dim, n_dim), args, dtype, device, offset=1), + } + + +def _make_attention_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + skv = _arg_int(args, "skv", seq) + n_heads = _arg_int(args, "n_heads", DEFAULT_N_HEADS) + n_kv_heads = _arg_int(args, "n_kv_heads", DEFAULT_N_KV_HEADS) + causal = bool(_arg_int(args, "causal", 1)) + use_padding = bool(_arg_int(args, "use_padding", 0)) + scale_mode = _arg_str(args, "scale_mode", "default") + + inputs: dict[str, Any] = { + "q": _floating_tensor((batch, n_heads, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0), + "k": _floating_tensor((batch, n_kv_heads, skv, DEFAULT_HEAD_DIM), args, dtype, device, 1), + "v": _floating_tensor((batch, n_kv_heads, skv, DEFAULT_HEAD_DIM), args, dtype, device, 2), + "causal": causal, } + if scale_mode == "zero": + inputs["scale"] = 0.0 + elif scale_mode == "custom": + inputs["scale"] = 0.05 + # else: scale_mode == "default" -> no scale kwarg (uses 1/sqrt(D)) + + if use_padding: + generator = _generator(args, device, offset=42) + key_padding_mask = torch.rand((batch, skv), generator=generator, device=device) > 0.3 + key_padding_mask[:, 0] = True + inputs["key_padding_mask"] = key_padding_mask + + return inputs + def _make_logp_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device @@ -132,6 +166,17 @@ def _make_linear_logp_inputs( } +def _make_batch_invariant_logp_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) + return { + "logits": _floating_tensor((batch, seq, vocab), args, dtype, device, offset=0), + "target_ids": _token_ids((batch, seq), vocab, args, device), + } + + def _make_rope_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: @@ -169,9 +214,10 @@ def _make_embedding_inputs( ) -> dict[str, Any]: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) + hidden_dim = _normalized_dim(args) return { "token_ids": _token_ids((batch, seq), vocab, args, device), - "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 0), + "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 0), } @@ -180,9 +226,10 @@ def _make_lm_head_inputs( ) -> dict[str, Any]: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) + hidden_dim = _normalized_dim(args) return { - "hidden": _floating_tensor((batch, seq, DEFAULT_HIDDEN), args, dtype, device, 0), - "weight": _floating_tensor((vocab, DEFAULT_HIDDEN), args, dtype, device, 1), + "hidden": _floating_tensor((batch, seq, hidden_dim), args, dtype, device, 0), + "weight": _floating_tensor((vocab, hidden_dim), args, dtype, device, 1), "bias": None, } diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 55a4a203..454f00ba 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -32,6 +32,53 @@ def _load_object(path: str) -> Any: OP_SPECS = { + "rms_norm": OperatorSpec( + name="rms_norm", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp", + "triton": "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp", + "cuda": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", + }, + grad_input_names=("x", "weight"), + ), + "attention": OperatorSpec( + name="attention", + op_class="attention", + gold_path="rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp", + "triton": ( + "rl_engine.kernels.ops.triton.attention.standard_attn." + "TritonBatchInvariantAttentionOp" + ), + "cuda": ( + "rl_engine.kernels.ops.cuda.attention.deterministic_attn." + "DeterministicAttentionOp" + ), + }, + grad_input_names=("q", "k", "v"), + ), + "cp_attention": OperatorSpec( + name="cp_attention", + op_class="attention", + gold_path=( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ), + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ), + }, + grad_input_names=("q", "k", "v"), + ), "logp": OperatorSpec( name="logp", op_class="logprob", @@ -57,6 +104,92 @@ def _load_object(path: str) -> Any: }, grad_input_names=("hidden", "lm_head_weight"), ), + "embedding": OperatorSpec( + name="embedding", + op_class="elementwise", + gold_path="rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp", + }, + grad_input_names=("weight",), + ), + "lm_head": OperatorSpec( + name="lm_head", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp", + }, + grad_input_names=("hidden", "weight"), + ), + "det_gemm": OperatorSpec( + name="det_gemm", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.matmul.det_gemm.NativeGemmOp", + gold_method="__call__", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.matmul.det_gemm.NativeGemmOp", + "cuda": "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp", + "triton": "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp", + }, + grad_input_names=("a", "b"), + ), + "rope": OperatorSpec( + name="rope", + op_class="elementwise", + gold_path="rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp", + "triton": "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op", + }, + 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", + gold_path="rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp." + "NativeBatchInvariantLogpOp", + gold_method="apply", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp." + "NativeBatchInvariantLogpOp", + "triton": "rl_engine.kernels.ops.triton.loss.batch_invariant_logp." + "TritonBatchInvariantLogpOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp." + "BatchInvariantLogpSM90Op", + }, + grad_input_names=("logits",), + ), } diff --git a/rl_engine/kernels/ops/cuda/attention/__init__.py b/rl_engine/kernels/ops/cuda/attention/__init__.py index 8d6addd9..2c3a1f1b 100644 --- a/rl_engine/kernels/ops/cuda/attention/__init__.py +++ b/rl_engine/kernels/ops/cuda/attention/__init__.py @@ -1,9 +1,49 @@ # File: rl_engine/kernels/ops/cuda/attention/__init__.py +from .cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunication, + AttentionCPCommunicationPlan, + AttentionCPCommunicationUnavailable, + AttentionCPMergedState, + AttentionCPPartialState, + AttentionParallelSpec, + CPCommunicationBackend, + CPCommunicationStatus, + CUDAAGRSAttentionCPCommunication, + P2PNCCLAttentionCPCommunication, + sort_attention_cp_partial_states, +) +from .deterministic_attn import DeterministicAttentionOp from .flash_attn import FlashAttentionOp +from .flashinfer_paged_attention import ( + FlashInferPagedAttentionConfig, + FlashInferQwen3PagedAttentionOp, + FlashInferRoPEFusionConfig, + FlashInferSplitKVPolicy, + FlashInferUnavailable, +) from .prefix_shared_attn import PrefixSharedAttentionOp __all__ = [ + "AttentionCPBlockMetadata", + "AttentionCPCommunication", + "AttentionCPCommunicationPlan", + "AttentionCPCommunicationUnavailable", + "AttentionCPMergedState", + "AttentionCPPartialState", + "AttentionParallelSpec", + "CPCommunicationBackend", + "CPCommunicationStatus", + "CUDAAGRSAttentionCPCommunication", + "P2PNCCLAttentionCPCommunication", + "DeterministicAttentionOp", "FlashAttentionOp", + "FlashInferPagedAttentionConfig", + "FlashInferQwen3PagedAttentionOp", + "FlashInferRoPEFusionConfig", + "FlashInferSplitKVPolicy", + "FlashInferUnavailable", "PrefixSharedAttentionOp", + "sort_attention_cp_partial_states", ] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index cab339fd..edd5efd5 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -6,6 +6,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, @@ -39,6 +41,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_FUSED_LOGP_SM90 = "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpSM90Op" CUDA_FUSED_LOGP_GENERIC = "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp" CUDA_DETERMINISTIC_LOGP = "rl_engine.kernels.ops.cuda.loss.logp.DeterministicLogpCUDAOp" + # Deterministic standard-softmax attention (issue #147); not FlashAttention. + CUDA_DETERMINISTIC_ATTENTION = ( + "rl_engine.kernels.ops.cuda.attention.deterministic_attn.DeterministicAttentionOp" + ) # AMD ROCm optimized stack ROCM_AITER = "rl_engine.kernels.ops.rocm.aiter.AiterOp" @@ -59,6 +65,25 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): TRITON_RATIO_KL = "rl_engine.kernels.ops.triton.loss.ratio_kl.TritonRatioKLOp" PYTORCH_RATIO_KL = "rl_engine.kernels.ops.pytorch.loss.ratio_kl.NativeRatioKLOp" + # Variable-length packing (pack-and-pad), [B,S,...] -> [Total_Active,...] + PYTORCH_PACK = "rl_engine.kernels.ops.pytorch.packing.pack.NativePackOp" + # Batch-invariant deterministic GEMM (WS1 #146) + CUDA_DET_GEMM = "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp" + TRITON_DET_GEMM = "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp" + # NON-deterministic reference (torch.matmul); reference/benchmark ONLY, + # intentionally excluded from det_gemm dispatch (cuBLAS breaks invariance). + PYTORCH_GEMM = "rl_engine.kernels.ops.pytorch.matmul.det_gemm.NativeGemmOp" + # Batch-invariant selected-logprob (WS1 #148: locked reduction order) + TRITON_BATCH_INVARIANT_LOGP = ( + "rl_engine.kernels.ops.triton.loss.batch_invariant_logp.TritonBatchInvariantLogpOp" + ) + PYTORCH_BATCH_INVARIANT_LOGP = ( + "rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp.NativeBatchInvariantLogpOp" + ) + CUDA_BATCH_INVARIANT_LOGP_SM90 = ( + "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90Op" + ) + # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" @@ -68,8 +93,14 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE = "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp" PYTORCH_NATIVE_MATMUL = "rl_engine.kernels.ops.pytorch.linear.matmul.NativeMatmulOp" PYTORCH_NATIVE_ROPE = "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp" + TRITON_ROPE = "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp" + 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. @@ -81,10 +112,18 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_KV_CACHE_ATTN = ( "rl_engine.kernels.ops.pytorch.attention.kv_cache.NativeKVCacheAttnOp" ) + # WS2 correctness-first context-parallel attention reference. It emulates + # CP prefill/chunked-prefill with fp32 attention-domain LSE merges. + 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 PYTORCH_NATIVE_EMBEDDING = "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp" + CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp" + CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" def resolve_logp_op_type( @@ -175,6 +214,27 @@ def __init__(self): 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, + # PR3 consumes attention-ready post-RoPE Q/K; RoPE execution + # and fused boundary validation remain in PR2/PR7 harnesses. + 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 = { @@ -206,22 +266,55 @@ def __init__(self): OpBackend.PYTORCH_NATIVE, ], "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN], - "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "attention": [ + OpBackend.CUDA_DETERMINISTIC_ATTENTION, + OpBackend.PYTORCH_NATIVE_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": [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.CUDA_DET_GEMM, OpBackend.TRITON_DET_GEMM], + "batch_invariant_logp": [ + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + ], "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.PYTORCH_NATIVE_ROPE], + "rope": [ + OpBackend.CUDA_ROPE_SM90, + OpBackend.TRITON_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": [ @@ -230,17 +323,27 @@ def __init__(self): OpBackend.TRITON_GENERIC, ], "attention": [OpBackend.PYTORCH_NATIVE_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.PYTORCH_NATIVE_ROPE], + "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], "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], + "batch_invariant_logp": [ + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + ], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "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], @@ -248,11 +351,18 @@ def __init__(self): "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "ws2_attention": [ + OpBackend.PYTORCH_CP_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP], "ratio_kl": [OpBackend.PYTORCH_RATIO_KL], + "pack": [OpBackend.PYTORCH_PACK], + "batch_invariant_logp": [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], @@ -279,7 +389,11 @@ 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 and rocm_attn_backend not in { + "native", + "pytorch", + "sdpa", + }: logger.warning( "Unknown RL_KERNEL_ROCM_ATTN_BACKEND=%s; using default ROCm attention priority.", rocm_attn_backend, @@ -315,24 +429,38 @@ def _adjust_priority_for_hardware(self): ll_list = self._priority_map["cuda"]["linear_logp"] if OpBackend.CUDA_FUSED_LINEAR_LOGP_SM90 not in ll_list: ll_list.insert(0, OpBackend.CUDA_FUSED_LINEAR_LOGP_SM90) + + # Batch-invariant logp SM90 kernel: same sm_90a TMA gating (Hopper only). + batch_inv_compiled = _EXT_AVAILABLE and hasattr(_C, "batch_invariant_logp_sm90") + if batch_inv_compiled and cc_major == 9: + bi_list = self._priority_map["cuda"]["batch_invariant_logp"] + if OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90 not in bi_list: + bi_list.insert(0, OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90) elif cc >= 90: logger.debug( f"SM{cc}: fused linear-logp SM90 kernel not compiled into _C; " "using generic linear-logp backend." ) + + sm90_embedding_compiled = _EXT_AVAILABLE and hasattr(_C, "embedding_sm90_forward") + if sm90_embedding_compiled and cc_major == 9: + embedding_list = self._priority_map["cuda"]["embedding"] + if OpBackend.CUDA_SM90_EMBEDDING not in embedding_list: + embedding_list.insert(0, OpBackend.CUDA_SM90_EMBEDDING) + + sm90_lm_head_compiled = _EXT_AVAILABLE and hasattr(_C, "lm_head_sm90_forward") + if sm90_lm_head_compiled and cc_major == 9: + lm_head_list = self._priority_map["cuda"]["lm_head"] + if OpBackend.CUDA_SM90_LM_HEAD not in lm_head_list: + lm_head_list.insert(0, OpBackend.CUDA_SM90_LM_HEAD) except Exception as e: logger.warning(f"Failed to probe device capability: {e}") - def get_op(self, op_type: str) -> Any: + def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: """Core distribution logic: Automatically select the best operator based on hardware and priority. """ - if device_ctx.is_rocm: - platform = "rocm" - elif device_ctx.device_type == "cuda": - platform = "cuda" - else: - platform = "cpu" + platform = self._platform_for_device(device) candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: @@ -356,6 +484,21 @@ def get_op(self, op_type: str) -> 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: + if device_ctx.is_rocm: + return "rocm" + if device_ctx.device_type == "cuda": + return "cuda" + return "cpu" + + 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_attention_op( self, contract: AttentionContract, @@ -377,7 +520,7 @@ def get_attention_op( platform = self._platform() op_type = "kv_cache_attention" if contract.mode is AttentionMode.DECODE else "attention" - candidates = self._priority_map.get(platform, {}).get(op_type, []) + candidates = self._priority_map.get(platform, {}).get("ws2_attention", []) rejected: list[str] = [] for backend in candidates: diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index b986fc15..47c94c40 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -22,6 +22,12 @@ 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 @@ -107,6 +113,8 @@ def _declared_cp_backend() -> AttentionBackendCapability: 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", ) @@ -254,6 +262,126 @@ def test_reduction_requires_fp32_accumulation(): 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"): @@ -460,6 +588,8 @@ def test_shared_prefix_pages_must_be_fully_populated(): 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()) @@ -473,7 +603,7 @@ def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): def test_undeclared_backend_capability_is_never_selected(): registry = KernelRegistry() platform = registry._platform() - registry._priority_map[platform]["attention"] = [OpBackend.PYTORCH_ATTN] + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_ATTN] with pytest.raises(RuntimeError, match="no AttentionBackendCapability declared"): registry.get_attention_op(_contract()) @@ -481,6 +611,8 @@ def test_undeclared_backend_capability_is_never_selected(): 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") @@ -497,6 +629,8 @@ def test_declared_compatible_backend_resolves_and_records_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"): From 9f46efff100e8b2067af979ec964c898bfdeeb12 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Wed, 12 Aug 2026 22:27:36 +0800 Subject: [PATCH 21/41] fix(attention): harden CP reference Split-KV validation --- docs/operators/attention.md | 64 +- rl_engine/kernels/attention_contract.py | 1472 +++++++++++++++++ .../ops/pytorch/attention/cp_attention.py | 486 +++++- tests/test_cp_attention.py | 213 +++ 4 files changed, 2198 insertions(+), 37 deletions(-) create mode 100644 rl_engine/kernels/attention_contract.py diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 988fad8c..bc7f4a38 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -85,20 +85,31 @@ Calling it (`__call__` -> `forward(...)`) computes in the input dtype; `forward_ the explicit fp32 golden path (NativeAttentionOp only). The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. -`kernel_registry.get_op("cp_attention")` resolves to -`DeterministicCPAttentionReferenceOp`, the WS2 correctness-first context-parallel -reference. It emulates CP prefill and chunked-prefill by splitting logical query -and KV sequence blocks, computing per-block `(out, lse)` partial states, and -merging them in fp32 by global KV block index. This path is not a production -fused backend; it defines the CP/LSE merge behavior that downstream fused paths -must match. Optional per-batch `query_position_offsets` / `key_position_offsets` -cover varlen causal-mask metadata while keeping the dense tensor layout. - -For Qwen3 WS2, `cp_attention` consumes post-QK-Norm, post-RoPE Q/K. It does not -call `NativeRoPEOp` internally and does not hide RoPE inside the CP merge. The -position offsets passed to CP attention must describe the same absolute token -positions used when RoPE was applied, so PR3 validates the post-RoPE Q/K boundary -while PR7 can later validate production fused `RoPE+Attention` kernels. +### WS2 CP-aware dispatch + +WS2 distributed callers use a separate contract-aware entry point, +`kernel_registry.get_attention_op(contract)`. It validates explicit TP/CP ownership, fixed +`(out, lse)` merge semantics, causal or packed-sequence offsets, and decode KV-cache identity +before selecting a backend. Legacy `get_op("attention")` behavior remains unchanged. + +Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic +CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as +a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md). + +Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow +`disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block +boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason. +Runtime-selected `auto` plans are diagnostic only unless both training and rollout export and +validate the same actual plan. + +The rank-aware drift benchmark can emit a CPU smoke artifact or a torchrun-friendly GPU report: + +```bash +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 \ + --cp-world-sizes 2 --kv-chunk-sizes none,1 --include-backward \ + --output artifacts/ws2-cp-attention-drift.json +``` ## Accuracy @@ -153,7 +164,6 @@ memory. ```bash python -m pytest tests/test_attention.py -v python -m pytest tests/test_cp_attention.py -v -python -m pytest tests/test_cp_attention_transformer_engine.py -v # optional TE oracle ``` Covers: `forward_fp32` vs an independent fp32 reference (bitwise), strict-fp32 under hostile @@ -163,30 +173,14 @@ invariance (slice + chunked, bitwise; padding is near-equality only, see below), gradient flow, registry dispatch, and a GPU-only LARGE Qwen3-8B real-shape smoke test. -`tests/test_cp_attention.py` covers the WS2 CP reference: CP=1 vs standard -attention, CP=2 prefill vs CP=1, post-RoPE Q/K input semantics with shared -global position metadata, chunked-prefill replay, global-position causal masking -across CP boundaries, order-independent LSE merge by global block index, -padding/all-masked stability, BF16 final-write behavior, input purity, argument -validation, and registry dispatch. -`make_operator_inputs("cp_attention", ...)` also emits a CP=2 chunked-prefill -synthetic case for local harnesses. -`tests/test_cp_attention_transformer_engine.py` optionally imports NVIDIA -Transformer Engine's context-parallel PyTorch correction helpers and checks that -RL-Kernel's fp32 `(out, lse)` merge matches those helpers; the test skips when -Transformer Engine is not installed. - ## Implementation Files - `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` — ground-truth reference -- `rl_engine/kernels/ops/pytorch/attention/cp_attention.py` — CP prefill/chunked reference - `rl_engine/kernels/ops/cuda/attention/deterministic_attn.py` — CUDA deterministic op - `csrc/cuda/attention/deterministic_attention.cu` — CUDA kernels - `rl_engine/kernels/registry.py` - `tests/test_attention.py` - `tests/test_deterministic_attention_cuda.py` -- `tests/test_cp_attention.py` -- `tests/test_cp_attention_transformer_engine.py` ## Fixed Reduction Order (CUDA Deterministic Backend) @@ -227,6 +221,10 @@ Hooks: - `forward(q, k, v, ...)` — main path (registry, #108 harness). Differentiable. - `forward_with_lse(q, k, v, ...)` — returns `(out, lse)` for LSE verification, debugging, and future KV-cache / training integration. +- `backward_reference(q, k, v, dout, ...)` — runs the deterministic training backward + validation path and returns `dq`, `dk`, `dv`, `out`, `lse`, and provenance. +- `compare_cp_attention_backward(q, k, v, dout, ...)` — compares CP=1 backward against + CP/chunked-prefill backward and emits whole-tensor plus per-logical-rank drift stats. ## Tolerance @@ -251,10 +249,6 @@ for measured peak memory at representative shapes. - First version: `D=128` only (Qwen3-8B alignment). - Supported dtypes: BF16, FP16. - Full materialization of scores/P limits practical sequence length. -- `cp_attention` is a PyTorch reference for CP prefill/chunked-prefill semantics, - not a distributed runtime or fused kernel. -- `cp_attention` consumes post-RoPE Q/K for Qwen3 WS2; RoPE execution and fused - `RoPE+Attention` backend alignment are outside PR3. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op). - No FP8, no multi-GPU / sequence-parallel. diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py new file mode 100644 index 00000000..eb4994b3 --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,1472 @@ +# 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) + + +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 _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 field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, field) + if values is None: + continue + normalized = _integer_tuple(values, field) + if not normalized or any(value < 0 for value in normalized): + raise AttentionContractError(f"{field} must contain non-negative positions") + object.__setattr__(self, 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 field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, field) + if offsets is not None and len(offsets) != batch_size: + raise AttentionContractError( + f"{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 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, field), bool): + raise AttentionContractError(f"{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", + "SplitKVExecutionPlan", + "SplitKVMode", + "SplitKVRuntimeCoordinate", + "SplitKVRuntimePlanEntry", + "SplitKVRuntimePlanSet", + "SplitKVSpec", + "validate_split_kv_alignment", + "validate_split_kv_plan_set_alignment", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 302ad73f..c334f49a 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -16,6 +16,13 @@ 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 @@ -45,6 +52,104 @@ def __post_init__(self) -> None: 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: @@ -97,6 +202,24 @@ class DeterministicCPAttentionReferenceOp: ``forward_fp32`` keeps the fp32 merged output. """ + op_class = "attention" + + @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, @@ -252,6 +375,108 @@ def forward_fp32_with_lse( output_dtype=torch.float32, ) + 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") + 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 + 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(device=out.device, 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, @@ -456,6 +681,143 @@ def _forward_impl( 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, @@ -476,8 +838,8 @@ def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) - 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 not overlap") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be gap-free and non-overlapping") previous_end = state.block_end @@ -546,8 +908,128 @@ def _kv_block_bounds( 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 + 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): + 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/tests/test_cp_attention.py b/tests/test_cp_attention.py index 03c861dd..196cfcaa 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -8,6 +8,7 @@ """ import contextlib +import json import math import pytest @@ -16,7 +17,9 @@ 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 @@ -26,6 +29,7 @@ _N_KV = 8 _HEAD_DIM = 128 _ATOL = 3.0e-6 +_GRAD_ATOL = 1.0e-5 @contextlib.contextmanager @@ -350,6 +354,203 @@ def test_cp2_chunked_gradients_match_cp1_reference(): 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, + ) + + def test_inputs_are_not_mutated(): op = DeterministicCPAttentionReferenceOp() q, k, v = _qkv(2, 6, 6, seed=7) @@ -419,5 +620,17 @@ def test_overlapping_partial_ranges_raise(): ) +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) From f98279793efce11129f7e9bcf487aad5bb2d0d89 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 17:22:05 +0800 Subject: [PATCH 22/41] fix(attention): validate CP reference numeric inputs Signed-off-by: lamentropetion <3051000145@qq.com> --- .../ops/pytorch/attention/cp_attention.py | 29 +++++++++++++++++-- tests/test_cp_attention.py | 28 ++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index c334f49a..6f80ea32 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -503,6 +503,7 @@ def local_partial_state( """ _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): @@ -598,9 +599,14 @@ def _forward_impl( kv_chunk_size: Optional[int], ) -> tuple[torch.Tensor, torch.Tensor]: _validate_qkv(q, k, v) - if cp_world_size < 1: + _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 kv_chunk_size < 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 @@ -850,10 +856,29 @@ def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: 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 _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: total = torch.tensor(0.0, device=tensors[0].device) for tensor in tensors: diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 196cfcaa..0db59caf 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -608,6 +608,34 @@ def test_invalid_gqa_and_mask_shapes_raise(): ) +@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) + + def test_overlapping_partial_ranges_raise(): out = torch.zeros(1, 1, 1, 1) lse = torch.zeros(1, 1, 1) From 3c8b734e196207f335b49568ed9847f57dd5efcd Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 17:28:58 +0800 Subject: [PATCH 23/41] fix(attention): bind backward gradient dtype and device Signed-off-by: lamentropetion <3051000145@qq.com> --- rl_engine/kernels/ops/pytorch/attention/cp_attention.py | 6 +++++- tests/test_cp_attention.py | 9 +++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 6f80ea32..897b23fd 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -406,6 +406,10 @@ def backward_reference( 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) @@ -424,7 +428,7 @@ def backward_reference( kv_chunk_size=kv_chunk_size, output_dtype=resolved_output_dtype, ) - torch.autograd.backward(out, dout.to(device=out.device, dtype=out.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") diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 0db59caf..8a84ac20 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -550,6 +550,15 @@ def test_backward_report_validates_dout_shape_and_dtype(): 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() From 88872f3491658acf6ce8bd372db76568a9f92aff Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 17:31:37 +0800 Subject: [PATCH 24/41] fix(attention): enforce FP32 CP merge state Signed-off-by: lamentropetion <3051000145@qq.com> --- .../ops/pytorch/attention/cp_attention.py | 17 ++++++++++++++++- tests/test_cp_attention.py | 17 +++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 897b23fd..adfcfb53 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -46,6 +46,10 @@ def __post_init__(self) -> None: 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: @@ -330,6 +334,8 @@ def forward_with_lse( ``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) out, lse = self._forward_impl( q, k, @@ -342,7 +348,7 @@ def forward_with_lse( cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size, ) - out = out.to(q.dtype if output_dtype is None else output_dtype) + out = out.to(resolved_output_dtype) return out, lse def forward_fp32_with_lse( @@ -415,6 +421,7 @@ def backward_reference( 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, @@ -883,6 +890,14 @@ def _validate_scale(scale: Optional[float]) -> None: 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: diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 8a84ac20..cc93874b 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -645,6 +645,23 @@ def test_boolean_parallelism_arguments_fail_closed(kwargs): 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) From 7f0e2914ef3d59ea2d3bbe33aa8ae48c9e0ec70f Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 21:19:13 +0800 Subject: [PATCH 25/41] style(attention): satisfy full PR lint Signed-off-by: lamentropetion <3051000145@qq.com> --- rl_engine/kernels/attention_contract.py | 70 +++++++++---------------- rl_engine/kernels/registry.py | 1 - tests/test_attention_contract.py | 4 +- 3 files changed, 26 insertions(+), 49 deletions(-) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index bb0104ad..206663e7 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -359,9 +359,7 @@ def __post_init__(self) -> None: or start < 0 or end <= start ): - raise AttentionContractError( - "Split-KV boundaries must satisfy 0 <= start < end" - ) + 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" @@ -486,14 +484,10 @@ def __post_init__(self) -> None: 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" - ) + 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" - ) + 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" @@ -647,13 +641,8 @@ def validate( "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("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" ) @@ -661,9 +650,7 @@ def validate( 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" - ) + raise AttentionContractError("Split-KV execution boundary escapes expected_kv_range") def to_dict(self) -> dict[str, Any]: return { @@ -709,9 +696,7 @@ def __post_init__(self) -> None: } 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" - ) + 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: @@ -821,17 +806,12 @@ def validate_split_kv_plan_set_alignment( ] if topology_mismatches: raise AttentionContractError( - "training/rollout Split-KV plan-set topology differs: " - + ", ".join(topology_mismatches) + "training/rollout Split-KV plan-set topology differs: " + ", ".join(topology_mismatches) ) - training_by_coordinate = { - entry.coordinate: entry for entry in training.entries - } + 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" - ) + 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] @@ -847,7 +827,7 @@ def validate_split_kv_plan_set_alignment( except AttentionContractError as exc: raise AttentionContractError( f"training/rollout Split-KV plan differs at {coordinate}: {exc}" - ) from exc + ) from exc def build_split_kv_runtime_plan_set( @@ -1148,14 +1128,16 @@ def __post_init__(self) -> None: 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 field in ("position_ids", "query_position_offsets", "key_position_offsets"): - values = getattr(self, field) + 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, field) + normalized = _integer_tuple(values, position_field) if not normalized or any(value < 0 for value in normalized): - raise AttentionContractError(f"{field} must contain non-negative positions") - object.__setattr__(self, field, 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") @@ -1259,11 +1241,11 @@ def __post_init__(self) -> None: "position_ids must describe the local query sequence or full local " "sequence length" ) - for field in ("query_position_offsets", "key_position_offsets"): - offsets = getattr(self.rope, field) + 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"{field} must contain one entry per logical batch entry" + 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: @@ -1409,7 +1391,7 @@ def __post_init__(self) -> None: 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 field in ( + for capability_field in ( "exports_attention_lse", "deterministic_cp_merge", "supports_packed_varlen", @@ -1421,8 +1403,8 @@ def __post_init__(self) -> None: "supports_split_kv_auto", "reports_actual_split_kv_plan", ): - if not isinstance(getattr(self, field), bool): - raise AttentionContractError(f"{field} must be a bool") + 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" @@ -1474,9 +1456,7 @@ def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: 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" - ) + 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) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index edd5efd5..6472bf08 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -519,7 +519,6 @@ def get_attention_op( requested_backend = requested_backend.strip().lower() platform = self._platform() - op_type = "kv_cache_attention" if contract.mode is AttentionMode.DECODE else "attention" candidates = self._priority_map.get(platform, {}).get("ws2_attention", []) rejected: list[str] = [] diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index 47c94c40..350545a1 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -321,9 +321,7 @@ def test_complete_split_kv_plan_set_covers_batch_tp_cp_and_owner_coordinates(): ) assert len(plan_set.entries) == 16 - assert plan_set.to_dict()["coverage"] == ( - "complete_batch_tp_cp_owner_cartesian_product" - ) + 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"] From 5b567e14b223ddb0263381d9690260575d1e224b Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 21:19:13 +0800 Subject: [PATCH 26/41] style(attention): satisfy full PR lint Signed-off-by: lamentropetion <3051000145@qq.com> --- rl_engine/kernels/attention_contract.py | 68 +++++++------------ .../ops/pytorch/attention/cp_attention.py | 18 +++-- 2 files changed, 32 insertions(+), 54 deletions(-) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index eb4994b3..1750476d 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -359,9 +359,7 @@ def __post_init__(self) -> None: or start < 0 or end <= start ): - raise AttentionContractError( - "Split-KV boundaries must satisfy 0 <= start < end" - ) + 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" @@ -486,14 +484,10 @@ def __post_init__(self) -> None: 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" - ) + 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" - ) + 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" @@ -647,13 +641,8 @@ def validate( "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("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" ) @@ -661,9 +650,7 @@ def validate( 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" - ) + raise AttentionContractError("Split-KV execution boundary escapes expected_kv_range") def to_dict(self) -> dict[str, Any]: return { @@ -709,9 +696,7 @@ def __post_init__(self) -> None: } 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" - ) + 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: @@ -821,17 +806,12 @@ def validate_split_kv_plan_set_alignment( ] if topology_mismatches: raise AttentionContractError( - "training/rollout Split-KV plan-set topology differs: " - + ", ".join(topology_mismatches) + "training/rollout Split-KV plan-set topology differs: " + ", ".join(topology_mismatches) ) - training_by_coordinate = { - entry.coordinate: entry for entry in training.entries - } + 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" - ) + 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] @@ -1075,14 +1055,16 @@ def __post_init__(self) -> None: 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 field in ("position_ids", "query_position_offsets", "key_position_offsets"): - values = getattr(self, field) + 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, field) + normalized = _integer_tuple(values, position_field) if not normalized or any(value < 0 for value in normalized): - raise AttentionContractError(f"{field} must contain non-negative positions") - object.__setattr__(self, field, 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") @@ -1186,11 +1168,11 @@ def __post_init__(self) -> None: "position_ids must describe the local query sequence or full local " "sequence length" ) - for field in ("query_position_offsets", "key_position_offsets"): - offsets = getattr(self.rope, field) + 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"{field} must contain one entry per logical batch entry" + 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: @@ -1336,7 +1318,7 @@ def __post_init__(self) -> None: 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 field in ( + for capability_field in ( "exports_attention_lse", "deterministic_cp_merge", "supports_packed_varlen", @@ -1348,8 +1330,8 @@ def __post_init__(self) -> None: "supports_split_kv_auto", "reports_actual_split_kv_plan", ): - if not isinstance(getattr(self, field), bool): - raise AttentionContractError(f"{field} must be a bool") + 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" @@ -1401,9 +1383,7 @@ def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: 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" - ) + 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) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index adfcfb53..e81ab8b7 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -465,9 +465,7 @@ def backward_reference( ], "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_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), @@ -611,7 +609,11 @@ def _forward_impl( ) -> 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: + 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) @@ -968,9 +970,7 @@ def split_kv_execution_plan_provenance( 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) - ): + for owner_cp_rank, (rank_start, rank_end) in enumerate(_split_bounds(length, cp_world_size)): if rank_start == rank_end: continue if kv_chunk_size is None: @@ -1007,9 +1007,7 @@ def build_reference_split_kv_runtime_plan_set( 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" - ) + 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: From c30be4b4f3cee7db247bf437ade94f26b9cf5e07 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 22:01:05 +0800 Subject: [PATCH 27/41] fix(attention): require runtime-verified cross-config binding --- .../ws2-attention-cross-config-integration.md | 34 +- ...config_qwen3_8b_megatron_tp2_cp2_vllm.json | 6 +- .../cross_config/adapters/__init__.py | 7 +- .../cross_config/adapters/_common.py | 18 + .../alignment/cross_config/adapters/knobs.py | 5 +- .../cross_config/adapters/megatron.py | 108 +++++- .../alignment/cross_config/adapters/vllm.py | 104 ++++- .../cross_config/attention_binding.py | 307 ++++++++++++++- tests/test_attention_cross_config_binding.py | 355 ++++++++++++++++-- 9 files changed, 871 insertions(+), 73 deletions(-) diff --git a/docs/design/ws2-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md index 06966ebe..c4d8d773 100644 --- a/docs/design/ws2-attention-cross-config-integration.md +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -27,9 +27,9 @@ tiers, in `rl_engine/alignment/cross_config/attention_binding.py`: | tier | fields | rule | failure | | --- | --- | --- | --- | -| `IDENTICAL` | checkpoint, model version, weight version, tokenizer, token ids, active mask, position ids, padding side, pre-update state, Q/KV heads, head dim, RoPE theta/scaling/rotary dim, QK-Norm, cached global token positions, KV sequence lengths | equal bit for bit | `comparable=False`; no drift number from the pair means anything | -| `SEMANTIC` | `reduction.merge`, `reduction.acc_dtype`, `reduction.order`, `reduction.downcast_at`, `export_lse`, cross-side determinism mode | 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, CP/TP world sizes, local sequence length | free to differ | none; recorded into provenance and measured | +| `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, 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: @@ -37,9 +37,13 @@ Two placements are load-bearing: 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. -* **`reduction.order` and `reduction.acc_dtype` are `SEMANTIC`.** This is the entire - WS2 claim: merge order and accumulation precision are decided by the contract, not - by whichever backend happens to be selected. +* **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 @@ -75,10 +79,16 @@ adapters: * `adapters/megatron.py` -- `MegatronProvenanceAdapter` (construction and distributed-context fingerprints, determinism probe, frozen-scope assertions) and `MegatronAttentionMaterializer`. -* `adapters/vllm.py` -- `VllmProvenanceAdapter` (adds `kv_page_size` from - `CacheConfig.block_size` and `split_kv_policy` from - `AttentionConfig.flash_attn_max_num_splits_for_cuda_graph`) and - `VllmRolloutMaterializer`. +* `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, 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. @@ -94,7 +104,9 @@ collapsing them onto the supported value: | `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 | -| `rollout.context_parallel_size>1` with `mode=decode` | `FALLBACK` | vLLM CP covers prefill only; recorded with the reason | +| 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 | ## Knobs 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 e74d3e6c..11f9cef1 100644 --- a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json +++ b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json @@ -21,7 +21,9 @@ "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; a decode-mode contract runs at CP=1." + "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": { @@ -30,7 +32,7 @@ }, "rollout": { "tensor_parallel_size": 2, - "context_parallel_size": 1, + "context_parallel_size": 2, "dtype": "bfloat16", "enable_prefix_caching": false, "enforce_eager": true, diff --git a/rl_engine/alignment/cross_config/adapters/__init__.py b/rl_engine/alignment/cross_config/adapters/__init__.py index c2db2134..f02b9b38 100644 --- a/rl_engine/alignment/cross_config/adapters/__init__.py +++ b/rl_engine/alignment/cross_config/adapters/__init__.py @@ -3,7 +3,11 @@ """Runtime adapters for the WS2 Qwen3-8B Megatron + vLLM cross-config target.""" -from rl_engine.alignment.cross_config.adapters._common import QWEN3_8B, Qwen3ModelSpec +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, @@ -21,6 +25,7 @@ __all__ = [ "MEGATRON_ATTENTION_BACKENDS", + "AttentionRuntimeReadback", "MegatronAttentionMaterializer", "MegatronProvenanceAdapter", "QWEN3_8B", diff --git a/rl_engine/alignment/cross_config/adapters/_common.py b/rl_engine/alignment/cross_config/adapters/_common.py index 33b4d91c..9a3c5384 100644 --- a/rl_engine/alignment/cross_config/adapters/_common.py +++ b/rl_engine/alignment/cross_config/adapters/_common.py @@ -9,6 +9,7 @@ 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, @@ -23,17 +24,20 @@ 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", ] @@ -107,6 +111,20 @@ def attention_dtype(value: Any, *, field: str) -> AttentionDType: ) 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.""" diff --git a/rl_engine/alignment/cross_config/adapters/knobs.py b/rl_engine/alignment/cross_config/adapters/knobs.py index c1cf0b84..33c29820 100644 --- a/rl_engine/alignment/cross_config/adapters/knobs.py +++ b/rl_engine/alignment/cross_config/adapters/knobs.py @@ -110,10 +110,11 @@ allowed_values=("unfused_rope_attention", "fused_rope_attention"), ), KnobDescriptor( - # vLLM: AttentionConfig.flash_attn_max_num_splits_for_cuda_graph + # 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",), + ("rollout", "training"), ), KnobDescriptor( # vLLM: CacheConfig.block_size -> AttentionContract.kv_cache.page_size diff --git a/rl_engine/alignment/cross_config/adapters/megatron.py b/rl_engine/alignment/cross_config/adapters/megatron.py index 8b96b159..8c9c9412 100644 --- a/rl_engine/alignment/cross_config/adapters/megatron.py +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -17,10 +17,10 @@ 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. It does not -launch ``torchrun``, initialize process groups, or execute attention. Binding a -constructed Megatron model to this contract is the next step and needs the 2-node -x 2-GPU environment that #239 fixes. +: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 @@ -32,6 +32,7 @@ from rl_engine.alignment.cross_config.adapters._common import ( QWEN3_8B, + AttentionRuntimeReadback, Qwen3ModelSpec, application, attention_dtype, @@ -39,6 +40,7 @@ build_sharding_spec, causal_offsets_for, flatten, + split_kv_spec, unsupported_reduction_reason, ) from rl_engine.alignment.cross_config.determinism import ( @@ -214,6 +216,7 @@ def __init__( 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 @@ -221,6 +224,7 @@ def __init__( self.cp_rank = cp_rank self.backend_id = backend_id self.provenance = provenance + self.runtime_readback = runtime_readback @property def implementation_fingerprint(self) -> str: @@ -272,6 +276,7 @@ def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: 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, ) @@ -326,14 +331,11 @@ def materialize( ) continue applications.append( - application( + self._runtime_application( descriptor, requested, - requested, - requested, - MaterializationStatus.APPLIED, - "bound to the training-side attention contract", - frozen_scope_violations=list(scope_violations), + contract=contract, + scope_violations=scope_violations, ) ) @@ -347,6 +349,9 @@ def materialize( "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: @@ -379,3 +384,86 @@ def materialize( 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 index a8de61c9..1e1ae185 100644 --- a/rl_engine/alignment/cross_config/adapters/vllm.py +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -16,7 +16,8 @@ ``init_batch_invariance()`` at worker startup. Like the Megatron adapter, nothing here imports ``vllm``; configs are duck-typed so -the module is importable anywhere. +the module is importable anywhere. Configured-only values remain ``UNOBSERVABLE``; +``APPLIED`` requires an explicit post-execution ``AttentionRuntimeReadback``. """ from __future__ import annotations @@ -28,6 +29,7 @@ from rl_engine.alignment.cross_config.adapters._common import ( QWEN3_8B, + AttentionRuntimeReadback, Qwen3ModelSpec, application, attention_dtype, @@ -35,6 +37,7 @@ 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 @@ -204,7 +207,7 @@ def kv_page_size(self) -> Optional[int]: @property def split_kv_policy(self) -> Optional[int]: - """The split-KV knob #235 PR5/PR7 needs; #236 has no field for it yet.""" + """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 @@ -218,7 +221,7 @@ def to_dict(self) -> dict[str, Any]: "distributed_context_fingerprint": self.distributed_context_fingerprint, "frozen_scope_violations": list(self.frozen_scope_violations()), "kv_page_size": self.kv_page_size, - "split_kv_policy": self.split_kv_policy, + "flash_attn_max_num_splits_for_cuda_graph": self.split_kv_policy, "determinism": self.determinism_probe().to_dict(), } @@ -238,6 +241,7 @@ def __init__( 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 @@ -246,6 +250,7 @@ def __init__( self.mode = mode self.backend_id = backend_id self.provenance = provenance + self.runtime_readback = runtime_readback @property def implementation_fingerprint(self) -> str: @@ -312,6 +317,7 @@ def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: 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, ) @@ -385,14 +391,11 @@ def materialize( ) continue applications.append( - application( + self._runtime_application( descriptor, requested, - requested, - requested, - MaterializationStatus.APPLIED, - "bound to the rollout-side attention contract", - frozen_scope_violations=list(scope_violations), + contract=contract, + scope_violations=scope_violations, ) ) @@ -408,6 +411,9 @@ def materialize( "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: @@ -440,3 +446,83 @@ def materialize( 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 index 48e382b8..4156d57e 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -40,21 +40,26 @@ 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 ( AttentionContract, + AttentionContractError, AttentionDType, AttentionMerge, AttentionRole, DowncastPoint, ReductionOrder, + SplitKVRuntimePlanSet, + validate_split_kv_plan_set_alignment, ) __all__ = [ "ATTENTION_LSE_DOMAIN", "AttentionBindingError", "AttentionBindingResult", + "AttentionRuntimeReadback", "BindingErrorCode", "BindingIssue", "BindingTier", @@ -63,8 +68,10 @@ "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", @@ -97,6 +104,51 @@ class BindingErrorCode(str, Enum): 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" + + +@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 + + 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") + + 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))) + + @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), + "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), + } #: Attention exports attention-domain LSE, never vocab-logprob LSE (#235). @@ -154,6 +206,29 @@ class BindingErrorCode(str, Enum): 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 @@ -181,14 +256,6 @@ class BindingErrorCode(str, Enum): "kv_cache.page_size", "kv_cache.prefix_cache_enabled", "kv_cache.block_table_shape", - "sharding.cp_world_size", - "sharding.tp_world_size", - "sharding.local_sequence_length", - # Supplied by the caller, not by the contract: #236 has no split-KV field yet, so - # the value comes from vLLM's flash_attn_max_num_splits_for_cuda_graph via the - # adapter. Recorded so split-KV differences are at least visible in provenance - # until #236 grows the field and it can move into the contract proper. - "split_kv_policy", ) @@ -233,7 +300,7 @@ class AttentionBindingResult: 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.v1" + schema_version: str = "cross_config.attention_binding.v2" def issues_by_code(self, code: BindingErrorCode) -> tuple[BindingIssue, ...]: return tuple(issue for issue in self.issues if issue.code is code) @@ -291,9 +358,6 @@ def _recorded_view( "mode": contract.mode.value, "backend_id": None, "reduction.engine": contract.reduction.engine.value, - "sharding.cp_world_size": contract.sharding.cp_world_size, - "sharding.tp_world_size": contract.sharding.tp_world_size, - "sharding.local_sequence_length": contract.sharding.local_sequence_length, } if rope is not None: view.update( @@ -322,6 +386,69 @@ def _recorded_view( 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. @@ -347,6 +474,8 @@ def bind_attention_contracts( 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: @@ -364,9 +493,13 @@ def bind_attention_contracts( ``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. - ``rollout_recorded_extra`` / ``training_recorded_extra`` carry materialization - facts that #236 does not yet model -- today that is ``split_kv_policy``. They are - merged into the recorded tier, never into identity or semantics. + 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: @@ -415,6 +548,26 @@ def bind_attention_contracts( 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) @@ -467,6 +620,79 @@ def bind_attention_contracts( ) ) + 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( @@ -513,6 +739,14 @@ def bind_attention_contracts( 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, @@ -535,6 +769,8 @@ def bind_attention_contracts( { "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, @@ -545,6 +781,47 @@ def bind_attention_contracts( ) +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", + ) + ) + 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, + ) + + def summarize_binding(result: AttentionBindingResult) -> str: """One-line human summary for CLI output and failure messages.""" diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 815daf12..6a99e6bb 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -11,6 +11,8 @@ from __future__ import annotations import json +from dataclasses import replace +from enum import Enum from pathlib import Path from types import SimpleNamespace @@ -19,6 +21,7 @@ from rl_engine.alignment.cross_config.adapters import ( QWEN3_8B, WS2_ATTENTION_KNOBS, + AttentionRuntimeReadback, MegatronAttentionMaterializer, MegatronProvenanceAdapter, VllmProvenanceAdapter, @@ -30,6 +33,7 @@ BindingErrorCode, BindingTier, bind_attention_contracts, + bind_attention_runtime_readbacks, first_blocking_issue, identity_fingerprint, summarize_binding, @@ -40,7 +44,17 @@ vllm_probe_from_env, ) from rl_engine.alignment.cross_config.schema import MaterializationStatus -from rl_engine.kernels.attention_contract import AttentionContractError, AttentionMode +from rl_engine.kernels.attention_contract import ( + AttentionContractError, + AttentionMode, + AttentionRole, + KVCacheSpec, + SplitKVExecutionPlan, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, + SplitKVSpec, + build_split_kv_runtime_plan_set, +) pytestmark = pytest.mark.unit @@ -50,13 +64,15 @@ "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": 1, + "rollout.context_parallel_size": 2, "rollout.dtype": "bf16", + "attention.split_kv_policy": 32, } @@ -86,15 +102,33 @@ def _contracts(): 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=kwargs.pop("rollout_contract", rollout), - training_contract=kwargs.pop("training_contract", training), + 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, ) @@ -104,7 +138,7 @@ def _bind(rollout_identity=None, training_identity=None, **kwargs): # -------------------------------------------------------------------------- -def test_matching_identity_binds_despite_different_materialization(): +def test_matching_identity_and_topology_bind_despite_different_materialization(): """The core claim of PR4: same identity + same reduction, different runtimes.""" result = _bind() @@ -112,10 +146,9 @@ def test_matching_identity_binds_despite_different_materialization(): assert result.comparable assert result.passed assert result.issues == () - # Training runs CP=2 full prefill, rollout runs CP=1 chunked prefill. Those - # differences are recorded, not rejected. + # 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 "sharding.cp_world_size" in result.recorded_differences assert result.recorded_differences["mode"] == { "rollout": "chunked_prefill", "training": "prefill", @@ -259,28 +292,133 @@ def test_batch_size_mismatch_is_not_comparable(): assert any(issue.field == "batch_size" for issue in result.issues) -def test_split_kv_policy_difference_is_recorded(): - """#236 has no split-KV field, so the adapter supplies it to the recorded tier.""" +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) + - result = _bind( - rollout_recorded_extra={"split_kv_policy": 8}, - training_recorded_extra={"split_kv_policy": None}, +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 result.passed - assert result.recorded_differences["split_kv_policy"] == { - "rollout": 8, - "training": None, - } + 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) -def test_matching_split_kv_policy_is_not_reported_as_a_difference(): - result = _bind( - rollout_recorded_extra={"split_kv_policy": 32}, - training_recorded_extra={"split_kv_policy": 32}, + 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) ) - assert "split_kv_policy" not in result.recorded_differences + +@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) # -------------------------------------------------------------------------- @@ -480,6 +618,175 @@ 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) + return AttentionRuntimeReadback( + contract=contract, + actual_knobs=dict(flat), + split_kv_plan_set=_plan_set(contract, backend=source), + source=source, + frozen_scope_verified=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, + ) + 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" + ) + + +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.""" @@ -535,6 +842,7 @@ def test_vllm_cp_falls_back_to_one_in_decode_and_says_why(): 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(): @@ -613,6 +921,7 @@ def test_vllm_provenance_reads_page_size_and_split_kv_policy(): 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(): From 59d68ebc9b83c22560e26e3430e007caab4ab742 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 22:36:52 +0800 Subject: [PATCH 28/41] fix(attention): type split-k runtime boundaries --- rl_engine/kernels/ops/pytorch/attention/cp_attention.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index e81ab8b7..0062dac4 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -973,6 +973,7 @@ def split_kv_execution_plan_provenance( 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 @@ -1019,6 +1020,7 @@ def build_reference_split_kv_runtime_plan_set( 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),) From 1bcb885914477f8425958d8038e2c476722599ba Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 22:41:18 +0800 Subject: [PATCH 29/41] fix(attention): fail closed without runtime contract --- .../cross_config/adapters/megatron.py | 12 +++++++++++ .../alignment/cross_config/adapters/vllm.py | 12 +++++++++++ tests/test_attention_cross_config_binding.py | 20 +++++++++++++++++++ 3 files changed, 44 insertions(+) diff --git a/rl_engine/alignment/cross_config/adapters/megatron.py b/rl_engine/alignment/cross_config/adapters/megatron.py index 8c9c9412..55903cc8 100644 --- a/rl_engine/alignment/cross_config/adapters/megatron.py +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -330,6 +330,18 @@ def materialize( ) ) 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, diff --git a/rl_engine/alignment/cross_config/adapters/vllm.py b/rl_engine/alignment/cross_config/adapters/vllm.py index 1e1ae185..f89ab232 100644 --- a/rl_engine/alignment/cross_config/adapters/vllm.py +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -374,6 +374,18 @@ def materialize( ) ) 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( diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 6a99e6bb..04f0bfdb 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -800,10 +800,30 @@ def test_arrival_merge_order_is_unsupported_not_silently_corrected(): 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}, From 48a4130662357e8a38708f7ad42d1e3294812c06 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 23:12:21 +0800 Subject: [PATCH 30/41] fix(attention): export only available CUDA operators --- .../kernels/ops/cuda/attention/__init__.py | 38 ------------------- 1 file changed, 38 deletions(-) diff --git a/rl_engine/kernels/ops/cuda/attention/__init__.py b/rl_engine/kernels/ops/cuda/attention/__init__.py index 2c3a1f1b..09775c8e 100644 --- a/rl_engine/kernels/ops/cuda/attention/__init__.py +++ b/rl_engine/kernels/ops/cuda/attention/__init__.py @@ -1,49 +1,11 @@ # File: rl_engine/kernels/ops/cuda/attention/__init__.py -from .cp_comm import ( - AttentionCPBlockMetadata, - AttentionCPCommunication, - AttentionCPCommunicationPlan, - AttentionCPCommunicationUnavailable, - AttentionCPMergedState, - AttentionCPPartialState, - AttentionParallelSpec, - CPCommunicationBackend, - CPCommunicationStatus, - CUDAAGRSAttentionCPCommunication, - P2PNCCLAttentionCPCommunication, - sort_attention_cp_partial_states, -) from .deterministic_attn import DeterministicAttentionOp from .flash_attn import FlashAttentionOp -from .flashinfer_paged_attention import ( - FlashInferPagedAttentionConfig, - FlashInferQwen3PagedAttentionOp, - FlashInferRoPEFusionConfig, - FlashInferSplitKVPolicy, - FlashInferUnavailable, -) from .prefix_shared_attn import PrefixSharedAttentionOp __all__ = [ - "AttentionCPBlockMetadata", - "AttentionCPCommunication", - "AttentionCPCommunicationPlan", - "AttentionCPCommunicationUnavailable", - "AttentionCPMergedState", - "AttentionCPPartialState", - "AttentionParallelSpec", - "CPCommunicationBackend", - "CPCommunicationStatus", - "CUDAAGRSAttentionCPCommunication", - "P2PNCCLAttentionCPCommunication", "DeterministicAttentionOp", "FlashAttentionOp", - "FlashInferPagedAttentionConfig", - "FlashInferQwen3PagedAttentionOp", - "FlashInferRoPEFusionConfig", - "FlashInferSplitKVPolicy", - "FlashInferUnavailable", "PrefixSharedAttentionOp", - "sort_attention_cp_partial_states", ] From 4de96a4e2cc822449bbefc40e384419aa2a93515 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 02:11:46 +0800 Subject: [PATCH 31/41] feat(attention): bind strict H100 QK norm and RoPE path --- .../ws2-attention-cross-config-integration.md | 46 ++++- .../cross_config/attention_binding.py | 105 +++++++++- rl_engine/kernels/attention_preprocess.py | 180 ++++++++++++++++++ rl_engine/kernels/ops/cuda/norm/rmsnorm.py | 15 ++ .../kernels/ops/cuda/rotary_embedding/rope.py | 39 +++- rl_engine/kernels/registry.py | 6 +- tests/test_attention_cross_config_binding.py | 59 ++++++ tests/test_attention_preprocess.py | 111 +++++++++++ tests/test_rms_norm.py | 15 +- 9 files changed, 558 insertions(+), 18 deletions(-) create mode 100644 rl_engine/kernels/attention_preprocess.py create mode 100644 tests/test_attention_preprocess.py diff --git a/docs/design/ws2-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md index c4d8d773..4bda3089 100644 --- a/docs/design/ws2-attention-cross-config-integration.md +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -28,7 +28,7 @@ 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, cross-side determinism | both sides equal **and** equal to the WS2 mandate | fail closed | +| `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: @@ -49,6 +49,42 @@ Two placements are load-bearing: 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 boundary + +The strict experiment does not use the PyTorch reference operators as an +executable option. `H100AttentionPreprocessor` applies these implementations in +the fixed Qwen3 order: + +1. `RMSNormCudaOp` on Q and K (`rlkernel.cuda.rmsnorm`) +2. `RoPESM90Op` with global `[S]` or per-batch `[B, S]` positions + (`rlkernel.cuda.rope_sm90`) + +There is no Megatron/vLLM-native fallback. The CUDA RoPE path accepts non-contiguous +global positions, including zigzag CP ownership. 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 backend ID, a runtime-native backend ID, or any +reported fallback. Printing a configured backend without executing it is not +accepted as evidence. + +The boundary starts at projected Q/K/V. Pre-attention model RMSNorm, QKV projection +GEMM, and projection-owned TP/SP All-Gather/Reduce-Scatter are not part of the +Attention operator experiment. The isolated H100 test must therefore capture or +reuse identical projected Q/K/V inputs. Those upstream operators must be aligned +separately before making an end-to-end Megatron-vLLM logprob claim. + ## Determinism is not one thing `rl_engine/alignment/cross_config/determinism.py` probes both sides and compares @@ -83,7 +119,8 @@ adapters: limits) and `VllmRolloutMaterializer`. * `AttentionRuntimeReadback` -- the explicit handoff from an executed engine. It carries the reconstructed actual contract, actual knob values, frozen-scope - verification, and the complete Split-KV runtime plan set. + 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 @@ -107,6 +144,7 @@ collapsing them onto the supported value: | 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/native/fallback QK-Norm or RoPE backend | binding failure | both sides must execute the RL-Kernel CUDA preprocessing path | ## Knobs @@ -139,7 +177,9 @@ retired rather than rewritten. Deliberately not in this PR: -* launching `torchrun`, initializing process groups, or executing attention; +* launching `torchrun`, initializing process groups, or executing core attention; +* pre-attention model RMSNorm, QKV projection GEMM, and projection-owned TP/SP + communication; isolated tests reuse identical projected Q/K/V; * decode-mode materialization, which needs the validated `KVCacheSpec` from #235 PR6 and is refused with that reference rather than stubbed; * Transformer Engine calls of any kind (PR4's TE plan is policy and provenance only); diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 4156d57e..1007f197 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -54,6 +54,7 @@ SplitKVRuntimePlanSet, validate_split_kv_plan_set_alignment, ) +from rl_engine.kernels.attention_preprocess import MANDATED_ATTENTION_PREPROCESS_BACKENDS __all__ = [ "ATTENTION_LSE_DOMAIN", @@ -108,6 +109,9 @@ class BindingErrorCode(str, Enum): 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" @dataclass(frozen=True) @@ -119,6 +123,8 @@ class AttentionRuntimeReadback: split_kv_plan_set: SplitKVRuntimePlanSet source: str frozen_scope_verified: bool + preprocess_backends: Mapping[str, str] = field(default_factory=dict) + preprocess_fallback: bool = False def __post_init__(self) -> None: if not isinstance(self.contract, AttentionContract): @@ -131,11 +137,25 @@ def __post_init__(self) -> None: 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") 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)), + ) @property def split_kv_fallback(self) -> bool: @@ -147,6 +167,10 @@ def to_dict(self) -> dict[str, Any]: "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, + }, "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), } @@ -253,6 +277,9 @@ def to_dict(self) -> dict[str, Any]: "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", @@ -300,7 +327,7 @@ class AttentionBindingResult: 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.v2" + 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) @@ -774,6 +801,22 @@ def bind_attention_contracts( "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, @@ -809,6 +852,7 @@ def bind_attention_runtime_readbacks( message=f"{side} runtime did not verify the frozen attention scope", ) ) + missing_scope_evidence.extend(_attention_preprocess_issues(side, readback)) return bind_attention_contracts( rollout_contract=rollout.contract, training_contract=training.contract, @@ -819,9 +863,68 @@ def bind_attention_runtime_readbacks( 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, + }, + training_recorded_extra={ + **{ + f"preprocess.{name}": backend + for name, backend in training.preprocess_backends.items() + }, + "preprocess.fallback": training.preprocess_fallback, + }, ) +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: + 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}" + ), + ) + ) + if readback.preprocess_fallback: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_FALLBACK, + tier=BindingTier.SEMANTIC, + field=f"{side}.preprocess.fallback", + message=f"{side} reported a QK-Norm or RoPE backend fallback", + ) + ) + return issues + + def summarize_binding(result: AttentionBindingResult) -> str: """One-line human summary for CLI output and failure messages.""" diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py new file mode 100644 index 00000000..196d2e18 --- /dev/null +++ b/rl_engine/kernels/attention_preprocess.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Strict H100 QK-Norm and RoPE handoff for WS2 Attention. + +This module intentionally has no runtime-native fallback. A caller either runs +the RL-Kernel CUDA operators and records their identities, or the experiment +fails before Attention executes. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Mapping + +import torch +from torch import Tensor + + +QK_RMSNORM_BACKEND_ID = "rlkernel.cuda.rmsnorm" +ROPE_BACKEND_ID = "rlkernel.cuda.rope_sm90" +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] + + 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, + } + + +class H100AttentionPreprocessor: + """Apply RL-Kernel CUDA QK-Norm then RoPE without silent fallback.""" + + def __init__(self, device: torch.device | str | int | None = None) -> 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() + actual_backends = { + "qk_rmsnorm": self.rmsnorm.backend_id, + "rope": self.rope.backend_id, + } + if actual_backends != dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS): + raise RuntimeError(f"unexpected Attention preprocess backends: {actual_backends}") + self.backend_ids = MappingProxyType(actual_backends) + + 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 = self.rmsnorm(q, q_weight, eps=eps) + k_norm = self.rmsnorm(k, k_weight, eps=eps) + return AttentionPreprocessResult( + q=self.rope(q_norm, positions, theta=theta), + k=self.rope(k_norm, positions, theta=theta), + backend_ids=self.backend_ids, + fallback=False, + device_capability=self.device_capability, + ) + + +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", +] diff --git a/rl_engine/kernels/ops/cuda/norm/rmsnorm.py b/rl_engine/kernels/ops/cuda/norm/rmsnorm.py index 76e33da8..7ac85b70 100644 --- a/rl_engine/kernels/ops/cuda/norm/rmsnorm.py +++ b/rl_engine/kernels/ops/cuda/norm/rmsnorm.py @@ -79,6 +79,21 @@ def rmsnorm_cuda(x, weight, eps=1e-6, mask=None): class RMSNormCudaOp: """CUDA RMSNorm wrapper compatible with the shared operator harness.""" + backend_id = "rlkernel.cuda.rmsnorm" + + def __init__(self): + required = ( + "rmsnorm_forward", + "rmsnorm_backward_dx", + "rmsnorm_backward_dw", + ) + missing = [name for name in required if not _EXT_AVAILABLE or not hasattr(_C, name)] + if missing: + raise RuntimeError( + "CUDA RMSNorm extension is incomplete; rebuild _C with rmsnorm.cu " + f"(missing: {', '.join(missing)})" + ) + def __call__(self, x, weight, *, eps=1e-6): return self.forward(x, weight, eps=eps) diff --git a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py index 0c1a7b73..028e44b9 100644 --- a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py +++ b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py @@ -18,7 +18,7 @@ def _build_cos_sin(positions: Tensor, half: int, theta: float, device: torch.device): - """fp32 cos/sin caches of shape [S, half], identical math to NativeRoPEOp.""" + """fp32 cos/sin rows, identical math to NativeRoPEOp.""" inv_freq = 1.0 / (theta ** (torch.arange(0, half, dtype=torch.float32, device=device) / half)) pos = positions.to(device=device, dtype=torch.float32).reshape(-1, 1) freqs = pos * inv_freq # [S, half] @@ -31,19 +31,39 @@ def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: D = x.shape[-1] if D % 2 != 0: raise ValueError(f"RoPE head_dim must be even, got {D}") - if positions.dim() != 1: - raise NotImplementedError( - "CUDA RoPE currently supports 1-D positions [S] (shared across batch)." - ) - S = positions.shape[0] + if positions.dim() not in (1, 2): + raise ValueError("positions must have shape [S] or [B, S]") + S = positions.shape[-1] + if S == 0: + raise ValueError("positions must not be empty") + if x.shape[-2] != S: + raise ValueError(f"x sequence length {x.shape[-2]} does not match positions length {S}") x_2d = x.contiguous().reshape(-1, D) n_rows = x_2d.shape[0] - if n_rows % S != 0: + if positions.dim() == 2: + batch = positions.shape[0] + if x.dim() < 3 or x.shape[0] != batch: + raise ValueError( + f"x batch size {x.shape[0]} does not match positions batch size {batch}" + ) + rows_per_token = n_rows // (batch * S) + if rows_per_token * batch * S != n_rows: + raise ValueError("x rows are incompatible with [B, S] positions") + # The CUDA kernel accepts one fp32 cos/sin row per flattened x row. + # Expanding positions preserves arbitrary global/zigzag indices while + # keeping the arithmetic inside the precompiled deterministic kernel. + kernel_positions = ( + positions[:, None, :].expand(batch, rows_per_token, S).contiguous().reshape(-1) + ) + else: + kernel_positions = positions + if n_rows % kernel_positions.numel() != 0: raise ValueError( - f"row count {n_rows} not divisible by seq length {S}; " + f"row count {n_rows} not divisible by position rows " + f"{kernel_positions.numel()}; " "expected a [..., S, D] contiguous layout." ) - cos, sin = _build_cos_sin(positions, D // 2, float(theta), x.device) + cos, sin = _build_cos_sin(kernel_positions, D // 2, float(theta), x.device) ctx.save_for_backward(cos, sin) out = _C.rope_apply_sm90(x_2d, cos, sin, 1.0) return out.reshape(x.shape) @@ -70,6 +90,7 @@ class RoPESM90Op: """ op_class = "elementwise" + backend_id = "rlkernel.cuda.rope_sm90" def __init__(self) -> None: if not _EXT_AVAILABLE or not hasattr(_C, "rope_apply_sm90"): diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 3c835deb..060234be 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -95,6 +95,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) + CUDA_RMS_NORM = "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp" PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" # Generic fallback @@ -348,7 +349,10 @@ def __init__(self): OpBackend.TRITON_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ], - "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], + "rms_norm": [ + OpBackend.CUDA_RMS_NORM, + OpBackend.PYTORCH_NATIVE_RMS_NORM, + ], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [ diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 04f0bfdb..12bb878e 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -55,6 +55,9 @@ SplitKVSpec, build_split_kv_runtime_plan_set, ) +from rl_engine.kernels.attention_preprocess import ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS, +) pytestmark = pytest.mark.unit @@ -626,6 +629,8 @@ def _readback(materializer, flat, *, source): split_kv_plan_set=_plan_set(contract, backend=source), source=source, frozen_scope_verified=True, + preprocess_backends=MANDATED_ATTENTION_PREPROCESS_BACKENDS, + preprocess_fallback=False, ) @@ -683,6 +688,8 @@ def test_runtime_readback_mismatch_is_a_fallback(): 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}, @@ -760,6 +767,58 @@ def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): 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( + ("backends", "fallback", "expected_code"), + [ + ( + {"rope": MANDATED_ATTENTION_PREPROCESS_BACKENDS["rope"]}, + False, + BindingErrorCode.ATTENTION_PREPROCESS_MISSING, + ), + ( + {**MANDATED_ATTENTION_PREPROCESS_BACKENDS, "qk_rmsnorm": "vllm.native"}, + False, + BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + ), + ( + dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), + True, + BindingErrorCode.ATTENTION_PREPROCESS_FALLBACK, + ), + ], +) +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, + ) + 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_entrypoint_rejects_unverified_frozen_scope(): diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py new file mode 100644 index 00000000..451ec979 --- /dev/null +++ b/tests/test_attention_preprocess.py @@ -0,0 +1,111 @@ +# 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 ( + H100AttentionPreprocessor, + MANDATED_ATTENTION_PREPROCESS_BACKENDS, +) +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_has_no_native_backend_option(): + 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 False + assert dict(result.backend_ids) == dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) + assert result.readback_fields() == { + "preprocess_backends": dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), + "preprocess_fallback": False, + } + 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_rms_norm.py b/tests/test_rms_norm.py index 43a9cf85..6572603e 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -5,14 +5,17 @@ import torch import torch.nn.functional as F -from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda +from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp, rmsnorm_cuda from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton try: from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE - _HAS_CUDA_RMSNORM = _EXT_AVAILABLE and hasattr(_C, "rmsnorm_forward") + _HAS_CUDA_RMSNORM = _EXT_AVAILABLE and all( + hasattr(_C, name) + for name in ("rmsnorm_forward", "rmsnorm_backward_dx", "rmsnorm_backward_dw") + ) except ImportError: # pragma: no cover - import can fail when the extension is not built. _HAS_CUDA_RMSNORM = False @@ -236,8 +239,12 @@ def test_registry_dispatches_rms_norm(): from rl_engine.kernels.registry import kernel_registry op = kernel_registry.get_op("rms_norm") - assert isinstance(op, NativeRMSNormOp) - assert hasattr(op, "forward") and hasattr(op, "forward_fp32") + if torch.cuda.is_available() and _HAS_CUDA_RMSNORM: + assert isinstance(op, RMSNormCudaOp) + assert hasattr(op, "forward") + else: + assert isinstance(op, NativeRMSNormOp) + assert hasattr(op, "forward") and hasattr(op, "forward_fp32") @requires_cuda From d9e04b90fb7c7654fbebdbe185092767f9fb3d7c Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 17:46:39 +0800 Subject: [PATCH 32/41] fix(attention): align projection collectives and runtime evidence --- .../ws2-attention-cross-config-integration.md | 46 ++-- docs/operators/attention.md | 6 + .../cross_config/attention_binding.py | 233 +++++++++++++++- rl_engine/kernels/attention_preprocess.py | 132 +++++++-- rl_engine/kernels/attention_projection.py | 256 ++++++++++++++++++ tests/test_attention_cross_config_binding.py | 103 ++++++- tests/test_attention_preprocess.py | 11 +- tests/test_attention_projection.py | 97 +++++++ 8 files changed, 830 insertions(+), 54 deletions(-) create mode 100644 rl_engine/kernels/attention_projection.py create mode 100644 tests/test_attention_projection.py diff --git a/docs/design/ws2-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md index 4bda3089..6ddd45a0 100644 --- a/docs/design/ws2-attention-cross-config-integration.md +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -49,18 +49,13 @@ Two placements are load-bearing: 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 boundary +## H100 Attention input and projection boundary -The strict experiment does not use the PyTorch reference operators as an -executable option. `H100AttentionPreprocessor` applies these implementations in -the fixed Qwen3 order: - -1. `RMSNormCudaOp` on Q and K (`rlkernel.cuda.rmsnorm`) -2. `RoPESM90Op` with global `[S]` or per-batch `[B, S]` positions - (`rlkernel.cuda.rope_sm90`) - -There is no Megatron/vLLM-native fallback. The CUDA RoPE path accepts non-contiguous -global positions, including zigzag CP ownership. The launcher passes the returned +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 @@ -75,15 +70,20 @@ readback = AttentionRuntimeReadback( ) ``` -Strict binding rejects a missing backend ID, a runtime-native backend ID, or any -reported fallback. Printing a configured backend without executing it is not -accepted as evidence. - -The boundary starts at projected Q/K/V. Pre-attention model RMSNorm, QKV projection -GEMM, and projection-owned TP/SP All-Gather/Reduce-Scatter are not part of the -Attention operator experiment. The isolated H100 test must therefore capture or -reuse identical projected Q/K/V inputs. Those upstream operators must be aligned -separately before making an end-to-end Megatron-vLLM logprob claim. +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 @@ -144,7 +144,7 @@ collapsing them onto the supported value: | 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/native/fallback QK-Norm or RoPE backend | binding failure | both sides must execute the RL-Kernel CUDA preprocessing path | +| 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 @@ -178,10 +178,8 @@ retired rather than rewritten. Deliberately not in this PR: * launching `torchrun`, initializing process groups, or executing core attention; -* pre-attention model RMSNorm, QKV projection GEMM, and projection-owned TP/SP - communication; isolated tests reuse identical projected Q/K/V; +* 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; -* Transformer Engine calls of any kind (PR4's TE plan is policy and provenance only); * distributed drift benchmarks and report artifacts (#235 PR5); * fused production backend alignment (#235 PR7) and backward (#235 PR8). diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 92d24e6c..592c6f95 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -13,6 +13,12 @@ friends) are validated against. This op covers **only** the softmax attention. Qwen3's QK-Norm and RoPE are applied *before* the call (see the chain), so the `q`, `k` passed in are already normalized and rotated. +For the WS2 Attention experiment, the measured boundary also includes QKV and `o_proj` +projections plus their TP/SP communication contracts. Those projections use native TE or +vLLM callables only after an H100 bitwise probe; otherwise both sides use the deterministic +`DetGemmOp` path with BF16 I/O, FP32 accumulation, ascending-K reduction, and Split-K disabled. +The model input RMSNorm and residual add remain outside this boundary. + ```text q --\ k ----softmax(QKᵀ/√d + mask)·V--> out diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 1007f197..9adba35d 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -54,7 +54,15 @@ SplitKVRuntimePlanSet, validate_split_kv_plan_set_alignment, ) -from rl_engine.kernels.attention_preprocess import MANDATED_ATTENTION_PREPROCESS_BACKENDS +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", @@ -112,6 +120,8 @@ class BindingErrorCode(str, Enum): 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" @dataclass(frozen=True) @@ -125,6 +135,10 @@ class AttentionRuntimeReadback: 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) def __post_init__(self) -> None: if not isinstance(self.contract, AttentionContract): @@ -146,6 +160,21 @@ def __post_init__(self) -> None: 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") + 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: @@ -156,6 +185,11 @@ def __post_init__(self) -> None: "preprocess_backends", MappingProxyType(dict(self.preprocess_backends)), ) + object.__setattr__( + self, + "projection_plans", + MappingProxyType(normalized_projection_plans), + ) @property def split_kv_fallback(self) -> bool: @@ -170,6 +204,12 @@ def to_dict(self) -> dict[str, Any]: "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() }, "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), } @@ -853,6 +893,82 @@ def bind_attention_runtime_readbacks( ) ) missing_scope_evidence.extend(_attention_preprocess_issues(side, readback)) + missing_scope_evidence.extend(_attention_projection_issues(side, readback)) + 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 in ("backend_id", "policy_id", "split_k", "reduction_order"): + if rollout_plan.get(field) != training_plan.get(field): + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"projection.{projection}.{field}", + rollout=rollout_plan.get(field), + training=training_plan.get(field), + message=( + "deterministic projection fallback evidence differs " + "between sides" + ), + ) + ) return bind_attention_contracts( rollout_contract=rollout.contract, training_contract=training.contract, @@ -869,6 +985,13 @@ def bind_attention_runtime_readbacks( 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, + **{ + f"projection.{projection}": dict(plan) + for projection, plan in rollout.projection_plans.items() + }, }, training_recorded_extra={ **{ @@ -876,6 +999,13 @@ def bind_attention_runtime_readbacks( 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, + **{ + f"projection.{projection}": dict(plan) + for projection, plan in training.projection_plans.items() + }, }, ) @@ -899,7 +1029,7 @@ def _attention_preprocess_issues( ), ) ) - elif actual != mandated: + elif actual != mandated and not actual.startswith("native."): issues.append( BindingIssue( code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, @@ -909,22 +1039,115 @@ def _attention_preprocess_issues( training=actual if side == "training" else None, message=( f"{side} executed {actual!r}; " - f"the H100 experiment requires {mandated!r}" + f"the H100 experiment requires {mandated!r} or a verified native backend" ), ) ) - if readback.preprocess_fallback: + 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 QK-Norm or RoPE backend 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, expected in fixed_fields.items(): + actual = plan.get(field) + if actual != expected: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.{field}", + rollout=actual if side == "rollout" else None, + training=actual if side == "training" else None, + message=f"{side} {projection} {field} 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.""" diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index 196d2e18..b72d6829 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -1,18 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Strict H100 QK-Norm and RoPE handoff for WS2 Attention. +"""Bitwise-bound QK-Norm and RoPE handoff for WS2 Attention. -This module intentionally has no runtime-native fallback. A caller either runs -the RL-Kernel CUDA operators and records their identities, or the experiment -fails before Attention executes. +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 from dataclasses import dataclass +import hashlib +import json from types import MappingProxyType -from typing import Any, Mapping +from typing import Any, Callable, Mapping import torch from torch import Tensor @@ -20,6 +24,9 @@ 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, @@ -37,6 +44,9 @@ class AttentionPreprocessResult: 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))) @@ -54,13 +64,31 @@ def readback_fields(self) -> dict[str, Any]: 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 RL-Kernel CUDA QK-Norm then RoPE without silent fallback.""" + """Apply native QK-Norm/RoPE when the H100 probe passes. - def __init__(self, device: torch.device | str | int | None = 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, + policy_id: str = PREPROCESS_POLICY_ID, + ) -> None: if not torch.cuda.is_available(): raise RuntimeError("H100AttentionPreprocessor requires an available CUDA runtime") @@ -89,13 +117,13 @@ def __init__(self, device: torch.device | str | int | None = None) -> None: self.rmsnorm = RMSNormCudaOp() self.rope = RoPESM90Op() - actual_backends = { - "qk_rmsnorm": self.rmsnorm.backend_id, - "rope": self.rope.backend_id, - } - if actual_backends != dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS): - raise RuntimeError(f"unexpected Attention preprocess backends: {actual_backends}") - self.backend_ids = MappingProxyType(actual_backends) + 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, @@ -130,17 +158,78 @@ def forward( theta: float = 1_000_000.0, ) -> AttentionPreprocessResult: _validate_inputs(q, k, q_weight, k_weight, positions, self.device) - q_norm = self.rmsnorm(q, q_weight, eps=eps) - k_norm = self.rmsnorm(k, k_weight, eps=eps) + 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_available = self.native_qk_norm is not None and self.native_rope is not None + if native_available: + try: + q_norm_native = self.native_qk_norm(q, q_weight, eps=eps) + k_norm_native = self.native_qk_norm(k, k_weight, eps=eps) + q_native = self.native_rope(q_norm_native, positions, theta=theta) + k_native = self.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=self.rope(q_norm, positions, theta=theta), - k=self.rope(k_norm, positions, theta=theta), - backend_ids=self.backend_ids, - fallback=False, + 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, @@ -177,4 +266,7 @@ def _validate_inputs( "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..c9bc3043 --- /dev/null +++ b/rl_engine/kernels/attention_projection.py @@ -0,0 +1,256 @@ +# 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 +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] = MappingProxyType({}) + + 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/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 12bb878e..abe55fa3 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -58,6 +58,11 @@ from rl_engine.kernels.attention_preprocess import ( MANDATED_ATTENTION_PREPROCESS_BACKENDS, ) +from rl_engine.kernels.attention_projection import ( + O_PROJ_COLLECTIVE_CONTRACT, + ProjectionPlan, + QKV_COLLECTIVE_CONTRACT, +) pytestmark = pytest.mark.unit @@ -623,6 +628,20 @@ def _statuses(materialization, 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), @@ -631,6 +650,7 @@ def _readback(materializer, flat, *, source): frozen_scope_verified=True, preprocess_backends=MANDATED_ATTENTION_PREPROCESS_BACKENDS, preprocess_fallback=False, + projection_plans=projection_plans, ) @@ -772,6 +792,60 @@ def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): ) +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"), [ @@ -781,14 +855,14 @@ def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): BindingErrorCode.ATTENTION_PREPROCESS_MISSING, ), ( - {**MANDATED_ATTENTION_PREPROCESS_BACKENDS, "qk_rmsnorm": "vllm.native"}, + {**MANDATED_ATTENTION_PREPROCESS_BACKENDS, "qk_rmsnorm": "unknown.backend"}, False, BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, ), ( dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), True, - BindingErrorCode.ATTENTION_PREPROCESS_FALLBACK, + BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, ), ], ) @@ -800,6 +874,7 @@ def test_strict_runtime_readback_rejects_unverified_preprocess_backend( rollout, preprocess_backends=backends, preprocess_fallback=fallback, + preprocess_fallback_reason=("test fallback" if fallback else None), ) training = _readback( MegatronAttentionMaterializer(), @@ -821,6 +896,30 @@ def test_strict_runtime_readback_rejects_unverified_preprocess_backend( 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() diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py index 451ec979..a2693d64 100644 --- a/tests/test_attention_preprocess.py +++ b/tests/test_attention_preprocess.py @@ -37,7 +37,7 @@ def _has_h100_preprocess() -> bool: ) -def test_h100_preprocessor_has_no_native_backend_option(): +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", @@ -83,11 +83,16 @@ def test_h100_preprocessor_executes_cuda_qk_norm_and_zigzag_rope(): q_ref = rope(norm(q, q_weight), positions) k_ref = rope(norm(k, k_weight), positions) - assert result.fallback is False + 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": False, + "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) diff --git a/tests/test_attention_projection.py b/tests/test_attention_projection.py new file mode 100644 index 00000000..146a86bf --- /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 ( + AttentionProjectionOp, + O_PROJ_COLLECTIVE_CONTRACT, + QKV_COLLECTIVE_CONTRACT, + 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" From 46c5692ca267e06a17aa04e715e6f1385f6a8c30 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 18:01:30 +0800 Subject: [PATCH 33/41] fix(attention): use portable projection plan default --- rl_engine/kernels/attention_projection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rl_engine/kernels/attention_projection.py b/rl_engine/kernels/attention_projection.py index c9bc3043..5e4dd775 100644 --- a/rl_engine/kernels/attention_projection.py +++ b/rl_engine/kernels/attention_projection.py @@ -13,7 +13,7 @@ import hashlib import json -from dataclasses import dataclass +from dataclasses import dataclass, field from types import MappingProxyType from typing import Any, Callable, Mapping @@ -91,7 +91,7 @@ class ProjectionPlan: reduction_order: str = "k_ascending" split_k: bool = False policy_id: str = PROJECTION_POLICY_ID - collective: Mapping[str, str] = MappingProxyType({}) + collective: Mapping[str, str] = field(default_factory=dict) def __post_init__(self) -> None: if self.projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: From f55a681dc0ec403f0ca734a2b59acac563ecb970 Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Sun, 16 Aug 2026 16:11:04 +0000 Subject: [PATCH 34/41] fix(attention): satisfy PR4 lint checks --- .../cross_config/attention_binding.py | 21 +++++++++---------- rl_engine/kernels/attention_preprocess.py | 12 ++++++----- rl_engine/kernels/attention_projection.py | 4 +--- tests/test_attention_cross_config_binding.py | 6 ++---- tests/test_attention_preprocess.py | 2 +- tests/test_attention_projection.py | 8 +++---- 6 files changed, 25 insertions(+), 28 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 9adba35d..d23329f6 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -948,21 +948,20 @@ def bind_attention_runtime_readbacks( rollout=rollout_fallback, training=training_fallback, message=( - "QKV/o_proj must use the same native-or-deterministic " - "path on both sides" + "QKV/o_proj must use the same native-or-deterministic " "path on both sides" ), ) ) if rollout_fallback and training_fallback: - for field in ("backend_id", "policy_id", "split_k", "reduction_order"): - if rollout_plan.get(field) != training_plan.get(field): + 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}", - rollout=rollout_plan.get(field), - training=training_plan.get(field), + 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" @@ -1091,17 +1090,17 @@ def _attention_projection_issues( ) ) continue - for field, expected in fixed_fields.items(): - actual = plan.get(field) + 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}", + 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} must be {expected!r}", + message=f"{side} {projection} {field_name} must be {expected!r}", ) ) collective = plan.get("collective") diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index b72d6829..1a3d4446 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -12,16 +12,15 @@ from __future__ import annotations -from dataclasses import dataclass 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" @@ -171,9 +170,12 @@ def forward( q_native = self.native_rope(q_norm_native, positions, theta=theta) k_native = self.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): + 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, diff --git a/rl_engine/kernels/attention_projection.py b/rl_engine/kernels/attention_projection.py index 5e4dd775..a7a91bfe 100644 --- a/rl_engine/kernels/attention_projection.py +++ b/rl_engine/kernels/attention_projection.py @@ -215,9 +215,7 @@ def split_qkv( 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]}" - ) + 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 diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index abe55fa3..a859c1e1 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -55,13 +55,11 @@ SplitKVSpec, build_split_kv_runtime_plan_set, ) -from rl_engine.kernels.attention_preprocess import ( - MANDATED_ATTENTION_PREPROCESS_BACKENDS, -) +from rl_engine.kernels.attention_preprocess import MANDATED_ATTENTION_PREPROCESS_BACKENDS from rl_engine.kernels.attention_projection import ( O_PROJ_COLLECTIVE_CONTRACT, - ProjectionPlan, QKV_COLLECTIVE_CONTRACT, + ProjectionPlan, ) pytestmark = pytest.mark.unit diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py index a2693d64..587c0100 100644 --- a/tests/test_attention_preprocess.py +++ b/tests/test_attention_preprocess.py @@ -7,8 +7,8 @@ import torch from rl_engine.kernels.attention_preprocess import ( - H100AttentionPreprocessor, 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 diff --git a/tests/test_attention_projection.py b/tests/test_attention_projection.py index 146a86bf..b4094351 100644 --- a/tests/test_attention_projection.py +++ b/tests/test_attention_projection.py @@ -6,9 +6,9 @@ import torch from rl_engine.kernels.attention_projection import ( - AttentionProjectionOp, O_PROJ_COLLECTIVE_CONTRACT, QKV_COLLECTIVE_CONTRACT, + AttentionProjectionOp, split_qkv, ) @@ -63,9 +63,9 @@ def test_projection_rejects_native_drift_and_records_reason(): 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) + 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" From 553d7993782a7a36bd98c02dceb7ff10a45672d1 Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Sun, 16 Aug 2026 16:26:58 +0000 Subject: [PATCH 35/41] fix(types): narrow optional attention ops --- rl_engine/kernels/attention_preprocess.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index 1a3d4446..e97b365f 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -162,13 +162,14 @@ def forward( q_det = self.rope(q_norm_det, positions, theta=theta) k_det = self.rope(k_norm_det, positions, theta=theta) - native_available = self.native_qk_norm is not None and self.native_rope is not None - if native_available: + 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 = self.native_qk_norm(q, q_weight, eps=eps) - k_norm_native = self.native_qk_norm(k, k_weight, eps=eps) - q_native = self.native_rope(q_norm_native, positions, theta=theta) - k_native = self.native_rope(k_norm_native, positions, theta=theta) + 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) From 62d73f8162c15bea2e6b507893d6e38c8575d59a Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Mon, 17 Aug 2026 18:47:36 +0800 Subject: [PATCH 36/41] feat(attention): bind shared strict core runtime evidence Signed-off-by: lamentropetion <3051000145@qq.com> --- .../cross_config/attention_binding.py | 114 ++++++++++++++++++ rl_engine/kernels/attention_contract.py | 4 + tests/test_attention_cross_config_binding.py | 82 +++++++++++++ 3 files changed, 200 insertions(+) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index d23329f6..d5f53df4 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -44,6 +44,7 @@ from typing import Any, Optional from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, AttentionContract, AttentionContractError, AttentionDType, @@ -122,6 +123,10 @@ class BindingErrorCode(str, Enum): 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_NATIVE_ARITHMETIC = "ATTENTION_NATIVE_ARITHMETIC" + ATTENTION_CORE_SPLIT_K = "ATTENTION_CORE_SPLIT_K" @dataclass(frozen=True) @@ -139,6 +144,10 @@ class AttentionRuntimeReadback: 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 + native_attention_arithmetic: bool = True + strict_split_kv_policy: str | None = None def __post_init__(self) -> None: if not isinstance(self.contract, AttentionContract): @@ -170,6 +179,20 @@ def __post_init__(self) -> None: 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 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") 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): @@ -211,6 +234,12 @@ def to_dict(self) -> dict[str, Any]: "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, + "native_attention_arithmetic": self.native_attention_arithmetic, + "split_kv_policy": self.strict_split_kv_policy, + }, "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), } @@ -894,6 +923,8 @@ def bind_attention_runtime_readbacks( ) 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( @@ -987,6 +1018,10 @@ def bind_attention_runtime_readbacks( "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.native_attention_arithmetic": rollout.native_attention_arithmetic, + "strict.split_kv_policy": rollout.strict_split_kv_policy, **{ f"projection.{projection}": dict(plan) for projection, plan in rollout.projection_plans.items() @@ -1001,6 +1036,10 @@ def bind_attention_runtime_readbacks( "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.native_attention_arithmetic": training.native_attention_arithmetic, + "strict.split_kv_policy": training.strict_split_kv_policy, **{ f"projection.{projection}": dict(plan) for projection, plan in training.projection_plans.items() @@ -1009,6 +1048,81 @@ def bind_attention_runtime_readbacks( ) +def _strict_attention_core_issues( + side: str, + readback: AttentionRuntimeReadback, +) -> list[BindingIssue]: + if not readback.strict_mode: + return [] + issues = [] + 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.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", + ) + ] + return [] + + def _attention_preprocess_issues( side: str, readback: AttentionRuntimeReadback, diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 206663e7..b4d85f7a 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -18,6 +18,9 @@ _EnumT = TypeVar("_EnumT", bound=Enum) +STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" + + class AttentionContractError(ValueError): """Raised when attention metadata does not describe a valid invocation.""" @@ -1521,6 +1524,7 @@ class AttentionDispatchResult: "SplitKVRuntimePlanEntry", "SplitKVRuntimePlanSet", "SplitKVSpec", + "STRICT_ATTENTION_CORE_ID", "validate_split_kv_alignment", "validate_split_kv_plan_set_alignment", ] diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index a859c1e1..c5c7f2d4 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -45,6 +45,7 @@ ) from rl_engine.alignment.cross_config.schema import MaterializationStatus from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, AttentionContractError, AttentionMode, AttentionRole, @@ -652,6 +653,23 @@ def _readback(materializer, flat, *, source): ) +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, + native_attention_arithmetic=False, + strict_split_kv_policy="disabled", + ) + + def test_configured_contract_without_runtime_readback_is_unobservable(): normalized = { "batch": {"size": 2}, @@ -790,6 +808,70 @@ 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, + ), + ( + {"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) + + def test_strict_runtime_readback_accepts_common_deterministic_preprocess_fallback(): rollout = replace( _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), From 683dd8b8cb84ec06f73270fd6df8a2efc7b7da1e Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Tue, 18 Aug 2026 09:40:17 +0000 Subject: [PATCH 37/41] feat(attention): reuse WS1 PR315 deterministic operators --- .../cuda/attention/deterministic_attention.cu | 54 +++++-- csrc/cuda/gemm/det_gemm_kernel.cu | 67 ++++++--- csrc/ops.cpp | 18 +++ rl_engine/kernels/ops/backward_runtime.py | 51 +++++++ .../ops/cuda/attention/deterministic_attn.py | 30 +++- rl_engine/kernels/ops/cuda/matmul/det_gemm.py | 76 +++++++++- rl_engine/kernels/ops/cuda/norm/rmsnorm.py | 40 ++++-- .../kernels/ops/cuda/rotary_embedding/rope.py | 119 ++++++++++------ rl_engine/kernels/ops/vjp_fp32.py | 133 ++++++++++++++++++ 9 files changed, 491 insertions(+), 97 deletions(-) create mode 100644 rl_engine/kernels/ops/backward_runtime.py create mode 100644 rl_engine/kernels/ops/vjp_fp32.py diff --git a/csrc/cuda/attention/deterministic_attention.cu b/csrc/cuda/attention/deterministic_attention.cu index 973b07a8..aaa70b42 100644 --- a/csrc/cuda/attention/deterministic_attention.cu +++ b/csrc/cuda/attention/deterministic_attention.cu @@ -147,9 +147,8 @@ __global__ void masked_softmax_lse_kernel( } } else { lse_val = row_max + logf(row_sum); - float inv_sum = 1.0f / row_sum; for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { - row[k] *= inv_sum; + row[k] /= row_sum; } } @@ -166,11 +165,11 @@ __global__ void masked_softmax_lse_kernel( constexpr int kPVTileQ = 16; constexpr int kPVTileD = 16; -template +template __global__ void pv_kernel( const float* __restrict__ P, // [B, Hq, Sq, Skv] - const scalar_t* __restrict__ V, // [B, Hkv, Skv, D] - scalar_t* __restrict__ out, // [B, Hq, Sq, D] + const input_t* __restrict__ V, // [B, Hkv, Skv, D] + output_t* __restrict__ out, // [B, Hq, Sq, D] int64_t B, int64_t Hq, int64_t Hkv, int64_t Sq, int64_t Skv, int64_t D) { @@ -185,7 +184,7 @@ __global__ void pv_kernel( const int kv_head = hq / (Hq / Hkv); const float* p_row = P + ((int64_t)b * Hq * Sq * Skv + (int64_t)hq * Sq * Skv + (int64_t)q_idx * Skv); - const scalar_t* v_base = V + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D); + const input_t* v_base = V + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D); float acc = 0.0f; for (int64_t k = 0; k < Skv; ++k) { @@ -193,7 +192,7 @@ __global__ void pv_kernel( } const int64_t out_idx = (int64_t)b * Hq * Sq * D + (int64_t)hq * Sq * D + (int64_t)q_idx * D + d_idx; - out[out_idx] = (scalar_t)acc; + out[out_idx] = (output_t)acc; } void check_deterministic_attention_inputs( @@ -257,13 +256,14 @@ void check_deterministic_attention_inputs( // out: [B, Hq, Sq, D] same dtype as q // lse: [B, Hq, Sq] FP32 // P: [B, Hq, Sq, Skv] FP32 (softmax probabilities, saved for backward) -std::vector deterministic_attention_forward( +std::vector deterministic_attention_forward_impl( torch::Tensor q, torch::Tensor k, torch::Tensor v, bool causal, double scale, - torch::optional key_padding_mask) { + torch::optional key_padding_mask, + bool output_fp32) { check_deterministic_attention_inputs(q, k, v, key_padding_mask); const at::cuda::OptionalCUDAGuard device_guard(at::device_of(q)); @@ -327,7 +327,9 @@ std::vector deterministic_attention_forward( } // --- Launch PV kernel --- - auto out = torch::empty_like(q_contig); + auto out = output_fp32 + ? torch::empty(q_contig.sizes(), q_contig.options().dtype(at::kFloat)) + : torch::empty_like(q_contig); { dim3 block(kPVTileD, kPVTileQ); dim3 grid( @@ -337,11 +339,19 @@ std::vector deterministic_attention_forward( AT_DISPATCH_FLOATING_TYPES_AND2( at::ScalarType::Half, at::ScalarType::BFloat16, q_contig.scalar_type(), "pv_kernel", [&] { - pv_kernel<<>>( - scores.data_ptr(), - v_contig.data_ptr(), - out.data_ptr(), - B, Hq, Hkv, Sq, Skv, D); + if (output_fp32) { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + } else { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + } C10_CUDA_KERNEL_LAUNCH_CHECK(); }); } @@ -349,6 +359,20 @@ std::vector deterministic_attention_forward( return {out, lse, scores}; } +std::vector deterministic_attention_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v, bool causal, double scale, + torch::optional key_padding_mask) { + return deterministic_attention_forward_impl( + q, k, v, causal, scale, key_padding_mask, false); +} + +std::vector deterministic_attention_forward_fp32( + torch::Tensor q, torch::Tensor k, torch::Tensor v, bool causal, double scale, + torch::optional key_padding_mask) { + return deterministic_attention_forward_impl( + q, k, v, causal, scale, key_padding_mask, true); +} + // =========================================================================== // BACKWARD // =========================================================================== diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index 4d9035fc..cd92fc9d 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/csrc/cuda/gemm/det_gemm_kernel.cu @@ -24,15 +24,29 @@ namespace { using nv_bf16 = __nv_bfloat16; +template +__device__ __forceinline__ output_t cast_output(float value); + +template <> +__device__ __forceinline__ nv_bf16 cast_output(float value) { + return __float2bfloat16(value); +} + +template <> +__device__ __forceinline__ float cast_output(float value) { + return value; +} + __host__ __device__ constexpr int cdiv(int a, int b) { return (a + b - 1) / b; } // Naive FP32 scalar kernel (fallback + ground truth). Batch-invariant by // construction: one thread = one output element, fixed ascending K loop. constexpr int NAIVE_TILE = 16; +template __global__ void det_gemm_naive(const nv_bf16* __restrict__ A, const nv_bf16* __restrict__ B, - nv_bf16* __restrict__ C, + output_t* __restrict__ C, int M, int N, int K) { const int row = blockIdx.y * NAIVE_TILE + threadIdx.y; const int col = blockIdx.x * NAIVE_TILE + threadIdx.x; @@ -40,14 +54,15 @@ __global__ void det_gemm_naive(const nv_bf16* __restrict__ A, 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] = cast_output(acc); } -void launch_naive(const nv_bf16* A, const nv_bf16* B, nv_bf16* C, +template +void launch_naive(const nv_bf16* A, const nv_bf16* B, output_t* C, int M, int N, int K, cudaStream_t stream) { dim3 block(NAIVE_TILE, NAIVE_TILE); dim3 grid(cdiv(N, NAIVE_TILE), cdiv(M, NAIVE_TILE)); - det_gemm_naive<<>>(A, B, C, M, N, K); + det_gemm_naive<<>>(A, B, C, M, N, K); } #if defined(RL_KERNEL_ENABLE_SM90) @@ -81,9 +96,10 @@ __device__ __forceinline__ void mma_m16n8k16(const uint32_t A[4], const uint32_t "f"(D[0]), "f"(D[1]), "f"(D[2]), "f"(D[3])); } +template __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap, const __grid_constant__ CUtensorMap bt_tmap, - nv_bf16* __restrict__ C, + output_t* __restrict__ C, int M, int N, int K) { const int tid = threadIdx.x; const int warp = tid / 32; @@ -186,18 +202,19 @@ __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] = cast_output(acc[mi][n][0]); + C[row * N + col + 1] = cast_output(acc[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] = cast_output(acc[mi][n][2]); + C[(row + 8) * N + col + 1] = cast_output(acc[mi][n][3]); } } } } -bool launch_sm90(const nv_bf16* A, const nv_bf16* Bt, nv_bf16* C, +template +bool launch_sm90(const nv_bf16* A, const nv_bf16* Bt, output_t* C, int M, int N, int K, cudaStream_t stream) { if (M % BM != 0 || N % BN != 0 || K % BK != 0) return false; // fall back @@ -207,11 +224,11 @@ bool launch_sm90(const nv_bf16* A, const nv_bf16* Bt, nv_bf16* C, const int smem = STAGES * (BM * BK + BN * BK) * sizeof(nv_bf16) + STAGES * 8; if (smem > 48 * 1024) - cudaFuncSetAttribute(det_gemm_sm90_kernel, + cudaFuncSetAttribute(det_gemm_sm90_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); dim3 grid(cdiv(N, BN), cdiv(M, BM)); - det_gemm_sm90_kernel<<>>(a_tmap, bt_tmap, C, M, N, K); + det_gemm_sm90_kernel<<>>(a_tmap, bt_tmap, C, M, N, K); return true; } #endif // RL_KERNEL_ENABLE_SM90 @@ -232,9 +249,11 @@ void check_in(const torch::Tensor& t, const char* n) { TORCH_CHECK(t.scalar_type() == torch::kBFloat16, n, " must be bf16"); } -torch::Tensor gemm_dispatch(const torch::Tensor& a, const torch::Tensor& b) { +torch::Tensor gemm_dispatch(const torch::Tensor& a, const torch::Tensor& b, + bool output_fp32 = false) { const int M = a.size(0), K = a.size(1), N = b.size(1); - auto c = torch::empty({M, N}, a.options()); + auto options = a.options().dtype(output_fp32 ? torch::kFloat32 : torch::kBFloat16); + auto c = torch::empty({M, N}, options); auto stream = at::cuda::getCurrentCUDAStream(); #if defined(RL_KERNEL_ENABLE_SM90) @@ -249,15 +268,21 @@ torch::Tensor gemm_dispatch(const torch::Tensor& a, const torch::Tensor& b) { a_use = torch::zeros({Mp, K}, a.options()); a_use.narrow(0, 0, M).copy_(a); } - torch::Tensor c_use = (Mp != M) ? torch::empty({Mp, N}, a.options()) : c; + torch::Tensor c_use = (Mp != M) ? torch::empty({Mp, N}, options) : c; auto bt = b.t().contiguous(); // [N,K] - if (launch_sm90(bf16(a_use), bf16(bt), bf16o(c_use), Mp, N, K, stream)) { + const bool launched = output_fp32 + ? launch_sm90(bf16(a_use), bf16(bt), c_use.data_ptr(), Mp, N, K, stream) + : launch_sm90(bf16(a_use), bf16(bt), bf16o(c_use), Mp, N, K, stream); + if (launched) { if (Mp != M) c.copy_(c_use.narrow(0, 0, M)); return c; } } #endif - launch_naive(bf16(a), bf16(b), bf16o(c), M, N, K, stream); + if (output_fp32) + launch_naive(bf16(a), bf16(b), c.data_ptr(), M, N, K, stream); + else + launch_naive(bf16(a), bf16(b), bf16o(c), M, N, K, stream); return c; } @@ -271,6 +296,14 @@ 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"); + return gemm_dispatch(a, b, true); +} + 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 eee328a4..58692de1 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -71,6 +71,7 @@ torch::Tensor lm_head_sm90_forward(torch::Tensor hidden, torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden, torch::Tensor weight, torch::optional bias); +torch::Tensor det_gemm_rowwise_fwd_fp32(torch::Tensor a, torch::Tensor b); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) @@ -90,6 +91,7 @@ torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torc // Batch-Invariant Deterministic GEMM Declarations torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b); +torch::Tensor det_gemm_fwd_fp32(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) @@ -241,6 +243,14 @@ std::vector deterministic_attention_forward( double scale, torch::optional key_padding_mask); +std::vector deterministic_attention_forward_fp32( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + bool causal, + double scale, + torch::optional key_padding_mask); + std::vector deterministic_attention_backward( torch::Tensor grad_output, torch::Tensor q, @@ -338,6 +348,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Single-card SM90 batch-invariant LM-head forward"); m.def("lm_head_sm90_forward_fp32", &lm_head_sm90_forward_fp32, "Single-card SM90 batch-invariant LM-head forward with fp32 output"); + m.def("det_gemm_rowwise_fwd_fp32", &det_gemm_rowwise_fwd_fp32, + "SM90 deterministic rowwise GEMM with FP32 inputs/accumulation/output"); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) @@ -360,6 +372,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // registry Batch-Invariant Deterministic GEMM m.def("det_gemm_fwd", &det_gemm_fwd, "Batch-invariant deterministic GEMM forward (C=A@B)"); + m.def("det_gemm_fwd_fp32", &det_gemm_fwd_fp32, + "Batch-invariant deterministic GEMM forward with FP32 output"); m.def("det_gemm_da", &det_gemm_da, "Batch-invariant deterministic GEMM backward dA (dC@B^T)"); m.def("det_gemm_db", &det_gemm_db, "Batch-invariant deterministic GEMM backward dB (A^T@dC)"); // registry RMSNorm @@ -378,6 +392,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "deterministic_attention_forward", &deterministic_attention_forward, "Deterministic standard softmax attention forward (out, lse)"); + m.def( + "deterministic_attention_forward_fp32", + &deterministic_attention_forward_fp32, + "Deterministic standard softmax attention forward with FP32 output"); m.def( "deterministic_attention_backward", &deterministic_attention_backward, diff --git a/rl_engine/kernels/ops/backward_runtime.py b/rl_engine/kernels/ops/backward_runtime.py new file mode 100644 index 00000000..5cc7d074 --- /dev/null +++ b/rl_engine/kernels/ops/backward_runtime.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime record of the kernel that actually executed a candidate backward.""" + +from __future__ import annotations + +from threading import Lock +from typing import Any + +_LOCK = Lock() +_EVENTS: dict[str, dict[str, Any]] = {} + + +def record_backward( + kind: str, + *, + kernel_id: str, + impl: str, + family: str, +) -> None: + with _LOCK: + previous = _EVENTS.get(kind) + count = 1 if previous is None else int(previous["execution_count"]) + 1 + kernel_ids = tuple(part for part in kernel_id.split("+") if part) + _EVENTS[kind] = { + "kind": kind, + "implementation_ids": list(kernel_ids), + "kernel_ids": list(kernel_ids), + "kernel_id": kernel_id, + "impl": impl, + "family": family, + "execution_count": count, + } + + +def snapshot_backward_runtime() -> dict[str, dict[str, Any]]: + with _LOCK: + return {key: dict(value) for key, value in _EVENTS.items()} + + +def reset_backward_runtime() -> None: + with _LOCK: + _EVENTS.clear() + + +__all__ = [ + "record_backward", + "reset_backward_runtime", + "snapshot_backward_runtime", +] diff --git a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py index 81f80a7f..01c51c99 100644 --- a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py @@ -33,13 +33,18 @@ def forward( causal: bool, scale: float, key_padding_mask: Optional[torch.Tensor], + output_fp32: bool, ) -> tuple[torch.Tensor, torch.Tensor]: q_c = q.contiguous() k_c = k.contiguous() v_c = v.contiguous() mask_c = key_padding_mask.contiguous() if key_padding_mask is not None else None - results = _C.deterministic_attention_forward(q_c, k_c, v_c, causal, float(scale), mask_c) + results = ( + _C.deterministic_attention_forward_fp32(q_c, k_c, v_c, causal, float(scale), mask_c) + if output_fp32 + else _C.deterministic_attention_forward(q_c, k_c, v_c, causal, float(scale), mask_c) + ) out, lse, P = results[0], results[1], results[2] ctx.save_for_backward(q_c, k_c, v_c, P, mask_c) @@ -54,6 +59,8 @@ def forward( def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): q_c, k_c, v_c, P, mask_c = ctx.saved_tensors + if grad_out.dtype != q_c.dtype: + grad_out = grad_out.to(q_c.dtype) dQ, dK, dV = _C.deterministic_attention_backward( grad_out.contiguous(), q_c, @@ -65,7 +72,7 @@ def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): mask_c, ) - return dQ, dK, dV, None, None, None + return dQ, dK, dV, None, None, None, None class DeterministicAttentionOp: @@ -130,10 +137,27 @@ def forward_with_lse( self._validate_inputs(q, k, v, key_padding_mask) resolved_scale = scale if scale is not None else (1.0 / math.sqrt(q.shape[-1])) out, lse = _DeterministicAttentionFn.apply( - q, k, v, causal, resolved_scale, key_padding_mask + q, k, v, causal, resolved_scale, key_padding_mask, False ) return out, lse + 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, + ) -> torch.Tensor: + self._validate_inputs(q, k, v, key_padding_mask) + resolved_scale = scale if scale is not None else (1.0 / math.sqrt(q.shape[-1])) + out, _lse = _DeterministicAttentionFn.apply( + q, k, v, causal, resolved_scale, key_padding_mask, True + ) + return out + @staticmethod def _validate_inputs( q: torch.Tensor, diff --git a/rl_engine/kernels/ops/cuda/matmul/det_gemm.py b/rl_engine/kernels/ops/cuda/matmul/det_gemm.py index 4778be90..c410fbb2 100644 --- a/rl_engine/kernels/ops/cuda/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/cuda/matmul/det_gemm.py @@ -9,22 +9,70 @@ """ import torch +from rl_engine.kernels.ops.backward_runtime import record_backward from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE from rl_engine.utils.logger import logger class _DetGemmFn(torch.autograd.Function): @staticmethod - def forward(ctx, a, b): + def forward(ctx, a, b, output_fp32=False): ctx.save_for_backward(a, b) + if output_fp32: + if not hasattr(_C, "det_gemm_fwd_fp32"): + raise RuntimeError("FP32 deterministic GEMM output requires the rebuilt extension") + return _C.det_gemm_fwd_fp32(a, b) return _C.det_gemm_fwd(a, b) @staticmethod def backward(ctx, grad_out): a, b = ctx.saved_tensors grad_out = grad_out.contiguous() + if grad_out.dtype != torch.bfloat16: + grad_out = grad_out.to(torch.bfloat16) da = _C.det_gemm_da(grad_out, b) if ctx.needs_input_grad[0] else None db = _C.det_gemm_db(a, grad_out) if ctx.needs_input_grad[1] else None + record_backward( + "det_gemm", + kernel_id="rl_engine._C.det_gemm_da+rl_engine._C.det_gemm_db", + impl="cuda_det_gemm", + family="cuda", + ) + return da, db, None + + +class _DetGemmAccumFn(torch.autograd.Function): + @staticmethod + def forward(ctx, a, b): + ctx.save_for_backward(a, b) + if not hasattr(_C, "det_gemm_rowwise_fwd_fp32"): + raise RuntimeError( + "FP32 rowwise deterministic GEMM requires the rebuilt SM90 extension" + ) + return _C.det_gemm_rowwise_fwd_fp32(a, b) + + @staticmethod + def backward(ctx, grad_out): + a, b = ctx.saved_tensors + grad_fp32 = grad_out.contiguous().float() + a_fp32 = a.contiguous().float() + b_fp32 = b.contiguous().float() + da = ( + _C.det_gemm_rowwise_fwd_fp32(grad_fp32, b_fp32.t().contiguous()).to(a.dtype) + if ctx.needs_input_grad[0] + else None + ) + db = ( + _C.det_gemm_rowwise_fwd_fp32(a_fp32.t().contiguous(), grad_fp32).to(b.dtype) + if ctx.needs_input_grad[1] + else None + ) + record_backward( + "det_gemm", + kernel_id=("rl_engine._C.det_gemm_rowwise_fwd_fp32"), + impl="cuda_rowwise_fp32_accum_det_gemm", + family="cuda", + ) return da, db @@ -51,9 +99,31 @@ def __call__(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: "DetGemmOp: compiled _C.det_gemm kernel unavailable; no " "batch-invariant fallback exists. Build the extension first." ) - return _DetGemmFn.apply(a.contiguous(), b.contiguous()) + return _DetGemmFn.apply(a.contiguous(), b.contiguous(), False) + + def forward_fp32(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16, "BF16 only" + assert a.is_cuda and b.is_cuda, "Inputs must be on CUDA device" + if not self.has_hardware_op: + raise RuntimeError("DetGemmOp: compiled CUDA extension unavailable") + return _DetGemmFn.apply(a.contiguous(), b.contiguous(), True) + + def forward_accum_fp32(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + if a.dtype not in (torch.bfloat16, torch.float32) or b.dtype not in ( + torch.bfloat16, + torch.float32, + ): + raise TypeError("FP32-accumulation GEMM requires BF16 or FP32 inputs") + assert a.is_cuda and b.is_cuda, "Inputs must be on CUDA device" + return _DetGemmAccumFn.apply(a.contiguous(), b.contiguous()) + + def parameter_vjp_contributions_fp32(self, *, a, b, grad_output): + del b + rows_a = a.float() + rows_g = grad_output.float() + return {"b": rows_a[:, :, None] * rows_g[:, None, :]} def deterministic_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """Functional entry. a:[M,K] bf16, b:[K,N] bf16 -> [M,N] bf16.""" - return _DetGemmFn.apply(a, b) + return _DetGemmFn.apply(a, b, False) diff --git a/rl_engine/kernels/ops/cuda/norm/rmsnorm.py b/rl_engine/kernels/ops/cuda/norm/rmsnorm.py index 7ac85b70..d4cefae1 100644 --- a/rl_engine/kernels/ops/cuda/norm/rmsnorm.py +++ b/rl_engine/kernels/ops/cuda/norm/rmsnorm.py @@ -1,6 +1,8 @@ import torch +from rl_engine.kernels.ops.backward_runtime import record_backward from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.vjp_fp32 import reduce_rows_fp32, rmsnorm_dweight_rows_fp32 class RMSNormCuda(torch.autograd.Function): @@ -62,7 +64,21 @@ def backward(ctx, grad_out): dx = _C.rmsnorm_backward_dx(dy, x, weight, rstd) - dw = _C.rmsnorm_backward_dw(dy, x, rstd, mask).to(weight.dtype) + # Explicit, shape-independent FP32 left fold. This is slower than + # the chunked extension but preserves the C2 Batch/Chunk reduction order. + rows = rmsnorm_dweight_rows_fp32(x, dy, rstd=rstd) + rows = rows * mask.to(dtype=rows.dtype).unsqueeze(-1) + dw = reduce_rows_fp32(rows).to(weight.dtype) + record_backward( + "rms_norm", + kernel_id=( + "rl_engine._C.rmsnorm_backward_dx" + "+rl_engine.kernels.ops.vjp_fp32.rmsnorm_dweight_rows_fp32" + "+rl_engine.kernels.ops.vjp_fp32.reduce_rows_fp32" + ), + impl="cuda_rmsnorm_dx_declared_fp32_rowfold_dw", + family="cuda", + ) return dx, dw, None, None @@ -79,20 +95,7 @@ def rmsnorm_cuda(x, weight, eps=1e-6, mask=None): class RMSNormCudaOp: """CUDA RMSNorm wrapper compatible with the shared operator harness.""" - backend_id = "rlkernel.cuda.rmsnorm" - - def __init__(self): - required = ( - "rmsnorm_forward", - "rmsnorm_backward_dx", - "rmsnorm_backward_dw", - ) - missing = [name for name in required if not _EXT_AVAILABLE or not hasattr(_C, name)] - if missing: - raise RuntimeError( - "CUDA RMSNorm extension is incomplete; rebuild _C with rmsnorm.cu " - f"(missing: {', '.join(missing)})" - ) + backward_impl = "cuda_rmsnorm_dx_declared_fp32_rowfold_dw" def __call__(self, x, weight, *, eps=1e-6): return self.forward(x, weight, eps=eps) @@ -102,3 +105,10 @@ def forward(self, x, weight, *, eps=1e-6): x_2d = x.contiguous().view(-1, hidden) y_2d = rmsnorm_cuda(x_2d, weight.contiguous(), eps=eps) return y_2d.view_as(x) + + def parameter_vjp_contributions_fp32(self, *, x, weight, grad_output, eps=1e-6): + del weight + x32 = x.float() + rstd = torch.rsqrt(x32.square().mean(dim=-1) + float(eps)) + rows = grad_output.float() * x32 * rstd.unsqueeze(-1) + return {"weight": rows} diff --git a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py index 028e44b9..9a764012 100644 --- a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py +++ b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py @@ -18,69 +18,96 @@ def _build_cos_sin(positions: Tensor, half: int, theta: float, device: torch.device): - """fp32 cos/sin rows, identical math to NativeRoPEOp.""" + """fp32 cos/sin caches of shape [S, half], identical math to NativeRoPEOp.""" inv_freq = 1.0 / (theta ** (torch.arange(0, half, dtype=torch.float32, device=device) / half)) pos = positions.to(device=device, dtype=torch.float32).reshape(-1, 1) freqs = pos * inv_freq # [S, half] return freqs.cos().contiguous(), freqs.sin().contiguous() -class _RoPEFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: - D = x.shape[-1] - if D % 2 != 0: - raise ValueError(f"RoPE head_dim must be even, got {D}") - if positions.dim() not in (1, 2): - raise ValueError("positions must have shape [S] or [B, S]") - S = positions.shape[-1] - if S == 0: - raise ValueError("positions must not be empty") - if x.shape[-2] != S: - raise ValueError(f"x sequence length {x.shape[-2]} does not match positions length {S}") +def _rope_table(x: Tensor, positions: Tensor, theta: float) -> tuple[Tensor, Tensor, Tensor]: + """Build (x_2d, cos, sin) for [S] or [B, S] positions. See Triton RoPE.""" + D = x.shape[-1] + if D % 2 != 0: + raise ValueError(f"RoPE head_dim must be even, got {D}") + if positions.dim() == 1: + table_len = int(positions.shape[0]) x_2d = x.contiguous().reshape(-1, D) - n_rows = x_2d.shape[0] - if positions.dim() == 2: - batch = positions.shape[0] - if x.dim() < 3 or x.shape[0] != batch: - raise ValueError( - f"x batch size {x.shape[0]} does not match positions batch size {batch}" - ) - rows_per_token = n_rows // (batch * S) - if rows_per_token * batch * S != n_rows: - raise ValueError("x rows are incompatible with [B, S] positions") - # The CUDA kernel accepts one fp32 cos/sin row per flattened x row. - # Expanding positions preserves arbitrary global/zigzag indices while - # keeping the arithmetic inside the precompiled deterministic kernel. - kernel_positions = ( - positions[:, None, :].expand(batch, rows_per_token, S).contiguous().reshape(-1) - ) - else: - kernel_positions = positions - if n_rows % kernel_positions.numel() != 0: + if x_2d.shape[0] % table_len != 0: raise ValueError( - f"row count {n_rows} not divisible by position rows " - f"{kernel_positions.numel()}; " + f"row count {x_2d.shape[0]} not divisible by seq length {table_len}; " "expected a [..., S, D] contiguous layout." ) - cos, sin = _build_cos_sin(kernel_positions, D // 2, float(theta), x.device) + cos, sin = _build_cos_sin(positions, D // 2, float(theta), x.device) + return x_2d, cos, sin + if positions.dim() != 2: + raise ValueError(f"positions must be [S] or [B, S], got shape {tuple(positions.shape)}") + batch, seq = positions.shape + if x.shape[0] != batch or x.shape[-2] != seq: + raise ValueError( + f"positions {tuple(positions.shape)} is incompatible with x {tuple(x.shape)}; " + "expected x [B, ..., S, D]" + ) + if x.dim() == 4: + x_2d = x.permute(1, 0, 2, 3).contiguous().reshape(-1, D) + elif x.dim() == 3: + x_2d = x.contiguous().reshape(-1, D) + else: + raise ValueError( + f"RoPE [B, S] positions require x [B, S, D] or [B, H, S, D], got {x.dim()}D" + ) + table_len = batch * seq + if x_2d.shape[0] % table_len != 0: + raise ValueError(f"row count {x_2d.shape[0]} not divisible by B*S={table_len}") + cos, sin = _build_cos_sin(positions.reshape(-1), D // 2, float(theta), x.device) + return x_2d, cos, sin + + +def _restore_rope(out_2d: Tensor, x: Tensor, positions: Tensor) -> Tensor: + if positions.dim() == 1 or x.dim() != 4: + return out_2d.reshape(x.shape) + heads, batch, seq, dim = x.shape[1], x.shape[0], x.shape[2], x.shape[3] + return out_2d.reshape(heads, batch, seq, dim).permute(1, 0, 2, 3).contiguous() + + +class _RoPEFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: + x_2d, cos, sin = _rope_table(x, positions, theta) ctx.save_for_backward(cos, sin) - out = _C.rope_apply_sm90(x_2d, cos, sin, 1.0) - return out.reshape(x.shape) + ctx.x_shape = tuple(x.shape) + ctx.pos_dim = positions.dim() + out_2d = _C.rope_apply_sm90(x_2d, cos, sin, 1.0) + return _restore_rope(out_2d, x, positions) @staticmethod def backward(ctx, grad_out: Tensor): cos, sin = ctx.saved_tensors grad_x = None if ctx.needs_input_grad[0]: - D = grad_out.shape[-1] - g_2d = grad_out.contiguous().reshape(-1, D) - # Inverse rotation: same kernel with the sine negated. - grad_x = _C.rope_apply_sm90(g_2d, cos, sin, -1.0).reshape(grad_out.shape) - # Inputs: x, positions, theta. + if ctx.pos_dim == 2 and len(ctx.x_shape) == 4: + g_2d = grad_out.permute(1, 0, 2, 3).contiguous().reshape(-1, ctx.x_shape[-1]) + out_2d = _C.rope_apply_sm90(g_2d, cos, sin, -1.0) + heads, batch, seq, dim = ( + ctx.x_shape[1], + ctx.x_shape[0], + ctx.x_shape[2], + ctx.x_shape[3], + ) + grad_x = out_2d.reshape(heads, batch, seq, dim).permute(1, 0, 2, 3).contiguous() + else: + g_2d = grad_out.contiguous().reshape(-1, grad_out.shape[-1]) + grad_x = _C.rope_apply_sm90(g_2d, cos, sin, -1.0).reshape(grad_out.shape) return grad_x, None, None +def _is_hopper(device: torch.device) -> bool: + try: + return torch.cuda.get_device_capability(device)[0] == 9 + except Exception: + return False + + class RoPESM90Op: """Custom CUDA RoPE op for SM90 (GPT-NeoX rotate-half), differentiable w.r.t. ``x``. @@ -90,7 +117,6 @@ class RoPESM90Op: """ op_class = "elementwise" - backend_id = "rlkernel.cuda.rope_sm90" def __init__(self) -> None: if not _EXT_AVAILABLE or not hasattr(_C, "rope_apply_sm90"): @@ -106,4 +132,9 @@ def __call__(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) def forward(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: if x.device.type != "cuda": raise RuntimeError(f"RoPESM90Op requires a CUDA tensor, got device '{x.device}'.") + if not _is_hopper(x.device): + raise RuntimeError( + "RoPESM90Op requires Hopper (SM90) CUDA; " + f"got compute capability {torch.cuda.get_device_capability(x.device)}" + ) return _RoPEFunction.apply(x, positions, theta) diff --git a/rl_engine/kernels/ops/vjp_fp32.py b/rl_engine/kernels/ops/vjp_fp32.py new file mode 100644 index 00000000..65acdfa6 --- /dev/null +++ b/rl_engine/kernels/ops/vjp_fp32.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Declared row-local FP32 VJPs. No batched torch.matmul / cuBLAS. + +Each output row is an independent GEMV or outer product. Parameter reductions +walk rows in the caller's order so C10 can re-aggregate by logical token. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import torch + +BACKWARD_IMPL = "row_local_fp32_vjp" + + +def row_local_linear_dx_fp32(grad_output: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """dX[t] = grad[t] @ weight, one GEMV per row.""" + + rows = grad_output.reshape(-1, grad_output.size(-1)).float() + weight_f = weight.float() + out_rows = torch.empty( + (rows.shape[0], weight_f.shape[1]), device=rows.device, dtype=torch.float32 + ) + weight_t = weight_f.t().contiguous() + for index in range(rows.shape[0]): + out_rows[index] = torch.mv(weight_t, rows[index]) + return out_rows.reshape(*grad_output.shape[:-1], weight_f.shape[1]) + + +def row_local_linear_dw_fp32(grad_output: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor: + """dW = sum_t outer(grad[t], hidden[t]) in physical row order.""" + + grad_rows = grad_output.reshape(-1, grad_output.size(-1)).float() + hidden_rows = hidden.reshape(-1, hidden.size(-1)).float() + if grad_rows.shape[0] != hidden_rows.shape[0]: + raise ValueError(f"grad rows {grad_rows.shape[0]} != hidden rows {hidden_rows.shape[0]}") + dweight = torch.zeros( + (grad_rows.shape[1], hidden_rows.shape[1]), + device=grad_rows.device, + dtype=torch.float32, + ) + for index in range(grad_rows.shape[0]): + dweight.addmm_(grad_rows[index].unsqueeze(1), hidden_rows[index].unsqueeze(0)) + return dweight + + +def row_local_bias_fp32(grad_output: torch.Tensor) -> torch.Tensor: + rows = grad_output.reshape(-1, grad_output.size(-1)).float() + acc = torch.zeros((rows.shape[1],), device=rows.device, dtype=torch.float32) + for index in range(rows.shape[0]): + acc = acc + rows[index] + return acc + + +def rmsnorm_dweight_rows_fp32( + x: torch.Tensor, + grad_output: torch.Tensor, + *, + rstd: torch.Tensor | None = None, + eps: float = 1e-6, +) -> torch.Tensor: + """Per-row dweight contributions, shape [..., H].""" + + x32 = x.float() + grad32 = grad_output.float() + if rstd is None: + rstd = torch.rsqrt(x32.square().mean(dim=-1) + float(eps)) + else: + rstd = rstd.float() + return grad32 * x32 * rstd.unsqueeze(-1) + + +def reduce_rows_fp32(rows: torch.Tensor) -> torch.Tensor: + """Left-fold dim 0 in FP32. Deterministic for a fixed row order.""" + + flat = rows.reshape(rows.shape[0], -1).float() + acc = torch.zeros((flat.shape[1],), device=flat.device, dtype=torch.float32) + for index in range(flat.shape[0]): + acc = acc + flat[index] + return acc.reshape(rows.shape[1:]) + + +def reduce_keyed_rows_fp32( + contributions: Mapping[tuple[str, int], torch.Tensor], +) -> torch.Tensor: + if not contributions: + raise RuntimeError("no logical-token contributions to reduce") + keys = sorted(contributions) + acc = contributions[keys[0]].float().clone() + for key in keys[1:]: + acc = acc + contributions[key].float() + return acc + + +def reduce_keyed_outers_fp32( + rows_g: Mapping[tuple[str, int], torch.Tensor], + rows_x: Mapping[tuple[str, int], torch.Tensor], +) -> torch.Tensor: + keys = sorted(set(rows_g) | set(rows_x)) + if not keys or set(rows_g) != set(rows_x): + raise RuntimeError("logical-token sets for outer-product VJP do not match") + first_g = rows_g[keys[0]].float() + first_x = rows_x[keys[0]].float() + acc = torch.outer(first_g, first_x) + for key in keys[1:]: + acc = acc + torch.outer(rows_g[key].float(), rows_x[key].float()) + return acc + + +def merge_keyed( + target: dict[tuple[str, int], torch.Tensor], + source: Mapping[tuple[str, int], torch.Tensor], +) -> None: + overlap = set(target) & set(source) + if overlap: + raise RuntimeError(f"logical token collision: {sorted(overlap)[:4]}") + target.update(source) + + +__all__ = [ + "BACKWARD_IMPL", + "merge_keyed", + "reduce_keyed_outers_fp32", + "reduce_keyed_rows_fp32", + "reduce_rows_fp32", + "rmsnorm_dweight_rows_fp32", + "row_local_bias_fp32", + "row_local_linear_dw_fp32", + "row_local_linear_dx_fp32", +] From 2e5cc49ff200e5a3d1ac306d5bf7ade307b01c0d Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Tue, 18 Aug 2026 10:47:50 +0000 Subject: [PATCH 38/41] feat(attention): add unified ablation matrix wrapper --- .../kernels/ops/pytorch/attention/__init__.py | 13 +- .../kernels/ops/pytorch/attention/ablation.py | 517 ++++++++++++++++++ rl_engine/kernels/registry.py | 38 ++ tests/test_attention_ablation.py | 144 +++++ 4 files changed, 711 insertions(+), 1 deletion(-) create mode 100644 rl_engine/kernels/ops/pytorch/attention/ablation.py create mode 100644 tests/test_attention_ablation.py 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..7d0d6dad --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -0,0 +1,517 @@ +# 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, + 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 + 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") + 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, + 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 + 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.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) + if cfg.deterministic and selected_id == "native": + raise AttentionContractError( + "deterministic=True cannot execute an unverified native Attention backend" + ) + + 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 = self._invoke(selected, q, k, v, call_kwargs, contract) + if cfg.validate: + self._validate_outputs(out, lse, q, contract) + + 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, + "communication_backend": cfg.communication_backend, + "split_kv": contract.split_kv.to_dict(), + "actual_split_kv": _actual_split_provenance(contract), + "reduction": _reduction_provenance(contract), + "contract_fingerprint": _contract_fingerprint(contract), + "return_lse": cfg.return_lse, + "return_gradients": cfg.return_gradients, + } + 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, + ) -> 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 + or 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() + 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]: + 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) + if isinstance(result, AttentionAblationResult): + out, lse = result.out, result.lse + elif isinstance(result, tuple) and len(result) == 2: + out, lse = result + else: + raise AttentionContractError( + "Attention backend must return (out, lse) or AttentionAblationResult" + ) + if not isinstance(out, Tensor) or not isinstance(lse, Tensor): + raise AttentionContractError("Attention backend returned non-tensor output or LSE") + return out, lse + + @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", + "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 _actual_split_provenance(contract: AttentionContract) -> dict[str, Any]: + plan = contract.split_kv.resolve( + contract.sharding.global_sequence_length, + backend=BACKEND_ID, + ) + 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/registry.py b/rl_engine/kernels/registry.py index 060234be..d9b4541c 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -185,6 +185,44 @@ 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_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-v1", + ), + 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", + ), ) diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py new file mode 100644 index 00000000..6d26f15e --- /dev/null +++ b/tests/test_attention_ablation.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.attention_contract import ( + 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.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()) + assert result.backend_id == "test.attention.backend" + assert torch.equal(result.out, q) + + +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 From a3c3d1268f47a9c2accf4793124a2b9ce8d91603 Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Tue, 18 Aug 2026 10:52:29 +0000 Subject: [PATCH 39/41] fix(attention): report selected core provenance --- .../kernels/ops/pytorch/attention/ablation.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index 7d0d6dad..0b0b5106 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -260,10 +260,16 @@ def apply( "semantic_operator": "attention", "backend_id": selected_id, "deterministic": cfg.deterministic, - "strict_core_id": cfg.strict_core_id if cfg.deterministic else None, + "strict_core_id": ( + cfg.strict_core_id + if cfg.deterministic and selected_id == BACKEND_ID + else None + ), + "core_id": selected_id, + "backend_deterministic": selected_id in {BACKEND_ID, REFERENCE_BACKEND_ID}, "communication_backend": cfg.communication_backend, "split_kv": contract.split_kv.to_dict(), - "actual_split_kv": _actual_split_provenance(contract), + "actual_split_kv": _actual_split_provenance(contract, backend=selected_id), "reduction": _reduction_provenance(contract), "contract_fingerprint": _contract_fingerprint(contract), "return_lse": cfg.return_lse, @@ -484,10 +490,14 @@ def _callable_backend_id(value: Any) -> str: return f"injected.{type(value).__module__}.{type(value).__qualname__}" -def _actual_split_provenance(contract: AttentionContract) -> dict[str, Any]: +def _actual_split_provenance( + contract: AttentionContract, + *, + backend: str = BACKEND_ID, +) -> dict[str, Any]: plan = contract.split_kv.resolve( contract.sharding.global_sequence_length, - backend=BACKEND_ID, + backend=backend, ) return plan.to_dict() From 43e67c5a37df7e2d4991f717d9b282f255275bc5 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 18 Aug 2026 21:47:09 +0800 Subject: [PATCH 40/41] feat(attention): enforce canonical bitwise ablation core --- .../cross_config/attention_binding.py | 33 +++++ rl_engine/kernels/attention_contract.py | 2 + .../kernels/ops/pytorch/attention/ablation.py | 46 ++++++- .../ops/pytorch/attention/cp_attention.py | 127 ++++++++++++++++-- rl_engine/kernels/registry.py | 3 +- tests/test_attention_ablation.py | 41 +++++- tests/test_attention_cross_config_binding.py | 6 + tests/test_cp_attention.py | 26 ++++ 8 files changed, 264 insertions(+), 20 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index d5f53df4..57a54205 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -45,6 +45,7 @@ from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, AttentionContract, AttentionContractError, AttentionDType, @@ -125,6 +126,7 @@ class BindingErrorCode(str, Enum): 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" @@ -146,6 +148,7 @@ class AttentionRuntimeReadback: 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 @@ -185,6 +188,10 @@ def __post_init__(self) -> None: 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 { @@ -237,6 +244,7 @@ def to_dict(self) -> dict[str, Any]: "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, }, @@ -1020,6 +1028,7 @@ def bind_attention_runtime_readbacks( "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, **{ @@ -1038,6 +1047,7 @@ def bind_attention_runtime_readbacks( "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, **{ @@ -1067,6 +1077,18 @@ def _strict_attention_core_issues( ), ) ) + 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( @@ -1120,6 +1142,17 @@ def _strict_attention_core_pair_issues( 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", + ) + ] return [] diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index b4d85f7a..8b5f588b 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -19,6 +19,7 @@ STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" +STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks" class AttentionContractError(ValueError): @@ -1525,6 +1526,7 @@ class AttentionDispatchResult: "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/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index 0b0b5106..de1a8d5c 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -24,6 +24,7 @@ from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, AttentionContract, AttentionContractError, AttentionDType, @@ -50,6 +51,7 @@ class AttentionAblationConfig: 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: @@ -64,6 +66,8 @@ def __post_init__(self) -> None: 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()) @@ -229,6 +233,15 @@ 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: raise AttentionContractError("dout is required when return_gradients=True") if dout is not None and dout.shape != q.shape: @@ -240,6 +253,16 @@ def apply( 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) @@ -262,14 +285,23 @@ def apply( "deterministic": cfg.deterministic, "strict_core_id": ( cfg.strict_core_id - if cfg.deterministic and selected_id == BACKEND_ID + if cfg.deterministic else None ), - "core_id": selected_id, - "backend_deterministic": selected_id in {BACKEND_ID, REFERENCE_BACKEND_ID}, + "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, backend=selected_id), + "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, @@ -330,7 +362,7 @@ def _reference_backend(self) -> Any: DeterministicCPAttentionReferenceOp, ) - self.reference = DeterministicCPAttentionReferenceOp() + self.reference = DeterministicCPAttentionReferenceOp(strict_bitwise=True) return self.reference @staticmethod @@ -461,6 +493,7 @@ def _resolve_config( "return_lse", "return_gradients", "strict_core_id", + "strict_schedule", "validate", ) } @@ -493,10 +526,11 @@ def _callable_backend_id(value: Any) -> str: def _actual_split_provenance( contract: AttentionContract, *, + total_kv_tokens: int, backend: str = BACKEND_ID, ) -> dict[str, Any]: plan = contract.split_kv.resolve( - contract.sharding.global_sequence_length, + total_kv_tokens, backend=backend, ) return plan.to_dict() diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 0062dac4..194c80d5 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -208,6 +208,11 @@ class DeterministicCPAttentionReferenceOp: 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, @@ -336,18 +341,31 @@ def forward_with_lse( resolved_output_dtype = q.dtype if output_dtype is None else output_dtype _validate_output_dtype(resolved_output_dtype) - 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, - ) + 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 @@ -381,6 +399,91 @@ def forward_fp32_with_lse( 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, diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index d9b4541c..edd1c535 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -198,6 +198,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "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, @@ -205,7 +206,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "rl_engine.kernels.ops.pytorch.attention.ablation.AttentionAblationOp" ), fallback_policy=OperatorFallbackPolicy.ERROR, - version_or_build_fingerprint="AttentionAblationOp-v1", + version_or_build_fingerprint="AttentionAblationOp-bitwise-v2", ), OperatorBackendDescriptor( semantic_op="attention", diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py index 6d26f15e..7023c6a5 100644 --- a/tests/test_attention_ablation.py +++ b/tests/test_attention_ablation.py @@ -6,6 +6,8 @@ import torch from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, AttentionContract, AttentionContractError, AttentionDType, @@ -75,6 +77,8 @@ def test_attention_wrapper_has_unified_result_and_provenance(): 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) @@ -88,7 +92,14 @@ 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()) + result = AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + backend=FakeBackend(), + deterministic=False, + ) assert result.backend_id == "test.attention.backend" assert torch.equal(result.out, q) @@ -142,3 +153,31 @@ def test_attention_backend_is_registered_for_pr230_semantic_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_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index c5c7f2d4..282e5f39 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -46,6 +46,7 @@ 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, @@ -665,6 +666,7 @@ def _strict_readback(materializer, flat, *, source): 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", ) @@ -819,6 +821,10 @@ def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): {"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, diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index cc93874b..4e58c7ba 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -688,3 +688,29 @@ def test_gapped_partial_ranges_raise(): 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) From fb58a1c2bd62d69db8e69c5b4777278c37b97f48 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 18 Aug 2026 23:25:30 +0800 Subject: [PATCH 41/41] fix(attention): fail closed without production CP backend --- .../cross_config/attention_binding.py | 85 ++++++++++++++ .../kernels/ops/pytorch/attention/ablation.py | 108 ++++++++++++++++-- tests/test_attention_ablation.py | 79 +++++++++++++ tests/test_attention_cross_config_binding.py | 51 +++++++++ 4 files changed, 311 insertions(+), 12 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 57a54205..458a386e 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -129,6 +129,8 @@ class BindingErrorCode(str, Enum): 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) @@ -151,6 +153,9 @@ class AttentionRuntimeReadback: 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): @@ -200,6 +205,14 @@ def __post_init__(self) -> None: "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): @@ -248,6 +261,11 @@ def to_dict(self) -> dict[str, Any]: "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(), } @@ -1031,6 +1049,9 @@ def bind_attention_runtime_readbacks( "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() @@ -1050,6 +1071,9 @@ def bind_attention_runtime_readbacks( "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() @@ -1065,6 +1089,45 @@ def _strict_attention_core_issues( 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( @@ -1153,6 +1216,28 @@ def _strict_attention_core_pair_issues( 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 [] diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index de1a8d5c..e1719085 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -148,6 +148,7 @@ def __init__( 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(): @@ -155,6 +156,10 @@ def __init__( 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__( @@ -248,7 +253,12 @@ def apply( 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) + 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" @@ -270,9 +280,17 @@ def apply( 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 = self._invoke(selected, q, k, v, call_kwargs, contract) + 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: @@ -307,6 +325,10 @@ def apply( "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, @@ -328,6 +350,8 @@ def _select_backend( 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) @@ -342,11 +366,15 @@ def _select_backend( 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 - or q.device.type != "cuda" - or torch.version.hip is not None - ): + 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 ( @@ -420,7 +448,7 @@ def _invoke( v: Tensor, kwargs: Mapping[str, Any], contract: AttentionContract, - ) -> tuple[Tensor, Tensor]: + ) -> tuple[Tensor, Tensor, dict[str, Any]]: method = getattr(backend, "forward_with_lse", None) if not callable(method): method = getattr(backend, "apply", None) @@ -436,17 +464,27 @@ def _invoke( "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: - raise AttentionContractError( - "Attention backend must return (out, lse) or AttentionAblationResult" - ) + 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 + return out, lse, backend_provenance @staticmethod def _backward( @@ -523,6 +561,52 @@ def _callable_backend_id(value: Any) -> str: 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, *, diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py index 7023c6a5..70dbd2ca 100644 --- a/tests/test_attention_ablation.py +++ b/tests/test_attention_ablation.py @@ -2,6 +2,8 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch @@ -104,6 +106,83 @@ 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(): + 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)) diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 282e5f39..c1c8d35d 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -651,6 +651,9 @@ def _readback(materializer, flat, *, source): preprocess_backends=MANDATED_ATTENTION_PREPROCESS_BACKENDS, preprocess_fallback=False, projection_plans=projection_plans, + actual_backend=f"reference.{source}", + communication_backend="none", + production_ready=False, ) @@ -669,6 +672,9 @@ def _strict_readback(materializer, flat, *, source): 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, ) @@ -876,6 +882,51 @@ def test_strict_runtime_readback_accepts_shared_no_split_k_core(): 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():