diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..a4c1fe89 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" @@ -76,6 +84,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/.gitignore b/.gitignore index edc4b005..c3b8d4a6 100644 --- a/.gitignore +++ b/.gitignore @@ -207,5 +207,8 @@ marimo/_static/ marimo/_lsp/ __marimo__/ +# Cross-configuration alignment local run artifacts +/runs/ + # Local dev notes (not for upstream) _dev_notes/ 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 c83db1ce..00000000 Binary files a/docs/assets/ws2-cross-config-before-after.png and /dev/null differ 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/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-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md new file mode 100644 index 00000000..4bda3089 --- /dev/null +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -0,0 +1,187 @@ +# WS2 Attention Cross-Configuration Integration + +Implements PR4 of [#235](https://github.com/RL-Align/RL-Kernel/issues/235): wiring the +CP attention path into the cross-configuration planner/runtime for the Qwen3-8B +TP=2 CP=2 BF16 target. + +Builds on [#236](https://github.com/RL-Align/RL-Kernel/pull/236) (attention contract +and dispatch metadata), [#238](https://github.com/RL-Align/RL-Kernel/pull/238) +(deterministic CP reference) and [#230](https://github.com/RL-Align/RL-Kernel/pull/230) +(cross-configuration framework). + +## What "bind to the same contract" means here + +The PR4 acceptance criteria say rollout and training descriptors must "bind to the +same semantic attention contract". Under the frozen deployment the two sides can +never produce identical `AttentionContract` instances: + +| | training (Megatron) | rollout (vLLM) | +| --- | --- | --- | +| mode | full-sequence prefill | chunked prefill, later decode | +| CP | `context_parallel_size`, whole forward | `prefill_context_parallel_size`, prefill only | +| KV | no paging | paged KV with a block table | +| backend vocabulary | `AttnBackend{flash,fused,unfused,local,auto}` | `AttentionBackendEnum` | + +Read literally, the criterion is unsatisfiable. It is therefore implemented as three +tiers, in `rl_engine/alignment/cross_config/attention_binding.py`: + +| tier | fields | rule | failure | +| --- | --- | --- | --- | +| `IDENTICAL` | checkpoint/model/token identity plus complete TP/CP GQA head and sequence ownership | equal bit for bit | `comparable=False`; no drift number from the pair means anything | +| `SEMANTIC` | reduction contract, dtype, exported LSE, first-class Split-KV request, complete actual Split-KV plan sets, CUDA QK-Norm/RoPE identity, cross-side determinism | both sides equal **and** equal to the WS2 mandate | fail closed | +| `RECORDED` | `mode`, `backend_id`, `reduction.engine`, RoPE materialization state and fusion boundary, KV-cache paging | free to differ | recorded into provenance and measured | + +Two placements are load-bearing: + +* **`reduction.engine` is `RECORDED`, not `SEMANTIC`.** Training may run the in-op + deterministic reference while rollout runs a Transformer Engine merge oracle. + Forcing them equal would defeat the oracle comparison that #235 PR2/PR3/PR5/PR6 + depend on. +* **TP/CP topology is `IDENTICAL`, not `RECORDED`.** TP selects local Qwen3 GQA head + ownership and CP selects local sequence ownership. Different topology is a + different local attention problem, not a backend detail. +* **`reduction.order`, `reduction.acc_dtype`, and actual Split-KV schedules are + `SEMANTIC`.** Both runtimes must export the complete batch x TP x CP x KV-owner + plan set, including logical boundaries, merge order, FP32 accumulation, final + downcast, and fallback state. Configured policy alone never passes strict binding. + +`comparable` and `passed` are separate flags. A pair with mismatched identity is not +comparable. A pair that is comparable but violates the reduction mandate is still +rejected -- the drift would be real but attributable to the wrong thing. + +## H100 Attention input 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 +them, because the two frameworks mean different things by "deterministic": + +| | Megatron `deterministic_mode` | vLLM `VLLM_BATCH_INVARIANT` | +| --- | --- | --- | +| `NCCL_ALGO` | asserts membership in a five-value set | hard-sets `allreduce:tree` | +| `NCCL_PROTO`, channels, threads | not managed | hard-set (`Simple`, `1`, `1`) | +| TF32 | **not managed at all** | disabled (`fp32_precision="ieee"`) | +| BF16 reduced-precision reduction | not managed | disabled | +| cuBLAS workspace / BLAS library | not managed | `:4096:8`, cuBLASLt | +| GEMM | cuBLAS / TE | Triton `matmul_persistent` | +| FlashAttention | forbidden | permitted | + +`NCCL_ALGO`, `NCCL_PROTO` and `CUBLAS_WORKSPACE_CONFIG` change arithmetic, so a +mismatch there is blocking. The remaining differences -- including the TF32 and +BF16-reduction asymmetry, which under a pure BF16 GEMM path does not fire -- are +recorded so the asymmetry appears in every artifact rather than being invisible. + +## Runtime adapters + +Before this PR the only `RuntimeMaterializer` in the repository was +`CpuSmokeMaterializer` over a synthetic CPU model, and every named scenario +(`S1`/`S2`/`S3`) was planning-only. This PR adds the first two framework-shaped +adapters: + +* `adapters/megatron.py` -- `MegatronProvenanceAdapter` (construction and + distributed-context fingerprints, determinism probe, frozen-scope assertions) and + `MegatronAttentionMaterializer`. +* `adapters/vllm.py` -- `VllmProvenanceAdapter` (including diagnostic vLLM split + limits) and `VllmRolloutMaterializer`. +* `AttentionRuntimeReadback` -- the explicit handoff from an executed engine. It + carries the reconstructed actual contract, actual knob values, frozen-scope + verification, executed CUDA QK-Norm/RoPE identities and fallback state, and the + complete Split-KV runtime plan set. + +Constructing a contract is not runtime verification. Without a readback, adapter +applications are `UNOBSERVABLE`; only matching values reconstructed from a real +Megatron or vLLM execution are `APPLIED`. `bind_attention_runtime_readbacks` is the +strict public entry point used after both framework launchers collect that evidence. + +Neither module imports `megatron` or `vllm`; configs are duck-typed, so the binding +rules are exercised on CPU in CI rather than only on a 2-node cluster. + +## Fail closed, never substitute + +`unsupported_reduction_reason` rejects requests that #236 cannot express, instead of +collapsing them onto the supported value: + +| request | status | why | +| --- | --- | --- | +| `attention.reduction_order=arrival` | `UNSUPPORTED` | the control group must stay distinguishable from the treatment | +| `attention.reduction_downcast_at=per_block` | `UNSUPPORTED` | `DowncastPoint` declares only `final_write` | +| `attention.reduction_engine=te_oracle` | `UNSUPPORTED` | the TE merge oracle lands in #235 PR2/PR3; PR4's TE plan is provenance only | +| `attention.reduction_acc_dtype=bf16` | `UNSUPPORTED` | the CP `(out, lse)` merge accumulates in FP32 | +| configured contract without runtime readback | `UNOBSERVABLE` | requested values do not prove what executed | +| `rollout.context_parallel_size>1` with effective decode CP=1 | `ERROR`/`FALLBACK` | strict TP=2/CP=2 acceptance rejects the topology change | +| missing/mismatched/fallback Split-KV plan set | binding failure | Split-KV provenance must cover every batch/TP/CP/owner coordinate | +| missing/native/fallback QK-Norm or RoPE backend | binding failure | both sides must execute the RL-Kernel CUDA preprocessing path | + +## Knobs + +`adapters/knobs.py` extends `V1_KNOBS` additively. Added: training-side +`tensor_parallel_size` / `context_parallel_size` / `deterministic_mode` / +`cp_comm_type`, `rollout.batch_invariant` / `rollout.kv_block_size`, and the +reduction axis (`acc_dtype`, `order`, `downcast_at`, `engine`) plus +`attention.fusion_boundary` and `attention.split_kv_policy`. + +`training.attention_backend` keeps its path but its value domain is replaced with +Megatron's `AttnBackend`; the HuggingFace names have no Megatron counterpart, so this +is a replacement rather than a mapping. + +Not done here, because both change `V1_KNOBS` itself and would break existing +cross-config tests: removing `training.sharding` (Megatron has no such concept, and +DP=1 makes it moot) and renaming `rollout.context_parallel_size` to reflect that it +binds to `prefill_context_parallel_size`. + +## Scenario + +`examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json` supersedes +`cross_config_s1_distributed_smoke.json` and +`cross_config_s3_qwen3_8b_tp4_cp4_bf16.json`, whose training sides used `sdpa` / +`flash_attention_2` and `sharding: fsdp` -- none of which exist under Megatron -- and +whose TP=4/CP=4 topology does not match the target. +`cross_config_s2_vllm_tp_vs_fsdp.json` has no Megatron-only counterpart and should be +retired rather than rewritten. + +## Out of scope + +Deliberately not in this PR: + +* launching `torchrun`, initializing process groups, or executing core attention; +* pre-attention model RMSNorm, 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); +* distributed drift benchmarks and report artifacts (#235 PR5); +* fused production backend alignment (#235 PR7) and backward (#235 PR8). diff --git a/docs/design/ws2-cp-attention-contract.md b/docs/design/ws2-cp-attention-contract.md new file mode 100644 index 00000000..95320101 --- /dev/null +++ b/docs/design/ws2-cp-attention-contract.md @@ -0,0 +1,227 @@ +# WS2 CP-Aware Attention Contract + +Status: PR1 contract and dispatch metadata + +Tracking and shared contracts: + +- [#235: CP-aware deterministic Attention](https://github.com/RL-Align/RL-Kernel/issues/235) +- [#83: WS2 roadmap](https://github.com/RL-Align/RL-Kernel/issues/83) +- [#108: WS1 numerical contract](https://github.com/RL-Align/RL-Kernel/issues/108) +- [#111: WS2 cross-config alignment](https://github.com/RL-Align/RL-Kernel/issues/111) +- [#207: cross-config logprob drift contract](https://github.com/RL-Align/RL-Kernel/issues/207) + +## Scope + +This contract describes the logical inputs and deterministic reduction semantics for standard +softmax Attention under tensor parallelism (TP) and context parallelism (CP). It lets runtime +dispatch reject a backend whose numerical semantics do not match the requested layout. + +This PR1 layer does not shard tensors, launch a collective, merge CP partial states, or implement +a fused kernel. The deterministic CP reference implementation and its distributed numerical tests +belong to later work in #235. + +## Contract Objects + +`rl_engine.kernels.attention_contract` defines: + +- `AttentionContract`: role, mode, dtype, causal metadata, sharding, reduction, and optional cache + identity; +- `ShardingSpec`: TP-local head ownership and CP block-to-token ownership; +- `ReductionSpec`: fixed `(out, lse)` merge semantics; +- `KVCacheSpec`: decode replay cache identity; +- `RoPESpec`: Qwen3 RoPE state, position identity, and fused/unfused boundary metadata; +- `AttentionBackendCapability`: the layouts and semantics a backend explicitly supports. + +Construction performs validation immediately. A structurally valid contract means that the +request is complete and internally consistent; it does not mean that an installed backend can +materialize it. + +`AttentionContract.batch_size` is the logical sequence count. For packed varlen input it must +equal `len(packed_sequence_offsets) - 1`; it is not the physical leading dimension of a flattened +token tensor. + +For full `prefill`, `query_sequence_length` equals the local sequence length described by +`ShardingSpec`. Chunked prefill and decode may use shorter query lengths than their available KV +context. + +## Qwen3-8B TP=2 CP=2 Example + +```python +from rl_engine.kernels.attention_contract import ( + AttentionContract, + ReductionSpec, + RoPESpec, + ShardingSpec, +) + +sharding = ShardingSpec( + tp_rank=0, + tp_world_size=2, + cp_rank=0, + cp_world_size=2, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=0, + local_q_heads=16, + local_kv_head_start=0, + local_kv_heads=4, + global_sequence_length=4096, + local_sequence_length=2048, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 2048), +) + +contract = AttentionContract( + role="infer", + mode="prefill", + dtype="bf16", + batch_size=1, + query_sequence_length=2048, + head_dim=128, + causal=True, + causal_offsets=(0,), + sharding=sharding, + reduction=ReductionSpec(), + rope=RoPESpec( + q_state="post_rope", + k_state="post_rope", + k_cache_state="post_rope", + theta=1.0e6, + rotary_dim=128, + query_position_offsets=(0,), + key_position_offsets=(0,), + cast_at="after_rope", + output_dtype="bf16", + fusion_boundary="unfused_rope_attention", + ), +) +``` + +The TP fields preserve the global Qwen3 GQA mapping: each rank owns 16 of 32 query heads and 4 of +8 KV heads. The CP fields map local tensor slices to stable logical global block ids. A rank that +owns non-contiguous blocks uses one global token start per block and one extra local boundary: + +```python +global_block_indices=(0, 3) +global_block_token_starts=(0, 3072) +local_block_offsets=(0, 1024, 2048) +``` + +This metadata is sufficient for a later implementation to restore logical global order without +using ring arrival order. + +## RoPE / Position Semantics + +RoPE is part of the attention contract because rollout can materialize +`RoPE+Attention` as a fused or cache-aware path while training may materialize +`RoPE -> Attention` as separate operators. PR1 does not execute the RoPE kernel, +but it records the metadata required to prove both materializations use the same +model semantics. + +`RoPESpec` records: + +- whether Q, K, and cached K are `pre_rope` or `post_rope`; +- `theta`, optional `rope_scaling`, and `rotary_dim`; +- dense `position_ids` or per-sequence `query_position_offsets` / + `key_position_offsets`; +- the RoPE cast point and output dtype; +- `fusion_boundary`, either `unfused_rope_attention` or `fused_rope_attention`. + +When RoPE metadata is present, construction validates that rotary dimensions fit +the attention head dimension and that offset metadata matches the logical batch +shape. Backends must declare RoPE support through `AttentionBackendCapability`; +a backend that cannot consume RoPE/position metadata or cannot support a fused +RoPE+Attention boundary is rejected before dispatch. + +## Reduction Semantics + +The only PR1 reduction contract is: + +```text +partial state: (out, attention-domain lse) +merge: online_softmax_lse +acc_dtype: fp32 +order: global_block_index +downcast_at: final_write +engine: in_op_reference +``` + +CP output is not a plain sum. A backend that cannot export attention-domain LSE or cannot merge +partial states in fixed logical order is incompatible with this contract. + +The acceptable output and selected-logprob drift thresholds remain owned by #108. This contract +does not introduce another tolerance table. When connected to the rollout/training chain, the +selected-token metric remains the #207 convention: + +```text +dlogp = training-side recomputed logp - rollout-side old logp +``` + +## Mode-Specific Metadata + +All causal calls provide `causal_offsets`. Packed varlen calls provide one causal offset per +packed sequence and validated `packed_sequence_offsets`. + +Decode additionally requires `KVCacheSpec` with: + +- one cache position and KV sequence length per logical sequence; +- a block/page table; +- the physical page size; +- global token positions for every logical cached token; +- a prefix-cache key and explicit shared-prefix page count when prefix caching is enabled. + +Within each logical sequence, global token positions must be strictly increasing. Block-table +padding must be trailing, the active page count must match `ceil(kv_seq_len / page_size)`, and a +sequence cannot repeat one physical page id. Different sequences may share physical pages for an +equivalent prefix only when those pages are declared by `shared_prefix_page_count`, use the same +leading page ids and logical positions, and are fully populated. Declared shared prefix pages are +read-only; all suffix pages are exclusive to one sequence, providing the contract boundary needed +for copy-on-write before divergent decode. When prefix caching is disabled, no active page may be +shared across sequences. Missing or inconsistent decode cache identity is an error at contract +construction time. + +Each `cache_positions` entry is the terminal logical position already present in that sequence's +KV cache, so it must equal the final corresponding `global_token_positions` entry. It is not the +next position to be written. + +## Contract-Aware Dispatch + +Legacy callers continue to use `KernelRegistry.get_op()`. WS2 callers use: + +```python +result = kernel_registry.get_attention_op(contract) +op = result.op +provenance = result.provenance +``` + +Dispatch considers only backends with an `AttentionBackendCapability`. It checks role, attention +mode, dtype, TP/CP degree, LSE export, deterministic CP merge, packed varlen, and KV-cache support. +When RoPE metadata is present, dispatch also checks whether the backend explicitly supports +RoPE/position metadata and fused RoPE+Attention boundaries. +An undeclared or incompatible backend is skipped with an explicit rejection reason. + +The current WS1 PyTorch Attention implementations support local reference math but do not export +attention-domain LSE or materialize deterministic CP merge. Strict WS2 requests therefore fail +clearly today. A later deterministic backend becomes selectable by registering a capability that +truthfully declares those features; no grid-planner branch or silent fallback is required. + +Successful dispatch provenance records: + +- requested and actual backend ids; +- platform and fallback status; +- prior candidate rejection reasons; +- the complete requested contract; +- the selected backend capability descriptor. + +## Validation + +Contract and dispatch behavior are covered by: + +```bash +python -m pytest tests/test_attention_contract.py -q +``` + +The tests include Qwen3 TP=2/CP=2 construction, GQA ownership errors, non-contiguous CP blocks, +packed varlen metadata, decode cache identity, undeclared backend rejection, no incompatible +fallback, RoPE metadata validation, and JSON-compatible provenance. diff --git a/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. diff --git a/docs/operators/attention.md b/docs/operators/attention.md index ebff9a58..92d24e6c 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -84,6 +84,59 @@ the inputs' device. 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 + +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). + +### WS2 deterministic CP reference + +### 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 +``` + +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 @@ -137,6 +190,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 @@ -194,6 +248,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 diff --git a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json new file mode 100644 index 00000000..11f9cef1 --- /dev/null +++ b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json @@ -0,0 +1,89 @@ +{ + "experiment_id": "ws2-qwen3-8b-attention-tp2-cp2", + "scenario_id": "qwen3_8b_megatron_tp2_cp2_vllm", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "scenario": { + "issue": "https://github.com/RL-Align/RL-Kernel/issues/235", + "pull_request": "PR4 -- cross-config integration", + "model": "Qwen3-8B dense", + "training_framework": "megatron", + "rollout_framework": "vllm", + "topology": "2 nodes x 2 GPUs, TP=2 CP=2 PP=1 DP=1, BF16, SM90", + "notes": [ + "Supersedes cross_config_s1_distributed_smoke.json and", + "cross_config_s3_qwen3_8b_tp4_cp4_bf16.json, whose training side used", + "HuggingFace attention backends and FSDP sharding. Neither exists in", + "Megatron, and DP=1 makes the sharding knob meaningless.", + "cross_config_s2_vllm_tp_vs_fsdp.json has no Megatron-only counterpart at", + "all and should be retired rather than rewritten.", + "rollout.context_parallel_size binds to vLLM", + "ParallelConfig.prefill_context_parallel_size and therefore applies to", + "prefill only; strict PR4 acceptance covers CP=2 prefill/chunked prefill.", + "A decode request that becomes CP=1 is a blocking fallback and is tested", + "separately by the PR6 logical KV replay harness." + ] + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "dtype": "bfloat16", + "enable_prefix_caching": false, + "enforce_eager": true, + "batch_invariant": true, + "kv_block_size": 16 + }, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "attention_backend": "unfused", + "compute_dtype": "bfloat16", + "deterministic_mode": true, + "cp_comm_type": "p2p", + "sharding": "unsharded" + }, + "attention": { + "reduction_acc_dtype": "fp32", + "reduction_order": "global_block_index", + "reduction_downcast_at": "final_write", + "reduction_engine": "in_op_reference", + "fusion_boundary": "unfused_rope_attention", + "split_kv_policy": 32 + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "training.context_parallel_size", + "values": [1, 2] + }, + { + "path": "training.tensor_parallel_size", + "values": [1, 2] + }, + { + "path": "attention.fusion_boundary", + "values": ["unfused_rope_attention", "fused_rope_attention"] + }, + { + "path": "training.cp_comm_type", + "values": ["p2p", "all_gather"] + }, + { + "path": "attention.reduction_order", + "values": ["global_block_index", "arrival"] + }, + { + "path": "attention.reduction_acc_dtype", + "values": ["fp32", "bf16"] + } + ] +} diff --git a/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/pyproject.toml b/pyproject.toml index b0b80a1a..2507e591 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/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/__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/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/_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/_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/adapters/__init__.py b/rl_engine/alignment/cross_config/adapters/__init__.py new file mode 100644 index 00000000..f02b9b38 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime adapters for the WS2 Qwen3-8B Megatron + vLLM cross-config target.""" + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + AttentionRuntimeReadback, + Qwen3ModelSpec, +) +from rl_engine.alignment.cross_config.adapters.knobs import ( + MEGATRON_ATTENTION_BACKENDS, + WS2_ATTENTION_KNOB_DESCRIPTORS, + WS2_ATTENTION_KNOBS, + WS2_ATTENTION_NORMALIZERS, +) +from rl_engine.alignment.cross_config.adapters.megatron import ( + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, +) +from rl_engine.alignment.cross_config.adapters.vllm import ( + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "AttentionRuntimeReadback", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", + "QWEN3_8B", + "Qwen3ModelSpec", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] diff --git a/rl_engine/alignment/cross_config/adapters/_common.py b/rl_engine/alignment/cross_config/adapters/_common.py new file mode 100644 index 00000000..9a3c5384 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/_common.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared pieces for the Megatron and vLLM WS2 attention adapters.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from rl_engine.alignment.cross_config.attention_binding import AttentionRuntimeReadback +from rl_engine.alignment.cross_config.runtime import KnobApplication +from rl_engine.alignment.cross_config.schema import ( + IsolationScope, + KnobDescriptor, + MaterializationStatus, +) +from rl_engine.kernels.attention_contract import ( + AttentionDType, + AttentionMerge, + DowncastPoint, + ReductionEngine, + ReductionOrder, + ReductionSpec, + ShardingSpec, + SplitKVSpec, +) + +__all__ = [ + "QWEN3_8B", + "Qwen3ModelSpec", + "AttentionRuntimeReadback", + "application", + "attention_dtype", + "build_reduction_spec", + "build_sharding_spec", + "causal_offsets_for", + "flatten", + "split_kv_spec", + "unsupported_reduction_reason", +] + + +@dataclass(frozen=True) +class Qwen3ModelSpec: + """Architecture constants for the frozen dense target. + + These are *not* knobs. #235/#239/#241 all fix Qwen3-8B dense, so they belong to + the scenario, and both sides must agree on them or the comparison is void. + """ + + name: str = "qwen3-8b" + hidden_size: int = 4096 + ffn_hidden_size: int = 12288 + num_layers: int = 36 + q_heads: int = 32 + kv_heads: int = 8 + head_dim: int = 128 + real_vocab_size: int = 151936 + rope_theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + qk_layernorm: bool = True + + def identity_fields(self) -> dict[str, Any]: + """The subset of :data:`IDENTITY_FIELDS` this spec is responsible for.""" + + return { + "q_heads": self.q_heads, + "kv_heads": self.kv_heads, + "head_dim": self.head_dim, + "rope_theta": self.rope_theta, + "rope_scaling": self.rope_scaling, + "rotary_dim": self.rotary_dim, + "qk_layernorm": self.qk_layernorm, + } + + +QWEN3_8B = Qwen3ModelSpec() + + +#: The planner normalizes dtype knobs to torch spellings (``bfloat16``), while +#: :class:`AttentionDType` uses short spellings (``bf16``). Passing a normalized knob +#: straight into the enum raises, so every adapter must translate here rather than +#: each inventing its own mapping. +_DTYPE_ALIASES: Mapping[str, AttentionDType] = { + "bf16": AttentionDType.BF16, + "bfloat16": AttentionDType.BF16, + "fp16": AttentionDType.FP16, + "float16": AttentionDType.FP16, + "half": AttentionDType.FP16, + "fp32": AttentionDType.FP32, + "float32": AttentionDType.FP32, + "float": AttentionDType.FP32, +} + + +def attention_dtype(value: Any, *, field: str) -> AttentionDType: + """Translate a normalized knob dtype into an :class:`AttentionDType`.""" + + if isinstance(value, AttentionDType): + return value + key = str(value).strip().lower().replace("torch.", "") + try: + return _DTYPE_ALIASES[key] + except KeyError as exc: + raise ValueError( + f"{field}={value!r} is not a supported attention dtype; " + f"expected one of {sorted(set(_DTYPE_ALIASES))}" + ) from exc + + +def split_kv_spec(flat: Mapping[str, Any]) -> SplitKVSpec: + """Build the first-class logical Split-KV request. + + The integer is a fixed logical KV chunk size in tokens. It is intentionally + not vLLM's ``flash_attn_max_num_splits_for_cuda_graph``: that setting is only + an upper bound and cannot prove which runtime boundaries executed. + """ + + split_size = flat.get("attention.split_kv_policy") + if split_size is None: + return SplitKVSpec.disabled() + return SplitKVSpec.fixed(int(split_size)) + + +def flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + """Flatten nested knob mappings into dotted paths.""" + + flat: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}{key}" + if isinstance(child, Mapping): + flat.update(flatten(child, f"{path}.")) + else: + flat[path] = child + return flat + + +def application( + descriptor: KnobDescriptor, + requested: Any, + materialized: Any, + actual: Any, + status: MaterializationStatus, + reason: str, + **evidence: Any, +) -> KnobApplication: + return KnobApplication( + path=descriptor.path, + requested=requested, + materialized=materialized, + actual=actual, + lifecycle=descriptor.lifecycle, + status=status, + evidence={"reason": reason, **evidence}, + critical=descriptor.critical, + ) + + +def unsupported_reduction_reason(flat: Mapping[str, Any]) -> str | None: + """Return why the requested reduction cannot be materialized, if it cannot. + + #236 declares single-member enums for merge order, downcast point and reduction + engine, so the alternative knob values exist only as control groups. Requesting + one must fail loudly rather than quietly collapse onto the supported value -- + silently substituting ``global_block_index`` for a requested ``arrival`` would + make the control group indistinguishable from the treatment. + """ + + order = flat.get("attention.reduction_order") + if order is not None and order != ReductionOrder.GLOBAL_BLOCK_INDEX.value: + return ( + f"attention.reduction_order={order!r} has no backend; #236 ReductionOrder " + f"declares only {ReductionOrder.GLOBAL_BLOCK_INDEX.value!r}" + ) + downcast = flat.get("attention.reduction_downcast_at") + if downcast is not None and downcast != DowncastPoint.FINAL_WRITE.value: + return ( + f"attention.reduction_downcast_at={downcast!r} has no backend; #236 " + f"DowncastPoint declares only {DowncastPoint.FINAL_WRITE.value!r}" + ) + engine = flat.get("attention.reduction_engine") + if engine is not None and engine != ReductionEngine.IN_OP_REFERENCE.value: + return ( + f"attention.reduction_engine={engine!r} has no backend; the Transformer " + "Engine merge oracle lands in #235 PR2/PR3, not here" + ) + acc_dtype = flat.get("attention.reduction_acc_dtype") + if ( + acc_dtype is not None + and attention_dtype(acc_dtype, field="attention.reduction_acc_dtype") + is not AttentionDType.FP32 + ): + return ( + f"attention.reduction_acc_dtype={acc_dtype!r} violates the WS2 mandate; " + "the CP (out, lse) merge accumulates in fp32" + ) + return None + + +def build_reduction_spec(flat: Mapping[str, Any]) -> ReductionSpec: + """Build the reduction spec, having already rejected unsupported requests.""" + + return ReductionSpec( + merge=AttentionMerge.ONLINE_SOFTMAX_LSE, + acc_dtype=AttentionDType.FP32, + order=ReductionOrder.GLOBAL_BLOCK_INDEX, + downcast_at=DowncastPoint.FINAL_WRITE, + engine=ReductionEngine.IN_OP_REFERENCE, + ) + + +def build_sharding_spec( + *, + model: Qwen3ModelSpec, + tp_rank: int, + tp_world_size: int, + cp_rank: int, + cp_world_size: int, + global_sequence_length: int, +) -> ShardingSpec: + """Build a CP/TP sharding spec for one rank of the frozen layout. + + TP splits heads, CP splits the sequence. The #239 rank layout fixes + ``rank = cp_rank * tp_world_size + tp_rank`` for a 2-node x 2-GPU deployment, + but nothing here depends on that mapping: ownership is derived from the ranks + themselves so the same builder serves CP=1 baselines. + """ + + if model.q_heads % tp_world_size or model.kv_heads % tp_world_size: + raise ValueError( + f"Qwen3 GQA heads ({model.q_heads}/{model.kv_heads}) must divide evenly " + f"across tp_world_size={tp_world_size}" + ) + if global_sequence_length % cp_world_size: + raise ValueError( + f"global_sequence_length={global_sequence_length} must divide evenly " + f"across cp_world_size={cp_world_size}" + ) + + local_q_heads = model.q_heads // tp_world_size + local_kv_heads = model.kv_heads // tp_world_size + local_sequence_length = global_sequence_length // cp_world_size + + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=model.q_heads, + global_kv_heads=model.kv_heads, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=global_sequence_length, + local_sequence_length=local_sequence_length, + # One contiguous CP block per rank. The merge order key is the global block + # index, never the arrival order of the CP exchange. + global_block_indices=(cp_rank,), + global_block_token_starts=(cp_rank * local_sequence_length,), + local_block_offsets=(0, local_sequence_length), + ) + + +def causal_offsets_for(sharding: ShardingSpec, batch_size: int) -> tuple[int, ...]: + """Causal offsets for one CP shard, one entry per batch entry. + + Under CP the local query block does not start at global position zero, so the + causal mask has to be shifted by the number of preceding global tokens. Taking + that from ``global_block_token_starts`` rather than recomputing + ``cp_rank * local_sequence_length`` keeps uneven CP splits correct. + """ + + offset = sharding.global_block_token_starts[0] + return (offset,) * batch_size + + +_PROCESS_SCOPES = (IsolationScope.PROCESS, IsolationScope.DISTRIBUTED_CONTEXT) diff --git a/rl_engine/alignment/cross_config/adapters/knobs.py b/rl_engine/alignment/cross_config/adapters/knobs.py new file mode 100644 index 00000000..33c29820 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/knobs.py @@ -0,0 +1,161 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 attention knobs for the Qwen3-8B TP=2 CP=2 Megatron + vLLM target. + +``V1_KNOBS`` was written against a HuggingFace/FSDP rollout-vs-training pair. Three +of its entries do not survive contact with the frozen Megatron + vLLM target: + +* ``training.sharding`` takes ``unsharded``/``fsdp``, neither of which exists in + Megatron, and is meaningless at DP=1 anyway; +* ``training.attention_backend`` takes HuggingFace names + (``flash_attention_2``/``sdpa``/``eager``/``model_default``) while Megatron's + ``AttnBackend`` is ``flash``/``fused``/``unfused``/``local``/``auto``; +* there is no training-side ``tensor_parallel_size`` or ``context_parallel_size`` + at all, so the target configuration cannot even be expressed. + +This module is deliberately **additive**: it extends ``V1_KNOBS`` rather than +editing it, and overrides only the normalizer for ``training.attention_backend``. +Deleting the two dead knobs changes ``V1_KNOBS`` itself and would break existing +cross-config tests, so it is left to a follow-up on the framework PR. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +from rl_engine.alignment.cross_config.planner import ( + _NORMALIZERS, + V1_KNOBS, + Normalizer, + _normalize_choice, + _positive_int, + _strict_bool, +) +from rl_engine.alignment.cross_config.schema import IsolationScope, KnobDescriptor + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] + + +#: ``megatron.core.transformer.enums.AttnBackend``. +MEGATRON_ATTENTION_BACKENDS: tuple[str, ...] = ( + "flash", + "fused", + "unfused", + "local", + "auto", +) + + +WS2_ATTENTION_KNOB_DESCRIPTORS: tuple[KnobDescriptor, ...] = ( + # -- training-side parallelism: the target configuration itself ------------ + KnobDescriptor( + "training.tensor_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "training.context_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + # -- determinism switches, one per framework ------------------------------ + KnobDescriptor( + "training.deterministic_mode", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "rollout.batch_invariant", + IsolationScope.PROCESS, + ("rollout",), + ), + # -- reduction knobs: the "turn the noise sources on and off" axis --------- + # These are what make drift attributable. ``reduction.order=arrival`` in + # particular is a control group, not a supported production value. + KnobDescriptor( + "attention.reduction_acc_dtype", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("fp32", "bf16"), + ), + KnobDescriptor( + "attention.reduction_order", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("global_block_index", "arrival"), + ), + KnobDescriptor( + "attention.reduction_downcast_at", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("final_write", "per_block"), + ), + KnobDescriptor( + "attention.reduction_engine", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("in_op_reference", "te_oracle"), + ), + # -- materialization knobs: differences the experiment measures ------------ + KnobDescriptor( + "attention.fusion_boundary", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("unfused_rope_attention", "fused_rope_attention"), + ), + KnobDescriptor( + # Shared logical KV chunk size. Runtime adapters must separately report + # the actual per-owner boundaries; a configured value is not evidence. + "attention.split_kv_policy", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + ), + KnobDescriptor( + # vLLM: CacheConfig.block_size -> AttentionContract.kv_cache.page_size + "rollout.kv_block_size", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + # The CP communication group cannot be reconfigured once built; it is bound to + # the distributed context, not merely to engine construction. + KnobDescriptor( + "training.cp_comm_type", + IsolationScope.DISTRIBUTED_CONTEXT, + ("training",), + allowed_values=("p2p", "all_gather", "a2a", "a2a+p2p"), + ), +) + + +WS2_ATTENTION_KNOBS: Mapping[str, KnobDescriptor] = { + **V1_KNOBS, + **{descriptor.path: descriptor for descriptor in WS2_ATTENTION_KNOB_DESCRIPTORS}, +} + + +WS2_ATTENTION_NORMALIZERS: Mapping[str, Normalizer] = { + **_NORMALIZERS, + # Replace, not map: the HuggingFace names have no Megatron counterpart. + "training.attention_backend": _normalize_choice(*MEGATRON_ATTENTION_BACKENDS), + "training.tensor_parallel_size": _positive_int, + "training.context_parallel_size": _positive_int, + "training.deterministic_mode": _strict_bool, + "rollout.batch_invariant": _strict_bool, + # AttentionDType values, not torch dtype names -- these feed ReductionSpec directly. + "attention.reduction_acc_dtype": _normalize_choice("fp32", "bf16"), + "attention.reduction_order": _normalize_choice("global_block_index", "arrival"), + "attention.reduction_downcast_at": _normalize_choice("final_write", "per_block"), + "attention.reduction_engine": _normalize_choice("in_op_reference", "te_oracle"), + "attention.fusion_boundary": _normalize_choice( + "unfused_rope_attention", "fused_rope_attention" + ), + "attention.split_kv_policy": _positive_int, + "rollout.kv_block_size": _positive_int, + "training.cp_comm_type": _normalize_choice("p2p", "all_gather", "a2a", "a2a+p2p"), +} diff --git a/rl_engine/alignment/cross_config/adapters/megatron.py b/rl_engine/alignment/cross_config/adapters/megatron.py new file mode 100644 index 00000000..55903cc8 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -0,0 +1,481 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Training-side (Megatron) runtime adapter for WS2 attention cross-config. + +Two things live here: + +``MegatronProvenanceAdapter`` + Read-only. Turns a Megatron config object into the construction and + distributed-context fingerprints the cross-config framework already expects, + plus the determinism probe. It never imports ``megatron`` -- every accessor is + duck-typed -- so this module is importable and testable on a laptop. + +``MegatronAttentionMaterializer`` + Implements the ``RuntimeMaterializer`` protocol. Before this PR the only + implementation in the repository was ``CpuSmokeMaterializer`` over a synthetic + CPU model, so nothing had ever materialized a real distributed runtime. + +Scope boundary: materialization builds and validates the training-side +:class:`AttentionContract` and reports what would be constructed. Without an +``AttentionRuntimeReadback`` it reports ``UNOBSERVABLE``, never ``APPLIED``. It +does not launch ``torchrun``, initialize process groups, or execute attention; +the 2-node x 2-GPU launcher must inject readback collected after execution. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + AttentionRuntimeReadback, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + split_kv_spec, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import ( + DeterminismProbe, + megatron_probe_from_config, +) +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "MEGATRON_CONSTRUCTION_KEYS", + "MEGATRON_DISTRIBUTED_KEYS", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", +] + + +#: ``TransformerConfig`` fields that change attention arithmetic. Hashed into the +#: construction fingerprint. Deliberately excludes MoE, Mamba, MLA and sparse +#: attention fields: the frozen target is Qwen3-8B dense, and those are asserted +#: off rather than recorded. +MEGATRON_CONSTRUCTION_KEYS: tuple[str, ...] = ( + "attention_backend", + "attention_softmax_in_fp32", + "apply_query_key_layer_scaling", + "apply_rope_fusion", + "masked_softmax_fusion", + "bias_activation_fusion", + "bias_dropout_fusion", + "gradient_accumulation_fusion", + "cross_entropy_loss_fusion", + "cross_entropy_fusion_impl", + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + "recompute_modules", + "rotary_base", + "rotary_percent", + "rotary_interleaved", + "rotary_scaling_factor", + "qk_layernorm", + "hidden_dropout", + "attention_dropout", + "params_dtype", + "bf16", + "fp16", + "fp8", + "deterministic_mode", +) + + +#: ``ModelParallelConfig`` fields that define the distributed context. +MEGATRON_DISTRIBUTED_KEYS: tuple[str, ...] = ( + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "virtual_pipeline_model_parallel_size", + "context_parallel_size", + "hierarchical_context_parallel_sizes", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "sequence_parallel", + "cp_comm_type", + "tp_comm_overlap", + "use_te_rng_tracker", +) + + +#: Fields that must hold these values for the frozen dense target. A mismatch is a +#: hard stop, not a recorded difference -- see the exclusion list in the WS2 scope. +MEGATRON_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "pipeline_model_parallel_size": 1, + "expert_model_parallel_size": 1, + "sequence_parallel": False, + "fp8": None, + "hidden_dropout": 0.0, + "attention_dropout": 0.0, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class MegatronProvenanceAdapter: + """Extract fingerprints and determinism evidence from a Megatron config. + + ``config`` may be a real ``TransformerConfig``/``ModelParallelConfig``, a merged + namespace, or a test double. Missing attributes read as ``None`` and are + recorded as such rather than raising: an absent field is itself provenance. + """ + + framework = "megatron" + + def __init__(self, config: Any, *, env: Optional[Mapping[str, str]] = None): + self.config = config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_CONSTRUCTION_KEYS} + + def distributed_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_DISTRIBUTED_KEYS} + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return megatron_probe_from_config(self.config, env=self.env) + + def frozen_scope_violations(self) -> tuple[str, ...]: + """Return the frozen-scope assertions this config violates.""" + + violations: list[str] = [] + for name, expected in MEGATRON_FROZEN_ASSERTIONS.items(): + actual = _value(self.config, name) + if actual is None: + # Not declared. Treated as unknown rather than as satisfied, because + # a silently-absent MoE or FP8 setting is exactly the case that would + # otherwise slip past a dense-only claim. + violations.append(f"{name} is not declared (expected {expected!r})") + elif actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "determinism": self.determinism_probe().to_dict(), + } + + +class MegatronAttentionMaterializer: + """Materialize the training-side attention runtime for the WS2 target.""" + + runtime_kind = "megatron_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + backend_id: str = "rlkernel.cp_attention_reference", + provenance: Optional[MegatronProvenanceAdapter] = None, + runtime_readback: Optional[AttentionRuntimeReadback] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.backend_id = backend_id + self.provenance = provenance + self.runtime_readback = runtime_readback + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + """Build the training-side contract. Raises on an unusable request.""" + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=attention_dtype( + flat.get("training.compute_dtype", "bf16"), field="training.compute_dtype" + ), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + split_kv=split_kv_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "training" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + if contract is None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + f"attention contract is unavailable: {blocked}", + ) + ) + continue + applications.append( + self._runtime_application( + descriptor, + requested, + contract=contract, + scope_violations=scope_violations, + ) + ) + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "megatron", + "attention_backend": flat.get("training.attention_backend"), + "compute_dtype": flat.get("training.compute_dtype"), + "deterministic_mode": flat.get("training.deterministic_mode"), + "cp_comm_type": flat.get("training.cp_comm_type"), + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "runtime_readback": ( + None if self.runtime_readback is None else self.runtime_readback.to_dict() + ), + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"training": side_config, "rollout": {}}, + topology={ + "training": { + "world_size": tp_world_size * cp_world_size, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": cp_world_size, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "rollout": {"world_size": 1}, + }, + scorer={ + "mode": "teacher_forcing", + "framework": "megatron", + "export_lse": True, + }, + operator_backends={ + "training": self.backend_id, + "rollout": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) + + def _runtime_application( + self, + descriptor: KnobDescriptor, + requested: Any, + *, + contract: AttentionContract, + scope_violations: tuple[str, ...], + ) -> KnobApplication: + readback = self.runtime_readback + if readback is None: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + ( + "configured in the training contract, but no Megatron runtime " + "readback was supplied" + ), + frozen_scope_violations=list(scope_violations), + ) + if scope_violations or not readback.frozen_scope_verified: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.UNOBSERVABLE, + "Megatron frozen-scope assertions were not all verified by runtime readback", + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) + if readback.split_kv_fallback: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "Megatron runtime reported a Split-KV fallback", + runtime_readback_source=readback.source, + ) + if readback.contract != contract: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "Megatron runtime contract differs from the requested contract", + runtime_readback_source=readback.source, + ) + if descriptor.path not in readback.actual_knobs: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "Megatron runtime readback does not expose this knob", + runtime_readback_source=readback.source, + ) + actual = readback.actual_knobs[descriptor.path] + status = ( + MaterializationStatus.APPLIED if actual == requested else MaterializationStatus.FALLBACK + ) + reason = ( + "verified from the executed Megatron runtime" + if status is MaterializationStatus.APPLIED + else "Megatron runtime value differs from the requested value" + ) + return application( + descriptor, + requested, + requested, + actual, + status, + reason, + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) diff --git a/rl_engine/alignment/cross_config/adapters/vllm.py b/rl_engine/alignment/cross_config/adapters/vllm.py new file mode 100644 index 00000000..f89ab232 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -0,0 +1,540 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rollout-side (vLLM) runtime adapter for WS2 attention cross-config. + +Mirrors :mod:`.megatron`, with three differences that come straight from what vLLM +actually is: + +* vLLM's context parallelism is ``prefill_context_parallel_size`` -- it applies to + prefill only, so a decode-mode contract must declare ``cp_world_size == 1`` + regardless of what the prefill knob says. +* ``CacheConfig.block_size`` is the paged-KV page size, and it feeds + ``KVCacheSpec.page_size`` directly rather than being invented here. +* Determinism comes from the ``VLLM_BATCH_INVARIANT`` environment variable rather + than from a config field, because vLLM applies it inside + ``init_batch_invariance()`` at worker startup. + +Like the Megatron adapter, nothing here imports ``vllm``; configs are duck-typed so +the module is importable anywhere. Configured-only values remain ``UNOBSERVABLE``; +``APPLIED`` requires an explicit post-execution ``AttentionRuntimeReadback``. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + AttentionRuntimeReadback, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + split_kv_spec, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import DeterminismProbe, vllm_probe_from_env +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "VLLM_ATTENTION_KEYS", + "VLLM_CACHE_KEYS", + "VLLM_FROZEN_ASSERTIONS", + "VLLM_MODEL_KEYS", + "VLLM_PARALLEL_KEYS", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", +] + + +VLLM_MODEL_KEYS: tuple[str, ...] = ( + "dtype", + "seed", + "quantization", + "enforce_eager", + "max_logprobs", + "disable_cascade_attn", + "max_model_len", +) + +VLLM_CACHE_KEYS: tuple[str, ...] = ( + "block_size", + "cache_dtype", + "enable_prefix_caching", + "prefix_caching_hash_algo", + "calculate_kv_scales", + "sliding_window", +) + +VLLM_ATTENTION_KEYS: tuple[str, ...] = ( + "backend", + "flash_attn_version", + "use_prefill_decode_attention", + "flash_attn_max_num_splits_for_cuda_graph", + "use_cudnn_prefill", + "disable_flashinfer_prefill", + "use_non_causal", +) + +VLLM_PARALLEL_KEYS: tuple[str, ...] = ( + "tensor_parallel_size", + "pipeline_parallel_size", + "prefill_context_parallel_size", + "data_parallel_size", +) + + +#: Frozen dense-target assertions on the rollout side. ``cache_dtype`` must stay +#: ``auto`` because an FP8 KV cache is a representation-drift problem tracked +#: separately, and ``disable_cascade_attn`` must stay ``True`` because cascade +#: attention changes the block-merge structure the contract pins down. +VLLM_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "quantization": None, + "cache_dtype": "auto", + "calculate_kv_scales": False, + "disable_cascade_attn": True, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "sliding_window": None, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class VllmProvenanceAdapter: + """Extract fingerprints and determinism evidence from vLLM configs.""" + + framework = "vllm" + + def __init__( + self, + *, + model_config: Any = None, + cache_config: Any = None, + attention_config: Any = None, + parallel_config: Any = None, + env: Optional[Mapping[str, str]] = None, + ): + self.model_config = model_config + self.cache_config = cache_config + self.attention_config = attention_config + self.parallel_config = parallel_config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + view: dict[str, Any] = {} + for prefix, config, keys in ( + ("model", self.model_config, VLLM_MODEL_KEYS), + ("cache", self.cache_config, VLLM_CACHE_KEYS), + ("attention", self.attention_config, VLLM_ATTENTION_KEYS), + ): + for name in keys: + view[f"{prefix}.{name}"] = _value(config, name) + return view + + def distributed_view(self) -> dict[str, Any]: + return { + f"parallel.{name}": _value(self.parallel_config, name) for name in VLLM_PARALLEL_KEYS + } + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return vllm_probe_from_env(self.env, model_config=self.model_config) + + def frozen_scope_violations(self) -> tuple[str, ...]: + sources = { + "quantization": self.model_config, + "disable_cascade_attn": self.model_config, + "cache_dtype": self.cache_config, + "calculate_kv_scales": self.cache_config, + "sliding_window": self.cache_config, + "pipeline_parallel_size": self.parallel_config, + "data_parallel_size": self.parallel_config, + } + violations: list[str] = [] + for name, expected in VLLM_FROZEN_ASSERTIONS.items(): + config = sources.get(name) + if config is None: + continue + actual = _value(config, name) + if actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + @property + def kv_page_size(self) -> Optional[int]: + """vLLM's paged-KV block size, which is the contract's ``page_size``.""" + + block_size = _value(self.cache_config, "block_size") + return int(block_size) if block_size is not None else None + + @property + def split_kv_policy(self) -> Optional[int]: + """Diagnostic vLLM maximum split count, not the logical chunk-size contract.""" + + splits = _value(self.attention_config, "flash_attn_max_num_splits_for_cuda_graph") + return int(splits) if splits is not None else None + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "kv_page_size": self.kv_page_size, + "flash_attn_max_num_splits_for_cuda_graph": self.split_kv_policy, + "determinism": self.determinism_probe().to_dict(), + } + + +class VllmRolloutMaterializer: + """Materialize the rollout-side attention runtime for the WS2 target.""" + + runtime_kind = "vllm_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + mode: AttentionMode = AttentionMode.CHUNKED_PREFILL, + backend_id: str = "vllm.flash_attn", + provenance: Optional[VllmProvenanceAdapter] = None, + runtime_readback: Optional[AttentionRuntimeReadback] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.mode = mode + self.backend_id = backend_id + self.provenance = provenance + self.runtime_readback = runtime_readback + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def effective_cp_world_size(self, flat: Mapping[str, Any]) -> int: + """CP applies to prefill only; decode always runs at CP=1.""" + + requested = int(flat.get("rollout.context_parallel_size", 1)) + if self.mode is AttentionMode.DECODE: + return 1 + return requested + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + if self.mode is AttentionMode.DECODE: + # Decode replay needs a validated KVCacheSpec (cache positions, page + # ownership, prefix-cache identity). That is #235 PR6's contract surface, + # and inventing a placeholder here would let an unvalidated decode case + # look bound. Fail instead. + raise AttentionContractError( + "decode-mode materialization requires KV-cache identity from #235 PR6; " + "this adapter covers prefill and chunked prefill" + ) + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + cp_world_size = self.effective_cp_world_size(flat) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank if cp_world_size > 1 else 0, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.FUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + # vLLM stores post-RoPE K in the cache; recorded, not asserted equal to + # the training side, because it is a materialization fact. + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.INFER, + mode=self.mode, + dtype=attention_dtype(flat.get("rollout.dtype", "bf16"), field="rollout.dtype"), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + split_kv=split_kv_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + requested_cp = int(flat.get("rollout.context_parallel_size", 1)) + effective_cp = self.effective_cp_world_size(flat) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "rollout" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + if contract is None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + f"attention contract is unavailable: {blocked}", + ) + ) + continue + if path == "rollout.context_parallel_size" and effective_cp != requested_cp: + applications.append( + application( + descriptor, + requested, + effective_cp, + effective_cp, + MaterializationStatus.FALLBACK, + ( + "vLLM context parallelism covers prefill only; a decode-mode " + f"contract runs at cp_world_size=1, not {requested_cp}" + ), + vllm_field="ParallelConfig.prefill_context_parallel_size", + ) + ) + continue + applications.append( + self._runtime_application( + descriptor, + requested, + contract=contract, + scope_violations=scope_violations, + ) + ) + + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "vllm", + "dtype": flat.get("rollout.dtype"), + "enforce_eager": flat.get("rollout.enforce_eager"), + "enable_prefix_caching": flat.get("rollout.enable_prefix_caching"), + "batch_invariant": flat.get("rollout.batch_invariant"), + "kv_block_size": flat.get("rollout.kv_block_size"), + "split_kv_policy": flat.get("attention.split_kv_policy"), + "attention_mode": self.mode.value, + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "runtime_readback": ( + None if self.runtime_readback is None else self.runtime_readback.to_dict() + ), + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"rollout": side_config, "training": {}}, + topology={ + "rollout": { + "world_size": tp_world_size * effective_cp, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": effective_cp, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "training": {"world_size": 1}, + }, + scorer={ + "mode": "rollout_logprob", + "framework": "vllm", + "export_lse": True, + }, + operator_backends={ + "rollout": self.backend_id, + "training": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) + + def _runtime_application( + self, + descriptor: KnobDescriptor, + requested: Any, + *, + contract: AttentionContract, + scope_violations: tuple[str, ...], + ) -> KnobApplication: + readback = self.runtime_readback + if readback is None: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "configured in the rollout contract, but no vLLM runtime readback was supplied", + frozen_scope_violations=list(scope_violations), + ) + if scope_violations or not readback.frozen_scope_verified: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.UNOBSERVABLE, + "vLLM frozen-scope assertions were not all verified by runtime readback", + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) + if readback.split_kv_fallback: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "vLLM runtime reported a Split-KV fallback", + runtime_readback_source=readback.source, + ) + if readback.contract != contract: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "vLLM runtime contract differs from the requested contract", + runtime_readback_source=readback.source, + ) + if descriptor.path not in readback.actual_knobs: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "vLLM runtime readback does not expose this knob", + runtime_readback_source=readback.source, + ) + actual = readback.actual_knobs[descriptor.path] + status = ( + MaterializationStatus.APPLIED if actual == requested else MaterializationStatus.FALLBACK + ) + reason = ( + "verified from the executed vLLM runtime" + if status is MaterializationStatus.APPLIED + else "vLLM runtime value differs from the requested value" + ) + return application( + descriptor, + requested, + requested, + actual, + status, + reason, + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) diff --git a/rl_engine/alignment/cross_config/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"] 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..1007f197 --- /dev/null +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -0,0 +1,957 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Three-tier binding between rollout-side and training-side attention contracts. + +Issue #235 PR4 requires that "rollout and training descriptors bind to the same +semantic attention contract". Under the frozen Megatron + vLLM deployment the two +sides can never produce *identical* :class:`AttentionContract` instances: training +runs full-sequence prefill over a CP-sharded sequence, while rollout runs vLLM +paged-KV chunked prefill and decode. Taking "same contract" literally would make +the target configuration permanently unbindable. + +This module therefore splits binding into three tiers: + +``IDENTICAL`` + Logical identity. Both sides must agree bit for bit, otherwise the pair is not + comparable at all and no drift number from it means anything. + +``SEMANTIC`` + The WS2 numerical claim: merge semantics, accumulation dtype, reduction order + and downcast point are decided by the contract, not by the implementation. + Both sides must carry the same values *and* those values must match the WS2 + mandate, otherwise the comparison fails closed. + +``RECORDED`` + Materialization facts that the two sides are expected to differ on -- attention + mode, RoPE fusion boundary, KV-cache paging, backend id, reduction engine. These + differences are exactly what the experiment measures, so they are recorded into + provenance rather than rejected. + +Deliberately *not* in ``SEMANTIC``: ``engine``. Training may run the in-op +deterministic reference while rollout runs a Transformer Engine merge oracle; forcing +those equal would defeat the purpose of the oracle comparison in #235 PR2/3/5/6. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from typing import Any, Optional + +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMerge, + AttentionRole, + DowncastPoint, + ReductionOrder, + SplitKVRuntimePlanSet, + validate_split_kv_plan_set_alignment, +) +from rl_engine.kernels.attention_preprocess import MANDATED_ATTENTION_PREPROCESS_BACKENDS + +__all__ = [ + "ATTENTION_LSE_DOMAIN", + "AttentionBindingError", + "AttentionBindingResult", + "AttentionRuntimeReadback", + "BindingErrorCode", + "BindingIssue", + "BindingTier", + "IDENTITY_FIELDS", + "NULLABLE_IDENTITY_FIELDS", + "RECORDED_FIELDS", + "SEMANTIC_CONTRACT_FIELDS", + "SEMANTIC_REDUCTION_FIELDS", + "TOPOLOGY_FIELDS", + "WS2_ATTENTION_REDUCTION_MANDATE", + "bind_attention_contracts", + "bind_attention_runtime_readbacks", + "first_blocking_issue", + "identity_fingerprint", + "summarize_binding", +] + + +class AttentionBindingError(ValueError): + """Raised when a caller supplies structurally unusable binding inputs.""" + + +class BindingTier(str, Enum): + """Which rule a field is governed by.""" + + IDENTICAL = "identical" + SEMANTIC = "semantic" + RECORDED = "recorded" + + +class BindingErrorCode(str, Enum): + """Stable, machine-readable reasons a binding is rejected. + + Callers branch on these; they are part of the artifact schema and must not be + renamed without a schema version bump. + """ + + IDENTITY_MISSING = "IDENTITY_MISSING" + IDENTITY_MISMATCH = "IDENTITY_MISMATCH" + REDUCTION_SEMANTIC_MISMATCH = "REDUCTION_SEMANTIC_MISMATCH" + REDUCTION_MANDATE_VIOLATION = "REDUCTION_MANDATE_VIOLATION" + LSE_NOT_EXPORTED = "LSE_NOT_EXPORTED" + ROLE_COLLISION = "ROLE_COLLISION" + DETERMINISM_INCOMPATIBLE = "DETERMINISM_INCOMPATIBLE" + TOPOLOGY_MISMATCH = "TOPOLOGY_MISMATCH" + SPLIT_KV_RUNTIME_MISSING = "SPLIT_KV_RUNTIME_MISSING" + SPLIT_KV_MISMATCH = "SPLIT_KV_MISMATCH" + SPLIT_KV_FALLBACK = "SPLIT_KV_FALLBACK" + ATTENTION_PREPROCESS_MISSING = "ATTENTION_PREPROCESS_MISSING" + ATTENTION_PREPROCESS_MISMATCH = "ATTENTION_PREPROCESS_MISMATCH" + ATTENTION_PREPROCESS_FALLBACK = "ATTENTION_PREPROCESS_FALLBACK" + + +@dataclass(frozen=True) +class AttentionRuntimeReadback: + """Actual attention contract and all-rank Split-KV evidence from one engine.""" + + contract: AttentionContract + actual_knobs: Mapping[str, Any] + split_kv_plan_set: SplitKVRuntimePlanSet + source: str + frozen_scope_verified: bool + preprocess_backends: Mapping[str, str] = field(default_factory=dict) + preprocess_fallback: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.contract, AttentionContract): + raise TypeError("runtime readback contract must be an AttentionContract") + if not isinstance(self.actual_knobs, Mapping): + raise TypeError("runtime readback actual_knobs must be a mapping") + if not isinstance(self.split_kv_plan_set, SplitKVRuntimePlanSet): + raise TypeError("runtime readback requires a complete SplitKVRuntimePlanSet") + if not isinstance(self.source, str) or not self.source.strip(): + raise ValueError("runtime readback source must be a non-empty string") + if not isinstance(self.frozen_scope_verified, bool): + raise TypeError("frozen_scope_verified must be a bool") + if not isinstance(self.preprocess_backends, Mapping): + raise TypeError("runtime readback preprocess_backends must be a mapping") + for name, backend in self.preprocess_backends.items(): + if not isinstance(name, str) or not name.strip(): + raise ValueError("preprocess backend names must be non-empty strings") + if not isinstance(backend, str) or not backend.strip(): + raise ValueError("preprocess backend IDs must be non-empty strings") + if not isinstance(self.preprocess_fallback, bool): + raise TypeError("preprocess_fallback must be a bool") + + 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: + return bool(_split_kv_fallbacks(self.split_kv_plan_set)) + + def to_dict(self) -> dict[str, Any]: + return { + "source": self.source, + "frozen_scope_verified": self.frozen_scope_verified, + "contract": self.contract.to_dict(), + "actual_knobs": dict(self.actual_knobs), + "attention_preprocess": { + "backends": dict(self.preprocess_backends), + "fallback": self.preprocess_fallback, + }, + "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), + } + + +#: Attention exports attention-domain LSE, never vocab-logprob LSE (#235). +#: Recorded explicitly so a future ``LogprobContract`` binding cannot be confused +#: with this one purely because both set ``export_lse=True``. +ATTENTION_LSE_DOMAIN = "attention" + + +#: Fields both sides must agree on bit for bit before any comparison is meaningful. +#: Sourced from #235 "Numerical Contract" preconditions plus the vime-owned rollout +#: provenance (weight version, sampling, padding) that the issue assumes but does +#: not enumerate. +IDENTITY_FIELDS: tuple[str, ...] = ( + "checkpoint_id", + "model_version", + "weight_version", + "tokenizer_fingerprint", + "token_ids_fingerprint", + "active_mask_fingerprint", + "position_ids_fingerprint", + "padding_side", + "pre_update_state", + # model semantics that decide what attention *means* + "q_heads", + "kv_heads", + "head_dim", + "rope_theta", + "rope_scaling", + "rotary_dim", + "qk_layernorm", + # batch composition: batch-invariance is a claim about results not changing with + # batch makeup, so two sides scoring different batches are not comparable at all + "batch_size", + # decode replay identity (#235 PR6) + "global_token_positions_fingerprint", + "kv_seq_lens_fingerprint", +) + + +#: Reduction fields that decide the numerical result. Both sides must carry the +#: same value, and that value must satisfy :data:`WS2_ATTENTION_REDUCTION_MANDATE`. +SEMANTIC_REDUCTION_FIELDS: tuple[str, ...] = ( + "merge", + "acc_dtype", + "order", + "downcast_at", +) + + +#: Contract fields outside ``ReductionSpec`` that still decide the numerical result. +#: ``dtype`` is here rather than in :data:`RECORDED_FIELDS` because comparing a BF16 +#: rollout against an FP16 training pass produces a real drift number attributable to +#: nothing. #235 PR5 does sweep BF16 against an FP32 reference; that sweep opts in via +#: ``allow_dtype_difference`` instead of loosening the default. +SEMANTIC_CONTRACT_FIELDS: tuple[str, ...] = ("dtype",) + + +#: Sharding fields that determine local GQA head and sequence ownership. These are +#: comparison preconditions, not harmless backend provenance: a TP/CP mismatch +#: means the two ranks did not evaluate the same local attention problem. +TOPOLOGY_FIELDS: tuple[str, ...] = ( + "tp_rank", + "tp_world_size", + "cp_rank", + "cp_world_size", + "global_q_heads", + "global_kv_heads", + "local_q_head_start", + "local_q_heads", + "local_kv_head_start", + "local_kv_heads", + "global_sequence_length", + "local_sequence_length", + "global_block_indices", + "global_block_token_starts", + "local_block_offsets", + "packed_sequence_offsets", +) + + +#: The WS2 mandate itself. ``#236`` currently declares single-member enums for +#: ``merge`` / ``order`` / ``downcast_at``, so those checks are tautological today; +#: they are written out anyway so that widening any of those enums later fails here +#: instead of silently admitting a non-conforming backend. +WS2_ATTENTION_REDUCTION_MANDATE: Mapping[str, str] = { + "merge": AttentionMerge.ONLINE_SOFTMAX_LSE.value, + "acc_dtype": AttentionDType.FP32.value, + "order": ReductionOrder.GLOBAL_BLOCK_INDEX.value, + "downcast_at": DowncastPoint.FINAL_WRITE.value, +} + + +#: Materialization facts the two sides are expected to differ on. Recorded into +#: provenance; never a rejection reason. +RECORDED_FIELDS: tuple[str, ...] = ( + "mode", + "backend_id", + "reduction.engine", + "rope.fusion_boundary", + "rope.q_state", + "rope.k_state", + "rope.k_cache_state", + "rope.cast_at", + "rope.output_dtype", + "preprocess.qk_rmsnorm", + "preprocess.rope", + "preprocess.fallback", + "kv_cache.page_size", + "kv_cache.prefix_cache_enabled", + "kv_cache.block_table_shape", +) + + +@dataclass(frozen=True) +class BindingIssue: + """One reason a binding is not comparable or not admissible.""" + + code: BindingErrorCode + tier: BindingTier + field: str + rollout: Any = None + training: Any = None + message: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "code": self.code.value, + "tier": self.tier.value, + "field": self.field, + "rollout": self.rollout, + "training": self.training, + "message": self.message, + } + + +@dataclass(frozen=True) +class AttentionBindingResult: + """Outcome of binding one rollout contract to one training contract. + + ``comparable`` and ``passed`` are deliberately separate. A pair whose identity + does not match is *not comparable* -- reporting a drift number for it would be + meaningless. A pair that is comparable but violates the reduction mandate *is* + comparable yet must still fail closed, because the whole WS2 claim is that + reduction order and accumulation precision come from the contract. + """ + + comparable: bool + passed: bool + issues: tuple[BindingIssue, ...] = () + identity_fingerprint: str = "" + reduction_fingerprint: str = "" + binding_fingerprint: str = "" + recorded_differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + provenance: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.attention_binding.v3" + + def issues_by_code(self, code: BindingErrorCode) -> tuple[BindingIssue, ...]: + return tuple(issue for issue in self.issues if issue.code is code) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "comparable": self.comparable, + "passed": self.passed, + "issues": [issue.to_dict() for issue in self.issues], + "identity_fingerprint": self.identity_fingerprint, + "reduction_fingerprint": self.reduction_fingerprint, + "binding_fingerprint": self.binding_fingerprint, + "recorded_differences": { + key: dict(value) for key, value in self.recorded_differences.items() + }, + "provenance": dict(self.provenance), + } + + +def _canonical_fingerprint(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def identity_fingerprint(identity: Mapping[str, Any]) -> str: + """Fingerprint only the declared :data:`IDENTITY_FIELDS`, in a fixed order. + + Extra keys in ``identity`` are ignored on purpose: callers pass whole + provenance bundles, and the fingerprint must not drift when an unrelated + diagnostic field is added. + """ + + return _canonical_fingerprint({name: identity.get(name) for name in IDENTITY_FIELDS}) + + +def _reduction_view(contract: AttentionContract) -> dict[str, Any]: + reduction = contract.reduction + return { + "merge": reduction.merge.value, + "acc_dtype": reduction.acc_dtype.value, + "order": reduction.order.value, + "downcast_at": reduction.downcast_at.value, + "engine": reduction.engine.value, + } + + +def _recorded_view( + contract: AttentionContract, + extra: Optional[Mapping[str, Any]] = None, +) -> dict[str, Any]: + rope = contract.rope + kv_cache = contract.kv_cache + view: dict[str, Any] = { + "mode": contract.mode.value, + "backend_id": None, + "reduction.engine": contract.reduction.engine.value, + } + if rope is not None: + view.update( + { + "rope.fusion_boundary": rope.fusion_boundary.value, + "rope.q_state": rope.q_state.value, + "rope.k_state": rope.k_state.value, + "rope.k_cache_state": rope.k_cache_state.value, + "rope.cast_at": rope.cast_at.value, + "rope.output_dtype": rope.output_dtype.value, + } + ) + if kv_cache is not None: + view.update( + { + "kv_cache.page_size": kv_cache.page_size, + "kv_cache.prefix_cache_enabled": kv_cache.prefix_cache_enabled, + "kv_cache.block_table_shape": [ + len(kv_cache.block_table), + max((len(row) for row in kv_cache.block_table), default=0), + ], + } + ) + if extra: + view.update(extra) + return view + + +def _topology_view(contract: AttentionContract) -> dict[str, Any]: + sharding = contract.sharding + return {name: getattr(sharding, name) for name in TOPOLOGY_FIELDS} + + +def _split_kv_fallbacks(plan_set: SplitKVRuntimePlanSet) -> list[dict[str, Any]]: + return [ + entry.to_dict() + for entry in plan_set.entries + if entry.execution.fallback + or entry.execution.actual_mode is None + or entry.execution.actual_mode is not entry.execution.requested_mode + or entry.execution.actual_split_size != entry.execution.requested_split_size + ] + + +def _split_kv_plan_contract_error( + contract: AttentionContract, + plan_set: SplitKVRuntimePlanSet, +) -> str | None: + sharding = contract.sharding + expected_topology = ( + contract.batch_size, + sharding.tp_world_size, + sharding.cp_world_size, + ) + actual_topology = ( + plan_set.batch_size, + plan_set.tp_world_size, + plan_set.cp_world_size, + ) + if actual_topology != expected_topology: + return ( + "Split-KV plan-set batch/TP/CP topology does not match the attention " + f"contract: actual={actual_topology}, expected={expected_topology}" + ) + if contract.mode.value in {"prefill", "chunked_prefill"}: + expected_totals = (sharding.global_sequence_length,) * contract.batch_size + if plan_set.total_kv_tokens != expected_totals: + return ( + "Split-KV plan-set KV lengths do not match the prefill attention " + f"contract: actual={plan_set.total_kv_tokens}, expected={expected_totals}" + ) + elif contract.kv_cache is not None: + expected_totals = contract.kv_cache.kv_seq_lens + if plan_set.total_kv_tokens != expected_totals: + return ( + "Split-KV plan-set KV lengths do not match decode KV-cache lengths: " + f"actual={plan_set.total_kv_tokens}, expected={expected_totals}" + ) + for entry in plan_set.entries: + execution = entry.execution + if ( + execution.requested_mode is not contract.split_kv.mode + or execution.requested_split_size != contract.split_kv.fixed_split_size + ): + return ( + "Split-KV runtime request does not match the first-class attention " + f"contract at {entry.coordinate}" + ) + return None + + +#: Identity fields where ``None`` is a real value rather than an omission. Qwen3-8B +#: applies no RoPE scaling, so ``rope_scaling=None`` must not read as "undeclared" -- +#: both sides still have to agree on it, which the equality pass below handles. +NULLABLE_IDENTITY_FIELDS: frozenset[str] = frozenset({"rope_scaling"}) + + +def _missing_identity_fields(identity: Mapping[str, Any]) -> tuple[str, ...]: + return tuple( + name + for name in IDENTITY_FIELDS + if name not in NULLABLE_IDENTITY_FIELDS and identity.get(name) is None + ) + + +def bind_attention_contracts( + *, + rollout_contract: AttentionContract, + training_contract: AttentionContract, + rollout_identity: Mapping[str, Any], + training_identity: Mapping[str, Any], + rollout_backend_id: str, + training_backend_id: str, + determinism_issues: Sequence[BindingIssue] = (), + require_full_identity: bool = True, + allow_dtype_difference: bool = False, + rollout_split_kv_plan_set: Optional[SplitKVRuntimePlanSet] = None, + training_split_kv_plan_set: Optional[SplitKVRuntimePlanSet] = None, + rollout_recorded_extra: Optional[Mapping[str, Any]] = None, + training_recorded_extra: Optional[Mapping[str, Any]] = None, +) -> AttentionBindingResult: + """Bind a rollout attention contract to a training attention contract. + + ``determinism_issues`` is threaded in from + :mod:`rl_engine.alignment.cross_config.determinism` rather than computed here, + so that this module stays free of framework probing and remains testable + without Megatron or vLLM present. + + ``require_full_identity`` exists for the single-GPU harness in #235 PR2, which + legitimately has no KV-cache or decode identity to declare. Distributed callers + must leave it at ``True``. + + ``allow_dtype_difference`` exists for the #235 PR5 sweep that deliberately scores + a BF16 path against an FP32 reference. It must stay ``False`` everywhere else. + + Strict binding requires complete actual Split-KV plan sets from both runtimes. + A configured policy is insufficient because auto-selection, graph capture, and + backend fallbacks can change the executed boundaries. The plan sets cover the + complete batch x TP x CP x KV-owner Cartesian product. + + ``rollout_recorded_extra`` / ``training_recorded_extra`` are diagnostic-only + backend facts. They can never make a semantic mismatch admissible. + """ + + if rollout_contract.role is not AttentionRole.INFER: + raise AttentionBindingError( + f"rollout_contract.role must be {AttentionRole.INFER.value!r}, " + f"got {rollout_contract.role.value!r}" + ) + if training_contract.role is not AttentionRole.TRAIN: + raise AttentionBindingError( + f"training_contract.role must be {AttentionRole.TRAIN.value!r}, " + f"got {training_contract.role.value!r}" + ) + + issues: list[BindingIssue] = [] + + # ---- tier 1: identity, bit for bit ------------------------------------- + if require_full_identity: + for side, identity in (("rollout", rollout_identity), ("training", training_identity)): + for name in _missing_identity_fields(identity): + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISSING, + tier=BindingTier.IDENTICAL, + field=f"{side}.{name}", + message=f"{side} identity does not declare {name!r}", + ) + ) + + for name in IDENTITY_FIELDS: + rollout_value = rollout_identity.get(name) + training_value = training_identity.get(name) + if rollout_value != training_value: + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISMATCH, + tier=BindingTier.IDENTICAL, + field=name, + rollout=rollout_value, + training=training_value, + message=( + f"{name!r} differs between sides; the pair is not comparable " + "and any drift computed from it is meaningless" + ), + ) + ) + + comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + + rollout_topology = _topology_view(rollout_contract) + training_topology = _topology_view(training_contract) + for name in TOPOLOGY_FIELDS: + if rollout_topology[name] != training_topology[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.TOPOLOGY_MISMATCH, + tier=BindingTier.IDENTICAL, + field=f"sharding.{name}", + rollout=rollout_topology[name], + training=training_topology[name], + message=( + f"sharding.{name} changes TP/CP ownership; the pair is not " + "the same local attention problem" + ), + ) + ) + + comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + + # ---- tier 2: reduction semantics, and the WS2 mandate ------------------- + rollout_reduction = _reduction_view(rollout_contract) + training_reduction = _reduction_view(training_contract) + + for name in SEMANTIC_REDUCTION_FIELDS: + if rollout_reduction[name] != training_reduction[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"reduction.{name!r} must be decided by the contract, not by the " + "backend; the two sides disagree" + ), + ) + ) + mandated = WS2_ATTENTION_REDUCTION_MANDATE[name] + for side, view in (("rollout", rollout_reduction), ("training", training_reduction)): + if view[name] != mandated: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_MANDATE_VIOLATION, + tier=BindingTier.SEMANTIC, + field=f"{side}.reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"WS2 requires reduction.{name} == {mandated!r}; " + f"{side} declares {view[name]!r}" + ), + ) + ) + + if not allow_dtype_difference and rollout_contract.dtype is not training_contract.dtype: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field="dtype", + rollout=rollout_contract.dtype.value, + training=training_contract.dtype.value, + message=( + "the two sides compute in different dtypes; the resulting drift is " + "not attributable. Pass allow_dtype_difference=True only for a " + "deliberate precision sweep" + ), + ) + ) + + if rollout_contract.split_kv != training_contract.split_kv: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field="split_kv", + rollout=rollout_contract.split_kv.to_dict(), + training=training_contract.split_kv.to_dict(), + message="training and rollout must request the same first-class Split-KV policy", + ) + ) + + for side, plan_set in ( + ("rollout", rollout_split_kv_plan_set), + ("training", training_split_kv_plan_set), + ): + if plan_set is None: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_RUNTIME_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + message=( + f"{side} did not report a complete actual Split-KV plan set; " + "configured policy alone is not runtime evidence" + ), + ) + ) + continue + contract = rollout_contract if side == "rollout" else training_contract + contract_error = _split_kv_plan_contract_error(contract, plan_set) + if contract_error is not None: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + rollout=plan_set.to_dict() if side == "rollout" else None, + training=plan_set.to_dict() if side == "training" else None, + message=contract_error, + ) + ) + fallbacks = _split_kv_fallbacks(plan_set) + if fallbacks: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_FALLBACK, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + rollout=fallbacks if side == "rollout" else None, + training=fallbacks if side == "training" else None, + message=f"{side} Split-KV runtime used an unknown or fallback plan", + ) + ) + + if rollout_split_kv_plan_set is not None and training_split_kv_plan_set is not None: + try: + validate_split_kv_plan_set_alignment( + training_split_kv_plan_set, + rollout_split_kv_plan_set, + ) + except AttentionContractError as exc: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field="split_kv_runtime_plan_set", + rollout=rollout_split_kv_plan_set.to_dict(), + training=training_split_kv_plan_set.to_dict(), + message=str(exc), + ) + ) + + for side, contract in (("rollout", rollout_contract), ("training", training_contract)): + if not contract.export_lse: + issues.append( + BindingIssue( + code=BindingErrorCode.LSE_NOT_EXPORTED, + tier=BindingTier.SEMANTIC, + field=f"{side}.export_lse", + message=( + "attention-domain LSE must be exported; without it the deterministic " + "CP merge cannot be validated" + ), + ) + ) + + if rollout_backend_id == training_backend_id and rollout_backend_id: + # Not an error, but worth surfacing: an identical backend on both sides means + # the experiment is not actually measuring a cross-implementation difference. + pass + + issues.extend(determinism_issues) + + # ---- tier 3: recorded differences -------------------------------------- + rollout_recorded = _recorded_view(rollout_contract, rollout_recorded_extra) + rollout_recorded["backend_id"] = rollout_backend_id + training_recorded = _recorded_view(training_contract, training_recorded_extra) + training_recorded["backend_id"] = training_backend_id + + recorded_differences: dict[str, dict[str, Any]] = {} + for name in RECORDED_FIELDS: + rollout_value = rollout_recorded.get(name) + training_value = training_recorded.get(name) + if rollout_value != training_value: + recorded_differences[name] = { + "rollout": rollout_value, + "training": training_value, + } + + identity_fp = identity_fingerprint(training_identity if comparable else rollout_identity) + reduction_fp = _canonical_fingerprint( + {name: training_reduction[name] for name in SEMANTIC_REDUCTION_FIELDS} + ) + passed = comparable and not any(issue.tier is BindingTier.SEMANTIC for issue in issues) + + provenance = { + "lse_domain": ATTENTION_LSE_DOMAIN, + "dtype": training_contract.dtype.value, + "split_kv_runtime": { + "rollout": ( + None if rollout_split_kv_plan_set is None else rollout_split_kv_plan_set.to_dict() + ), + "training": ( + None if training_split_kv_plan_set is None else training_split_kv_plan_set.to_dict() + ), + }, + "rollout": { + "contract": rollout_contract.to_dict(), + "backend_id": rollout_backend_id, + "recorded": rollout_recorded, + }, + "training": { + "contract": training_contract.to_dict(), + "backend_id": training_backend_id, + "recorded": training_recorded, + }, + } + + return AttentionBindingResult( + comparable=comparable, + passed=passed, + issues=tuple(issues), + identity_fingerprint=identity_fp, + reduction_fingerprint=reduction_fp, + binding_fingerprint=_canonical_fingerprint( + { + "identity": identity_fp, + "reduction": reduction_fp, + "topology": training_topology, + "split_kv": provenance["split_kv_runtime"], + "lse_domain": ATTENTION_LSE_DOMAIN, + "rollout_backend": rollout_backend_id, + "training_backend": training_backend_id, + "attention_preprocess": { + "rollout": { + name: rollout_recorded.get(f"preprocess.{name}") + for name in ( + *MANDATED_ATTENTION_PREPROCESS_BACKENDS, + "fallback", + ) + }, + "training": { + name: training_recorded.get(f"preprocess.{name}") + for name in ( + *MANDATED_ATTENTION_PREPROCESS_BACKENDS, + "fallback", + ) + }, + }, + } + ), + recorded_differences=recorded_differences, + provenance=provenance, + ) + + +def bind_attention_runtime_readbacks( + *, + rollout: AttentionRuntimeReadback, + training: AttentionRuntimeReadback, + rollout_identity: Mapping[str, Any], + training_identity: Mapping[str, Any], + rollout_backend_id: str, + training_backend_id: str, + determinism_issues: Sequence[BindingIssue] = (), +) -> AttentionBindingResult: + """Strict public handoff from executed framework runtimes to PR4 binding. + + The Megatron/vLLM launchers remain environment-owned. Once both launchers have + reconstructed their actual contracts and all-rank Split-KV reports, this entry + point performs the complete comparison without accepting configured-only data. + """ + + missing_scope_evidence = [] + for side, readback in (("rollout", rollout), ("training", training)): + if not readback.frozen_scope_verified: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"{side}.frozen_scope_verified", + message=f"{side} runtime did not verify the frozen attention scope", + ) + ) + missing_scope_evidence.extend(_attention_preprocess_issues(side, readback)) + return bind_attention_contracts( + rollout_contract=rollout.contract, + training_contract=training.contract, + rollout_identity=rollout_identity, + training_identity=training_identity, + rollout_backend_id=rollout_backend_id, + training_backend_id=training_backend_id, + determinism_issues=tuple(determinism_issues) + tuple(missing_scope_evidence), + rollout_split_kv_plan_set=rollout.split_kv_plan_set, + training_split_kv_plan_set=training.split_kv_plan_set, + rollout_recorded_extra={ + **{ + f"preprocess.{name}": backend + for name, backend in rollout.preprocess_backends.items() + }, + "preprocess.fallback": rollout.preprocess_fallback, + }, + 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.""" + + 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/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/determinism.py b/rl_engine/alignment/cross_config/determinism.py new file mode 100644 index 00000000..be81654e --- /dev/null +++ b/rl_engine/alignment/cross_config/determinism.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-side determinism probing for the Megatron + vLLM cross-config target. + +Both frameworks ship a "make this deterministic" switch, but they mean different +things by it, and neither knows the other exists: + +``Megatron`` ``ModelParallelConfig.deterministic_mode`` + Asserts ``NCCL_ALGO`` is one of five values, forbids FlashAttention and fused + cross-entropy, calls ``torch.use_deterministic_algorithms(True)``, and requires + ``NVTE_ALLOW_NONDETERMINISTIC_ALGO == 0``. It does **not** touch TF32, BF16 + reduced-precision reduction, cuBLAS workspace, NCCL protocol, or NCCL channel + counts. + +``vLLM`` ``VLLM_BATCH_INVARIANT`` + Replaces ``aten::mm/addmm/matmul/linear/bmm``, ``log_softmax``/``softmax``, + ``mean.dim`` and ``rms_norm`` with Triton kernels, disables TF32 and BF16/FP16 + reduced-precision reduction, pins cuBLAS workspace and the BLAS library, and + hard-sets ten NCCL environment variables. + +So a run can have both switches on and still be comparing two different notions of +determinism. This module makes that difference explicit and, where it changes the +numerics, blocking. It never imports Megatron or vLLM: probes are built from plain +mappings so the logic is testable on any machine. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, Optional + +from rl_engine.alignment.cross_config.attention_binding import ( + BindingErrorCode, + BindingIssue, + BindingTier, +) + +__all__ = [ + "COMPARED_NCCL_KEYS", + "DeterminismProbe", + "DeterminismReport", + "compare_determinism", + "megatron_probe_from_config", + "vllm_probe_from_env", +] + + +#: Environment keys whose value can change a reduction result. Compared across +#: sides; a difference is reported, and a difference in the *arithmetic* subset is +#: blocking. Ordering is fixed so the fingerprint is stable. +COMPARED_NCCL_KEYS: tuple[str, ...] = ( + "NCCL_ALGO", + "NCCL_PROTO", + "NCCL_MIN_NCHANNELS", + "NCCL_MAX_NCHANNELS", + "NCCL_NTHREADS", + "NCCL_SOCKET_NTHREADS", + "NCCL_COLLNET_ENABLE", + "NCCL_NVLS_ENABLE", + "NCCL_P2P_NET_DISABLE", + "NCCL_LAUNCH_MODE", + "CUBLAS_WORKSPACE_CONFIG", +) + + +#: The subset above that changes arithmetic rather than only scheduling. A mismatch +#: here fails the binding closed; a mismatch in the remainder is recorded only. +_ARITHMETIC_NCCL_KEYS: frozenset[str] = frozenset( + {"NCCL_ALGO", "NCCL_PROTO", "CUBLAS_WORKSPACE_CONFIG"} +) + + +@dataclass(frozen=True) +class DeterminismProbe: + """What one side actually has switched on. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are tri-state on + purpose: ``None`` means "the framework does not manage this", which is exactly + Megatron's situation and is itself the finding. + """ + + side: str + framework: str + mode_flag: str + enabled: bool + env: Mapping[str, Any] = field(default_factory=dict) + tf32_disabled: Optional[bool] = None + bf16_reduced_precision_reduction: Optional[bool] = None + forbids_flash_attention: Optional[bool] = None + evidence: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_probe.v1" + + def __post_init__(self) -> None: + if self.side not in ("rollout", "training"): + raise ValueError("side must be 'rollout' or 'training'") + if not self.framework: + raise ValueError("framework must not be empty") + object.__setattr__(self, "env", dict(self.env)) + object.__setattr__(self, "evidence", dict(self.evidence)) + + @property + def env_fingerprint(self) -> str: + payload = {key: self.env.get(key) for key in COMPARED_NCCL_KEYS} + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "side": self.side, + "framework": self.framework, + "mode_flag": self.mode_flag, + "enabled": self.enabled, + "env": {key: self.env.get(key) for key in COMPARED_NCCL_KEYS}, + "env_fingerprint": self.env_fingerprint, + "tf32_disabled": self.tf32_disabled, + "bf16_reduced_precision_reduction": self.bf16_reduced_precision_reduction, + "forbids_flash_attention": self.forbids_flash_attention, + "evidence": dict(self.evidence), + } + + +@dataclass(frozen=True) +class DeterminismReport: + """Cross-side comparison result.""" + + rollout: DeterminismProbe + training: DeterminismProbe + issues: tuple[BindingIssue, ...] = () + differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_report.v1" + + @property + def compatible(self) -> bool: + return not self.issues + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "compatible": self.compatible, + "rollout": self.rollout.to_dict(), + "training": self.training.to_dict(), + "issues": [issue.to_dict() for issue in self.issues], + "differences": {key: dict(value) for key, value in self.differences.items()}, + } + + +def megatron_probe_from_config( + config: Any, + env: Optional[Mapping[str, str]] = None, +) -> DeterminismProbe: + """Build a training-side probe from a Megatron config object. + + ``config`` is duck-typed (anything exposing ``deterministic_mode`` and + optionally ``attention_backend`` / ``cross_entropy_loss_fusion``) so this works + against a real ``ModelParallelConfig``, a test double, or a plain namespace, + and so importing this module never requires Megatron. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are reported as + ``None`` because Megatron does not manage them -- a ``grep`` for ``allow_tf32`` + and ``fp32_precision`` across ``megatron/`` returns nothing. That asymmetry + against vLLM is the point of :func:`compare_determinism`. + """ + + environ = dict(env or {}) + enabled = bool(getattr(config, "deterministic_mode", False)) + return DeterminismProbe( + side="training", + framework="megatron", + mode_flag="deterministic_mode", + enabled=enabled, + env={key: environ.get(key) for key in COMPARED_NCCL_KEYS}, + tf32_disabled=None, + bf16_reduced_precision_reduction=None, + forbids_flash_attention=enabled, + evidence={ + "nvte_allow_nondeterministic_algo": environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO"), + "cross_entropy_loss_fusion": getattr(config, "cross_entropy_loss_fusion", None), + "attention_backend": _enum_value(getattr(config, "attention_backend", None)), + "tensor_model_parallel_size": getattr(config, "tensor_model_parallel_size", None), + "context_parallel_size": getattr(config, "context_parallel_size", None), + "sequence_parallel": getattr(config, "sequence_parallel", None), + "manages_tf32": False, + "manages_bf16_reduced_precision_reduction": False, + }, + ) + + +def vllm_probe_from_env( + env: Mapping[str, str], + *, + model_config: Any = None, +) -> DeterminismProbe: + """Build a rollout-side probe from the vLLM process environment. + + ``VLLM_BATCH_INVARIANT`` is read from ``env`` rather than ``vllm.envs`` so the + probe can be constructed from a remote worker's reported environment, which is + how vime's Ray actors expose it. + """ + + enabled = str(env.get("VLLM_BATCH_INVARIANT", "0")).strip() in ("1", "true", "True") + return DeterminismProbe( + side="rollout", + framework="vllm", + mode_flag="VLLM_BATCH_INVARIANT", + enabled=enabled, + env={key: env.get(key) for key in COMPARED_NCCL_KEYS}, + # vLLM sets both to "ieee"/disabled inside init_batch_invariance(). + tf32_disabled=enabled or None, + bf16_reduced_precision_reduction=(False if enabled else None), + forbids_flash_attention=False, + evidence={ + "vllm_allreduce_use_symm_mem": env.get("VLLM_ALLREDUCE_USE_SYMM_MEM"), + "vllm_use_aot_compile": env.get("VLLM_USE_AOT_COMPILE"), + "enforce_eager": getattr(model_config, "enforce_eager", None), + "disable_cascade_attn": getattr(model_config, "disable_cascade_attn", None), + "quantization": getattr(model_config, "quantization", None), + "manages_tf32": True, + "manages_bf16_reduced_precision_reduction": True, + }, + ) + + +def _enum_value(value: Any) -> Any: + return getattr(value, "value", value) + + +def compare_determinism( + *, + rollout: DeterminismProbe, + training: DeterminismProbe, +) -> DeterminismReport: + """Compare two probes and produce blocking issues plus recorded differences.""" + + if rollout.side != "rollout" or training.side != "training": + raise ValueError("compare_determinism expects one rollout probe and one training probe") + + issues: list[BindingIssue] = [] + differences: dict[str, dict[str, Any]] = {} + + for probe in (rollout, training): + if not probe.enabled: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"{probe.side}.{probe.mode_flag}", + rollout=rollout.enabled, + training=training.enabled, + message=( + f"{probe.framework} {probe.mode_flag} is not enabled; the " + f"{probe.side} side is not batch-invariant and cannot anchor a " + "cross-config comparison" + ), + ) + ) + + for key in COMPARED_NCCL_KEYS: + rollout_value = rollout.env.get(key) + training_value = training.env.get(key) + if rollout_value == training_value: + continue + differences[key] = {"rollout": rollout_value, "training": training_value} + if key in _ARITHMETIC_NCCL_KEYS: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"env.{key}", + rollout=rollout_value, + training=training_value, + message=( + f"{key} differs between sides; the two sides would reduce with " + "different arithmetic and the resulting drift is not attributable" + ), + ) + ) + + # Megatron reports None for these because it does not manage them at all. That is + # recorded rather than blocking: under a pure BF16 GEMM path TF32 does not fire, and + # forcing Megatron to manage it is out of scope for this PR. It is surfaced so the + # asymmetry appears in every artifact instead of being invisible. + for name in ("tf32_disabled", "bf16_reduced_precision_reduction"): + rollout_value = getattr(rollout, name) + training_value = getattr(training, name) + if rollout_value != training_value: + differences[name] = { + "rollout": rollout_value, + "training": training_value, + "note": ( + "megatron does not manage this setting; vllm sets it inside " + "init_batch_invariance()" + ), + } + + return DeterminismReport( + rollout=rollout, + training=training, + issues=tuple(issues), + differences=differences, + ) diff --git a/rl_engine/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/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/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/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", +] 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/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/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", +] 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 diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py new file mode 100644 index 00000000..206663e7 --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,1526 @@ +# 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 build_split_kv_runtime_plan_set( + total_kv_tokens: Iterable[int], + *, + tp_world_size: int, + cp_world_size: int, + split_kv: SplitKVSpec, + backend: str = "contract_reference", +) -> SplitKVRuntimePlanSet: + """Build a complete owner-local plan set for contract tests and adapters.""" + + totals = _integer_tuple(total_kv_tokens, "total_kv_tokens") + if not totals or any(total < cp_world_size for total in totals): + raise AttentionContractError( + "contract plan sets require at least one KV token per CP owner" + ) + tp_world_size = _positive_int(tp_world_size, "tp_world_size") + cp_world_size = _positive_int(cp_world_size, "cp_world_size") + if not isinstance(split_kv, SplitKVSpec): + raise AttentionContractError("split_kv must be a SplitKVSpec") + + entries: list[SplitKVRuntimePlanEntry] = [] + for batch_index, total in enumerate(totals): + base = total // cp_world_size + remainder = total % cp_world_size + owner_ranges: list[tuple[int, int]] = [] + start = 0 + for owner_cp_rank in range(cp_world_size): + end = start + base + (1 if owner_cp_rank < remainder else 0) + owner_ranges.append((start, end)) + start = end + for tp_rank in range(tp_world_size): + for cp_rank in range(cp_world_size): + for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + local_total = owner_end - owner_start + local = split_kv.resolve(local_total, backend=backend) + execution = SplitKVExecutionPlan( + requested_mode=local.requested_mode, + requested_split_size=local.requested_split_size, + actual_mode=local.actual_mode, + actual_split_size=local.actual_split_size, + boundaries=tuple( + (owner_start + start, owner_start + end) + for start, end in local.boundaries + ), + merge_order=local.merge_order, + acc_dtype=local.acc_dtype, + downcast_at=local.downcast_at, + backend=local.backend, + source=local.source, + fallback=local.fallback, + fallback_reason=local.fallback_reason, + ) + entries.append( + SplitKVRuntimePlanEntry( + coordinate=SplitKVRuntimeCoordinate( + batch_index=batch_index, + tp_rank=tp_rank, + cp_rank=cp_rank, + owner_cp_rank=owner_cp_rank, + ), + expected_kv_range=(owner_start, owner_end), + execution=execution, + ) + ) + return SplitKVRuntimePlanSet( + batch_size=len(totals), + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals, + entries=tuple(entries), + ) + + +def _format_split_kv_coordinates( + coordinates: Iterable[SplitKVRuntimeCoordinate], +) -> list[dict[str, int]]: + return [coordinate.to_dict() for coordinate in sorted(coordinates)] + + +@dataclass(frozen=True) +class KVCacheSpec: + """Logical identity of the paged/block KV cache used for replay.""" + + cache_positions: tuple[int, ...] + kv_seq_lens: tuple[int, ...] + block_table: tuple[tuple[int, ...], ...] + global_token_positions: tuple[int, ...] + page_size: int + prefix_cache_enabled: bool = False + prefix_cache_key: str | None = None + shared_prefix_page_count: int = 0 + + def __post_init__(self) -> None: + cache_positions = _integer_tuple(self.cache_positions, "cache_positions") + kv_seq_lens = _integer_tuple(self.kv_seq_lens, "kv_seq_lens") + page_size = _positive_int(self.page_size, "page_size") + global_token_positions = _integer_tuple( + self.global_token_positions, "global_token_positions" + ) + if not cache_positions or any(position < 0 for position in cache_positions): + raise AttentionContractError("cache_positions must contain non-negative positions") + if not kv_seq_lens or any(length <= 0 for length in kv_seq_lens): + raise AttentionContractError("kv_seq_lens must contain positive sequence lengths") + if len(cache_positions) != len(kv_seq_lens): + raise AttentionContractError( + "cache_positions must contain one entry per kv_seq_lens entry" + ) + if not global_token_positions or any(position < 0 for position in global_token_positions): + raise AttentionContractError( + "global_token_positions must contain non-negative positions" + ) + if len(global_token_positions) != sum(kv_seq_lens): + raise AttentionContractError( + "global_token_positions must describe every logical cached token; " + f"expected {sum(kv_seq_lens)}, got {len(global_token_positions)}" + ) + token_offset = 0 + sequence_position_rows: list[tuple[int, ...]] = [] + for sequence_index, sequence_length in enumerate(kv_seq_lens): + sequence_positions = global_token_positions[ + token_offset : token_offset + sequence_length + ] + if any( + left >= right + for left, right in zip(sequence_positions, sequence_positions[1:], strict=False) + ): + raise AttentionContractError( + "global_token_positions must be strictly increasing within each sequence; " + f"sequence {sequence_index} is invalid" + ) + sequence_position_rows.append(sequence_positions) + token_offset += sequence_length + + for sequence_index, (cache_position, sequence_positions) in enumerate( + zip(cache_positions, sequence_position_rows, strict=True) + ): + terminal_position = sequence_positions[-1] + if cache_position != terminal_position: + raise AttentionContractError( + "cache_positions must equal the terminal global token position for each " + f"sequence; sequence {sequence_index} expected {terminal_position}, " + f"got {cache_position}" + ) + + try: + block_table = tuple(tuple(row) for row in self.block_table) + except TypeError as exc: + raise AttentionContractError( + "block_table must be a two-dimensional integer table" + ) from exc + if len(block_table) != len(kv_seq_lens) or any(not row for row in block_table): + raise AttentionContractError( + "block_table must contain one non-empty row per kv_seq_lens entry" + ) + active_block_rows: list[tuple[int, ...]] = [] + for row_index, (row, sequence_length) in enumerate( + zip(block_table, kv_seq_lens, strict=True) + ): + row_active_blocks: list[int] = [] + saw_padding = False + for column_index, block in enumerate(row): + if isinstance(block, bool) or not isinstance(block, int) or block < -1: + raise AttentionContractError( + "block_table entries must be integer block ids or -1 padding; " + f"got block_table[{row_index}][{column_index}]={block!r}" + ) + if block == -1: + saw_padding = True + continue + if saw_padding: + raise AttentionContractError( + "block_table -1 padding must be trailing; " + f"row {row_index} contains an active block after padding" + ) + row_active_blocks.append(block) + + expected_blocks = (sequence_length + page_size - 1) // page_size + if len(row_active_blocks) != expected_blocks: + raise AttentionContractError( + "block_table active page count must match kv_seq_lens and page_size; " + f"row {row_index} expected {expected_blocks}, got {len(row_active_blocks)}" + ) + if len(set(row_active_blocks)) != len(row_active_blocks): + raise AttentionContractError( + f"block_table row {row_index} contains duplicate active page ids" + ) + active_block_rows.append(tuple(row_active_blocks)) + + if not isinstance(self.prefix_cache_enabled, bool): + raise AttentionContractError("prefix_cache_enabled must be a bool") + shared_prefix_page_count = _non_negative_int( + self.shared_prefix_page_count, "shared_prefix_page_count" + ) + if self.prefix_cache_enabled and not self.prefix_cache_key: + raise AttentionContractError( + "prefix_cache_key is required when prefix_cache_enabled=True" + ) + if not self.prefix_cache_enabled and self.prefix_cache_key is not None: + raise AttentionContractError( + "prefix_cache_key must be None when prefix_cache_enabled=False" + ) + if not self.prefix_cache_enabled and shared_prefix_page_count != 0: + raise AttentionContractError( + "shared_prefix_page_count must be 0 when prefix_cache_enabled=False" + ) + + shared_prefix_pages: tuple[int, ...] = () + if shared_prefix_page_count > 0: + if any(len(row_blocks) < shared_prefix_page_count for row_blocks in active_block_rows): + raise AttentionContractError( + "shared_prefix_page_count exceeds an active block-table row" + ) + shared_prefix_token_count = shared_prefix_page_count * page_size + if any(length < shared_prefix_token_count for length in kv_seq_lens): + raise AttentionContractError( + "shared prefix pages must be fully populated and read-only" + ) + + shared_prefix_pages = active_block_rows[0][:shared_prefix_page_count] + shared_prefix_positions = sequence_position_rows[0][:shared_prefix_token_count] + for sequence_index, (row_blocks, positions) in enumerate( + zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 + ): + if row_blocks[:shared_prefix_page_count] != shared_prefix_pages: + raise AttentionContractError( + "shared prefix page ids must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + if positions[:shared_prefix_token_count] != shared_prefix_positions: + raise AttentionContractError( + "shared prefix token positions must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + + exclusive_page_owners: dict[int, int] = {} + shared_prefix_page_ids = set(shared_prefix_pages) + for sequence_index, row_blocks in enumerate(active_block_rows): + for page_id in row_blocks[shared_prefix_page_count:]: + if page_id in shared_prefix_page_ids: + raise AttentionContractError( + "a writable suffix page cannot alias a read-only shared prefix page" + ) + previous_owner = exclusive_page_owners.get(page_id) + if previous_owner is not None: + raise AttentionContractError( + "active pages may be shared across sequences only when declared as " + "read-only prefix pages; " + f"page {page_id} is used by sequences {previous_owner} and " + f"{sequence_index}" + ) + exclusive_page_owners[page_id] = sequence_index + + object.__setattr__(self, "cache_positions", cache_positions) + object.__setattr__(self, "kv_seq_lens", kv_seq_lens) + object.__setattr__(self, "block_table", block_table) + object.__setattr__(self, "global_token_positions", global_token_positions) + + +@dataclass(frozen=True) +class RoPESpec: + """Qwen3 RoPE identity for attention inputs and KV-cache replay. + + The CP attention reference consumes attention inputs. This object records + whether those inputs are pre- or post-RoPE and pins the position/cache + metadata needed to compare fused ``RoPE+Attention`` and unfused + ``RoPE -> Attention`` materializations. + """ + + q_state: RoPEState = RoPEState.POST_ROPE + k_state: RoPEState = RoPEState.POST_ROPE + k_cache_state: RoPEState = RoPEState.POST_ROPE + theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + position_ids: tuple[int, ...] | None = None + query_position_offsets: tuple[int, ...] | None = None + key_position_offsets: tuple[int, ...] | None = None + cast_at: RoPECastPoint = RoPECastPoint.AFTER_ROPE + output_dtype: AttentionDType = AttentionDType.BF16 + fusion_boundary: RoPEFusionBoundary = RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION + + def __post_init__(self) -> None: + object.__setattr__(self, "q_state", _enum_value(RoPEState, self.q_state, "q_state")) + object.__setattr__(self, "k_state", _enum_value(RoPEState, self.k_state, "k_state")) + object.__setattr__( + self, "k_cache_state", _enum_value(RoPEState, self.k_cache_state, "k_cache_state") + ) + if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + theta = float(self.theta) + if theta <= 0.0: + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + object.__setattr__(self, "theta", theta) + _positive_int(self.rotary_dim, "rotary_dim") + if self.rope_scaling is not None and ( + not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() + ): + raise AttentionContractError("rope_scaling must be a non-empty string when provided") + for position_field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, position_field) + if values is None: + continue + normalized = _integer_tuple(values, position_field) + if not normalized or any(value < 0 for value in normalized): + raise AttentionContractError( + f"{position_field} must contain non-negative positions" + ) + object.__setattr__(self, position_field, normalized) + object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) + object.__setattr__( + self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") + ) + object.__setattr__( + self, + "fusion_boundary", + _enum_value(RoPEFusionBoundary, self.fusion_boundary, "fusion_boundary"), + ) + + +@dataclass(frozen=True) +class AttentionContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: AttentionRole + mode: AttentionMode + dtype: AttentionDType + batch_size: int + query_sequence_length: int + head_dim: int + causal: bool + causal_offsets: tuple[int, ...] | None + sharding: ShardingSpec + reduction: ReductionSpec + split_kv: SplitKVSpec = field(default_factory=SplitKVSpec.disabled) + kv_cache: KVCacheSpec | None = None + rope: RoPESpec | None = None + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role")) + object.__setattr__(self, "mode", _enum_value(AttentionMode, self.mode, "mode")) + object.__setattr__(self, "dtype", _enum_value(AttentionDType, self.dtype, "dtype")) + batch_size = _positive_int(self.batch_size, "batch_size") + query_sequence_length = _positive_int(self.query_sequence_length, "query_sequence_length") + _positive_int(self.head_dim, "head_dim") + if not isinstance(self.sharding, ShardingSpec): + raise AttentionContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise AttentionContractError("reduction must be a ReductionSpec") + if not isinstance(self.split_kv, SplitKVSpec): + raise AttentionContractError("split_kv must be a SplitKVSpec") + if ( + self.mode is AttentionMode.PREFILL + and query_sequence_length != self.sharding.local_sequence_length + ): + raise AttentionContractError( + "prefill query_sequence_length must equal sharding.local_sequence_length; " + f"got {query_sequence_length} and {self.sharding.local_sequence_length}" + ) + if self.sharding.packed_sequence_offsets is not None: + packed_sequence_count = len(self.sharding.packed_sequence_offsets) - 1 + if packed_sequence_count != batch_size: + raise AttentionContractError( + "packed sequence count must equal logical batch_size; " + f"got {packed_sequence_count} packed sequences and batch_size={batch_size}" + ) + if not isinstance(self.causal, bool): + raise AttentionContractError("causal must be a bool") + if self.causal: + if self.causal_offsets is None: + raise AttentionContractError("causal_offsets are required for causal attention") + if self.causal_offsets is not None: + causal_offsets = _integer_tuple(self.causal_offsets, "causal_offsets") + if not causal_offsets or any(offset < 0 for offset in causal_offsets): + raise AttentionContractError("causal_offsets must contain non-negative offsets") + if self.sharding.packed_sequence_offsets is not None: + expected_causal_offsets = batch_size + offset_owner = "packed sequence" + else: + expected_causal_offsets = batch_size + offset_owner = "batch entry" + if len(causal_offsets) != expected_causal_offsets: + raise AttentionContractError( + f"causal_offsets must contain one entry per {offset_owner}" + ) + object.__setattr__(self, "causal_offsets", causal_offsets) + + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise AttentionContractError( + "export_lse must be True for the WS2 attention-domain LSE contract" + ) + + if self.mode is AttentionMode.DECODE and self.kv_cache is None: + raise AttentionContractError("kv_cache metadata is required for decode attention") + if self.kv_cache is not None and not isinstance(self.kv_cache, KVCacheSpec): + raise AttentionContractError("kv_cache must be a KVCacheSpec when provided") + if self.rope is not None: + if not isinstance(self.rope, RoPESpec): + raise AttentionContractError("rope must be a RoPESpec when provided") + if self.rope.rotary_dim > self.head_dim: + raise AttentionContractError( + f"rotary_dim={self.rope.rotary_dim} must not exceed head_dim={self.head_dim}" + ) + if self.rope.position_ids is not None and len(self.rope.position_ids) not in { + query_sequence_length, + self.sharding.local_sequence_length, + }: + raise AttentionContractError( + "position_ids must describe the local query sequence or full local " + "sequence length" + ) + for position_field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, position_field) + if offsets is not None and len(offsets) != batch_size: + raise AttentionContractError( + f"{position_field} must contain one entry per logical batch entry" + ) + if self.mode is AttentionMode.DECODE and self.kv_cache is not None: + if len(self.kv_cache.kv_seq_lens) != batch_size: + raise AttentionContractError( + "decode kv_seq_lens must contain one entry per batch entry" + ) + if len(self.kv_cache.cache_positions) != batch_size: + raise AttentionContractError( + "decode cache_positions must contain one entry per batch entry" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "global_q_heads": self.sharding.global_q_heads, + "global_kv_heads": self.sharding.global_kv_heads, + "local_q_head_start": self.sharding.local_q_head_start, + "local_q_heads": self.sharding.local_q_heads, + "local_kv_head_start": self.sharding.local_kv_head_start, + "local_kv_heads": self.sharding.local_kv_heads, + "global_sequence_length": self.sharding.global_sequence_length, + "local_sequence_length": self.sharding.local_sequence_length, + "global_block_indices": list(self.sharding.global_block_indices), + "global_block_token_starts": list(self.sharding.global_block_token_starts), + "local_block_offsets": list(self.sharding.local_block_offsets), + "packed_sequence_offsets": ( + list(self.sharding.packed_sequence_offsets) + if self.sharding.packed_sequence_offsets is not None + else None + ), + } + reduction = { + "merge": self.reduction.merge.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + } + kv_cache = None + if self.kv_cache is not None: + kv_cache = { + "cache_positions": list(self.kv_cache.cache_positions), + "kv_seq_lens": list(self.kv_cache.kv_seq_lens), + "block_table": [list(row) for row in self.kv_cache.block_table], + "global_token_positions": list(self.kv_cache.global_token_positions), + "page_size": self.kv_cache.page_size, + "prefix_cache_enabled": self.kv_cache.prefix_cache_enabled, + "prefix_cache_key": self.kv_cache.prefix_cache_key, + "shared_prefix_page_count": self.kv_cache.shared_prefix_page_count, + } + rope = None + if self.rope is not None: + rope = { + "q_state": self.rope.q_state.value, + "k_state": self.rope.k_state.value, + "k_cache_state": self.rope.k_cache_state.value, + "theta": self.rope.theta, + "rotary_dim": self.rope.rotary_dim, + "rope_scaling": self.rope.rope_scaling, + "position_ids": ( + list(self.rope.position_ids) if self.rope.position_ids is not None else None + ), + "query_position_offsets": ( + list(self.rope.query_position_offsets) + if self.rope.query_position_offsets is not None + else None + ), + "key_position_offsets": ( + list(self.rope.key_position_offsets) + if self.rope.key_position_offsets is not None + else None + ), + "cast_at": self.rope.cast_at.value, + "output_dtype": self.rope.output_dtype.value, + "fusion_boundary": self.rope.fusion_boundary.value, + } + return { + "semantic_operator": "standard_softmax_attention", + "role": self.role.value, + "mode": self.mode.value, + "dtype": self.dtype.value, + "batch_size": self.batch_size, + "query_sequence_length": self.query_sequence_length, + "head_dim": self.head_dim, + "causal": self.causal, + "causal_offsets": ( + list(self.causal_offsets) if self.causal_offsets is not None else None + ), + "export_lse": self.export_lse, + "lse_domain": "attention", + "sharding": sharding, + "reduction": reduction, + "split_kv": self.split_kv.to_dict(), + "kv_cache": kv_cache, + "rope": rope, + } + + +@dataclass(frozen=True) +class AttentionBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[AttentionRole] + modes: frozenset[AttentionMode] + dtypes: frozenset[AttentionDType] + cp_world_sizes: tuple[int, ...] + tp_world_sizes: tuple[int, ...] | None = None + exports_attention_lse: bool = False + deterministic_cp_merge: bool = False + supports_packed_varlen: bool = False + supports_kv_cache: bool = False + supports_rope_metadata: bool = False + supports_fused_rope_attention: bool = False + supports_split_kv_disabled: bool = True + supports_split_kv_fixed: bool = False + supports_split_kv_auto: bool = False + reports_actual_split_kv_plan: bool = False + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise AttentionContractError("backend_id must be a non-empty string") + roles = frozenset(_enum_value(AttentionRole, value, "roles") for value in self.roles) + modes = frozenset(_enum_value(AttentionMode, value, "modes") for value in self.modes) + dtypes = frozenset(_enum_value(AttentionDType, value, "dtypes") for value in self.dtypes) + if not roles or not modes or not dtypes: + raise AttentionContractError("backend roles, modes, and dtypes must not be empty") + cp_world_sizes = _integer_tuple(self.cp_world_sizes, "cp_world_sizes") + if not cp_world_sizes or any(size <= 0 for size in cp_world_sizes): + raise AttentionContractError("cp_world_sizes must contain positive values") + if len(set(cp_world_sizes)) != len(cp_world_sizes): + raise AttentionContractError("cp_world_sizes must not contain duplicates") + tp_world_sizes = None + if self.tp_world_sizes is not None: + tp_world_sizes = _integer_tuple(self.tp_world_sizes, "tp_world_sizes") + if not tp_world_sizes or any(size <= 0 for size in tp_world_sizes): + raise AttentionContractError("tp_world_sizes must contain positive values") + if len(set(tp_world_sizes)) != len(tp_world_sizes): + raise AttentionContractError("tp_world_sizes must not contain duplicates") + for capability_field in ( + "exports_attention_lse", + "deterministic_cp_merge", + "supports_packed_varlen", + "supports_kv_cache", + "supports_rope_metadata", + "supports_fused_rope_attention", + "supports_split_kv_disabled", + "supports_split_kv_fixed", + "supports_split_kv_auto", + "reports_actual_split_kv_plan", + ): + if not isinstance(getattr(self, capability_field), bool): + raise AttentionContractError(f"{capability_field} must be a bool") + if self.implementation_kind not in {"production", "reference", "deterministic"}: + raise AttentionContractError( + "implementation_kind must be production, reference, or deterministic" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "modes", modes) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + + def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.mode not in self.modes: + reasons.append(f"mode={contract.mode.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if contract.export_lse and not self.exports_attention_lse: + reasons.append("attention-domain LSE export is unsupported") + if cp_size > 1 and not self.deterministic_cp_merge: + reasons.append("deterministic CP (out, lse) merge is unsupported") + if ( + contract.sharding.packed_sequence_offsets is not None + and not self.supports_packed_varlen + ): + reasons.append("packed varlen layout is unsupported") + if contract.kv_cache is not None and not self.supports_kv_cache: + reasons.append("KV-cache identity materialization is unsupported") + if contract.rope is not None and not self.supports_rope_metadata: + reasons.append("RoPE/position metadata is unsupported") + if ( + contract.rope is not None + and contract.rope.fusion_boundary is RoPEFusionBoundary.FUSED_ROPE_ATTENTION + and not self.supports_fused_rope_attention + ): + reasons.append("fused RoPE+Attention boundary is unsupported") + split_support = { + SplitKVMode.DISABLED: self.supports_split_kv_disabled, + SplitKVMode.FIXED: self.supports_split_kv_fixed, + SplitKVMode.AUTO: self.supports_split_kv_auto, + } + if not split_support[contract.split_kv.mode]: + reasons.append(f"Split-KV policy={contract.split_kv.mode.value} is unsupported") + if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan: + reasons.append("actual Split-KV execution-plan provenance is unsupported") + return tuple(reasons) + + def supports(self, contract: AttentionContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "modes": sorted(mode.value for mode in self.modes), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes), + "exports_attention_lse": self.exports_attention_lse, + "deterministic_cp_merge": self.deterministic_cp_merge, + "supports_packed_varlen": self.supports_packed_varlen, + "supports_kv_cache": self.supports_kv_cache, + "supports_rope_metadata": self.supports_rope_metadata, + "supports_fused_rope_attention": self.supports_fused_rope_attention, + "supports_split_kv_disabled": self.supports_split_kv_disabled, + "supports_split_kv_fixed": self.supports_split_kv_fixed, + "supports_split_kv_auto": self.supports_split_kv_auto, + "reports_actual_split_kv_plan": self.reports_actual_split_kv_plan, + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class AttentionDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: AttentionBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "AttentionContract", + "AttentionContractError", + "AttentionBackendCapability", + "AttentionDispatchResult", + "AttentionDType", + "AttentionMerge", + "AttentionMode", + "AttentionRole", + "DowncastPoint", + "KVCacheSpec", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "RoPECastPoint", + "RoPEFusionBoundary", + "RoPESpec", + "RoPEState", + "ShardingSpec", + "build_split_kv_runtime_plan_set", + "SplitKVExecutionPlan", + "SplitKVMode", + "SplitKVRuntimeCoordinate", + "SplitKVRuntimePlanEntry", + "SplitKVRuntimePlanSet", + "SplitKVSpec", + "validate_split_kv_alignment", + "validate_split_kv_plan_set_alignment", +] diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py new file mode 100644 index 00000000..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/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 835ee0e4..0156907b 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -29,6 +29,7 @@ def make_operator_inputs( "matmul": _make_matmul_inputs, "det_gemm": _make_det_gemm_inputs, "attention": _make_attention_inputs, + "cp_attention": _make_cp_attention_inputs, "logp": _make_logp_inputs, "linear_logp": _make_linear_logp_inputs, "batch_invariant_logp": _make_batch_invariant_logp_inputs, @@ -53,6 +54,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "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}", @@ -139,6 +141,26 @@ def _make_attention_inputs( return 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/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 08925021..454f00ba 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -63,6 +63,22 @@ def _load_object(path: str) -> Any: }, 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", 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", +] 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/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py new file mode 100644 index 00000000..0062dac4 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -0,0 +1,1079 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic context-parallel attention reference. + +This module is the correctness-first WS2 reference for CP-aware standard +softmax attention. It intentionally stays in PyTorch and uses fp32 partial +states so fused CUDA/Triton backends can validate their CP/LSE merge semantics +against a small, inspectable implementation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional, Sequence + +import torch + +from rl_engine.kernels.attention_contract import ( + SplitKVExecutionPlan, + SplitKVMode, + SplitKVRuntimeCoordinate, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + +@dataclass(frozen=True) +class AttentionPartialState: + """One KV block's attention state before deterministic LSE merge. + + ``out`` is already normalized within the local KV block and has shape + ``[B, Hq, Sq, D]``. ``lse`` is the local attention-domain log-sum-exp with + shape ``[B, Hq, Sq]``. ``block_start`` / ``block_end`` are logical global KV + positions and define the canonical merge order. + """ + + out: torch.Tensor + lse: torch.Tensor + block_start: int + block_end: int + + def __post_init__(self) -> None: + if self.out.ndim != 4: + raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") + if self.lse.shape != self.out.shape[:3]: + raise ValueError("partial attention lse must have shape [B, Hq, Sq]") + if self.out.device != self.lse.device: + raise ValueError("partial attention out/lse must be on the same device") + if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: + raise ValueError("partial attention out/lse must remain FP32 before merge") + if self.block_start < 0: + raise ValueError("block_start must be non-negative") + if self.block_end < self.block_start: + raise ValueError("block_end must be >= block_start") + + +@dataclass(frozen=True) +class AttentionBackwardGradients: + """Training-side gradients emitted by the CP attention backward reference.""" + + dq: torch.Tensor + dk: torch.Tensor + dv: torch.Tensor + + +@dataclass(frozen=True) +class AttentionBackwardPathResult: + """One materialized CP attention backward path.""" + + name: str + out: torch.Tensor + lse: torch.Tensor + gradients: AttentionBackwardGradients + provenance: dict[str, object] + + +@dataclass(frozen=True) +class GradientDriftStats: + """Shape-aware absolute drift summary for backward validation reports.""" + + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + def to_dict(self) -> dict[str, object]: + return { + "max_abs": self.max_abs, + "mean_abs": self.mean_abs, + "p95_abs": self.p95_abs, + "p99_abs": self.p99_abs, + "active_count": self.active_count, + } + + +@dataclass(frozen=True) +class AttentionBackwardRankDrift: + """Backward drift for one logical CP rank's sequence ownership.""" + + rank: int + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + + def to_dict(self) -> dict[str, object]: + return { + "rank": self.rank, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + } + + +@dataclass(frozen=True) +class AttentionBackwardPathDrift: + """Candidate-vs-reference backward drift for one CP path.""" + + candidate_name: str + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + out: GradientDriftStats + lse: GradientDriftStats + per_rank: tuple[AttentionBackwardRankDrift, ...] + provenance: dict[str, object] + + def to_dict(self) -> dict[str, object]: + return { + "candidate_name": self.candidate_name, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + "out": self.out.to_dict(), + "lse": self.lse.to_dict(), + "per_rank": [item.to_dict() for item in self.per_rank], + "provenance": self.provenance, + } + + +@dataclass(frozen=True) +class AttentionBackwardComparisonReport: + """Structured PR8 report for CP attention gradient drift validation.""" + + reference_name: str + drifts: tuple[AttentionBackwardPathDrift, ...] + + def to_dict(self) -> dict[str, object]: + return { + "reference_name": self.reference_name, + "drifts": [drift.to_dict() for drift in self.drifts], + } + + +def merge_attention_partial_states( + states: Sequence[AttentionPartialState], +) -> AttentionPartialState: + """Merge CP/chunk partial states in logical block order. + + The merge is the online-softmax/LSE merge used by attention, not a plain + sum. The input order is deliberately ignored: states are sorted by logical + ``block_start`` so the result depends on global block indices rather than + arrival order. + """ + + if not states: + raise ValueError("at least one attention partial state is required") + + ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) + _validate_merge_shapes_and_ranges(ordered) + + merged = ordered[0] + merged_out = merged.out.float() + merged_lse = merged.lse.float() + for state in ordered[1:]: + merged_out, merged_lse = _merge_two_states( + merged_out, + merged_lse, + state.out.float(), + state.lse.float(), + ) + + return AttentionPartialState( + out=merged_out, + lse=merged_lse, + block_start=ordered[0].block_start, + block_end=ordered[-1].block_end, + ) + + +class DeterministicCPAttentionReferenceOp: + """Correctness-first CP attention reference for prefill and chunked prefill. + + The reference consumes attention-ready Q/K. For Qwen3 WS2 this means Q/K + have already passed QK-Norm and RoPE unless an outer contract explicitly + marks them as pre-RoPE. RoPE is intentionally kept outside this CP merge + implementation so fused and unfused ``RoPE+Attention`` paths can compare the + same post-RoPE Q/K boundary before validating CP communication. + + The op emulates CP by splitting query and KV sequence dimensions into + logical CP shards. Each query shard computes one partial attention state per + KV block, then merges those states in fixed global-block order using fp32 + LSE arithmetic. ``forward`` returns the input dtype after the final write; + ``forward_fp32`` keeps the fp32 merged output. + """ + + op_class = "attention" + + @staticmethod + def split_kv_execution_plans( + total_kv_tokens: int, + *, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> list[dict[str, object]]: + """Export the actual logical Split-KV plan before execution.""" + + return split_kv_execution_plan_provenance( + total_kv_tokens, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ) + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + return self.forward( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and final input-dtype write.""" + + out, _ = self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=q.dtype, + ) + return out + + def forward_fp32( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and fp32 output.""" + + out, _ = self.forward_fp32_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(out, lse)`` for the CP reference path. + + ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 + until the final write, then downcast to ``output_dtype``. When omitted, + ``output_dtype`` defaults to the input dtype. + """ + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + out = out.to(resolved_output_dtype) + return out, lse + + def forward_fp32_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return fp32 ``(out, lse)`` for the CP reference path.""" + + return self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=torch.float32, + ) + + def backward_reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, + name: Optional[str] = None, + ) -> AttentionBackwardPathResult: + """Run the deterministic training-side backward validation path. + + The semantic backward input is ``dout`` plus the forward attention state + produced from the same Q/K/V, masks, position offsets, CP world, and KV + block order. The reference keeps the softmax/merge math in fp32 and + records the final-write dtype in provenance; decode backward is + intentionally out of scope for PR8. + """ + + _validate_qkv(q, k, v) + if dout.shape != q.shape: + raise ValueError("dout must have shape [B, Hq, Sq, D], matching q") + if not torch.is_floating_point(dout) or torch.is_complex(dout): + raise ValueError("dout must be a real floating-point tensor") + if dout.device != q.device: + raise ValueError("dout must be on the same device as q, k, and v") + if dout.dtype != q.dtype: + raise ValueError("dout must have the same dtype as q") + q_leaf = q.detach().clone().requires_grad_(True) + k_leaf = k.detach().clone().requires_grad_(True) + v_leaf = v.detach().clone().requires_grad_(True) + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) + out, lse = self.forward_with_lse( + q_leaf, + k_leaf, + v_leaf, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=resolved_output_dtype, + ) + torch.autograd.backward(out, dout.to(dtype=out.dtype)) + if q_leaf.grad is None or k_leaf.grad is None or v_leaf.grad is None: + raise RuntimeError("CP attention backward did not produce dq/dk/dv") + + return AttentionBackwardPathResult( + name=name + or _backward_path_name(cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size), + out=out.detach(), + lse=lse.detach(), + gradients=AttentionBackwardGradients( + dq=q_leaf.grad.detach(), + dk=k_leaf.grad.detach(), + dv=v_leaf.grad.detach(), + ), + provenance={ + "attention_mode": "prefill" if kv_chunk_size is None else "chunked_prefill", + "gradient_mode": "training_backward", + "gradient_inputs": ["q", "k", "v"], + "gradient_outputs": ["out"], + "saved_forward_state": [ + "out", + "attention_lse", + "causal_mask", + "key_padding_mask", + "query_position_offsets", + "key_position_offsets", + "global_block_index", + ], + "cp_world_size": cp_world_size, + "kv_chunk_size": kv_chunk_size, + "requested_split_kv_policy": ("disabled" if kv_chunk_size is None else "fixed"), + "requested_split_kv_size": kv_chunk_size, + "actual_split_kv_plans": split_kv_execution_plan_provenance( + k.size(2), + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_backward_reference", + ), + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "output_dtype": str(resolved_output_dtype).replace("torch.", ""), + "q_dtype": str(q.dtype).replace("torch.", ""), + "k_dtype": str(k.dtype).replace("torch.", ""), + "v_dtype": str(v.dtype).replace("torch.", ""), + "dout_dtype": str(dout.dtype).replace("torch.", ""), + "te_backward_oracle": "not_used", + "decode_backward": "not_supported", + }, + ) + + def local_partial_state( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + q_start: int, + k_start: int, + total_kv_len: int, + total_query_len: Optional[int] = None, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + ) -> AttentionPartialState: + """Compute one query shard against one logical KV block. + + ``query_position_offsets`` and ``key_position_offsets`` are optional + per-batch-row base positions. They let the reference express varlen or + packed metadata while retaining the dense [B, H, S, D] tensor layout. + For post-RoPE Q/K, these offsets must describe the same absolute token + positions used when RoPE was applied. + """ + + _validate_qkv(q, k, v) + _validate_scale(scale) + if q_start < 0 or k_start < 0: + raise ValueError("q_start and k_start must be non-negative") + if total_kv_len < k_start + k.size(2): + raise ValueError("total_kv_len must cover the local KV block") + if total_query_len is None: + total_query_len = q.size(2) + if total_query_len < q_start + q.size(2): + raise ValueError("total_query_len must cover the local query block") + if key_padding_mask is not None: + if key_padding_mask.shape != (q.size(0), k.size(2)): + raise ValueError("local key_padding_mask must have shape [B, local_skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("local key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + q.size(0), + q.device, + default=total_kv_len - total_query_len, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + q.size(0), + q.device, + default=0, + name="key_position_offsets", + ) + + ctx = NativeAttentionOp._strict_fp32_math(q.device.type) + with ctx: + qf = q.float() + kf = k.float() + vf = v.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hkv != hq: + repeat = hq // hkv + kf = kf.repeat_interleave(repeat, dim=1) + vf = vf.repeat_interleave(repeat, dim=1) + + if skv == 0: + zero_dep = _zero_dependency(qf, kf, vf) + return AttentionPartialState( + out=torch.zeros(q.size(0), hq, sq, dim, device=q.device, dtype=torch.float32) + + zero_dep, + lse=torch.full( + (q.size(0), hq, sq), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep, + block_start=k_start, + block_end=k_start, + ) + + scale_value = scale if scale is not None else (1.0 / math.sqrt(dim)) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * scale_value + if causal: + query_base = query_offsets[:, None] + q_start + key_base = key_offsets[:, None] + k_start + q_pos = torch.arange(sq, device=q.device, dtype=torch.long) + query_base + k_pos = torch.arange(skv, device=q.device, dtype=torch.long) + key_base + causal_mask = k_pos[:, None, :] > q_pos[:, :, None] + scores = scores.masked_fill(causal_mask[:, None, :, :], float("-inf")) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + + lse = torch.logsumexp(scores, dim=-1) + finite_lse = torch.isfinite(lse) + weights = torch.exp(scores - lse.unsqueeze(-1)) + weights = torch.where(finite_lse.unsqueeze(-1), weights, torch.zeros_like(weights)) + out = torch.matmul(weights, vf) + return AttentionPartialState( + out=out, + lse=lse, + block_start=k_start, + block_end=k_start + skv, + ) + + def _forward_impl( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + cp_world_size: int, + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + _validate_qkv(q, k, v) + _validate_scale(scale) + if ( + isinstance(cp_world_size, bool) + or not isinstance(cp_world_size, int) + or cp_world_size < 1 + ): + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and ( + isinstance(kv_chunk_size, bool) + or not isinstance(kv_chunk_size, int) + or kv_chunk_size < 1 + ): + raise ValueError("kv_chunk_size must be >= 1 when provided") + + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + + q_bounds = _split_bounds(sq, cp_world_size) + kv_bounds = _kv_block_bounds(skv, cp_world_size, kv_chunk_size) + out_chunks: list[torch.Tensor] = [] + lse_chunks: list[torch.Tensor] = [] + for q_start, q_end in q_bounds: + if q_start == q_end: + continue + q_block = q[:, :, q_start:q_end, :] + states = [ + self.local_partial_state( + q_block, + k[:, :, k_start:k_end, :], + v[:, :, k_start:k_end, :], + q_start=q_start, + k_start=k_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None if key_padding_mask is None else key_padding_mask[:, k_start:k_end] + ), + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + ) + for k_start, k_end in kv_bounds + if k_start != k_end + ] + if states: + merged = merge_attention_partial_states(states) + out_chunks.append(merged.out) + lse_chunks.append(merged.lse) + else: + zero_dep = _zero_dependency(q_block.float(), k.float(), v.float()) + out_chunks.append( + torch.zeros(batch, hq, q_end - q_start, dim, device=q.device) + zero_dep + ) + lse_chunks.append( + torch.full( + (batch, hq, q_end - q_start), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep + ) + + if not out_chunks: + zero_dep = _zero_dependency(q.float(), k.float(), v.float()) + return ( + torch.empty(batch, hq, 0, dim, device=q.device, dtype=torch.float32) + zero_dep, + torch.empty(batch, hq, 0, device=q.device, dtype=torch.float32) + zero_dep, + ) + return torch.cat(out_chunks, dim=2), torch.cat(lse_chunks, dim=2) + + +def compare_cp_attention_backward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + candidate_cp_world_size: int = 2, + candidate_kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, +) -> AttentionBackwardComparisonReport: + """Compare CP=1 backward with a CP/chunked-prefill candidate. + + The report includes whole-tensor ``dq/dk/dv`` drift and per-logical-CP-rank + slices. It is a validation/reporting helper, not a separate production + backward kernel. + """ + + op = DeterministicCPAttentionReferenceOp() + reference = op.backward_reference( + q, + k, + v, + dout, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=1, + kv_chunk_size=None, + output_dtype=output_dtype, + name="cp1_backward_reference", + ) + candidate = op.backward_reference( + q, + k, + v, + dout, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=candidate_cp_world_size, + kv_chunk_size=candidate_kv_chunk_size, + output_dtype=output_dtype, + ) + return AttentionBackwardComparisonReport( + reference_name=reference.name, + drifts=(_compare_backward_path(candidate, reference),), + ) + + +def _compare_backward_path( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, +) -> AttentionBackwardPathDrift: + cp_world_size = _provenance_int(candidate.provenance, "cp_world_size") + return AttentionBackwardPathDrift( + candidate_name=candidate.name, + dq=_drift_stats(candidate.gradients.dq, reference.gradients.dq), + dk=_drift_stats(candidate.gradients.dk, reference.gradients.dk), + dv=_drift_stats(candidate.gradients.dv, reference.gradients.dv), + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + per_rank=_per_rank_backward_drifts(candidate, reference, cp_world_size), + provenance=candidate.provenance, + ) + + +def _per_rank_backward_drifts( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, + cp_world_size: int, +) -> tuple[AttentionBackwardRankDrift, ...]: + q_bounds = _split_bounds(candidate.gradients.dq.size(2), cp_world_size) + kv_bounds = _split_bounds(candidate.gradients.dk.size(2), cp_world_size) + per_rank = [] + for rank, ((q_start, q_end), (kv_start, kv_end)) in enumerate(zip(q_bounds, kv_bounds)): + per_rank.append( + AttentionBackwardRankDrift( + rank=rank, + dq=_drift_stats( + candidate.gradients.dq[:, :, q_start:q_end, :], + reference.gradients.dq[:, :, q_start:q_end, :], + ), + dk=_drift_stats( + candidate.gradients.dk[:, :, kv_start:kv_end, :], + reference.gradients.dk[:, :, kv_start:kv_end, :], + ), + dv=_drift_stats( + candidate.gradients.dv[:, :, kv_start:kv_end, :], + reference.gradients.dv[:, :, kv_start:kv_end, :], + ), + ) + ) + return tuple(per_rank) + + +def _drift_stats(candidate: torch.Tensor, reference: torch.Tensor) -> GradientDriftStats: + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match " + f"reference shape {tuple(reference.shape)}" + ) + diff = (candidate.float() - reference.float()).abs().reshape(-1) + active_count = int(diff.numel()) + if active_count == 0: + return GradientDriftStats(0.0, 0.0, 0.0, 0.0, 0) + return GradientDriftStats( + max_abs=float(diff.max().item()), + mean_abs=float(diff.mean().item()), + p95_abs=float(torch.quantile(diff, 0.95).item()), + p99_abs=float(torch.quantile(diff, 0.99).item()), + active_count=active_count, + ) + + +def _backward_path_name(*, cp_world_size: int, kv_chunk_size: Optional[int]) -> str: + prefix = f"cp{cp_world_size}" + if kv_chunk_size is None: + return f"{prefix}_backward" + return f"{prefix}_chunked_backward" + + +def _provenance_int(provenance: dict[str, object], key: str) -> int: + value = provenance[key] + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"provenance field {key!r} must be an int") + return value + + +def _merge_two_states( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + merged_lse = torch.logaddexp(lse_a, lse_b) + finite = torch.isfinite(merged_lse) + weight_a = torch.where(finite, torch.exp(lse_a - merged_lse), torch.zeros_like(merged_lse)) + weight_b = torch.where(finite, torch.exp(lse_b - merged_lse), torch.zeros_like(merged_lse)) + merged_out = weight_a.unsqueeze(-1) * out_a + weight_b.unsqueeze(-1) * out_b + return merged_out, merged_lse + + +def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) -> None: + first = states[0] + previous_end = first.block_end + for state in states[1:]: + if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: + raise ValueError("all partial states must have matching out/lse shapes") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be gap-free and non-overlapping") + previous_end = state.block_end + + +def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q, k, and v must have shape [B, H, S, D]") + if k.shape != v.shape: + raise ValueError("k and v must have the same shape") + if q.size(0) != k.size(0) or q.size(3) != k.size(3): + raise ValueError("q, k, and v must share batch size and head dim") + if q.size(1) < 1 or k.size(1) < 1 or q.size(3) < 1: + raise ValueError("q, k, and v must have positive head counts and head dim") + if not all(torch.is_floating_point(tensor) for tensor in (q, k, v)) or any( + torch.is_complex(tensor) for tensor in (q, k, v) + ): + raise ValueError("q, k, and v must be real floating-point tensors") + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError("q, k, and v must have the same dtype") + if q.device != k.device or q.device != v.device: + raise ValueError("q, k, and v must be on the same device") + if q.size(1) % k.size(1) != 0: + raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") + + +def _validate_scale(scale: Optional[float]) -> None: + if scale is None: + return + if isinstance(scale, bool) or not isinstance(scale, (int, float)): + raise ValueError("scale must be a positive finite number") + if not math.isfinite(float(scale)) or float(scale) <= 0: + raise ValueError("scale must be a positive finite number") + + +def _validate_output_dtype(output_dtype: torch.dtype) -> None: + if not isinstance(output_dtype, torch.dtype): + raise ValueError("output_dtype must be a real floating-point torch dtype") + probe = torch.empty((), dtype=output_dtype) + if not torch.is_floating_point(probe) or torch.is_complex(probe): + raise ValueError("output_dtype must be a real floating-point torch dtype") + + +def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: + total = torch.tensor(0.0, device=tensors[0].device) + for tensor in tensors: + total = total + tensor.sum() + return total * 0.0 + + +def _normalize_position_offsets( + offsets: Optional[torch.Tensor], + batch: int, + device: torch.device, + *, + default: int, + name: str, +) -> torch.Tensor: + if offsets is None: + return torch.full((batch,), default, dtype=torch.long, device=device) + if offsets.ndim != 1 or offsets.numel() != batch: + raise ValueError(f"{name} must have shape [B]") + if torch.is_floating_point(offsets) or torch.is_complex(offsets) or offsets.dtype == torch.bool: + raise ValueError(f"{name} must contain integer positions") + return offsets.to(device=device, dtype=torch.long) + + +def _split_bounds(length: int, parts: int) -> list[tuple[int, int]]: + base, extra = divmod(length, parts) + bounds: list[tuple[int, int]] = [] + start = 0 + for index in range(parts): + width = base + (1 if index < extra else 0) + end = start + width + bounds.append((start, end)) + start = end + return bounds + + +def _kv_block_bounds( + length: int, + cp_world_size: int, + kv_chunk_size: Optional[int], +) -> list[tuple[int, int]]: + bounds: list[tuple[int, int]] = [] + for start, end in _split_bounds(length, cp_world_size): + if kv_chunk_size is None: + bounds.append((start, end)) + continue + cursor = start + while cursor < end: + chunk_end = min(cursor + kv_chunk_size, end) + bounds.append((cursor, chunk_end)) + cursor = chunk_end + return bounds + + +def split_kv_execution_plan_provenance( + length: int, + *, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str, +) -> list[dict[str, object]]: + """Return the actual backend-local Split-KV plan for every CP owner.""" + + if length < 1: + raise ValueError("Split-KV sequence length must be >= 1") + if cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + result: list[dict[str, object]] = [] + for owner_cp_rank, (rank_start, rank_end) in enumerate(_split_bounds(length, cp_world_size)): + if rank_start == rank_end: + continue + boundaries: tuple[tuple[int, int], ...] + if kv_chunk_size is None: + boundaries = ((rank_start, rank_end),) + mode = SplitKVMode.DISABLED + else: + boundaries = tuple( + (start, min(start + kv_chunk_size, rank_end)) + for start in range(rank_start, rank_end, kv_chunk_size) + ) + mode = SplitKVMode.FIXED + plan = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + result.append({"owner_cp_rank": owner_cp_rank, **plan.to_dict()}) + return result + + +def build_reference_split_kv_runtime_plan_set( + total_kv_tokens: Sequence[int], + *, + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str = "deterministic_cp_reference", +) -> SplitKVRuntimePlanSet: + """Build complete per-batch/TP/CP/owner plans for the reference path.""" + + totals = tuple(total_kv_tokens) + if not totals or any(total < cp_world_size for total in totals): + raise ValueError("reference runtime plan sets require at least one KV token per CP owner") + if tp_world_size < 1 or cp_world_size < 1: + raise ValueError("TP and CP world sizes must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + + entries: list[SplitKVRuntimePlanEntry] = [] + for batch_index, total in enumerate(totals): + owner_ranges = _split_bounds(total, cp_world_size) + for tp_rank in range(tp_world_size): + for cp_rank in range(cp_world_size): + for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + boundaries: tuple[tuple[int, int], ...] + if kv_chunk_size is None: + mode = SplitKVMode.DISABLED + boundaries = ((owner_start, owner_end),) + else: + mode = SplitKVMode.FIXED + boundaries = tuple( + (start, min(start + kv_chunk_size, owner_end)) + for start in range(owner_start, owner_end, kv_chunk_size) + ) + execution = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + entries.append( + SplitKVRuntimePlanEntry( + coordinate=SplitKVRuntimeCoordinate( + batch_index=batch_index, + tp_rank=tp_rank, + cp_rank=cp_rank, + owner_cp_rank=owner_cp_rank, + ), + expected_kv_range=(owner_start, owner_end), + execution=execution, + ) + ) + return SplitKVRuntimePlanSet( + batch_size=len(totals), + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals, + entries=tuple(entries), + ) + + +CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp + +__all__ = [ + "AttentionBackwardComparisonReport", + "AttentionBackwardGradients", + "AttentionBackwardPathDrift", + "AttentionBackwardPathResult", + "AttentionBackwardRankDrift", + "AttentionPartialState", + "build_reference_split_kv_runtime_plan_set", + "CPAttentionReferenceOp", + "DeterministicCPAttentionReferenceOp", + "GradientDriftStats", + "compare_cp_attention_backward", + "merge_attention_partial_states", + "split_kv_execution_plan_provenance", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index efde5c25..060234be 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -1,6 +1,8 @@ # 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 @@ -8,6 +10,21 @@ import torch +from rl_engine.kernels.attention_contract import ( + AttentionBackendCapability, + AttentionContract, + AttentionContractError, + AttentionDispatchResult, + AttentionDType, + AttentionMode, + AttentionRole, +) +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 @@ -18,9 +35,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): @@ -76,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 @@ -103,6 +123,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 @@ -111,6 +137,57 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" +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, *, @@ -159,14 +236,66 @@ 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()) + + # 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", + ), + 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 = { "cuda": { @@ -201,6 +330,12 @@ def __init__(self): OpBackend.CUDA_DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION, ], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], + "ws2_attention": [ + OpBackend.PYTORCH_CP_ATTENTION, + OpBackend.CUDA_DETERMINISTIC_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "linear_logp": [ @@ -214,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": [ @@ -249,6 +387,11 @@ def __init__(self): OpBackend.TRITON_GENERIC, ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], + "ws2_attention": [ + OpBackend.PYTORCH_CP_ATTENTION, + OpBackend.PYTORCH_NATIVE_ATTENTION, + ], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], @@ -273,6 +416,11 @@ 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], @@ -317,7 +465,8 @@ def _adjust_priority_from_env(self): ) 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: @@ -339,8 +488,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"] @@ -358,7 +505,6 @@ def _adjust_priority_for_hardware(self): 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"] @@ -370,8 +516,8 @@ def _adjust_priority_for_hardware(self): 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}") + except Exception as exc: + logger.warning(f"Failed to probe device capability: {exc}") def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: """Core distribution logic: Automatically select the best operator @@ -383,7 +529,6 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: for backend in candidates: if backend.name in self._instance_cache: return self._instance_cache[backend.name] - if backend.name in self._failed_backends: continue @@ -393,8 +538,8 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> 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) @@ -416,24 +561,144 @@ def _platform_for_device(self, device: torch.device | str | None) -> str: return resolved.type return "cpu" - def _load_backend(self, backend: OpBackend) -> Optional[Type]: - """Dynamic loading technique: Import modules only when needed - and check environment dependencies. + 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() + candidates = self._priority_map.get(platform, {}).get("ws2_attention", []) + rejected: list[str] = [] + + for backend in candidates: + capability = self._attention_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no AttentionBackendCapability declared") + continue + incompatibilities = list(capability.incompatibilities(contract)) + policy_mismatch = self._attention_policy_mismatch(requested_backend, capability) + if policy_mismatch is not None: + incompatibilities.append(policy_mismatch) + if incompatibilities: + rejected.append(f"{backend.name}: " + "; ".join(incompatibilities)) + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + continue + + 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]: + """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() diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py new file mode 100644 index 00000000..350545a1 --- /dev/null +++ b/tests/test_attention_contract.py @@ -0,0 +1,683 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 Attention CP contract and contract-aware dispatch tests (issue #235).""" + +from __future__ import annotations + +import json +from dataclasses import replace + +import pytest + +from rl_engine.kernels.attention_contract import ( + AttentionBackendCapability, + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMode, + AttentionRole, + KVCacheSpec, + ReductionSpec, + RoPEFusionBoundary, + RoPESpec, + ShardingSpec, + SplitKVExecutionPlan, + SplitKVRuntimePlanSet, + SplitKVSpec, + build_split_kv_runtime_plan_set, + validate_split_kv_alignment, + validate_split_kv_plan_set_alignment, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + + +def _sharding( + *, + tp_rank: int = 0, + tp_world_size: int = 2, + cp_rank: int = 0, + cp_world_size: int = 2, + global_sequence_length: int = 4096, + local_sequence_length: int = 2048, + global_block_indices: tuple[int, ...] = (0,), + global_block_token_starts: tuple[int, ...] = (0,), + local_block_offsets: tuple[int, ...] = (0, 2048), + packed_sequence_offsets: tuple[int, ...] | None = None, +) -> ShardingSpec: + local_q_heads = 32 // tp_world_size + local_kv_heads = 8 // tp_world_size + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=global_sequence_length, + local_sequence_length=local_sequence_length, + global_block_indices=global_block_indices, + global_block_token_starts=global_block_token_starts, + local_block_offsets=local_block_offsets, + packed_sequence_offsets=packed_sequence_offsets, + ) + + +def _contract( + *, + role: str = "infer", + mode: str = "prefill", + sharding: ShardingSpec | None = None, + kv_cache: KVCacheSpec | None = None, + causal_offsets: tuple[int, ...] = (0,), + batch_size: int = 1, + query_sequence_length: int | None = None, + rope: RoPESpec | None = None, +) -> AttentionContract: + resolved_sharding = sharding or _sharding() + return AttentionContract( + role=role, + mode=mode, + dtype="bf16", + batch_size=batch_size, + query_sequence_length=( + query_sequence_length + if query_sequence_length is not None + else (1 if mode == "decode" else resolved_sharding.local_sequence_length) + ), + head_dim=128, + causal=True, + causal_offsets=causal_offsets, + sharding=resolved_sharding, + reduction=ReductionSpec(), + kv_cache=kv_cache, + rope=rope, + ) + + +def _declared_cp_backend() -> AttentionBackendCapability: + return AttentionBackendCapability( + backend_id="test-deterministic-cp-attention", + roles=frozenset({AttentionRole.TRAIN, AttentionRole.INFER}), + modes=frozenset( + {AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL, AttentionMode.DECODE} + ), + dtypes=frozenset({AttentionDType.BF16}), + tp_world_sizes=(2,), + cp_world_sizes=(1, 2), + exports_attention_lse=True, + deterministic_cp_merge=True, + supports_packed_varlen=True, + supports_kv_cache=True, + supports_split_kv_fixed=True, + reports_actual_split_kv_plan=True, + implementation_kind="deterministic", + ) + + +def test_qwen3_tp2_cp2_contract_is_representable_and_serializable(): + contract = _contract() + + assert contract.sharding.local_q_heads == 16 + assert contract.sharding.local_kv_heads == 4 + assert contract.sharding.tp_world_size == 2 + assert contract.sharding.cp_world_size == 2 + assert contract.reduction.acc_dtype is AttentionDType.FP32 + assert contract.to_dict()["reduction"] == { + "merge": "online_softmax_lse", + "acc_dtype": "fp32", + "order": "global_block_index", + "downcast_at": "final_write", + "engine": "in_op_reference", + } + json.dumps(contract.to_dict()) + + +def test_rope_metadata_is_part_of_attention_contract_provenance(): + rope = RoPESpec( + q_state="post_rope", + k_state="post_rope", + k_cache_state="post_rope", + theta=1.0e6, + rotary_dim=128, + query_position_offsets=(0,), + key_position_offsets=(0,), + cast_at="after_rope", + output_dtype="bf16", + fusion_boundary="unfused_rope_attention", + ) + + contract = _contract(rope=rope) + payload = contract.to_dict() + + assert payload["rope"] == { + "q_state": "post_rope", + "k_state": "post_rope", + "k_cache_state": "post_rope", + "theta": 1.0e6, + "rotary_dim": 128, + "rope_scaling": None, + "position_ids": None, + "query_position_offsets": [0], + "key_position_offsets": [0], + "cast_at": "after_rope", + "output_dtype": "bf16", + "fusion_boundary": "unfused_rope_attention", + } + json.dumps(payload) + + +def test_rope_position_metadata_is_validated_against_contract_shape(): + with pytest.raises(AttentionContractError, match="rotary_dim=256"): + _contract(rope=RoPESpec(rotary_dim=256)) + + with pytest.raises(AttentionContractError, match="query_position_offsets"): + _contract(batch_size=2, causal_offsets=(0, 0), rope=RoPESpec(query_position_offsets=(0,))) + + with pytest.raises(AttentionContractError, match="position_ids"): + _contract(rope=RoPESpec(position_ids=(0, 1, 2))) + + valid = _contract(rope=RoPESpec(position_ids=tuple(range(2048)))) + assert valid.rope is not None + assert valid.rope.position_ids == tuple(range(2048)) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("tp_rank", 2, "tp_rank=2"), + ("cp_rank", 2, "cp_rank=2"), + ("global_block_indices", (), "must not be empty"), + ("global_block_indices", (1, 0), "strictly increasing"), + ], +) +def test_invalid_rank_and_cp_order_metadata_fail_loudly(field, value, message): + values = { + "tp_rank": 0, + "tp_world_size": 2, + "cp_rank": 0, + "cp_world_size": 2, + "global_q_heads": 32, + "global_kv_heads": 8, + "local_q_head_start": 0, + "local_q_heads": 16, + "local_kv_head_start": 0, + "local_kv_heads": 4, + "global_sequence_length": 4096, + "local_sequence_length": 2048, + "global_block_indices": (0,), + "global_block_token_starts": (0,), + "local_block_offsets": (0, 2048), + } + values[field] = value + + with pytest.raises(AttentionContractError, match=message): + ShardingSpec(**values) + + +def test_tp_local_heads_must_preserve_global_gqa_mapping(): + with pytest.raises(AttentionContractError, match="local TP head counts"): + replace(_sharding(), local_q_heads=7) + + with pytest.raises(AttentionContractError, match="head starts"): + replace(_sharding(tp_rank=1), local_q_head_start=0) + + +def test_sequence_range_and_packed_offsets_are_validated(): + with pytest.raises(AttentionContractError, match="exceeds global_sequence_length"): + _sharding(global_block_token_starts=(4000,)) + + with pytest.raises(AttentionContractError, match="final packed_sequence_offsets"): + _sharding(packed_sequence_offsets=(0, 512)) + + sharding = _sharding(packed_sequence_offsets=(0, 512, 2048)) + assert sharding.packed_sequence_offsets == (0, 512, 2048) + + +def test_non_contiguous_cp_blocks_have_explicit_global_and_local_offsets(): + sharding = _sharding( + global_block_indices=(0, 3), + global_block_token_starts=(0, 3072), + local_block_offsets=(0, 1024, 2048), + ) + + assert sharding.global_block_indices == (0, 3) + assert sharding.global_block_token_starts == (0, 3072) + assert sharding.local_block_offsets == (0, 1024, 2048) + + with pytest.raises(AttentionContractError, match="non-overlapping and ordered"): + _sharding( + global_block_indices=(0, 1), + global_block_token_starts=(0, 512), + local_block_offsets=(0, 1024, 2048), + ) + + +def test_reduction_requires_fp32_accumulation(): + with pytest.raises(AttentionContractError, match="must be fp32"): + ReductionSpec(acc_dtype="bf16") + + +def test_split_kv_policy_is_a_first_class_strict_contract(): + contract = _contract() + assert contract.to_dict()["split_kv"] == { + "mode": "disabled", + "fixed_split_size": None, + "strict_consistency": True, + } + + fixed = replace(contract, split_kv=SplitKVSpec.fixed(128)) + assert fixed.to_dict()["split_kv"]["mode"] == "fixed" + assert fixed.to_dict()["split_kv"]["fixed_split_size"] == 128 + + with pytest.raises(AttentionContractError, match="auto Split-KV"): + SplitKVSpec.auto(strict_consistency=True) + + +def test_split_kv_execution_plan_records_actual_logical_schedule(): + plan = SplitKVSpec.fixed(4).resolve(10, backend="training-reference") + + assert plan.actual_split_count == 3 + assert plan.to_dict()["actual_split_boundaries"] == [[0, 4], [4, 8], [8, 10]] + assert plan.to_dict()["split_kv_merge_order"] == "global_block_index" + assert plan.to_dict()["split_kv_accum_dtype"] == "fp32" + assert plan.to_dict()["split_kv_downcast_at"] == "final_write" + + +def test_strict_split_kv_alignment_rejects_unknown_or_mismatched_actual_plan(): + training = SplitKVSpec.fixed(4).resolve(10, backend="training") + rollout = SplitKVSpec.fixed(4).resolve(10, backend="rollout") + validate_split_kv_alignment(training, rollout) + + unknown = SplitKVSpec.auto().resolve(10, backend="rollout") + with pytest.raises(AttentionContractError, match="actual runtime plans"): + validate_split_kv_alignment(training, unknown) + + mismatched = SplitKVSpec.fixed(5).resolve(10, backend="rollout") + with pytest.raises(AttentionContractError, match="differ"): + validate_split_kv_alignment(training, mismatched) + + with pytest.raises(AttentionContractError, match="contiguous"): + SplitKVExecutionPlan( + requested_mode="fixed", + requested_split_size=4, + actual_mode="fixed", + actual_split_size=4, + boundaries=((0, 4), (5, 10)), + ) + + +def test_complete_split_kv_plan_set_covers_batch_tp_cp_and_owner_coordinates(): + plan_set = build_split_kv_runtime_plan_set( + (8, 10), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(2), + backend="training-reference", + ) + + assert len(plan_set.entries) == 16 + assert plan_set.to_dict()["coverage"] == ("complete_batch_tp_cp_owner_cartesian_product") + assert { + tuple(entry["expected_kv_range"]) + for entry in plan_set.to_dict()["entries"] + if entry["batch_index"] == 0 + } == {(0, 4), (4, 8)} + + +def test_split_kv_plan_set_alignment_rejects_missing_and_mismatched_rank_plans(): + training = build_split_kv_runtime_plan_set( + (8,), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(2), + backend="training", + ) + rollout = build_split_kv_runtime_plan_set( + (8,), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(2), + backend="rollout", + ) + validate_split_kv_plan_set_alignment(training, rollout) + + with pytest.raises(AttentionContractError, match="coordinate coverage is incomplete"): + SplitKVRuntimePlanSet( + batch_size=training.batch_size, + tp_world_size=training.tp_world_size, + cp_world_size=training.cp_world_size, + total_kv_tokens=training.total_kv_tokens, + entries=training.entries[:-1], + ) + + mismatched = build_split_kv_runtime_plan_set( + (8,), + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(1), + backend="rollout", + ) + with pytest.raises(AttentionContractError, match="plan differs"): + validate_split_kv_plan_set_alignment(training, mismatched) + + +def test_backend_must_support_policy_and_actual_plan_provenance(): + fixed = replace(_contract(), split_kv=SplitKVSpec.fixed(128)) + capability = replace( + _declared_cp_backend(), + supports_split_kv_fixed=False, + reports_actual_split_kv_plan=False, + ) + + assert capability.incompatibilities(fixed)[-2:] == ( + "Split-KV policy=fixed is unsupported", + "actual Split-KV execution-plan provenance is unsupported", + ) + + +def test_causal_attention_requires_explicit_offset(): + contract = _contract() + with pytest.raises(AttentionContractError, match="causal_offsets are required"): + replace(contract, causal_offsets=None) + + +def test_full_prefill_query_length_must_match_local_sequence_length(): + with pytest.raises(AttentionContractError, match="prefill query_sequence_length must equal"): + _contract(mode="prefill", query_sequence_length=1024) + + chunked = _contract(mode="chunked_prefill", query_sequence_length=512) + decode = _contract( + mode="decode", + query_sequence_length=1, + kv_cache=KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + ), + ) + assert chunked.query_sequence_length == 512 + assert decode.query_sequence_length == 1 + + +def test_decode_requires_complete_kv_cache_identity(): + with pytest.raises(AttentionContractError, match="kv_cache metadata is required"): + _contract(mode="decode") + + cache = KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1, -1),), + global_token_positions=tuple(range(17)), + page_size=16, + prefix_cache_enabled=True, + prefix_cache_key="prefix:sample-0", + ) + contract = _contract(mode="decode", kv_cache=cache) + assert contract.to_dict()["kv_cache"]["block_table"] == [[0, 1, -1]] + + +def test_prefix_cache_key_is_required_only_when_prefix_cache_is_enabled(): + with pytest.raises(AttentionContractError, match="prefix_cache_key is required"): + KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + prefix_cache_enabled=True, + ) + + +def test_cache_positions_must_match_kv_sequence_count(): + with pytest.raises(AttentionContractError, match="one entry per kv_seq_lens"): + KVCacheSpec( + cache_positions=(1,), + kv_seq_lens=(2, 2), + block_table=((0,), (1,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + ) + + +def test_cache_position_must_match_terminal_global_token_position(): + with pytest.raises(AttentionContractError, match="terminal global token position"): + KVCacheSpec( + cache_positions=(999,), + kv_seq_lens=(17,), + block_table=((0, 1),), + global_token_positions=tuple(range(17)), + page_size=16, + ) + + +@pytest.mark.parametrize("positions", [(7, 6), (7, 7)]) +def test_kv_cache_positions_must_be_strictly_increasing_per_sequence(positions): + with pytest.raises(AttentionContractError, match="strictly increasing"): + KVCacheSpec( + cache_positions=(7,), + kv_seq_lens=(2,), + block_table=((0,),), + global_token_positions=positions, + page_size=2, + ) + + +@pytest.mark.parametrize( + ("block_table", "message"), + [ + ((0, -1, 1), "padding must be trailing"), + ((0, 0, -1), "duplicate active page ids"), + ((0, -1, -1), "active page count"), + ], +) +def test_kv_cache_block_table_page_mapping_is_validated(block_table, message): + with pytest.raises(AttentionContractError, match=message): + KVCacheSpec( + cache_positions=(16,), + kv_seq_lens=(17,), + block_table=(block_table,), + global_token_positions=tuple(range(17)), + page_size=16, + ) + + +def test_prefix_pages_may_be_shared_across_sequences(): + cache = KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + assert cache.block_table == ((3,), (3,)) + assert cache.shared_prefix_page_count == 1 + + +def test_non_prefix_cache_rejects_cross_sequence_page_sharing(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=False, + ) + + +def test_prefix_cache_requires_explicit_shared_page_count(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(1, 1), + kv_seq_lens=(2, 2), + block_table=((3,), (3,)), + global_token_positions=(0, 1, 0, 1), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=0, + ) + + +def test_prefix_cache_rejects_shared_writable_suffix_pages(): + with pytest.raises(AttentionContractError, match="only when declared"): + KVCacheSpec( + cache_positions=(3, 3), + kv_seq_lens=(4, 4), + block_table=((3, 4), (3, 4)), + global_token_positions=(0, 1, 2, 3, 0, 1, 2, 3), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + +def test_shared_prefix_identity_must_match_pages_and_positions(): + with pytest.raises(AttentionContractError, match="page ids must match"): + KVCacheSpec( + cache_positions=(3, 3), + kv_seq_lens=(4, 4), + block_table=((3, 4), (5, 6)), + global_token_positions=(0, 1, 2, 3, 0, 1, 2, 3), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + with pytest.raises(AttentionContractError, match="token positions must match"): + KVCacheSpec( + cache_positions=(3, 13), + kv_seq_lens=(4, 4), + block_table=((3, 4), (3, 5)), + global_token_positions=(0, 1, 2, 3, 10, 11, 12, 13), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="shared-prefix", + shared_prefix_page_count=1, + ) + + +def test_shared_prefix_pages_must_be_fully_populated(): + with pytest.raises(AttentionContractError, match="fully populated and read-only"): + KVCacheSpec( + cache_positions=(0, 0), + kv_seq_lens=(1, 1), + block_table=((3,), (3,)), + global_token_positions=(0, 0), + page_size=2, + prefix_cache_enabled=True, + prefix_cache_key="partial-prefix-page", + shared_prefix_page_count=1, + ) + + +def test_current_ws1_backend_rejects_strict_cp_contract_without_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION] + + with pytest.raises(RuntimeError) as exc_info: + registry.get_attention_op(_contract()) + + message = str(exc_info.value) + assert "CP=2 is unsupported" in message + assert "attention-domain LSE export is unsupported" in message + assert "deterministic CP (out, lse) merge is unsupported" in message + + +def test_undeclared_backend_capability_is_never_selected(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_ATTN] + + with pytest.raises(RuntimeError, match="no AttentionBackendCapability declared"): + registry.get_attention_op(_contract()) + + +def test_declared_compatible_backend_resolves_and_records_provenance(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION] + registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend() + + result = registry.get_attention_op(_contract(), requested_backend="deterministic") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-cp-attention" + assert result.provenance["requested_backend"] == "deterministic" + assert result.provenance["actual_backend"] == "test-deterministic-cp-attention" + assert result.provenance["fallback"] is False + assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 + assert result.provenance["contract"]["sharding"]["cp_world_size"] == 2 + json.dumps(result.provenance) + + +def test_requested_stable_backend_id_is_enforced(): + registry = KernelRegistry() + platform = registry._platform() + registry._priority_map[platform]["ws2_attention"] = [OpBackend.PYTORCH_NATIVE_ATTENTION] + registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend() + + with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"): + registry.get_attention_op(_contract(), requested_backend="another-backend") + + result = registry.get_attention_op( + _contract(), requested_backend="test-deterministic-cp-attention" + ) + assert result.provenance["actual_backend"] == "test-deterministic-cp-attention" + + +def test_packed_layout_requires_declared_backend_support(): + capability = replace(_declared_cp_backend(), supports_packed_varlen=False) + contract = _contract( + sharding=_sharding(packed_sequence_offsets=(0, 512, 2048)), + causal_offsets=(0, 0), + batch_size=2, + ) + + assert capability.incompatibilities(contract) == ("packed varlen layout is unsupported",) + + +def test_rope_contract_requires_declared_backend_support(): + contract = _contract(rope=RoPESpec()) + capability = _declared_cp_backend() + + assert capability.incompatibilities(contract) == ("RoPE/position metadata is unsupported",) + + supported = replace(capability, supports_rope_metadata=True) + assert supported.incompatibilities(contract) == () + + +def test_fused_rope_attention_boundary_requires_declared_backend_support(): + contract = _contract(rope=RoPESpec(fusion_boundary=RoPEFusionBoundary.FUSED_ROPE_ATTENTION)) + capability = replace(_declared_cp_backend(), supports_rope_metadata=True) + + assert capability.incompatibilities(contract) == ( + "fused RoPE+Attention boundary is unsupported", + ) + + supported = replace(capability, supports_fused_rope_attention=True) + assert supported.incompatibilities(contract) == () + + +def test_packed_sequence_count_must_match_logical_batch_size(): + sharding = _sharding(packed_sequence_offsets=(0, 512, 2048)) + + with pytest.raises(AttentionContractError, match="must equal logical batch_size"): + _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=1) + + contract = _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=2) + assert contract.batch_size == 2 diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py new file mode 100644 index 00000000..12bb878e --- /dev/null +++ b/tests/test_attention_cross_config_binding.py @@ -0,0 +1,1059 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for #235 PR4: rollout/training attention contract binding. + +Every test here runs on CPU without Megatron or vLLM installed. That is the point: +the binding rules are contract logic, and contract logic that can only be exercised +on a 2-node x 2-GPU cluster would never be exercised. +""" + +from __future__ import annotations + +import json +from dataclasses import replace +from enum import Enum +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from rl_engine.alignment.cross_config.adapters import ( + QWEN3_8B, + WS2_ATTENTION_KNOBS, + AttentionRuntimeReadback, + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) +from rl_engine.alignment.cross_config.attention_binding import ( + ATTENTION_LSE_DOMAIN, + AttentionBindingError, + BindingErrorCode, + BindingTier, + bind_attention_contracts, + bind_attention_runtime_readbacks, + first_blocking_issue, + identity_fingerprint, + summarize_binding, +) +from rl_engine.alignment.cross_config.determinism import ( + compare_determinism, + megatron_probe_from_config, + vllm_probe_from_env, +) +from rl_engine.alignment.cross_config.schema import MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContractError, + AttentionMode, + AttentionRole, + KVCacheSpec, + SplitKVExecutionPlan, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, + SplitKVSpec, + build_split_kv_runtime_plan_set, +) +from rl_engine.kernels.attention_preprocess import ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS, +) + +pytestmark = pytest.mark.unit + + +TRAINING_KNOBS = { + "batch.size": 2, + "training.tensor_parallel_size": 2, + "training.context_parallel_size": 2, + "training.compute_dtype": "bf16", + "attention.split_kv_policy": 32, +} + +ROLLOUT_KNOBS = { + "batch.size": 2, + "rollout.tensor_parallel_size": 2, + "rollout.context_parallel_size": 2, + "rollout.dtype": "bf16", + "attention.split_kv_policy": 32, +} + + +def _identity(**overrides): + identity = { + "checkpoint_id": "qwen3-8b", + "model_version": "v1", + "weight_version": 7, + "tokenizer_fingerprint": "tokenizer-abc", + "token_ids_fingerprint": "tokens-abc", + "active_mask_fingerprint": "mask-abc", + "position_ids_fingerprint": "pos-abc", + "padding_side": "right", + "pre_update_state": "pre_update", + "batch_size": 2, + "global_token_positions_fingerprint": "gtp-abc", + "kv_seq_lens_fingerprint": "kvlen-abc", + } + identity.update(QWEN3_8B.identity_fields()) + identity.update(overrides) + return identity + + +def _contracts(): + training = MegatronAttentionMaterializer().build_contract(TRAINING_KNOBS) + rollout = VllmRolloutMaterializer().build_contract(ROLLOUT_KNOBS) + return rollout, training + + +def _plan_set(contract, *, backend): + return build_split_kv_runtime_plan_set( + (contract.sharding.global_sequence_length,) * contract.batch_size, + tp_world_size=contract.sharding.tp_world_size, + cp_world_size=contract.sharding.cp_world_size, + split_kv=contract.split_kv, + backend=backend, + ) + + +def _bind(rollout_identity=None, training_identity=None, **kwargs): + rollout, training = _contracts() + rollout = kwargs.pop("rollout_contract", rollout) + training = kwargs.pop("training_contract", training) + return bind_attention_contracts( + rollout_contract=rollout, + training_contract=training, + rollout_identity=rollout_identity if rollout_identity is not None else _identity(), + training_identity=training_identity if training_identity is not None else _identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + rollout_split_kv_plan_set=kwargs.pop( + "rollout_split_kv_plan_set", _plan_set(rollout, backend="vllm.readback") + ), + training_split_kv_plan_set=kwargs.pop( + "training_split_kv_plan_set", _plan_set(training, backend="megatron.readback") + ), + **kwargs, + ) + + +# -------------------------------------------------------------------------- +# tier 1: identity +# -------------------------------------------------------------------------- + + +def test_matching_identity_and_topology_bind_despite_different_materialization(): + """The core claim of PR4: same identity + same reduction, different runtimes.""" + + result = _bind() + + assert result.comparable + assert result.passed + assert result.issues == () + # Attention mode is a framework materialization difference, while both sides + # execute the same TP=2/CP=2 local ownership and Split-K schedule. + assert "mode" in result.recorded_differences + assert result.recorded_differences["mode"] == { + "rollout": "chunked_prefill", + "training": "prefill", + } + + +def test_weight_version_mismatch_is_not_comparable(): + result = _bind(rollout_identity=_identity(weight_version=6)) + + assert not result.comparable + assert not result.passed + codes = {issue.code for issue in result.issues} + assert BindingErrorCode.IDENTITY_MISMATCH in codes + blocking = first_blocking_issue(result) + assert blocking is not None and blocking.tier is BindingTier.IDENTICAL + assert "NOT COMPARABLE" in summarize_binding(result) + + +def test_rope_theta_mismatch_is_not_comparable(): + """RoPE math constants are identity, not materialization.""" + + result = _bind(training_identity=_identity(rope_theta=10000.0)) + + assert not result.comparable + assert any(issue.field == "rope_theta" for issue in result.issues) + + +def test_null_rope_scaling_is_a_value_not_an_omission(): + """Qwen3-8B applies no RoPE scaling; ``None`` must not read as undeclared.""" + + result = _bind() + + assert not result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + + +def test_missing_identity_field_is_reported_per_side(): + identity = _identity() + del identity["padding_side"] + result = _bind(rollout_identity=identity, training_identity=identity) + + missing = result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + assert {issue.field for issue in missing} == { + "rollout.padding_side", + "training.padding_side", + } + assert not result.comparable + + +def test_single_gpu_harness_may_waive_full_identity(): + """#235 PR2 has no KV-cache identity to declare; it opts out explicitly.""" + + identity = _identity() + del identity["global_token_positions_fingerprint"] + del identity["kv_seq_lens_fingerprint"] + + strict = _bind(rollout_identity=identity, training_identity=identity) + waived = _bind( + rollout_identity=identity, + training_identity=identity, + require_full_identity=False, + ) + + assert not strict.comparable + assert waived.comparable and waived.passed + + +def test_identity_fingerprint_ignores_undeclared_extra_keys(): + base = _identity() + decorated = dict(base, diagnostic_note="added later") + + assert identity_fingerprint(base) == identity_fingerprint(decorated) + + +# -------------------------------------------------------------------------- +# tier 2: reduction semantics +# -------------------------------------------------------------------------- + + +def test_reduction_semantics_are_bound_and_fingerprinted(): + result = _bind() + + reduction = result.provenance["training"]["contract"]["reduction"] + assert reduction["merge"] == "online_softmax_lse" + assert reduction["acc_dtype"] == "fp32" + assert reduction["order"] == "global_block_index" + assert reduction["downcast_at"] == "final_write" + assert result.reduction_fingerprint + + +def test_reduction_engine_difference_is_recorded_not_rejected(): + """A TE merge oracle on one side must not fail the binding.""" + + from rl_engine.alignment.cross_config.attention_binding import ( + RECORDED_FIELDS, + SEMANTIC_REDUCTION_FIELDS, + ) + + assert "reduction.engine" in RECORDED_FIELDS + assert "engine" not in SEMANTIC_REDUCTION_FIELDS + + +def test_lse_domain_is_recorded_as_attention_domain(): + """#235: attention exports attention-domain LSE, not vocab-logprob LSE.""" + + result = _bind() + + assert result.provenance["lse_domain"] == ATTENTION_LSE_DOMAIN == "attention" + + +def test_mixed_dtypes_fail_closed(): + """BF16 rollout against FP16 training produces an unattributable number.""" + + rollout = VllmRolloutMaterializer().build_contract( + {**ROLLOUT_KNOBS, "rollout.dtype": "float16"} + ) + result = _bind(rollout_contract=rollout) + + assert result.comparable # identity is fine + assert not result.passed + assert any(issue.field == "dtype" for issue in result.issues) + + +def test_precision_sweep_may_opt_into_mixed_dtypes(): + """#235 PR5 sweeps BF16 against an FP32 reference; it says so explicitly.""" + + training = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "float32"} + ) + result = _bind(training_contract=training, allow_dtype_difference=True) + + assert result.passed + assert result.provenance["dtype"] == "fp32" + + +def test_batch_size_mismatch_is_not_comparable(): + """Batch invariance is a claim about batch makeup, so it belongs to identity.""" + + result = _bind(rollout_identity=_identity(batch_size=4)) + + assert not result.comparable + assert any(issue.field == "batch_size" for issue in result.issues) + + +def test_missing_split_kv_runtime_evidence_fails_closed(): + result = _bind(rollout_split_kv_plan_set=None) + + assert result.comparable + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_RUNTIME_MISSING) + + +def test_split_kv_requested_policy_mismatch_fails_closed(): + rollout, training = _contracts() + rollout = replace(rollout, split_kv=SplitKVSpec.fixed(16)) + result = _bind(rollout_contract=rollout) + + assert result.comparable + assert not result.passed + assert any(issue.field == "split_kv" for issue in result.issues) + + +def test_split_kv_runtime_boundary_mismatch_fails_closed(): + rollout, training = _contracts() + mismatched_rollout = build_split_kv_runtime_plan_set( + (rollout.sharding.global_sequence_length,) * rollout.batch_size, + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(16), + backend="vllm.readback", + ) + result = _bind(rollout_split_kv_plan_set=mismatched_rollout) + + assert not result.passed + issues = result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + assert any(issue.field == "split_kv_runtime_plan_set" for issue in issues) + + +def test_split_kv_plan_set_must_match_its_own_contract_topology(): + rollout, _ = _contracts() + wrong_topology = build_split_kv_runtime_plan_set( + (rollout.sharding.global_sequence_length,) * rollout.batch_size, + tp_world_size=1, + cp_world_size=2, + split_kv=rollout.split_kv, + backend="vllm.readback", + ) + + result = _bind(rollout_split_kv_plan_set=wrong_topology) + + assert not result.passed + assert any( + issue.field == "rollout.split_kv_runtime_plan_set" + for issue in result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + ) + + +@pytest.mark.parametrize( + ("field_name", "corrupt_value"), + [ + ("merge_order", Enum("BadOrder", {"ARRIVAL": "arrival"}).ARRIVAL), + ("acc_dtype", Enum("BadDType", {"BF16": "bf16"}).BF16), + ("downcast_at", Enum("BadDowncast", {"PER_BLOCK": "per_block"}).PER_BLOCK), + ], +) +def test_split_kv_runtime_merge_semantic_corruption_fails_closed(field_name, corrupt_value): + rollout, _ = _contracts() + corrupted = _plan_set(rollout, backend="vllm.readback") + # Runtime reports are deserialized at this boundary. Simulate a corrupted + # report after construction to prove binding compares the actual fields. + object.__setattr__(corrupted.entries[0].execution, field_name, corrupt_value) + + result = _bind(rollout_split_kv_plan_set=corrupted) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + + +def test_split_kv_runtime_fallback_fails_closed(): + rollout, _ = _contracts() + plan_set = _plan_set(rollout, backend="vllm.readback") + fallback_entries = [] + for entry in plan_set.entries: + execution = entry.execution + fallback_entries.append( + SplitKVRuntimePlanEntry( + coordinate=entry.coordinate, + expected_kv_range=entry.expected_kv_range, + execution=SplitKVExecutionPlan( + requested_mode=execution.requested_mode, + requested_split_size=execution.requested_split_size, + actual_mode=execution.actual_mode, + actual_split_size=execution.actual_split_size, + boundaries=execution.boundaries, + merge_order=execution.merge_order, + acc_dtype=execution.acc_dtype, + downcast_at=execution.downcast_at, + backend=execution.backend, + source="runtime_fallback", + fallback=True, + fallback_reason="backend substituted a runtime plan", + ), + ) + ) + fallback = SplitKVRuntimePlanSet( + batch_size=plan_set.batch_size, + tp_world_size=plan_set.tp_world_size, + cp_world_size=plan_set.cp_world_size, + total_kv_tokens=plan_set.total_kv_tokens, + entries=tuple(fallback_entries), + ) + + result = _bind(rollout_split_kv_plan_set=fallback) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_FALLBACK) + + +@pytest.mark.parametrize( + "rollout_overrides", + [ + {"rollout.tensor_parallel_size": 1}, + {"rollout.context_parallel_size": 1}, + ], +) +def test_tp_or_cp_topology_mismatch_is_not_comparable(rollout_overrides): + rollout = VllmRolloutMaterializer().build_contract({**ROLLOUT_KNOBS, **rollout_overrides}) + result = _bind(rollout_contract=rollout) + + assert not result.comparable + assert result.issues_by_code(BindingErrorCode.TOPOLOGY_MISMATCH) + + +# -------------------------------------------------------------------------- +# role and input validation +# -------------------------------------------------------------------------- + + +def test_swapped_roles_are_rejected_outright(): + rollout, training = _contracts() + + with pytest.raises(AttentionBindingError): + bind_attention_contracts( + rollout_contract=training, + training_contract=rollout, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="a", + training_backend_id="b", + ) + + +# -------------------------------------------------------------------------- +# determinism cross-check +# -------------------------------------------------------------------------- + + +def _megatron_env(): + return {"NCCL_ALGO": "Tree", "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0"} + + +def _vllm_env(**overrides): + env = { + "VLLM_BATCH_INVARIANT": "1", + "NCCL_ALGO": "allreduce:tree", + "NCCL_PROTO": "Simple", + "NCCL_MIN_NCHANNELS": "1", + "NCCL_MAX_NCHANNELS": "1", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + } + env.update(overrides) + return env + + +def test_nccl_algo_mismatch_blocks_the_binding(): + """Megatron asserts NCCL_ALGO; vLLM hard-sets a different value.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert not report.compatible + fields = {issue.field for issue in report.issues} + assert "env.NCCL_ALGO" in fields + assert "env.NCCL_PROTO" in fields + + +def test_matching_nccl_settings_are_compatible(): + shared = {"NCCL_ALGO": "allreduce:tree", "NCCL_PROTO": "Simple"} + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=dict(shared) + ) + rollout = vllm_probe_from_env({**_vllm_env(**shared), "CUBLAS_WORKSPACE_CONFIG": None}) + + report = compare_determinism(rollout=rollout, training=training) + + assert report.compatible, [issue.to_dict() for issue in report.issues] + + +def test_determinism_switch_off_on_either_side_blocks(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=False), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env(VLLM_BATCH_INVARIANT="0")) + + report = compare_determinism(rollout=rollout, training=training) + + fields = {issue.field for issue in report.issues} + assert "training.deterministic_mode" in fields + assert "rollout.VLLM_BATCH_INVARIANT" in fields + + +def test_tf32_asymmetry_is_recorded_not_blocking(): + """Megatron does not manage TF32 at all; vLLM disables it. Record the gap.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert training.tf32_disabled is None + assert rollout.tf32_disabled is True + assert "tf32_disabled" in report.differences + assert not any(issue.field == "tf32_disabled" for issue in report.issues) + + +def test_determinism_issues_flow_into_the_binding(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + report = compare_determinism(rollout=rollout, training=training) + + result = _bind(determinism_issues=report.issues) + + assert result.comparable # identity is fine + assert not result.passed # but the reduction environment is not + assert result.issues_by_code(BindingErrorCode.DETERMINISM_INCOMPATIBLE) + assert "FAILED CLOSED" in summarize_binding(result) + + +# -------------------------------------------------------------------------- +# sharding derived from the frozen #239 layout +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_cp_shards_cover_the_global_sequence_without_overlap(cp_rank): + contract = MegatronAttentionMaterializer( + cp_rank=cp_rank, global_sequence_length=4096 + ).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert sharding.local_sequence_length == 2048 + assert sharding.global_block_indices == (cp_rank,) + assert sharding.global_block_token_starts == (cp_rank * 2048,) + # The causal offset must be the number of preceding *global* tokens, otherwise + # rank 1 would mask as if its shard started at position zero. + assert contract.causal_offsets == (cp_rank * 2048, cp_rank * 2048) + + +def test_tp_head_shards_split_qwen3_gqa_evenly(): + contract = MegatronAttentionMaterializer(tp_rank=1).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert (sharding.global_q_heads, sharding.global_kv_heads) == (32, 8) + assert (sharding.local_q_heads, sharding.local_kv_heads) == (16, 4) + assert (sharding.local_q_head_start, sharding.local_kv_head_start) == (16, 4) + + +@pytest.mark.parametrize("tp_world_size", [2, 4, 8]) +def test_supported_tp_degrees_shard_qwen3_gqa(tp_world_size): + """Qwen3-8B has 32 Q heads and 8 KV heads, so TP in {2, 4, 8} all divide.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": tp_world_size} + ) + + sharding = contract.sharding + assert sharding.local_q_heads == 32 // tp_world_size + assert sharding.local_kv_heads == 8 // tp_world_size + + +@pytest.mark.parametrize( + ("knob_value", "expected"), + [("bfloat16", "bf16"), ("float16", "fp16"), ("float32", "fp32"), ("fp16", "fp16")], +) +def test_planner_normalized_dtypes_reach_the_contract(knob_value, expected): + """The planner emits torch spellings; AttentionDType uses short ones.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": knob_value} + ) + + assert contract.dtype.value == expected + + +def test_unknown_dtype_is_rejected_with_the_offending_field(): + with pytest.raises(ValueError, match="training.compute_dtype"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "int8"} + ) + + +def test_indivisible_tp_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": 3} + ) + + +def test_indivisible_cp_sequence_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer(global_sequence_length=4097).build_contract(TRAINING_KNOBS) + + +# -------------------------------------------------------------------------- +# materialization: fail closed rather than silently substitute +# -------------------------------------------------------------------------- + + +def _statuses(materialization, path): + return [app.status for app in materialization.applications if app.path == path] + + +def _readback(materializer, flat, *, source): + contract = materializer.build_contract(flat) + return AttentionRuntimeReadback( + contract=contract, + actual_knobs=dict(flat), + split_kv_plan_set=_plan_set(contract, backend=source), + source=source, + frozen_scope_verified=True, + preprocess_backends=MANDATED_ATTENTION_PREPROCESS_BACKENDS, + preprocess_fallback=False, + ) + + +def test_configured_contract_without_runtime_readback_is_unobservable(): + normalized = { + "batch": {"size": 2}, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "compute_dtype": "bf16", + }, + "attention": {"split_kv_policy": 32}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert materialization.applications + assert {app.status for app in materialization.applications} == { + MaterializationStatus.UNOBSERVABLE + } + assert all(app.actual is None for app in materialization.applications) + + +@pytest.mark.parametrize( + ("materializer_type", "flat", "source"), + [ + (MegatronAttentionMaterializer, TRAINING_KNOBS, "megatron.runtime_readback"), + (VllmRolloutMaterializer, ROLLOUT_KNOBS, "vllm.runtime_readback"), + ], +) +def test_runtime_readback_can_verify_materialized_knobs(materializer_type, flat, source): + configured = materializer_type() + readback = _readback(configured, flat, source=source) + materializer = materializer_type(runtime_readback=readback) + normalized = {} + for path, value in flat.items(): + section, key = path.split(".", 1) + normalized.setdefault(section, {})[key] = value + + materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) + + assert materialization.applications + assert {app.status for app in materialization.applications} == {MaterializationStatus.APPLIED} + side = "training" if materializer_type is MegatronAttentionMaterializer else "rollout" + assert materialization.binding.side_configs[side]["runtime_readback"]["source"] == source + + +def test_runtime_readback_mismatch_is_a_fallback(): + configured = MegatronAttentionMaterializer() + readback = _readback(configured, TRAINING_KNOBS, source="megatron.runtime_readback") + actual = dict(readback.actual_knobs) + actual["training.context_parallel_size"] = 1 + mismatched = AttentionRuntimeReadback( + contract=readback.contract, + actual_knobs=actual, + split_kv_plan_set=readback.split_kv_plan_set, + source=readback.source, + frozen_scope_verified=True, + preprocess_backends=readback.preprocess_backends, + preprocess_fallback=readback.preprocess_fallback, + ) + normalized = { + "batch": {"size": 2}, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "compute_dtype": "bf16", + }, + "attention": {"split_kv_policy": 32}, + } + materialization = MegatronAttentionMaterializer(runtime_readback=mismatched).materialize( + normalized, WS2_ATTENTION_KNOBS + ) + + assert _statuses(materialization, "training.context_parallel_size") == [ + MaterializationStatus.FALLBACK + ] + + +def test_decode_split_kv_plan_set_must_match_kv_cache_lengths(): + rollout, _ = _contracts() + kv_cache = KVCacheSpec( + cache_positions=(3, 5), + kv_seq_lens=(4, 6), + block_table=((0, 1, -1), (2, 3, 4)), + global_token_positions=tuple(range(4)) + tuple(range(6)), + page_size=2, + ) + decode = replace( + rollout, + role=AttentionRole.INFER, + mode=AttentionMode.DECODE, + query_sequence_length=1, + causal_offsets=(3, 5), + kv_cache=kv_cache, + ) + wrong_lengths = build_split_kv_runtime_plan_set( + (4, 8), + tp_world_size=2, + cp_world_size=2, + split_kv=decode.split_kv, + backend="vllm.decode.readback", + ) + + with pytest.raises(ValueError, match="KV-cache lengths"): + AttentionRuntimeReadback( + contract=decode, + actual_knobs={}, + split_kv_plan_set=wrong_lengths, + source="vllm.decode.readback", + frozen_scope_verified=True, + ) + + +def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): + rollout_materializer = VllmRolloutMaterializer() + training_materializer = MegatronAttentionMaterializer() + rollout = _readback(rollout_materializer, ROLLOUT_KNOBS, source="vllm.runtime_readback") + training = _readback( + training_materializer, + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.passed + assert result.provenance["split_kv_runtime"]["rollout"]["coverage"] == ( + "complete_batch_tp_cp_owner_cartesian_product" + ) + assert result.provenance["rollout"]["recorded"]["preprocess.qk_rmsnorm"] == ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS["qk_rmsnorm"] + ) + + +@pytest.mark.parametrize( + ("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(): + rollout_materializer = VllmRolloutMaterializer() + training_materializer = MegatronAttentionMaterializer() + rollout = _readback(rollout_materializer, ROLLOUT_KNOBS, source="vllm.runtime_readback") + rollout = replace(rollout, frozen_scope_verified=False) + training = _readback( + training_materializer, + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.comparable + assert not result.passed + assert any(issue.field == "rollout.frozen_scope_verified" for issue in result.issues) + + +def test_arrival_merge_order_is_unsupported_not_silently_corrected(): + """The control group must stay distinguishable from the treatment.""" + + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_order": "arrival"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_order") == [ + MaterializationStatus.UNSUPPORTED + ] + assert _statuses(materialization, "training.tensor_parallel_size") == [ + MaterializationStatus.ERROR + ] + assert materialization.binding.side_configs["training"]["contract"] is None + assert "arrival" in materialization.binding.side_configs["training"]["contract_error"] + + +def test_unsupported_reduction_invalidates_vllm_contract_applications(): + normalized = { + "batch": {"size": 2}, + "rollout": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_order": "arrival"}, + } + materialization = VllmRolloutMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_order") == [ + MaterializationStatus.UNSUPPORTED + ] + assert _statuses(materialization, "rollout.tensor_parallel_size") == [ + MaterializationStatus.ERROR + ] + assert materialization.binding.side_configs["rollout"]["contract"] is None + + +def test_bf16_reduction_accumulation_is_unsupported(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_acc_dtype": "bf16"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_acc_dtype") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_te_oracle_engine_is_unsupported_until_pr2_pr3(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_engine": "te_oracle"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_engine") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_vllm_cp_falls_back_to_one_in_decode_and_says_why(): + materializer = VllmRolloutMaterializer(mode=AttentionMode.DECODE) + normalized = { + "batch": {"size": 2}, + "rollout": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + + assert materializer.effective_cp_world_size({"rollout.context_parallel_size": 2}) == 1 + + materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) + contract_error = materialization.binding.side_configs["rollout"]["contract_error"] + assert "#235 PR6" in contract_error + assert MaterializationStatus.APPLIED not in {app.status for app in materialization.applications} + + +def test_decode_contract_is_refused_without_kv_cache_identity(): + with pytest.raises(AttentionContractError, match="PR6"): + VllmRolloutMaterializer(mode=AttentionMode.DECODE).build_contract(ROLLOUT_KNOBS) + + +def test_materializers_expose_distinct_implementation_fingerprints(): + megatron = MegatronAttentionMaterializer().implementation_fingerprint + vllm = VllmRolloutMaterializer().implementation_fingerprint + + assert megatron and vllm and megatron != vllm + + +def test_runtime_binding_reports_the_frozen_topology(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + binding = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS).binding + + topology = binding.topology["training"] + assert topology["tensor_parallel_size"] == 2 + assert topology["context_parallel_size"] == 2 + assert topology["world_size"] == 4 + assert topology["pipeline_parallel_size"] == 1 + assert topology["data_parallel_size"] == 1 + + +# -------------------------------------------------------------------------- +# provenance adapters +# -------------------------------------------------------------------------- + + +def test_megatron_provenance_flags_undeclared_frozen_scope_fields(): + adapter = MegatronProvenanceAdapter(SimpleNamespace(deterministic_mode=True)) + + violations = adapter.frozen_scope_violations() + + # Nothing is declared, so every assertion reads as unknown rather than as met. + assert any("expert_model_parallel_size" in text for text in violations) + assert any("fp8" in text for text in violations) + + +def test_megatron_provenance_accepts_a_conforming_dense_config(): + adapter = MegatronProvenanceAdapter( + SimpleNamespace( + deterministic_mode=True, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + sequence_parallel=False, + fp8=None, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + ) + + assert adapter.frozen_scope_violations() == ("fp8 is not declared (expected None)",) + + +def test_megatron_construction_fingerprint_tracks_fusion_changes(): + base = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=False) + fused = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=True) + + assert ( + MegatronProvenanceAdapter(base).construction_fingerprint + != MegatronProvenanceAdapter(fused).construction_fingerprint + ) + + +def test_vllm_provenance_reads_page_size_and_split_kv_policy(): + adapter = VllmProvenanceAdapter( + cache_config=SimpleNamespace(block_size=16, cache_dtype="auto"), + attention_config=SimpleNamespace(flash_attn_max_num_splits_for_cuda_graph=32), + ) + + assert adapter.kv_page_size == 16 + assert adapter.split_kv_policy == 32 + assert adapter.to_dict()["flash_attn_max_num_splits_for_cuda_graph"] == 32 + + +def test_vllm_provenance_flags_fp8_kv_cache_and_cascade_attention(): + adapter = VllmProvenanceAdapter( + model_config=SimpleNamespace(quantization=None, disable_cascade_attn=False), + cache_config=SimpleNamespace( + cache_dtype="fp8", calculate_kv_scales=False, sliding_window=None + ), + parallel_config=SimpleNamespace(pipeline_parallel_size=1, data_parallel_size=1), + ) + + violations = adapter.frozen_scope_violations() + + assert any("cache_dtype" in text for text in violations) + assert any("disable_cascade_attn" in text for text in violations) + + +# -------------------------------------------------------------------------- +# scenario config +# -------------------------------------------------------------------------- + + +SCENARIO = ( + Path(__file__).resolve().parents[1] + / "examples" + / "cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json" +) + + +def test_scenario_uses_megatron_vocabulary_only(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + training = config["baseline"]["training"] + + assert training["attention_backend"] in {"flash", "fused", "unfused", "local", "auto"} + assert training["tensor_parallel_size"] == 2 + assert training["context_parallel_size"] == 2 + assert config["baseline"]["rollout"]["batch_invariant"] is True + + +def test_scenario_knob_paths_all_exist(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + + def paths(mapping, prefix=""): + for key, value in mapping.items(): + path = f"{prefix}{key}" + if isinstance(value, dict): + yield from paths(value, f"{path}.") + else: + yield path + + declared = set(paths(config["baseline"])) + unknown = declared - set(WS2_ATTENTION_KNOBS) + assert not unknown, f"scenario declares unknown knobs: {sorted(unknown)}" + + for intervention in config["interventions"]: + assert intervention["path"] in WS2_ATTENTION_KNOBS diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py new file mode 100644 index 00000000..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_cp_attention.py b/tests/test_cp_attention.py new file mode 100644 index 00000000..cc93874b --- /dev/null +++ b/tests/test_cp_attention.py @@ -0,0 +1,690 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the WS2 deterministic CP attention reference. + +The implementation is a correctness-first prefill/chunked-prefill reference: +local KV blocks produce ``(out, lse)`` partial states and CP merges those states +with fp32 online-softmax arithmetic in logical global-block order. +""" + +import contextlib +import json +import math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + compare_cp_attention_backward, + merge_attention_partial_states, + split_kv_execution_plan_provenance, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp +from rl_engine.kernels.registry import kernel_registry + +_N_HEADS = 32 +_N_KV = 8 +_HEAD_DIM = 128 +_ATOL = 3.0e-6 +_GRAD_ATOL = 1.0e-5 + + +@contextlib.contextmanager +def _single_thread(): + prev = torch.get_num_threads() + torch.set_num_threads(1) + try: + yield + finally: + torch.set_num_threads(prev) + + +def _qkv( + batch, + sq, + skv, + *, + seed, + dtype=torch.float32, + heads=_N_HEADS, + kv_heads=_N_KV, + dim=_HEAD_DIM, +): + gen = torch.Generator().manual_seed(seed) + q = torch.randn(batch, heads, sq, dim, generator=gen, dtype=dtype) + k = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + v = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + return q, k, v + + +def _full_lse(q, k, *, causal, scale=None, key_padding_mask=None): + qf, kf = q.float(), k.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hq % hkv != 0: + raise ValueError("invalid GQA shape") + if hq != hkv: + kf = kf.repeat_interleave(hq // hkv, dim=1) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * ( + scale if scale is not None else 1.0 / math.sqrt(dim) + ) + if causal: + query_pos = torch.arange(skv - sq, skv) + key_pos = torch.arange(skv) + scores = scores.masked_fill( + (key_pos[None, :] > query_pos[:, None])[None, None, :, :], + float("-inf"), + ) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + return torch.logsumexp(scores, dim=-1) + + +def test_cp1_matches_native_attention_and_exports_lse(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 8, 8, seed=1) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + want = native.forward_fp32(q, k, v, causal=True) + want_lse = _full_lse(q, k, causal=True) + + torch.testing.assert_close(out, want, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse, want_lse, atol=_ATOL, rtol=0.0) + assert lse.dtype == torch.float32 + assert lse.shape == q.shape[:3] + + +def test_cp2_prefill_matches_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 9, 9, seed=2) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + out2, lse2 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=2) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + +def test_cp2_consumes_post_rope_qk_with_shared_global_position_metadata(): + op = DeterministicCPAttentionReferenceOp() + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 7, 7, seed=14, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([17, 103], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + + assert not torch.equal(q, pre_rope_q.float()) + assert not torch.equal(k, pre_rope_k.float()) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=1, + ) + out2, lse2 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + +def test_chunked_prefill_replay_matches_unchunked_cp2(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 10, 10, seed=3) + + with _single_thread(): + unchunked_out, unchunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + ) + chunked_out, chunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=3, + ) + + torch.testing.assert_close(chunked_out, unchunked_out, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(chunked_lse, unchunked_lse, atol=_ATOL, rtol=0.0) + + +def test_causal_mask_uses_global_positions_across_cp_boundary(): + op = DeterministicCPAttentionReferenceOp() + batch, heads, kv_heads, seq, dim = 1, 2, 1, 5, 3 + q = torch.zeros(batch, heads, seq, dim) + k = torch.zeros(batch, kv_heads, seq, dim) + v = torch.arange(seq * dim, dtype=torch.float32).reshape(1, 1, seq, dim) + out = op.forward_fp32(q, k, v, causal=True, cp_world_size=2) + + expected = torch.stack([v[0, 0, : index + 1].mean(dim=0) for index in range(seq)]) + expected = expected.reshape(1, 1, seq, dim).repeat(1, heads, 1, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + + +def test_position_offsets_apply_varlen_causal_metadata_per_batch_row(): + op = DeterministicCPAttentionReferenceOp() + q = torch.zeros(2, 2, 2, 1) + k = torch.zeros(2, 1, 4, 1) + v = torch.arange(8, dtype=torch.float32).reshape(2, 1, 4, 1) + + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=torch.tensor([0, 11]), + key_position_offsets=torch.tensor([0, 10]), + cp_world_size=2, + kv_chunk_size=1, + ) + + expected = torch.tensor([0.0, 0.5, 4.5, 5.0]).reshape(2, 1, 2, 1).repeat(1, 2, 1, 1) + expected_lse = torch.log(torch.tensor([1.0, 2.0, 2.0, 3.0])).reshape(2, 1, 2) + expected_lse = expected_lse.repeat(1, 2, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(lse, expected_lse, atol=1.0e-6, rtol=0.0) + + +def test_merge_order_uses_global_block_index_not_arrival_order(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 6, 6, seed=4) + first = op.local_partial_state( + q, + k[:, :, :3], + v[:, :, :3], + q_start=0, + k_start=0, + total_kv_len=6, + causal=True, + ) + second = op.local_partial_state( + q, + k[:, :, 3:], + v[:, :, 3:], + q_start=0, + k_start=3, + total_kv_len=6, + causal=True, + ) + + forward = merge_attention_partial_states([first, second]) + reversed_arrival = merge_attention_partial_states([second, first]) + assert torch.equal(forward.out, reversed_arrival.out) + assert torch.equal(forward.lse, reversed_arrival.lse) + + +def test_key_padding_mask_and_all_masked_rows_are_stable(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 6, 6, seed=5) + mask = torch.tensor( + [ + [True, True, True, False, False, False], + [False, False, False, False, False, False], + ], + dtype=torch.bool, + ) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=False, + key_padding_mask=mask, + cp_world_size=2, + kv_chunk_size=2, + ) + want = native.forward_fp32(q, k, v, causal=False, key_padding_mask=mask) + + torch.testing.assert_close(out[:1], want[:1], atol=_ATOL, rtol=0.0) + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert torch.isneginf(lse[1]).all() + assert torch.isfinite(out).all() + + +def test_empty_query_and_empty_kv_edges_are_stable(): + op = DeterministicCPAttentionReferenceOp() + q_empty = torch.randn(1, 2, 0, 4, requires_grad=True) + k_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + v_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + out, lse = op.forward_fp32_with_lse(q_empty, k_empty, v_empty, cp_world_size=2) + assert out.shape == (1, 2, 0, 4) + assert lse.shape == (1, 2, 0) + assert out.requires_grad + out.sum().backward() + assert torch.equal(q_empty.grad, torch.zeros_like(q_empty)) + assert torch.equal(k_empty.grad, torch.zeros_like(k_empty)) + assert torch.equal(v_empty.grad, torch.zeros_like(v_empty)) + + q = torch.randn(1, 2, 3, 4) + out, lse = op.forward_fp32_with_lse(q, k_empty, v_empty, causal=False, cp_world_size=4) + assert torch.equal(out, torch.zeros_like(out)) + assert torch.isneginf(lse).all() + + +def test_empty_kv_backward_returns_zero_grads(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 2, 3, 4, requires_grad=True) + k = torch.randn(1, 1, 0, 4, requires_grad=True) + v = torch.randn(1, 1, 0, 4, requires_grad=True) + + out = op.forward_fp32(q, k, v, causal=False, cp_world_size=4) + assert out.requires_grad + out.sum().backward() + + assert torch.equal(q.grad, torch.zeros_like(q)) + assert torch.equal(k.grad, torch.zeros_like(k)) + assert torch.equal(v.grad, torch.zeros_like(v)) + + +def test_bf16_forward_uses_fp32_merge_then_final_write(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 8, 8, seed=6, dtype=torch.bfloat16) + + out, lse = op.forward_with_lse(q, k, v, causal=True, cp_world_size=2, kv_chunk_size=2) + fp32_out, fp32_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + + assert out.dtype == torch.bfloat16 + assert lse.dtype == torch.float32 + assert torch.equal(out, fp32_out.to(torch.bfloat16)) + assert torch.equal(lse, fp32_lse) + + +def test_cp2_chunked_gradients_match_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 5, 5, seed=12, heads=4, kv_heads=2, dim=8) + q_ref = q.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + v_ref = v.detach().clone().requires_grad_(True) + q_cp = q.detach().clone().requires_grad_(True) + k_cp = k.detach().clone().requires_grad_(True) + v_cp = v.detach().clone().requires_grad_(True) + gen = torch.Generator().manual_seed(13) + dy = torch.randn(1, 4, 5, 8, generator=gen) + + with _single_thread(): + out_ref = op.forward_fp32(q_ref, k_ref, v_ref, causal=True, cp_world_size=1) + out_cp = op.forward_fp32( + q_cp, + k_cp, + v_cp, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + out_ref.backward(dy) + out_cp.backward(dy) + + torch.testing.assert_close(out_cp, out_ref, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(q_cp.grad, q_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(k_cp.grad, k_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(v_cp.grad, v_ref.grad, atol=1.0e-5, rtol=0.0) + + +def test_backward_report_cp2_prefill_matches_cp1_reference(): + q, k, v = _qkv(1, 5, 5, seed=15, heads=4, kv_heads=2, dim=8) + dout = torch.randn(1, 4, 5, 8, generator=torch.Generator().manual_seed(16)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + output_dtype=torch.float32, + ) + + assert report.reference_name == "cp1_backward_reference" + drift = report.drifts[0] + assert drift.candidate_name == "cp2_backward" + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + assert drift.out.max_abs <= _ATOL + assert drift.lse.max_abs <= _ATOL + assert len(drift.per_rank) == 2 + assert drift.per_rank[0].dq.active_count > 0 + assert drift.per_rank[1].dk.active_count > 0 + assert drift.provenance["saved_forward_state"][0] == "out" + assert drift.provenance["merge_order"] == "global_block_index" + assert drift.provenance["te_backward_oracle"] == "not_used" + assert drift.provenance["decode_backward"] == "not_supported" + json.dumps(report.to_dict()) + + +def test_backward_report_cp2_chunked_prefill_matches_cp1_reference(): + q, k, v = _qkv(1, 6, 6, seed=17, heads=4, kv_heads=2, dim=8) + dout = torch.randn(1, 4, 6, 8, generator=torch.Generator().manual_seed(18)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.float32, + ) + + drift = report.drifts[0] + assert drift.candidate_name == "cp2_chunked_backward" + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + assert drift.provenance["attention_mode"] == "chunked_prefill" + assert drift.provenance["kv_chunk_size"] == 2 + assert drift.provenance["requested_split_kv_policy"] == "fixed" + assert drift.provenance["actual_split_kv_plans"] == [ + { + "owner_cp_rank": 0, + "requested_split_kv_policy": "fixed", + "requested_split_kv_size": 2, + "actual_split_kv_policy": "fixed", + "actual_split_kv_size": 2, + "actual_split_kv_count": 2, + "actual_split_boundaries": [[0, 2], [2, 3]], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "deterministic_cp_backward_reference", + "split_kv_plan_source": "reference_execution", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + }, + { + "owner_cp_rank": 1, + "requested_split_kv_policy": "fixed", + "requested_split_kv_size": 2, + "actual_split_kv_policy": "fixed", + "actual_split_kv_size": 2, + "actual_split_kv_count": 2, + "actual_split_boundaries": [[3, 5], [5, 6]], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "deterministic_cp_backward_reference", + "split_kv_plan_source": "reference_execution", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + }, + ] + + +def test_split_kv_plan_never_crosses_cp_owner_boundaries(): + plans = split_kv_execution_plan_provenance( + 10, + cp_world_size=3, + kv_chunk_size=3, + backend="test-reference", + ) + + assert [plan["actual_split_boundaries"] for plan in plans] == [ + [[0, 3], [3, 4]], + [[4, 7]], + [[7, 10]], + ] + assert [plan["owner_cp_rank"] for plan in plans] == [0, 1, 2] + + +def test_backward_report_preserves_post_rope_position_metadata(): + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 5, 5, seed=19, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([23, 101], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + dout = torch.randn(2, 4, 5, 8, generator=torch.Generator().manual_seed(20)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.float32, + ) + + drift = report.drifts[0] + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + + +def test_qwen3_8b_local_tp2_cp2_bf16_backward_report_smoke(): + # Qwen3-8B global Hq/Hkv is 32/8. A TP=2 local shard owns 16/4 heads. + q, k, v = _qkv( + 1, + 4, + 4, + seed=21, + dtype=torch.bfloat16, + heads=16, + kv_heads=4, + dim=_HEAD_DIM, + ) + dout = torch.randn( + 1, + 16, + 4, + _HEAD_DIM, + generator=torch.Generator().manual_seed(22), + dtype=torch.bfloat16, + ) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.bfloat16, + ) + + drift = report.drifts[0] + assert drift.provenance["q_dtype"] == "bfloat16" + assert drift.provenance["output_dtype"] == "bfloat16" + assert drift.provenance["downcast_at"] == "final_write" + assert drift.dq.max_abs <= 5.0e-2 + assert drift.dk.max_abs <= 5.0e-2 + assert drift.dv.max_abs <= 5.0e-2 + + +def test_backward_report_validates_dout_shape_and_dtype(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=23, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="dout must have shape"): + op.backward_reference(q, k, v, torch.randn(1, 4, 3, 8), cp_world_size=2) + + with pytest.raises(ValueError, match="dout must be a real floating-point tensor"): + op.backward_reference( + q, + k, + v, + torch.ones(1, 4, 4, 8, dtype=torch.long), + cp_world_size=2, + ) + + with pytest.raises(ValueError, match="dout must have the same dtype"): + op.backward_reference( + q.to(torch.bfloat16), + k.to(torch.bfloat16), + v.to(torch.bfloat16), + torch.ones_like(q), + cp_world_size=2, + ) + + +def test_inputs_are_not_mutated(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 6, 6, seed=7) + mask = torch.ones(2, 6, dtype=torch.bool) + qc, kc, vc, mc = q.clone(), k.clone(), v.clone(), mask.clone() + + op.forward_fp32_with_lse(q, k, v, causal=True, key_padding_mask=mask, cp_world_size=2) + + assert torch.equal(q, qc) + assert torch.equal(k, kc) + assert torch.equal(v, vc) + assert torch.equal(mask, mc) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"cp_world_size": 0}, "cp_world_size"), + ({"cp_world_size": 2, "kv_chunk_size": 0}, "kv_chunk_size"), + ], +) +def test_invalid_parallelism_arguments_raise(kwargs, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=8) + with pytest.raises(ValueError, match=message): + op.forward_fp32_with_lse(q, k, v, causal=True, **kwargs) + + +def test_invalid_gqa_and_mask_shapes_raise(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 6, 4, _HEAD_DIM) + k = torch.randn(1, 4, 4, _HEAD_DIM) + v = torch.randn(1, 4, 4, _HEAD_DIM) + with pytest.raises(ValueError, match="not divisible"): + op.forward_fp32_with_lse(q, k, v, causal=True) + + q, k, v = _qkv(1, 4, 4, seed=9) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 3, dtype=torch.bool)) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 4)) + with pytest.raises(ValueError, match="query_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + query_position_offsets=torch.ones(2, dtype=torch.long), + ) + with pytest.raises(ValueError, match="key_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + key_position_offsets=torch.ones(1, dtype=torch.float32), + ) + + +@pytest.mark.parametrize("scale", [0.0, -1.0, float("nan"), float("inf"), True, "bad"]) +def test_invalid_scale_fails_before_attention_math(scale): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=24, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="scale must be a positive finite number"): + op.forward_fp32_with_lse(q, k, v, scale=scale) + + +def test_qkv_dtype_and_floating_contract_fails_closed(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=25, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="same dtype"): + op.forward_fp32_with_lse(q, k.to(torch.bfloat16), v) + with pytest.raises(ValueError, match="real floating-point"): + op.forward_fp32_with_lse(q.to(torch.long), k.to(torch.long), v.to(torch.long)) + + +@pytest.mark.parametrize("kwargs", [{"cp_world_size": True}, {"kv_chunk_size": True}]) +def test_boolean_parallelism_arguments_fail_closed(kwargs): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=26, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError): + op.forward_fp32_with_lse(q, k, v, **kwargs) + + +@pytest.mark.parametrize("output_dtype", [torch.long, torch.complex64, "fp32"]) +def test_nonfloating_output_dtype_fails_closed(output_dtype): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=27, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="output_dtype must be a real floating-point"): + op.forward_with_lse(q, k, v, output_dtype=output_dtype) + + +def test_partial_states_must_remain_fp32_and_colocated(): + out = torch.zeros(1, 1, 1, 1, dtype=torch.bfloat16) + lse = torch.zeros(1, 1, 1, dtype=torch.float32) + + with pytest.raises(ValueError, match="must remain FP32"): + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=1) + + +def test_overlapping_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="overlap"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=3), + AttentionPartialState(out=out, lse=lse, block_start=2, block_end=4), + ] + ) + + +def test_gapped_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="gap-free"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=2), + AttentionPartialState(out=out, lse=lse, block_start=3, block_end=4), + ] + ) + + +def test_registry_dispatches_cp_attention_reference(): + assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) diff --git a/tests/test_cp_attention_transformer_engine.py b/tests/test_cp_attention_transformer_engine.py new file mode 100644 index 00000000..d98e31a1 --- /dev/null +++ b/tests/test_cp_attention_transformer_engine.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Optional Transformer Engine oracle tests for CP attention merging.""" + +from __future__ import annotations + +import importlib + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + merge_attention_partial_states, +) + + +def _te_context_parallel_module(): + try: + return importlib.import_module( + "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" + ) + except (ImportError, OSError, RuntimeError) as exc: + pytest.skip(f"Transformer Engine context-parallel attention is unavailable: {exc}") + + +def test_cp_attention_merge_matches_transformer_engine_corrections(): + te_cp = _te_context_parallel_module() + gen = torch.Generator().manual_seed(238) + out_a = torch.randn(2, 3, 5, 4, generator=gen) + out_b = torch.randn(2, 3, 5, 4, generator=gen) + lse_a = torch.randn(2, 3, 5, generator=gen) + lse_b = torch.randn(2, 3, 5, generator=gen) + + ours = merge_attention_partial_states( + [ + AttentionPartialState(out=out_b, lse=lse_b, block_start=5, block_end=9), + AttentionPartialState(out=out_a, lse=lse_a, block_start=0, block_end=5), + ] + ) + + te_lse = lse_a.clone() + te_cp.flash_attn_fwd_softmax_lse_correction(te_lse, lse_b) + te_out = te_cp.flash_attn_fwd_out_correction_init(out_a.clone(), te_lse, lse_a, seq_dim=2) + te_cp.flash_attn_fwd_out_correction(te_out, out_b, te_lse, lse_b, seq_dim=2) + + torch.testing.assert_close(ours.lse, te_lse, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(ours.out, te_out, atol=1.0e-6, rtol=0.0) 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_operator_inputs.py b/tests/test_operator_inputs.py index 4f742734..ca919d89 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -37,6 +37,7 @@ def _args(**overrides): "matmul", "det_gemm", "attention", + "cp_attention", "logp", "linear_logp", "batch_invariant_logp", 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 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 5eb75cbd..5dfef6aa 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(): @@ -29,6 +41,51 @@ def test_logprob_bfloat16_tolerance_covers_observed_reference_drift(): 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 + + def test_attention_bfloat16_tolerance_matches_contract(): contract = load_contract() tolerance = contract["accuracy"]["default"]["attention"]["bfloat16"]