diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e9e9c91d..a4c1fe89 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -84,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/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..06966ebe --- /dev/null +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -0,0 +1,135 @@ +# WS2 Attention Cross-Configuration Integration + +Implements PR4 of [#235](https://github.com/RL-Align/RL-Kernel/issues/235): wiring the +CP attention path into the cross-configuration planner/runtime for the Qwen3-8B +TP=2 CP=2 BF16 target. + +Builds on [#236](https://github.com/RL-Align/RL-Kernel/pull/236) (attention contract +and dispatch metadata), [#238](https://github.com/RL-Align/RL-Kernel/pull/238) +(deterministic CP reference) and [#230](https://github.com/RL-Align/RL-Kernel/pull/230) +(cross-configuration framework). + +## What "bind to the same contract" means here + +The PR4 acceptance criteria say rollout and training descriptors must "bind to the +same semantic attention contract". Under the frozen deployment the two sides can +never produce identical `AttentionContract` instances: + +| | training (Megatron) | rollout (vLLM) | +| --- | --- | --- | +| mode | full-sequence prefill | chunked prefill, later decode | +| CP | `context_parallel_size`, whole forward | `prefill_context_parallel_size`, prefill only | +| KV | no paging | paged KV with a block table | +| backend vocabulary | `AttnBackend{flash,fused,unfused,local,auto}` | `AttentionBackendEnum` | + +Read literally, the criterion is unsatisfiable. It is therefore implemented as three +tiers, in `rl_engine/alignment/cross_config/attention_binding.py`: + +| tier | fields | rule | failure | +| --- | --- | --- | --- | +| `IDENTICAL` | checkpoint, model version, weight version, tokenizer, token ids, active mask, position ids, padding side, pre-update state, Q/KV heads, head dim, RoPE theta/scaling/rotary dim, QK-Norm, cached global token positions, KV sequence lengths | equal bit for bit | `comparable=False`; no drift number from the pair means anything | +| `SEMANTIC` | `reduction.merge`, `reduction.acc_dtype`, `reduction.order`, `reduction.downcast_at`, `export_lse`, cross-side determinism mode | both sides equal **and** equal to the WS2 mandate | fail closed | +| `RECORDED` | `mode`, `backend_id`, `reduction.engine`, RoPE materialization state and fusion boundary, KV-cache paging, CP/TP world sizes, local sequence length | free to differ | none; recorded into provenance and measured | + +Two placements are load-bearing: + +* **`reduction.engine` is `RECORDED`, not `SEMANTIC`.** Training may run the in-op + deterministic reference while rollout runs a Transformer Engine merge oracle. + Forcing them equal would defeat the oracle comparison that #235 PR2/PR3/PR5/PR6 + depend on. +* **`reduction.order` and `reduction.acc_dtype` are `SEMANTIC`.** This is the entire + WS2 claim: merge order and accumulation precision are decided by the contract, not + by whichever backend happens to be selected. + +`comparable` and `passed` are separate flags. A pair with mismatched identity is not +comparable. A pair that is comparable but violates the reduction mandate is still +rejected -- the drift would be real but attributable to the wrong thing. + +## Determinism is not one thing + +`rl_engine/alignment/cross_config/determinism.py` probes both sides and compares +them, because the two frameworks mean different things by "deterministic": + +| | Megatron `deterministic_mode` | vLLM `VLLM_BATCH_INVARIANT` | +| --- | --- | --- | +| `NCCL_ALGO` | asserts membership in a five-value set | hard-sets `allreduce:tree` | +| `NCCL_PROTO`, channels, threads | not managed | hard-set (`Simple`, `1`, `1`) | +| TF32 | **not managed at all** | disabled (`fp32_precision="ieee"`) | +| BF16 reduced-precision reduction | not managed | disabled | +| cuBLAS workspace / BLAS library | not managed | `:4096:8`, cuBLASLt | +| GEMM | cuBLAS / TE | Triton `matmul_persistent` | +| FlashAttention | forbidden | permitted | + +`NCCL_ALGO`, `NCCL_PROTO` and `CUBLAS_WORKSPACE_CONFIG` change arithmetic, so a +mismatch there is blocking. The remaining differences -- including the TF32 and +BF16-reduction asymmetry, which under a pure BF16 GEMM path does not fire -- are +recorded so the asymmetry appears in every artifact rather than being invisible. + +## Runtime adapters + +Before this PR the only `RuntimeMaterializer` in the repository was +`CpuSmokeMaterializer` over a synthetic CPU model, and every named scenario +(`S1`/`S2`/`S3`) was planning-only. This PR adds the first two framework-shaped +adapters: + +* `adapters/megatron.py` -- `MegatronProvenanceAdapter` (construction and + distributed-context fingerprints, determinism probe, frozen-scope assertions) and + `MegatronAttentionMaterializer`. +* `adapters/vllm.py` -- `VllmProvenanceAdapter` (adds `kv_page_size` from + `CacheConfig.block_size` and `split_kv_policy` from + `AttentionConfig.flash_attn_max_num_splits_for_cuda_graph`) and + `VllmRolloutMaterializer`. + +Neither module imports `megatron` or `vllm`; configs are duck-typed, so the binding +rules are exercised on CPU in CI rather than only on a 2-node cluster. + +## Fail closed, never substitute + +`unsupported_reduction_reason` rejects requests that #236 cannot express, instead of +collapsing them onto the supported value: + +| request | status | why | +| --- | --- | --- | +| `attention.reduction_order=arrival` | `UNSUPPORTED` | the control group must stay distinguishable from the treatment | +| `attention.reduction_downcast_at=per_block` | `UNSUPPORTED` | `DowncastPoint` declares only `final_write` | +| `attention.reduction_engine=te_oracle` | `UNSUPPORTED` | the TE merge oracle lands in #235 PR2/PR3; PR4's TE plan is provenance only | +| `attention.reduction_acc_dtype=bf16` | `UNSUPPORTED` | the CP `(out, lse)` merge accumulates in FP32 | +| `rollout.context_parallel_size>1` with `mode=decode` | `FALLBACK` | vLLM CP covers prefill only; recorded with the reason | + +## Knobs + +`adapters/knobs.py` extends `V1_KNOBS` additively. Added: training-side +`tensor_parallel_size` / `context_parallel_size` / `deterministic_mode` / +`cp_comm_type`, `rollout.batch_invariant` / `rollout.kv_block_size`, and the +reduction axis (`acc_dtype`, `order`, `downcast_at`, `engine`) plus +`attention.fusion_boundary` and `attention.split_kv_policy`. + +`training.attention_backend` keeps its path but its value domain is replaced with +Megatron's `AttnBackend`; the HuggingFace names have no Megatron counterpart, so this +is a replacement rather than a mapping. + +Not done here, because both change `V1_KNOBS` itself and would break existing +cross-config tests: removing `training.sharding` (Megatron has no such concept, and +DP=1 makes it moot) and renaming `rollout.context_parallel_size` to reflect that it +binds to `prefill_context_parallel_size`. + +## Scenario + +`examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json` supersedes +`cross_config_s1_distributed_smoke.json` and +`cross_config_s3_qwen3_8b_tp4_cp4_bf16.json`, whose training sides used `sdpa` / +`flash_attention_2` and `sharding: fsdp` -- none of which exist under Megatron -- and +whose TP=4/CP=4 topology does not match the target. +`cross_config_s2_vllm_tp_vs_fsdp.json` has no Megatron-only counterpart and should be +retired rather than rewritten. + +## Out of scope + +Deliberately not in this PR: + +* launching `torchrun`, initializing process groups, or executing attention; +* decode-mode materialization, which needs the validated `KVCacheSpec` from #235 PR6 + and is refused with that reference rather than stubbed; +* Transformer Engine calls of any kind (PR4's TE plan is policy and provenance only); +* distributed drift benchmarks and report artifacts (#235 PR5); +* fused production backend alignment (#235 PR7) and backward (#235 PR8). diff --git a/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/operators/attention.md b/docs/operators/attention.md index ebff9a58..52374483 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -85,6 +85,34 @@ Calling it (`__call__` -> `forward(...)`) computes in the input dtype; `forward_ the explicit fp32 golden path (NativeAttentionOp only). The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. +### 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 + +`kernel_registry.get_op("cp_attention")` resolves to +`DeterministicCPAttentionReferenceOp`, the WS2 correctness-first context-parallel +reference. It emulates CP prefill and chunked-prefill by splitting logical query +and KV sequence blocks, computing per-block `(out, lse)` partial states, and +merging them in fp32 by global KV block index. This path is not a production +fused backend; it defines the CP/LSE merge behavior that downstream fused paths +must match. Optional per-batch `query_position_offsets` / `key_position_offsets` +cover varlen causal-mask metadata while keeping the dense tensor layout. + +For Qwen3 WS2, `cp_attention` consumes post-QK-Norm, post-RoPE Q/K. It does not +call `NativeRoPEOp` internally and does not hide RoPE inside the CP merge. The +position offsets passed to CP attention must describe the same absolute token +positions used when RoPE was applied, so PR3 validates the post-RoPE Q/K boundary +while PR7 can later validate production fused `RoPE+Attention` kernels. + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -137,6 +165,8 @@ memory. ```bash python -m pytest tests/test_attention.py -v +python -m pytest tests/test_cp_attention.py -v +python -m pytest tests/test_cp_attention_transformer_engine.py -v # optional TE oracle ``` Covers: `forward_fp32` vs an independent fp32 reference (bitwise), strict-fp32 under hostile @@ -146,14 +176,30 @@ invariance (slice + chunked, bitwise; padding is near-equality only, see below), gradient flow, registry dispatch, and a GPU-only LARGE Qwen3-8B real-shape smoke test. +`tests/test_cp_attention.py` covers the WS2 CP reference: CP=1 vs standard +attention, CP=2 prefill vs CP=1, post-RoPE Q/K input semantics with shared +global position metadata, chunked-prefill replay, global-position causal masking +across CP boundaries, order-independent LSE merge by global block index, +padding/all-masked stability, BF16 final-write behavior, input purity, argument +validation, and registry dispatch. +`make_operator_inputs("cp_attention", ...)` also emits a CP=2 chunked-prefill +synthetic case for local harnesses. +`tests/test_cp_attention_transformer_engine.py` optionally imports NVIDIA +Transformer Engine's context-parallel PyTorch correction helpers and checks that +RL-Kernel's fp32 `(out, lse)` merge matches those helpers; the test skips when +Transformer Engine is not installed. + ## Implementation Files - `rl_engine/kernels/ops/pytorch/attention/standard_attn.py` — ground-truth reference +- `rl_engine/kernels/ops/pytorch/attention/cp_attention.py` — CP prefill/chunked reference - `rl_engine/kernels/ops/cuda/attention/deterministic_attn.py` — CUDA deterministic op - `csrc/cuda/attention/deterministic_attention.cu` — CUDA kernels - `rl_engine/kernels/registry.py` - `tests/test_attention.py` - `tests/test_deterministic_attention_cuda.py` +- `tests/test_cp_attention.py` +- `tests/test_cp_attention_transformer_engine.py` ## Fixed Reduction Order (CUDA Deterministic Backend) @@ -218,6 +264,10 @@ for measured peak memory at representative shapes. - First version: `D=128` only (Qwen3-8B alignment). - Supported dtypes: BF16, FP16. - Full materialization of scores/P limits practical sequence length. +- `cp_attention` is a PyTorch reference for CP prefill/chunked-prefill semantics, + not a distributed runtime or fused kernel. +- `cp_attention` consumes post-RoPE Q/K for Qwen3 WS2; RoPE execution and fused + `RoPE+Attention` backend alignment are outside PR3. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op). - No FP8, no multi-GPU / sequence-parallel. diff --git a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json new file mode 100644 index 00000000..e74d3e6c --- /dev/null +++ b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json @@ -0,0 +1,87 @@ +{ + "experiment_id": "ws2-qwen3-8b-attention-tp2-cp2", + "scenario_id": "qwen3_8b_megatron_tp2_cp2_vllm", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "scenario": { + "issue": "https://github.com/RL-Align/RL-Kernel/issues/235", + "pull_request": "PR4 -- cross-config integration", + "model": "Qwen3-8B dense", + "training_framework": "megatron", + "rollout_framework": "vllm", + "topology": "2 nodes x 2 GPUs, TP=2 CP=2 PP=1 DP=1, BF16, SM90", + "notes": [ + "Supersedes cross_config_s1_distributed_smoke.json and", + "cross_config_s3_qwen3_8b_tp4_cp4_bf16.json, whose training side used", + "HuggingFace attention backends and FSDP sharding. Neither exists in", + "Megatron, and DP=1 makes the sharding knob meaningless.", + "cross_config_s2_vllm_tp_vs_fsdp.json has no Megatron-only counterpart at", + "all and should be retired rather than rewritten.", + "rollout.context_parallel_size binds to vLLM", + "ParallelConfig.prefill_context_parallel_size and therefore applies to", + "prefill only; a decode-mode contract runs at CP=1." + ] + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 1, + "dtype": "bfloat16", + "enable_prefix_caching": false, + "enforce_eager": true, + "batch_invariant": true, + "kv_block_size": 16 + }, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "attention_backend": "unfused", + "compute_dtype": "bfloat16", + "deterministic_mode": true, + "cp_comm_type": "p2p", + "sharding": "unsharded" + }, + "attention": { + "reduction_acc_dtype": "fp32", + "reduction_order": "global_block_index", + "reduction_downcast_at": "final_write", + "reduction_engine": "in_op_reference", + "fusion_boundary": "unfused_rope_attention", + "split_kv_policy": 32 + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "training.context_parallel_size", + "values": [1, 2] + }, + { + "path": "training.tensor_parallel_size", + "values": [1, 2] + }, + { + "path": "attention.fusion_boundary", + "values": ["unfused_rope_attention", "fused_rope_attention"] + }, + { + "path": "training.cp_comm_type", + "values": ["p2p", "all_gather"] + }, + { + "path": "attention.reduction_order", + "values": ["global_block_index", "arrival"] + }, + { + "path": "attention.reduction_acc_dtype", + "values": ["fp32", "bf16"] + } + ] +} diff --git a/rl_engine/alignment/cross_config/adapters/__init__.py b/rl_engine/alignment/cross_config/adapters/__init__.py new file mode 100644 index 00000000..c2db2134 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/__init__.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime adapters for the WS2 Qwen3-8B Megatron + vLLM cross-config target.""" + +from rl_engine.alignment.cross_config.adapters._common import QWEN3_8B, Qwen3ModelSpec +from rl_engine.alignment.cross_config.adapters.knobs import ( + MEGATRON_ATTENTION_BACKENDS, + WS2_ATTENTION_KNOB_DESCRIPTORS, + WS2_ATTENTION_KNOBS, + WS2_ATTENTION_NORMALIZERS, +) +from rl_engine.alignment.cross_config.adapters.megatron import ( + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, +) +from rl_engine.alignment.cross_config.adapters.vllm import ( + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", + "QWEN3_8B", + "Qwen3ModelSpec", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] diff --git a/rl_engine/alignment/cross_config/adapters/_common.py b/rl_engine/alignment/cross_config/adapters/_common.py new file mode 100644 index 00000000..33b4d91c --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/_common.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared pieces for the Megatron and vLLM WS2 attention adapters.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from rl_engine.alignment.cross_config.runtime import KnobApplication +from rl_engine.alignment.cross_config.schema import ( + IsolationScope, + KnobDescriptor, + MaterializationStatus, +) +from rl_engine.kernels.attention_contract import ( + AttentionDType, + AttentionMerge, + DowncastPoint, + ReductionEngine, + ReductionOrder, + ReductionSpec, + ShardingSpec, +) + +__all__ = [ + "QWEN3_8B", + "Qwen3ModelSpec", + "application", + "attention_dtype", + "build_reduction_spec", + "build_sharding_spec", + "causal_offsets_for", + "flatten", + "unsupported_reduction_reason", +] + + +@dataclass(frozen=True) +class Qwen3ModelSpec: + """Architecture constants for the frozen dense target. + + These are *not* knobs. #235/#239/#241 all fix Qwen3-8B dense, so they belong to + the scenario, and both sides must agree on them or the comparison is void. + """ + + name: str = "qwen3-8b" + hidden_size: int = 4096 + ffn_hidden_size: int = 12288 + num_layers: int = 36 + q_heads: int = 32 + kv_heads: int = 8 + head_dim: int = 128 + real_vocab_size: int = 151936 + rope_theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + qk_layernorm: bool = True + + def identity_fields(self) -> dict[str, Any]: + """The subset of :data:`IDENTITY_FIELDS` this spec is responsible for.""" + + return { + "q_heads": self.q_heads, + "kv_heads": self.kv_heads, + "head_dim": self.head_dim, + "rope_theta": self.rope_theta, + "rope_scaling": self.rope_scaling, + "rotary_dim": self.rotary_dim, + "qk_layernorm": self.qk_layernorm, + } + + +QWEN3_8B = Qwen3ModelSpec() + + +#: The planner normalizes dtype knobs to torch spellings (``bfloat16``), while +#: :class:`AttentionDType` uses short spellings (``bf16``). Passing a normalized knob +#: straight into the enum raises, so every adapter must translate here rather than +#: each inventing its own mapping. +_DTYPE_ALIASES: Mapping[str, AttentionDType] = { + "bf16": AttentionDType.BF16, + "bfloat16": AttentionDType.BF16, + "fp16": AttentionDType.FP16, + "float16": AttentionDType.FP16, + "half": AttentionDType.FP16, + "fp32": AttentionDType.FP32, + "float32": AttentionDType.FP32, + "float": AttentionDType.FP32, +} + + +def attention_dtype(value: Any, *, field: str) -> AttentionDType: + """Translate a normalized knob dtype into an :class:`AttentionDType`.""" + + if isinstance(value, AttentionDType): + return value + key = str(value).strip().lower().replace("torch.", "") + try: + return _DTYPE_ALIASES[key] + except KeyError as exc: + raise ValueError( + f"{field}={value!r} is not a supported attention dtype; " + f"expected one of {sorted(set(_DTYPE_ALIASES))}" + ) from exc + + +def flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + """Flatten nested knob mappings into dotted paths.""" + + flat: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}{key}" + if isinstance(child, Mapping): + flat.update(flatten(child, f"{path}.")) + else: + flat[path] = child + return flat + + +def application( + descriptor: KnobDescriptor, + requested: Any, + materialized: Any, + actual: Any, + status: MaterializationStatus, + reason: str, + **evidence: Any, +) -> KnobApplication: + return KnobApplication( + path=descriptor.path, + requested=requested, + materialized=materialized, + actual=actual, + lifecycle=descriptor.lifecycle, + status=status, + evidence={"reason": reason, **evidence}, + critical=descriptor.critical, + ) + + +def unsupported_reduction_reason(flat: Mapping[str, Any]) -> str | None: + """Return why the requested reduction cannot be materialized, if it cannot. + + #236 declares single-member enums for merge order, downcast point and reduction + engine, so the alternative knob values exist only as control groups. Requesting + one must fail loudly rather than quietly collapse onto the supported value -- + silently substituting ``global_block_index`` for a requested ``arrival`` would + make the control group indistinguishable from the treatment. + """ + + order = flat.get("attention.reduction_order") + if order is not None and order != ReductionOrder.GLOBAL_BLOCK_INDEX.value: + return ( + f"attention.reduction_order={order!r} has no backend; #236 ReductionOrder " + f"declares only {ReductionOrder.GLOBAL_BLOCK_INDEX.value!r}" + ) + downcast = flat.get("attention.reduction_downcast_at") + if downcast is not None and downcast != DowncastPoint.FINAL_WRITE.value: + return ( + f"attention.reduction_downcast_at={downcast!r} has no backend; #236 " + f"DowncastPoint declares only {DowncastPoint.FINAL_WRITE.value!r}" + ) + engine = flat.get("attention.reduction_engine") + if engine is not None and engine != ReductionEngine.IN_OP_REFERENCE.value: + return ( + f"attention.reduction_engine={engine!r} has no backend; the Transformer " + "Engine merge oracle lands in #235 PR2/PR3, not here" + ) + acc_dtype = flat.get("attention.reduction_acc_dtype") + if ( + acc_dtype is not None + and attention_dtype(acc_dtype, field="attention.reduction_acc_dtype") + is not AttentionDType.FP32 + ): + return ( + f"attention.reduction_acc_dtype={acc_dtype!r} violates the WS2 mandate; " + "the CP (out, lse) merge accumulates in fp32" + ) + return None + + +def build_reduction_spec(flat: Mapping[str, Any]) -> ReductionSpec: + """Build the reduction spec, having already rejected unsupported requests.""" + + return ReductionSpec( + merge=AttentionMerge.ONLINE_SOFTMAX_LSE, + acc_dtype=AttentionDType.FP32, + order=ReductionOrder.GLOBAL_BLOCK_INDEX, + downcast_at=DowncastPoint.FINAL_WRITE, + engine=ReductionEngine.IN_OP_REFERENCE, + ) + + +def build_sharding_spec( + *, + model: Qwen3ModelSpec, + tp_rank: int, + tp_world_size: int, + cp_rank: int, + cp_world_size: int, + global_sequence_length: int, +) -> ShardingSpec: + """Build a CP/TP sharding spec for one rank of the frozen layout. + + TP splits heads, CP splits the sequence. The #239 rank layout fixes + ``rank = cp_rank * tp_world_size + tp_rank`` for a 2-node x 2-GPU deployment, + but nothing here depends on that mapping: ownership is derived from the ranks + themselves so the same builder serves CP=1 baselines. + """ + + if model.q_heads % tp_world_size or model.kv_heads % tp_world_size: + raise ValueError( + f"Qwen3 GQA heads ({model.q_heads}/{model.kv_heads}) must divide evenly " + f"across tp_world_size={tp_world_size}" + ) + if global_sequence_length % cp_world_size: + raise ValueError( + f"global_sequence_length={global_sequence_length} must divide evenly " + f"across cp_world_size={cp_world_size}" + ) + + local_q_heads = model.q_heads // tp_world_size + local_kv_heads = model.kv_heads // tp_world_size + local_sequence_length = global_sequence_length // cp_world_size + + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=model.q_heads, + global_kv_heads=model.kv_heads, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=global_sequence_length, + local_sequence_length=local_sequence_length, + # One contiguous CP block per rank. The merge order key is the global block + # index, never the arrival order of the CP exchange. + global_block_indices=(cp_rank,), + global_block_token_starts=(cp_rank * local_sequence_length,), + local_block_offsets=(0, local_sequence_length), + ) + + +def causal_offsets_for(sharding: ShardingSpec, batch_size: int) -> tuple[int, ...]: + """Causal offsets for one CP shard, one entry per batch entry. + + Under CP the local query block does not start at global position zero, so the + causal mask has to be shifted by the number of preceding global tokens. Taking + that from ``global_block_token_starts`` rather than recomputing + ``cp_rank * local_sequence_length`` keeps uneven CP splits correct. + """ + + offset = sharding.global_block_token_starts[0] + return (offset,) * batch_size + + +_PROCESS_SCOPES = (IsolationScope.PROCESS, IsolationScope.DISTRIBUTED_CONTEXT) diff --git a/rl_engine/alignment/cross_config/adapters/knobs.py b/rl_engine/alignment/cross_config/adapters/knobs.py new file mode 100644 index 00000000..c1cf0b84 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/knobs.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 attention knobs for the Qwen3-8B TP=2 CP=2 Megatron + vLLM target. + +``V1_KNOBS`` was written against a HuggingFace/FSDP rollout-vs-training pair. Three +of its entries do not survive contact with the frozen Megatron + vLLM target: + +* ``training.sharding`` takes ``unsharded``/``fsdp``, neither of which exists in + Megatron, and is meaningless at DP=1 anyway; +* ``training.attention_backend`` takes HuggingFace names + (``flash_attention_2``/``sdpa``/``eager``/``model_default``) while Megatron's + ``AttnBackend`` is ``flash``/``fused``/``unfused``/``local``/``auto``; +* there is no training-side ``tensor_parallel_size`` or ``context_parallel_size`` + at all, so the target configuration cannot even be expressed. + +This module is deliberately **additive**: it extends ``V1_KNOBS`` rather than +editing it, and overrides only the normalizer for ``training.attention_backend``. +Deleting the two dead knobs changes ``V1_KNOBS`` itself and would break existing +cross-config tests, so it is left to a follow-up on the framework PR. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +from rl_engine.alignment.cross_config.planner import ( + _NORMALIZERS, + V1_KNOBS, + Normalizer, + _normalize_choice, + _positive_int, + _strict_bool, +) +from rl_engine.alignment.cross_config.schema import IsolationScope, KnobDescriptor + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] + + +#: ``megatron.core.transformer.enums.AttnBackend``. +MEGATRON_ATTENTION_BACKENDS: tuple[str, ...] = ( + "flash", + "fused", + "unfused", + "local", + "auto", +) + + +WS2_ATTENTION_KNOB_DESCRIPTORS: tuple[KnobDescriptor, ...] = ( + # -- training-side parallelism: the target configuration itself ------------ + KnobDescriptor( + "training.tensor_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "training.context_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + # -- determinism switches, one per framework ------------------------------ + KnobDescriptor( + "training.deterministic_mode", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "rollout.batch_invariant", + IsolationScope.PROCESS, + ("rollout",), + ), + # -- reduction knobs: the "turn the noise sources on and off" axis --------- + # These are what make drift attributable. ``reduction.order=arrival`` in + # particular is a control group, not a supported production value. + KnobDescriptor( + "attention.reduction_acc_dtype", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("fp32", "bf16"), + ), + KnobDescriptor( + "attention.reduction_order", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("global_block_index", "arrival"), + ), + KnobDescriptor( + "attention.reduction_downcast_at", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("final_write", "per_block"), + ), + KnobDescriptor( + "attention.reduction_engine", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("in_op_reference", "te_oracle"), + ), + # -- materialization knobs: differences the experiment measures ------------ + KnobDescriptor( + "attention.fusion_boundary", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("unfused_rope_attention", "fused_rope_attention"), + ), + KnobDescriptor( + # vLLM: AttentionConfig.flash_attn_max_num_splits_for_cuda_graph + "attention.split_kv_policy", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + KnobDescriptor( + # vLLM: CacheConfig.block_size -> AttentionContract.kv_cache.page_size + "rollout.kv_block_size", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + # The CP communication group cannot be reconfigured once built; it is bound to + # the distributed context, not merely to engine construction. + KnobDescriptor( + "training.cp_comm_type", + IsolationScope.DISTRIBUTED_CONTEXT, + ("training",), + allowed_values=("p2p", "all_gather", "a2a", "a2a+p2p"), + ), +) + + +WS2_ATTENTION_KNOBS: Mapping[str, KnobDescriptor] = { + **V1_KNOBS, + **{descriptor.path: descriptor for descriptor in WS2_ATTENTION_KNOB_DESCRIPTORS}, +} + + +WS2_ATTENTION_NORMALIZERS: Mapping[str, Normalizer] = { + **_NORMALIZERS, + # Replace, not map: the HuggingFace names have no Megatron counterpart. + "training.attention_backend": _normalize_choice(*MEGATRON_ATTENTION_BACKENDS), + "training.tensor_parallel_size": _positive_int, + "training.context_parallel_size": _positive_int, + "training.deterministic_mode": _strict_bool, + "rollout.batch_invariant": _strict_bool, + # AttentionDType values, not torch dtype names -- these feed ReductionSpec directly. + "attention.reduction_acc_dtype": _normalize_choice("fp32", "bf16"), + "attention.reduction_order": _normalize_choice("global_block_index", "arrival"), + "attention.reduction_downcast_at": _normalize_choice("final_write", "per_block"), + "attention.reduction_engine": _normalize_choice("in_op_reference", "te_oracle"), + "attention.fusion_boundary": _normalize_choice( + "unfused_rope_attention", "fused_rope_attention" + ), + "attention.split_kv_policy": _positive_int, + "rollout.kv_block_size": _positive_int, + "training.cp_comm_type": _normalize_choice("p2p", "all_gather", "a2a", "a2a+p2p"), +} diff --git a/rl_engine/alignment/cross_config/adapters/megatron.py b/rl_engine/alignment/cross_config/adapters/megatron.py new file mode 100644 index 00000000..8b96b159 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -0,0 +1,381 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Training-side (Megatron) runtime adapter for WS2 attention cross-config. + +Two things live here: + +``MegatronProvenanceAdapter`` + Read-only. Turns a Megatron config object into the construction and + distributed-context fingerprints the cross-config framework already expects, + plus the determinism probe. It never imports ``megatron`` -- every accessor is + duck-typed -- so this module is importable and testable on a laptop. + +``MegatronAttentionMaterializer`` + Implements the ``RuntimeMaterializer`` protocol. Before this PR the only + implementation in the repository was ``CpuSmokeMaterializer`` over a synthetic + CPU model, so nothing had ever materialized a real distributed runtime. + +Scope boundary: materialization builds and validates the training-side +:class:`AttentionContract` and reports what would be constructed. It does not +launch ``torchrun``, initialize process groups, or execute attention. Binding a +constructed Megatron model to this contract is the next step and needs the 2-node +x 2-GPU environment that #239 fixes. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import ( + DeterminismProbe, + megatron_probe_from_config, +) +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "MEGATRON_CONSTRUCTION_KEYS", + "MEGATRON_DISTRIBUTED_KEYS", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", +] + + +#: ``TransformerConfig`` fields that change attention arithmetic. Hashed into the +#: construction fingerprint. Deliberately excludes MoE, Mamba, MLA and sparse +#: attention fields: the frozen target is Qwen3-8B dense, and those are asserted +#: off rather than recorded. +MEGATRON_CONSTRUCTION_KEYS: tuple[str, ...] = ( + "attention_backend", + "attention_softmax_in_fp32", + "apply_query_key_layer_scaling", + "apply_rope_fusion", + "masked_softmax_fusion", + "bias_activation_fusion", + "bias_dropout_fusion", + "gradient_accumulation_fusion", + "cross_entropy_loss_fusion", + "cross_entropy_fusion_impl", + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + "recompute_modules", + "rotary_base", + "rotary_percent", + "rotary_interleaved", + "rotary_scaling_factor", + "qk_layernorm", + "hidden_dropout", + "attention_dropout", + "params_dtype", + "bf16", + "fp16", + "fp8", + "deterministic_mode", +) + + +#: ``ModelParallelConfig`` fields that define the distributed context. +MEGATRON_DISTRIBUTED_KEYS: tuple[str, ...] = ( + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "virtual_pipeline_model_parallel_size", + "context_parallel_size", + "hierarchical_context_parallel_sizes", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "sequence_parallel", + "cp_comm_type", + "tp_comm_overlap", + "use_te_rng_tracker", +) + + +#: Fields that must hold these values for the frozen dense target. A mismatch is a +#: hard stop, not a recorded difference -- see the exclusion list in the WS2 scope. +MEGATRON_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "pipeline_model_parallel_size": 1, + "expert_model_parallel_size": 1, + "sequence_parallel": False, + "fp8": None, + "hidden_dropout": 0.0, + "attention_dropout": 0.0, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class MegatronProvenanceAdapter: + """Extract fingerprints and determinism evidence from a Megatron config. + + ``config`` may be a real ``TransformerConfig``/``ModelParallelConfig``, a merged + namespace, or a test double. Missing attributes read as ``None`` and are + recorded as such rather than raising: an absent field is itself provenance. + """ + + framework = "megatron" + + def __init__(self, config: Any, *, env: Optional[Mapping[str, str]] = None): + self.config = config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_CONSTRUCTION_KEYS} + + def distributed_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_DISTRIBUTED_KEYS} + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return megatron_probe_from_config(self.config, env=self.env) + + def frozen_scope_violations(self) -> tuple[str, ...]: + """Return the frozen-scope assertions this config violates.""" + + violations: list[str] = [] + for name, expected in MEGATRON_FROZEN_ASSERTIONS.items(): + actual = _value(self.config, name) + if actual is None: + # Not declared. Treated as unknown rather than as satisfied, because + # a silently-absent MoE or FP8 setting is exactly the case that would + # otherwise slip past a dense-only claim. + violations.append(f"{name} is not declared (expected {expected!r})") + elif actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "determinism": self.determinism_probe().to_dict(), + } + + +class MegatronAttentionMaterializer: + """Materialize the training-side attention runtime for the WS2 target.""" + + runtime_kind = "megatron_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + backend_id: str = "rlkernel.cp_attention_reference", + provenance: Optional[MegatronProvenanceAdapter] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.backend_id = backend_id + self.provenance = provenance + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + """Build the training-side contract. Raises on an unusable request.""" + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=attention_dtype( + flat.get("training.compute_dtype", "bf16"), field="training.compute_dtype" + ), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "training" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + applications.append( + application( + descriptor, + requested, + requested, + requested, + MaterializationStatus.APPLIED, + "bound to the training-side attention contract", + frozen_scope_violations=list(scope_violations), + ) + ) + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "megatron", + "attention_backend": flat.get("training.attention_backend"), + "compute_dtype": flat.get("training.compute_dtype"), + "deterministic_mode": flat.get("training.deterministic_mode"), + "cp_comm_type": flat.get("training.cp_comm_type"), + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"training": side_config, "rollout": {}}, + topology={ + "training": { + "world_size": tp_world_size * cp_world_size, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": cp_world_size, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "rollout": {"world_size": 1}, + }, + scorer={ + "mode": "teacher_forcing", + "framework": "megatron", + "export_lse": True, + }, + operator_backends={ + "training": self.backend_id, + "rollout": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) diff --git a/rl_engine/alignment/cross_config/adapters/vllm.py b/rl_engine/alignment/cross_config/adapters/vllm.py new file mode 100644 index 00000000..a8de61c9 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -0,0 +1,442 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rollout-side (vLLM) runtime adapter for WS2 attention cross-config. + +Mirrors :mod:`.megatron`, with three differences that come straight from what vLLM +actually is: + +* vLLM's context parallelism is ``prefill_context_parallel_size`` -- it applies to + prefill only, so a decode-mode contract must declare ``cp_world_size == 1`` + regardless of what the prefill knob says. +* ``CacheConfig.block_size`` is the paged-KV page size, and it feeds + ``KVCacheSpec.page_size`` directly rather than being invented here. +* Determinism comes from the ``VLLM_BATCH_INVARIANT`` environment variable rather + than from a config field, because vLLM applies it inside + ``init_batch_invariance()`` at worker startup. + +Like the Megatron adapter, nothing here imports ``vllm``; configs are duck-typed so +the module is importable anywhere. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import DeterminismProbe, vllm_probe_from_env +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "VLLM_ATTENTION_KEYS", + "VLLM_CACHE_KEYS", + "VLLM_FROZEN_ASSERTIONS", + "VLLM_MODEL_KEYS", + "VLLM_PARALLEL_KEYS", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", +] + + +VLLM_MODEL_KEYS: tuple[str, ...] = ( + "dtype", + "seed", + "quantization", + "enforce_eager", + "max_logprobs", + "disable_cascade_attn", + "max_model_len", +) + +VLLM_CACHE_KEYS: tuple[str, ...] = ( + "block_size", + "cache_dtype", + "enable_prefix_caching", + "prefix_caching_hash_algo", + "calculate_kv_scales", + "sliding_window", +) + +VLLM_ATTENTION_KEYS: tuple[str, ...] = ( + "backend", + "flash_attn_version", + "use_prefill_decode_attention", + "flash_attn_max_num_splits_for_cuda_graph", + "use_cudnn_prefill", + "disable_flashinfer_prefill", + "use_non_causal", +) + +VLLM_PARALLEL_KEYS: tuple[str, ...] = ( + "tensor_parallel_size", + "pipeline_parallel_size", + "prefill_context_parallel_size", + "data_parallel_size", +) + + +#: Frozen dense-target assertions on the rollout side. ``cache_dtype`` must stay +#: ``auto`` because an FP8 KV cache is a representation-drift problem tracked +#: separately, and ``disable_cascade_attn`` must stay ``True`` because cascade +#: attention changes the block-merge structure the contract pins down. +VLLM_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "quantization": None, + "cache_dtype": "auto", + "calculate_kv_scales": False, + "disable_cascade_attn": True, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "sliding_window": None, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class VllmProvenanceAdapter: + """Extract fingerprints and determinism evidence from vLLM configs.""" + + framework = "vllm" + + def __init__( + self, + *, + model_config: Any = None, + cache_config: Any = None, + attention_config: Any = None, + parallel_config: Any = None, + env: Optional[Mapping[str, str]] = None, + ): + self.model_config = model_config + self.cache_config = cache_config + self.attention_config = attention_config + self.parallel_config = parallel_config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + view: dict[str, Any] = {} + for prefix, config, keys in ( + ("model", self.model_config, VLLM_MODEL_KEYS), + ("cache", self.cache_config, VLLM_CACHE_KEYS), + ("attention", self.attention_config, VLLM_ATTENTION_KEYS), + ): + for name in keys: + view[f"{prefix}.{name}"] = _value(config, name) + return view + + def distributed_view(self) -> dict[str, Any]: + return { + f"parallel.{name}": _value(self.parallel_config, name) for name in VLLM_PARALLEL_KEYS + } + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return vllm_probe_from_env(self.env, model_config=self.model_config) + + def frozen_scope_violations(self) -> tuple[str, ...]: + sources = { + "quantization": self.model_config, + "disable_cascade_attn": self.model_config, + "cache_dtype": self.cache_config, + "calculate_kv_scales": self.cache_config, + "sliding_window": self.cache_config, + "pipeline_parallel_size": self.parallel_config, + "data_parallel_size": self.parallel_config, + } + violations: list[str] = [] + for name, expected in VLLM_FROZEN_ASSERTIONS.items(): + config = sources.get(name) + if config is None: + continue + actual = _value(config, name) + if actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + @property + def kv_page_size(self) -> Optional[int]: + """vLLM's paged-KV block size, which is the contract's ``page_size``.""" + + block_size = _value(self.cache_config, "block_size") + return int(block_size) if block_size is not None else None + + @property + def split_kv_policy(self) -> Optional[int]: + """The split-KV knob #235 PR5/PR7 needs; #236 has no field for it yet.""" + + splits = _value(self.attention_config, "flash_attn_max_num_splits_for_cuda_graph") + return int(splits) if splits is not None else None + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "kv_page_size": self.kv_page_size, + "split_kv_policy": self.split_kv_policy, + "determinism": self.determinism_probe().to_dict(), + } + + +class VllmRolloutMaterializer: + """Materialize the rollout-side attention runtime for the WS2 target.""" + + runtime_kind = "vllm_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + mode: AttentionMode = AttentionMode.CHUNKED_PREFILL, + backend_id: str = "vllm.flash_attn", + provenance: Optional[VllmProvenanceAdapter] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.mode = mode + self.backend_id = backend_id + self.provenance = provenance + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def effective_cp_world_size(self, flat: Mapping[str, Any]) -> int: + """CP applies to prefill only; decode always runs at CP=1.""" + + requested = int(flat.get("rollout.context_parallel_size", 1)) + if self.mode is AttentionMode.DECODE: + return 1 + return requested + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + if self.mode is AttentionMode.DECODE: + # Decode replay needs a validated KVCacheSpec (cache positions, page + # ownership, prefix-cache identity). That is #235 PR6's contract surface, + # and inventing a placeholder here would let an unvalidated decode case + # look bound. Fail instead. + raise AttentionContractError( + "decode-mode materialization requires KV-cache identity from #235 PR6; " + "this adapter covers prefill and chunked prefill" + ) + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + cp_world_size = self.effective_cp_world_size(flat) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank if cp_world_size > 1 else 0, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.FUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + # vLLM stores post-RoPE K in the cache; recorded, not asserted equal to + # the training side, because it is a materialization fact. + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.INFER, + mode=self.mode, + dtype=attention_dtype(flat.get("rollout.dtype", "bf16"), field="rollout.dtype"), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + requested_cp = int(flat.get("rollout.context_parallel_size", 1)) + effective_cp = self.effective_cp_world_size(flat) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "rollout" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + if path == "rollout.context_parallel_size" and effective_cp != requested_cp: + applications.append( + application( + descriptor, + requested, + effective_cp, + effective_cp, + MaterializationStatus.FALLBACK, + ( + "vLLM context parallelism covers prefill only; a decode-mode " + f"contract runs at cp_world_size=1, not {requested_cp}" + ), + vllm_field="ParallelConfig.prefill_context_parallel_size", + ) + ) + continue + applications.append( + application( + descriptor, + requested, + requested, + requested, + MaterializationStatus.APPLIED, + "bound to the rollout-side attention contract", + frozen_scope_violations=list(scope_violations), + ) + ) + + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "vllm", + "dtype": flat.get("rollout.dtype"), + "enforce_eager": flat.get("rollout.enforce_eager"), + "enable_prefix_caching": flat.get("rollout.enable_prefix_caching"), + "batch_invariant": flat.get("rollout.batch_invariant"), + "kv_block_size": flat.get("rollout.kv_block_size"), + "split_kv_policy": flat.get("attention.split_kv_policy"), + "attention_mode": self.mode.value, + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"rollout": side_config, "training": {}}, + topology={ + "rollout": { + "world_size": tp_world_size * effective_cp, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": effective_cp, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "training": {"world_size": 1}, + }, + scorer={ + "mode": "rollout_logprob", + "framework": "vllm", + "export_lse": True, + }, + operator_backends={ + "rollout": self.backend_id, + "training": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py new file mode 100644 index 00000000..48e382b8 --- /dev/null +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -0,0 +1,577 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Three-tier binding between rollout-side and training-side attention contracts. + +Issue #235 PR4 requires that "rollout and training descriptors bind to the same +semantic attention contract". Under the frozen Megatron + vLLM deployment the two +sides can never produce *identical* :class:`AttentionContract` instances: training +runs full-sequence prefill over a CP-sharded sequence, while rollout runs vLLM +paged-KV chunked prefill and decode. Taking "same contract" literally would make +the target configuration permanently unbindable. + +This module therefore splits binding into three tiers: + +``IDENTICAL`` + Logical identity. Both sides must agree bit for bit, otherwise the pair is not + comparable at all and no drift number from it means anything. + +``SEMANTIC`` + The WS2 numerical claim: merge semantics, accumulation dtype, reduction order + and downcast point are decided by the contract, not by the implementation. + Both sides must carry the same values *and* those values must match the WS2 + mandate, otherwise the comparison fails closed. + +``RECORDED`` + Materialization facts that the two sides are expected to differ on -- attention + mode, RoPE fusion boundary, KV-cache paging, backend id, reduction engine. These + differences are exactly what the experiment measures, so they are recorded into + provenance rather than rejected. + +Deliberately *not* in ``SEMANTIC``: ``engine``. Training may run the in-op +deterministic reference while rollout runs a Transformer Engine merge oracle; forcing +those equal would defeat the purpose of the oracle comparison in #235 PR2/3/5/6. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Optional + +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionDType, + AttentionMerge, + AttentionRole, + DowncastPoint, + ReductionOrder, +) + +__all__ = [ + "ATTENTION_LSE_DOMAIN", + "AttentionBindingError", + "AttentionBindingResult", + "BindingErrorCode", + "BindingIssue", + "BindingTier", + "IDENTITY_FIELDS", + "NULLABLE_IDENTITY_FIELDS", + "RECORDED_FIELDS", + "SEMANTIC_CONTRACT_FIELDS", + "SEMANTIC_REDUCTION_FIELDS", + "WS2_ATTENTION_REDUCTION_MANDATE", + "bind_attention_contracts", + "first_blocking_issue", + "identity_fingerprint", + "summarize_binding", +] + + +class AttentionBindingError(ValueError): + """Raised when a caller supplies structurally unusable binding inputs.""" + + +class BindingTier(str, Enum): + """Which rule a field is governed by.""" + + IDENTICAL = "identical" + SEMANTIC = "semantic" + RECORDED = "recorded" + + +class BindingErrorCode(str, Enum): + """Stable, machine-readable reasons a binding is rejected. + + Callers branch on these; they are part of the artifact schema and must not be + renamed without a schema version bump. + """ + + IDENTITY_MISSING = "IDENTITY_MISSING" + IDENTITY_MISMATCH = "IDENTITY_MISMATCH" + REDUCTION_SEMANTIC_MISMATCH = "REDUCTION_SEMANTIC_MISMATCH" + REDUCTION_MANDATE_VIOLATION = "REDUCTION_MANDATE_VIOLATION" + LSE_NOT_EXPORTED = "LSE_NOT_EXPORTED" + ROLE_COLLISION = "ROLE_COLLISION" + DETERMINISM_INCOMPATIBLE = "DETERMINISM_INCOMPATIBLE" + + +#: Attention exports attention-domain LSE, never vocab-logprob LSE (#235). +#: Recorded explicitly so a future ``LogprobContract`` binding cannot be confused +#: with this one purely because both set ``export_lse=True``. +ATTENTION_LSE_DOMAIN = "attention" + + +#: Fields both sides must agree on bit for bit before any comparison is meaningful. +#: Sourced from #235 "Numerical Contract" preconditions plus the vime-owned rollout +#: provenance (weight version, sampling, padding) that the issue assumes but does +#: not enumerate. +IDENTITY_FIELDS: tuple[str, ...] = ( + "checkpoint_id", + "model_version", + "weight_version", + "tokenizer_fingerprint", + "token_ids_fingerprint", + "active_mask_fingerprint", + "position_ids_fingerprint", + "padding_side", + "pre_update_state", + # model semantics that decide what attention *means* + "q_heads", + "kv_heads", + "head_dim", + "rope_theta", + "rope_scaling", + "rotary_dim", + "qk_layernorm", + # 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",) + + +#: The WS2 mandate itself. ``#236`` currently declares single-member enums for +#: ``merge`` / ``order`` / ``downcast_at``, so those checks are tautological today; +#: they are written out anyway so that widening any of those enums later fails here +#: instead of silently admitting a non-conforming backend. +WS2_ATTENTION_REDUCTION_MANDATE: Mapping[str, str] = { + "merge": AttentionMerge.ONLINE_SOFTMAX_LSE.value, + "acc_dtype": AttentionDType.FP32.value, + "order": ReductionOrder.GLOBAL_BLOCK_INDEX.value, + "downcast_at": DowncastPoint.FINAL_WRITE.value, +} + + +#: Materialization facts the two sides are expected to differ on. Recorded into +#: provenance; never a rejection reason. +RECORDED_FIELDS: tuple[str, ...] = ( + "mode", + "backend_id", + "reduction.engine", + "rope.fusion_boundary", + "rope.q_state", + "rope.k_state", + "rope.k_cache_state", + "rope.cast_at", + "rope.output_dtype", + "kv_cache.page_size", + "kv_cache.prefix_cache_enabled", + "kv_cache.block_table_shape", + "sharding.cp_world_size", + "sharding.tp_world_size", + "sharding.local_sequence_length", + # Supplied by the caller, not by the contract: #236 has no split-KV field yet, so + # the value comes from vLLM's flash_attn_max_num_splits_for_cuda_graph via the + # adapter. Recorded so split-KV differences are at least visible in provenance + # until #236 grows the field and it can move into the contract proper. + "split_kv_policy", +) + + +@dataclass(frozen=True) +class BindingIssue: + """One reason a binding is not comparable or not admissible.""" + + code: BindingErrorCode + tier: BindingTier + field: str + rollout: Any = None + training: Any = None + message: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "code": self.code.value, + "tier": self.tier.value, + "field": self.field, + "rollout": self.rollout, + "training": self.training, + "message": self.message, + } + + +@dataclass(frozen=True) +class AttentionBindingResult: + """Outcome of binding one rollout contract to one training contract. + + ``comparable`` and ``passed`` are deliberately separate. A pair whose identity + does not match is *not comparable* -- reporting a drift number for it would be + meaningless. A pair that is comparable but violates the reduction mandate *is* + comparable yet must still fail closed, because the whole WS2 claim is that + reduction order and accumulation precision come from the contract. + """ + + comparable: bool + passed: bool + issues: tuple[BindingIssue, ...] = () + identity_fingerprint: str = "" + reduction_fingerprint: str = "" + binding_fingerprint: str = "" + recorded_differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + provenance: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.attention_binding.v1" + + def issues_by_code(self, code: BindingErrorCode) -> tuple[BindingIssue, ...]: + return tuple(issue for issue in self.issues if issue.code is code) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "comparable": self.comparable, + "passed": self.passed, + "issues": [issue.to_dict() for issue in self.issues], + "identity_fingerprint": self.identity_fingerprint, + "reduction_fingerprint": self.reduction_fingerprint, + "binding_fingerprint": self.binding_fingerprint, + "recorded_differences": { + key: dict(value) for key, value in self.recorded_differences.items() + }, + "provenance": dict(self.provenance), + } + + +def _canonical_fingerprint(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def identity_fingerprint(identity: Mapping[str, Any]) -> str: + """Fingerprint only the declared :data:`IDENTITY_FIELDS`, in a fixed order. + + Extra keys in ``identity`` are ignored on purpose: callers pass whole + provenance bundles, and the fingerprint must not drift when an unrelated + diagnostic field is added. + """ + + return _canonical_fingerprint({name: identity.get(name) for name in IDENTITY_FIELDS}) + + +def _reduction_view(contract: AttentionContract) -> dict[str, Any]: + reduction = contract.reduction + return { + "merge": reduction.merge.value, + "acc_dtype": reduction.acc_dtype.value, + "order": reduction.order.value, + "downcast_at": reduction.downcast_at.value, + "engine": reduction.engine.value, + } + + +def _recorded_view( + contract: AttentionContract, + 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, + "sharding.cp_world_size": contract.sharding.cp_world_size, + "sharding.tp_world_size": contract.sharding.tp_world_size, + "sharding.local_sequence_length": contract.sharding.local_sequence_length, + } + if rope is not None: + view.update( + { + "rope.fusion_boundary": rope.fusion_boundary.value, + "rope.q_state": rope.q_state.value, + "rope.k_state": rope.k_state.value, + "rope.k_cache_state": rope.k_cache_state.value, + "rope.cast_at": rope.cast_at.value, + "rope.output_dtype": rope.output_dtype.value, + } + ) + if kv_cache is not None: + view.update( + { + "kv_cache.page_size": kv_cache.page_size, + "kv_cache.prefix_cache_enabled": kv_cache.prefix_cache_enabled, + "kv_cache.block_table_shape": [ + len(kv_cache.block_table), + max((len(row) for row in kv_cache.block_table), default=0), + ], + } + ) + if extra: + view.update(extra) + return view + + +#: Identity fields where ``None`` is a real value rather than an omission. Qwen3-8B +#: applies no RoPE scaling, so ``rope_scaling=None`` must not read as "undeclared" -- +#: both sides still have to agree on it, which the equality pass below handles. +NULLABLE_IDENTITY_FIELDS: frozenset[str] = frozenset({"rope_scaling"}) + + +def _missing_identity_fields(identity: Mapping[str, Any]) -> tuple[str, ...]: + return tuple( + name + for name in IDENTITY_FIELDS + if name not in NULLABLE_IDENTITY_FIELDS and identity.get(name) is None + ) + + +def bind_attention_contracts( + *, + rollout_contract: AttentionContract, + training_contract: AttentionContract, + rollout_identity: Mapping[str, Any], + training_identity: Mapping[str, Any], + rollout_backend_id: str, + training_backend_id: str, + determinism_issues: Sequence[BindingIssue] = (), + require_full_identity: bool = True, + allow_dtype_difference: bool = False, + rollout_recorded_extra: Optional[Mapping[str, Any]] = None, + training_recorded_extra: Optional[Mapping[str, Any]] = None, +) -> AttentionBindingResult: + """Bind a rollout attention contract to a training attention contract. + + ``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. + + ``rollout_recorded_extra`` / ``training_recorded_extra`` carry materialization + facts that #236 does not yet model -- today that is ``split_kv_policy``. They are + merged into the recorded tier, never into identity or semantics. + """ + + if rollout_contract.role is not AttentionRole.INFER: + raise AttentionBindingError( + f"rollout_contract.role must be {AttentionRole.INFER.value!r}, " + f"got {rollout_contract.role.value!r}" + ) + if training_contract.role is not AttentionRole.TRAIN: + raise AttentionBindingError( + f"training_contract.role must be {AttentionRole.TRAIN.value!r}, " + f"got {training_contract.role.value!r}" + ) + + issues: list[BindingIssue] = [] + + # ---- tier 1: identity, bit for bit ------------------------------------- + if require_full_identity: + for side, identity in (("rollout", rollout_identity), ("training", training_identity)): + for name in _missing_identity_fields(identity): + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISSING, + tier=BindingTier.IDENTICAL, + field=f"{side}.{name}", + message=f"{side} identity does not declare {name!r}", + ) + ) + + for name in IDENTITY_FIELDS: + rollout_value = rollout_identity.get(name) + training_value = training_identity.get(name) + if rollout_value != training_value: + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISMATCH, + tier=BindingTier.IDENTICAL, + field=name, + rollout=rollout_value, + training=training_value, + message=( + f"{name!r} differs between sides; the pair is not comparable " + "and any drift computed from it is meaningless" + ), + ) + ) + + comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + + # ---- tier 2: reduction semantics, and the WS2 mandate ------------------- + rollout_reduction = _reduction_view(rollout_contract) + training_reduction = _reduction_view(training_contract) + + for name in SEMANTIC_REDUCTION_FIELDS: + if rollout_reduction[name] != training_reduction[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"reduction.{name!r} must be decided by the contract, not by the " + "backend; the two sides disagree" + ), + ) + ) + mandated = WS2_ATTENTION_REDUCTION_MANDATE[name] + for side, view in (("rollout", rollout_reduction), ("training", training_reduction)): + if view[name] != mandated: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_MANDATE_VIOLATION, + tier=BindingTier.SEMANTIC, + field=f"{side}.reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"WS2 requires reduction.{name} == {mandated!r}; " + f"{side} declares {view[name]!r}" + ), + ) + ) + + if not allow_dtype_difference and rollout_contract.dtype is not training_contract.dtype: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field="dtype", + rollout=rollout_contract.dtype.value, + training=training_contract.dtype.value, + message=( + "the two sides compute in different dtypes; the resulting drift is " + "not attributable. Pass allow_dtype_difference=True only for a " + "deliberate precision sweep" + ), + ) + ) + + for side, contract in (("rollout", rollout_contract), ("training", training_contract)): + if not contract.export_lse: + issues.append( + 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, + "rollout": { + "contract": rollout_contract.to_dict(), + "backend_id": rollout_backend_id, + "recorded": rollout_recorded, + }, + "training": { + "contract": training_contract.to_dict(), + "backend_id": training_backend_id, + "recorded": training_recorded, + }, + } + + return AttentionBindingResult( + comparable=comparable, + passed=passed, + issues=tuple(issues), + identity_fingerprint=identity_fp, + reduction_fingerprint=reduction_fp, + binding_fingerprint=_canonical_fingerprint( + { + "identity": identity_fp, + "reduction": reduction_fp, + "lse_domain": ATTENTION_LSE_DOMAIN, + "rollout_backend": rollout_backend_id, + "training_backend": training_backend_id, + } + ), + recorded_differences=recorded_differences, + provenance=provenance, + ) + + +def summarize_binding(result: AttentionBindingResult) -> str: + """One-line human summary for CLI output and failure messages.""" + + if result.passed: + return ( + f"attention binding OK " + f"(identity={result.identity_fingerprint[:12]}, " + f"{len(result.recorded_differences)} recorded difference(s))" + ) + if not result.comparable: + fields = ", ".join( + sorted({issue.field for issue in result.issues if issue.tier is BindingTier.IDENTICAL}) + ) + return f"attention binding NOT COMPARABLE; identity problems: {fields}" + fields = ", ".join( + sorted({issue.field for issue in result.issues if issue.tier is BindingTier.SEMANTIC}) + ) + return f"attention binding FAILED CLOSED; semantic problems: {fields}" + + +def first_blocking_issue( + result: AttentionBindingResult, +) -> Optional[BindingIssue]: + """Return the issue a caller should report, preferring identity over semantics.""" + + for tier in (BindingTier.IDENTICAL, BindingTier.SEMANTIC): + for issue in result.issues: + if issue.tier is tier: + return issue + return None diff --git a/rl_engine/alignment/cross_config/determinism.py b/rl_engine/alignment/cross_config/determinism.py new file mode 100644 index 00000000..be81654e --- /dev/null +++ b/rl_engine/alignment/cross_config/determinism.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-side determinism probing for the Megatron + vLLM cross-config target. + +Both frameworks ship a "make this deterministic" switch, but they mean different +things by it, and neither knows the other exists: + +``Megatron`` ``ModelParallelConfig.deterministic_mode`` + Asserts ``NCCL_ALGO`` is one of five values, forbids FlashAttention and fused + cross-entropy, calls ``torch.use_deterministic_algorithms(True)``, and requires + ``NVTE_ALLOW_NONDETERMINISTIC_ALGO == 0``. It does **not** touch TF32, BF16 + reduced-precision reduction, cuBLAS workspace, NCCL protocol, or NCCL channel + counts. + +``vLLM`` ``VLLM_BATCH_INVARIANT`` + Replaces ``aten::mm/addmm/matmul/linear/bmm``, ``log_softmax``/``softmax``, + ``mean.dim`` and ``rms_norm`` with Triton kernels, disables TF32 and BF16/FP16 + reduced-precision reduction, pins cuBLAS workspace and the BLAS library, and + hard-sets ten NCCL environment variables. + +So a run can have both switches on and still be comparing two different notions of +determinism. This module makes that difference explicit and, where it changes the +numerics, blocking. It never imports Megatron or vLLM: probes are built from plain +mappings so the logic is testable on any machine. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, Optional + +from rl_engine.alignment.cross_config.attention_binding import ( + BindingErrorCode, + BindingIssue, + BindingTier, +) + +__all__ = [ + "COMPARED_NCCL_KEYS", + "DeterminismProbe", + "DeterminismReport", + "compare_determinism", + "megatron_probe_from_config", + "vllm_probe_from_env", +] + + +#: Environment keys whose value can change a reduction result. Compared across +#: sides; a difference is reported, and a difference in the *arithmetic* subset is +#: blocking. Ordering is fixed so the fingerprint is stable. +COMPARED_NCCL_KEYS: tuple[str, ...] = ( + "NCCL_ALGO", + "NCCL_PROTO", + "NCCL_MIN_NCHANNELS", + "NCCL_MAX_NCHANNELS", + "NCCL_NTHREADS", + "NCCL_SOCKET_NTHREADS", + "NCCL_COLLNET_ENABLE", + "NCCL_NVLS_ENABLE", + "NCCL_P2P_NET_DISABLE", + "NCCL_LAUNCH_MODE", + "CUBLAS_WORKSPACE_CONFIG", +) + + +#: The subset above that changes arithmetic rather than only scheduling. A mismatch +#: here fails the binding closed; a mismatch in the remainder is recorded only. +_ARITHMETIC_NCCL_KEYS: frozenset[str] = frozenset( + {"NCCL_ALGO", "NCCL_PROTO", "CUBLAS_WORKSPACE_CONFIG"} +) + + +@dataclass(frozen=True) +class DeterminismProbe: + """What one side actually has switched on. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are tri-state on + purpose: ``None`` means "the framework does not manage this", which is exactly + Megatron's situation and is itself the finding. + """ + + side: str + framework: str + mode_flag: str + enabled: bool + env: Mapping[str, Any] = field(default_factory=dict) + tf32_disabled: Optional[bool] = None + bf16_reduced_precision_reduction: Optional[bool] = None + forbids_flash_attention: Optional[bool] = None + evidence: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_probe.v1" + + def __post_init__(self) -> None: + if self.side not in ("rollout", "training"): + raise ValueError("side must be 'rollout' or 'training'") + if not self.framework: + raise ValueError("framework must not be empty") + object.__setattr__(self, "env", dict(self.env)) + object.__setattr__(self, "evidence", dict(self.evidence)) + + @property + def env_fingerprint(self) -> str: + payload = {key: self.env.get(key) for key in COMPARED_NCCL_KEYS} + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "side": self.side, + "framework": self.framework, + "mode_flag": self.mode_flag, + "enabled": self.enabled, + "env": {key: self.env.get(key) for key in COMPARED_NCCL_KEYS}, + "env_fingerprint": self.env_fingerprint, + "tf32_disabled": self.tf32_disabled, + "bf16_reduced_precision_reduction": self.bf16_reduced_precision_reduction, + "forbids_flash_attention": self.forbids_flash_attention, + "evidence": dict(self.evidence), + } + + +@dataclass(frozen=True) +class DeterminismReport: + """Cross-side comparison result.""" + + rollout: DeterminismProbe + training: DeterminismProbe + issues: tuple[BindingIssue, ...] = () + differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_report.v1" + + @property + def compatible(self) -> bool: + return not self.issues + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "compatible": self.compatible, + "rollout": self.rollout.to_dict(), + "training": self.training.to_dict(), + "issues": [issue.to_dict() for issue in self.issues], + "differences": {key: dict(value) for key, value in self.differences.items()}, + } + + +def megatron_probe_from_config( + config: Any, + env: Optional[Mapping[str, str]] = None, +) -> DeterminismProbe: + """Build a training-side probe from a Megatron config object. + + ``config`` is duck-typed (anything exposing ``deterministic_mode`` and + optionally ``attention_backend`` / ``cross_entropy_loss_fusion``) so this works + against a real ``ModelParallelConfig``, a test double, or a plain namespace, + and so importing this module never requires Megatron. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are reported as + ``None`` because Megatron does not manage them -- a ``grep`` for ``allow_tf32`` + and ``fp32_precision`` across ``megatron/`` returns nothing. That asymmetry + against vLLM is the point of :func:`compare_determinism`. + """ + + environ = dict(env or {}) + enabled = bool(getattr(config, "deterministic_mode", False)) + return DeterminismProbe( + side="training", + framework="megatron", + mode_flag="deterministic_mode", + enabled=enabled, + env={key: environ.get(key) for key in COMPARED_NCCL_KEYS}, + tf32_disabled=None, + bf16_reduced_precision_reduction=None, + forbids_flash_attention=enabled, + evidence={ + "nvte_allow_nondeterministic_algo": environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO"), + "cross_entropy_loss_fusion": getattr(config, "cross_entropy_loss_fusion", None), + "attention_backend": _enum_value(getattr(config, "attention_backend", None)), + "tensor_model_parallel_size": getattr(config, "tensor_model_parallel_size", None), + "context_parallel_size": getattr(config, "context_parallel_size", None), + "sequence_parallel": getattr(config, "sequence_parallel", None), + "manages_tf32": False, + "manages_bf16_reduced_precision_reduction": False, + }, + ) + + +def vllm_probe_from_env( + env: Mapping[str, str], + *, + model_config: Any = None, +) -> DeterminismProbe: + """Build a rollout-side probe from the vLLM process environment. + + ``VLLM_BATCH_INVARIANT`` is read from ``env`` rather than ``vllm.envs`` so the + probe can be constructed from a remote worker's reported environment, which is + how vime's Ray actors expose it. + """ + + enabled = str(env.get("VLLM_BATCH_INVARIANT", "0")).strip() in ("1", "true", "True") + return DeterminismProbe( + side="rollout", + framework="vllm", + mode_flag="VLLM_BATCH_INVARIANT", + enabled=enabled, + env={key: env.get(key) for key in COMPARED_NCCL_KEYS}, + # vLLM sets both to "ieee"/disabled inside init_batch_invariance(). + tf32_disabled=enabled or None, + bf16_reduced_precision_reduction=(False if enabled else None), + forbids_flash_attention=False, + evidence={ + "vllm_allreduce_use_symm_mem": env.get("VLLM_ALLREDUCE_USE_SYMM_MEM"), + "vllm_use_aot_compile": env.get("VLLM_USE_AOT_COMPILE"), + "enforce_eager": getattr(model_config, "enforce_eager", None), + "disable_cascade_attn": getattr(model_config, "disable_cascade_attn", None), + "quantization": getattr(model_config, "quantization", None), + "manages_tf32": True, + "manages_bf16_reduced_precision_reduction": True, + }, + ) + + +def _enum_value(value: Any) -> Any: + return getattr(value, "value", value) + + +def compare_determinism( + *, + rollout: DeterminismProbe, + training: DeterminismProbe, +) -> DeterminismReport: + """Compare two probes and produce blocking issues plus recorded differences.""" + + if rollout.side != "rollout" or training.side != "training": + raise ValueError("compare_determinism expects one rollout probe and one training probe") + + issues: list[BindingIssue] = [] + differences: dict[str, dict[str, Any]] = {} + + for probe in (rollout, training): + if not probe.enabled: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"{probe.side}.{probe.mode_flag}", + rollout=rollout.enabled, + training=training.enabled, + message=( + f"{probe.framework} {probe.mode_flag} is not enabled; the " + f"{probe.side} side is not batch-invariant and cannot anchor a " + "cross-config comparison" + ), + ) + ) + + for key in COMPARED_NCCL_KEYS: + rollout_value = rollout.env.get(key) + training_value = training.env.get(key) + if rollout_value == training_value: + continue + differences[key] = {"rollout": rollout_value, "training": training_value} + if key in _ARITHMETIC_NCCL_KEYS: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"env.{key}", + rollout=rollout_value, + training=training_value, + message=( + f"{key} differs between sides; the two sides would reduce with " + "different arithmetic and the resulting drift is not attributable" + ), + ) + ) + + # Megatron reports None for these because it does not manage them at all. That is + # recorded rather than blocking: under a pure BF16 GEMM path TF32 does not fire, and + # forcing Megatron to manage it is out of scope for this PR. It is surfaced so the + # asymmetry appears in every artifact instead of being invisible. + for name in ("tf32_disabled", "bf16_reduced_precision_reduction"): + rollout_value = getattr(rollout, name) + training_value = getattr(training, name) + if rollout_value != training_value: + differences[name] = { + "rollout": rollout_value, + "training": training_value, + "note": ( + "megatron does not manage this setting; vllm sets it inside " + "init_batch_invariance()" + ), + } + + return DeterminismReport( + rollout=rollout, + training=training, + issues=tuple(issues), + differences=differences, + ) diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py new file mode 100644 index 00000000..7deeddf0 --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,865 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for context-parallel standard softmax attention. + +The objects in this module describe a distributed attention invocation. They +do not shard tensors, launch collectives, or implement the ``(out, lse)`` +merge. Keeping description and materialization separate lets dispatch reject +an incompatible backend before any numerically different path is launched. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Iterable, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +class AttentionContractError(ValueError): + """Raised when attention metadata does not describe a valid invocation.""" + + +class AttentionRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class AttentionMode(str, Enum): + PREFILL = "prefill" + CHUNKED_PREFILL = "chunked_prefill" + DECODE = "decode" + + +class AttentionDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class AttentionMerge(str, Enum): + ONLINE_SOFTMAX_LSE = "online_softmax_lse" + + +class ReductionOrder(str, Enum): + GLOBAL_BLOCK_INDEX = "global_block_index" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +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 KVCacheSpec: + """Logical identity of the paged/block KV cache used for replay.""" + + cache_positions: tuple[int, ...] + kv_seq_lens: tuple[int, ...] + block_table: tuple[tuple[int, ...], ...] + global_token_positions: tuple[int, ...] + page_size: int + prefix_cache_enabled: bool = False + prefix_cache_key: str | None = None + shared_prefix_page_count: int = 0 + + def __post_init__(self) -> None: + cache_positions = _integer_tuple(self.cache_positions, "cache_positions") + kv_seq_lens = _integer_tuple(self.kv_seq_lens, "kv_seq_lens") + page_size = _positive_int(self.page_size, "page_size") + global_token_positions = _integer_tuple( + self.global_token_positions, "global_token_positions" + ) + if not cache_positions or any(position < 0 for position in cache_positions): + raise AttentionContractError("cache_positions must contain non-negative positions") + if not kv_seq_lens or any(length <= 0 for length in kv_seq_lens): + raise AttentionContractError("kv_seq_lens must contain positive sequence lengths") + if len(cache_positions) != len(kv_seq_lens): + raise AttentionContractError( + "cache_positions must contain one entry per kv_seq_lens entry" + ) + if not global_token_positions or any(position < 0 for position in global_token_positions): + raise AttentionContractError( + "global_token_positions must contain non-negative positions" + ) + if len(global_token_positions) != sum(kv_seq_lens): + raise AttentionContractError( + "global_token_positions must describe every logical cached token; " + f"expected {sum(kv_seq_lens)}, got {len(global_token_positions)}" + ) + token_offset = 0 + sequence_position_rows: list[tuple[int, ...]] = [] + for sequence_index, sequence_length in enumerate(kv_seq_lens): + sequence_positions = global_token_positions[ + token_offset : token_offset + sequence_length + ] + if any( + left >= right + for left, right in zip(sequence_positions, sequence_positions[1:], strict=False) + ): + raise AttentionContractError( + "global_token_positions must be strictly increasing within each sequence; " + f"sequence {sequence_index} is invalid" + ) + sequence_position_rows.append(sequence_positions) + token_offset += sequence_length + + for sequence_index, (cache_position, sequence_positions) in enumerate( + zip(cache_positions, sequence_position_rows, strict=True) + ): + terminal_position = sequence_positions[-1] + if cache_position != terminal_position: + raise AttentionContractError( + "cache_positions must equal the terminal global token position for each " + f"sequence; sequence {sequence_index} expected {terminal_position}, " + f"got {cache_position}" + ) + + try: + block_table = tuple(tuple(row) for row in self.block_table) + except TypeError as exc: + raise AttentionContractError( + "block_table must be a two-dimensional integer table" + ) from exc + if len(block_table) != len(kv_seq_lens) or any(not row for row in block_table): + raise AttentionContractError( + "block_table must contain one non-empty row per kv_seq_lens entry" + ) + active_block_rows: list[tuple[int, ...]] = [] + for row_index, (row, sequence_length) in enumerate( + zip(block_table, kv_seq_lens, strict=True) + ): + row_active_blocks: list[int] = [] + saw_padding = False + for column_index, block in enumerate(row): + if isinstance(block, bool) or not isinstance(block, int) or block < -1: + raise AttentionContractError( + "block_table entries must be integer block ids or -1 padding; " + f"got block_table[{row_index}][{column_index}]={block!r}" + ) + if block == -1: + saw_padding = True + continue + if saw_padding: + raise AttentionContractError( + "block_table -1 padding must be trailing; " + f"row {row_index} contains an active block after padding" + ) + row_active_blocks.append(block) + + expected_blocks = (sequence_length + page_size - 1) // page_size + if len(row_active_blocks) != expected_blocks: + raise AttentionContractError( + "block_table active page count must match kv_seq_lens and page_size; " + f"row {row_index} expected {expected_blocks}, got {len(row_active_blocks)}" + ) + if len(set(row_active_blocks)) != len(row_active_blocks): + raise AttentionContractError( + f"block_table row {row_index} contains duplicate active page ids" + ) + active_block_rows.append(tuple(row_active_blocks)) + + if not isinstance(self.prefix_cache_enabled, bool): + raise AttentionContractError("prefix_cache_enabled must be a bool") + shared_prefix_page_count = _non_negative_int( + self.shared_prefix_page_count, "shared_prefix_page_count" + ) + if self.prefix_cache_enabled and not self.prefix_cache_key: + raise AttentionContractError( + "prefix_cache_key is required when prefix_cache_enabled=True" + ) + if not self.prefix_cache_enabled and self.prefix_cache_key is not None: + raise AttentionContractError( + "prefix_cache_key must be None when prefix_cache_enabled=False" + ) + if not self.prefix_cache_enabled and shared_prefix_page_count != 0: + raise AttentionContractError( + "shared_prefix_page_count must be 0 when prefix_cache_enabled=False" + ) + + shared_prefix_pages: tuple[int, ...] = () + if shared_prefix_page_count > 0: + if any(len(row_blocks) < shared_prefix_page_count for row_blocks in active_block_rows): + raise AttentionContractError( + "shared_prefix_page_count exceeds an active block-table row" + ) + shared_prefix_token_count = shared_prefix_page_count * page_size + if any(length < shared_prefix_token_count for length in kv_seq_lens): + raise AttentionContractError( + "shared prefix pages must be fully populated and read-only" + ) + + shared_prefix_pages = active_block_rows[0][:shared_prefix_page_count] + shared_prefix_positions = sequence_position_rows[0][:shared_prefix_token_count] + for sequence_index, (row_blocks, positions) in enumerate( + zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 + ): + if row_blocks[:shared_prefix_page_count] != shared_prefix_pages: + raise AttentionContractError( + "shared prefix page ids must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + if positions[:shared_prefix_token_count] != shared_prefix_positions: + raise AttentionContractError( + "shared prefix token positions must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + + exclusive_page_owners: dict[int, int] = {} + shared_prefix_page_ids = set(shared_prefix_pages) + for sequence_index, row_blocks in enumerate(active_block_rows): + for page_id in row_blocks[shared_prefix_page_count:]: + if page_id in shared_prefix_page_ids: + raise AttentionContractError( + "a writable suffix page cannot alias a read-only shared prefix page" + ) + previous_owner = exclusive_page_owners.get(page_id) + if previous_owner is not None: + raise AttentionContractError( + "active pages may be shared across sequences only when declared as " + "read-only prefix pages; " + f"page {page_id} is used by sequences {previous_owner} and " + f"{sequence_index}" + ) + exclusive_page_owners[page_id] = sequence_index + + object.__setattr__(self, "cache_positions", cache_positions) + object.__setattr__(self, "kv_seq_lens", kv_seq_lens) + object.__setattr__(self, "block_table", block_table) + object.__setattr__(self, "global_token_positions", global_token_positions) + + +@dataclass(frozen=True) +class RoPESpec: + """Qwen3 RoPE identity for attention inputs and KV-cache replay. + + The CP attention reference consumes attention inputs. This object records + whether those inputs are pre- or post-RoPE and pins the position/cache + metadata needed to compare fused ``RoPE+Attention`` and unfused + ``RoPE -> Attention`` materializations. + """ + + q_state: RoPEState = RoPEState.POST_ROPE + k_state: RoPEState = RoPEState.POST_ROPE + k_cache_state: RoPEState = RoPEState.POST_ROPE + theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + position_ids: tuple[int, ...] | None = None + query_position_offsets: tuple[int, ...] | None = None + key_position_offsets: tuple[int, ...] | None = None + cast_at: RoPECastPoint = RoPECastPoint.AFTER_ROPE + output_dtype: AttentionDType = AttentionDType.BF16 + fusion_boundary: RoPEFusionBoundary = RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION + + def __post_init__(self) -> None: + object.__setattr__(self, "q_state", _enum_value(RoPEState, self.q_state, "q_state")) + object.__setattr__(self, "k_state", _enum_value(RoPEState, self.k_state, "k_state")) + object.__setattr__( + self, "k_cache_state", _enum_value(RoPEState, self.k_cache_state, "k_cache_state") + ) + if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + theta = float(self.theta) + if theta <= 0.0: + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + object.__setattr__(self, "theta", theta) + _positive_int(self.rotary_dim, "rotary_dim") + if self.rope_scaling is not None and ( + not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() + ): + raise AttentionContractError("rope_scaling must be a non-empty string when provided") + for field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, field) + if values is None: + continue + normalized = _integer_tuple(values, field) + if not normalized or any(value < 0 for value in normalized): + raise AttentionContractError(f"{field} must contain non-negative positions") + object.__setattr__(self, field, normalized) + object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) + object.__setattr__( + self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") + ) + object.__setattr__( + self, + "fusion_boundary", + _enum_value(RoPEFusionBoundary, self.fusion_boundary, "fusion_boundary"), + ) + + +@dataclass(frozen=True) +class AttentionContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: AttentionRole + mode: AttentionMode + dtype: AttentionDType + batch_size: int + query_sequence_length: int + head_dim: int + causal: bool + causal_offsets: tuple[int, ...] | None + sharding: ShardingSpec + reduction: ReductionSpec + 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 ( + self.mode is AttentionMode.PREFILL + and query_sequence_length != self.sharding.local_sequence_length + ): + raise AttentionContractError( + "prefill query_sequence_length must equal sharding.local_sequence_length; " + f"got {query_sequence_length} and {self.sharding.local_sequence_length}" + ) + if self.sharding.packed_sequence_offsets is not None: + packed_sequence_count = len(self.sharding.packed_sequence_offsets) - 1 + if packed_sequence_count != batch_size: + raise AttentionContractError( + "packed sequence count must equal logical batch_size; " + f"got {packed_sequence_count} packed sequences and batch_size={batch_size}" + ) + if not isinstance(self.causal, bool): + raise AttentionContractError("causal must be a bool") + if self.causal: + if self.causal_offsets is None: + raise AttentionContractError("causal_offsets are required for causal attention") + if self.causal_offsets is not None: + causal_offsets = _integer_tuple(self.causal_offsets, "causal_offsets") + if not causal_offsets or any(offset < 0 for offset in causal_offsets): + raise AttentionContractError("causal_offsets must contain non-negative offsets") + if self.sharding.packed_sequence_offsets is not None: + expected_causal_offsets = batch_size + offset_owner = "packed sequence" + else: + expected_causal_offsets = batch_size + offset_owner = "batch entry" + if len(causal_offsets) != expected_causal_offsets: + raise AttentionContractError( + f"causal_offsets must contain one entry per {offset_owner}" + ) + object.__setattr__(self, "causal_offsets", causal_offsets) + + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise AttentionContractError( + "export_lse must be True for the WS2 attention-domain LSE contract" + ) + + if self.mode is AttentionMode.DECODE and self.kv_cache is None: + raise AttentionContractError("kv_cache metadata is required for decode attention") + if self.kv_cache is not None and not isinstance(self.kv_cache, KVCacheSpec): + raise AttentionContractError("kv_cache must be a KVCacheSpec when provided") + if self.rope is not None: + if not isinstance(self.rope, RoPESpec): + raise AttentionContractError("rope must be a RoPESpec when provided") + if self.rope.rotary_dim > self.head_dim: + raise AttentionContractError( + f"rotary_dim={self.rope.rotary_dim} must not exceed head_dim={self.head_dim}" + ) + if self.rope.position_ids is not None and len(self.rope.position_ids) not in { + query_sequence_length, + self.sharding.local_sequence_length, + }: + raise AttentionContractError( + "position_ids must describe the local query sequence or full local " + "sequence length" + ) + for field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, field) + if offsets is not None and len(offsets) != batch_size: + raise AttentionContractError( + f"{field} must contain one entry per logical batch entry" + ) + if self.mode is AttentionMode.DECODE and self.kv_cache is not None: + if len(self.kv_cache.kv_seq_lens) != batch_size: + raise AttentionContractError( + "decode kv_seq_lens must contain one entry per batch entry" + ) + if len(self.kv_cache.cache_positions) != batch_size: + raise AttentionContractError( + "decode cache_positions must contain one entry per batch entry" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "global_q_heads": self.sharding.global_q_heads, + "global_kv_heads": self.sharding.global_kv_heads, + "local_q_head_start": self.sharding.local_q_head_start, + "local_q_heads": self.sharding.local_q_heads, + "local_kv_head_start": self.sharding.local_kv_head_start, + "local_kv_heads": self.sharding.local_kv_heads, + "global_sequence_length": self.sharding.global_sequence_length, + "local_sequence_length": self.sharding.local_sequence_length, + "global_block_indices": list(self.sharding.global_block_indices), + "global_block_token_starts": list(self.sharding.global_block_token_starts), + "local_block_offsets": list(self.sharding.local_block_offsets), + "packed_sequence_offsets": ( + list(self.sharding.packed_sequence_offsets) + if self.sharding.packed_sequence_offsets is not None + else None + ), + } + reduction = { + "merge": self.reduction.merge.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + } + kv_cache = None + if self.kv_cache is not None: + kv_cache = { + "cache_positions": list(self.kv_cache.cache_positions), + "kv_seq_lens": list(self.kv_cache.kv_seq_lens), + "block_table": [list(row) for row in self.kv_cache.block_table], + "global_token_positions": list(self.kv_cache.global_token_positions), + "page_size": self.kv_cache.page_size, + "prefix_cache_enabled": self.kv_cache.prefix_cache_enabled, + "prefix_cache_key": self.kv_cache.prefix_cache_key, + "shared_prefix_page_count": self.kv_cache.shared_prefix_page_count, + } + rope = None + if self.rope is not None: + rope = { + "q_state": self.rope.q_state.value, + "k_state": self.rope.k_state.value, + "k_cache_state": self.rope.k_cache_state.value, + "theta": self.rope.theta, + "rotary_dim": self.rope.rotary_dim, + "rope_scaling": self.rope.rope_scaling, + "position_ids": ( + list(self.rope.position_ids) if self.rope.position_ids is not None else None + ), + "query_position_offsets": ( + list(self.rope.query_position_offsets) + if self.rope.query_position_offsets is not None + else None + ), + "key_position_offsets": ( + list(self.rope.key_position_offsets) + if self.rope.key_position_offsets is not None + else None + ), + "cast_at": self.rope.cast_at.value, + "output_dtype": self.rope.output_dtype.value, + "fusion_boundary": self.rope.fusion_boundary.value, + } + return { + "semantic_operator": "standard_softmax_attention", + "role": self.role.value, + "mode": self.mode.value, + "dtype": self.dtype.value, + "batch_size": self.batch_size, + "query_sequence_length": self.query_sequence_length, + "head_dim": self.head_dim, + "causal": self.causal, + "causal_offsets": ( + list(self.causal_offsets) if self.causal_offsets is not None else None + ), + "export_lse": self.export_lse, + "lse_domain": "attention", + "sharding": sharding, + "reduction": reduction, + "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 + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise AttentionContractError("backend_id must be a non-empty string") + roles = frozenset(_enum_value(AttentionRole, value, "roles") for value in self.roles) + modes = frozenset(_enum_value(AttentionMode, value, "modes") for value in self.modes) + dtypes = frozenset(_enum_value(AttentionDType, value, "dtypes") for value in self.dtypes) + if not roles or not modes or not dtypes: + raise AttentionContractError("backend roles, modes, and dtypes must not be empty") + cp_world_sizes = _integer_tuple(self.cp_world_sizes, "cp_world_sizes") + if not cp_world_sizes or any(size <= 0 for size in cp_world_sizes): + raise AttentionContractError("cp_world_sizes must contain positive values") + if len(set(cp_world_sizes)) != len(cp_world_sizes): + raise AttentionContractError("cp_world_sizes must not contain duplicates") + tp_world_sizes = None + if self.tp_world_sizes is not None: + tp_world_sizes = _integer_tuple(self.tp_world_sizes, "tp_world_sizes") + if not tp_world_sizes or any(size <= 0 for size in tp_world_sizes): + raise AttentionContractError("tp_world_sizes must contain positive values") + if len(set(tp_world_sizes)) != len(tp_world_sizes): + raise AttentionContractError("tp_world_sizes must not contain duplicates") + for field in ( + "exports_attention_lse", + "deterministic_cp_merge", + "supports_packed_varlen", + "supports_kv_cache", + "supports_rope_metadata", + "supports_fused_rope_attention", + ): + if not isinstance(getattr(self, field), bool): + raise AttentionContractError(f"{field} must be a bool") + if self.implementation_kind not in {"production", "reference", "deterministic"}: + raise AttentionContractError( + "implementation_kind must be production, reference, or deterministic" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "modes", modes) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + + def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.mode not in self.modes: + reasons.append(f"mode={contract.mode.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if contract.export_lse and not self.exports_attention_lse: + reasons.append("attention-domain LSE export is unsupported") + if cp_size > 1 and not self.deterministic_cp_merge: + reasons.append("deterministic CP (out, lse) merge is unsupported") + if ( + contract.sharding.packed_sequence_offsets is not None + and not self.supports_packed_varlen + ): + reasons.append("packed varlen layout is unsupported") + if contract.kv_cache is not None and not self.supports_kv_cache: + reasons.append("KV-cache identity materialization is unsupported") + if contract.rope is not None and not self.supports_rope_metadata: + reasons.append("RoPE/position metadata is unsupported") + if ( + contract.rope is not None + and contract.rope.fusion_boundary is RoPEFusionBoundary.FUSED_ROPE_ATTENTION + and not self.supports_fused_rope_attention + ): + reasons.append("fused RoPE+Attention boundary is unsupported") + 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, + "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", +] diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 75a2ab7f..c9777cf7 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/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py new file mode 100644 index 00000000..302ad73f --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -0,0 +1,553 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic context-parallel attention reference. + +This module is the correctness-first WS2 reference for CP-aware standard +softmax attention. It intentionally stays in PyTorch and uses fp32 partial +states so fused CUDA/Triton backends can validate their CP/LSE merge semantics +against a small, inspectable implementation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional, Sequence + +import torch + +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + +@dataclass(frozen=True) +class AttentionPartialState: + """One KV block's attention state before deterministic LSE merge. + + ``out`` is already normalized within the local KV block and has shape + ``[B, Hq, Sq, D]``. ``lse`` is the local attention-domain log-sum-exp with + shape ``[B, Hq, Sq]``. ``block_start`` / ``block_end`` are logical global KV + positions and define the canonical merge order. + """ + + out: torch.Tensor + lse: torch.Tensor + block_start: int + block_end: int + + def __post_init__(self) -> None: + if self.out.ndim != 4: + raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") + if self.lse.shape != self.out.shape[:3]: + raise ValueError("partial attention lse must have shape [B, Hq, Sq]") + if self.block_start < 0: + raise ValueError("block_start must be non-negative") + if self.block_end < self.block_start: + raise ValueError("block_end must be >= block_start") + + +def merge_attention_partial_states( + states: Sequence[AttentionPartialState], +) -> AttentionPartialState: + """Merge CP/chunk partial states in logical block order. + + The merge is the online-softmax/LSE merge used by attention, not a plain + sum. The input order is deliberately ignored: states are sorted by logical + ``block_start`` so the result depends on global block indices rather than + arrival order. + """ + + if not states: + raise ValueError("at least one attention partial state is required") + + ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) + _validate_merge_shapes_and_ranges(ordered) + + merged = ordered[0] + merged_out = merged.out.float() + merged_lse = merged.lse.float() + for state in ordered[1:]: + merged_out, merged_lse = _merge_two_states( + merged_out, + merged_lse, + state.out.float(), + state.lse.float(), + ) + + return AttentionPartialState( + out=merged_out, + lse=merged_lse, + block_start=ordered[0].block_start, + block_end=ordered[-1].block_end, + ) + + +class DeterministicCPAttentionReferenceOp: + """Correctness-first CP attention reference for prefill and chunked prefill. + + The 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. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + return self.forward( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and final input-dtype write.""" + + out, _ = self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=q.dtype, + ) + return out + + def forward_fp32( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and fp32 output.""" + + out, _ = self.forward_fp32_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(out, lse)`` for the CP reference path. + + ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 + until the final write, then downcast to ``output_dtype``. When omitted, + ``output_dtype`` defaults to the input dtype. + """ + + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + out = out.to(q.dtype if output_dtype is None else output_dtype) + return out, lse + + def forward_fp32_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return fp32 ``(out, lse)`` for the CP reference path.""" + + return self.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=torch.float32, + ) + + def local_partial_state( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + q_start: int, + k_start: int, + total_kv_len: int, + total_query_len: Optional[int] = None, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + ) -> AttentionPartialState: + """Compute one query shard against one logical KV block. + + ``query_position_offsets`` and ``key_position_offsets`` are optional + per-batch-row base positions. They let the reference express varlen or + packed metadata while retaining the dense [B, H, S, D] tensor layout. + For post-RoPE Q/K, these offsets must describe the same absolute token + positions used when RoPE was applied. + """ + + _validate_qkv(q, k, v) + if q_start < 0 or k_start < 0: + raise ValueError("q_start and k_start must be non-negative") + if total_kv_len < k_start + k.size(2): + raise ValueError("total_kv_len must cover the local KV block") + if total_query_len is None: + total_query_len = q.size(2) + if total_query_len < q_start + q.size(2): + raise ValueError("total_query_len must cover the local query block") + if key_padding_mask is not None: + if key_padding_mask.shape != (q.size(0), k.size(2)): + raise ValueError("local key_padding_mask must have shape [B, local_skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("local key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + q.size(0), + q.device, + default=total_kv_len - total_query_len, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + q.size(0), + q.device, + default=0, + name="key_position_offsets", + ) + + ctx = NativeAttentionOp._strict_fp32_math(q.device.type) + with ctx: + qf = q.float() + kf = k.float() + vf = v.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hkv != hq: + repeat = hq // hkv + kf = kf.repeat_interleave(repeat, dim=1) + vf = vf.repeat_interleave(repeat, dim=1) + + if skv == 0: + zero_dep = _zero_dependency(qf, kf, vf) + return AttentionPartialState( + out=torch.zeros(q.size(0), hq, sq, dim, device=q.device, dtype=torch.float32) + + zero_dep, + lse=torch.full( + (q.size(0), hq, sq), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep, + block_start=k_start, + block_end=k_start, + ) + + scale_value = scale if scale is not None else (1.0 / math.sqrt(dim)) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * scale_value + if causal: + query_base = query_offsets[:, None] + q_start + key_base = key_offsets[:, None] + k_start + q_pos = torch.arange(sq, device=q.device, dtype=torch.long) + query_base + k_pos = torch.arange(skv, device=q.device, dtype=torch.long) + key_base + causal_mask = k_pos[:, None, :] > q_pos[:, :, None] + scores = scores.masked_fill(causal_mask[:, None, :, :], float("-inf")) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + + lse = torch.logsumexp(scores, dim=-1) + finite_lse = torch.isfinite(lse) + weights = torch.exp(scores - lse.unsqueeze(-1)) + weights = torch.where(finite_lse.unsqueeze(-1), weights, torch.zeros_like(weights)) + out = torch.matmul(weights, vf) + return AttentionPartialState( + out=out, + lse=lse, + block_start=k_start, + block_end=k_start + skv, + ) + + def _forward_impl( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + cp_world_size: int, + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + _validate_qkv(q, k, v) + if cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + + q_bounds = _split_bounds(sq, cp_world_size) + kv_bounds = _kv_block_bounds(skv, cp_world_size, kv_chunk_size) + out_chunks: list[torch.Tensor] = [] + lse_chunks: list[torch.Tensor] = [] + for q_start, q_end in q_bounds: + if q_start == q_end: + continue + q_block = q[:, :, q_start:q_end, :] + states = [ + self.local_partial_state( + q_block, + k[:, :, k_start:k_end, :], + v[:, :, k_start:k_end, :], + q_start=q_start, + k_start=k_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None if key_padding_mask is None else key_padding_mask[:, k_start:k_end] + ), + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + ) + for k_start, k_end in kv_bounds + if k_start != k_end + ] + if states: + merged = merge_attention_partial_states(states) + out_chunks.append(merged.out) + lse_chunks.append(merged.lse) + else: + zero_dep = _zero_dependency(q_block.float(), k.float(), v.float()) + out_chunks.append( + torch.zeros(batch, hq, q_end - q_start, dim, device=q.device) + zero_dep + ) + lse_chunks.append( + torch.full( + (batch, hq, q_end - q_start), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep + ) + + if not out_chunks: + zero_dep = _zero_dependency(q.float(), k.float(), v.float()) + return ( + torch.empty(batch, hq, 0, dim, device=q.device, dtype=torch.float32) + zero_dep, + torch.empty(batch, hq, 0, device=q.device, dtype=torch.float32) + zero_dep, + ) + return torch.cat(out_chunks, dim=2), torch.cat(lse_chunks, dim=2) + + +def _merge_two_states( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + merged_lse = torch.logaddexp(lse_a, lse_b) + finite = torch.isfinite(merged_lse) + weight_a = torch.where(finite, torch.exp(lse_a - merged_lse), torch.zeros_like(merged_lse)) + weight_b = torch.where(finite, torch.exp(lse_b - merged_lse), torch.zeros_like(merged_lse)) + merged_out = weight_a.unsqueeze(-1) * out_a + weight_b.unsqueeze(-1) * out_b + return merged_out, merged_lse + + +def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) -> None: + first = states[0] + previous_end = first.block_end + for state in states[1:]: + if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: + raise ValueError("all partial states must have matching out/lse shapes") + if state.block_start < previous_end: + raise ValueError("partial state block ranges must not overlap") + previous_end = state.block_end + + +def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q, k, and v must have shape [B, H, S, D]") + if k.shape != v.shape: + raise ValueError("k and v must have the same shape") + if q.size(0) != k.size(0) or q.size(3) != k.size(3): + raise ValueError("q, k, and v must share batch size and head dim") + if q.size(1) % k.size(1) != 0: + raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") + + +def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: + total = torch.tensor(0.0, device=tensors[0].device) + for tensor in tensors: + total = total + tensor.sum() + return total * 0.0 + + +def _normalize_position_offsets( + offsets: Optional[torch.Tensor], + batch: int, + device: torch.device, + *, + default: int, + name: str, +) -> torch.Tensor: + if offsets is None: + return torch.full((batch,), default, dtype=torch.long, device=device) + if offsets.ndim != 1 or offsets.numel() != batch: + raise ValueError(f"{name} must have shape [B]") + if torch.is_floating_point(offsets) or torch.is_complex(offsets) or offsets.dtype == torch.bool: + raise ValueError(f"{name} must contain integer positions") + return offsets.to(device=device, dtype=torch.long) + + +def _split_bounds(length: int, parts: int) -> list[tuple[int, int]]: + base, extra = divmod(length, parts) + bounds: list[tuple[int, int]] = [] + start = 0 + for index in range(parts): + width = base + (1 if index < extra else 0) + end = start + width + bounds.append((start, end)) + start = end + return bounds + + +def _kv_block_bounds( + length: int, + cp_world_size: int, + kv_chunk_size: Optional[int], +) -> list[tuple[int, int]]: + bounds: list[tuple[int, int]] = [] + for start, end in _split_bounds(length, cp_world_size): + if kv_chunk_size is None: + bounds.append((start, end)) + continue + cursor = start + while cursor < end: + chunk_end = min(cursor + kv_chunk_size, end) + bounds.append((cursor, chunk_end)) + cursor = chunk_end + return bounds + + +__all__ = [ + "AttentionPartialState", + "DeterministicCPAttentionReferenceOp", + "merge_attention_partial_states", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 2d5c470f..8cfad139 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -8,6 +8,15 @@ from enum import Enum, EnumMeta from typing import Any, Dict, Optional, Set, Type +from rl_engine.kernels.attention_contract import ( + AttentionBackendCapability, + AttentionContract, + AttentionContractError, + AttentionDispatchResult, + AttentionDType, + AttentionMode, + AttentionRole, +) from rl_engine.kernels.semantic_registry import ( OperatorBackendDescriptor, OperatorFallbackPolicy, @@ -107,6 +116,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 @@ -219,6 +234,39 @@ def __init__(self): 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", + ), + } + self._priority_map = { "cuda": { "logp": [ @@ -252,6 +300,7 @@ def __init__(self): OpBackend.CUDA_DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION, ], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "linear_logp": [ @@ -291,6 +340,7 @@ def __init__(self): OpBackend.TRITON_GENERIC, ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], @@ -315,6 +365,7 @@ def __init__(self): "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], @@ -432,6 +483,120 @@ def get_op(self, op_type: str) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + def get_attention_op( + self, + contract: AttentionContract, + *, + requested_backend: str = "deterministic", + ) -> AttentionDispatchResult: + """Resolve only a backend that explicitly supports the WS2 contract. + + This entry point is intentionally separate from legacy ``get_op`` so + existing callers retain their current behavior while WS2 callers cannot + silently fall back to a backend with different distributed semantics. + """ + + if not isinstance(contract, AttentionContract): + raise AttentionContractError("contract must be an AttentionContract") + if not isinstance(requested_backend, str) or not requested_backend.strip(): + raise AttentionContractError("requested_backend must be a non-empty string") + requested_backend = requested_backend.strip().lower() + + platform = self._platform() + op_type = "kv_cache_attention" if contract.mode is AttentionMode.DECODE else "attention" + candidates = self._priority_map.get(platform, {}).get(op_type, []) + rejected: list[str] = [] + + for backend in candidates: + capability = self._attention_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no AttentionBackendCapability declared") + continue + incompatibilities = list(capability.incompatibilities(contract)) + policy_mismatch = self._attention_policy_mismatch(requested_backend, capability) + if policy_mismatch is not None: + incompatibilities.append(policy_mismatch) + if incompatibilities: + rejected.append(f"{backend.name}: " + "; ".join(incompatibilities)) + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + continue + + provenance = { + "requested_backend": requested_backend, + "actual_backend": capability.backend_id, + "backend_enum": backend.name, + "platform": platform, + "fallback": bool(rejected), + "prior_rejections": list(rejected), + "contract": contract.to_dict(), + "capability": capability.to_dict(), + } + return AttentionDispatchResult( + op=op, + capability=capability, + provenance=provenance, + ) + + details = " | ".join(rejected) if rejected else "no candidates registered" + requested = contract.to_dict() + raise RuntimeError( + "No attention backend supports the requested WS2 contract on " + f"{platform}: role={requested['role']}, mode={requested['mode']}, " + f"dtype={requested['dtype']}, TP={contract.sharding.tp_world_size}, " + f"CP={contract.sharding.cp_world_size}. Rejections: {details}" + ) + + @staticmethod + def _attention_policy_mismatch( + requested_backend: str, + capability: AttentionBackendCapability, + ) -> str | None: + if requested_backend == "auto": + return None + if requested_backend in {"production", "reference", "deterministic"}: + if capability.implementation_kind == requested_backend: + return None + return ( + f"implementation_kind={capability.implementation_kind} does not satisfy " + f"requested_backend={requested_backend}" + ) + if capability.backend_id == requested_backend: + return None + return ( + f"backend_id={capability.backend_id} does not match " + f"requested_backend={requested_backend}" + ) + + def _platform(self) -> str: + if device_ctx.is_rocm: + return "rocm" + if device_ctx.device_type == "cuda": + return "cuda" + return "cpu" + + def _get_or_create_backend(self, backend: OpBackend) -> Any | None: + if backend.name in self._instance_cache: + return self._instance_cache[backend.name] + if backend.name in self._failed_backends: + return None + + op_class = self._load_backend(backend) + if op_class is None: + self._failed_backends.add(backend.name) + return None + try: + op = op_class() + except Exception as exc: + logger.error(f"Failed to instantiate {backend.name}: {exc}") + self._failed_backends.add(backend.name) + return None + self._instance_cache[backend.name] = op + return op + def _load_backend(self, backend: OpBackend) -> Optional[Type]: """Import a legacy backend and distinguish wrapper bugs from absence.""" diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py new file mode 100644 index 00000000..b986fc15 --- /dev/null +++ b/tests/test_attention_contract.py @@ -0,0 +1,551 @@ +# 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, +) +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, + 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_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() + + 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]["attention"] = [OpBackend.PYTORCH_ATTN] + + with pytest.raises(RuntimeError, match="no AttentionBackendCapability declared"): + registry.get_attention_op(_contract()) + + +def test_declared_compatible_backend_resolves_and_records_provenance(): + registry = KernelRegistry() + registry._attention_capabilities[OpBackend.PYTORCH_NATIVE_ATTENTION] = _declared_cp_backend() + + result = registry.get_attention_op(_contract(), requested_backend="deterministic") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-cp-attention" + assert result.provenance["requested_backend"] == "deterministic" + assert result.provenance["actual_backend"] == "test-deterministic-cp-attention" + assert result.provenance["fallback"] is False + assert result.provenance["contract"]["sharding"]["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() + 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..815daf12 --- /dev/null +++ b/tests/test_attention_cross_config_binding.py @@ -0,0 +1,671 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for #235 PR4: rollout/training attention contract binding. + +Every test here runs on CPU without Megatron or vLLM installed. That is the point: +the binding rules are contract logic, and contract logic that can only be exercised +on a 2-node x 2-GPU cluster would never be exercised. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from rl_engine.alignment.cross_config.adapters import ( + QWEN3_8B, + WS2_ATTENTION_KNOBS, + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) +from rl_engine.alignment.cross_config.attention_binding import ( + ATTENTION_LSE_DOMAIN, + AttentionBindingError, + BindingErrorCode, + BindingTier, + bind_attention_contracts, + first_blocking_issue, + identity_fingerprint, + summarize_binding, +) +from rl_engine.alignment.cross_config.determinism import ( + compare_determinism, + megatron_probe_from_config, + vllm_probe_from_env, +) +from rl_engine.alignment.cross_config.schema import MaterializationStatus +from rl_engine.kernels.attention_contract import AttentionContractError, AttentionMode + +pytestmark = pytest.mark.unit + + +TRAINING_KNOBS = { + "batch.size": 2, + "training.tensor_parallel_size": 2, + "training.context_parallel_size": 2, + "training.compute_dtype": "bf16", +} + +ROLLOUT_KNOBS = { + "batch.size": 2, + "rollout.tensor_parallel_size": 2, + "rollout.context_parallel_size": 1, + "rollout.dtype": "bf16", +} + + +def _identity(**overrides): + identity = { + "checkpoint_id": "qwen3-8b", + "model_version": "v1", + "weight_version": 7, + "tokenizer_fingerprint": "tokenizer-abc", + "token_ids_fingerprint": "tokens-abc", + "active_mask_fingerprint": "mask-abc", + "position_ids_fingerprint": "pos-abc", + "padding_side": "right", + "pre_update_state": "pre_update", + "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 _bind(rollout_identity=None, training_identity=None, **kwargs): + rollout, training = _contracts() + return bind_attention_contracts( + rollout_contract=kwargs.pop("rollout_contract", rollout), + training_contract=kwargs.pop("training_contract", training), + rollout_identity=rollout_identity if rollout_identity is not None else _identity(), + training_identity=training_identity if training_identity is not None else _identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + **kwargs, + ) + + +# -------------------------------------------------------------------------- +# tier 1: identity +# -------------------------------------------------------------------------- + + +def test_matching_identity_binds_despite_different_materialization(): + """The core claim of PR4: same identity + same reduction, different runtimes.""" + + result = _bind() + + assert result.comparable + assert result.passed + assert result.issues == () + # Training runs CP=2 full prefill, rollout runs CP=1 chunked prefill. Those + # differences are recorded, not rejected. + assert "mode" in result.recorded_differences + assert "sharding.cp_world_size" in result.recorded_differences + assert result.recorded_differences["mode"] == { + "rollout": "chunked_prefill", + "training": "prefill", + } + + +def test_weight_version_mismatch_is_not_comparable(): + result = _bind(rollout_identity=_identity(weight_version=6)) + + assert not result.comparable + assert not result.passed + codes = {issue.code for issue in result.issues} + assert BindingErrorCode.IDENTITY_MISMATCH in codes + blocking = first_blocking_issue(result) + assert blocking is not None and blocking.tier is BindingTier.IDENTICAL + assert "NOT COMPARABLE" in summarize_binding(result) + + +def test_rope_theta_mismatch_is_not_comparable(): + """RoPE math constants are identity, not materialization.""" + + result = _bind(training_identity=_identity(rope_theta=10000.0)) + + assert not result.comparable + assert any(issue.field == "rope_theta" for issue in result.issues) + + +def test_null_rope_scaling_is_a_value_not_an_omission(): + """Qwen3-8B applies no RoPE scaling; ``None`` must not read as undeclared.""" + + result = _bind() + + assert not result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + + +def test_missing_identity_field_is_reported_per_side(): + identity = _identity() + del identity["padding_side"] + result = _bind(rollout_identity=identity, training_identity=identity) + + missing = result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + assert {issue.field for issue in missing} == { + "rollout.padding_side", + "training.padding_side", + } + assert not result.comparable + + +def test_single_gpu_harness_may_waive_full_identity(): + """#235 PR2 has no KV-cache identity to declare; it opts out explicitly.""" + + identity = _identity() + del identity["global_token_positions_fingerprint"] + del identity["kv_seq_lens_fingerprint"] + + strict = _bind(rollout_identity=identity, training_identity=identity) + waived = _bind( + rollout_identity=identity, + training_identity=identity, + require_full_identity=False, + ) + + assert not strict.comparable + assert waived.comparable and waived.passed + + +def test_identity_fingerprint_ignores_undeclared_extra_keys(): + base = _identity() + decorated = dict(base, diagnostic_note="added later") + + assert identity_fingerprint(base) == identity_fingerprint(decorated) + + +# -------------------------------------------------------------------------- +# tier 2: reduction semantics +# -------------------------------------------------------------------------- + + +def test_reduction_semantics_are_bound_and_fingerprinted(): + result = _bind() + + reduction = result.provenance["training"]["contract"]["reduction"] + assert reduction["merge"] == "online_softmax_lse" + assert reduction["acc_dtype"] == "fp32" + assert reduction["order"] == "global_block_index" + assert reduction["downcast_at"] == "final_write" + assert result.reduction_fingerprint + + +def test_reduction_engine_difference_is_recorded_not_rejected(): + """A TE merge oracle on one side must not fail the binding.""" + + from rl_engine.alignment.cross_config.attention_binding import ( + RECORDED_FIELDS, + SEMANTIC_REDUCTION_FIELDS, + ) + + assert "reduction.engine" in RECORDED_FIELDS + assert "engine" not in SEMANTIC_REDUCTION_FIELDS + + +def test_lse_domain_is_recorded_as_attention_domain(): + """#235: attention exports attention-domain LSE, not vocab-logprob LSE.""" + + result = _bind() + + assert result.provenance["lse_domain"] == ATTENTION_LSE_DOMAIN == "attention" + + +def test_mixed_dtypes_fail_closed(): + """BF16 rollout against FP16 training produces an unattributable number.""" + + rollout = VllmRolloutMaterializer().build_contract( + {**ROLLOUT_KNOBS, "rollout.dtype": "float16"} + ) + result = _bind(rollout_contract=rollout) + + assert result.comparable # identity is fine + assert not result.passed + assert any(issue.field == "dtype" for issue in result.issues) + + +def test_precision_sweep_may_opt_into_mixed_dtypes(): + """#235 PR5 sweeps BF16 against an FP32 reference; it says so explicitly.""" + + training = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "float32"} + ) + result = _bind(training_contract=training, allow_dtype_difference=True) + + assert result.passed + assert result.provenance["dtype"] == "fp32" + + +def test_batch_size_mismatch_is_not_comparable(): + """Batch invariance is a claim about batch makeup, so it belongs to identity.""" + + result = _bind(rollout_identity=_identity(batch_size=4)) + + assert not result.comparable + assert any(issue.field == "batch_size" for issue in result.issues) + + +def test_split_kv_policy_difference_is_recorded(): + """#236 has no split-KV field, so the adapter supplies it to the recorded tier.""" + + result = _bind( + rollout_recorded_extra={"split_kv_policy": 8}, + training_recorded_extra={"split_kv_policy": None}, + ) + + assert result.passed + assert result.recorded_differences["split_kv_policy"] == { + "rollout": 8, + "training": None, + } + + +def test_matching_split_kv_policy_is_not_reported_as_a_difference(): + result = _bind( + rollout_recorded_extra={"split_kv_policy": 32}, + training_recorded_extra={"split_kv_policy": 32}, + ) + + assert "split_kv_policy" not in result.recorded_differences + + +# -------------------------------------------------------------------------- +# role and input validation +# -------------------------------------------------------------------------- + + +def test_swapped_roles_are_rejected_outright(): + rollout, training = _contracts() + + with pytest.raises(AttentionBindingError): + bind_attention_contracts( + rollout_contract=training, + training_contract=rollout, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="a", + training_backend_id="b", + ) + + +# -------------------------------------------------------------------------- +# determinism cross-check +# -------------------------------------------------------------------------- + + +def _megatron_env(): + return {"NCCL_ALGO": "Tree", "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0"} + + +def _vllm_env(**overrides): + env = { + "VLLM_BATCH_INVARIANT": "1", + "NCCL_ALGO": "allreduce:tree", + "NCCL_PROTO": "Simple", + "NCCL_MIN_NCHANNELS": "1", + "NCCL_MAX_NCHANNELS": "1", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + } + env.update(overrides) + return env + + +def test_nccl_algo_mismatch_blocks_the_binding(): + """Megatron asserts NCCL_ALGO; vLLM hard-sets a different value.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert not report.compatible + fields = {issue.field for issue in report.issues} + assert "env.NCCL_ALGO" in fields + assert "env.NCCL_PROTO" in fields + + +def test_matching_nccl_settings_are_compatible(): + shared = {"NCCL_ALGO": "allreduce:tree", "NCCL_PROTO": "Simple"} + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=dict(shared) + ) + rollout = vllm_probe_from_env({**_vllm_env(**shared), "CUBLAS_WORKSPACE_CONFIG": None}) + + report = compare_determinism(rollout=rollout, training=training) + + assert report.compatible, [issue.to_dict() for issue in report.issues] + + +def test_determinism_switch_off_on_either_side_blocks(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=False), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env(VLLM_BATCH_INVARIANT="0")) + + report = compare_determinism(rollout=rollout, training=training) + + fields = {issue.field for issue in report.issues} + assert "training.deterministic_mode" in fields + assert "rollout.VLLM_BATCH_INVARIANT" in fields + + +def test_tf32_asymmetry_is_recorded_not_blocking(): + """Megatron does not manage TF32 at all; vLLM disables it. Record the gap.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert training.tf32_disabled is None + assert rollout.tf32_disabled is True + assert "tf32_disabled" in report.differences + assert not any(issue.field == "tf32_disabled" for issue in report.issues) + + +def test_determinism_issues_flow_into_the_binding(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + report = compare_determinism(rollout=rollout, training=training) + + result = _bind(determinism_issues=report.issues) + + assert result.comparable # identity is fine + assert not result.passed # but the reduction environment is not + assert result.issues_by_code(BindingErrorCode.DETERMINISM_INCOMPATIBLE) + assert "FAILED CLOSED" in summarize_binding(result) + + +# -------------------------------------------------------------------------- +# sharding derived from the frozen #239 layout +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_cp_shards_cover_the_global_sequence_without_overlap(cp_rank): + contract = MegatronAttentionMaterializer( + cp_rank=cp_rank, global_sequence_length=4096 + ).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert sharding.local_sequence_length == 2048 + assert sharding.global_block_indices == (cp_rank,) + assert sharding.global_block_token_starts == (cp_rank * 2048,) + # The causal offset must be the number of preceding *global* tokens, otherwise + # rank 1 would mask as if its shard started at position zero. + assert contract.causal_offsets == (cp_rank * 2048, cp_rank * 2048) + + +def test_tp_head_shards_split_qwen3_gqa_evenly(): + contract = MegatronAttentionMaterializer(tp_rank=1).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert (sharding.global_q_heads, sharding.global_kv_heads) == (32, 8) + assert (sharding.local_q_heads, sharding.local_kv_heads) == (16, 4) + assert (sharding.local_q_head_start, sharding.local_kv_head_start) == (16, 4) + + +@pytest.mark.parametrize("tp_world_size", [2, 4, 8]) +def test_supported_tp_degrees_shard_qwen3_gqa(tp_world_size): + """Qwen3-8B has 32 Q heads and 8 KV heads, so TP in {2, 4, 8} all divide.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": tp_world_size} + ) + + sharding = contract.sharding + assert sharding.local_q_heads == 32 // tp_world_size + assert sharding.local_kv_heads == 8 // tp_world_size + + +@pytest.mark.parametrize( + ("knob_value", "expected"), + [("bfloat16", "bf16"), ("float16", "fp16"), ("float32", "fp32"), ("fp16", "fp16")], +) +def test_planner_normalized_dtypes_reach_the_contract(knob_value, expected): + """The planner emits torch spellings; AttentionDType uses short ones.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": knob_value} + ) + + assert contract.dtype.value == expected + + +def test_unknown_dtype_is_rejected_with_the_offending_field(): + with pytest.raises(ValueError, match="training.compute_dtype"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "int8"} + ) + + +def test_indivisible_tp_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": 3} + ) + + +def test_indivisible_cp_sequence_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer(global_sequence_length=4097).build_contract(TRAINING_KNOBS) + + +# -------------------------------------------------------------------------- +# materialization: fail closed rather than silently substitute +# -------------------------------------------------------------------------- + + +def _statuses(materialization, path): + return [app.status for app in materialization.applications if app.path == path] + + +def test_arrival_merge_order_is_unsupported_not_silently_corrected(): + """The control group must stay distinguishable from the treatment.""" + + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_order": "arrival"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_order") == [ + MaterializationStatus.UNSUPPORTED + ] + assert materialization.binding.side_configs["training"]["contract"] is None + assert "arrival" in materialization.binding.side_configs["training"]["contract_error"] + + +def test_bf16_reduction_accumulation_is_unsupported(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_acc_dtype": "bf16"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_acc_dtype") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_te_oracle_engine_is_unsupported_until_pr2_pr3(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_engine": "te_oracle"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_engine") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_vllm_cp_falls_back_to_one_in_decode_and_says_why(): + materializer = VllmRolloutMaterializer(mode=AttentionMode.DECODE) + normalized = { + "batch": {"size": 2}, + "rollout": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + + assert materializer.effective_cp_world_size({"rollout.context_parallel_size": 2}) == 1 + + materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) + contract_error = materialization.binding.side_configs["rollout"]["contract_error"] + assert "#235 PR6" in contract_error + + +def test_decode_contract_is_refused_without_kv_cache_identity(): + with pytest.raises(AttentionContractError, match="PR6"): + VllmRolloutMaterializer(mode=AttentionMode.DECODE).build_contract(ROLLOUT_KNOBS) + + +def test_materializers_expose_distinct_implementation_fingerprints(): + megatron = MegatronAttentionMaterializer().implementation_fingerprint + vllm = VllmRolloutMaterializer().implementation_fingerprint + + assert megatron and vllm and megatron != vllm + + +def test_runtime_binding_reports_the_frozen_topology(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + binding = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS).binding + + topology = binding.topology["training"] + assert topology["tensor_parallel_size"] == 2 + assert topology["context_parallel_size"] == 2 + assert topology["world_size"] == 4 + assert topology["pipeline_parallel_size"] == 1 + assert topology["data_parallel_size"] == 1 + + +# -------------------------------------------------------------------------- +# provenance adapters +# -------------------------------------------------------------------------- + + +def test_megatron_provenance_flags_undeclared_frozen_scope_fields(): + adapter = MegatronProvenanceAdapter(SimpleNamespace(deterministic_mode=True)) + + violations = adapter.frozen_scope_violations() + + # Nothing is declared, so every assertion reads as unknown rather than as met. + assert any("expert_model_parallel_size" in text for text in violations) + assert any("fp8" in text for text in violations) + + +def test_megatron_provenance_accepts_a_conforming_dense_config(): + adapter = MegatronProvenanceAdapter( + SimpleNamespace( + deterministic_mode=True, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + sequence_parallel=False, + fp8=None, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + ) + + assert adapter.frozen_scope_violations() == ("fp8 is not declared (expected None)",) + + +def test_megatron_construction_fingerprint_tracks_fusion_changes(): + base = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=False) + fused = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=True) + + assert ( + MegatronProvenanceAdapter(base).construction_fingerprint + != MegatronProvenanceAdapter(fused).construction_fingerprint + ) + + +def test_vllm_provenance_reads_page_size_and_split_kv_policy(): + adapter = VllmProvenanceAdapter( + cache_config=SimpleNamespace(block_size=16, cache_dtype="auto"), + attention_config=SimpleNamespace(flash_attn_max_num_splits_for_cuda_graph=32), + ) + + assert adapter.kv_page_size == 16 + assert adapter.split_kv_policy == 32 + + +def test_vllm_provenance_flags_fp8_kv_cache_and_cascade_attention(): + adapter = VllmProvenanceAdapter( + model_config=SimpleNamespace(quantization=None, disable_cascade_attn=False), + cache_config=SimpleNamespace( + cache_dtype="fp8", calculate_kv_scales=False, sliding_window=None + ), + parallel_config=SimpleNamespace(pipeline_parallel_size=1, data_parallel_size=1), + ) + + violations = adapter.frozen_scope_violations() + + assert any("cache_dtype" in text for text in violations) + assert any("disable_cascade_attn" in text for text in violations) + + +# -------------------------------------------------------------------------- +# scenario config +# -------------------------------------------------------------------------- + + +SCENARIO = ( + Path(__file__).resolve().parents[1] + / "examples" + / "cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json" +) + + +def test_scenario_uses_megatron_vocabulary_only(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + training = config["baseline"]["training"] + + assert training["attention_backend"] in {"flash", "fused", "unfused", "local", "auto"} + assert training["tensor_parallel_size"] == 2 + assert training["context_parallel_size"] == 2 + assert config["baseline"]["rollout"]["batch_invariant"] is True + + +def test_scenario_knob_paths_all_exist(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + + def paths(mapping, prefix=""): + for key, value in mapping.items(): + path = f"{prefix}{key}" + if isinstance(value, dict): + yield from paths(value, f"{path}.") + else: + yield path + + declared = set(paths(config["baseline"])) + unknown = declared - set(WS2_ATTENTION_KNOBS) + assert not unknown, f"scenario declares unknown knobs: {sorted(unknown)}" + + for intervention in config["interventions"]: + assert intervention["path"] in WS2_ATTENTION_KNOBS diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py new file mode 100644 index 00000000..03c861dd --- /dev/null +++ b/tests/test_cp_attention.py @@ -0,0 +1,423 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the WS2 deterministic CP attention reference. + +The implementation is a correctness-first prefill/chunked-prefill reference: +local KV blocks produce ``(out, lse)`` partial states and CP merges those states +with fp32 online-softmax arithmetic in logical global-block order. +""" + +import contextlib +import math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + merge_attention_partial_states, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.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 + + +@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_inputs_are_not_mutated(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 6, 6, seed=7) + mask = torch.ones(2, 6, dtype=torch.bool) + qc, kc, vc, mc = q.clone(), k.clone(), v.clone(), mask.clone() + + op.forward_fp32_with_lse(q, k, v, causal=True, key_padding_mask=mask, cp_world_size=2) + + assert torch.equal(q, qc) + assert torch.equal(k, kc) + assert torch.equal(v, vc) + assert torch.equal(mask, mc) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"cp_world_size": 0}, "cp_world_size"), + ({"cp_world_size": 2, "kv_chunk_size": 0}, "kv_chunk_size"), + ], +) +def test_invalid_parallelism_arguments_raise(kwargs, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=8) + with pytest.raises(ValueError, match=message): + op.forward_fp32_with_lse(q, k, v, causal=True, **kwargs) + + +def test_invalid_gqa_and_mask_shapes_raise(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 6, 4, _HEAD_DIM) + k = torch.randn(1, 4, 4, _HEAD_DIM) + v = torch.randn(1, 4, 4, _HEAD_DIM) + with pytest.raises(ValueError, match="not divisible"): + op.forward_fp32_with_lse(q, k, v, causal=True) + + q, k, v = _qkv(1, 4, 4, seed=9) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 3, dtype=torch.bool)) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 4)) + with pytest.raises(ValueError, match="query_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + query_position_offsets=torch.ones(2, dtype=torch.long), + ) + with pytest.raises(ValueError, match="key_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + key_position_offsets=torch.ones(1, dtype=torch.float32), + ) + + +def test_overlapping_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="overlap"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=3), + AttentionPartialState(out=out, lse=lse, block_start=2, block_end=4), + ] + ) + + +def test_registry_dispatches_cp_attention_reference(): + assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) diff --git a/tests/test_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_operator_inputs.py b/tests/test_operator_inputs.py index 3f2b4863..66395e8c 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -36,6 +36,7 @@ def _args(**overrides): "rms_norm", "matmul", "attention", + "cp_attention", "logp", "linear_logp", "batch_invariant_logp",