[WS2][Cross-config][Attention] Integrate the Qwen3-8B TP=2 CP=2 BF16 Attention setting to cross-config alignment framework

.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:

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 |

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).

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.