From f2c6acb5b59f3051eea646bedc831a6b99fd5628 Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 4 Aug 2026 22:20:11 +0800 Subject: [PATCH 01/17] feat(alignment): bind rollout and training attention contracts (#235 PR4) Wire the CP attention path into the cross-configuration planner/runtime for the Qwen3-8B TP=2 CP=2 BF16 target. The PR4 criterion "rollout and training descriptors bind to the same semantic attention contract" cannot hold literally: training runs full-sequence prefill over a CP-sharded sequence while rollout runs vLLM paged-KV chunked prefill, so the two AttentionContract instances always differ. Binding is therefore split into three tiers -- identity must match bit for bit, reduction semantics must match each other and the WS2 mandate, and materialization differences are recorded and measured rather than rejected. reduction.engine stays in the recorded tier so a Transformer Engine merge oracle on one side does not fail the binding; reduction.order and acc_dtype stay in the semantic tier because that is the WS2 claim. Also adds the first two framework-shaped RuntimeMaterializer implementations. Before this the only one was CpuSmokeMaterializer over a synthetic CPU model, and every named scenario was planning-only. Neither adapter imports megatron or vllm, so the binding rules run on CPU in CI. Determinism is probed on both sides and compared, because the two frameworks mean different things by it: Megatron asserts NCCL_ALGO and leaves TF32 and BF16 reduced-precision reduction unmanaged, while vLLM hard-sets ten NCCL variables and disables both. Mismatches in NCCL_ALGO, NCCL_PROTO and CUBLAS_WORKSPACE_CONFIG are blocking; the rest are recorded. Fixes a latent break on the way: the planner normalizes dtype knobs to torch spellings (bfloat16) while AttentionDType uses short ones (bf16), so passing a normalized knob into the enum raised. Stacked on #236 (attention contract) and #238 (deterministic CP reference), on top of #230 (cross-configuration framework). Part of #235 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Q3Ar3z9fHEBFQQHddSEMaw --- .../ws2-attention-cross-config-integration.md | 135 ++++ ...config_qwen3_8b_megatron_tp2_cp2_vllm.json | 87 +++ .../cross_config/adapters/__init__.py | 33 + .../cross_config/adapters/_common.py | 263 ++++++++ .../alignment/cross_config/adapters/knobs.py | 160 +++++ .../cross_config/adapters/megatron.py | 381 +++++++++++ .../alignment/cross_config/adapters/vllm.py | 442 +++++++++++++ .../cross_config/attention_binding.py | 528 +++++++++++++++ .../alignment/cross_config/determinism.py | 305 +++++++++ tests/test_attention_cross_config_binding.py | 612 ++++++++++++++++++ 10 files changed, 2946 insertions(+) create mode 100644 docs/design/ws2-attention-cross-config-integration.md create mode 100644 examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json create mode 100644 rl_engine/alignment/cross_config/adapters/__init__.py create mode 100644 rl_engine/alignment/cross_config/adapters/_common.py create mode 100644 rl_engine/alignment/cross_config/adapters/knobs.py create mode 100644 rl_engine/alignment/cross_config/adapters/megatron.py create mode 100644 rl_engine/alignment/cross_config/adapters/vllm.py create mode 100644 rl_engine/alignment/cross_config/attention_binding.py create mode 100644 rl_engine/alignment/cross_config/determinism.py create mode 100644 tests/test_attention_cross_config_binding.py diff --git a/docs/design/ws2-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md new file mode 100644 index 00000000..06966ebe --- /dev/null +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -0,0 +1,135 @@ +# WS2 Attention Cross-Configuration Integration + +Implements PR4 of [#235](https://github.com/RL-Align/RL-Kernel/issues/235): wiring the +CP attention path into the cross-configuration planner/runtime for the Qwen3-8B +TP=2 CP=2 BF16 target. + +Builds on [#236](https://github.com/RL-Align/RL-Kernel/pull/236) (attention contract +and dispatch metadata), [#238](https://github.com/RL-Align/RL-Kernel/pull/238) +(deterministic CP reference) and [#230](https://github.com/RL-Align/RL-Kernel/pull/230) +(cross-configuration framework). + +## What "bind to the same contract" means here + +The PR4 acceptance criteria say rollout and training descriptors must "bind to the +same semantic attention contract". Under the frozen deployment the two sides can +never produce identical `AttentionContract` instances: + +| | training (Megatron) | rollout (vLLM) | +| --- | --- | --- | +| mode | full-sequence prefill | chunked prefill, later decode | +| CP | `context_parallel_size`, whole forward | `prefill_context_parallel_size`, prefill only | +| KV | no paging | paged KV with a block table | +| backend vocabulary | `AttnBackend{flash,fused,unfused,local,auto}` | `AttentionBackendEnum` | + +Read literally, the criterion is unsatisfiable. It is therefore implemented as three +tiers, in `rl_engine/alignment/cross_config/attention_binding.py`: + +| tier | fields | rule | failure | +| --- | --- | --- | --- | +| `IDENTICAL` | checkpoint, model version, weight version, tokenizer, token ids, active mask, position ids, padding side, pre-update state, Q/KV heads, head dim, RoPE theta/scaling/rotary dim, QK-Norm, cached global token positions, KV sequence lengths | equal bit for bit | `comparable=False`; no drift number from the pair means anything | +| `SEMANTIC` | `reduction.merge`, `reduction.acc_dtype`, `reduction.order`, `reduction.downcast_at`, `export_lse`, cross-side determinism mode | both sides equal **and** equal to the WS2 mandate | fail closed | +| `RECORDED` | `mode`, `backend_id`, `reduction.engine`, RoPE materialization state and fusion boundary, KV-cache paging, CP/TP world sizes, local sequence length | free to differ | none; recorded into provenance and measured | + +Two placements are load-bearing: + +* **`reduction.engine` is `RECORDED`, not `SEMANTIC`.** Training may run the in-op + deterministic reference while rollout runs a Transformer Engine merge oracle. + Forcing them equal would defeat the oracle comparison that #235 PR2/PR3/PR5/PR6 + depend on. +* **`reduction.order` and `reduction.acc_dtype` are `SEMANTIC`.** This is the entire + WS2 claim: merge order and accumulation precision are decided by the contract, not + by whichever backend happens to be selected. + +`comparable` and `passed` are separate flags. A pair with mismatched identity is not +comparable. A pair that is comparable but violates the reduction mandate is still +rejected -- the drift would be real but attributable to the wrong thing. + +## Determinism is not one thing + +`rl_engine/alignment/cross_config/determinism.py` probes both sides and compares +them, because the two frameworks mean different things by "deterministic": + +| | Megatron `deterministic_mode` | vLLM `VLLM_BATCH_INVARIANT` | +| --- | --- | --- | +| `NCCL_ALGO` | asserts membership in a five-value set | hard-sets `allreduce:tree` | +| `NCCL_PROTO`, channels, threads | not managed | hard-set (`Simple`, `1`, `1`) | +| TF32 | **not managed at all** | disabled (`fp32_precision="ieee"`) | +| BF16 reduced-precision reduction | not managed | disabled | +| cuBLAS workspace / BLAS library | not managed | `:4096:8`, cuBLASLt | +| GEMM | cuBLAS / TE | Triton `matmul_persistent` | +| FlashAttention | forbidden | permitted | + +`NCCL_ALGO`, `NCCL_PROTO` and `CUBLAS_WORKSPACE_CONFIG` change arithmetic, so a +mismatch there is blocking. The remaining differences -- including the TF32 and +BF16-reduction asymmetry, which under a pure BF16 GEMM path does not fire -- are +recorded so the asymmetry appears in every artifact rather than being invisible. + +## Runtime adapters + +Before this PR the only `RuntimeMaterializer` in the repository was +`CpuSmokeMaterializer` over a synthetic CPU model, and every named scenario +(`S1`/`S2`/`S3`) was planning-only. This PR adds the first two framework-shaped +adapters: + +* `adapters/megatron.py` -- `MegatronProvenanceAdapter` (construction and + distributed-context fingerprints, determinism probe, frozen-scope assertions) and + `MegatronAttentionMaterializer`. +* `adapters/vllm.py` -- `VllmProvenanceAdapter` (adds `kv_page_size` from + `CacheConfig.block_size` and `split_kv_policy` from + `AttentionConfig.flash_attn_max_num_splits_for_cuda_graph`) and + `VllmRolloutMaterializer`. + +Neither module imports `megatron` or `vllm`; configs are duck-typed, so the binding +rules are exercised on CPU in CI rather than only on a 2-node cluster. + +## Fail closed, never substitute + +`unsupported_reduction_reason` rejects requests that #236 cannot express, instead of +collapsing them onto the supported value: + +| request | status | why | +| --- | --- | --- | +| `attention.reduction_order=arrival` | `UNSUPPORTED` | the control group must stay distinguishable from the treatment | +| `attention.reduction_downcast_at=per_block` | `UNSUPPORTED` | `DowncastPoint` declares only `final_write` | +| `attention.reduction_engine=te_oracle` | `UNSUPPORTED` | the TE merge oracle lands in #235 PR2/PR3; PR4's TE plan is provenance only | +| `attention.reduction_acc_dtype=bf16` | `UNSUPPORTED` | the CP `(out, lse)` merge accumulates in FP32 | +| `rollout.context_parallel_size>1` with `mode=decode` | `FALLBACK` | vLLM CP covers prefill only; recorded with the reason | + +## Knobs + +`adapters/knobs.py` extends `V1_KNOBS` additively. Added: training-side +`tensor_parallel_size` / `context_parallel_size` / `deterministic_mode` / +`cp_comm_type`, `rollout.batch_invariant` / `rollout.kv_block_size`, and the +reduction axis (`acc_dtype`, `order`, `downcast_at`, `engine`) plus +`attention.fusion_boundary` and `attention.split_kv_policy`. + +`training.attention_backend` keeps its path but its value domain is replaced with +Megatron's `AttnBackend`; the HuggingFace names have no Megatron counterpart, so this +is a replacement rather than a mapping. + +Not done here, because both change `V1_KNOBS` itself and would break existing +cross-config tests: removing `training.sharding` (Megatron has no such concept, and +DP=1 makes it moot) and renaming `rollout.context_parallel_size` to reflect that it +binds to `prefill_context_parallel_size`. + +## Scenario + +`examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json` supersedes +`cross_config_s1_distributed_smoke.json` and +`cross_config_s3_qwen3_8b_tp4_cp4_bf16.json`, whose training sides used `sdpa` / +`flash_attention_2` and `sharding: fsdp` -- none of which exist under Megatron -- and +whose TP=4/CP=4 topology does not match the target. +`cross_config_s2_vllm_tp_vs_fsdp.json` has no Megatron-only counterpart and should be +retired rather than rewritten. + +## Out of scope + +Deliberately not in this PR: + +* launching `torchrun`, initializing process groups, or executing attention; +* decode-mode materialization, which needs the validated `KVCacheSpec` from #235 PR6 + and is refused with that reference rather than stubbed; +* Transformer Engine calls of any kind (PR4's TE plan is policy and provenance only); +* distributed drift benchmarks and report artifacts (#235 PR5); +* fused production backend alignment (#235 PR7) and backward (#235 PR8). diff --git a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json new file mode 100644 index 00000000..e74d3e6c --- /dev/null +++ b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json @@ -0,0 +1,87 @@ +{ + "experiment_id": "ws2-qwen3-8b-attention-tp2-cp2", + "scenario_id": "qwen3_8b_megatron_tp2_cp2_vllm", + "contract_source": "ws1", + "contract_version": "current", + "strategy": "one_at_a_time", + "strict_fallback": true, + "scenario": { + "issue": "https://github.com/RL-Align/RL-Kernel/issues/235", + "pull_request": "PR4 -- cross-config integration", + "model": "Qwen3-8B dense", + "training_framework": "megatron", + "rollout_framework": "vllm", + "topology": "2 nodes x 2 GPUs, TP=2 CP=2 PP=1 DP=1, BF16, SM90", + "notes": [ + "Supersedes cross_config_s1_distributed_smoke.json and", + "cross_config_s3_qwen3_8b_tp4_cp4_bf16.json, whose training side used", + "HuggingFace attention backends and FSDP sharding. Neither exists in", + "Megatron, and DP=1 makes the sharding knob meaningless.", + "cross_config_s2_vllm_tp_vs_fsdp.json has no Megatron-only counterpart at", + "all and should be retired rather than rewritten.", + "rollout.context_parallel_size binds to vLLM", + "ParallelConfig.prefill_context_parallel_size and therefore applies to", + "prefill only; a decode-mode contract runs at CP=1." + ] + }, + "baseline": { + "batch": { + "size": 2 + }, + "rollout": { + "tensor_parallel_size": 2, + "context_parallel_size": 1, + "dtype": "bfloat16", + "enable_prefix_caching": false, + "enforce_eager": true, + "batch_invariant": true, + "kv_block_size": 16 + }, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "attention_backend": "unfused", + "compute_dtype": "bfloat16", + "deterministic_mode": true, + "cp_comm_type": "p2p", + "sharding": "unsharded" + }, + "attention": { + "reduction_acc_dtype": "fp32", + "reduction_order": "global_block_index", + "reduction_downcast_at": "final_write", + "reduction_engine": "in_op_reference", + "fusion_boundary": "unfused_rope_attention", + "split_kv_policy": 32 + }, + "logp": { + "backend": "native" + } + }, + "interventions": [ + { + "path": "training.context_parallel_size", + "values": [1, 2] + }, + { + "path": "training.tensor_parallel_size", + "values": [1, 2] + }, + { + "path": "attention.fusion_boundary", + "values": ["unfused_rope_attention", "fused_rope_attention"] + }, + { + "path": "training.cp_comm_type", + "values": ["p2p", "all_gather"] + }, + { + "path": "attention.reduction_order", + "values": ["global_block_index", "arrival"] + }, + { + "path": "attention.reduction_acc_dtype", + "values": ["fp32", "bf16"] + } + ] +} diff --git a/rl_engine/alignment/cross_config/adapters/__init__.py b/rl_engine/alignment/cross_config/adapters/__init__.py new file mode 100644 index 00000000..c2db2134 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/__init__.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime adapters for the WS2 Qwen3-8B Megatron + vLLM cross-config target.""" + +from rl_engine.alignment.cross_config.adapters._common import QWEN3_8B, Qwen3ModelSpec +from rl_engine.alignment.cross_config.adapters.knobs import ( + MEGATRON_ATTENTION_BACKENDS, + WS2_ATTENTION_KNOB_DESCRIPTORS, + WS2_ATTENTION_KNOBS, + WS2_ATTENTION_NORMALIZERS, +) +from rl_engine.alignment.cross_config.adapters.megatron import ( + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, +) +from rl_engine.alignment.cross_config.adapters.vllm import ( + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", + "QWEN3_8B", + "Qwen3ModelSpec", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] diff --git a/rl_engine/alignment/cross_config/adapters/_common.py b/rl_engine/alignment/cross_config/adapters/_common.py new file mode 100644 index 00000000..33b4d91c --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/_common.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared pieces for the Megatron and vLLM WS2 attention adapters.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from rl_engine.alignment.cross_config.runtime import KnobApplication +from rl_engine.alignment.cross_config.schema import ( + IsolationScope, + KnobDescriptor, + MaterializationStatus, +) +from rl_engine.kernels.attention_contract import ( + AttentionDType, + AttentionMerge, + DowncastPoint, + ReductionEngine, + ReductionOrder, + ReductionSpec, + ShardingSpec, +) + +__all__ = [ + "QWEN3_8B", + "Qwen3ModelSpec", + "application", + "attention_dtype", + "build_reduction_spec", + "build_sharding_spec", + "causal_offsets_for", + "flatten", + "unsupported_reduction_reason", +] + + +@dataclass(frozen=True) +class Qwen3ModelSpec: + """Architecture constants for the frozen dense target. + + These are *not* knobs. #235/#239/#241 all fix Qwen3-8B dense, so they belong to + the scenario, and both sides must agree on them or the comparison is void. + """ + + name: str = "qwen3-8b" + hidden_size: int = 4096 + ffn_hidden_size: int = 12288 + num_layers: int = 36 + q_heads: int = 32 + kv_heads: int = 8 + head_dim: int = 128 + real_vocab_size: int = 151936 + rope_theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + qk_layernorm: bool = True + + def identity_fields(self) -> dict[str, Any]: + """The subset of :data:`IDENTITY_FIELDS` this spec is responsible for.""" + + return { + "q_heads": self.q_heads, + "kv_heads": self.kv_heads, + "head_dim": self.head_dim, + "rope_theta": self.rope_theta, + "rope_scaling": self.rope_scaling, + "rotary_dim": self.rotary_dim, + "qk_layernorm": self.qk_layernorm, + } + + +QWEN3_8B = Qwen3ModelSpec() + + +#: The planner normalizes dtype knobs to torch spellings (``bfloat16``), while +#: :class:`AttentionDType` uses short spellings (``bf16``). Passing a normalized knob +#: straight into the enum raises, so every adapter must translate here rather than +#: each inventing its own mapping. +_DTYPE_ALIASES: Mapping[str, AttentionDType] = { + "bf16": AttentionDType.BF16, + "bfloat16": AttentionDType.BF16, + "fp16": AttentionDType.FP16, + "float16": AttentionDType.FP16, + "half": AttentionDType.FP16, + "fp32": AttentionDType.FP32, + "float32": AttentionDType.FP32, + "float": AttentionDType.FP32, +} + + +def attention_dtype(value: Any, *, field: str) -> AttentionDType: + """Translate a normalized knob dtype into an :class:`AttentionDType`.""" + + if isinstance(value, AttentionDType): + return value + key = str(value).strip().lower().replace("torch.", "") + try: + return _DTYPE_ALIASES[key] + except KeyError as exc: + raise ValueError( + f"{field}={value!r} is not a supported attention dtype; " + f"expected one of {sorted(set(_DTYPE_ALIASES))}" + ) from exc + + +def flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: + """Flatten nested knob mappings into dotted paths.""" + + flat: dict[str, Any] = {} + for key, child in value.items(): + path = f"{prefix}{key}" + if isinstance(child, Mapping): + flat.update(flatten(child, f"{path}.")) + else: + flat[path] = child + return flat + + +def application( + descriptor: KnobDescriptor, + requested: Any, + materialized: Any, + actual: Any, + status: MaterializationStatus, + reason: str, + **evidence: Any, +) -> KnobApplication: + return KnobApplication( + path=descriptor.path, + requested=requested, + materialized=materialized, + actual=actual, + lifecycle=descriptor.lifecycle, + status=status, + evidence={"reason": reason, **evidence}, + critical=descriptor.critical, + ) + + +def unsupported_reduction_reason(flat: Mapping[str, Any]) -> str | None: + """Return why the requested reduction cannot be materialized, if it cannot. + + #236 declares single-member enums for merge order, downcast point and reduction + engine, so the alternative knob values exist only as control groups. Requesting + one must fail loudly rather than quietly collapse onto the supported value -- + silently substituting ``global_block_index`` for a requested ``arrival`` would + make the control group indistinguishable from the treatment. + """ + + order = flat.get("attention.reduction_order") + if order is not None and order != ReductionOrder.GLOBAL_BLOCK_INDEX.value: + return ( + f"attention.reduction_order={order!r} has no backend; #236 ReductionOrder " + f"declares only {ReductionOrder.GLOBAL_BLOCK_INDEX.value!r}" + ) + downcast = flat.get("attention.reduction_downcast_at") + if downcast is not None and downcast != DowncastPoint.FINAL_WRITE.value: + return ( + f"attention.reduction_downcast_at={downcast!r} has no backend; #236 " + f"DowncastPoint declares only {DowncastPoint.FINAL_WRITE.value!r}" + ) + engine = flat.get("attention.reduction_engine") + if engine is not None and engine != ReductionEngine.IN_OP_REFERENCE.value: + return ( + f"attention.reduction_engine={engine!r} has no backend; the Transformer " + "Engine merge oracle lands in #235 PR2/PR3, not here" + ) + acc_dtype = flat.get("attention.reduction_acc_dtype") + if ( + acc_dtype is not None + and attention_dtype(acc_dtype, field="attention.reduction_acc_dtype") + is not AttentionDType.FP32 + ): + return ( + f"attention.reduction_acc_dtype={acc_dtype!r} violates the WS2 mandate; " + "the CP (out, lse) merge accumulates in fp32" + ) + return None + + +def build_reduction_spec(flat: Mapping[str, Any]) -> ReductionSpec: + """Build the reduction spec, having already rejected unsupported requests.""" + + return ReductionSpec( + merge=AttentionMerge.ONLINE_SOFTMAX_LSE, + acc_dtype=AttentionDType.FP32, + order=ReductionOrder.GLOBAL_BLOCK_INDEX, + downcast_at=DowncastPoint.FINAL_WRITE, + engine=ReductionEngine.IN_OP_REFERENCE, + ) + + +def build_sharding_spec( + *, + model: Qwen3ModelSpec, + tp_rank: int, + tp_world_size: int, + cp_rank: int, + cp_world_size: int, + global_sequence_length: int, +) -> ShardingSpec: + """Build a CP/TP sharding spec for one rank of the frozen layout. + + TP splits heads, CP splits the sequence. The #239 rank layout fixes + ``rank = cp_rank * tp_world_size + tp_rank`` for a 2-node x 2-GPU deployment, + but nothing here depends on that mapping: ownership is derived from the ranks + themselves so the same builder serves CP=1 baselines. + """ + + if model.q_heads % tp_world_size or model.kv_heads % tp_world_size: + raise ValueError( + f"Qwen3 GQA heads ({model.q_heads}/{model.kv_heads}) must divide evenly " + f"across tp_world_size={tp_world_size}" + ) + if global_sequence_length % cp_world_size: + raise ValueError( + f"global_sequence_length={global_sequence_length} must divide evenly " + f"across cp_world_size={cp_world_size}" + ) + + local_q_heads = model.q_heads // tp_world_size + local_kv_heads = model.kv_heads // tp_world_size + local_sequence_length = global_sequence_length // cp_world_size + + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=model.q_heads, + global_kv_heads=model.kv_heads, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=global_sequence_length, + local_sequence_length=local_sequence_length, + # One contiguous CP block per rank. The merge order key is the global block + # index, never the arrival order of the CP exchange. + global_block_indices=(cp_rank,), + global_block_token_starts=(cp_rank * local_sequence_length,), + local_block_offsets=(0, local_sequence_length), + ) + + +def causal_offsets_for(sharding: ShardingSpec, batch_size: int) -> tuple[int, ...]: + """Causal offsets for one CP shard, one entry per batch entry. + + Under CP the local query block does not start at global position zero, so the + causal mask has to be shifted by the number of preceding global tokens. Taking + that from ``global_block_token_starts`` rather than recomputing + ``cp_rank * local_sequence_length`` keeps uneven CP splits correct. + """ + + offset = sharding.global_block_token_starts[0] + return (offset,) * batch_size + + +_PROCESS_SCOPES = (IsolationScope.PROCESS, IsolationScope.DISTRIBUTED_CONTEXT) diff --git a/rl_engine/alignment/cross_config/adapters/knobs.py b/rl_engine/alignment/cross_config/adapters/knobs.py new file mode 100644 index 00000000..c1cf0b84 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/knobs.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 attention knobs for the Qwen3-8B TP=2 CP=2 Megatron + vLLM target. + +``V1_KNOBS`` was written against a HuggingFace/FSDP rollout-vs-training pair. Three +of its entries do not survive contact with the frozen Megatron + vLLM target: + +* ``training.sharding`` takes ``unsharded``/``fsdp``, neither of which exists in + Megatron, and is meaningless at DP=1 anyway; +* ``training.attention_backend`` takes HuggingFace names + (``flash_attention_2``/``sdpa``/``eager``/``model_default``) while Megatron's + ``AttnBackend`` is ``flash``/``fused``/``unfused``/``local``/``auto``; +* there is no training-side ``tensor_parallel_size`` or ``context_parallel_size`` + at all, so the target configuration cannot even be expressed. + +This module is deliberately **additive**: it extends ``V1_KNOBS`` rather than +editing it, and overrides only the normalizer for ``training.attention_backend``. +Deleting the two dead knobs changes ``V1_KNOBS`` itself and would break existing +cross-config tests, so it is left to a follow-up on the framework PR. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +from rl_engine.alignment.cross_config.planner import ( + _NORMALIZERS, + V1_KNOBS, + Normalizer, + _normalize_choice, + _positive_int, + _strict_bool, +) +from rl_engine.alignment.cross_config.schema import IsolationScope, KnobDescriptor + +__all__ = [ + "MEGATRON_ATTENTION_BACKENDS", + "WS2_ATTENTION_KNOBS", + "WS2_ATTENTION_KNOB_DESCRIPTORS", + "WS2_ATTENTION_NORMALIZERS", +] + + +#: ``megatron.core.transformer.enums.AttnBackend``. +MEGATRON_ATTENTION_BACKENDS: tuple[str, ...] = ( + "flash", + "fused", + "unfused", + "local", + "auto", +) + + +WS2_ATTENTION_KNOB_DESCRIPTORS: tuple[KnobDescriptor, ...] = ( + # -- training-side parallelism: the target configuration itself ------------ + KnobDescriptor( + "training.tensor_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "training.context_parallel_size", + IsolationScope.PROCESS, + ("training",), + ), + # -- determinism switches, one per framework ------------------------------ + KnobDescriptor( + "training.deterministic_mode", + IsolationScope.PROCESS, + ("training",), + ), + KnobDescriptor( + "rollout.batch_invariant", + IsolationScope.PROCESS, + ("rollout",), + ), + # -- reduction knobs: the "turn the noise sources on and off" axis --------- + # These are what make drift attributable. ``reduction.order=arrival`` in + # particular is a control group, not a supported production value. + KnobDescriptor( + "attention.reduction_acc_dtype", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("fp32", "bf16"), + ), + KnobDescriptor( + "attention.reduction_order", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("global_block_index", "arrival"), + ), + KnobDescriptor( + "attention.reduction_downcast_at", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("final_write", "per_block"), + ), + KnobDescriptor( + "attention.reduction_engine", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("in_op_reference", "te_oracle"), + ), + # -- materialization knobs: differences the experiment measures ------------ + KnobDescriptor( + "attention.fusion_boundary", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout", "training"), + allowed_values=("unfused_rope_attention", "fused_rope_attention"), + ), + KnobDescriptor( + # vLLM: AttentionConfig.flash_attn_max_num_splits_for_cuda_graph + "attention.split_kv_policy", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + KnobDescriptor( + # vLLM: CacheConfig.block_size -> AttentionContract.kv_cache.page_size + "rollout.kv_block_size", + IsolationScope.ENGINE_CONSTRUCTION, + ("rollout",), + ), + # The CP communication group cannot be reconfigured once built; it is bound to + # the distributed context, not merely to engine construction. + KnobDescriptor( + "training.cp_comm_type", + IsolationScope.DISTRIBUTED_CONTEXT, + ("training",), + allowed_values=("p2p", "all_gather", "a2a", "a2a+p2p"), + ), +) + + +WS2_ATTENTION_KNOBS: Mapping[str, KnobDescriptor] = { + **V1_KNOBS, + **{descriptor.path: descriptor for descriptor in WS2_ATTENTION_KNOB_DESCRIPTORS}, +} + + +WS2_ATTENTION_NORMALIZERS: Mapping[str, Normalizer] = { + **_NORMALIZERS, + # Replace, not map: the HuggingFace names have no Megatron counterpart. + "training.attention_backend": _normalize_choice(*MEGATRON_ATTENTION_BACKENDS), + "training.tensor_parallel_size": _positive_int, + "training.context_parallel_size": _positive_int, + "training.deterministic_mode": _strict_bool, + "rollout.batch_invariant": _strict_bool, + # AttentionDType values, not torch dtype names -- these feed ReductionSpec directly. + "attention.reduction_acc_dtype": _normalize_choice("fp32", "bf16"), + "attention.reduction_order": _normalize_choice("global_block_index", "arrival"), + "attention.reduction_downcast_at": _normalize_choice("final_write", "per_block"), + "attention.reduction_engine": _normalize_choice("in_op_reference", "te_oracle"), + "attention.fusion_boundary": _normalize_choice( + "unfused_rope_attention", "fused_rope_attention" + ), + "attention.split_kv_policy": _positive_int, + "rollout.kv_block_size": _positive_int, + "training.cp_comm_type": _normalize_choice("p2p", "all_gather", "a2a", "a2a+p2p"), +} diff --git a/rl_engine/alignment/cross_config/adapters/megatron.py b/rl_engine/alignment/cross_config/adapters/megatron.py new file mode 100644 index 00000000..8b96b159 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -0,0 +1,381 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Training-side (Megatron) runtime adapter for WS2 attention cross-config. + +Two things live here: + +``MegatronProvenanceAdapter`` + Read-only. Turns a Megatron config object into the construction and + distributed-context fingerprints the cross-config framework already expects, + plus the determinism probe. It never imports ``megatron`` -- every accessor is + duck-typed -- so this module is importable and testable on a laptop. + +``MegatronAttentionMaterializer`` + Implements the ``RuntimeMaterializer`` protocol. Before this PR the only + implementation in the repository was ``CpuSmokeMaterializer`` over a synthetic + CPU model, so nothing had ever materialized a real distributed runtime. + +Scope boundary: materialization builds and validates the training-side +:class:`AttentionContract` and reports what would be constructed. It does not +launch ``torchrun``, initialize process groups, or execute attention. Binding a +constructed Megatron model to this contract is the next step and needs the 2-node +x 2-GPU environment that #239 fixes. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import ( + DeterminismProbe, + megatron_probe_from_config, +) +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "MEGATRON_CONSTRUCTION_KEYS", + "MEGATRON_DISTRIBUTED_KEYS", + "MegatronAttentionMaterializer", + "MegatronProvenanceAdapter", +] + + +#: ``TransformerConfig`` fields that change attention arithmetic. Hashed into the +#: construction fingerprint. Deliberately excludes MoE, Mamba, MLA and sparse +#: attention fields: the frozen target is Qwen3-8B dense, and those are asserted +#: off rather than recorded. +MEGATRON_CONSTRUCTION_KEYS: tuple[str, ...] = ( + "attention_backend", + "attention_softmax_in_fp32", + "apply_query_key_layer_scaling", + "apply_rope_fusion", + "masked_softmax_fusion", + "bias_activation_fusion", + "bias_dropout_fusion", + "gradient_accumulation_fusion", + "cross_entropy_loss_fusion", + "cross_entropy_fusion_impl", + "recompute_granularity", + "recompute_method", + "recompute_num_layers", + "recompute_modules", + "rotary_base", + "rotary_percent", + "rotary_interleaved", + "rotary_scaling_factor", + "qk_layernorm", + "hidden_dropout", + "attention_dropout", + "params_dtype", + "bf16", + "fp16", + "fp8", + "deterministic_mode", +) + + +#: ``ModelParallelConfig`` fields that define the distributed context. +MEGATRON_DISTRIBUTED_KEYS: tuple[str, ...] = ( + "tensor_model_parallel_size", + "pipeline_model_parallel_size", + "virtual_pipeline_model_parallel_size", + "context_parallel_size", + "hierarchical_context_parallel_sizes", + "expert_model_parallel_size", + "expert_tensor_parallel_size", + "sequence_parallel", + "cp_comm_type", + "tp_comm_overlap", + "use_te_rng_tracker", +) + + +#: Fields that must hold these values for the frozen dense target. A mismatch is a +#: hard stop, not a recorded difference -- see the exclusion list in the WS2 scope. +MEGATRON_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "pipeline_model_parallel_size": 1, + "expert_model_parallel_size": 1, + "sequence_parallel": False, + "fp8": None, + "hidden_dropout": 0.0, + "attention_dropout": 0.0, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class MegatronProvenanceAdapter: + """Extract fingerprints and determinism evidence from a Megatron config. + + ``config`` may be a real ``TransformerConfig``/``ModelParallelConfig``, a merged + namespace, or a test double. Missing attributes read as ``None`` and are + recorded as such rather than raising: an absent field is itself provenance. + """ + + framework = "megatron" + + def __init__(self, config: Any, *, env: Optional[Mapping[str, str]] = None): + self.config = config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_CONSTRUCTION_KEYS} + + def distributed_view(self) -> dict[str, Any]: + return {name: _value(self.config, name) for name in MEGATRON_DISTRIBUTED_KEYS} + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return megatron_probe_from_config(self.config, env=self.env) + + def frozen_scope_violations(self) -> tuple[str, ...]: + """Return the frozen-scope assertions this config violates.""" + + violations: list[str] = [] + for name, expected in MEGATRON_FROZEN_ASSERTIONS.items(): + actual = _value(self.config, name) + if actual is None: + # Not declared. Treated as unknown rather than as satisfied, because + # a silently-absent MoE or FP8 setting is exactly the case that would + # otherwise slip past a dense-only claim. + violations.append(f"{name} is not declared (expected {expected!r})") + elif actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "determinism": self.determinism_probe().to_dict(), + } + + +class MegatronAttentionMaterializer: + """Materialize the training-side attention runtime for the WS2 target.""" + + runtime_kind = "megatron_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + backend_id: str = "rlkernel.cp_attention_reference", + provenance: Optional[MegatronProvenanceAdapter] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.backend_id = backend_id + self.provenance = provenance + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + """Build the training-side contract. Raises on an unusable request.""" + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=attention_dtype( + flat.get("training.compute_dtype", "bf16"), field="training.compute_dtype" + ), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "training" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + applications.append( + application( + descriptor, + requested, + requested, + requested, + MaterializationStatus.APPLIED, + "bound to the training-side attention contract", + frozen_scope_violations=list(scope_violations), + ) + ) + + tp_world_size = int(flat.get("training.tensor_parallel_size", 1)) + cp_world_size = int(flat.get("training.context_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "megatron", + "attention_backend": flat.get("training.attention_backend"), + "compute_dtype": flat.get("training.compute_dtype"), + "deterministic_mode": flat.get("training.deterministic_mode"), + "cp_comm_type": flat.get("training.cp_comm_type"), + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"training": side_config, "rollout": {}}, + topology={ + "training": { + "world_size": tp_world_size * cp_world_size, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": cp_world_size, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "rollout": {"world_size": 1}, + }, + scorer={ + "mode": "teacher_forcing", + "framework": "megatron", + "export_lse": True, + }, + operator_backends={ + "training": self.backend_id, + "rollout": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) diff --git a/rl_engine/alignment/cross_config/adapters/vllm.py b/rl_engine/alignment/cross_config/adapters/vllm.py new file mode 100644 index 00000000..a8de61c9 --- /dev/null +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -0,0 +1,442 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rollout-side (vLLM) runtime adapter for WS2 attention cross-config. + +Mirrors :mod:`.megatron`, with three differences that come straight from what vLLM +actually is: + +* vLLM's context parallelism is ``prefill_context_parallel_size`` -- it applies to + prefill only, so a decode-mode contract must declare ``cp_world_size == 1`` + regardless of what the prefill knob says. +* ``CacheConfig.block_size`` is the paged-KV page size, and it feeds + ``KVCacheSpec.page_size`` directly rather than being invented here. +* Determinism comes from the ``VLLM_BATCH_INVARIANT`` environment variable rather + than from a config field, because vLLM applies it inside + ``init_batch_invariance()`` at worker startup. + +Like the Megatron adapter, nothing here imports ``vllm``; configs are duck-typed so +the module is importable anywhere. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from typing import Any, Optional + +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + Qwen3ModelSpec, + application, + attention_dtype, + build_reduction_spec, + build_sharding_spec, + causal_offsets_for, + flatten, + unsupported_reduction_reason, +) +from rl_engine.alignment.cross_config.determinism import DeterminismProbe, vllm_probe_from_env +from rl_engine.alignment.cross_config.runtime import ( + AdapterMaterialization, + KnobApplication, + RuntimeBinding, +) +from rl_engine.alignment.cross_config.schema import KnobDescriptor, MaterializationStatus +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionMode, + AttentionRole, + RoPEFusionBoundary, + RoPESpec, + RoPEState, +) +from rl_engine.kernels.semantic_registry import implementation_fingerprint + +__all__ = [ + "VLLM_ATTENTION_KEYS", + "VLLM_CACHE_KEYS", + "VLLM_FROZEN_ASSERTIONS", + "VLLM_MODEL_KEYS", + "VLLM_PARALLEL_KEYS", + "VllmProvenanceAdapter", + "VllmRolloutMaterializer", +] + + +VLLM_MODEL_KEYS: tuple[str, ...] = ( + "dtype", + "seed", + "quantization", + "enforce_eager", + "max_logprobs", + "disable_cascade_attn", + "max_model_len", +) + +VLLM_CACHE_KEYS: tuple[str, ...] = ( + "block_size", + "cache_dtype", + "enable_prefix_caching", + "prefix_caching_hash_algo", + "calculate_kv_scales", + "sliding_window", +) + +VLLM_ATTENTION_KEYS: tuple[str, ...] = ( + "backend", + "flash_attn_version", + "use_prefill_decode_attention", + "flash_attn_max_num_splits_for_cuda_graph", + "use_cudnn_prefill", + "disable_flashinfer_prefill", + "use_non_causal", +) + +VLLM_PARALLEL_KEYS: tuple[str, ...] = ( + "tensor_parallel_size", + "pipeline_parallel_size", + "prefill_context_parallel_size", + "data_parallel_size", +) + + +#: Frozen dense-target assertions on the rollout side. ``cache_dtype`` must stay +#: ``auto`` because an FP8 KV cache is a representation-drift problem tracked +#: separately, and ``disable_cascade_attn`` must stay ``True`` because cascade +#: attention changes the block-merge structure the contract pins down. +VLLM_FROZEN_ASSERTIONS: Mapping[str, Any] = { + "quantization": None, + "cache_dtype": "auto", + "calculate_kv_scales": False, + "disable_cascade_attn": True, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + "sliding_window": None, +} + + +def _value(config: Any, name: str) -> Any: + raw = getattr(config, name, None) + return getattr(raw, "value", raw) if raw is not None else None + + +def _fingerprint(payload: Mapping[str, Any]) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +class VllmProvenanceAdapter: + """Extract fingerprints and determinism evidence from vLLM configs.""" + + framework = "vllm" + + def __init__( + self, + *, + model_config: Any = None, + cache_config: Any = None, + attention_config: Any = None, + parallel_config: Any = None, + env: Optional[Mapping[str, str]] = None, + ): + self.model_config = model_config + self.cache_config = cache_config + self.attention_config = attention_config + self.parallel_config = parallel_config + self.env = dict(env or {}) + + def construction_view(self) -> dict[str, Any]: + view: dict[str, Any] = {} + for prefix, config, keys in ( + ("model", self.model_config, VLLM_MODEL_KEYS), + ("cache", self.cache_config, VLLM_CACHE_KEYS), + ("attention", self.attention_config, VLLM_ATTENTION_KEYS), + ): + for name in keys: + view[f"{prefix}.{name}"] = _value(config, name) + return view + + def distributed_view(self) -> dict[str, Any]: + return { + f"parallel.{name}": _value(self.parallel_config, name) for name in VLLM_PARALLEL_KEYS + } + + @property + def construction_fingerprint(self) -> str: + return _fingerprint(self.construction_view()) + + @property + def distributed_context_fingerprint(self) -> str: + return _fingerprint(self.distributed_view()) + + def determinism_probe(self) -> DeterminismProbe: + return vllm_probe_from_env(self.env, model_config=self.model_config) + + def frozen_scope_violations(self) -> tuple[str, ...]: + sources = { + "quantization": self.model_config, + "disable_cascade_attn": self.model_config, + "cache_dtype": self.cache_config, + "calculate_kv_scales": self.cache_config, + "sliding_window": self.cache_config, + "pipeline_parallel_size": self.parallel_config, + "data_parallel_size": self.parallel_config, + } + violations: list[str] = [] + for name, expected in VLLM_FROZEN_ASSERTIONS.items(): + config = sources.get(name) + if config is None: + continue + actual = _value(config, name) + if actual != expected: + violations.append(f"{name}={actual!r} (expected {expected!r})") + return tuple(violations) + + @property + def kv_page_size(self) -> Optional[int]: + """vLLM's paged-KV block size, which is the contract's ``page_size``.""" + + block_size = _value(self.cache_config, "block_size") + return int(block_size) if block_size is not None else None + + @property + def split_kv_policy(self) -> Optional[int]: + """The split-KV knob #235 PR5/PR7 needs; #236 has no field for it yet.""" + + splits = _value(self.attention_config, "flash_attn_max_num_splits_for_cuda_graph") + return int(splits) if splits is not None else None + + def to_dict(self) -> dict[str, Any]: + return { + "framework": self.framework, + "construction": self.construction_view(), + "distributed_context": self.distributed_view(), + "construction_fingerprint": self.construction_fingerprint, + "distributed_context_fingerprint": self.distributed_context_fingerprint, + "frozen_scope_violations": list(self.frozen_scope_violations()), + "kv_page_size": self.kv_page_size, + "split_kv_policy": self.split_kv_policy, + "determinism": self.determinism_probe().to_dict(), + } + + +class VllmRolloutMaterializer: + """Materialize the rollout-side attention runtime for the WS2 target.""" + + runtime_kind = "vllm_attention" + + def __init__( + self, + *, + model: Qwen3ModelSpec = QWEN3_8B, + global_sequence_length: int = 4096, + tp_rank: int = 0, + cp_rank: int = 0, + mode: AttentionMode = AttentionMode.CHUNKED_PREFILL, + backend_id: str = "vllm.flash_attn", + provenance: Optional[VllmProvenanceAdapter] = None, + ): + self.model = model + self.global_sequence_length = global_sequence_length + self.tp_rank = tp_rank + self.cp_rank = cp_rank + self.mode = mode + self.backend_id = backend_id + self.provenance = provenance + + @property + def implementation_fingerprint(self) -> str: + return implementation_fingerprint( + type(self), + instance=self, + entrypoints=("materialize", "build_contract"), + ) + + def effective_cp_world_size(self, flat: Mapping[str, Any]) -> int: + """CP applies to prefill only; decode always runs at CP=1.""" + + requested = int(flat.get("rollout.context_parallel_size", 1)) + if self.mode is AttentionMode.DECODE: + return 1 + return requested + + def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: + if self.mode is AttentionMode.DECODE: + # Decode replay needs a validated KVCacheSpec (cache positions, page + # ownership, prefix-cache identity). That is #235 PR6's contract surface, + # and inventing a placeholder here would let an unvalidated decode case + # look bound. Fail instead. + raise AttentionContractError( + "decode-mode materialization requires KV-cache identity from #235 PR6; " + "this adapter covers prefill and chunked prefill" + ) + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + cp_world_size = self.effective_cp_world_size(flat) + sharding = build_sharding_spec( + model=self.model, + tp_rank=self.tp_rank, + tp_world_size=tp_world_size, + cp_rank=self.cp_rank if cp_world_size > 1 else 0, + cp_world_size=cp_world_size, + global_sequence_length=self.global_sequence_length, + ) + fusion = RoPEFusionBoundary( + flat.get( + "attention.fusion_boundary", + RoPEFusionBoundary.FUSED_ROPE_ATTENTION.value, + ) + ) + rope = RoPESpec( + q_state=RoPEState.POST_ROPE, + k_state=RoPEState.POST_ROPE, + # vLLM stores post-RoPE K in the cache; recorded, not asserted equal to + # the training side, because it is a materialization fact. + k_cache_state=RoPEState.POST_ROPE, + theta=self.model.rope_theta, + rotary_dim=self.model.rotary_dim, + rope_scaling=self.model.rope_scaling, + fusion_boundary=fusion, + ) + batch_size = int(flat.get("batch.size", 1)) + return AttentionContract( + role=AttentionRole.INFER, + mode=self.mode, + dtype=attention_dtype(flat.get("rollout.dtype", "bf16"), field="rollout.dtype"), + batch_size=batch_size, + query_sequence_length=sharding.local_sequence_length, + head_dim=self.model.head_dim, + causal=True, + causal_offsets=causal_offsets_for(sharding, batch_size), + sharding=sharding, + reduction=build_reduction_spec(flat), + rope=rope, + export_lse=True, + ) + + def materialize( + self, + normalized: Mapping[str, Any], + descriptors: Mapping[str, KnobDescriptor], + ) -> AdapterMaterialization: + flat = flatten(normalized) + applications: list[KnobApplication] = [] + + blocked = unsupported_reduction_reason(flat) + scope_violations = ( + self.provenance.frozen_scope_violations() if self.provenance is not None else () + ) + + contract: AttentionContract | None = None + contract_error: str | None = None + if blocked is None: + try: + contract = self.build_contract(flat) + except (AttentionContractError, ValueError) as exc: + contract_error = str(exc) + + requested_cp = int(flat.get("rollout.context_parallel_size", 1)) + effective_cp = self.effective_cp_world_size(flat) + + for path, requested in flat.items(): + descriptor = descriptors.get(path) + if descriptor is None or "rollout" not in descriptor.targets: + continue + if blocked is not None and path.startswith("attention.reduction"): + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.UNSUPPORTED, + blocked, + ) + ) + continue + if contract_error is not None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + contract_error, + ) + ) + continue + if path == "rollout.context_parallel_size" and effective_cp != requested_cp: + applications.append( + application( + descriptor, + requested, + effective_cp, + effective_cp, + MaterializationStatus.FALLBACK, + ( + "vLLM context parallelism covers prefill only; a decode-mode " + f"contract runs at cp_world_size=1, not {requested_cp}" + ), + vllm_field="ParallelConfig.prefill_context_parallel_size", + ) + ) + continue + applications.append( + application( + descriptor, + requested, + requested, + requested, + MaterializationStatus.APPLIED, + "bound to the rollout-side attention contract", + frozen_scope_violations=list(scope_violations), + ) + ) + + tp_world_size = int(flat.get("rollout.tensor_parallel_size", 1)) + side_config: dict[str, Any] = { + "framework": "vllm", + "dtype": flat.get("rollout.dtype"), + "enforce_eager": flat.get("rollout.enforce_eager"), + "enable_prefix_caching": flat.get("rollout.enable_prefix_caching"), + "batch_invariant": flat.get("rollout.batch_invariant"), + "kv_block_size": flat.get("rollout.kv_block_size"), + "split_kv_policy": flat.get("attention.split_kv_policy"), + "attention_mode": self.mode.value, + "contract": contract.to_dict() if contract is not None else None, + "contract_error": contract_error or blocked, + "frozen_scope_violations": list(scope_violations), + } + if self.provenance is not None: + side_config["provenance"] = self.provenance.to_dict() + + return AdapterMaterialization( + applications=tuple(applications), + binding=RuntimeBinding( + batch_size=int(flat.get("batch.size", 1)), + side_configs={"rollout": side_config, "training": {}}, + topology={ + "rollout": { + "world_size": tp_world_size * effective_cp, + "tensor_parallel_size": tp_world_size, + "context_parallel_size": effective_cp, + "pipeline_parallel_size": 1, + "data_parallel_size": 1, + }, + "training": {"world_size": 1}, + }, + scorer={ + "mode": "rollout_logprob", + "framework": "vllm", + "export_lse": True, + }, + operator_backends={ + "rollout": self.backend_id, + "training": self.backend_id, + }, + runtime_kind=self.runtime_kind, + ), + ) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py new file mode 100644 index 00000000..fcd4c203 --- /dev/null +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -0,0 +1,528 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Three-tier binding between rollout-side and training-side attention contracts. + +Issue #235 PR4 requires that "rollout and training descriptors bind to the same +semantic attention contract". Under the frozen Megatron + vLLM deployment the two +sides can never produce *identical* :class:`AttentionContract` instances: training +runs full-sequence prefill over a CP-sharded sequence, while rollout runs vLLM +paged-KV chunked prefill and decode. Taking "same contract" literally would make +the target configuration permanently unbindable. + +This module therefore splits binding into three tiers: + +``IDENTICAL`` + Logical identity. Both sides must agree bit for bit, otherwise the pair is not + comparable at all and no drift number from it means anything. + +``SEMANTIC`` + The WS2 numerical claim: merge semantics, accumulation dtype, reduction order + and downcast point are decided by the contract, not by the implementation. + Both sides must carry the same values *and* those values must match the WS2 + mandate, otherwise the comparison fails closed. + +``RECORDED`` + Materialization facts that the two sides are expected to differ on -- attention + mode, RoPE fusion boundary, KV-cache paging, backend id, reduction engine. These + differences are exactly what the experiment measures, so they are recorded into + provenance rather than rejected. + +Deliberately *not* in ``SEMANTIC``: ``engine``. Training may run the in-op +deterministic reference while rollout runs a Transformer Engine merge oracle; forcing +those equal would defeat the purpose of the oracle comparison in #235 PR2/3/5/6. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Optional + +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionDType, + AttentionMerge, + AttentionRole, + DowncastPoint, + ReductionOrder, +) + +__all__ = [ + "ATTENTION_LSE_DOMAIN", + "AttentionBindingError", + "AttentionBindingResult", + "BindingErrorCode", + "BindingIssue", + "BindingTier", + "IDENTITY_FIELDS", + "NULLABLE_IDENTITY_FIELDS", + "RECORDED_FIELDS", + "SEMANTIC_REDUCTION_FIELDS", + "WS2_ATTENTION_REDUCTION_MANDATE", + "bind_attention_contracts", + "first_blocking_issue", + "identity_fingerprint", + "summarize_binding", +] + + +class AttentionBindingError(ValueError): + """Raised when a caller supplies structurally unusable binding inputs.""" + + +class BindingTier(str, Enum): + """Which rule a field is governed by.""" + + IDENTICAL = "identical" + SEMANTIC = "semantic" + RECORDED = "recorded" + + +class BindingErrorCode(str, Enum): + """Stable, machine-readable reasons a binding is rejected. + + Callers branch on these; they are part of the artifact schema and must not be + renamed without a schema version bump. + """ + + IDENTITY_MISSING = "IDENTITY_MISSING" + IDENTITY_MISMATCH = "IDENTITY_MISMATCH" + REDUCTION_SEMANTIC_MISMATCH = "REDUCTION_SEMANTIC_MISMATCH" + REDUCTION_MANDATE_VIOLATION = "REDUCTION_MANDATE_VIOLATION" + LSE_NOT_EXPORTED = "LSE_NOT_EXPORTED" + ROLE_COLLISION = "ROLE_COLLISION" + DETERMINISM_INCOMPATIBLE = "DETERMINISM_INCOMPATIBLE" + + +#: Attention exports attention-domain LSE, never vocab-logprob LSE (#235). +#: Recorded explicitly so a future ``LogprobContract`` binding cannot be confused +#: with this one purely because both set ``export_lse=True``. +ATTENTION_LSE_DOMAIN = "attention" + + +#: Fields both sides must agree on bit for bit before any comparison is meaningful. +#: Sourced from #235 "Numerical Contract" preconditions plus the vime-owned rollout +#: provenance (weight version, sampling, padding) that the issue assumes but does +#: not enumerate. +IDENTITY_FIELDS: tuple[str, ...] = ( + "checkpoint_id", + "model_version", + "weight_version", + "tokenizer_fingerprint", + "token_ids_fingerprint", + "active_mask_fingerprint", + "position_ids_fingerprint", + "padding_side", + "pre_update_state", + # model semantics that decide what attention *means* + "q_heads", + "kv_heads", + "head_dim", + "rope_theta", + "rope_scaling", + "rotary_dim", + "qk_layernorm", + # decode replay identity (#235 PR6) + "global_token_positions_fingerprint", + "kv_seq_lens_fingerprint", +) + + +#: Reduction fields that decide the numerical result. Both sides must carry the +#: same value, and that value must satisfy :data:`WS2_ATTENTION_REDUCTION_MANDATE`. +SEMANTIC_REDUCTION_FIELDS: tuple[str, ...] = ( + "merge", + "acc_dtype", + "order", + "downcast_at", +) + + +#: The WS2 mandate itself. ``#236`` currently declares single-member enums for +#: ``merge`` / ``order`` / ``downcast_at``, so those checks are tautological today; +#: they are written out anyway so that widening any of those enums later fails here +#: instead of silently admitting a non-conforming backend. +WS2_ATTENTION_REDUCTION_MANDATE: Mapping[str, str] = { + "merge": AttentionMerge.ONLINE_SOFTMAX_LSE.value, + "acc_dtype": AttentionDType.FP32.value, + "order": ReductionOrder.GLOBAL_BLOCK_INDEX.value, + "downcast_at": DowncastPoint.FINAL_WRITE.value, +} + + +#: Materialization facts the two sides are expected to differ on. Recorded into +#: provenance; never a rejection reason. +RECORDED_FIELDS: tuple[str, ...] = ( + "mode", + "backend_id", + "reduction.engine", + "rope.fusion_boundary", + "rope.q_state", + "rope.k_state", + "rope.k_cache_state", + "rope.cast_at", + "rope.output_dtype", + "kv_cache.page_size", + "kv_cache.prefix_cache_enabled", + "kv_cache.block_table_shape", + "sharding.cp_world_size", + "sharding.tp_world_size", + "sharding.local_sequence_length", +) + + +@dataclass(frozen=True) +class BindingIssue: + """One reason a binding is not comparable or not admissible.""" + + code: BindingErrorCode + tier: BindingTier + field: str + rollout: Any = None + training: Any = None + message: str = "" + + def to_dict(self) -> dict[str, Any]: + return { + "code": self.code.value, + "tier": self.tier.value, + "field": self.field, + "rollout": self.rollout, + "training": self.training, + "message": self.message, + } + + +@dataclass(frozen=True) +class AttentionBindingResult: + """Outcome of binding one rollout contract to one training contract. + + ``comparable`` and ``passed`` are deliberately separate. A pair whose identity + does not match is *not comparable* -- reporting a drift number for it would be + meaningless. A pair that is comparable but violates the reduction mandate *is* + comparable yet must still fail closed, because the whole WS2 claim is that + reduction order and accumulation precision come from the contract. + """ + + comparable: bool + passed: bool + issues: tuple[BindingIssue, ...] = () + identity_fingerprint: str = "" + reduction_fingerprint: str = "" + binding_fingerprint: str = "" + recorded_differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + provenance: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.attention_binding.v1" + + def issues_by_code(self, code: BindingErrorCode) -> tuple[BindingIssue, ...]: + return tuple(issue for issue in self.issues if issue.code is code) + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "comparable": self.comparable, + "passed": self.passed, + "issues": [issue.to_dict() for issue in self.issues], + "identity_fingerprint": self.identity_fingerprint, + "reduction_fingerprint": self.reduction_fingerprint, + "binding_fingerprint": self.binding_fingerprint, + "recorded_differences": { + key: dict(value) for key, value in self.recorded_differences.items() + }, + "provenance": dict(self.provenance), + } + + +def _canonical_fingerprint(payload: Any) -> str: + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + +def identity_fingerprint(identity: Mapping[str, Any]) -> str: + """Fingerprint only the declared :data:`IDENTITY_FIELDS`, in a fixed order. + + Extra keys in ``identity`` are ignored on purpose: callers pass whole + provenance bundles, and the fingerprint must not drift when an unrelated + diagnostic field is added. + """ + + return _canonical_fingerprint({name: identity.get(name) for name in IDENTITY_FIELDS}) + + +def _reduction_view(contract: AttentionContract) -> dict[str, Any]: + reduction = contract.reduction + return { + "merge": reduction.merge.value, + "acc_dtype": reduction.acc_dtype.value, + "order": reduction.order.value, + "downcast_at": reduction.downcast_at.value, + "engine": reduction.engine.value, + } + + +def _recorded_view(contract: AttentionContract) -> dict[str, Any]: + rope = contract.rope + kv_cache = contract.kv_cache + view: dict[str, Any] = { + "mode": contract.mode.value, + "backend_id": None, + "reduction.engine": contract.reduction.engine.value, + "sharding.cp_world_size": contract.sharding.cp_world_size, + "sharding.tp_world_size": contract.sharding.tp_world_size, + "sharding.local_sequence_length": contract.sharding.local_sequence_length, + } + if rope is not None: + view.update( + { + "rope.fusion_boundary": rope.fusion_boundary.value, + "rope.q_state": rope.q_state.value, + "rope.k_state": rope.k_state.value, + "rope.k_cache_state": rope.k_cache_state.value, + "rope.cast_at": rope.cast_at.value, + "rope.output_dtype": rope.output_dtype.value, + } + ) + if kv_cache is not None: + view.update( + { + "kv_cache.page_size": kv_cache.page_size, + "kv_cache.prefix_cache_enabled": kv_cache.prefix_cache_enabled, + "kv_cache.block_table_shape": [ + len(kv_cache.block_table), + max((len(row) for row in kv_cache.block_table), default=0), + ], + } + ) + return view + + +#: Identity fields where ``None`` is a real value rather than an omission. Qwen3-8B +#: applies no RoPE scaling, so ``rope_scaling=None`` must not read as "undeclared" -- +#: both sides still have to agree on it, which the equality pass below handles. +NULLABLE_IDENTITY_FIELDS: frozenset[str] = frozenset({"rope_scaling"}) + + +def _missing_identity_fields(identity: Mapping[str, Any]) -> tuple[str, ...]: + return tuple( + name + for name in IDENTITY_FIELDS + if name not in NULLABLE_IDENTITY_FIELDS and identity.get(name) is None + ) + + +def bind_attention_contracts( + *, + rollout_contract: AttentionContract, + training_contract: AttentionContract, + rollout_identity: Mapping[str, Any], + training_identity: Mapping[str, Any], + rollout_backend_id: str, + training_backend_id: str, + determinism_issues: Sequence[BindingIssue] = (), + require_full_identity: bool = True, +) -> AttentionBindingResult: + """Bind a rollout attention contract to a training attention contract. + + ``determinism_issues`` is threaded in from + :mod:`rl_engine.alignment.cross_config.determinism` rather than computed here, + so that this module stays free of framework probing and remains testable + without Megatron or vLLM present. + + ``require_full_identity`` exists for the single-GPU harness in #235 PR2, which + legitimately has no KV-cache or decode identity to declare. Distributed callers + must leave it at ``True``. + """ + + if rollout_contract.role is not AttentionRole.INFER: + raise AttentionBindingError( + f"rollout_contract.role must be {AttentionRole.INFER.value!r}, " + f"got {rollout_contract.role.value!r}" + ) + if training_contract.role is not AttentionRole.TRAIN: + raise AttentionBindingError( + f"training_contract.role must be {AttentionRole.TRAIN.value!r}, " + f"got {training_contract.role.value!r}" + ) + + issues: list[BindingIssue] = [] + + # ---- tier 1: identity, bit for bit ------------------------------------- + if require_full_identity: + for side, identity in (("rollout", rollout_identity), ("training", training_identity)): + for name in _missing_identity_fields(identity): + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISSING, + tier=BindingTier.IDENTICAL, + field=f"{side}.{name}", + message=f"{side} identity does not declare {name!r}", + ) + ) + + for name in IDENTITY_FIELDS: + rollout_value = rollout_identity.get(name) + training_value = training_identity.get(name) + if rollout_value != training_value: + issues.append( + BindingIssue( + code=BindingErrorCode.IDENTITY_MISMATCH, + tier=BindingTier.IDENTICAL, + field=name, + rollout=rollout_value, + training=training_value, + message=( + f"{name!r} differs between sides; the pair is not comparable " + "and any drift computed from it is meaningless" + ), + ) + ) + + comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + + # ---- tier 2: reduction semantics, and the WS2 mandate ------------------- + rollout_reduction = _reduction_view(rollout_contract) + training_reduction = _reduction_view(training_contract) + + for name in SEMANTIC_REDUCTION_FIELDS: + if rollout_reduction[name] != training_reduction[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"reduction.{name!r} must be decided by the contract, not by the " + "backend; the two sides disagree" + ), + ) + ) + mandated = WS2_ATTENTION_REDUCTION_MANDATE[name] + for side, view in (("rollout", rollout_reduction), ("training", training_reduction)): + if view[name] != mandated: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_MANDATE_VIOLATION, + tier=BindingTier.SEMANTIC, + field=f"{side}.reduction.{name}", + rollout=rollout_reduction[name], + training=training_reduction[name], + message=( + f"WS2 requires reduction.{name} == {mandated!r}; " + f"{side} declares {view[name]!r}" + ), + ) + ) + + for side, contract in (("rollout", rollout_contract), ("training", training_contract)): + if not contract.export_lse: + issues.append( + BindingIssue( + code=BindingErrorCode.LSE_NOT_EXPORTED, + tier=BindingTier.SEMANTIC, + field=f"{side}.export_lse", + message=( + "attention-domain LSE must be exported; without it the deterministic " + "CP merge cannot be validated" + ), + ) + ) + + if rollout_backend_id == training_backend_id and rollout_backend_id: + # Not an error, but worth surfacing: an identical backend on both sides means + # the experiment is not actually measuring a cross-implementation difference. + pass + + issues.extend(determinism_issues) + + # ---- tier 3: recorded differences -------------------------------------- + rollout_recorded = _recorded_view(rollout_contract) + rollout_recorded["backend_id"] = rollout_backend_id + training_recorded = _recorded_view(training_contract) + training_recorded["backend_id"] = training_backend_id + + recorded_differences: dict[str, dict[str, Any]] = {} + for name in RECORDED_FIELDS: + rollout_value = rollout_recorded.get(name) + training_value = training_recorded.get(name) + if rollout_value != training_value: + recorded_differences[name] = { + "rollout": rollout_value, + "training": training_value, + } + + identity_fp = identity_fingerprint(training_identity if comparable else rollout_identity) + reduction_fp = _canonical_fingerprint( + {name: training_reduction[name] for name in SEMANTIC_REDUCTION_FIELDS} + ) + passed = comparable and not any(issue.tier is BindingTier.SEMANTIC for issue in issues) + + provenance = { + "lse_domain": ATTENTION_LSE_DOMAIN, + "rollout": { + "contract": rollout_contract.to_dict(), + "backend_id": rollout_backend_id, + "recorded": rollout_recorded, + }, + "training": { + "contract": training_contract.to_dict(), + "backend_id": training_backend_id, + "recorded": training_recorded, + }, + } + + return AttentionBindingResult( + comparable=comparable, + passed=passed, + issues=tuple(issues), + identity_fingerprint=identity_fp, + reduction_fingerprint=reduction_fp, + binding_fingerprint=_canonical_fingerprint( + { + "identity": identity_fp, + "reduction": reduction_fp, + "lse_domain": ATTENTION_LSE_DOMAIN, + "rollout_backend": rollout_backend_id, + "training_backend": training_backend_id, + } + ), + recorded_differences=recorded_differences, + provenance=provenance, + ) + + +def summarize_binding(result: AttentionBindingResult) -> str: + """One-line human summary for CLI output and failure messages.""" + + if result.passed: + return ( + f"attention binding OK " + f"(identity={result.identity_fingerprint[:12]}, " + f"{len(result.recorded_differences)} recorded difference(s))" + ) + if not result.comparable: + fields = ", ".join( + sorted({issue.field for issue in result.issues if issue.tier is BindingTier.IDENTICAL}) + ) + return f"attention binding NOT COMPARABLE; identity problems: {fields}" + fields = ", ".join( + sorted({issue.field for issue in result.issues if issue.tier is BindingTier.SEMANTIC}) + ) + return f"attention binding FAILED CLOSED; semantic problems: {fields}" + + +def first_blocking_issue( + result: AttentionBindingResult, +) -> Optional[BindingIssue]: + """Return the issue a caller should report, preferring identity over semantics.""" + + for tier in (BindingTier.IDENTICAL, BindingTier.SEMANTIC): + for issue in result.issues: + if issue.tier is tier: + return issue + return None diff --git a/rl_engine/alignment/cross_config/determinism.py b/rl_engine/alignment/cross_config/determinism.py new file mode 100644 index 00000000..be81654e --- /dev/null +++ b/rl_engine/alignment/cross_config/determinism.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-side determinism probing for the Megatron + vLLM cross-config target. + +Both frameworks ship a "make this deterministic" switch, but they mean different +things by it, and neither knows the other exists: + +``Megatron`` ``ModelParallelConfig.deterministic_mode`` + Asserts ``NCCL_ALGO`` is one of five values, forbids FlashAttention and fused + cross-entropy, calls ``torch.use_deterministic_algorithms(True)``, and requires + ``NVTE_ALLOW_NONDETERMINISTIC_ALGO == 0``. It does **not** touch TF32, BF16 + reduced-precision reduction, cuBLAS workspace, NCCL protocol, or NCCL channel + counts. + +``vLLM`` ``VLLM_BATCH_INVARIANT`` + Replaces ``aten::mm/addmm/matmul/linear/bmm``, ``log_softmax``/``softmax``, + ``mean.dim`` and ``rms_norm`` with Triton kernels, disables TF32 and BF16/FP16 + reduced-precision reduction, pins cuBLAS workspace and the BLAS library, and + hard-sets ten NCCL environment variables. + +So a run can have both switches on and still be comparing two different notions of +determinism. This module makes that difference explicit and, where it changes the +numerics, blocking. It never imports Megatron or vLLM: probes are built from plain +mappings so the logic is testable on any machine. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import Any, Optional + +from rl_engine.alignment.cross_config.attention_binding import ( + BindingErrorCode, + BindingIssue, + BindingTier, +) + +__all__ = [ + "COMPARED_NCCL_KEYS", + "DeterminismProbe", + "DeterminismReport", + "compare_determinism", + "megatron_probe_from_config", + "vllm_probe_from_env", +] + + +#: Environment keys whose value can change a reduction result. Compared across +#: sides; a difference is reported, and a difference in the *arithmetic* subset is +#: blocking. Ordering is fixed so the fingerprint is stable. +COMPARED_NCCL_KEYS: tuple[str, ...] = ( + "NCCL_ALGO", + "NCCL_PROTO", + "NCCL_MIN_NCHANNELS", + "NCCL_MAX_NCHANNELS", + "NCCL_NTHREADS", + "NCCL_SOCKET_NTHREADS", + "NCCL_COLLNET_ENABLE", + "NCCL_NVLS_ENABLE", + "NCCL_P2P_NET_DISABLE", + "NCCL_LAUNCH_MODE", + "CUBLAS_WORKSPACE_CONFIG", +) + + +#: The subset above that changes arithmetic rather than only scheduling. A mismatch +#: here fails the binding closed; a mismatch in the remainder is recorded only. +_ARITHMETIC_NCCL_KEYS: frozenset[str] = frozenset( + {"NCCL_ALGO", "NCCL_PROTO", "CUBLAS_WORKSPACE_CONFIG"} +) + + +@dataclass(frozen=True) +class DeterminismProbe: + """What one side actually has switched on. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are tri-state on + purpose: ``None`` means "the framework does not manage this", which is exactly + Megatron's situation and is itself the finding. + """ + + side: str + framework: str + mode_flag: str + enabled: bool + env: Mapping[str, Any] = field(default_factory=dict) + tf32_disabled: Optional[bool] = None + bf16_reduced_precision_reduction: Optional[bool] = None + forbids_flash_attention: Optional[bool] = None + evidence: Mapping[str, Any] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_probe.v1" + + def __post_init__(self) -> None: + if self.side not in ("rollout", "training"): + raise ValueError("side must be 'rollout' or 'training'") + if not self.framework: + raise ValueError("framework must not be empty") + object.__setattr__(self, "env", dict(self.env)) + object.__setattr__(self, "evidence", dict(self.evidence)) + + @property + def env_fingerprint(self) -> str: + payload = {key: self.env.get(key) for key in COMPARED_NCCL_KEYS} + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=str) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "side": self.side, + "framework": self.framework, + "mode_flag": self.mode_flag, + "enabled": self.enabled, + "env": {key: self.env.get(key) for key in COMPARED_NCCL_KEYS}, + "env_fingerprint": self.env_fingerprint, + "tf32_disabled": self.tf32_disabled, + "bf16_reduced_precision_reduction": self.bf16_reduced_precision_reduction, + "forbids_flash_attention": self.forbids_flash_attention, + "evidence": dict(self.evidence), + } + + +@dataclass(frozen=True) +class DeterminismReport: + """Cross-side comparison result.""" + + rollout: DeterminismProbe + training: DeterminismProbe + issues: tuple[BindingIssue, ...] = () + differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + schema_version: str = "cross_config.determinism_report.v1" + + @property + def compatible(self) -> bool: + return not self.issues + + def to_dict(self) -> dict[str, Any]: + return { + "schema_version": self.schema_version, + "compatible": self.compatible, + "rollout": self.rollout.to_dict(), + "training": self.training.to_dict(), + "issues": [issue.to_dict() for issue in self.issues], + "differences": {key: dict(value) for key, value in self.differences.items()}, + } + + +def megatron_probe_from_config( + config: Any, + env: Optional[Mapping[str, str]] = None, +) -> DeterminismProbe: + """Build a training-side probe from a Megatron config object. + + ``config`` is duck-typed (anything exposing ``deterministic_mode`` and + optionally ``attention_backend`` / ``cross_entropy_loss_fusion``) so this works + against a real ``ModelParallelConfig``, a test double, or a plain namespace, + and so importing this module never requires Megatron. + + ``tf32_disabled`` and ``bf16_reduced_precision_reduction`` are reported as + ``None`` because Megatron does not manage them -- a ``grep`` for ``allow_tf32`` + and ``fp32_precision`` across ``megatron/`` returns nothing. That asymmetry + against vLLM is the point of :func:`compare_determinism`. + """ + + environ = dict(env or {}) + enabled = bool(getattr(config, "deterministic_mode", False)) + return DeterminismProbe( + side="training", + framework="megatron", + mode_flag="deterministic_mode", + enabled=enabled, + env={key: environ.get(key) for key in COMPARED_NCCL_KEYS}, + tf32_disabled=None, + bf16_reduced_precision_reduction=None, + forbids_flash_attention=enabled, + evidence={ + "nvte_allow_nondeterministic_algo": environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO"), + "cross_entropy_loss_fusion": getattr(config, "cross_entropy_loss_fusion", None), + "attention_backend": _enum_value(getattr(config, "attention_backend", None)), + "tensor_model_parallel_size": getattr(config, "tensor_model_parallel_size", None), + "context_parallel_size": getattr(config, "context_parallel_size", None), + "sequence_parallel": getattr(config, "sequence_parallel", None), + "manages_tf32": False, + "manages_bf16_reduced_precision_reduction": False, + }, + ) + + +def vllm_probe_from_env( + env: Mapping[str, str], + *, + model_config: Any = None, +) -> DeterminismProbe: + """Build a rollout-side probe from the vLLM process environment. + + ``VLLM_BATCH_INVARIANT`` is read from ``env`` rather than ``vllm.envs`` so the + probe can be constructed from a remote worker's reported environment, which is + how vime's Ray actors expose it. + """ + + enabled = str(env.get("VLLM_BATCH_INVARIANT", "0")).strip() in ("1", "true", "True") + return DeterminismProbe( + side="rollout", + framework="vllm", + mode_flag="VLLM_BATCH_INVARIANT", + enabled=enabled, + env={key: env.get(key) for key in COMPARED_NCCL_KEYS}, + # vLLM sets both to "ieee"/disabled inside init_batch_invariance(). + tf32_disabled=enabled or None, + bf16_reduced_precision_reduction=(False if enabled else None), + forbids_flash_attention=False, + evidence={ + "vllm_allreduce_use_symm_mem": env.get("VLLM_ALLREDUCE_USE_SYMM_MEM"), + "vllm_use_aot_compile": env.get("VLLM_USE_AOT_COMPILE"), + "enforce_eager": getattr(model_config, "enforce_eager", None), + "disable_cascade_attn": getattr(model_config, "disable_cascade_attn", None), + "quantization": getattr(model_config, "quantization", None), + "manages_tf32": True, + "manages_bf16_reduced_precision_reduction": True, + }, + ) + + +def _enum_value(value: Any) -> Any: + return getattr(value, "value", value) + + +def compare_determinism( + *, + rollout: DeterminismProbe, + training: DeterminismProbe, +) -> DeterminismReport: + """Compare two probes and produce blocking issues plus recorded differences.""" + + if rollout.side != "rollout" or training.side != "training": + raise ValueError("compare_determinism expects one rollout probe and one training probe") + + issues: list[BindingIssue] = [] + differences: dict[str, dict[str, Any]] = {} + + for probe in (rollout, training): + if not probe.enabled: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"{probe.side}.{probe.mode_flag}", + rollout=rollout.enabled, + training=training.enabled, + message=( + f"{probe.framework} {probe.mode_flag} is not enabled; the " + f"{probe.side} side is not batch-invariant and cannot anchor a " + "cross-config comparison" + ), + ) + ) + + for key in COMPARED_NCCL_KEYS: + rollout_value = rollout.env.get(key) + training_value = training.env.get(key) + if rollout_value == training_value: + continue + differences[key] = {"rollout": rollout_value, "training": training_value} + if key in _ARITHMETIC_NCCL_KEYS: + issues.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"env.{key}", + rollout=rollout_value, + training=training_value, + message=( + f"{key} differs between sides; the two sides would reduce with " + "different arithmetic and the resulting drift is not attributable" + ), + ) + ) + + # Megatron reports None for these because it does not manage them at all. That is + # recorded rather than blocking: under a pure BF16 GEMM path TF32 does not fire, and + # forcing Megatron to manage it is out of scope for this PR. It is surfaced so the + # asymmetry appears in every artifact instead of being invisible. + for name in ("tf32_disabled", "bf16_reduced_precision_reduction"): + rollout_value = getattr(rollout, name) + training_value = getattr(training, name) + if rollout_value != training_value: + differences[name] = { + "rollout": rollout_value, + "training": training_value, + "note": ( + "megatron does not manage this setting; vllm sets it inside " + "init_batch_invariance()" + ), + } + + return DeterminismReport( + rollout=rollout, + training=training, + issues=tuple(issues), + differences=differences, + ) diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py new file mode 100644 index 00000000..058dfa44 --- /dev/null +++ b/tests/test_attention_cross_config_binding.py @@ -0,0 +1,612 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tests for #235 PR4: rollout/training attention contract binding. + +Every test here runs on CPU without Megatron or vLLM installed. That is the point: +the binding rules are contract logic, and contract logic that can only be exercised +on a 2-node x 2-GPU cluster would never be exercised. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from rl_engine.alignment.cross_config.adapters import ( + QWEN3_8B, + WS2_ATTENTION_KNOBS, + MegatronAttentionMaterializer, + MegatronProvenanceAdapter, + VllmProvenanceAdapter, + VllmRolloutMaterializer, +) +from rl_engine.alignment.cross_config.attention_binding import ( + ATTENTION_LSE_DOMAIN, + AttentionBindingError, + BindingErrorCode, + BindingTier, + bind_attention_contracts, + first_blocking_issue, + identity_fingerprint, + summarize_binding, +) +from rl_engine.alignment.cross_config.determinism import ( + compare_determinism, + megatron_probe_from_config, + vllm_probe_from_env, +) +from rl_engine.alignment.cross_config.schema import MaterializationStatus +from rl_engine.kernels.attention_contract import AttentionContractError, AttentionMode + +pytestmark = pytest.mark.unit + + +TRAINING_KNOBS = { + "batch.size": 2, + "training.tensor_parallel_size": 2, + "training.context_parallel_size": 2, + "training.compute_dtype": "bf16", +} + +ROLLOUT_KNOBS = { + "batch.size": 2, + "rollout.tensor_parallel_size": 2, + "rollout.context_parallel_size": 1, + "rollout.dtype": "bf16", +} + + +def _identity(**overrides): + identity = { + "checkpoint_id": "qwen3-8b", + "model_version": "v1", + "weight_version": 7, + "tokenizer_fingerprint": "tokenizer-abc", + "token_ids_fingerprint": "tokens-abc", + "active_mask_fingerprint": "mask-abc", + "position_ids_fingerprint": "pos-abc", + "padding_side": "right", + "pre_update_state": "pre_update", + "global_token_positions_fingerprint": "gtp-abc", + "kv_seq_lens_fingerprint": "kvlen-abc", + } + identity.update(QWEN3_8B.identity_fields()) + identity.update(overrides) + return identity + + +def _contracts(): + training = MegatronAttentionMaterializer().build_contract(TRAINING_KNOBS) + rollout = VllmRolloutMaterializer().build_contract(ROLLOUT_KNOBS) + return rollout, training + + +def _bind(rollout_identity=None, training_identity=None, **kwargs): + rollout, training = _contracts() + return bind_attention_contracts( + rollout_contract=kwargs.pop("rollout_contract", rollout), + training_contract=kwargs.pop("training_contract", training), + rollout_identity=rollout_identity if rollout_identity is not None else _identity(), + training_identity=training_identity if training_identity is not None else _identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + **kwargs, + ) + + +# -------------------------------------------------------------------------- +# tier 1: identity +# -------------------------------------------------------------------------- + + +def test_matching_identity_binds_despite_different_materialization(): + """The core claim of PR4: same identity + same reduction, different runtimes.""" + + result = _bind() + + assert result.comparable + assert result.passed + assert result.issues == () + # Training runs CP=2 full prefill, rollout runs CP=1 chunked prefill. Those + # differences are recorded, not rejected. + assert "mode" in result.recorded_differences + assert "sharding.cp_world_size" in result.recorded_differences + assert result.recorded_differences["mode"] == { + "rollout": "chunked_prefill", + "training": "prefill", + } + + +def test_weight_version_mismatch_is_not_comparable(): + result = _bind(rollout_identity=_identity(weight_version=6)) + + assert not result.comparable + assert not result.passed + codes = {issue.code for issue in result.issues} + assert BindingErrorCode.IDENTITY_MISMATCH in codes + blocking = first_blocking_issue(result) + assert blocking is not None and blocking.tier is BindingTier.IDENTICAL + assert "NOT COMPARABLE" in summarize_binding(result) + + +def test_rope_theta_mismatch_is_not_comparable(): + """RoPE math constants are identity, not materialization.""" + + result = _bind(training_identity=_identity(rope_theta=10000.0)) + + assert not result.comparable + assert any(issue.field == "rope_theta" for issue in result.issues) + + +def test_null_rope_scaling_is_a_value_not_an_omission(): + """Qwen3-8B applies no RoPE scaling; ``None`` must not read as undeclared.""" + + result = _bind() + + assert not result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + + +def test_missing_identity_field_is_reported_per_side(): + identity = _identity() + del identity["padding_side"] + result = _bind(rollout_identity=identity, training_identity=identity) + + missing = result.issues_by_code(BindingErrorCode.IDENTITY_MISSING) + assert {issue.field for issue in missing} == { + "rollout.padding_side", + "training.padding_side", + } + assert not result.comparable + + +def test_single_gpu_harness_may_waive_full_identity(): + """#235 PR2 has no KV-cache identity to declare; it opts out explicitly.""" + + identity = _identity() + del identity["global_token_positions_fingerprint"] + del identity["kv_seq_lens_fingerprint"] + + strict = _bind(rollout_identity=identity, training_identity=identity) + waived = _bind( + rollout_identity=identity, + training_identity=identity, + require_full_identity=False, + ) + + assert not strict.comparable + assert waived.comparable and waived.passed + + +def test_identity_fingerprint_ignores_undeclared_extra_keys(): + base = _identity() + decorated = dict(base, diagnostic_note="added later") + + assert identity_fingerprint(base) == identity_fingerprint(decorated) + + +# -------------------------------------------------------------------------- +# tier 2: reduction semantics +# -------------------------------------------------------------------------- + + +def test_reduction_semantics_are_bound_and_fingerprinted(): + result = _bind() + + reduction = result.provenance["training"]["contract"]["reduction"] + assert reduction["merge"] == "online_softmax_lse" + assert reduction["acc_dtype"] == "fp32" + assert reduction["order"] == "global_block_index" + assert reduction["downcast_at"] == "final_write" + assert result.reduction_fingerprint + + +def test_reduction_engine_difference_is_recorded_not_rejected(): + """A TE merge oracle on one side must not fail the binding.""" + + from rl_engine.alignment.cross_config.attention_binding import ( + RECORDED_FIELDS, + SEMANTIC_REDUCTION_FIELDS, + ) + + assert "reduction.engine" in RECORDED_FIELDS + assert "engine" not in SEMANTIC_REDUCTION_FIELDS + + +def test_lse_domain_is_recorded_as_attention_domain(): + """#235: attention exports attention-domain LSE, not vocab-logprob LSE.""" + + result = _bind() + + assert result.provenance["lse_domain"] == ATTENTION_LSE_DOMAIN == "attention" + + +# -------------------------------------------------------------------------- +# role and input validation +# -------------------------------------------------------------------------- + + +def test_swapped_roles_are_rejected_outright(): + rollout, training = _contracts() + + with pytest.raises(AttentionBindingError): + bind_attention_contracts( + rollout_contract=training, + training_contract=rollout, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="a", + training_backend_id="b", + ) + + +# -------------------------------------------------------------------------- +# determinism cross-check +# -------------------------------------------------------------------------- + + +def _megatron_env(): + return {"NCCL_ALGO": "Tree", "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0"} + + +def _vllm_env(**overrides): + env = { + "VLLM_BATCH_INVARIANT": "1", + "NCCL_ALGO": "allreduce:tree", + "NCCL_PROTO": "Simple", + "NCCL_MIN_NCHANNELS": "1", + "NCCL_MAX_NCHANNELS": "1", + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + } + env.update(overrides) + return env + + +def test_nccl_algo_mismatch_blocks_the_binding(): + """Megatron asserts NCCL_ALGO; vLLM hard-sets a different value.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert not report.compatible + fields = {issue.field for issue in report.issues} + assert "env.NCCL_ALGO" in fields + assert "env.NCCL_PROTO" in fields + + +def test_matching_nccl_settings_are_compatible(): + shared = {"NCCL_ALGO": "allreduce:tree", "NCCL_PROTO": "Simple"} + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=dict(shared) + ) + rollout = vllm_probe_from_env({**_vllm_env(**shared), "CUBLAS_WORKSPACE_CONFIG": None}) + + report = compare_determinism(rollout=rollout, training=training) + + assert report.compatible, [issue.to_dict() for issue in report.issues] + + +def test_determinism_switch_off_on_either_side_blocks(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=False), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env(VLLM_BATCH_INVARIANT="0")) + + report = compare_determinism(rollout=rollout, training=training) + + fields = {issue.field for issue in report.issues} + assert "training.deterministic_mode" in fields + assert "rollout.VLLM_BATCH_INVARIANT" in fields + + +def test_tf32_asymmetry_is_recorded_not_blocking(): + """Megatron does not manage TF32 at all; vLLM disables it. Record the gap.""" + + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + + report = compare_determinism(rollout=rollout, training=training) + + assert training.tf32_disabled is None + assert rollout.tf32_disabled is True + assert "tf32_disabled" in report.differences + assert not any(issue.field == "tf32_disabled" for issue in report.issues) + + +def test_determinism_issues_flow_into_the_binding(): + training = megatron_probe_from_config( + SimpleNamespace(deterministic_mode=True), env=_megatron_env() + ) + rollout = vllm_probe_from_env(_vllm_env()) + report = compare_determinism(rollout=rollout, training=training) + + result = _bind(determinism_issues=report.issues) + + assert result.comparable # identity is fine + assert not result.passed # but the reduction environment is not + assert result.issues_by_code(BindingErrorCode.DETERMINISM_INCOMPATIBLE) + assert "FAILED CLOSED" in summarize_binding(result) + + +# -------------------------------------------------------------------------- +# sharding derived from the frozen #239 layout +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("cp_rank", [0, 1]) +def test_cp_shards_cover_the_global_sequence_without_overlap(cp_rank): + contract = MegatronAttentionMaterializer( + cp_rank=cp_rank, global_sequence_length=4096 + ).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert sharding.local_sequence_length == 2048 + assert sharding.global_block_indices == (cp_rank,) + assert sharding.global_block_token_starts == (cp_rank * 2048,) + # The causal offset must be the number of preceding *global* tokens, otherwise + # rank 1 would mask as if its shard started at position zero. + assert contract.causal_offsets == (cp_rank * 2048, cp_rank * 2048) + + +def test_tp_head_shards_split_qwen3_gqa_evenly(): + contract = MegatronAttentionMaterializer(tp_rank=1).build_contract(TRAINING_KNOBS) + + sharding = contract.sharding + assert (sharding.global_q_heads, sharding.global_kv_heads) == (32, 8) + assert (sharding.local_q_heads, sharding.local_kv_heads) == (16, 4) + assert (sharding.local_q_head_start, sharding.local_kv_head_start) == (16, 4) + + +@pytest.mark.parametrize("tp_world_size", [2, 4, 8]) +def test_supported_tp_degrees_shard_qwen3_gqa(tp_world_size): + """Qwen3-8B has 32 Q heads and 8 KV heads, so TP in {2, 4, 8} all divide.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": tp_world_size} + ) + + sharding = contract.sharding + assert sharding.local_q_heads == 32 // tp_world_size + assert sharding.local_kv_heads == 8 // tp_world_size + + +@pytest.mark.parametrize( + ("knob_value", "expected"), + [("bfloat16", "bf16"), ("float16", "fp16"), ("float32", "fp32"), ("fp16", "fp16")], +) +def test_planner_normalized_dtypes_reach_the_contract(knob_value, expected): + """The planner emits torch spellings; AttentionDType uses short ones.""" + + contract = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": knob_value} + ) + + assert contract.dtype.value == expected + + +def test_unknown_dtype_is_rejected_with_the_offending_field(): + with pytest.raises(ValueError, match="training.compute_dtype"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "int8"} + ) + + +def test_indivisible_tp_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.tensor_parallel_size": 3} + ) + + +def test_indivisible_cp_sequence_is_rejected(): + with pytest.raises(ValueError, match="divide evenly"): + MegatronAttentionMaterializer(global_sequence_length=4097).build_contract(TRAINING_KNOBS) + + +# -------------------------------------------------------------------------- +# materialization: fail closed rather than silently substitute +# -------------------------------------------------------------------------- + + +def _statuses(materialization, path): + return [app.status for app in materialization.applications if app.path == path] + + +def test_arrival_merge_order_is_unsupported_not_silently_corrected(): + """The control group must stay distinguishable from the treatment.""" + + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_order": "arrival"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_order") == [ + MaterializationStatus.UNSUPPORTED + ] + assert materialization.binding.side_configs["training"]["contract"] is None + assert "arrival" in materialization.binding.side_configs["training"]["contract_error"] + + +def test_bf16_reduction_accumulation_is_unsupported(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_acc_dtype": "bf16"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_acc_dtype") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_te_oracle_engine_is_unsupported_until_pr2_pr3(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_engine": "te_oracle"}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_engine") == [ + MaterializationStatus.UNSUPPORTED + ] + + +def test_vllm_cp_falls_back_to_one_in_decode_and_says_why(): + materializer = VllmRolloutMaterializer(mode=AttentionMode.DECODE) + normalized = { + "batch": {"size": 2}, + "rollout": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + + assert materializer.effective_cp_world_size({"rollout.context_parallel_size": 2}) == 1 + + materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) + contract_error = materialization.binding.side_configs["rollout"]["contract_error"] + assert "#235 PR6" in contract_error + + +def test_decode_contract_is_refused_without_kv_cache_identity(): + with pytest.raises(AttentionContractError, match="PR6"): + VllmRolloutMaterializer(mode=AttentionMode.DECODE).build_contract(ROLLOUT_KNOBS) + + +def test_materializers_expose_distinct_implementation_fingerprints(): + megatron = MegatronAttentionMaterializer().implementation_fingerprint + vllm = VllmRolloutMaterializer().implementation_fingerprint + + assert megatron and vllm and megatron != vllm + + +def test_runtime_binding_reports_the_frozen_topology(): + normalized = { + "batch": {"size": 2}, + "training": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + } + binding = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS).binding + + topology = binding.topology["training"] + assert topology["tensor_parallel_size"] == 2 + assert topology["context_parallel_size"] == 2 + assert topology["world_size"] == 4 + assert topology["pipeline_parallel_size"] == 1 + assert topology["data_parallel_size"] == 1 + + +# -------------------------------------------------------------------------- +# provenance adapters +# -------------------------------------------------------------------------- + + +def test_megatron_provenance_flags_undeclared_frozen_scope_fields(): + adapter = MegatronProvenanceAdapter(SimpleNamespace(deterministic_mode=True)) + + violations = adapter.frozen_scope_violations() + + # Nothing is declared, so every assertion reads as unknown rather than as met. + assert any("expert_model_parallel_size" in text for text in violations) + assert any("fp8" in text for text in violations) + + +def test_megatron_provenance_accepts_a_conforming_dense_config(): + adapter = MegatronProvenanceAdapter( + SimpleNamespace( + deterministic_mode=True, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + sequence_parallel=False, + fp8=None, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + ) + + assert adapter.frozen_scope_violations() == ("fp8 is not declared (expected None)",) + + +def test_megatron_construction_fingerprint_tracks_fusion_changes(): + base = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=False) + fused = SimpleNamespace(deterministic_mode=True, apply_rope_fusion=True) + + assert ( + MegatronProvenanceAdapter(base).construction_fingerprint + != MegatronProvenanceAdapter(fused).construction_fingerprint + ) + + +def test_vllm_provenance_reads_page_size_and_split_kv_policy(): + adapter = VllmProvenanceAdapter( + cache_config=SimpleNamespace(block_size=16, cache_dtype="auto"), + attention_config=SimpleNamespace(flash_attn_max_num_splits_for_cuda_graph=32), + ) + + assert adapter.kv_page_size == 16 + assert adapter.split_kv_policy == 32 + + +def test_vllm_provenance_flags_fp8_kv_cache_and_cascade_attention(): + adapter = VllmProvenanceAdapter( + model_config=SimpleNamespace(quantization=None, disable_cascade_attn=False), + cache_config=SimpleNamespace( + cache_dtype="fp8", calculate_kv_scales=False, sliding_window=None + ), + parallel_config=SimpleNamespace(pipeline_parallel_size=1, data_parallel_size=1), + ) + + violations = adapter.frozen_scope_violations() + + assert any("cache_dtype" in text for text in violations) + assert any("disable_cascade_attn" in text for text in violations) + + +# -------------------------------------------------------------------------- +# scenario config +# -------------------------------------------------------------------------- + + +SCENARIO = ( + Path(__file__).resolve().parents[1] + / "examples" + / "cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json" +) + + +def test_scenario_uses_megatron_vocabulary_only(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + training = config["baseline"]["training"] + + assert training["attention_backend"] in {"flash", "fused", "unfused", "local", "auto"} + assert training["tensor_parallel_size"] == 2 + assert training["context_parallel_size"] == 2 + assert config["baseline"]["rollout"]["batch_invariant"] is True + + +def test_scenario_knob_paths_all_exist(): + config = json.loads(SCENARIO.read_text(encoding="utf-8")) + + def paths(mapping, prefix=""): + for key, value in mapping.items(): + path = f"{prefix}{key}" + if isinstance(value, dict): + yield from paths(value, f"{path}.") + else: + yield path + + declared = set(paths(config["baseline"])) + unknown = declared - set(WS2_ATTENTION_KNOBS) + assert not unknown, f"scenario declares unknown knobs: {sorted(unknown)}" + + for intervention in config["interventions"]: + assert intervention["path"] in WS2_ATTENTION_KNOBS From 4ad305b0829b21da0bbb610b325148c337ab7e3b Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Tue, 4 Aug 2026 22:23:29 +0800 Subject: [PATCH 02/17] fix(alignment): bind dtype, batch size and split-KV policy (#235 PR4) Three fields could differ between the two sides without the binding noticing. dtype was in no tier at all, so a BF16 rollout could bind to an FP16 training pass and produce a drift number attributable to nothing. It joins the semantic tier, with allow_dtype_difference for the #235 PR5 sweep that deliberately scores BF16 against an FP32 reference. batch_size was likewise unchecked. Batch invariance is a claim about results not changing with batch makeup, so two sides scoring different batches are not comparable and it belongs to identity. split_kv_policy has no field in the #236 contract, so it only reached side_configs and never took part in binding. Callers now pass it through rollout_recorded_extra / training_recorded_extra so the difference is at least visible in provenance; it can move into the contract once #236 grows the field. Part of #235 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Q3Ar3z9fHEBFQQHddSEMaw --- .../cross_config/attention_binding.py | 55 ++++++++++++++++- tests/test_attention_cross_config_binding.py | 59 +++++++++++++++++++ 2 files changed, 111 insertions(+), 3 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index fcd4c203..48e382b8 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -61,6 +61,7 @@ "IDENTITY_FIELDS", "NULLABLE_IDENTITY_FIELDS", "RECORDED_FIELDS", + "SEMANTIC_CONTRACT_FIELDS", "SEMANTIC_REDUCTION_FIELDS", "WS2_ATTENTION_REDUCTION_MANDATE", "bind_attention_contracts", @@ -126,6 +127,9 @@ class BindingErrorCode(str, Enum): "rope_scaling", "rotary_dim", "qk_layernorm", + # batch composition: batch-invariance is a claim about results not changing with + # batch makeup, so two sides scoring different batches are not comparable at all + "batch_size", # decode replay identity (#235 PR6) "global_token_positions_fingerprint", "kv_seq_lens_fingerprint", @@ -142,6 +146,14 @@ class BindingErrorCode(str, Enum): ) +#: Contract fields outside ``ReductionSpec`` that still decide the numerical result. +#: ``dtype`` is here rather than in :data:`RECORDED_FIELDS` because comparing a BF16 +#: rollout against an FP16 training pass produces a real drift number attributable to +#: nothing. #235 PR5 does sweep BF16 against an FP32 reference; that sweep opts in via +#: ``allow_dtype_difference`` instead of loosening the default. +SEMANTIC_CONTRACT_FIELDS: tuple[str, ...] = ("dtype",) + + #: The WS2 mandate itself. ``#236`` currently declares single-member enums for #: ``merge`` / ``order`` / ``downcast_at``, so those checks are tautological today; #: they are written out anyway so that widening any of those enums later fails here @@ -172,6 +184,11 @@ class BindingErrorCode(str, Enum): "sharding.cp_world_size", "sharding.tp_world_size", "sharding.local_sequence_length", + # Supplied by the caller, not by the contract: #236 has no split-KV field yet, so + # the value comes from vLLM's flash_attn_max_num_splits_for_cuda_graph via the + # adapter. Recorded so split-KV differences are at least visible in provenance + # until #236 grows the field and it can move into the contract proper. + "split_kv_policy", ) @@ -264,7 +281,10 @@ def _reduction_view(contract: AttentionContract) -> dict[str, Any]: } -def _recorded_view(contract: AttentionContract) -> dict[str, Any]: +def _recorded_view( + contract: AttentionContract, + extra: Optional[Mapping[str, Any]] = None, +) -> dict[str, Any]: rope = contract.rope kv_cache = contract.kv_cache view: dict[str, Any] = { @@ -297,6 +317,8 @@ def _recorded_view(contract: AttentionContract) -> dict[str, Any]: ], } ) + if extra: + view.update(extra) return view @@ -324,6 +346,9 @@ def bind_attention_contracts( training_backend_id: str, determinism_issues: Sequence[BindingIssue] = (), require_full_identity: bool = True, + allow_dtype_difference: bool = False, + rollout_recorded_extra: Optional[Mapping[str, Any]] = None, + training_recorded_extra: Optional[Mapping[str, Any]] = None, ) -> AttentionBindingResult: """Bind a rollout attention contract to a training attention contract. @@ -335,6 +360,13 @@ def bind_attention_contracts( ``require_full_identity`` exists for the single-GPU harness in #235 PR2, which legitimately has no KV-cache or decode identity to declare. Distributed callers must leave it at ``True``. + + ``allow_dtype_difference`` exists for the #235 PR5 sweep that deliberately scores + a BF16 path against an FP32 reference. It must stay ``False`` everywhere else. + + ``rollout_recorded_extra`` / ``training_recorded_extra`` carry materialization + facts that #236 does not yet model -- today that is ``split_kv_policy``. They are + merged into the recorded tier, never into identity or semantics. """ if rollout_contract.role is not AttentionRole.INFER: @@ -419,6 +451,22 @@ def bind_attention_contracts( ) ) + if not allow_dtype_difference and rollout_contract.dtype is not training_contract.dtype: + issues.append( + BindingIssue( + code=BindingErrorCode.REDUCTION_SEMANTIC_MISMATCH, + tier=BindingTier.SEMANTIC, + field="dtype", + rollout=rollout_contract.dtype.value, + training=training_contract.dtype.value, + message=( + "the two sides compute in different dtypes; the resulting drift is " + "not attributable. Pass allow_dtype_difference=True only for a " + "deliberate precision sweep" + ), + ) + ) + for side, contract in (("rollout", rollout_contract), ("training", training_contract)): if not contract.export_lse: issues.append( @@ -441,9 +489,9 @@ def bind_attention_contracts( issues.extend(determinism_issues) # ---- tier 3: recorded differences -------------------------------------- - rollout_recorded = _recorded_view(rollout_contract) + rollout_recorded = _recorded_view(rollout_contract, rollout_recorded_extra) rollout_recorded["backend_id"] = rollout_backend_id - training_recorded = _recorded_view(training_contract) + training_recorded = _recorded_view(training_contract, training_recorded_extra) training_recorded["backend_id"] = training_backend_id recorded_differences: dict[str, dict[str, Any]] = {} @@ -464,6 +512,7 @@ def bind_attention_contracts( provenance = { "lse_domain": ATTENTION_LSE_DOMAIN, + "dtype": training_contract.dtype.value, "rollout": { "contract": rollout_contract.to_dict(), "backend_id": rollout_backend_id, diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 058dfa44..815daf12 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -71,6 +71,7 @@ def _identity(**overrides): "position_ids_fingerprint": "pos-abc", "padding_side": "right", "pre_update_state": "pre_update", + "batch_size": 2, "global_token_positions_fingerprint": "gtp-abc", "kv_seq_lens_fingerprint": "kvlen-abc", } @@ -224,6 +225,64 @@ def test_lse_domain_is_recorded_as_attention_domain(): assert result.provenance["lse_domain"] == ATTENTION_LSE_DOMAIN == "attention" +def test_mixed_dtypes_fail_closed(): + """BF16 rollout against FP16 training produces an unattributable number.""" + + rollout = VllmRolloutMaterializer().build_contract( + {**ROLLOUT_KNOBS, "rollout.dtype": "float16"} + ) + result = _bind(rollout_contract=rollout) + + assert result.comparable # identity is fine + assert not result.passed + assert any(issue.field == "dtype" for issue in result.issues) + + +def test_precision_sweep_may_opt_into_mixed_dtypes(): + """#235 PR5 sweeps BF16 against an FP32 reference; it says so explicitly.""" + + training = MegatronAttentionMaterializer().build_contract( + {**TRAINING_KNOBS, "training.compute_dtype": "float32"} + ) + result = _bind(training_contract=training, allow_dtype_difference=True) + + assert result.passed + assert result.provenance["dtype"] == "fp32" + + +def test_batch_size_mismatch_is_not_comparable(): + """Batch invariance is a claim about batch makeup, so it belongs to identity.""" + + result = _bind(rollout_identity=_identity(batch_size=4)) + + assert not result.comparable + assert any(issue.field == "batch_size" for issue in result.issues) + + +def test_split_kv_policy_difference_is_recorded(): + """#236 has no split-KV field, so the adapter supplies it to the recorded tier.""" + + result = _bind( + rollout_recorded_extra={"split_kv_policy": 8}, + training_recorded_extra={"split_kv_policy": None}, + ) + + assert result.passed + assert result.recorded_differences["split_kv_policy"] == { + "rollout": 8, + "training": None, + } + + +def test_matching_split_kv_policy_is_not_reported_as_a_difference(): + result = _bind( + rollout_recorded_extra={"split_kv_policy": 32}, + training_recorded_extra={"split_kv_policy": 32}, + ) + + assert "split_kv_policy" not in result.recorded_differences + + # -------------------------------------------------------------------------- # role and input validation # -------------------------------------------------------------------------- From c30be4b4f3cee7db247bf437ade94f26b9cf5e07 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 22:01:05 +0800 Subject: [PATCH 03/17] fix(attention): require runtime-verified cross-config binding --- .../ws2-attention-cross-config-integration.md | 34 +- ...config_qwen3_8b_megatron_tp2_cp2_vllm.json | 6 +- .../cross_config/adapters/__init__.py | 7 +- .../cross_config/adapters/_common.py | 18 + .../alignment/cross_config/adapters/knobs.py | 5 +- .../cross_config/adapters/megatron.py | 108 +++++- .../alignment/cross_config/adapters/vllm.py | 104 ++++- .../cross_config/attention_binding.py | 307 ++++++++++++++- tests/test_attention_cross_config_binding.py | 355 ++++++++++++++++-- 9 files changed, 871 insertions(+), 73 deletions(-) diff --git a/docs/design/ws2-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md index 06966ebe..c4d8d773 100644 --- a/docs/design/ws2-attention-cross-config-integration.md +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -27,9 +27,9 @@ tiers, in `rl_engine/alignment/cross_config/attention_binding.py`: | tier | fields | rule | failure | | --- | --- | --- | --- | -| `IDENTICAL` | checkpoint, model version, weight version, tokenizer, token ids, active mask, position ids, padding side, pre-update state, Q/KV heads, head dim, RoPE theta/scaling/rotary dim, QK-Norm, cached global token positions, KV sequence lengths | equal bit for bit | `comparable=False`; no drift number from the pair means anything | -| `SEMANTIC` | `reduction.merge`, `reduction.acc_dtype`, `reduction.order`, `reduction.downcast_at`, `export_lse`, cross-side determinism mode | both sides equal **and** equal to the WS2 mandate | fail closed | -| `RECORDED` | `mode`, `backend_id`, `reduction.engine`, RoPE materialization state and fusion boundary, KV-cache paging, CP/TP world sizes, local sequence length | free to differ | none; recorded into provenance and measured | +| `IDENTICAL` | checkpoint/model/token identity plus complete TP/CP GQA head and sequence ownership | equal bit for bit | `comparable=False`; no drift number from the pair means anything | +| `SEMANTIC` | reduction contract, dtype, exported LSE, first-class Split-KV request, complete actual Split-KV plan sets, cross-side determinism | both sides equal **and** equal to the WS2 mandate | fail closed | +| `RECORDED` | `mode`, `backend_id`, `reduction.engine`, RoPE materialization state and fusion boundary, KV-cache paging | free to differ | recorded into provenance and measured | Two placements are load-bearing: @@ -37,9 +37,13 @@ Two placements are load-bearing: deterministic reference while rollout runs a Transformer Engine merge oracle. Forcing them equal would defeat the oracle comparison that #235 PR2/PR3/PR5/PR6 depend on. -* **`reduction.order` and `reduction.acc_dtype` are `SEMANTIC`.** This is the entire - WS2 claim: merge order and accumulation precision are decided by the contract, not - by whichever backend happens to be selected. +* **TP/CP topology is `IDENTICAL`, not `RECORDED`.** TP selects local Qwen3 GQA head + ownership and CP selects local sequence ownership. Different topology is a + different local attention problem, not a backend detail. +* **`reduction.order`, `reduction.acc_dtype`, and actual Split-KV schedules are + `SEMANTIC`.** Both runtimes must export the complete batch x TP x CP x KV-owner + plan set, including logical boundaries, merge order, FP32 accumulation, final + downcast, and fallback state. Configured policy alone never passes strict binding. `comparable` and `passed` are separate flags. A pair with mismatched identity is not comparable. A pair that is comparable but violates the reduction mandate is still @@ -75,10 +79,16 @@ adapters: * `adapters/megatron.py` -- `MegatronProvenanceAdapter` (construction and distributed-context fingerprints, determinism probe, frozen-scope assertions) and `MegatronAttentionMaterializer`. -* `adapters/vllm.py` -- `VllmProvenanceAdapter` (adds `kv_page_size` from - `CacheConfig.block_size` and `split_kv_policy` from - `AttentionConfig.flash_attn_max_num_splits_for_cuda_graph`) and - `VllmRolloutMaterializer`. +* `adapters/vllm.py` -- `VllmProvenanceAdapter` (including diagnostic vLLM split + limits) and `VllmRolloutMaterializer`. +* `AttentionRuntimeReadback` -- the explicit handoff from an executed engine. It + carries the reconstructed actual contract, actual knob values, frozen-scope + verification, and the complete Split-KV runtime plan set. + +Constructing a contract is not runtime verification. Without a readback, adapter +applications are `UNOBSERVABLE`; only matching values reconstructed from a real +Megatron or vLLM execution are `APPLIED`. `bind_attention_runtime_readbacks` is the +strict public entry point used after both framework launchers collect that evidence. Neither module imports `megatron` or `vllm`; configs are duck-typed, so the binding rules are exercised on CPU in CI rather than only on a 2-node cluster. @@ -94,7 +104,9 @@ collapsing them onto the supported value: | `attention.reduction_downcast_at=per_block` | `UNSUPPORTED` | `DowncastPoint` declares only `final_write` | | `attention.reduction_engine=te_oracle` | `UNSUPPORTED` | the TE merge oracle lands in #235 PR2/PR3; PR4's TE plan is provenance only | | `attention.reduction_acc_dtype=bf16` | `UNSUPPORTED` | the CP `(out, lse)` merge accumulates in FP32 | -| `rollout.context_parallel_size>1` with `mode=decode` | `FALLBACK` | vLLM CP covers prefill only; recorded with the reason | +| configured contract without runtime readback | `UNOBSERVABLE` | requested values do not prove what executed | +| `rollout.context_parallel_size>1` with effective decode CP=1 | `ERROR`/`FALLBACK` | strict TP=2/CP=2 acceptance rejects the topology change | +| missing/mismatched/fallback Split-KV plan set | binding failure | Split-KV provenance must cover every batch/TP/CP/owner coordinate | ## Knobs diff --git a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json index e74d3e6c..11f9cef1 100644 --- a/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json +++ b/examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json @@ -21,7 +21,9 @@ "all and should be retired rather than rewritten.", "rollout.context_parallel_size binds to vLLM", "ParallelConfig.prefill_context_parallel_size and therefore applies to", - "prefill only; a decode-mode contract runs at CP=1." + "prefill only; strict PR4 acceptance covers CP=2 prefill/chunked prefill.", + "A decode request that becomes CP=1 is a blocking fallback and is tested", + "separately by the PR6 logical KV replay harness." ] }, "baseline": { @@ -30,7 +32,7 @@ }, "rollout": { "tensor_parallel_size": 2, - "context_parallel_size": 1, + "context_parallel_size": 2, "dtype": "bfloat16", "enable_prefix_caching": false, "enforce_eager": true, diff --git a/rl_engine/alignment/cross_config/adapters/__init__.py b/rl_engine/alignment/cross_config/adapters/__init__.py index c2db2134..f02b9b38 100644 --- a/rl_engine/alignment/cross_config/adapters/__init__.py +++ b/rl_engine/alignment/cross_config/adapters/__init__.py @@ -3,7 +3,11 @@ """Runtime adapters for the WS2 Qwen3-8B Megatron + vLLM cross-config target.""" -from rl_engine.alignment.cross_config.adapters._common import QWEN3_8B, Qwen3ModelSpec +from rl_engine.alignment.cross_config.adapters._common import ( + QWEN3_8B, + AttentionRuntimeReadback, + Qwen3ModelSpec, +) from rl_engine.alignment.cross_config.adapters.knobs import ( MEGATRON_ATTENTION_BACKENDS, WS2_ATTENTION_KNOB_DESCRIPTORS, @@ -21,6 +25,7 @@ __all__ = [ "MEGATRON_ATTENTION_BACKENDS", + "AttentionRuntimeReadback", "MegatronAttentionMaterializer", "MegatronProvenanceAdapter", "QWEN3_8B", diff --git a/rl_engine/alignment/cross_config/adapters/_common.py b/rl_engine/alignment/cross_config/adapters/_common.py index 33b4d91c..9a3c5384 100644 --- a/rl_engine/alignment/cross_config/adapters/_common.py +++ b/rl_engine/alignment/cross_config/adapters/_common.py @@ -9,6 +9,7 @@ from dataclasses import dataclass from typing import Any +from rl_engine.alignment.cross_config.attention_binding import AttentionRuntimeReadback from rl_engine.alignment.cross_config.runtime import KnobApplication from rl_engine.alignment.cross_config.schema import ( IsolationScope, @@ -23,17 +24,20 @@ ReductionOrder, ReductionSpec, ShardingSpec, + SplitKVSpec, ) __all__ = [ "QWEN3_8B", "Qwen3ModelSpec", + "AttentionRuntimeReadback", "application", "attention_dtype", "build_reduction_spec", "build_sharding_spec", "causal_offsets_for", "flatten", + "split_kv_spec", "unsupported_reduction_reason", ] @@ -107,6 +111,20 @@ def attention_dtype(value: Any, *, field: str) -> AttentionDType: ) from exc +def split_kv_spec(flat: Mapping[str, Any]) -> SplitKVSpec: + """Build the first-class logical Split-KV request. + + The integer is a fixed logical KV chunk size in tokens. It is intentionally + not vLLM's ``flash_attn_max_num_splits_for_cuda_graph``: that setting is only + an upper bound and cannot prove which runtime boundaries executed. + """ + + split_size = flat.get("attention.split_kv_policy") + if split_size is None: + return SplitKVSpec.disabled() + return SplitKVSpec.fixed(int(split_size)) + + def flatten(value: Mapping[str, Any], prefix: str = "") -> dict[str, Any]: """Flatten nested knob mappings into dotted paths.""" diff --git a/rl_engine/alignment/cross_config/adapters/knobs.py b/rl_engine/alignment/cross_config/adapters/knobs.py index c1cf0b84..33c29820 100644 --- a/rl_engine/alignment/cross_config/adapters/knobs.py +++ b/rl_engine/alignment/cross_config/adapters/knobs.py @@ -110,10 +110,11 @@ allowed_values=("unfused_rope_attention", "fused_rope_attention"), ), KnobDescriptor( - # vLLM: AttentionConfig.flash_attn_max_num_splits_for_cuda_graph + # Shared logical KV chunk size. Runtime adapters must separately report + # the actual per-owner boundaries; a configured value is not evidence. "attention.split_kv_policy", IsolationScope.ENGINE_CONSTRUCTION, - ("rollout",), + ("rollout", "training"), ), KnobDescriptor( # vLLM: CacheConfig.block_size -> AttentionContract.kv_cache.page_size diff --git a/rl_engine/alignment/cross_config/adapters/megatron.py b/rl_engine/alignment/cross_config/adapters/megatron.py index 8b96b159..8c9c9412 100644 --- a/rl_engine/alignment/cross_config/adapters/megatron.py +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -17,10 +17,10 @@ CPU model, so nothing had ever materialized a real distributed runtime. Scope boundary: materialization builds and validates the training-side -:class:`AttentionContract` and reports what would be constructed. It does not -launch ``torchrun``, initialize process groups, or execute attention. Binding a -constructed Megatron model to this contract is the next step and needs the 2-node -x 2-GPU environment that #239 fixes. +:class:`AttentionContract` and reports what would be constructed. Without an +``AttentionRuntimeReadback`` it reports ``UNOBSERVABLE``, never ``APPLIED``. It +does not launch ``torchrun``, initialize process groups, or execute attention; +the 2-node x 2-GPU launcher must inject readback collected after execution. """ from __future__ import annotations @@ -32,6 +32,7 @@ from rl_engine.alignment.cross_config.adapters._common import ( QWEN3_8B, + AttentionRuntimeReadback, Qwen3ModelSpec, application, attention_dtype, @@ -39,6 +40,7 @@ build_sharding_spec, causal_offsets_for, flatten, + split_kv_spec, unsupported_reduction_reason, ) from rl_engine.alignment.cross_config.determinism import ( @@ -214,6 +216,7 @@ def __init__( cp_rank: int = 0, backend_id: str = "rlkernel.cp_attention_reference", provenance: Optional[MegatronProvenanceAdapter] = None, + runtime_readback: Optional[AttentionRuntimeReadback] = None, ): self.model = model self.global_sequence_length = global_sequence_length @@ -221,6 +224,7 @@ def __init__( self.cp_rank = cp_rank self.backend_id = backend_id self.provenance = provenance + self.runtime_readback = runtime_readback @property def implementation_fingerprint(self) -> str: @@ -272,6 +276,7 @@ def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: causal_offsets=causal_offsets_for(sharding, batch_size), sharding=sharding, reduction=build_reduction_spec(flat), + split_kv=split_kv_spec(flat), rope=rope, export_lse=True, ) @@ -326,14 +331,11 @@ def materialize( ) continue applications.append( - application( + self._runtime_application( descriptor, requested, - requested, - requested, - MaterializationStatus.APPLIED, - "bound to the training-side attention contract", - frozen_scope_violations=list(scope_violations), + contract=contract, + scope_violations=scope_violations, ) ) @@ -347,6 +349,9 @@ def materialize( "cp_comm_type": flat.get("training.cp_comm_type"), "contract": contract.to_dict() if contract is not None else None, "contract_error": contract_error or blocked, + "runtime_readback": ( + None if self.runtime_readback is None else self.runtime_readback.to_dict() + ), "frozen_scope_violations": list(scope_violations), } if self.provenance is not None: @@ -379,3 +384,86 @@ def materialize( runtime_kind=self.runtime_kind, ), ) + + def _runtime_application( + self, + descriptor: KnobDescriptor, + requested: Any, + *, + contract: AttentionContract, + scope_violations: tuple[str, ...], + ) -> KnobApplication: + readback = self.runtime_readback + if readback is None: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + ( + "configured in the training contract, but no Megatron runtime " + "readback was supplied" + ), + frozen_scope_violations=list(scope_violations), + ) + if scope_violations or not readback.frozen_scope_verified: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.UNOBSERVABLE, + "Megatron frozen-scope assertions were not all verified by runtime readback", + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) + if readback.split_kv_fallback: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "Megatron runtime reported a Split-KV fallback", + runtime_readback_source=readback.source, + ) + if readback.contract != contract: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "Megatron runtime contract differs from the requested contract", + runtime_readback_source=readback.source, + ) + if descriptor.path not in readback.actual_knobs: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "Megatron runtime readback does not expose this knob", + runtime_readback_source=readback.source, + ) + actual = readback.actual_knobs[descriptor.path] + status = ( + MaterializationStatus.APPLIED if actual == requested else MaterializationStatus.FALLBACK + ) + reason = ( + "verified from the executed Megatron runtime" + if status is MaterializationStatus.APPLIED + else "Megatron runtime value differs from the requested value" + ) + return application( + descriptor, + requested, + requested, + actual, + status, + reason, + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) diff --git a/rl_engine/alignment/cross_config/adapters/vllm.py b/rl_engine/alignment/cross_config/adapters/vllm.py index a8de61c9..1e1ae185 100644 --- a/rl_engine/alignment/cross_config/adapters/vllm.py +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -16,7 +16,8 @@ ``init_batch_invariance()`` at worker startup. Like the Megatron adapter, nothing here imports ``vllm``; configs are duck-typed so -the module is importable anywhere. +the module is importable anywhere. Configured-only values remain ``UNOBSERVABLE``; +``APPLIED`` requires an explicit post-execution ``AttentionRuntimeReadback``. """ from __future__ import annotations @@ -28,6 +29,7 @@ from rl_engine.alignment.cross_config.adapters._common import ( QWEN3_8B, + AttentionRuntimeReadback, Qwen3ModelSpec, application, attention_dtype, @@ -35,6 +37,7 @@ build_sharding_spec, causal_offsets_for, flatten, + split_kv_spec, unsupported_reduction_reason, ) from rl_engine.alignment.cross_config.determinism import DeterminismProbe, vllm_probe_from_env @@ -204,7 +207,7 @@ def kv_page_size(self) -> Optional[int]: @property def split_kv_policy(self) -> Optional[int]: - """The split-KV knob #235 PR5/PR7 needs; #236 has no field for it yet.""" + """Diagnostic vLLM maximum split count, not the logical chunk-size contract.""" splits = _value(self.attention_config, "flash_attn_max_num_splits_for_cuda_graph") return int(splits) if splits is not None else None @@ -218,7 +221,7 @@ def to_dict(self) -> dict[str, Any]: "distributed_context_fingerprint": self.distributed_context_fingerprint, "frozen_scope_violations": list(self.frozen_scope_violations()), "kv_page_size": self.kv_page_size, - "split_kv_policy": self.split_kv_policy, + "flash_attn_max_num_splits_for_cuda_graph": self.split_kv_policy, "determinism": self.determinism_probe().to_dict(), } @@ -238,6 +241,7 @@ def __init__( mode: AttentionMode = AttentionMode.CHUNKED_PREFILL, backend_id: str = "vllm.flash_attn", provenance: Optional[VllmProvenanceAdapter] = None, + runtime_readback: Optional[AttentionRuntimeReadback] = None, ): self.model = model self.global_sequence_length = global_sequence_length @@ -246,6 +250,7 @@ def __init__( self.mode = mode self.backend_id = backend_id self.provenance = provenance + self.runtime_readback = runtime_readback @property def implementation_fingerprint(self) -> str: @@ -312,6 +317,7 @@ def build_contract(self, flat: Mapping[str, Any]) -> AttentionContract: causal_offsets=causal_offsets_for(sharding, batch_size), sharding=sharding, reduction=build_reduction_spec(flat), + split_kv=split_kv_spec(flat), rope=rope, export_lse=True, ) @@ -385,14 +391,11 @@ def materialize( ) continue applications.append( - application( + self._runtime_application( descriptor, requested, - requested, - requested, - MaterializationStatus.APPLIED, - "bound to the rollout-side attention contract", - frozen_scope_violations=list(scope_violations), + contract=contract, + scope_violations=scope_violations, ) ) @@ -408,6 +411,9 @@ def materialize( "attention_mode": self.mode.value, "contract": contract.to_dict() if contract is not None else None, "contract_error": contract_error or blocked, + "runtime_readback": ( + None if self.runtime_readback is None else self.runtime_readback.to_dict() + ), "frozen_scope_violations": list(scope_violations), } if self.provenance is not None: @@ -440,3 +446,83 @@ def materialize( runtime_kind=self.runtime_kind, ), ) + + def _runtime_application( + self, + descriptor: KnobDescriptor, + requested: Any, + *, + contract: AttentionContract, + scope_violations: tuple[str, ...], + ) -> KnobApplication: + readback = self.runtime_readback + if readback is None: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "configured in the rollout contract, but no vLLM runtime readback was supplied", + frozen_scope_violations=list(scope_violations), + ) + if scope_violations or not readback.frozen_scope_verified: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.UNOBSERVABLE, + "vLLM frozen-scope assertions were not all verified by runtime readback", + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) + if readback.split_kv_fallback: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "vLLM runtime reported a Split-KV fallback", + runtime_readback_source=readback.source, + ) + if readback.contract != contract: + return application( + descriptor, + requested, + requested, + readback.actual_knobs.get(descriptor.path), + MaterializationStatus.FALLBACK, + "vLLM runtime contract differs from the requested contract", + runtime_readback_source=readback.source, + ) + if descriptor.path not in readback.actual_knobs: + return application( + descriptor, + requested, + requested, + None, + MaterializationStatus.UNOBSERVABLE, + "vLLM runtime readback does not expose this knob", + runtime_readback_source=readback.source, + ) + actual = readback.actual_knobs[descriptor.path] + status = ( + MaterializationStatus.APPLIED if actual == requested else MaterializationStatus.FALLBACK + ) + reason = ( + "verified from the executed vLLM runtime" + if status is MaterializationStatus.APPLIED + else "vLLM runtime value differs from the requested value" + ) + return application( + descriptor, + requested, + requested, + actual, + status, + reason, + runtime_readback_source=readback.source, + frozen_scope_violations=list(scope_violations), + ) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 48e382b8..4156d57e 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -40,21 +40,26 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from enum import Enum +from types import MappingProxyType from typing import Any, Optional from rl_engine.kernels.attention_contract import ( AttentionContract, + AttentionContractError, AttentionDType, AttentionMerge, AttentionRole, DowncastPoint, ReductionOrder, + SplitKVRuntimePlanSet, + validate_split_kv_plan_set_alignment, ) __all__ = [ "ATTENTION_LSE_DOMAIN", "AttentionBindingError", "AttentionBindingResult", + "AttentionRuntimeReadback", "BindingErrorCode", "BindingIssue", "BindingTier", @@ -63,8 +68,10 @@ "RECORDED_FIELDS", "SEMANTIC_CONTRACT_FIELDS", "SEMANTIC_REDUCTION_FIELDS", + "TOPOLOGY_FIELDS", "WS2_ATTENTION_REDUCTION_MANDATE", "bind_attention_contracts", + "bind_attention_runtime_readbacks", "first_blocking_issue", "identity_fingerprint", "summarize_binding", @@ -97,6 +104,51 @@ class BindingErrorCode(str, Enum): LSE_NOT_EXPORTED = "LSE_NOT_EXPORTED" ROLE_COLLISION = "ROLE_COLLISION" DETERMINISM_INCOMPATIBLE = "DETERMINISM_INCOMPATIBLE" + TOPOLOGY_MISMATCH = "TOPOLOGY_MISMATCH" + SPLIT_KV_RUNTIME_MISSING = "SPLIT_KV_RUNTIME_MISSING" + SPLIT_KV_MISMATCH = "SPLIT_KV_MISMATCH" + SPLIT_KV_FALLBACK = "SPLIT_KV_FALLBACK" + + +@dataclass(frozen=True) +class AttentionRuntimeReadback: + """Actual attention contract and all-rank Split-KV evidence from one engine.""" + + contract: AttentionContract + actual_knobs: Mapping[str, Any] + split_kv_plan_set: SplitKVRuntimePlanSet + source: str + frozen_scope_verified: bool + + def __post_init__(self) -> None: + if not isinstance(self.contract, AttentionContract): + raise TypeError("runtime readback contract must be an AttentionContract") + if not isinstance(self.actual_knobs, Mapping): + raise TypeError("runtime readback actual_knobs must be a mapping") + if not isinstance(self.split_kv_plan_set, SplitKVRuntimePlanSet): + raise TypeError("runtime readback requires a complete SplitKVRuntimePlanSet") + if not isinstance(self.source, str) or not self.source.strip(): + raise ValueError("runtime readback source must be a non-empty string") + if not isinstance(self.frozen_scope_verified, bool): + raise TypeError("frozen_scope_verified must be a bool") + + plan_error = _split_kv_plan_contract_error(self.contract, self.split_kv_plan_set) + if plan_error is not None: + raise ValueError(plan_error) + object.__setattr__(self, "actual_knobs", MappingProxyType(dict(self.actual_knobs))) + + @property + def split_kv_fallback(self) -> bool: + return bool(_split_kv_fallbacks(self.split_kv_plan_set)) + + def to_dict(self) -> dict[str, Any]: + return { + "source": self.source, + "frozen_scope_verified": self.frozen_scope_verified, + "contract": self.contract.to_dict(), + "actual_knobs": dict(self.actual_knobs), + "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), + } #: Attention exports attention-domain LSE, never vocab-logprob LSE (#235). @@ -154,6 +206,29 @@ class BindingErrorCode(str, Enum): SEMANTIC_CONTRACT_FIELDS: tuple[str, ...] = ("dtype",) +#: Sharding fields that determine local GQA head and sequence ownership. These are +#: comparison preconditions, not harmless backend provenance: a TP/CP mismatch +#: means the two ranks did not evaluate the same local attention problem. +TOPOLOGY_FIELDS: tuple[str, ...] = ( + "tp_rank", + "tp_world_size", + "cp_rank", + "cp_world_size", + "global_q_heads", + "global_kv_heads", + "local_q_head_start", + "local_q_heads", + "local_kv_head_start", + "local_kv_heads", + "global_sequence_length", + "local_sequence_length", + "global_block_indices", + "global_block_token_starts", + "local_block_offsets", + "packed_sequence_offsets", +) + + #: The WS2 mandate itself. ``#236`` currently declares single-member enums for #: ``merge`` / ``order`` / ``downcast_at``, so those checks are tautological today; #: they are written out anyway so that widening any of those enums later fails here @@ -181,14 +256,6 @@ class BindingErrorCode(str, Enum): "kv_cache.page_size", "kv_cache.prefix_cache_enabled", "kv_cache.block_table_shape", - "sharding.cp_world_size", - "sharding.tp_world_size", - "sharding.local_sequence_length", - # Supplied by the caller, not by the contract: #236 has no split-KV field yet, so - # the value comes from vLLM's flash_attn_max_num_splits_for_cuda_graph via the - # adapter. Recorded so split-KV differences are at least visible in provenance - # until #236 grows the field and it can move into the contract proper. - "split_kv_policy", ) @@ -233,7 +300,7 @@ class AttentionBindingResult: binding_fingerprint: str = "" recorded_differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) provenance: Mapping[str, Any] = field(default_factory=dict) - schema_version: str = "cross_config.attention_binding.v1" + schema_version: str = "cross_config.attention_binding.v2" def issues_by_code(self, code: BindingErrorCode) -> tuple[BindingIssue, ...]: return tuple(issue for issue in self.issues if issue.code is code) @@ -291,9 +358,6 @@ def _recorded_view( "mode": contract.mode.value, "backend_id": None, "reduction.engine": contract.reduction.engine.value, - "sharding.cp_world_size": contract.sharding.cp_world_size, - "sharding.tp_world_size": contract.sharding.tp_world_size, - "sharding.local_sequence_length": contract.sharding.local_sequence_length, } if rope is not None: view.update( @@ -322,6 +386,69 @@ def _recorded_view( return view +def _topology_view(contract: AttentionContract) -> dict[str, Any]: + sharding = contract.sharding + return {name: getattr(sharding, name) for name in TOPOLOGY_FIELDS} + + +def _split_kv_fallbacks(plan_set: SplitKVRuntimePlanSet) -> list[dict[str, Any]]: + return [ + entry.to_dict() + for entry in plan_set.entries + if entry.execution.fallback + or entry.execution.actual_mode is None + or entry.execution.actual_mode is not entry.execution.requested_mode + or entry.execution.actual_split_size != entry.execution.requested_split_size + ] + + +def _split_kv_plan_contract_error( + contract: AttentionContract, + plan_set: SplitKVRuntimePlanSet, +) -> str | None: + sharding = contract.sharding + expected_topology = ( + contract.batch_size, + sharding.tp_world_size, + sharding.cp_world_size, + ) + actual_topology = ( + plan_set.batch_size, + plan_set.tp_world_size, + plan_set.cp_world_size, + ) + if actual_topology != expected_topology: + return ( + "Split-KV plan-set batch/TP/CP topology does not match the attention " + f"contract: actual={actual_topology}, expected={expected_topology}" + ) + if contract.mode.value in {"prefill", "chunked_prefill"}: + expected_totals = (sharding.global_sequence_length,) * contract.batch_size + if plan_set.total_kv_tokens != expected_totals: + return ( + "Split-KV plan-set KV lengths do not match the prefill attention " + f"contract: actual={plan_set.total_kv_tokens}, expected={expected_totals}" + ) + elif contract.kv_cache is not None: + expected_totals = contract.kv_cache.kv_seq_lens + if plan_set.total_kv_tokens != expected_totals: + return ( + "Split-KV plan-set KV lengths do not match decode KV-cache lengths: " + f"actual={plan_set.total_kv_tokens}, expected={expected_totals}" + ) + for entry in plan_set.entries: + execution = entry.execution + if ( + execution.requested_mode is not contract.split_kv.mode + or execution.requested_split_size != contract.split_kv.fixed_split_size + ): + return ( + "Split-KV runtime request does not match the first-class attention " + f"contract at {entry.coordinate}" + ) + return None + + #: Identity fields where ``None`` is a real value rather than an omission. Qwen3-8B #: applies no RoPE scaling, so ``rope_scaling=None`` must not read as "undeclared" -- #: both sides still have to agree on it, which the equality pass below handles. @@ -347,6 +474,8 @@ def bind_attention_contracts( determinism_issues: Sequence[BindingIssue] = (), require_full_identity: bool = True, allow_dtype_difference: bool = False, + rollout_split_kv_plan_set: Optional[SplitKVRuntimePlanSet] = None, + training_split_kv_plan_set: Optional[SplitKVRuntimePlanSet] = None, rollout_recorded_extra: Optional[Mapping[str, Any]] = None, training_recorded_extra: Optional[Mapping[str, Any]] = None, ) -> AttentionBindingResult: @@ -364,9 +493,13 @@ def bind_attention_contracts( ``allow_dtype_difference`` exists for the #235 PR5 sweep that deliberately scores a BF16 path against an FP32 reference. It must stay ``False`` everywhere else. - ``rollout_recorded_extra`` / ``training_recorded_extra`` carry materialization - facts that #236 does not yet model -- today that is ``split_kv_policy``. They are - merged into the recorded tier, never into identity or semantics. + Strict binding requires complete actual Split-KV plan sets from both runtimes. + A configured policy is insufficient because auto-selection, graph capture, and + backend fallbacks can change the executed boundaries. The plan sets cover the + complete batch x TP x CP x KV-owner Cartesian product. + + ``rollout_recorded_extra`` / ``training_recorded_extra`` are diagnostic-only + backend facts. They can never make a semantic mismatch admissible. """ if rollout_contract.role is not AttentionRole.INFER: @@ -415,6 +548,26 @@ def bind_attention_contracts( comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + rollout_topology = _topology_view(rollout_contract) + training_topology = _topology_view(training_contract) + for name in TOPOLOGY_FIELDS: + if rollout_topology[name] != training_topology[name]: + issues.append( + BindingIssue( + code=BindingErrorCode.TOPOLOGY_MISMATCH, + tier=BindingTier.IDENTICAL, + field=f"sharding.{name}", + rollout=rollout_topology[name], + training=training_topology[name], + message=( + f"sharding.{name} changes TP/CP ownership; the pair is not " + "the same local attention problem" + ), + ) + ) + + comparable = not any(issue.tier is BindingTier.IDENTICAL for issue in issues) + # ---- tier 2: reduction semantics, and the WS2 mandate ------------------- rollout_reduction = _reduction_view(rollout_contract) training_reduction = _reduction_view(training_contract) @@ -467,6 +620,79 @@ def bind_attention_contracts( ) ) + if rollout_contract.split_kv != training_contract.split_kv: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field="split_kv", + rollout=rollout_contract.split_kv.to_dict(), + training=training_contract.split_kv.to_dict(), + message="training and rollout must request the same first-class Split-KV policy", + ) + ) + + for side, plan_set in ( + ("rollout", rollout_split_kv_plan_set), + ("training", training_split_kv_plan_set), + ): + if plan_set is None: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_RUNTIME_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + message=( + f"{side} did not report a complete actual Split-KV plan set; " + "configured policy alone is not runtime evidence" + ), + ) + ) + continue + contract = rollout_contract if side == "rollout" else training_contract + contract_error = _split_kv_plan_contract_error(contract, plan_set) + if contract_error is not None: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + rollout=plan_set.to_dict() if side == "rollout" else None, + training=plan_set.to_dict() if side == "training" else None, + message=contract_error, + ) + ) + fallbacks = _split_kv_fallbacks(plan_set) + if fallbacks: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_FALLBACK, + tier=BindingTier.SEMANTIC, + field=f"{side}.split_kv_runtime_plan_set", + rollout=fallbacks if side == "rollout" else None, + training=fallbacks if side == "training" else None, + message=f"{side} Split-KV runtime used an unknown or fallback plan", + ) + ) + + if rollout_split_kv_plan_set is not None and training_split_kv_plan_set is not None: + try: + validate_split_kv_plan_set_alignment( + training_split_kv_plan_set, + rollout_split_kv_plan_set, + ) + except AttentionContractError as exc: + issues.append( + BindingIssue( + code=BindingErrorCode.SPLIT_KV_MISMATCH, + tier=BindingTier.SEMANTIC, + field="split_kv_runtime_plan_set", + rollout=rollout_split_kv_plan_set.to_dict(), + training=training_split_kv_plan_set.to_dict(), + message=str(exc), + ) + ) + for side, contract in (("rollout", rollout_contract), ("training", training_contract)): if not contract.export_lse: issues.append( @@ -513,6 +739,14 @@ def bind_attention_contracts( provenance = { "lse_domain": ATTENTION_LSE_DOMAIN, "dtype": training_contract.dtype.value, + "split_kv_runtime": { + "rollout": ( + None if rollout_split_kv_plan_set is None else rollout_split_kv_plan_set.to_dict() + ), + "training": ( + None if training_split_kv_plan_set is None else training_split_kv_plan_set.to_dict() + ), + }, "rollout": { "contract": rollout_contract.to_dict(), "backend_id": rollout_backend_id, @@ -535,6 +769,8 @@ def bind_attention_contracts( { "identity": identity_fp, "reduction": reduction_fp, + "topology": training_topology, + "split_kv": provenance["split_kv_runtime"], "lse_domain": ATTENTION_LSE_DOMAIN, "rollout_backend": rollout_backend_id, "training_backend": training_backend_id, @@ -545,6 +781,47 @@ def bind_attention_contracts( ) +def bind_attention_runtime_readbacks( + *, + rollout: AttentionRuntimeReadback, + training: AttentionRuntimeReadback, + rollout_identity: Mapping[str, Any], + training_identity: Mapping[str, Any], + rollout_backend_id: str, + training_backend_id: str, + determinism_issues: Sequence[BindingIssue] = (), +) -> AttentionBindingResult: + """Strict public handoff from executed framework runtimes to PR4 binding. + + The Megatron/vLLM launchers remain environment-owned. Once both launchers have + reconstructed their actual contracts and all-rank Split-KV reports, this entry + point performs the complete comparison without accepting configured-only data. + """ + + missing_scope_evidence = [] + for side, readback in (("rollout", rollout), ("training", training)): + if not readback.frozen_scope_verified: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.DETERMINISM_INCOMPATIBLE, + tier=BindingTier.SEMANTIC, + field=f"{side}.frozen_scope_verified", + message=f"{side} runtime did not verify the frozen attention scope", + ) + ) + return bind_attention_contracts( + rollout_contract=rollout.contract, + training_contract=training.contract, + rollout_identity=rollout_identity, + training_identity=training_identity, + rollout_backend_id=rollout_backend_id, + training_backend_id=training_backend_id, + determinism_issues=tuple(determinism_issues) + tuple(missing_scope_evidence), + rollout_split_kv_plan_set=rollout.split_kv_plan_set, + training_split_kv_plan_set=training.split_kv_plan_set, + ) + + def summarize_binding(result: AttentionBindingResult) -> str: """One-line human summary for CLI output and failure messages.""" diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 815daf12..6a99e6bb 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -11,6 +11,8 @@ from __future__ import annotations import json +from dataclasses import replace +from enum import Enum from pathlib import Path from types import SimpleNamespace @@ -19,6 +21,7 @@ from rl_engine.alignment.cross_config.adapters import ( QWEN3_8B, WS2_ATTENTION_KNOBS, + AttentionRuntimeReadback, MegatronAttentionMaterializer, MegatronProvenanceAdapter, VllmProvenanceAdapter, @@ -30,6 +33,7 @@ BindingErrorCode, BindingTier, bind_attention_contracts, + bind_attention_runtime_readbacks, first_blocking_issue, identity_fingerprint, summarize_binding, @@ -40,7 +44,17 @@ vllm_probe_from_env, ) from rl_engine.alignment.cross_config.schema import MaterializationStatus -from rl_engine.kernels.attention_contract import AttentionContractError, AttentionMode +from rl_engine.kernels.attention_contract import ( + AttentionContractError, + AttentionMode, + AttentionRole, + KVCacheSpec, + SplitKVExecutionPlan, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, + SplitKVSpec, + build_split_kv_runtime_plan_set, +) pytestmark = pytest.mark.unit @@ -50,13 +64,15 @@ "training.tensor_parallel_size": 2, "training.context_parallel_size": 2, "training.compute_dtype": "bf16", + "attention.split_kv_policy": 32, } ROLLOUT_KNOBS = { "batch.size": 2, "rollout.tensor_parallel_size": 2, - "rollout.context_parallel_size": 1, + "rollout.context_parallel_size": 2, "rollout.dtype": "bf16", + "attention.split_kv_policy": 32, } @@ -86,15 +102,33 @@ def _contracts(): return rollout, training +def _plan_set(contract, *, backend): + return build_split_kv_runtime_plan_set( + (contract.sharding.global_sequence_length,) * contract.batch_size, + tp_world_size=contract.sharding.tp_world_size, + cp_world_size=contract.sharding.cp_world_size, + split_kv=contract.split_kv, + backend=backend, + ) + + def _bind(rollout_identity=None, training_identity=None, **kwargs): rollout, training = _contracts() + rollout = kwargs.pop("rollout_contract", rollout) + training = kwargs.pop("training_contract", training) return bind_attention_contracts( - rollout_contract=kwargs.pop("rollout_contract", rollout), - training_contract=kwargs.pop("training_contract", training), + rollout_contract=rollout, + training_contract=training, rollout_identity=rollout_identity if rollout_identity is not None else _identity(), training_identity=training_identity if training_identity is not None else _identity(), rollout_backend_id="vllm.flash_attn", training_backend_id="rlkernel.cp_attention_reference", + rollout_split_kv_plan_set=kwargs.pop( + "rollout_split_kv_plan_set", _plan_set(rollout, backend="vllm.readback") + ), + training_split_kv_plan_set=kwargs.pop( + "training_split_kv_plan_set", _plan_set(training, backend="megatron.readback") + ), **kwargs, ) @@ -104,7 +138,7 @@ def _bind(rollout_identity=None, training_identity=None, **kwargs): # -------------------------------------------------------------------------- -def test_matching_identity_binds_despite_different_materialization(): +def test_matching_identity_and_topology_bind_despite_different_materialization(): """The core claim of PR4: same identity + same reduction, different runtimes.""" result = _bind() @@ -112,10 +146,9 @@ def test_matching_identity_binds_despite_different_materialization(): assert result.comparable assert result.passed assert result.issues == () - # Training runs CP=2 full prefill, rollout runs CP=1 chunked prefill. Those - # differences are recorded, not rejected. + # Attention mode is a framework materialization difference, while both sides + # execute the same TP=2/CP=2 local ownership and Split-K schedule. assert "mode" in result.recorded_differences - assert "sharding.cp_world_size" in result.recorded_differences assert result.recorded_differences["mode"] == { "rollout": "chunked_prefill", "training": "prefill", @@ -259,28 +292,133 @@ def test_batch_size_mismatch_is_not_comparable(): assert any(issue.field == "batch_size" for issue in result.issues) -def test_split_kv_policy_difference_is_recorded(): - """#236 has no split-KV field, so the adapter supplies it to the recorded tier.""" +def test_missing_split_kv_runtime_evidence_fails_closed(): + result = _bind(rollout_split_kv_plan_set=None) + + assert result.comparable + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_RUNTIME_MISSING) + + +def test_split_kv_requested_policy_mismatch_fails_closed(): + rollout, training = _contracts() + rollout = replace(rollout, split_kv=SplitKVSpec.fixed(16)) + result = _bind(rollout_contract=rollout) + + assert result.comparable + assert not result.passed + assert any(issue.field == "split_kv" for issue in result.issues) + - result = _bind( - rollout_recorded_extra={"split_kv_policy": 8}, - training_recorded_extra={"split_kv_policy": None}, +def test_split_kv_runtime_boundary_mismatch_fails_closed(): + rollout, training = _contracts() + mismatched_rollout = build_split_kv_runtime_plan_set( + (rollout.sharding.global_sequence_length,) * rollout.batch_size, + tp_world_size=2, + cp_world_size=2, + split_kv=SplitKVSpec.fixed(16), + backend="vllm.readback", ) + result = _bind(rollout_split_kv_plan_set=mismatched_rollout) - assert result.passed - assert result.recorded_differences["split_kv_policy"] == { - "rollout": 8, - "training": None, - } + assert not result.passed + issues = result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + assert any(issue.field == "split_kv_runtime_plan_set" for issue in issues) + + +def test_split_kv_plan_set_must_match_its_own_contract_topology(): + rollout, _ = _contracts() + wrong_topology = build_split_kv_runtime_plan_set( + (rollout.sharding.global_sequence_length,) * rollout.batch_size, + tp_world_size=1, + cp_world_size=2, + split_kv=rollout.split_kv, + backend="vllm.readback", + ) + result = _bind(rollout_split_kv_plan_set=wrong_topology) -def test_matching_split_kv_policy_is_not_reported_as_a_difference(): - result = _bind( - rollout_recorded_extra={"split_kv_policy": 32}, - training_recorded_extra={"split_kv_policy": 32}, + assert not result.passed + assert any( + issue.field == "rollout.split_kv_runtime_plan_set" + for issue in result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) ) - assert "split_kv_policy" not in result.recorded_differences + +@pytest.mark.parametrize( + ("field_name", "corrupt_value"), + [ + ("merge_order", Enum("BadOrder", {"ARRIVAL": "arrival"}).ARRIVAL), + ("acc_dtype", Enum("BadDType", {"BF16": "bf16"}).BF16), + ("downcast_at", Enum("BadDowncast", {"PER_BLOCK": "per_block"}).PER_BLOCK), + ], +) +def test_split_kv_runtime_merge_semantic_corruption_fails_closed(field_name, corrupt_value): + rollout, _ = _contracts() + corrupted = _plan_set(rollout, backend="vllm.readback") + # Runtime reports are deserialized at this boundary. Simulate a corrupted + # report after construction to prove binding compares the actual fields. + object.__setattr__(corrupted.entries[0].execution, field_name, corrupt_value) + + result = _bind(rollout_split_kv_plan_set=corrupted) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_MISMATCH) + + +def test_split_kv_runtime_fallback_fails_closed(): + rollout, _ = _contracts() + plan_set = _plan_set(rollout, backend="vllm.readback") + fallback_entries = [] + for entry in plan_set.entries: + execution = entry.execution + fallback_entries.append( + SplitKVRuntimePlanEntry( + coordinate=entry.coordinate, + expected_kv_range=entry.expected_kv_range, + execution=SplitKVExecutionPlan( + requested_mode=execution.requested_mode, + requested_split_size=execution.requested_split_size, + actual_mode=execution.actual_mode, + actual_split_size=execution.actual_split_size, + boundaries=execution.boundaries, + merge_order=execution.merge_order, + acc_dtype=execution.acc_dtype, + downcast_at=execution.downcast_at, + backend=execution.backend, + source="runtime_fallback", + fallback=True, + fallback_reason="backend substituted a runtime plan", + ), + ) + ) + fallback = SplitKVRuntimePlanSet( + batch_size=plan_set.batch_size, + tp_world_size=plan_set.tp_world_size, + cp_world_size=plan_set.cp_world_size, + total_kv_tokens=plan_set.total_kv_tokens, + entries=tuple(fallback_entries), + ) + + result = _bind(rollout_split_kv_plan_set=fallback) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.SPLIT_KV_FALLBACK) + + +@pytest.mark.parametrize( + "rollout_overrides", + [ + {"rollout.tensor_parallel_size": 1}, + {"rollout.context_parallel_size": 1}, + ], +) +def test_tp_or_cp_topology_mismatch_is_not_comparable(rollout_overrides): + rollout = VllmRolloutMaterializer().build_contract({**ROLLOUT_KNOBS, **rollout_overrides}) + result = _bind(rollout_contract=rollout) + + assert not result.comparable + assert result.issues_by_code(BindingErrorCode.TOPOLOGY_MISMATCH) # -------------------------------------------------------------------------- @@ -480,6 +618,175 @@ def _statuses(materialization, path): return [app.status for app in materialization.applications if app.path == path] +def _readback(materializer, flat, *, source): + contract = materializer.build_contract(flat) + return AttentionRuntimeReadback( + contract=contract, + actual_knobs=dict(flat), + split_kv_plan_set=_plan_set(contract, backend=source), + source=source, + frozen_scope_verified=True, + ) + + +def test_configured_contract_without_runtime_readback_is_unobservable(): + normalized = { + "batch": {"size": 2}, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "compute_dtype": "bf16", + }, + "attention": {"split_kv_policy": 32}, + } + materialization = MegatronAttentionMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert materialization.applications + assert {app.status for app in materialization.applications} == { + MaterializationStatus.UNOBSERVABLE + } + assert all(app.actual is None for app in materialization.applications) + + +@pytest.mark.parametrize( + ("materializer_type", "flat", "source"), + [ + (MegatronAttentionMaterializer, TRAINING_KNOBS, "megatron.runtime_readback"), + (VllmRolloutMaterializer, ROLLOUT_KNOBS, "vllm.runtime_readback"), + ], +) +def test_runtime_readback_can_verify_materialized_knobs(materializer_type, flat, source): + configured = materializer_type() + readback = _readback(configured, flat, source=source) + materializer = materializer_type(runtime_readback=readback) + normalized = {} + for path, value in flat.items(): + section, key = path.split(".", 1) + normalized.setdefault(section, {})[key] = value + + materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) + + assert materialization.applications + assert {app.status for app in materialization.applications} == {MaterializationStatus.APPLIED} + side = "training" if materializer_type is MegatronAttentionMaterializer else "rollout" + assert materialization.binding.side_configs[side]["runtime_readback"]["source"] == source + + +def test_runtime_readback_mismatch_is_a_fallback(): + configured = MegatronAttentionMaterializer() + readback = _readback(configured, TRAINING_KNOBS, source="megatron.runtime_readback") + actual = dict(readback.actual_knobs) + actual["training.context_parallel_size"] = 1 + mismatched = AttentionRuntimeReadback( + contract=readback.contract, + actual_knobs=actual, + split_kv_plan_set=readback.split_kv_plan_set, + source=readback.source, + frozen_scope_verified=True, + ) + normalized = { + "batch": {"size": 2}, + "training": { + "tensor_parallel_size": 2, + "context_parallel_size": 2, + "compute_dtype": "bf16", + }, + "attention": {"split_kv_policy": 32}, + } + materialization = MegatronAttentionMaterializer(runtime_readback=mismatched).materialize( + normalized, WS2_ATTENTION_KNOBS + ) + + assert _statuses(materialization, "training.context_parallel_size") == [ + MaterializationStatus.FALLBACK + ] + + +def test_decode_split_kv_plan_set_must_match_kv_cache_lengths(): + rollout, _ = _contracts() + kv_cache = KVCacheSpec( + cache_positions=(3, 5), + kv_seq_lens=(4, 6), + block_table=((0, 1, -1), (2, 3, 4)), + global_token_positions=tuple(range(4)) + tuple(range(6)), + page_size=2, + ) + decode = replace( + rollout, + role=AttentionRole.INFER, + mode=AttentionMode.DECODE, + query_sequence_length=1, + causal_offsets=(3, 5), + kv_cache=kv_cache, + ) + wrong_lengths = build_split_kv_runtime_plan_set( + (4, 8), + tp_world_size=2, + cp_world_size=2, + split_kv=decode.split_kv, + backend="vllm.decode.readback", + ) + + with pytest.raises(ValueError, match="KV-cache lengths"): + AttentionRuntimeReadback( + contract=decode, + actual_knobs={}, + split_kv_plan_set=wrong_lengths, + source="vllm.decode.readback", + frozen_scope_verified=True, + ) + + +def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): + rollout_materializer = VllmRolloutMaterializer() + training_materializer = MegatronAttentionMaterializer() + rollout = _readback(rollout_materializer, ROLLOUT_KNOBS, source="vllm.runtime_readback") + training = _readback( + training_materializer, + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.passed + assert result.provenance["split_kv_runtime"]["rollout"]["coverage"] == ( + "complete_batch_tp_cp_owner_cartesian_product" + ) + + +def test_strict_runtime_readback_entrypoint_rejects_unverified_frozen_scope(): + rollout_materializer = VllmRolloutMaterializer() + training_materializer = MegatronAttentionMaterializer() + rollout = _readback(rollout_materializer, ROLLOUT_KNOBS, source="vllm.runtime_readback") + rollout = replace(rollout, frozen_scope_verified=False) + training = _readback( + training_materializer, + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.comparable + assert not result.passed + assert any(issue.field == "rollout.frozen_scope_verified" for issue in result.issues) + + def test_arrival_merge_order_is_unsupported_not_silently_corrected(): """The control group must stay distinguishable from the treatment.""" @@ -535,6 +842,7 @@ def test_vllm_cp_falls_back_to_one_in_decode_and_says_why(): materialization = materializer.materialize(normalized, WS2_ATTENTION_KNOBS) contract_error = materialization.binding.side_configs["rollout"]["contract_error"] assert "#235 PR6" in contract_error + assert MaterializationStatus.APPLIED not in {app.status for app in materialization.applications} def test_decode_contract_is_refused_without_kv_cache_identity(): @@ -613,6 +921,7 @@ def test_vllm_provenance_reads_page_size_and_split_kv_policy(): assert adapter.kv_page_size == 16 assert adapter.split_kv_policy == 32 + assert adapter.to_dict()["flash_attn_max_num_splits_for_cuda_graph"] == 32 def test_vllm_provenance_flags_fp8_kv_cache_and_cascade_attention(): From 59d68ebc9b83c22560e26e3430e007caab4ab742 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 22:36:52 +0800 Subject: [PATCH 04/17] fix(attention): type split-k runtime boundaries --- rl_engine/kernels/ops/pytorch/attention/cp_attention.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index e81ab8b7..0062dac4 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -973,6 +973,7 @@ def split_kv_execution_plan_provenance( for owner_cp_rank, (rank_start, rank_end) in enumerate(_split_bounds(length, cp_world_size)): if rank_start == rank_end: continue + boundaries: tuple[tuple[int, int], ...] if kv_chunk_size is None: boundaries = ((rank_start, rank_end),) mode = SplitKVMode.DISABLED @@ -1019,6 +1020,7 @@ def build_reference_split_kv_runtime_plan_set( for tp_rank in range(tp_world_size): for cp_rank in range(cp_world_size): for owner_cp_rank, (owner_start, owner_end) in enumerate(owner_ranges): + boundaries: tuple[tuple[int, int], ...] if kv_chunk_size is None: mode = SplitKVMode.DISABLED boundaries = ((owner_start, owner_end),) From 1bcb885914477f8425958d8038e2c476722599ba Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 22:41:18 +0800 Subject: [PATCH 05/17] fix(attention): fail closed without runtime contract --- .../cross_config/adapters/megatron.py | 12 +++++++++++ .../alignment/cross_config/adapters/vllm.py | 12 +++++++++++ tests/test_attention_cross_config_binding.py | 20 +++++++++++++++++++ 3 files changed, 44 insertions(+) diff --git a/rl_engine/alignment/cross_config/adapters/megatron.py b/rl_engine/alignment/cross_config/adapters/megatron.py index 8c9c9412..55903cc8 100644 --- a/rl_engine/alignment/cross_config/adapters/megatron.py +++ b/rl_engine/alignment/cross_config/adapters/megatron.py @@ -330,6 +330,18 @@ def materialize( ) ) continue + if contract is None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + f"attention contract is unavailable: {blocked}", + ) + ) + continue applications.append( self._runtime_application( descriptor, diff --git a/rl_engine/alignment/cross_config/adapters/vllm.py b/rl_engine/alignment/cross_config/adapters/vllm.py index 1e1ae185..f89ab232 100644 --- a/rl_engine/alignment/cross_config/adapters/vllm.py +++ b/rl_engine/alignment/cross_config/adapters/vllm.py @@ -374,6 +374,18 @@ def materialize( ) ) continue + if contract is None: + applications.append( + application( + descriptor, + requested, + None, + None, + MaterializationStatus.ERROR, + f"attention contract is unavailable: {blocked}", + ) + ) + continue if path == "rollout.context_parallel_size" and effective_cp != requested_cp: applications.append( application( diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 6a99e6bb..04f0bfdb 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -800,10 +800,30 @@ def test_arrival_merge_order_is_unsupported_not_silently_corrected(): assert _statuses(materialization, "attention.reduction_order") == [ MaterializationStatus.UNSUPPORTED ] + assert _statuses(materialization, "training.tensor_parallel_size") == [ + MaterializationStatus.ERROR + ] assert materialization.binding.side_configs["training"]["contract"] is None assert "arrival" in materialization.binding.side_configs["training"]["contract_error"] +def test_unsupported_reduction_invalidates_vllm_contract_applications(): + normalized = { + "batch": {"size": 2}, + "rollout": {"tensor_parallel_size": 2, "context_parallel_size": 2}, + "attention": {"reduction_order": "arrival"}, + } + materialization = VllmRolloutMaterializer().materialize(normalized, WS2_ATTENTION_KNOBS) + + assert _statuses(materialization, "attention.reduction_order") == [ + MaterializationStatus.UNSUPPORTED + ] + assert _statuses(materialization, "rollout.tensor_parallel_size") == [ + MaterializationStatus.ERROR + ] + assert materialization.binding.side_configs["rollout"]["contract"] is None + + def test_bf16_reduction_accumulation_is_unsupported(): normalized = { "batch": {"size": 2}, From 48a4130662357e8a38708f7ad42d1e3294812c06 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Thu, 13 Aug 2026 23:12:21 +0800 Subject: [PATCH 06/17] fix(attention): export only available CUDA operators --- .../kernels/ops/cuda/attention/__init__.py | 38 ------------------- 1 file changed, 38 deletions(-) diff --git a/rl_engine/kernels/ops/cuda/attention/__init__.py b/rl_engine/kernels/ops/cuda/attention/__init__.py index 2c3a1f1b..09775c8e 100644 --- a/rl_engine/kernels/ops/cuda/attention/__init__.py +++ b/rl_engine/kernels/ops/cuda/attention/__init__.py @@ -1,49 +1,11 @@ # File: rl_engine/kernels/ops/cuda/attention/__init__.py -from .cp_comm import ( - AttentionCPBlockMetadata, - AttentionCPCommunication, - AttentionCPCommunicationPlan, - AttentionCPCommunicationUnavailable, - AttentionCPMergedState, - AttentionCPPartialState, - AttentionParallelSpec, - CPCommunicationBackend, - CPCommunicationStatus, - CUDAAGRSAttentionCPCommunication, - P2PNCCLAttentionCPCommunication, - sort_attention_cp_partial_states, -) from .deterministic_attn import DeterministicAttentionOp from .flash_attn import FlashAttentionOp -from .flashinfer_paged_attention import ( - FlashInferPagedAttentionConfig, - FlashInferQwen3PagedAttentionOp, - FlashInferRoPEFusionConfig, - FlashInferSplitKVPolicy, - FlashInferUnavailable, -) from .prefix_shared_attn import PrefixSharedAttentionOp __all__ = [ - "AttentionCPBlockMetadata", - "AttentionCPCommunication", - "AttentionCPCommunicationPlan", - "AttentionCPCommunicationUnavailable", - "AttentionCPMergedState", - "AttentionCPPartialState", - "AttentionParallelSpec", - "CPCommunicationBackend", - "CPCommunicationStatus", - "CUDAAGRSAttentionCPCommunication", - "P2PNCCLAttentionCPCommunication", "DeterministicAttentionOp", "FlashAttentionOp", - "FlashInferPagedAttentionConfig", - "FlashInferQwen3PagedAttentionOp", - "FlashInferRoPEFusionConfig", - "FlashInferSplitKVPolicy", - "FlashInferUnavailable", "PrefixSharedAttentionOp", - "sort_attention_cp_partial_states", ] From 4de96a4e2cc822449bbefc40e384419aa2a93515 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 02:11:46 +0800 Subject: [PATCH 07/17] feat(attention): bind strict H100 QK norm and RoPE path --- .../ws2-attention-cross-config-integration.md | 46 ++++- .../cross_config/attention_binding.py | 105 +++++++++- rl_engine/kernels/attention_preprocess.py | 180 ++++++++++++++++++ rl_engine/kernels/ops/cuda/norm/rmsnorm.py | 15 ++ .../kernels/ops/cuda/rotary_embedding/rope.py | 39 +++- rl_engine/kernels/registry.py | 6 +- tests/test_attention_cross_config_binding.py | 59 ++++++ tests/test_attention_preprocess.py | 111 +++++++++++ tests/test_rms_norm.py | 15 +- 9 files changed, 558 insertions(+), 18 deletions(-) create mode 100644 rl_engine/kernels/attention_preprocess.py create mode 100644 tests/test_attention_preprocess.py diff --git a/docs/design/ws2-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md index c4d8d773..4bda3089 100644 --- a/docs/design/ws2-attention-cross-config-integration.md +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -28,7 +28,7 @@ tiers, in `rl_engine/alignment/cross_config/attention_binding.py`: | tier | fields | rule | failure | | --- | --- | --- | --- | | `IDENTICAL` | checkpoint/model/token identity plus complete TP/CP GQA head and sequence ownership | equal bit for bit | `comparable=False`; no drift number from the pair means anything | -| `SEMANTIC` | reduction contract, dtype, exported LSE, first-class Split-KV request, complete actual Split-KV plan sets, cross-side determinism | both sides equal **and** equal to the WS2 mandate | fail closed | +| `SEMANTIC` | reduction contract, dtype, exported LSE, first-class Split-KV request, complete actual Split-KV plan sets, CUDA QK-Norm/RoPE identity, cross-side determinism | both sides equal **and** equal to the WS2 mandate | fail closed | | `RECORDED` | `mode`, `backend_id`, `reduction.engine`, RoPE materialization state and fusion boundary, KV-cache paging | free to differ | recorded into provenance and measured | Two placements are load-bearing: @@ -49,6 +49,42 @@ Two placements are load-bearing: comparable. A pair that is comparable but violates the reduction mandate is still rejected -- the drift would be real but attributable to the wrong thing. +## H100 Attention input boundary + +The strict experiment does not use the PyTorch reference operators as an +executable option. `H100AttentionPreprocessor` applies these implementations in +the fixed Qwen3 order: + +1. `RMSNormCudaOp` on Q and K (`rlkernel.cuda.rmsnorm`) +2. `RoPESM90Op` with global `[S]` or per-batch `[B, S]` positions + (`rlkernel.cuda.rope_sm90`) + +There is no Megatron/vLLM-native fallback. The CUDA RoPE path accepts non-contiguous +global positions, including zigzag CP ownership. The launcher passes the returned +backend evidence into `AttentionRuntimeReadback`: + +```python +from rl_engine.kernels.attention_preprocess import H100AttentionPreprocessor + +prepared = H100AttentionPreprocessor(device)( + q, k, q_norm_weight, k_norm_weight, position_ids +) +readback = AttentionRuntimeReadback( + # contract, knobs, Split-KV plan set, source, and scope fields omitted here + **prepared.readback_fields(), +) +``` + +Strict binding rejects a missing backend ID, a runtime-native backend ID, or any +reported fallback. Printing a configured backend without executing it is not +accepted as evidence. + +The boundary starts at projected Q/K/V. Pre-attention model RMSNorm, QKV projection +GEMM, and projection-owned TP/SP All-Gather/Reduce-Scatter are not part of the +Attention operator experiment. The isolated H100 test must therefore capture or +reuse identical projected Q/K/V inputs. Those upstream operators must be aligned +separately before making an end-to-end Megatron-vLLM logprob claim. + ## Determinism is not one thing `rl_engine/alignment/cross_config/determinism.py` probes both sides and compares @@ -83,7 +119,8 @@ adapters: limits) and `VllmRolloutMaterializer`. * `AttentionRuntimeReadback` -- the explicit handoff from an executed engine. It carries the reconstructed actual contract, actual knob values, frozen-scope - verification, and the complete Split-KV runtime plan set. + verification, executed CUDA QK-Norm/RoPE identities and fallback state, and the + complete Split-KV runtime plan set. Constructing a contract is not runtime verification. Without a readback, adapter applications are `UNOBSERVABLE`; only matching values reconstructed from a real @@ -107,6 +144,7 @@ collapsing them onto the supported value: | configured contract without runtime readback | `UNOBSERVABLE` | requested values do not prove what executed | | `rollout.context_parallel_size>1` with effective decode CP=1 | `ERROR`/`FALLBACK` | strict TP=2/CP=2 acceptance rejects the topology change | | missing/mismatched/fallback Split-KV plan set | binding failure | Split-KV provenance must cover every batch/TP/CP/owner coordinate | +| missing/native/fallback QK-Norm or RoPE backend | binding failure | both sides must execute the RL-Kernel CUDA preprocessing path | ## Knobs @@ -139,7 +177,9 @@ retired rather than rewritten. Deliberately not in this PR: -* launching `torchrun`, initializing process groups, or executing attention; +* launching `torchrun`, initializing process groups, or executing core attention; +* pre-attention model RMSNorm, QKV projection GEMM, and projection-owned TP/SP + communication; isolated tests reuse identical projected Q/K/V; * decode-mode materialization, which needs the validated `KVCacheSpec` from #235 PR6 and is refused with that reference rather than stubbed; * Transformer Engine calls of any kind (PR4's TE plan is policy and provenance only); diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 4156d57e..1007f197 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -54,6 +54,7 @@ SplitKVRuntimePlanSet, validate_split_kv_plan_set_alignment, ) +from rl_engine.kernels.attention_preprocess import MANDATED_ATTENTION_PREPROCESS_BACKENDS __all__ = [ "ATTENTION_LSE_DOMAIN", @@ -108,6 +109,9 @@ class BindingErrorCode(str, Enum): SPLIT_KV_RUNTIME_MISSING = "SPLIT_KV_RUNTIME_MISSING" SPLIT_KV_MISMATCH = "SPLIT_KV_MISMATCH" SPLIT_KV_FALLBACK = "SPLIT_KV_FALLBACK" + ATTENTION_PREPROCESS_MISSING = "ATTENTION_PREPROCESS_MISSING" + ATTENTION_PREPROCESS_MISMATCH = "ATTENTION_PREPROCESS_MISMATCH" + ATTENTION_PREPROCESS_FALLBACK = "ATTENTION_PREPROCESS_FALLBACK" @dataclass(frozen=True) @@ -119,6 +123,8 @@ class AttentionRuntimeReadback: split_kv_plan_set: SplitKVRuntimePlanSet source: str frozen_scope_verified: bool + preprocess_backends: Mapping[str, str] = field(default_factory=dict) + preprocess_fallback: bool = False def __post_init__(self) -> None: if not isinstance(self.contract, AttentionContract): @@ -131,11 +137,25 @@ def __post_init__(self) -> None: raise ValueError("runtime readback source must be a non-empty string") if not isinstance(self.frozen_scope_verified, bool): raise TypeError("frozen_scope_verified must be a bool") + if not isinstance(self.preprocess_backends, Mapping): + raise TypeError("runtime readback preprocess_backends must be a mapping") + for name, backend in self.preprocess_backends.items(): + if not isinstance(name, str) or not name.strip(): + raise ValueError("preprocess backend names must be non-empty strings") + if not isinstance(backend, str) or not backend.strip(): + raise ValueError("preprocess backend IDs must be non-empty strings") + if not isinstance(self.preprocess_fallback, bool): + raise TypeError("preprocess_fallback must be a bool") plan_error = _split_kv_plan_contract_error(self.contract, self.split_kv_plan_set) if plan_error is not None: raise ValueError(plan_error) object.__setattr__(self, "actual_knobs", MappingProxyType(dict(self.actual_knobs))) + object.__setattr__( + self, + "preprocess_backends", + MappingProxyType(dict(self.preprocess_backends)), + ) @property def split_kv_fallback(self) -> bool: @@ -147,6 +167,10 @@ def to_dict(self) -> dict[str, Any]: "frozen_scope_verified": self.frozen_scope_verified, "contract": self.contract.to_dict(), "actual_knobs": dict(self.actual_knobs), + "attention_preprocess": { + "backends": dict(self.preprocess_backends), + "fallback": self.preprocess_fallback, + }, "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), } @@ -253,6 +277,9 @@ def to_dict(self) -> dict[str, Any]: "rope.k_cache_state", "rope.cast_at", "rope.output_dtype", + "preprocess.qk_rmsnorm", + "preprocess.rope", + "preprocess.fallback", "kv_cache.page_size", "kv_cache.prefix_cache_enabled", "kv_cache.block_table_shape", @@ -300,7 +327,7 @@ class AttentionBindingResult: binding_fingerprint: str = "" recorded_differences: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) provenance: Mapping[str, Any] = field(default_factory=dict) - schema_version: str = "cross_config.attention_binding.v2" + schema_version: str = "cross_config.attention_binding.v3" def issues_by_code(self, code: BindingErrorCode) -> tuple[BindingIssue, ...]: return tuple(issue for issue in self.issues if issue.code is code) @@ -774,6 +801,22 @@ def bind_attention_contracts( "lse_domain": ATTENTION_LSE_DOMAIN, "rollout_backend": rollout_backend_id, "training_backend": training_backend_id, + "attention_preprocess": { + "rollout": { + name: rollout_recorded.get(f"preprocess.{name}") + for name in ( + *MANDATED_ATTENTION_PREPROCESS_BACKENDS, + "fallback", + ) + }, + "training": { + name: training_recorded.get(f"preprocess.{name}") + for name in ( + *MANDATED_ATTENTION_PREPROCESS_BACKENDS, + "fallback", + ) + }, + }, } ), recorded_differences=recorded_differences, @@ -809,6 +852,7 @@ def bind_attention_runtime_readbacks( message=f"{side} runtime did not verify the frozen attention scope", ) ) + missing_scope_evidence.extend(_attention_preprocess_issues(side, readback)) return bind_attention_contracts( rollout_contract=rollout.contract, training_contract=training.contract, @@ -819,9 +863,68 @@ def bind_attention_runtime_readbacks( determinism_issues=tuple(determinism_issues) + tuple(missing_scope_evidence), rollout_split_kv_plan_set=rollout.split_kv_plan_set, training_split_kv_plan_set=training.split_kv_plan_set, + rollout_recorded_extra={ + **{ + f"preprocess.{name}": backend + for name, backend in rollout.preprocess_backends.items() + }, + "preprocess.fallback": rollout.preprocess_fallback, + }, + training_recorded_extra={ + **{ + f"preprocess.{name}": backend + for name, backend in training.preprocess_backends.items() + }, + "preprocess.fallback": training.preprocess_fallback, + }, ) +def _attention_preprocess_issues( + side: str, + readback: AttentionRuntimeReadback, +) -> list[BindingIssue]: + issues: list[BindingIssue] = [] + for name, mandated in MANDATED_ATTENTION_PREPROCESS_BACKENDS.items(): + actual = readback.preprocess_backends.get(name) + if actual is None: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.preprocess.{name}", + message=( + f"{side} did not report the executed {name} backend; " + "runtime-native execution cannot validate the Attention input boundary" + ), + ) + ) + elif actual != mandated: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.preprocess.{name}", + rollout=actual if side == "rollout" else None, + training=actual if side == "training" else None, + message=( + f"{side} executed {actual!r}; " + f"the H100 experiment requires {mandated!r}" + ), + ) + ) + if readback.preprocess_fallback: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_FALLBACK, + tier=BindingTier.SEMANTIC, + field=f"{side}.preprocess.fallback", + message=f"{side} reported a QK-Norm or RoPE backend fallback", + ) + ) + return issues + + def summarize_binding(result: AttentionBindingResult) -> str: """One-line human summary for CLI output and failure messages.""" diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py new file mode 100644 index 00000000..196d2e18 --- /dev/null +++ b/rl_engine/kernels/attention_preprocess.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Strict H100 QK-Norm and RoPE handoff for WS2 Attention. + +This module intentionally has no runtime-native fallback. A caller either runs +the RL-Kernel CUDA operators and records their identities, or the experiment +fails before Attention executes. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Mapping + +import torch +from torch import Tensor + + +QK_RMSNORM_BACKEND_ID = "rlkernel.cuda.rmsnorm" +ROPE_BACKEND_ID = "rlkernel.cuda.rope_sm90" +MANDATED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( + { + "qk_rmsnorm": QK_RMSNORM_BACKEND_ID, + "rope": ROPE_BACKEND_ID, + } +) + + +@dataclass(frozen=True) +class AttentionPreprocessResult: + """Post-QK-Norm, post-RoPE tensors plus executed backend evidence.""" + + q: Tensor + k: Tensor + backend_ids: Mapping[str, str] + fallback: bool + device_capability: tuple[int, int] + + def __post_init__(self) -> None: + object.__setattr__(self, "backend_ids", MappingProxyType(dict(self.backend_ids))) + + def evidence(self) -> dict[str, Any]: + return { + "backends": dict(self.backend_ids), + "fallback": self.fallback, + "device_capability": list(self.device_capability), + } + + def readback_fields(self) -> dict[str, Any]: + """Keyword fields consumed by ``AttentionRuntimeReadback``.""" + + return { + "preprocess_backends": dict(self.backend_ids), + "preprocess_fallback": self.fallback, + } + + +class H100AttentionPreprocessor: + """Apply RL-Kernel CUDA QK-Norm then RoPE without silent fallback.""" + + def __init__(self, device: torch.device | str | int | None = None) -> None: + if not torch.cuda.is_available(): + raise RuntimeError("H100AttentionPreprocessor requires an available CUDA runtime") + + current_device = torch.cuda.current_device() + self.device = torch.device("cuda", current_device) + if device is not None: + self.device = ( + torch.device("cuda", device) if isinstance(device, int) else torch.device(device) + ) + if self.device.type != "cuda": + raise RuntimeError(f"H100AttentionPreprocessor requires CUDA, got {self.device}") + if self.device.index is None: + self.device = torch.device("cuda", current_device) + + capability = torch.cuda.get_device_capability(self.device) + self.device_capability: tuple[int, int] = (int(capability[0]), int(capability[1])) + if self.device_capability[0] != 9: + raise RuntimeError( + "H100AttentionPreprocessor requires Hopper SM90; " + f"got sm_{self.device_capability[0]}{self.device_capability[1]}" + ) + + # Import only after the hardware gate so CPU tools can inspect the module. + from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op + + self.rmsnorm = RMSNormCudaOp() + self.rope = RoPESM90Op() + actual_backends = { + "qk_rmsnorm": self.rmsnorm.backend_id, + "rope": self.rope.backend_id, + } + if actual_backends != dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS): + raise RuntimeError(f"unexpected Attention preprocess backends: {actual_backends}") + self.backend_ids = MappingProxyType(actual_backends) + + def __call__( + self, + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + *, + eps: float = 1.0e-6, + theta: float = 1_000_000.0, + ) -> AttentionPreprocessResult: + return self.forward( + q, + k, + q_weight, + k_weight, + positions, + eps=eps, + theta=theta, + ) + + def forward( + self, + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + *, + eps: float = 1.0e-6, + theta: float = 1_000_000.0, + ) -> AttentionPreprocessResult: + _validate_inputs(q, k, q_weight, k_weight, positions, self.device) + q_norm = self.rmsnorm(q, q_weight, eps=eps) + k_norm = self.rmsnorm(k, k_weight, eps=eps) + return AttentionPreprocessResult( + q=self.rope(q_norm, positions, theta=theta), + k=self.rope(k_norm, positions, theta=theta), + backend_ids=self.backend_ids, + fallback=False, + device_capability=self.device_capability, + ) + + +def _validate_inputs( + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + device: torch.device, +) -> None: + if q.dim() != 4 or k.dim() != 4: + raise ValueError("q and k must use [B, H, S, D] layout") + if q.shape[0] != k.shape[0] or q.shape[-2:] != k.shape[-2:]: + raise ValueError("q and k must have the same batch, sequence, and head dimensions") + if q.dtype is not torch.bfloat16 or k.dtype is not torch.bfloat16: + raise TypeError("the frozen H100 Attention experiment requires BF16 q and k") + if q.device != device or k.device != device: + raise ValueError(f"q and k must both be on the configured device {device}") + for name, weight in (("q_weight", q_weight), ("k_weight", k_weight)): + if weight.shape != (q.shape[-1],): + raise ValueError(f"{name} must have shape ({q.shape[-1]},)") + if weight.device != device or weight.dtype is not torch.bfloat16: + raise ValueError(f"{name} must be BF16 on {device}") + if positions.device != device: + raise ValueError(f"positions must be on {device}") + if positions.dtype not in (torch.int32, torch.int64): + raise TypeError("positions must use int32 or int64 global token indices") + expected = (q.shape[-2],) if positions.dim() == 1 else (q.shape[0], q.shape[-2]) + if positions.dim() not in (1, 2) or tuple(positions.shape) != expected: + raise ValueError(f"positions must have shape [S] or [B, S], expected {expected}") + + +__all__ = [ + "AttentionPreprocessResult", + "H100AttentionPreprocessor", + "MANDATED_ATTENTION_PREPROCESS_BACKENDS", + "QK_RMSNORM_BACKEND_ID", + "ROPE_BACKEND_ID", +] diff --git a/rl_engine/kernels/ops/cuda/norm/rmsnorm.py b/rl_engine/kernels/ops/cuda/norm/rmsnorm.py index 76e33da8..7ac85b70 100644 --- a/rl_engine/kernels/ops/cuda/norm/rmsnorm.py +++ b/rl_engine/kernels/ops/cuda/norm/rmsnorm.py @@ -79,6 +79,21 @@ def rmsnorm_cuda(x, weight, eps=1e-6, mask=None): class RMSNormCudaOp: """CUDA RMSNorm wrapper compatible with the shared operator harness.""" + backend_id = "rlkernel.cuda.rmsnorm" + + def __init__(self): + required = ( + "rmsnorm_forward", + "rmsnorm_backward_dx", + "rmsnorm_backward_dw", + ) + missing = [name for name in required if not _EXT_AVAILABLE or not hasattr(_C, name)] + if missing: + raise RuntimeError( + "CUDA RMSNorm extension is incomplete; rebuild _C with rmsnorm.cu " + f"(missing: {', '.join(missing)})" + ) + def __call__(self, x, weight, *, eps=1e-6): return self.forward(x, weight, eps=eps) diff --git a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py index 0c1a7b73..028e44b9 100644 --- a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py +++ b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py @@ -18,7 +18,7 @@ def _build_cos_sin(positions: Tensor, half: int, theta: float, device: torch.device): - """fp32 cos/sin caches of shape [S, half], identical math to NativeRoPEOp.""" + """fp32 cos/sin rows, identical math to NativeRoPEOp.""" inv_freq = 1.0 / (theta ** (torch.arange(0, half, dtype=torch.float32, device=device) / half)) pos = positions.to(device=device, dtype=torch.float32).reshape(-1, 1) freqs = pos * inv_freq # [S, half] @@ -31,19 +31,39 @@ def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: D = x.shape[-1] if D % 2 != 0: raise ValueError(f"RoPE head_dim must be even, got {D}") - if positions.dim() != 1: - raise NotImplementedError( - "CUDA RoPE currently supports 1-D positions [S] (shared across batch)." - ) - S = positions.shape[0] + if positions.dim() not in (1, 2): + raise ValueError("positions must have shape [S] or [B, S]") + S = positions.shape[-1] + if S == 0: + raise ValueError("positions must not be empty") + if x.shape[-2] != S: + raise ValueError(f"x sequence length {x.shape[-2]} does not match positions length {S}") x_2d = x.contiguous().reshape(-1, D) n_rows = x_2d.shape[0] - if n_rows % S != 0: + if positions.dim() == 2: + batch = positions.shape[0] + if x.dim() < 3 or x.shape[0] != batch: + raise ValueError( + f"x batch size {x.shape[0]} does not match positions batch size {batch}" + ) + rows_per_token = n_rows // (batch * S) + if rows_per_token * batch * S != n_rows: + raise ValueError("x rows are incompatible with [B, S] positions") + # The CUDA kernel accepts one fp32 cos/sin row per flattened x row. + # Expanding positions preserves arbitrary global/zigzag indices while + # keeping the arithmetic inside the precompiled deterministic kernel. + kernel_positions = ( + positions[:, None, :].expand(batch, rows_per_token, S).contiguous().reshape(-1) + ) + else: + kernel_positions = positions + if n_rows % kernel_positions.numel() != 0: raise ValueError( - f"row count {n_rows} not divisible by seq length {S}; " + f"row count {n_rows} not divisible by position rows " + f"{kernel_positions.numel()}; " "expected a [..., S, D] contiguous layout." ) - cos, sin = _build_cos_sin(positions, D // 2, float(theta), x.device) + cos, sin = _build_cos_sin(kernel_positions, D // 2, float(theta), x.device) ctx.save_for_backward(cos, sin) out = _C.rope_apply_sm90(x_2d, cos, sin, 1.0) return out.reshape(x.shape) @@ -70,6 +90,7 @@ class RoPESM90Op: """ op_class = "elementwise" + backend_id = "rlkernel.cuda.rope_sm90" def __init__(self) -> None: if not _EXT_AVAILABLE or not hasattr(_C, "rope_apply_sm90"): diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 3c835deb..060234be 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -95,6 +95,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) + CUDA_RMS_NORM = "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp" PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" # Generic fallback @@ -348,7 +349,10 @@ def __init__(self): OpBackend.TRITON_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, ], - "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], + "rms_norm": [ + OpBackend.CUDA_RMS_NORM, + OpBackend.PYTORCH_NATIVE_RMS_NORM, + ], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [ diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 04f0bfdb..12bb878e 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -55,6 +55,9 @@ SplitKVSpec, build_split_kv_runtime_plan_set, ) +from rl_engine.kernels.attention_preprocess import ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS, +) pytestmark = pytest.mark.unit @@ -626,6 +629,8 @@ def _readback(materializer, flat, *, source): split_kv_plan_set=_plan_set(contract, backend=source), source=source, frozen_scope_verified=True, + preprocess_backends=MANDATED_ATTENTION_PREPROCESS_BACKENDS, + preprocess_fallback=False, ) @@ -683,6 +688,8 @@ def test_runtime_readback_mismatch_is_a_fallback(): split_kv_plan_set=readback.split_kv_plan_set, source=readback.source, frozen_scope_verified=True, + preprocess_backends=readback.preprocess_backends, + preprocess_fallback=readback.preprocess_fallback, ) normalized = { "batch": {"size": 2}, @@ -760,6 +767,58 @@ def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): assert result.provenance["split_kv_runtime"]["rollout"]["coverage"] == ( "complete_batch_tp_cp_owner_cartesian_product" ) + assert result.provenance["rollout"]["recorded"]["preprocess.qk_rmsnorm"] == ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS["qk_rmsnorm"] + ) + + +@pytest.mark.parametrize( + ("backends", "fallback", "expected_code"), + [ + ( + {"rope": MANDATED_ATTENTION_PREPROCESS_BACKENDS["rope"]}, + False, + BindingErrorCode.ATTENTION_PREPROCESS_MISSING, + ), + ( + {**MANDATED_ATTENTION_PREPROCESS_BACKENDS, "qk_rmsnorm": "vllm.native"}, + False, + BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + ), + ( + dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), + True, + BindingErrorCode.ATTENTION_PREPROCESS_FALLBACK, + ), + ], +) +def test_strict_runtime_readback_rejects_unverified_preprocess_backend( + backends, fallback, expected_code +): + rollout = _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback") + rollout = replace( + rollout, + preprocess_backends=backends, + preprocess_fallback=fallback, + ) + training = _readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.comparable + assert not result.passed + assert result.issues_by_code(expected_code) def test_strict_runtime_readback_entrypoint_rejects_unverified_frozen_scope(): diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py new file mode 100644 index 00000000..451ec979 --- /dev/null +++ b/tests/test_attention_preprocess.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.attention_preprocess import ( + H100AttentionPreprocessor, + MANDATED_ATTENTION_PREPROCESS_BACKENDS, +) +from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp + + +def _has_h100_preprocess() -> bool: + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + return False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + required = ( + "rmsnorm_forward", + "rmsnorm_backward_dx", + "rmsnorm_backward_dw", + "rope_apply_sm90", + ) + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in required)) + except ImportError: + return False + + +requires_h100_preprocess = pytest.mark.skipif( + not _has_h100_preprocess(), + reason="Hopper with compiled RMSNorm and RoPE CUDA kernels is required", +) + + +def test_h100_preprocessor_has_no_native_backend_option(): + assert dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) == { + "qk_rmsnorm": "rlkernel.cuda.rmsnorm", + "rope": "rlkernel.cuda.rope_sm90", + } + + +def test_h100_preprocessor_fails_before_dispatch_without_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + with pytest.raises(RuntimeError, match="requires an available CUDA runtime"): + H100AttentionPreprocessor() + + +def test_h100_preprocessor_rejects_non_hopper_cuda(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (8, 0)) + with pytest.raises(RuntimeError, match="requires Hopper SM90"): + H100AttentionPreprocessor() + + +def _inputs(): + torch.manual_seed(7) + device = torch.device("cuda") + q = torch.randn(2, 4, 8, 128, device=device, dtype=torch.bfloat16) + k = torch.randn(2, 2, 8, 128, device=device, dtype=torch.bfloat16) + q_weight = torch.randn(128, device=device, dtype=torch.bfloat16) + k_weight = torch.randn(128, device=device, dtype=torch.bfloat16) + positions = torch.tensor( + [[0, 7, 2, 9, 4, 11, 6, 13], [100, 107, 102, 109, 104, 111, 106, 113]], + device=device, + dtype=torch.int64, + ) + return q, k, q_weight, k_weight, positions + + +@requires_h100_preprocess +def test_h100_preprocessor_executes_cuda_qk_norm_and_zigzag_rope(): + q, k, q_weight, k_weight, positions = _inputs() + result = H100AttentionPreprocessor()(q, k, q_weight, k_weight, positions) + + norm = NativeRMSNormOp() + rope = NativeRoPEOp() + q_ref = rope(norm(q, q_weight), positions) + k_ref = rope(norm(k, k_weight), positions) + + assert result.fallback is False + assert dict(result.backend_ids) == dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) + assert result.readback_fields() == { + "preprocess_backends": dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), + "preprocess_fallback": False, + } + torch.testing.assert_close(result.q.float(), q_ref.float(), atol=2e-2, rtol=2e-2) + torch.testing.assert_close(result.k.float(), k_ref.float(), atol=2e-2, rtol=2e-2) + + +@requires_h100_preprocess +def test_h100_preprocessor_is_bitwise_batch_invariant_for_2d_positions(): + q, k, q_weight, k_weight, positions = _inputs() + op = H100AttentionPreprocessor() + full = op(q, k, q_weight, k_weight, positions) + + for batch_index in range(q.shape[0]): + single = op( + q[batch_index : batch_index + 1], + k[batch_index : batch_index + 1], + q_weight, + k_weight, + positions[batch_index : batch_index + 1], + ) + assert torch.equal(full.q[batch_index], single.q[0]) + assert torch.equal(full.k[batch_index], single.k[0]) diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index 43a9cf85..6572603e 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -5,14 +5,17 @@ import torch import torch.nn.functional as F -from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda +from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp, rmsnorm_cuda from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton try: from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE - _HAS_CUDA_RMSNORM = _EXT_AVAILABLE and hasattr(_C, "rmsnorm_forward") + _HAS_CUDA_RMSNORM = _EXT_AVAILABLE and all( + hasattr(_C, name) + for name in ("rmsnorm_forward", "rmsnorm_backward_dx", "rmsnorm_backward_dw") + ) except ImportError: # pragma: no cover - import can fail when the extension is not built. _HAS_CUDA_RMSNORM = False @@ -236,8 +239,12 @@ def test_registry_dispatches_rms_norm(): from rl_engine.kernels.registry import kernel_registry op = kernel_registry.get_op("rms_norm") - assert isinstance(op, NativeRMSNormOp) - assert hasattr(op, "forward") and hasattr(op, "forward_fp32") + if torch.cuda.is_available() and _HAS_CUDA_RMSNORM: + assert isinstance(op, RMSNormCudaOp) + assert hasattr(op, "forward") + else: + assert isinstance(op, NativeRMSNormOp) + assert hasattr(op, "forward") and hasattr(op, "forward_fp32") @requires_cuda From d9e04b90fb7c7654fbebdbe185092767f9fb3d7c Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 17:46:39 +0800 Subject: [PATCH 08/17] fix(attention): align projection collectives and runtime evidence --- .../ws2-attention-cross-config-integration.md | 46 ++-- docs/operators/attention.md | 6 + .../cross_config/attention_binding.py | 233 +++++++++++++++- rl_engine/kernels/attention_preprocess.py | 132 +++++++-- rl_engine/kernels/attention_projection.py | 256 ++++++++++++++++++ tests/test_attention_cross_config_binding.py | 103 ++++++- tests/test_attention_preprocess.py | 11 +- tests/test_attention_projection.py | 97 +++++++ 8 files changed, 830 insertions(+), 54 deletions(-) create mode 100644 rl_engine/kernels/attention_projection.py create mode 100644 tests/test_attention_projection.py diff --git a/docs/design/ws2-attention-cross-config-integration.md b/docs/design/ws2-attention-cross-config-integration.md index 4bda3089..6ddd45a0 100644 --- a/docs/design/ws2-attention-cross-config-integration.md +++ b/docs/design/ws2-attention-cross-config-integration.md @@ -49,18 +49,13 @@ Two placements are load-bearing: comparable. A pair that is comparable but violates the reduction mandate is still rejected -- the drift would be real but attributable to the wrong thing. -## H100 Attention input boundary +## H100 Attention input and projection boundary -The strict experiment does not use the PyTorch reference operators as an -executable option. `H100AttentionPreprocessor` applies these implementations in -the fixed Qwen3 order: - -1. `RMSNormCudaOp` on Q and K (`rlkernel.cuda.rmsnorm`) -2. `RoPESM90Op` with global `[S]` or per-batch `[B, S]` positions - (`rlkernel.cuda.rope_sm90`) - -There is no Megatron/vLLM-native fallback. The CUDA RoPE path accepts non-contiguous -global positions, including zigzag CP ownership. The launcher passes the returned +Megatron/TE and vLLM/FlashInfer are the first-choice implementations. +`H100AttentionPreprocessor` runs a same-input H100 bitwise probe against the +deterministic RL-Kernel path. An unavailable native callable, a native exception, +or a failed probe switches both sides to `RMSNormCudaOp` + `RoPESM90Op` and +records the fallback reason and probe ID. The launcher passes the returned backend evidence into `AttentionRuntimeReadback`: ```python @@ -75,15 +70,20 @@ readback = AttentionRuntimeReadback( ) ``` -Strict binding rejects a missing backend ID, a runtime-native backend ID, or any -reported fallback. Printing a configured backend without executing it is not -accepted as evidence. - -The boundary starts at projected Q/K/V. Pre-attention model RMSNorm, QKV projection -GEMM, and projection-owned TP/SP All-Gather/Reduce-Scatter are not part of the -Attention operator experiment. The isolated H100 test must therefore capture or -reuse identical projected Q/K/V inputs. Those upstream operators must be aligned -separately before making an end-to-end Megatron-vLLM logprob claim. +Strict binding rejects a missing or unknown backend and rejects mixed native / +fallback execution. If both sides fall back, they must report the same deterministic +backend IDs and policy ID. Printing a configured backend without executing the +probe is not evidence. + +The Attention boundary includes QKV projection, Q/K RMSNorm, RoPE, core +attention, KV-cache access, CP `(Out, LSE)` communication/merge, and o_proj. +`AttentionProjectionOp` freezes QKV/o_proj to BF16 input and output, FP32 +accumulation, ascending-K reduction, and Split-K disabled. Native projection +callables are accepted only after a bitwise probe against `DetGemmOp`; otherwise +both sides use the deterministic fallback. Its collective contract records QKV +column-parallel plus backward TP all-reduce, o_proj row-parallel partial output, +and the SP all-gather/reduce-scatter directions. The model input RMSNorm and +residual add remain outside this Attention experiment. ## Determinism is not one thing @@ -144,7 +144,7 @@ collapsing them onto the supported value: | configured contract without runtime readback | `UNOBSERVABLE` | requested values do not prove what executed | | `rollout.context_parallel_size>1` with effective decode CP=1 | `ERROR`/`FALLBACK` | strict TP=2/CP=2 acceptance rejects the topology change | | missing/mismatched/fallback Split-KV plan set | binding failure | Split-KV provenance must cover every batch/TP/CP/owner coordinate | -| missing/native/fallback QK-Norm or RoPE backend | binding failure | both sides must execute the RL-Kernel CUDA preprocessing path | +| missing/unknown QK-Norm or RoPE backend, or mixed native/fallback sides | binding failure | both sides must execute the same verified native policy or the common RL-Kernel CUDA fallback | ## Knobs @@ -178,10 +178,8 @@ retired rather than rewritten. Deliberately not in this PR: * launching `torchrun`, initializing process groups, or executing core attention; -* pre-attention model RMSNorm, QKV projection GEMM, and projection-owned TP/SP - communication; isolated tests reuse identical projected Q/K/V; +* pre-attention model RMSNorm and residual add; * decode-mode materialization, which needs the validated `KVCacheSpec` from #235 PR6 and is refused with that reference rather than stubbed; -* Transformer Engine calls of any kind (PR4's TE plan is policy and provenance only); * distributed drift benchmarks and report artifacts (#235 PR5); * fused production backend alignment (#235 PR7) and backward (#235 PR8). diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 92d24e6c..592c6f95 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -13,6 +13,12 @@ friends) are validated against. This op covers **only** the softmax attention. Qwen3's QK-Norm and RoPE are applied *before* the call (see the chain), so the `q`, `k` passed in are already normalized and rotated. +For the WS2 Attention experiment, the measured boundary also includes QKV and `o_proj` +projections plus their TP/SP communication contracts. Those projections use native TE or +vLLM callables only after an H100 bitwise probe; otherwise both sides use the deterministic +`DetGemmOp` path with BF16 I/O, FP32 accumulation, ascending-K reduction, and Split-K disabled. +The model input RMSNorm and residual add remain outside this boundary. + ```text q --\ k ----softmax(QKᵀ/√d + mask)·V--> out diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 1007f197..9adba35d 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -54,7 +54,15 @@ SplitKVRuntimePlanSet, validate_split_kv_plan_set_alignment, ) -from rl_engine.kernels.attention_preprocess import MANDATED_ATTENTION_PREPROCESS_BACKENDS +from rl_engine.kernels.attention_preprocess import ( + MANDATED_ATTENTION_PREPROCESS_BACKENDS, + PREPROCESS_POLICY_ID, +) +from rl_engine.kernels.attention_projection import ( + O_PROJ_COLLECTIVE_CONTRACT, + PROJECTION_POLICY_ID, + QKV_COLLECTIVE_CONTRACT, +) __all__ = [ "ATTENTION_LSE_DOMAIN", @@ -112,6 +120,8 @@ class BindingErrorCode(str, Enum): ATTENTION_PREPROCESS_MISSING = "ATTENTION_PREPROCESS_MISSING" ATTENTION_PREPROCESS_MISMATCH = "ATTENTION_PREPROCESS_MISMATCH" ATTENTION_PREPROCESS_FALLBACK = "ATTENTION_PREPROCESS_FALLBACK" + ATTENTION_PROJECTION_MISSING = "ATTENTION_PROJECTION_MISSING" + ATTENTION_PROJECTION_MISMATCH = "ATTENTION_PROJECTION_MISMATCH" @dataclass(frozen=True) @@ -125,6 +135,10 @@ class AttentionRuntimeReadback: frozen_scope_verified: bool preprocess_backends: Mapping[str, str] = field(default_factory=dict) preprocess_fallback: bool = False + preprocess_fallback_reason: str | None = None + preprocess_probe_id: str = "" + preprocess_policy_id: str = PREPROCESS_POLICY_ID + projection_plans: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) def __post_init__(self) -> None: if not isinstance(self.contract, AttentionContract): @@ -146,6 +160,21 @@ def __post_init__(self) -> None: raise ValueError("preprocess backend IDs must be non-empty strings") if not isinstance(self.preprocess_fallback, bool): raise TypeError("preprocess_fallback must be a bool") + if self.preprocess_fallback and not self.preprocess_fallback_reason: + raise ValueError( + "preprocess_fallback_reason is required when preprocess_fallback is true" + ) + if not isinstance(self.preprocess_probe_id, str): + raise TypeError("preprocess_probe_id must be a string") + if not isinstance(self.preprocess_policy_id, str) or not self.preprocess_policy_id.strip(): + raise ValueError("preprocess_policy_id must be a non-empty string") + if not isinstance(self.projection_plans, Mapping): + raise TypeError("projection_plans must be a mapping") + normalized_projection_plans: dict[str, Mapping[str, Any]] = {} + for name, plan in self.projection_plans.items(): + if not isinstance(name, str) or not isinstance(plan, Mapping): + raise TypeError("projection_plans must map projection names to mappings") + normalized_projection_plans[name] = MappingProxyType(dict(plan)) plan_error = _split_kv_plan_contract_error(self.contract, self.split_kv_plan_set) if plan_error is not None: @@ -156,6 +185,11 @@ def __post_init__(self) -> None: "preprocess_backends", MappingProxyType(dict(self.preprocess_backends)), ) + object.__setattr__( + self, + "projection_plans", + MappingProxyType(normalized_projection_plans), + ) @property def split_kv_fallback(self) -> bool: @@ -170,6 +204,12 @@ def to_dict(self) -> dict[str, Any]: "attention_preprocess": { "backends": dict(self.preprocess_backends), "fallback": self.preprocess_fallback, + "fallback_reason": self.preprocess_fallback_reason, + "probe_id": self.preprocess_probe_id, + "policy_id": self.preprocess_policy_id, + }, + "attention_projections": { + name: dict(plan) for name, plan in self.projection_plans.items() }, "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), } @@ -853,6 +893,82 @@ def bind_attention_runtime_readbacks( ) ) missing_scope_evidence.extend(_attention_preprocess_issues(side, readback)) + missing_scope_evidence.extend(_attention_projection_issues(side, readback)) + if rollout.preprocess_fallback != training.preprocess_fallback: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field="preprocess.fallback", + rollout=rollout.preprocess_fallback, + training=training.preprocess_fallback, + message=( + "both runtimes must either pass the native H100 bitwise probe or " + "use the same deterministic preprocess fallback" + ), + ) + ) + if rollout.preprocess_policy_id != training.preprocess_policy_id: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field="preprocess.policy_id", + rollout=rollout.preprocess_policy_id, + training=training.preprocess_policy_id, + message="QK-Norm/RoPE policy IDs differ between runtimes", + ) + ) + if rollout.preprocess_fallback and training.preprocess_fallback: + for name in MANDATED_ATTENTION_PREPROCESS_BACKENDS: + rollout_backend = rollout.preprocess_backends.get(name) + training_backend = training.preprocess_backends.get(name) + if rollout_backend != training_backend: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"preprocess.{name}", + rollout=rollout_backend, + training=training_backend, + message="fallback sides must use the same deterministic preprocess backend", + ) + ) + for projection in ("qkv", "o_proj"): + rollout_plan = rollout.projection_plans.get(projection, {}) + training_plan = training.projection_plans.get(projection, {}) + rollout_fallback = bool(rollout_plan.get("fallback", False)) + training_fallback = bool(training_plan.get("fallback", False)) + if rollout_fallback != training_fallback: + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"projection.{projection}.fallback", + rollout=rollout_fallback, + training=training_fallback, + message=( + "QKV/o_proj must use the same native-or-deterministic " + "path on both sides" + ), + ) + ) + if rollout_fallback and training_fallback: + for field in ("backend_id", "policy_id", "split_k", "reduction_order"): + if rollout_plan.get(field) != training_plan.get(field): + missing_scope_evidence.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"projection.{projection}.{field}", + rollout=rollout_plan.get(field), + training=training_plan.get(field), + message=( + "deterministic projection fallback evidence differs " + "between sides" + ), + ) + ) return bind_attention_contracts( rollout_contract=rollout.contract, training_contract=training.contract, @@ -869,6 +985,13 @@ def bind_attention_runtime_readbacks( for name, backend in rollout.preprocess_backends.items() }, "preprocess.fallback": rollout.preprocess_fallback, + "preprocess.fallback_reason": rollout.preprocess_fallback_reason, + "preprocess.probe_id": rollout.preprocess_probe_id, + "preprocess.policy_id": rollout.preprocess_policy_id, + **{ + f"projection.{projection}": dict(plan) + for projection, plan in rollout.projection_plans.items() + }, }, training_recorded_extra={ **{ @@ -876,6 +999,13 @@ def bind_attention_runtime_readbacks( for name, backend in training.preprocess_backends.items() }, "preprocess.fallback": training.preprocess_fallback, + "preprocess.fallback_reason": training.preprocess_fallback_reason, + "preprocess.probe_id": training.preprocess_probe_id, + "preprocess.policy_id": training.preprocess_policy_id, + **{ + f"projection.{projection}": dict(plan) + for projection, plan in training.projection_plans.items() + }, }, ) @@ -899,7 +1029,7 @@ def _attention_preprocess_issues( ), ) ) - elif actual != mandated: + elif actual != mandated and not actual.startswith("native."): issues.append( BindingIssue( code=BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, @@ -909,22 +1039,115 @@ def _attention_preprocess_issues( training=actual if side == "training" else None, message=( f"{side} executed {actual!r}; " - f"the H100 experiment requires {mandated!r}" + f"the H100 experiment requires {mandated!r} or a verified native backend" ), ) ) - if readback.preprocess_fallback: + if readback.preprocess_fallback and any( + readback.preprocess_backends.get(name) != mandated + for name, mandated in MANDATED_ATTENTION_PREPROCESS_BACKENDS.items() + ): issues.append( BindingIssue( code=BindingErrorCode.ATTENTION_PREPROCESS_FALLBACK, tier=BindingTier.SEMANTIC, field=f"{side}.preprocess.fallback", - message=f"{side} reported a QK-Norm or RoPE backend fallback", + message=( + f"{side} reported a fallback but did not use the common deterministic " + "QK-Norm/RoPE backends" + ), ) ) return issues +def _attention_projection_issues( + side: str, + readback: AttentionRuntimeReadback, +) -> list[BindingIssue]: + issues: list[BindingIssue] = [] + expected_collectives = { + "qkv": QKV_COLLECTIVE_CONTRACT.to_dict(), + "o_proj": O_PROJ_COLLECTIVE_CONTRACT.to_dict(), + } + fixed_fields = { + "input_dtype": "torch.bfloat16", + "weight_dtype": "torch.bfloat16", + "output_dtype": "torch.bfloat16", + "accumulation_dtype": "torch.float32", + "reduction_order": "k_ascending", + "split_k": False, + "policy_id": PROJECTION_POLICY_ID, + } + for projection, expected_collective in expected_collectives.items(): + plan = readback.projection_plans.get(projection) + if plan is None: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}", + message=f"{side} did not report the executed {projection} projection plan", + ) + ) + continue + for field, expected in fixed_fields.items(): + actual = plan.get(field) + if actual != expected: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.{field}", + rollout=actual if side == "rollout" else None, + training=actual if side == "training" else None, + message=f"{side} {projection} {field} must be {expected!r}", + ) + ) + collective = plan.get("collective") + if not isinstance(collective, Mapping) or dict(collective) != expected_collective: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.collective", + message=f"{side} {projection} TP/SP collective directions are invalid", + ) + ) + backend_id = plan.get("backend_id") + if not isinstance(backend_id, str) or not backend_id.strip(): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.backend_id", + message=f"{side} {projection} backend identity is missing", + ) + ) + if not isinstance(plan.get("probe_id"), str) or not plan.get("probe_id"): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.probe_id", + message=f"{side} {projection} bitwise probe identity is missing", + ) + ) + if plan.get("fallback"): + if backend_id != "rlkernel.cuda.det_gemm" or not plan.get("fallback_reason"): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, + tier=BindingTier.SEMANTIC, + field=f"{side}.projection.{projection}.fallback", + message=( + f"{side} {projection} fallback must execute DetGemmOp and record why" + ), + ) + ) + return issues + + def summarize_binding(result: AttentionBindingResult) -> str: """One-line human summary for CLI output and failure messages.""" diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index 196d2e18..b72d6829 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -1,18 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Strict H100 QK-Norm and RoPE handoff for WS2 Attention. +"""Bitwise-bound QK-Norm and RoPE handoff for WS2 Attention. -This module intentionally has no runtime-native fallback. A caller either runs -the RL-Kernel CUDA operators and records their identities, or the experiment -fails before Attention executes. +Megatron/TE and vLLM/FlashInfer remain the first-choice implementations. They +are admitted only after the same-input H100 probe is bitwise identical to the +deterministic RL-Kernel path. A failed probe (or an unavailable native backend) +selects the common RL-Kernel path for both sides and records why the fallback +was taken. The caller must pass this readback to the cross-config binder. """ from __future__ import annotations from dataclasses import dataclass +import hashlib +import json from types import MappingProxyType -from typing import Any, Mapping +from typing import Any, Callable, Mapping import torch from torch import Tensor @@ -20,6 +24,9 @@ QK_RMSNORM_BACKEND_ID = "rlkernel.cuda.rmsnorm" ROPE_BACKEND_ID = "rlkernel.cuda.rope_sm90" +NATIVE_QK_RMSNORM_BACKEND_ID = "native.qk_rmsnorm" +NATIVE_ROPE_BACKEND_ID = "native.rope" +PREPROCESS_POLICY_ID = "ws2.attention.preprocess.v2" MANDATED_ATTENTION_PREPROCESS_BACKENDS: Mapping[str, str] = MappingProxyType( { "qk_rmsnorm": QK_RMSNORM_BACKEND_ID, @@ -37,6 +44,9 @@ class AttentionPreprocessResult: backend_ids: Mapping[str, str] fallback: bool device_capability: tuple[int, int] + fallback_reason: str | None = None + probe_id: str = "" + policy_id: str = PREPROCESS_POLICY_ID def __post_init__(self) -> None: object.__setattr__(self, "backend_ids", MappingProxyType(dict(self.backend_ids))) @@ -54,13 +64,31 @@ def readback_fields(self) -> dict[str, Any]: return { "preprocess_backends": dict(self.backend_ids), "preprocess_fallback": self.fallback, + "preprocess_fallback_reason": self.fallback_reason, + "preprocess_probe_id": self.probe_id, + "preprocess_policy_id": self.policy_id, } class H100AttentionPreprocessor: - """Apply RL-Kernel CUDA QK-Norm then RoPE without silent fallback.""" + """Apply native QK-Norm/RoPE when the H100 probe passes. - def __init__(self, device: torch.device | str | int | None = None) -> None: + ``native_qk_norm`` and ``native_rope`` are framework-owned callables. They + are intentionally injected instead of importing TE/vLLM here, so the same + policy can be used by both runtimes. The deterministic callables default to + RL-Kernel's CUDA operators and are always run to establish the probe oracle. + """ + + def __init__( + self, + device: torch.device | str | int | None = None, + *, + native_qk_norm: Callable[..., Tensor] | None = None, + native_rope: Callable[..., Tensor] | None = None, + native_qk_norm_backend_id: str = NATIVE_QK_RMSNORM_BACKEND_ID, + native_rope_backend_id: str = NATIVE_ROPE_BACKEND_ID, + policy_id: str = PREPROCESS_POLICY_ID, + ) -> None: if not torch.cuda.is_available(): raise RuntimeError("H100AttentionPreprocessor requires an available CUDA runtime") @@ -89,13 +117,13 @@ def __init__(self, device: torch.device | str | int | None = None) -> None: self.rmsnorm = RMSNormCudaOp() self.rope = RoPESM90Op() - actual_backends = { - "qk_rmsnorm": self.rmsnorm.backend_id, - "rope": self.rope.backend_id, - } - if actual_backends != dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS): - raise RuntimeError(f"unexpected Attention preprocess backends: {actual_backends}") - self.backend_ids = MappingProxyType(actual_backends) + if not isinstance(policy_id, str) or not policy_id.strip(): + raise ValueError("policy_id must be a non-empty string") + self.native_qk_norm = native_qk_norm + self.native_rope = native_rope + self.native_qk_norm_backend_id = native_qk_norm_backend_id + self.native_rope_backend_id = native_rope_backend_id + self.policy_id = policy_id def __call__( self, @@ -130,17 +158,78 @@ def forward( theta: float = 1_000_000.0, ) -> AttentionPreprocessResult: _validate_inputs(q, k, q_weight, k_weight, positions, self.device) - q_norm = self.rmsnorm(q, q_weight, eps=eps) - k_norm = self.rmsnorm(k, k_weight, eps=eps) + q_norm_det = self.rmsnorm(q, q_weight, eps=eps) + k_norm_det = self.rmsnorm(k, k_weight, eps=eps) + q_det = self.rope(q_norm_det, positions, theta=theta) + k_det = self.rope(k_norm_det, positions, theta=theta) + + native_available = self.native_qk_norm is not None and self.native_rope is not None + if native_available: + try: + q_norm_native = self.native_qk_norm(q, q_weight, eps=eps) + k_norm_native = self.native_qk_norm(k, k_weight, eps=eps) + q_native = self.native_rope(q_norm_native, positions, theta=theta) + k_native = self.native_rope(k_norm_native, positions, theta=theta) + probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) + if torch.equal(q_norm_native, q_norm_det) and torch.equal( + k_norm_native, k_norm_det + ) and torch.equal(q_native, q_det) and torch.equal(k_native, k_det): + return AttentionPreprocessResult( + q=q_native, + k=k_native, + backend_ids=MappingProxyType( + { + "qk_rmsnorm": self.native_qk_norm_backend_id, + "rope": self.native_rope_backend_id, + } + ), + fallback=False, + device_capability=self.device_capability, + probe_id=probe_id, + policy_id=self.policy_id, + ) + fallback_reason = "native_preprocess_bitwise_probe_failed" + except Exception as exc: # framework backend failures must fail over together + fallback_reason = f"native_preprocess_unavailable:{type(exc).__name__}" + probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) + else: + fallback_reason = "native_preprocess_not_supplied" + probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) return AttentionPreprocessResult( - q=self.rope(q_norm, positions, theta=theta), - k=self.rope(k_norm, positions, theta=theta), - backend_ids=self.backend_ids, - fallback=False, + q=q_det, + k=k_det, + backend_ids=MANDATED_ATTENTION_PREPROCESS_BACKENDS, + fallback=True, device_capability=self.device_capability, + fallback_reason=fallback_reason, + probe_id=probe_id, + policy_id=self.policy_id, ) +def _probe_id( + q: Tensor, + k: Tensor, + q_weight: Tensor, + k_weight: Tensor, + positions: Tensor, + eps: float, + theta: float, +) -> str: + payload = { + "q_shape": list(q.shape), + "k_shape": list(k.shape), + "q_dtype": str(q.dtype), + "k_dtype": str(k.dtype), + "weight_dtype": str(q_weight.dtype), + "positions_shape": list(positions.shape), + "positions_sha256": hashlib.sha256(positions.detach().cpu().numpy().tobytes()).hexdigest(), + "eps": float(eps), + "theta": float(theta), + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16] + + def _validate_inputs( q: Tensor, k: Tensor, @@ -177,4 +266,7 @@ def _validate_inputs( "MANDATED_ATTENTION_PREPROCESS_BACKENDS", "QK_RMSNORM_BACKEND_ID", "ROPE_BACKEND_ID", + "NATIVE_QK_RMSNORM_BACKEND_ID", + "NATIVE_ROPE_BACKEND_ID", + "PREPROCESS_POLICY_ID", ] diff --git a/rl_engine/kernels/attention_projection.py b/rl_engine/kernels/attention_projection.py new file mode 100644 index 00000000..c9bc3043 --- /dev/null +++ b/rl_engine/kernels/attention_projection.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""QKV and output-projection boundaries for the WS2 Attention experiment. + +The framework still owns the native TE/vLLM implementation. This wrapper only +freezes the semantics that must be shared by training and inference: BF16 I/O, +FP32 accumulation, ascending K reduction, and no Split-K. A native callable is +accepted only when its result is bitwise equal to the deterministic fallback. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Callable, Mapping + +import torch +from torch import Tensor + +from rl_engine.kernels.ops.cuda.matmul.det_gemm import DetGemmOp + +ProjectionCallable = Callable[[Tensor, Tensor], Tensor] + +QKV_PROJECTION = "qkv" +O_PROJ_PROJECTION = "o_proj" +PROJECTION_POLICY_ID = "ws2.attention.projection.v1" + + +@dataclass(frozen=True) +class ProjectionCollectiveContract: + """TP/SP directions fixed by the Attention table.""" + + projection: str + tp_forward: str + tp_backward: str + sp_forward: str + sp_backward: str + reduction_forward: str + reduction_backward: str + + def __post_init__(self) -> None: + if self.projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: + raise ValueError(f"unsupported projection {self.projection!r}") + + def to_dict(self) -> dict[str, str]: + return { + "projection": self.projection, + "tp_forward": self.tp_forward, + "tp_backward": self.tp_backward, + "sp_forward": self.sp_forward, + "sp_backward": self.sp_backward, + "reduction_forward": self.reduction_forward, + "reduction_backward": self.reduction_backward, + } + + +QKV_COLLECTIVE_CONTRACT = ProjectionCollectiveContract( + projection=QKV_PROJECTION, + tp_forward="column_parallel", + tp_backward="all_reduce", + sp_forward="all_gather", + sp_backward="reduce_scatter", + reduction_forward="none", + reduction_backward="none", +) +O_PROJ_COLLECTIVE_CONTRACT = ProjectionCollectiveContract( + projection=O_PROJ_PROJECTION, + tp_forward="row_parallel", + tp_backward="none", + sp_forward="reduce_scatter", + sp_backward="all_gather", + reduction_forward="all_reduce", + reduction_backward="none", +) + + +@dataclass(frozen=True) +class ProjectionPlan: + projection: str + backend_id: str + fallback: bool + fallback_reason: str | None + probe_id: str + input_dtype: str = "torch.bfloat16" + weight_dtype: str = "torch.bfloat16" + output_dtype: str = "torch.bfloat16" + accumulation_dtype: str = "torch.float32" + reduction_order: str = "k_ascending" + split_k: bool = False + policy_id: str = PROJECTION_POLICY_ID + collective: Mapping[str, str] = MappingProxyType({}) + + def __post_init__(self) -> None: + if self.projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: + raise ValueError(f"unsupported projection {self.projection!r}") + if self.input_dtype != "torch.bfloat16" or self.weight_dtype != "torch.bfloat16": + raise ValueError("Attention projections require BF16 input and weight") + if self.output_dtype != "torch.bfloat16" or self.accumulation_dtype != "torch.float32": + raise ValueError("Attention projections require FP32 accumulation and BF16 output") + if self.reduction_order != "k_ascending" or self.split_k: + raise ValueError( + "Attention projections require ascending K reduction with Split-K disabled" + ) + object.__setattr__(self, "collective", MappingProxyType(dict(self.collective))) + + def to_dict(self) -> dict[str, Any]: + return { + "projection": self.projection, + "backend_id": self.backend_id, + "fallback": self.fallback, + "fallback_reason": self.fallback_reason, + "probe_id": self.probe_id, + "input_dtype": self.input_dtype, + "weight_dtype": self.weight_dtype, + "output_dtype": self.output_dtype, + "accumulation_dtype": self.accumulation_dtype, + "reduction_order": self.reduction_order, + "split_k": self.split_k, + "policy_id": self.policy_id, + "collective": dict(self.collective), + } + + +@dataclass(frozen=True) +class ProjectionResult: + output: Tensor + plan: ProjectionPlan + + def to_readback(self) -> dict[str, Any]: + return self.plan.to_dict() + + +class AttentionProjectionOp: + """Native-first projection wrapper with a deterministic common fallback.""" + + def __init__( + self, + projection: str, + *, + native: ProjectionCallable | None = None, + native_backend_id: str | None = None, + deterministic: ProjectionCallable | None = None, + policy_id: str = PROJECTION_POLICY_ID, + ) -> None: + if projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: + raise ValueError(f"unsupported projection {projection!r}") + self.projection = projection + self.native = native + self.native_backend_id = native_backend_id or f"native.{projection}" + self.deterministic = deterministic or DetGemmOp() + self.policy_id = policy_id + self.collective = ( + QKV_COLLECTIVE_CONTRACT if projection == QKV_PROJECTION else O_PROJ_COLLECTIVE_CONTRACT + ) + + def __call__(self, x: Tensor, weight: Tensor) -> ProjectionResult: + _validate_projection_inputs(x, weight) + deterministic_out = self.deterministic(x, weight) + if deterministic_out.dtype is not torch.bfloat16: + deterministic_out = deterministic_out.to(torch.bfloat16) + probe_id = _probe_id(x, weight) + + if self.native is not None: + try: + native_out = self.native(x, weight) + if native_out.dtype is not torch.bfloat16: + native_out = native_out.to(torch.bfloat16) + if torch.equal(native_out, deterministic_out): + return ProjectionResult( + native_out, + ProjectionPlan( + projection=self.projection, + backend_id=self.native_backend_id, + fallback=False, + fallback_reason=None, + probe_id=probe_id, + policy_id=self.policy_id, + collective=self.collective.to_dict(), + ), + ) + reason = "native_projection_bitwise_probe_failed" + except Exception as exc: # framework backend failure: use common fallback + reason = f"native_projection_unavailable:{type(exc).__name__}" + else: + reason = "native_projection_not_supplied" + + return ProjectionResult( + deterministic_out, + ProjectionPlan( + projection=self.projection, + backend_id="rlkernel.cuda.det_gemm", + fallback=True, + fallback_reason=reason, + probe_id=probe_id, + policy_id=self.policy_id, + collective=self.collective.to_dict(), + ), + ) + + +def split_qkv( + projected_qkv: Tensor, + q_heads: int, + kv_heads: int, + head_dim: int, +) -> tuple[Tensor, Tensor, Tensor]: + """Split a [Q, K, V] projection in the fixed contiguous Q/K/V order.""" + + if projected_qkv.dim() != 2: + raise ValueError("projected QKV must be [tokens, features]") + q_width = q_heads * head_dim + kv_width = kv_heads * head_dim + expected = q_width + kv_width + kv_width + if projected_qkv.shape[-1] != expected: + raise ValueError( + f"projected QKV width must be {expected}, got {projected_qkv.shape[-1]}" + ) + q, k, v = projected_qkv.split((q_width, kv_width, kv_width), dim=-1) + return q, k, v + + +def _validate_projection_inputs(x: Tensor, weight: Tensor) -> None: + if x.dim() != 2 or weight.dim() != 2: + raise ValueError("projection inputs must be [tokens, K] and [K, N]") + if x.shape[-1] != weight.shape[0]: + raise ValueError("projection K dimensions must match") + if x.dtype is not torch.bfloat16 or weight.dtype is not torch.bfloat16: + raise TypeError("Attention projections require BF16 inputs and weights") + if x.device != weight.device: + raise ValueError("projection inputs must be on the same device") + + +def _probe_id(x: Tensor, weight: Tensor) -> str: + payload = { + "x_shape": list(x.shape), + "weight_shape": list(weight.shape), + "device": str(x.device), + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16] + + +__all__ = [ + "AttentionProjectionOp", + "O_PROJ_COLLECTIVE_CONTRACT", + "O_PROJ_PROJECTION", + "PROJECTION_POLICY_ID", + "ProjectionCollectiveContract", + "ProjectionPlan", + "ProjectionResult", + "QKV_COLLECTIVE_CONTRACT", + "QKV_PROJECTION", + "split_qkv", +] diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 12bb878e..abe55fa3 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -58,6 +58,11 @@ from rl_engine.kernels.attention_preprocess import ( MANDATED_ATTENTION_PREPROCESS_BACKENDS, ) +from rl_engine.kernels.attention_projection import ( + O_PROJ_COLLECTIVE_CONTRACT, + ProjectionPlan, + QKV_COLLECTIVE_CONTRACT, +) pytestmark = pytest.mark.unit @@ -623,6 +628,20 @@ def _statuses(materialization, path): def _readback(materializer, flat, *, source): contract = materializer.build_contract(flat) + projection_plans = { + name: ProjectionPlan( + projection=name, + backend_id="rlkernel.cuda.det_gemm", + fallback=True, + fallback_reason="native projection probe failed", + probe_id=f"{source}-{name}", + collective=collective.to_dict(), + ).to_dict() + for name, collective in ( + ("qkv", QKV_COLLECTIVE_CONTRACT), + ("o_proj", O_PROJ_COLLECTIVE_CONTRACT), + ) + } return AttentionRuntimeReadback( contract=contract, actual_knobs=dict(flat), @@ -631,6 +650,7 @@ def _readback(materializer, flat, *, source): frozen_scope_verified=True, preprocess_backends=MANDATED_ATTENTION_PREPROCESS_BACKENDS, preprocess_fallback=False, + projection_plans=projection_plans, ) @@ -772,6 +792,60 @@ def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): ) +def test_strict_runtime_readback_accepts_common_deterministic_preprocess_fallback(): + rollout = replace( + _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + preprocess_fallback=True, + preprocess_fallback_reason="native probe failed", + preprocess_probe_id="rollout-probe", + ) + training = replace( + _readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ), + preprocess_fallback=True, + preprocess_fallback_reason="native probe failed", + preprocess_probe_id="training-probe", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="rlkernel.cp_attention_reference", + ) + + assert result.passed + + +def test_strict_runtime_readback_accepts_distinct_verified_native_preprocess_backends(): + rollout = replace( + _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + preprocess_backends={"qk_rmsnorm": "native.vllm.rmsnorm", "rope": "native.vllm.rope"}, + preprocess_probe_id="rollout-probe", + ) + training = replace( + _readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ), + preprocess_backends={"qk_rmsnorm": "native.te.rmsnorm", "rope": "native.te.rope"}, + preprocess_probe_id="training-probe", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="megatron.te", + ) + + assert result.passed + + @pytest.mark.parametrize( ("backends", "fallback", "expected_code"), [ @@ -781,14 +855,14 @@ def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): BindingErrorCode.ATTENTION_PREPROCESS_MISSING, ), ( - {**MANDATED_ATTENTION_PREPROCESS_BACKENDS, "qk_rmsnorm": "vllm.native"}, + {**MANDATED_ATTENTION_PREPROCESS_BACKENDS, "qk_rmsnorm": "unknown.backend"}, False, BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, ), ( dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), True, - BindingErrorCode.ATTENTION_PREPROCESS_FALLBACK, + BindingErrorCode.ATTENTION_PREPROCESS_MISMATCH, ), ], ) @@ -800,6 +874,7 @@ def test_strict_runtime_readback_rejects_unverified_preprocess_backend( rollout, preprocess_backends=backends, preprocess_fallback=fallback, + preprocess_fallback_reason=("test fallback" if fallback else None), ) training = _readback( MegatronAttentionMaterializer(), @@ -821,6 +896,30 @@ def test_strict_runtime_readback_rejects_unverified_preprocess_backend( assert result.issues_by_code(expected_code) +def test_strict_runtime_readback_rejects_projection_split_k_or_missing_plan(): + rollout = _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback") + plans = {name: dict(plan) for name, plan in rollout.projection_plans.items()} + plans["qkv"]["split_k"] = True + plans.pop("o_proj") + rollout = replace(rollout, projection_plans=plans) + training = _readback( + MegatronAttentionMaterializer(), TRAINING_KNOBS, source="megatron.runtime_readback" + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="vllm.flash_attn", + training_backend_id="megatron.te", + ) + + assert not result.passed + assert result.issues_by_code(BindingErrorCode.ATTENTION_PROJECTION_MISMATCH) + assert result.issues_by_code(BindingErrorCode.ATTENTION_PROJECTION_MISSING) + + def test_strict_runtime_readback_entrypoint_rejects_unverified_frozen_scope(): rollout_materializer = VllmRolloutMaterializer() training_materializer = MegatronAttentionMaterializer() diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py index 451ec979..a2693d64 100644 --- a/tests/test_attention_preprocess.py +++ b/tests/test_attention_preprocess.py @@ -37,7 +37,7 @@ def _has_h100_preprocess() -> bool: ) -def test_h100_preprocessor_has_no_native_backend_option(): +def test_h100_preprocessor_uses_common_backend_ids_for_fallback(): assert dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) == { "qk_rmsnorm": "rlkernel.cuda.rmsnorm", "rope": "rlkernel.cuda.rope_sm90", @@ -83,11 +83,16 @@ def test_h100_preprocessor_executes_cuda_qk_norm_and_zigzag_rope(): q_ref = rope(norm(q, q_weight), positions) k_ref = rope(norm(k, k_weight), positions) - assert result.fallback is False + assert result.fallback is True assert dict(result.backend_ids) == dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS) + assert result.fallback_reason == "native_preprocess_not_supplied" + assert result.probe_id assert result.readback_fields() == { "preprocess_backends": dict(MANDATED_ATTENTION_PREPROCESS_BACKENDS), - "preprocess_fallback": False, + "preprocess_fallback": True, + "preprocess_fallback_reason": "native_preprocess_not_supplied", + "preprocess_probe_id": result.probe_id, + "preprocess_policy_id": result.policy_id, } torch.testing.assert_close(result.q.float(), q_ref.float(), atol=2e-2, rtol=2e-2) torch.testing.assert_close(result.k.float(), k_ref.float(), atol=2e-2, rtol=2e-2) diff --git a/tests/test_attention_projection.py b/tests/test_attention_projection.py new file mode 100644 index 00000000..146a86bf --- /dev/null +++ b/tests/test_attention_projection.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.attention_projection import ( + AttentionProjectionOp, + O_PROJ_COLLECTIVE_CONTRACT, + QKV_COLLECTIVE_CONTRACT, + split_qkv, +) + + +def _deterministic(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + return torch.mm(x.float(), weight.float()).to(torch.bfloat16) + + +def _inputs(): + torch.manual_seed(19) + return ( + torch.randn(7, 8, dtype=torch.bfloat16), + torch.randn(8, 12, dtype=torch.bfloat16), + ) + + +@pytest.mark.parametrize( + ("projection", "collective"), + [("qkv", QKV_COLLECTIVE_CONTRACT), ("o_proj", O_PROJ_COLLECTIVE_CONTRACT)], +) +def test_projection_falls_back_to_common_deterministic_path(projection, collective): + x, weight = _inputs() + result = AttentionProjectionOp(projection, deterministic=_deterministic)(x, weight) + + assert torch.equal(result.output, _deterministic(x, weight)) + assert result.plan.backend_id == "rlkernel.cuda.det_gemm" + assert result.plan.fallback is True + assert result.plan.fallback_reason == "native_projection_not_supplied" + assert result.plan.split_k is False + assert result.plan.accumulation_dtype == "torch.float32" + assert dict(result.plan.collective) == collective.to_dict() + + +def test_projection_accepts_native_only_after_bitwise_probe(): + x, weight = _inputs() + result = AttentionProjectionOp( + "qkv", + native=_deterministic, + native_backend_id="megatron.te.qkv", + deterministic=_deterministic, + )(x, weight) + + assert torch.equal(result.output, _deterministic(x, weight)) + assert result.plan.backend_id == "megatron.te.qkv" + assert result.plan.fallback is False + assert result.plan.fallback_reason is None + + +def test_projection_rejects_native_drift_and_records_reason(): + x, weight = _inputs() + + def drifting_native(a, b): + return (_deterministic(a, b).float() + 1.0).to(torch.bfloat16) + + result = AttentionProjectionOp( + "o_proj", native=drifting_native, deterministic=_deterministic + )(x, weight) + + assert result.plan.fallback is True + assert result.plan.fallback_reason == "native_projection_bitwise_probe_failed" + assert torch.equal(result.output, _deterministic(x, weight)) + + +def test_split_qkv_is_fixed_contiguous_q_k_v_order(): + projected = torch.arange(2 * 16, dtype=torch.bfloat16).reshape(2, 16) + q, k, v = split_qkv(projected, q_heads=2, kv_heads=1, head_dim=4) + + assert q.shape == (2, 8) + assert k.shape == (2, 4) + assert v.shape == (2, 4) + assert torch.equal(torch.cat((q, k, v), dim=-1), projected) + + +def test_projection_requires_bf16_and_compatible_k(): + x, weight = _inputs() + with pytest.raises(TypeError, match="BF16"): + AttentionProjectionOp("qkv", deterministic=_deterministic)(x.float(), weight) + with pytest.raises(ValueError, match="K dimensions"): + AttentionProjectionOp("qkv", deterministic=_deterministic)(x, weight[:-1]) + + +def test_o_proj_collective_contract_includes_sp_scatter_gather_and_tp_reduction(): + assert O_PROJ_COLLECTIVE_CONTRACT.sp_forward == "reduce_scatter" + assert O_PROJ_COLLECTIVE_CONTRACT.sp_backward == "all_gather" + assert O_PROJ_COLLECTIVE_CONTRACT.reduction_forward == "all_reduce" + assert O_PROJ_COLLECTIVE_CONTRACT.reduction_backward == "none" From 46c5692ca267e06a17aa04e715e6f1385f6a8c30 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Sun, 16 Aug 2026 18:01:30 +0800 Subject: [PATCH 09/17] fix(attention): use portable projection plan default --- rl_engine/kernels/attention_projection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rl_engine/kernels/attention_projection.py b/rl_engine/kernels/attention_projection.py index c9bc3043..5e4dd775 100644 --- a/rl_engine/kernels/attention_projection.py +++ b/rl_engine/kernels/attention_projection.py @@ -13,7 +13,7 @@ import hashlib import json -from dataclasses import dataclass +from dataclasses import dataclass, field from types import MappingProxyType from typing import Any, Callable, Mapping @@ -91,7 +91,7 @@ class ProjectionPlan: reduction_order: str = "k_ascending" split_k: bool = False policy_id: str = PROJECTION_POLICY_ID - collective: Mapping[str, str] = MappingProxyType({}) + collective: Mapping[str, str] = field(default_factory=dict) def __post_init__(self) -> None: if self.projection not in {QKV_PROJECTION, O_PROJ_PROJECTION}: From f55a681dc0ec403f0ca734a2b59acac563ecb970 Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Sun, 16 Aug 2026 16:11:04 +0000 Subject: [PATCH 10/17] fix(attention): satisfy PR4 lint checks --- .../cross_config/attention_binding.py | 21 +++++++++---------- rl_engine/kernels/attention_preprocess.py | 12 ++++++----- rl_engine/kernels/attention_projection.py | 4 +--- tests/test_attention_cross_config_binding.py | 6 ++---- tests/test_attention_preprocess.py | 2 +- tests/test_attention_projection.py | 8 +++---- 6 files changed, 25 insertions(+), 28 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 9adba35d..d23329f6 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -948,21 +948,20 @@ def bind_attention_runtime_readbacks( rollout=rollout_fallback, training=training_fallback, message=( - "QKV/o_proj must use the same native-or-deterministic " - "path on both sides" + "QKV/o_proj must use the same native-or-deterministic " "path on both sides" ), ) ) if rollout_fallback and training_fallback: - for field in ("backend_id", "policy_id", "split_k", "reduction_order"): - if rollout_plan.get(field) != training_plan.get(field): + for field_name in ("backend_id", "policy_id", "split_k", "reduction_order"): + if rollout_plan.get(field_name) != training_plan.get(field_name): missing_scope_evidence.append( BindingIssue( code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, tier=BindingTier.SEMANTIC, - field=f"projection.{projection}.{field}", - rollout=rollout_plan.get(field), - training=training_plan.get(field), + field=f"projection.{projection}.{field_name}", + rollout=rollout_plan.get(field_name), + training=training_plan.get(field_name), message=( "deterministic projection fallback evidence differs " "between sides" @@ -1091,17 +1090,17 @@ def _attention_projection_issues( ) ) continue - for field, expected in fixed_fields.items(): - actual = plan.get(field) + for field_name, expected in fixed_fields.items(): + actual = plan.get(field_name) if actual != expected: issues.append( BindingIssue( code=BindingErrorCode.ATTENTION_PROJECTION_MISMATCH, tier=BindingTier.SEMANTIC, - field=f"{side}.projection.{projection}.{field}", + field=f"{side}.projection.{projection}.{field_name}", rollout=actual if side == "rollout" else None, training=actual if side == "training" else None, - message=f"{side} {projection} {field} must be {expected!r}", + message=f"{side} {projection} {field_name} must be {expected!r}", ) ) collective = plan.get("collective") diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index b72d6829..1a3d4446 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -12,16 +12,15 @@ from __future__ import annotations -from dataclasses import dataclass import hashlib import json +from dataclasses import dataclass from types import MappingProxyType from typing import Any, Callable, Mapping import torch from torch import Tensor - QK_RMSNORM_BACKEND_ID = "rlkernel.cuda.rmsnorm" ROPE_BACKEND_ID = "rlkernel.cuda.rope_sm90" NATIVE_QK_RMSNORM_BACKEND_ID = "native.qk_rmsnorm" @@ -171,9 +170,12 @@ def forward( q_native = self.native_rope(q_norm_native, positions, theta=theta) k_native = self.native_rope(k_norm_native, positions, theta=theta) probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) - if torch.equal(q_norm_native, q_norm_det) and torch.equal( - k_norm_native, k_norm_det - ) and torch.equal(q_native, q_det) and torch.equal(k_native, k_det): + if ( + torch.equal(q_norm_native, q_norm_det) + and torch.equal(k_norm_native, k_norm_det) + and torch.equal(q_native, q_det) + and torch.equal(k_native, k_det) + ): return AttentionPreprocessResult( q=q_native, k=k_native, diff --git a/rl_engine/kernels/attention_projection.py b/rl_engine/kernels/attention_projection.py index 5e4dd775..a7a91bfe 100644 --- a/rl_engine/kernels/attention_projection.py +++ b/rl_engine/kernels/attention_projection.py @@ -215,9 +215,7 @@ def split_qkv( kv_width = kv_heads * head_dim expected = q_width + kv_width + kv_width if projected_qkv.shape[-1] != expected: - raise ValueError( - f"projected QKV width must be {expected}, got {projected_qkv.shape[-1]}" - ) + raise ValueError(f"projected QKV width must be {expected}, got {projected_qkv.shape[-1]}") q, k, v = projected_qkv.split((q_width, kv_width, kv_width), dim=-1) return q, k, v diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index abe55fa3..a859c1e1 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -55,13 +55,11 @@ SplitKVSpec, build_split_kv_runtime_plan_set, ) -from rl_engine.kernels.attention_preprocess import ( - MANDATED_ATTENTION_PREPROCESS_BACKENDS, -) +from rl_engine.kernels.attention_preprocess import MANDATED_ATTENTION_PREPROCESS_BACKENDS from rl_engine.kernels.attention_projection import ( O_PROJ_COLLECTIVE_CONTRACT, - ProjectionPlan, QKV_COLLECTIVE_CONTRACT, + ProjectionPlan, ) pytestmark = pytest.mark.unit diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py index a2693d64..587c0100 100644 --- a/tests/test_attention_preprocess.py +++ b/tests/test_attention_preprocess.py @@ -7,8 +7,8 @@ import torch from rl_engine.kernels.attention_preprocess import ( - H100AttentionPreprocessor, MANDATED_ATTENTION_PREPROCESS_BACKENDS, + H100AttentionPreprocessor, ) from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp diff --git a/tests/test_attention_projection.py b/tests/test_attention_projection.py index 146a86bf..b4094351 100644 --- a/tests/test_attention_projection.py +++ b/tests/test_attention_projection.py @@ -6,9 +6,9 @@ import torch from rl_engine.kernels.attention_projection import ( - AttentionProjectionOp, O_PROJ_COLLECTIVE_CONTRACT, QKV_COLLECTIVE_CONTRACT, + AttentionProjectionOp, split_qkv, ) @@ -63,9 +63,9 @@ def test_projection_rejects_native_drift_and_records_reason(): def drifting_native(a, b): return (_deterministic(a, b).float() + 1.0).to(torch.bfloat16) - result = AttentionProjectionOp( - "o_proj", native=drifting_native, deterministic=_deterministic - )(x, weight) + result = AttentionProjectionOp("o_proj", native=drifting_native, deterministic=_deterministic)( + x, weight + ) assert result.plan.fallback is True assert result.plan.fallback_reason == "native_projection_bitwise_probe_failed" From 553d7993782a7a36bd98c02dceb7ff10a45672d1 Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Sun, 16 Aug 2026 16:26:58 +0000 Subject: [PATCH 11/17] fix(types): narrow optional attention ops --- rl_engine/kernels/attention_preprocess.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/rl_engine/kernels/attention_preprocess.py b/rl_engine/kernels/attention_preprocess.py index 1a3d4446..e97b365f 100644 --- a/rl_engine/kernels/attention_preprocess.py +++ b/rl_engine/kernels/attention_preprocess.py @@ -162,13 +162,14 @@ def forward( q_det = self.rope(q_norm_det, positions, theta=theta) k_det = self.rope(k_norm_det, positions, theta=theta) - native_available = self.native_qk_norm is not None and self.native_rope is not None - if native_available: + native_qk_norm = self.native_qk_norm + native_rope = self.native_rope + if native_qk_norm is not None and native_rope is not None: try: - q_norm_native = self.native_qk_norm(q, q_weight, eps=eps) - k_norm_native = self.native_qk_norm(k, k_weight, eps=eps) - q_native = self.native_rope(q_norm_native, positions, theta=theta) - k_native = self.native_rope(k_norm_native, positions, theta=theta) + q_norm_native = native_qk_norm(q, q_weight, eps=eps) + k_norm_native = native_qk_norm(k, k_weight, eps=eps) + q_native = native_rope(q_norm_native, positions, theta=theta) + k_native = native_rope(k_norm_native, positions, theta=theta) probe_id = _probe_id(q, k, q_weight, k_weight, positions, eps, theta) if ( torch.equal(q_norm_native, q_norm_det) From 62d73f8162c15bea2e6b507893d6e38c8575d59a Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Mon, 17 Aug 2026 18:47:36 +0800 Subject: [PATCH 12/17] feat(attention): bind shared strict core runtime evidence Signed-off-by: lamentropetion <3051000145@qq.com> --- .../cross_config/attention_binding.py | 114 ++++++++++++++++++ rl_engine/kernels/attention_contract.py | 4 + tests/test_attention_cross_config_binding.py | 82 +++++++++++++ 3 files changed, 200 insertions(+) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index d23329f6..d5f53df4 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -44,6 +44,7 @@ from typing import Any, Optional from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, AttentionContract, AttentionContractError, AttentionDType, @@ -122,6 +123,10 @@ class BindingErrorCode(str, Enum): ATTENTION_PREPROCESS_FALLBACK = "ATTENTION_PREPROCESS_FALLBACK" ATTENTION_PROJECTION_MISSING = "ATTENTION_PROJECTION_MISSING" ATTENTION_PROJECTION_MISMATCH = "ATTENTION_PROJECTION_MISMATCH" + ATTENTION_CORE_MISSING = "ATTENTION_CORE_MISSING" + ATTENTION_CORE_MISMATCH = "ATTENTION_CORE_MISMATCH" + ATTENTION_NATIVE_ARITHMETIC = "ATTENTION_NATIVE_ARITHMETIC" + ATTENTION_CORE_SPLIT_K = "ATTENTION_CORE_SPLIT_K" @dataclass(frozen=True) @@ -139,6 +144,10 @@ class AttentionRuntimeReadback: preprocess_probe_id: str = "" preprocess_policy_id: str = PREPROCESS_POLICY_ID projection_plans: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) + strict_mode: bool = False + strict_core_id: str | None = None + native_attention_arithmetic: bool = True + strict_split_kv_policy: str | None = None def __post_init__(self) -> None: if not isinstance(self.contract, AttentionContract): @@ -170,6 +179,20 @@ def __post_init__(self) -> None: raise ValueError("preprocess_policy_id must be a non-empty string") if not isinstance(self.projection_plans, Mapping): raise TypeError("projection_plans must be a mapping") + if not isinstance(self.strict_mode, bool): + raise TypeError("strict_mode must be a bool") + if self.strict_core_id is not None and ( + not isinstance(self.strict_core_id, str) or not self.strict_core_id.strip() + ): + raise ValueError("strict_core_id must be a non-empty string when provided") + if not isinstance(self.native_attention_arithmetic, bool): + raise TypeError("native_attention_arithmetic must be a bool") + if self.strict_split_kv_policy is not None and self.strict_split_kv_policy not in { + "disabled", + "fixed", + "auto", + }: + raise ValueError("strict_split_kv_policy must be disabled, fixed, or auto") normalized_projection_plans: dict[str, Mapping[str, Any]] = {} for name, plan in self.projection_plans.items(): if not isinstance(name, str) or not isinstance(plan, Mapping): @@ -211,6 +234,12 @@ def to_dict(self) -> dict[str, Any]: "attention_projections": { name: dict(plan) for name, plan in self.projection_plans.items() }, + "strict_attention": { + "enabled": self.strict_mode, + "core_id": self.strict_core_id, + "native_attention_arithmetic": self.native_attention_arithmetic, + "split_kv_policy": self.strict_split_kv_policy, + }, "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), } @@ -894,6 +923,8 @@ def bind_attention_runtime_readbacks( ) missing_scope_evidence.extend(_attention_preprocess_issues(side, readback)) missing_scope_evidence.extend(_attention_projection_issues(side, readback)) + missing_scope_evidence.extend(_strict_attention_core_issues(side, readback)) + missing_scope_evidence.extend(_strict_attention_core_pair_issues(rollout, training)) if rollout.preprocess_fallback != training.preprocess_fallback: missing_scope_evidence.append( BindingIssue( @@ -987,6 +1018,10 @@ def bind_attention_runtime_readbacks( "preprocess.fallback_reason": rollout.preprocess_fallback_reason, "preprocess.probe_id": rollout.preprocess_probe_id, "preprocess.policy_id": rollout.preprocess_policy_id, + "strict.enabled": rollout.strict_mode, + "strict.core_id": rollout.strict_core_id, + "strict.native_attention_arithmetic": rollout.native_attention_arithmetic, + "strict.split_kv_policy": rollout.strict_split_kv_policy, **{ f"projection.{projection}": dict(plan) for projection, plan in rollout.projection_plans.items() @@ -1001,6 +1036,10 @@ def bind_attention_runtime_readbacks( "preprocess.fallback_reason": training.preprocess_fallback_reason, "preprocess.probe_id": training.preprocess_probe_id, "preprocess.policy_id": training.preprocess_policy_id, + "strict.enabled": training.strict_mode, + "strict.core_id": training.strict_core_id, + "strict.native_attention_arithmetic": training.native_attention_arithmetic, + "strict.split_kv_policy": training.strict_split_kv_policy, **{ f"projection.{projection}": dict(plan) for projection, plan in training.projection_plans.items() @@ -1009,6 +1048,81 @@ def bind_attention_runtime_readbacks( ) +def _strict_attention_core_issues( + side: str, + readback: AttentionRuntimeReadback, +) -> list[BindingIssue]: + if not readback.strict_mode: + return [] + issues = [] + if readback.strict_core_id != STRICT_ATTENTION_CORE_ID: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.core_id", + message=( + f"{side} strict Attention did not execute the shared core " + f"{STRICT_ATTENTION_CORE_ID!r}" + ), + ) + ) + if readback.native_attention_arithmetic: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_NATIVE_ARITHMETIC, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.native_attention_arithmetic", + message=f"{side} strict Attention entered native TE/FlashInfer arithmetic", + ) + ) + if ( + readback.strict_split_kv_policy != "disabled" + or readback.contract.split_kv.mode.value != "disabled" + ): + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_SPLIT_K, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.split_kv_policy", + message=( + f"{side} strict Attention did not prove Split-KV disabled " + "in both runtime evidence and AttentionContract" + ), + ) + ) + return issues + + +def _strict_attention_core_pair_issues( + rollout: AttentionRuntimeReadback, + training: AttentionRuntimeReadback, +) -> list[BindingIssue]: + if rollout.strict_mode != training.strict_mode: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="strict.enabled", + rollout=rollout.strict_mode, + training=training.strict_mode, + message="training and rollout must use the same strict Attention mode", + ) + ] + if rollout.strict_mode and rollout.strict_core_id != training.strict_core_id: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="strict.core_id", + rollout=rollout.strict_core_id, + training=training.strict_core_id, + message="training and rollout executed different Attention cores", + ) + ] + return [] + + def _attention_preprocess_issues( side: str, readback: AttentionRuntimeReadback, diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 206663e7..b4d85f7a 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -18,6 +18,9 @@ _EnumT = TypeVar("_EnumT", bound=Enum) +STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" + + class AttentionContractError(ValueError): """Raised when attention metadata does not describe a valid invocation.""" @@ -1521,6 +1524,7 @@ class AttentionDispatchResult: "SplitKVRuntimePlanEntry", "SplitKVRuntimePlanSet", "SplitKVSpec", + "STRICT_ATTENTION_CORE_ID", "validate_split_kv_alignment", "validate_split_kv_plan_set_alignment", ] diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index a859c1e1..c5c7f2d4 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -45,6 +45,7 @@ ) from rl_engine.alignment.cross_config.schema import MaterializationStatus from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, AttentionContractError, AttentionMode, AttentionRole, @@ -652,6 +653,23 @@ def _readback(materializer, flat, *, source): ) +def _strict_readback(materializer, flat, *, source): + readback = _readback(materializer, flat, source=source) + contract = replace(readback.contract, split_kv=SplitKVSpec.disabled()) + actual_knobs = dict(readback.actual_knobs) + actual_knobs["attention.split_kv_policy"] = "disabled" + return replace( + readback, + contract=contract, + actual_knobs=actual_knobs, + split_kv_plan_set=_plan_set(contract, backend=source), + strict_mode=True, + strict_core_id=STRICT_ATTENTION_CORE_ID, + native_attention_arithmetic=False, + strict_split_kv_policy="disabled", + ) + + def test_configured_contract_without_runtime_readback_is_unobservable(): normalized = { "batch": {"size": 2}, @@ -790,6 +808,70 @@ def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): ) +@pytest.mark.parametrize( + ("changes", "expected_code"), + [ + ( + {"native_attention_arithmetic": True}, + BindingErrorCode.ATTENTION_NATIVE_ARITHMETIC, + ), + ( + {"strict_core_id": "different.core"}, + BindingErrorCode.ATTENTION_CORE_MISSING, + ), + ( + {"strict_split_kv_policy": "fixed"}, + BindingErrorCode.ATTENTION_CORE_SPLIT_K, + ), + ], +) +def test_strict_runtime_readback_rejects_non_shared_attention_arithmetic(changes, expected_code): + rollout = replace( + _strict_readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + **changes, + ) + training = _strict_readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="rlkernel.cuda.deterministic_attention", + training_backend_id="rlkernel.cuda.deterministic_attention", + ) + + assert not result.passed + assert result.issues_by_code(expected_code) + + +def test_strict_runtime_readback_accepts_shared_no_split_k_core(): + rollout = _strict_readback( + VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback" + ) + training = _strict_readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="rlkernel.cuda.deterministic_attention", + training_backend_id="rlkernel.cuda.deterministic_attention", + ) + + assert result.passed + assert result.provenance["rollout"]["recorded"]["strict.core_id"] == (STRICT_ATTENTION_CORE_ID) + + def test_strict_runtime_readback_accepts_common_deterministic_preprocess_fallback(): rollout = replace( _readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), From 683dd8b8cb84ec06f73270fd6df8a2efc7b7da1e Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Tue, 18 Aug 2026 09:40:17 +0000 Subject: [PATCH 13/17] feat(attention): reuse WS1 PR315 deterministic operators --- .../cuda/attention/deterministic_attention.cu | 54 +++++-- csrc/cuda/gemm/det_gemm_kernel.cu | 67 ++++++--- csrc/ops.cpp | 18 +++ rl_engine/kernels/ops/backward_runtime.py | 51 +++++++ .../ops/cuda/attention/deterministic_attn.py | 30 +++- rl_engine/kernels/ops/cuda/matmul/det_gemm.py | 76 +++++++++- rl_engine/kernels/ops/cuda/norm/rmsnorm.py | 40 ++++-- .../kernels/ops/cuda/rotary_embedding/rope.py | 119 ++++++++++------ rl_engine/kernels/ops/vjp_fp32.py | 133 ++++++++++++++++++ 9 files changed, 491 insertions(+), 97 deletions(-) create mode 100644 rl_engine/kernels/ops/backward_runtime.py create mode 100644 rl_engine/kernels/ops/vjp_fp32.py diff --git a/csrc/cuda/attention/deterministic_attention.cu b/csrc/cuda/attention/deterministic_attention.cu index 973b07a8..aaa70b42 100644 --- a/csrc/cuda/attention/deterministic_attention.cu +++ b/csrc/cuda/attention/deterministic_attention.cu @@ -147,9 +147,8 @@ __global__ void masked_softmax_lse_kernel( } } else { lse_val = row_max + logf(row_sum); - float inv_sum = 1.0f / row_sum; for (int k = threadIdx.x; k < Skv; k += kSoftmaxThreads) { - row[k] *= inv_sum; + row[k] /= row_sum; } } @@ -166,11 +165,11 @@ __global__ void masked_softmax_lse_kernel( constexpr int kPVTileQ = 16; constexpr int kPVTileD = 16; -template +template __global__ void pv_kernel( const float* __restrict__ P, // [B, Hq, Sq, Skv] - const scalar_t* __restrict__ V, // [B, Hkv, Skv, D] - scalar_t* __restrict__ out, // [B, Hq, Sq, D] + const input_t* __restrict__ V, // [B, Hkv, Skv, D] + output_t* __restrict__ out, // [B, Hq, Sq, D] int64_t B, int64_t Hq, int64_t Hkv, int64_t Sq, int64_t Skv, int64_t D) { @@ -185,7 +184,7 @@ __global__ void pv_kernel( const int kv_head = hq / (Hq / Hkv); const float* p_row = P + ((int64_t)b * Hq * Sq * Skv + (int64_t)hq * Sq * Skv + (int64_t)q_idx * Skv); - const scalar_t* v_base = V + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D); + const input_t* v_base = V + ((int64_t)b * Hkv * Skv * D + (int64_t)kv_head * Skv * D); float acc = 0.0f; for (int64_t k = 0; k < Skv; ++k) { @@ -193,7 +192,7 @@ __global__ void pv_kernel( } const int64_t out_idx = (int64_t)b * Hq * Sq * D + (int64_t)hq * Sq * D + (int64_t)q_idx * D + d_idx; - out[out_idx] = (scalar_t)acc; + out[out_idx] = (output_t)acc; } void check_deterministic_attention_inputs( @@ -257,13 +256,14 @@ void check_deterministic_attention_inputs( // out: [B, Hq, Sq, D] same dtype as q // lse: [B, Hq, Sq] FP32 // P: [B, Hq, Sq, Skv] FP32 (softmax probabilities, saved for backward) -std::vector deterministic_attention_forward( +std::vector deterministic_attention_forward_impl( torch::Tensor q, torch::Tensor k, torch::Tensor v, bool causal, double scale, - torch::optional key_padding_mask) { + torch::optional key_padding_mask, + bool output_fp32) { check_deterministic_attention_inputs(q, k, v, key_padding_mask); const at::cuda::OptionalCUDAGuard device_guard(at::device_of(q)); @@ -327,7 +327,9 @@ std::vector deterministic_attention_forward( } // --- Launch PV kernel --- - auto out = torch::empty_like(q_contig); + auto out = output_fp32 + ? torch::empty(q_contig.sizes(), q_contig.options().dtype(at::kFloat)) + : torch::empty_like(q_contig); { dim3 block(kPVTileD, kPVTileQ); dim3 grid( @@ -337,11 +339,19 @@ std::vector deterministic_attention_forward( AT_DISPATCH_FLOATING_TYPES_AND2( at::ScalarType::Half, at::ScalarType::BFloat16, q_contig.scalar_type(), "pv_kernel", [&] { - pv_kernel<<>>( - scores.data_ptr(), - v_contig.data_ptr(), - out.data_ptr(), - B, Hq, Hkv, Sq, Skv, D); + if (output_fp32) { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + } else { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + } C10_CUDA_KERNEL_LAUNCH_CHECK(); }); } @@ -349,6 +359,20 @@ std::vector deterministic_attention_forward( return {out, lse, scores}; } +std::vector deterministic_attention_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v, bool causal, double scale, + torch::optional key_padding_mask) { + return deterministic_attention_forward_impl( + q, k, v, causal, scale, key_padding_mask, false); +} + +std::vector deterministic_attention_forward_fp32( + torch::Tensor q, torch::Tensor k, torch::Tensor v, bool causal, double scale, + torch::optional key_padding_mask) { + return deterministic_attention_forward_impl( + q, k, v, causal, scale, key_padding_mask, true); +} + // =========================================================================== // BACKWARD // =========================================================================== diff --git a/csrc/cuda/gemm/det_gemm_kernel.cu b/csrc/cuda/gemm/det_gemm_kernel.cu index 4d9035fc..cd92fc9d 100644 --- a/csrc/cuda/gemm/det_gemm_kernel.cu +++ b/csrc/cuda/gemm/det_gemm_kernel.cu @@ -24,15 +24,29 @@ namespace { using nv_bf16 = __nv_bfloat16; +template +__device__ __forceinline__ output_t cast_output(float value); + +template <> +__device__ __forceinline__ nv_bf16 cast_output(float value) { + return __float2bfloat16(value); +} + +template <> +__device__ __forceinline__ float cast_output(float value) { + return value; +} + __host__ __device__ constexpr int cdiv(int a, int b) { return (a + b - 1) / b; } // Naive FP32 scalar kernel (fallback + ground truth). Batch-invariant by // construction: one thread = one output element, fixed ascending K loop. constexpr int NAIVE_TILE = 16; +template __global__ void det_gemm_naive(const nv_bf16* __restrict__ A, const nv_bf16* __restrict__ B, - nv_bf16* __restrict__ C, + output_t* __restrict__ C, int M, int N, int K) { const int row = blockIdx.y * NAIVE_TILE + threadIdx.y; const int col = blockIdx.x * NAIVE_TILE + threadIdx.x; @@ -40,14 +54,15 @@ __global__ void det_gemm_naive(const nv_bf16* __restrict__ A, float acc = 0.0f; for (int k = 0; k < K; ++k) acc += __bfloat162float(A[row * K + k]) * __bfloat162float(B[k * N + col]); - C[row * N + col] = __float2bfloat16(acc); + C[row * N + col] = cast_output(acc); } -void launch_naive(const nv_bf16* A, const nv_bf16* B, nv_bf16* C, +template +void launch_naive(const nv_bf16* A, const nv_bf16* B, output_t* C, int M, int N, int K, cudaStream_t stream) { dim3 block(NAIVE_TILE, NAIVE_TILE); dim3 grid(cdiv(N, NAIVE_TILE), cdiv(M, NAIVE_TILE)); - det_gemm_naive<<>>(A, B, C, M, N, K); + det_gemm_naive<<>>(A, B, C, M, N, K); } #if defined(RL_KERNEL_ENABLE_SM90) @@ -81,9 +96,10 @@ __device__ __forceinline__ void mma_m16n8k16(const uint32_t A[4], const uint32_t "f"(D[0]), "f"(D[1]), "f"(D[2]), "f"(D[3])); } +template __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap, const __grid_constant__ CUtensorMap bt_tmap, - nv_bf16* __restrict__ C, + output_t* __restrict__ C, int M, int N, int K) { const int tid = threadIdx.x; const int warp = tid / 32; @@ -186,18 +202,19 @@ __global__ void det_gemm_sm90_kernel(const __grid_constant__ CUtensorMap a_tmap, for (int n = 0; n < N_TILES; ++n) { const int col = col_base + n * MMA_N + (lane % 4) * 2; if (row < M && col + 1 < N) { - C[row * N + col + 0] = __float2bfloat16(acc[mi][n][0]); - C[row * N + col + 1] = __float2bfloat16(acc[mi][n][1]); + C[row * N + col + 0] = cast_output(acc[mi][n][0]); + C[row * N + col + 1] = cast_output(acc[mi][n][1]); } if (row + 8 < M && col + 1 < N) { - C[(row + 8) * N + col + 0] = __float2bfloat16(acc[mi][n][2]); - C[(row + 8) * N + col + 1] = __float2bfloat16(acc[mi][n][3]); + C[(row + 8) * N + col + 0] = cast_output(acc[mi][n][2]); + C[(row + 8) * N + col + 1] = cast_output(acc[mi][n][3]); } } } } -bool launch_sm90(const nv_bf16* A, const nv_bf16* Bt, nv_bf16* C, +template +bool launch_sm90(const nv_bf16* A, const nv_bf16* Bt, output_t* C, int M, int N, int K, cudaStream_t stream) { if (M % BM != 0 || N % BN != 0 || K % BK != 0) return false; // fall back @@ -207,11 +224,11 @@ bool launch_sm90(const nv_bf16* A, const nv_bf16* Bt, nv_bf16* C, const int smem = STAGES * (BM * BK + BN * BK) * sizeof(nv_bf16) + STAGES * 8; if (smem > 48 * 1024) - cudaFuncSetAttribute(det_gemm_sm90_kernel, + cudaFuncSetAttribute(det_gemm_sm90_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); dim3 grid(cdiv(N, BN), cdiv(M, BM)); - det_gemm_sm90_kernel<<>>(a_tmap, bt_tmap, C, M, N, K); + det_gemm_sm90_kernel<<>>(a_tmap, bt_tmap, C, M, N, K); return true; } #endif // RL_KERNEL_ENABLE_SM90 @@ -232,9 +249,11 @@ void check_in(const torch::Tensor& t, const char* n) { TORCH_CHECK(t.scalar_type() == torch::kBFloat16, n, " must be bf16"); } -torch::Tensor gemm_dispatch(const torch::Tensor& a, const torch::Tensor& b) { +torch::Tensor gemm_dispatch(const torch::Tensor& a, const torch::Tensor& b, + bool output_fp32 = false) { const int M = a.size(0), K = a.size(1), N = b.size(1); - auto c = torch::empty({M, N}, a.options()); + auto options = a.options().dtype(output_fp32 ? torch::kFloat32 : torch::kBFloat16); + auto c = torch::empty({M, N}, options); auto stream = at::cuda::getCurrentCUDAStream(); #if defined(RL_KERNEL_ENABLE_SM90) @@ -249,15 +268,21 @@ torch::Tensor gemm_dispatch(const torch::Tensor& a, const torch::Tensor& b) { a_use = torch::zeros({Mp, K}, a.options()); a_use.narrow(0, 0, M).copy_(a); } - torch::Tensor c_use = (Mp != M) ? torch::empty({Mp, N}, a.options()) : c; + torch::Tensor c_use = (Mp != M) ? torch::empty({Mp, N}, options) : c; auto bt = b.t().contiguous(); // [N,K] - if (launch_sm90(bf16(a_use), bf16(bt), bf16o(c_use), Mp, N, K, stream)) { + const bool launched = output_fp32 + ? launch_sm90(bf16(a_use), bf16(bt), c_use.data_ptr(), Mp, N, K, stream) + : launch_sm90(bf16(a_use), bf16(bt), bf16o(c_use), Mp, N, K, stream); + if (launched) { if (Mp != M) c.copy_(c_use.narrow(0, 0, M)); return c; } } #endif - launch_naive(bf16(a), bf16(b), bf16o(c), M, N, K, stream); + if (output_fp32) + launch_naive(bf16(a), bf16(b), c.data_ptr(), M, N, K, stream); + else + launch_naive(bf16(a), bf16(b), bf16o(c), M, N, K, stream); return c; } @@ -271,6 +296,14 @@ torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b) { return gemm_dispatch(a, b); } +torch::Tensor det_gemm_fwd_fp32(torch::Tensor a, torch::Tensor b) { + check_in(a, "A"); check_in(b, "B"); + a = a.contiguous(); b = b.contiguous(); + TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "det_gemm_fwd_fp32: expect 2D [M,K]@[K,N]"); + TORCH_CHECK(b.size(0) == a.size(1), "det_gemm_fwd_fp32: K mismatch"); + return gemm_dispatch(a, b, true); +} + torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b) { check_in(dc, "dC"); check_in(b, "B"); dc = dc.contiguous(); diff --git a/csrc/ops.cpp b/csrc/ops.cpp index eee328a4..58692de1 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -71,6 +71,7 @@ torch::Tensor lm_head_sm90_forward(torch::Tensor hidden, torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden, torch::Tensor weight, torch::optional bias); +torch::Tensor det_gemm_rowwise_fwd_fp32(torch::Tensor a, torch::Tensor b); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) @@ -90,6 +91,7 @@ torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torc // Batch-Invariant Deterministic GEMM Declarations torch::Tensor det_gemm_fwd(torch::Tensor a, torch::Tensor b); +torch::Tensor det_gemm_fwd_fp32(torch::Tensor a, torch::Tensor b); torch::Tensor det_gemm_da(torch::Tensor dc, torch::Tensor b); torch::Tensor det_gemm_db(torch::Tensor a, torch::Tensor dc); // SiLU / SwiGLU Declarations (elementwise activation, general CUDA) @@ -241,6 +243,14 @@ std::vector deterministic_attention_forward( double scale, torch::optional key_padding_mask); +std::vector deterministic_attention_forward_fp32( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + bool causal, + double scale, + torch::optional key_padding_mask); + std::vector deterministic_attention_backward( torch::Tensor grad_output, torch::Tensor q, @@ -338,6 +348,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Single-card SM90 batch-invariant LM-head forward"); m.def("lm_head_sm90_forward_fp32", &lm_head_sm90_forward_fp32, "Single-card SM90 batch-invariant LM-head forward with fp32 output"); + m.def("det_gemm_rowwise_fwd_fp32", &det_gemm_rowwise_fwd_fp32, + "SM90 deterministic rowwise GEMM with FP32 inputs/accumulation/output"); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) @@ -360,6 +372,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // registry Batch-Invariant Deterministic GEMM m.def("det_gemm_fwd", &det_gemm_fwd, "Batch-invariant deterministic GEMM forward (C=A@B)"); + m.def("det_gemm_fwd_fp32", &det_gemm_fwd_fp32, + "Batch-invariant deterministic GEMM forward with FP32 output"); m.def("det_gemm_da", &det_gemm_da, "Batch-invariant deterministic GEMM backward dA (dC@B^T)"); m.def("det_gemm_db", &det_gemm_db, "Batch-invariant deterministic GEMM backward dB (A^T@dC)"); // registry RMSNorm @@ -378,6 +392,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "deterministic_attention_forward", &deterministic_attention_forward, "Deterministic standard softmax attention forward (out, lse)"); + m.def( + "deterministic_attention_forward_fp32", + &deterministic_attention_forward_fp32, + "Deterministic standard softmax attention forward with FP32 output"); m.def( "deterministic_attention_backward", &deterministic_attention_backward, diff --git a/rl_engine/kernels/ops/backward_runtime.py b/rl_engine/kernels/ops/backward_runtime.py new file mode 100644 index 00000000..5cc7d074 --- /dev/null +++ b/rl_engine/kernels/ops/backward_runtime.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime record of the kernel that actually executed a candidate backward.""" + +from __future__ import annotations + +from threading import Lock +from typing import Any + +_LOCK = Lock() +_EVENTS: dict[str, dict[str, Any]] = {} + + +def record_backward( + kind: str, + *, + kernel_id: str, + impl: str, + family: str, +) -> None: + with _LOCK: + previous = _EVENTS.get(kind) + count = 1 if previous is None else int(previous["execution_count"]) + 1 + kernel_ids = tuple(part for part in kernel_id.split("+") if part) + _EVENTS[kind] = { + "kind": kind, + "implementation_ids": list(kernel_ids), + "kernel_ids": list(kernel_ids), + "kernel_id": kernel_id, + "impl": impl, + "family": family, + "execution_count": count, + } + + +def snapshot_backward_runtime() -> dict[str, dict[str, Any]]: + with _LOCK: + return {key: dict(value) for key, value in _EVENTS.items()} + + +def reset_backward_runtime() -> None: + with _LOCK: + _EVENTS.clear() + + +__all__ = [ + "record_backward", + "reset_backward_runtime", + "snapshot_backward_runtime", +] diff --git a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py index 81f80a7f..01c51c99 100644 --- a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py @@ -33,13 +33,18 @@ def forward( causal: bool, scale: float, key_padding_mask: Optional[torch.Tensor], + output_fp32: bool, ) -> tuple[torch.Tensor, torch.Tensor]: q_c = q.contiguous() k_c = k.contiguous() v_c = v.contiguous() mask_c = key_padding_mask.contiguous() if key_padding_mask is not None else None - results = _C.deterministic_attention_forward(q_c, k_c, v_c, causal, float(scale), mask_c) + results = ( + _C.deterministic_attention_forward_fp32(q_c, k_c, v_c, causal, float(scale), mask_c) + if output_fp32 + else _C.deterministic_attention_forward(q_c, k_c, v_c, causal, float(scale), mask_c) + ) out, lse, P = results[0], results[1], results[2] ctx.save_for_backward(q_c, k_c, v_c, P, mask_c) @@ -54,6 +59,8 @@ def forward( def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): q_c, k_c, v_c, P, mask_c = ctx.saved_tensors + if grad_out.dtype != q_c.dtype: + grad_out = grad_out.to(q_c.dtype) dQ, dK, dV = _C.deterministic_attention_backward( grad_out.contiguous(), q_c, @@ -65,7 +72,7 @@ def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): mask_c, ) - return dQ, dK, dV, None, None, None + return dQ, dK, dV, None, None, None, None class DeterministicAttentionOp: @@ -130,10 +137,27 @@ def forward_with_lse( self._validate_inputs(q, k, v, key_padding_mask) resolved_scale = scale if scale is not None else (1.0 / math.sqrt(q.shape[-1])) out, lse = _DeterministicAttentionFn.apply( - q, k, v, causal, resolved_scale, key_padding_mask + q, k, v, causal, resolved_scale, key_padding_mask, False ) return out, lse + def forward_fp32( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + self._validate_inputs(q, k, v, key_padding_mask) + resolved_scale = scale if scale is not None else (1.0 / math.sqrt(q.shape[-1])) + out, _lse = _DeterministicAttentionFn.apply( + q, k, v, causal, resolved_scale, key_padding_mask, True + ) + return out + @staticmethod def _validate_inputs( q: torch.Tensor, diff --git a/rl_engine/kernels/ops/cuda/matmul/det_gemm.py b/rl_engine/kernels/ops/cuda/matmul/det_gemm.py index 4778be90..c410fbb2 100644 --- a/rl_engine/kernels/ops/cuda/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/cuda/matmul/det_gemm.py @@ -9,22 +9,70 @@ """ import torch +from rl_engine.kernels.ops.backward_runtime import record_backward from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE from rl_engine.utils.logger import logger class _DetGemmFn(torch.autograd.Function): @staticmethod - def forward(ctx, a, b): + def forward(ctx, a, b, output_fp32=False): ctx.save_for_backward(a, b) + if output_fp32: + if not hasattr(_C, "det_gemm_fwd_fp32"): + raise RuntimeError("FP32 deterministic GEMM output requires the rebuilt extension") + return _C.det_gemm_fwd_fp32(a, b) return _C.det_gemm_fwd(a, b) @staticmethod def backward(ctx, grad_out): a, b = ctx.saved_tensors grad_out = grad_out.contiguous() + if grad_out.dtype != torch.bfloat16: + grad_out = grad_out.to(torch.bfloat16) da = _C.det_gemm_da(grad_out, b) if ctx.needs_input_grad[0] else None db = _C.det_gemm_db(a, grad_out) if ctx.needs_input_grad[1] else None + record_backward( + "det_gemm", + kernel_id="rl_engine._C.det_gemm_da+rl_engine._C.det_gemm_db", + impl="cuda_det_gemm", + family="cuda", + ) + return da, db, None + + +class _DetGemmAccumFn(torch.autograd.Function): + @staticmethod + def forward(ctx, a, b): + ctx.save_for_backward(a, b) + if not hasattr(_C, "det_gemm_rowwise_fwd_fp32"): + raise RuntimeError( + "FP32 rowwise deterministic GEMM requires the rebuilt SM90 extension" + ) + return _C.det_gemm_rowwise_fwd_fp32(a, b) + + @staticmethod + def backward(ctx, grad_out): + a, b = ctx.saved_tensors + grad_fp32 = grad_out.contiguous().float() + a_fp32 = a.contiguous().float() + b_fp32 = b.contiguous().float() + da = ( + _C.det_gemm_rowwise_fwd_fp32(grad_fp32, b_fp32.t().contiguous()).to(a.dtype) + if ctx.needs_input_grad[0] + else None + ) + db = ( + _C.det_gemm_rowwise_fwd_fp32(a_fp32.t().contiguous(), grad_fp32).to(b.dtype) + if ctx.needs_input_grad[1] + else None + ) + record_backward( + "det_gemm", + kernel_id=("rl_engine._C.det_gemm_rowwise_fwd_fp32"), + impl="cuda_rowwise_fp32_accum_det_gemm", + family="cuda", + ) return da, db @@ -51,9 +99,31 @@ def __call__(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: "DetGemmOp: compiled _C.det_gemm kernel unavailable; no " "batch-invariant fallback exists. Build the extension first." ) - return _DetGemmFn.apply(a.contiguous(), b.contiguous()) + return _DetGemmFn.apply(a.contiguous(), b.contiguous(), False) + + def forward_fp32(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16, "BF16 only" + assert a.is_cuda and b.is_cuda, "Inputs must be on CUDA device" + if not self.has_hardware_op: + raise RuntimeError("DetGemmOp: compiled CUDA extension unavailable") + return _DetGemmFn.apply(a.contiguous(), b.contiguous(), True) + + def forward_accum_fp32(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + if a.dtype not in (torch.bfloat16, torch.float32) or b.dtype not in ( + torch.bfloat16, + torch.float32, + ): + raise TypeError("FP32-accumulation GEMM requires BF16 or FP32 inputs") + assert a.is_cuda and b.is_cuda, "Inputs must be on CUDA device" + return _DetGemmAccumFn.apply(a.contiguous(), b.contiguous()) + + def parameter_vjp_contributions_fp32(self, *, a, b, grad_output): + del b + rows_a = a.float() + rows_g = grad_output.float() + return {"b": rows_a[:, :, None] * rows_g[:, None, :]} def deterministic_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """Functional entry. a:[M,K] bf16, b:[K,N] bf16 -> [M,N] bf16.""" - return _DetGemmFn.apply(a, b) + return _DetGemmFn.apply(a, b, False) diff --git a/rl_engine/kernels/ops/cuda/norm/rmsnorm.py b/rl_engine/kernels/ops/cuda/norm/rmsnorm.py index 7ac85b70..d4cefae1 100644 --- a/rl_engine/kernels/ops/cuda/norm/rmsnorm.py +++ b/rl_engine/kernels/ops/cuda/norm/rmsnorm.py @@ -1,6 +1,8 @@ import torch +from rl_engine.kernels.ops.backward_runtime import record_backward from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.vjp_fp32 import reduce_rows_fp32, rmsnorm_dweight_rows_fp32 class RMSNormCuda(torch.autograd.Function): @@ -62,7 +64,21 @@ def backward(ctx, grad_out): dx = _C.rmsnorm_backward_dx(dy, x, weight, rstd) - dw = _C.rmsnorm_backward_dw(dy, x, rstd, mask).to(weight.dtype) + # Explicit, shape-independent FP32 left fold. This is slower than + # the chunked extension but preserves the C2 Batch/Chunk reduction order. + rows = rmsnorm_dweight_rows_fp32(x, dy, rstd=rstd) + rows = rows * mask.to(dtype=rows.dtype).unsqueeze(-1) + dw = reduce_rows_fp32(rows).to(weight.dtype) + record_backward( + "rms_norm", + kernel_id=( + "rl_engine._C.rmsnorm_backward_dx" + "+rl_engine.kernels.ops.vjp_fp32.rmsnorm_dweight_rows_fp32" + "+rl_engine.kernels.ops.vjp_fp32.reduce_rows_fp32" + ), + impl="cuda_rmsnorm_dx_declared_fp32_rowfold_dw", + family="cuda", + ) return dx, dw, None, None @@ -79,20 +95,7 @@ def rmsnorm_cuda(x, weight, eps=1e-6, mask=None): class RMSNormCudaOp: """CUDA RMSNorm wrapper compatible with the shared operator harness.""" - backend_id = "rlkernel.cuda.rmsnorm" - - def __init__(self): - required = ( - "rmsnorm_forward", - "rmsnorm_backward_dx", - "rmsnorm_backward_dw", - ) - missing = [name for name in required if not _EXT_AVAILABLE or not hasattr(_C, name)] - if missing: - raise RuntimeError( - "CUDA RMSNorm extension is incomplete; rebuild _C with rmsnorm.cu " - f"(missing: {', '.join(missing)})" - ) + backward_impl = "cuda_rmsnorm_dx_declared_fp32_rowfold_dw" def __call__(self, x, weight, *, eps=1e-6): return self.forward(x, weight, eps=eps) @@ -102,3 +105,10 @@ def forward(self, x, weight, *, eps=1e-6): x_2d = x.contiguous().view(-1, hidden) y_2d = rmsnorm_cuda(x_2d, weight.contiguous(), eps=eps) return y_2d.view_as(x) + + def parameter_vjp_contributions_fp32(self, *, x, weight, grad_output, eps=1e-6): + del weight + x32 = x.float() + rstd = torch.rsqrt(x32.square().mean(dim=-1) + float(eps)) + rows = grad_output.float() * x32 * rstd.unsqueeze(-1) + return {"weight": rows} diff --git a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py index 028e44b9..9a764012 100644 --- a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py +++ b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py @@ -18,69 +18,96 @@ def _build_cos_sin(positions: Tensor, half: int, theta: float, device: torch.device): - """fp32 cos/sin rows, identical math to NativeRoPEOp.""" + """fp32 cos/sin caches of shape [S, half], identical math to NativeRoPEOp.""" inv_freq = 1.0 / (theta ** (torch.arange(0, half, dtype=torch.float32, device=device) / half)) pos = positions.to(device=device, dtype=torch.float32).reshape(-1, 1) freqs = pos * inv_freq # [S, half] return freqs.cos().contiguous(), freqs.sin().contiguous() -class _RoPEFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: - D = x.shape[-1] - if D % 2 != 0: - raise ValueError(f"RoPE head_dim must be even, got {D}") - if positions.dim() not in (1, 2): - raise ValueError("positions must have shape [S] or [B, S]") - S = positions.shape[-1] - if S == 0: - raise ValueError("positions must not be empty") - if x.shape[-2] != S: - raise ValueError(f"x sequence length {x.shape[-2]} does not match positions length {S}") +def _rope_table(x: Tensor, positions: Tensor, theta: float) -> tuple[Tensor, Tensor, Tensor]: + """Build (x_2d, cos, sin) for [S] or [B, S] positions. See Triton RoPE.""" + D = x.shape[-1] + if D % 2 != 0: + raise ValueError(f"RoPE head_dim must be even, got {D}") + if positions.dim() == 1: + table_len = int(positions.shape[0]) x_2d = x.contiguous().reshape(-1, D) - n_rows = x_2d.shape[0] - if positions.dim() == 2: - batch = positions.shape[0] - if x.dim() < 3 or x.shape[0] != batch: - raise ValueError( - f"x batch size {x.shape[0]} does not match positions batch size {batch}" - ) - rows_per_token = n_rows // (batch * S) - if rows_per_token * batch * S != n_rows: - raise ValueError("x rows are incompatible with [B, S] positions") - # The CUDA kernel accepts one fp32 cos/sin row per flattened x row. - # Expanding positions preserves arbitrary global/zigzag indices while - # keeping the arithmetic inside the precompiled deterministic kernel. - kernel_positions = ( - positions[:, None, :].expand(batch, rows_per_token, S).contiguous().reshape(-1) - ) - else: - kernel_positions = positions - if n_rows % kernel_positions.numel() != 0: + if x_2d.shape[0] % table_len != 0: raise ValueError( - f"row count {n_rows} not divisible by position rows " - f"{kernel_positions.numel()}; " + f"row count {x_2d.shape[0]} not divisible by seq length {table_len}; " "expected a [..., S, D] contiguous layout." ) - cos, sin = _build_cos_sin(kernel_positions, D // 2, float(theta), x.device) + cos, sin = _build_cos_sin(positions, D // 2, float(theta), x.device) + return x_2d, cos, sin + if positions.dim() != 2: + raise ValueError(f"positions must be [S] or [B, S], got shape {tuple(positions.shape)}") + batch, seq = positions.shape + if x.shape[0] != batch or x.shape[-2] != seq: + raise ValueError( + f"positions {tuple(positions.shape)} is incompatible with x {tuple(x.shape)}; " + "expected x [B, ..., S, D]" + ) + if x.dim() == 4: + x_2d = x.permute(1, 0, 2, 3).contiguous().reshape(-1, D) + elif x.dim() == 3: + x_2d = x.contiguous().reshape(-1, D) + else: + raise ValueError( + f"RoPE [B, S] positions require x [B, S, D] or [B, H, S, D], got {x.dim()}D" + ) + table_len = batch * seq + if x_2d.shape[0] % table_len != 0: + raise ValueError(f"row count {x_2d.shape[0]} not divisible by B*S={table_len}") + cos, sin = _build_cos_sin(positions.reshape(-1), D // 2, float(theta), x.device) + return x_2d, cos, sin + + +def _restore_rope(out_2d: Tensor, x: Tensor, positions: Tensor) -> Tensor: + if positions.dim() == 1 or x.dim() != 4: + return out_2d.reshape(x.shape) + heads, batch, seq, dim = x.shape[1], x.shape[0], x.shape[2], x.shape[3] + return out_2d.reshape(heads, batch, seq, dim).permute(1, 0, 2, 3).contiguous() + + +class _RoPEFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: + x_2d, cos, sin = _rope_table(x, positions, theta) ctx.save_for_backward(cos, sin) - out = _C.rope_apply_sm90(x_2d, cos, sin, 1.0) - return out.reshape(x.shape) + ctx.x_shape = tuple(x.shape) + ctx.pos_dim = positions.dim() + out_2d = _C.rope_apply_sm90(x_2d, cos, sin, 1.0) + return _restore_rope(out_2d, x, positions) @staticmethod def backward(ctx, grad_out: Tensor): cos, sin = ctx.saved_tensors grad_x = None if ctx.needs_input_grad[0]: - D = grad_out.shape[-1] - g_2d = grad_out.contiguous().reshape(-1, D) - # Inverse rotation: same kernel with the sine negated. - grad_x = _C.rope_apply_sm90(g_2d, cos, sin, -1.0).reshape(grad_out.shape) - # Inputs: x, positions, theta. + if ctx.pos_dim == 2 and len(ctx.x_shape) == 4: + g_2d = grad_out.permute(1, 0, 2, 3).contiguous().reshape(-1, ctx.x_shape[-1]) + out_2d = _C.rope_apply_sm90(g_2d, cos, sin, -1.0) + heads, batch, seq, dim = ( + ctx.x_shape[1], + ctx.x_shape[0], + ctx.x_shape[2], + ctx.x_shape[3], + ) + grad_x = out_2d.reshape(heads, batch, seq, dim).permute(1, 0, 2, 3).contiguous() + else: + g_2d = grad_out.contiguous().reshape(-1, grad_out.shape[-1]) + grad_x = _C.rope_apply_sm90(g_2d, cos, sin, -1.0).reshape(grad_out.shape) return grad_x, None, None +def _is_hopper(device: torch.device) -> bool: + try: + return torch.cuda.get_device_capability(device)[0] == 9 + except Exception: + return False + + class RoPESM90Op: """Custom CUDA RoPE op for SM90 (GPT-NeoX rotate-half), differentiable w.r.t. ``x``. @@ -90,7 +117,6 @@ class RoPESM90Op: """ op_class = "elementwise" - backend_id = "rlkernel.cuda.rope_sm90" def __init__(self) -> None: if not _EXT_AVAILABLE or not hasattr(_C, "rope_apply_sm90"): @@ -106,4 +132,9 @@ def __call__(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) def forward(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: if x.device.type != "cuda": raise RuntimeError(f"RoPESM90Op requires a CUDA tensor, got device '{x.device}'.") + if not _is_hopper(x.device): + raise RuntimeError( + "RoPESM90Op requires Hopper (SM90) CUDA; " + f"got compute capability {torch.cuda.get_device_capability(x.device)}" + ) return _RoPEFunction.apply(x, positions, theta) diff --git a/rl_engine/kernels/ops/vjp_fp32.py b/rl_engine/kernels/ops/vjp_fp32.py new file mode 100644 index 00000000..65acdfa6 --- /dev/null +++ b/rl_engine/kernels/ops/vjp_fp32.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Declared row-local FP32 VJPs. No batched torch.matmul / cuBLAS. + +Each output row is an independent GEMV or outer product. Parameter reductions +walk rows in the caller's order so C10 can re-aggregate by logical token. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import torch + +BACKWARD_IMPL = "row_local_fp32_vjp" + + +def row_local_linear_dx_fp32(grad_output: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """dX[t] = grad[t] @ weight, one GEMV per row.""" + + rows = grad_output.reshape(-1, grad_output.size(-1)).float() + weight_f = weight.float() + out_rows = torch.empty( + (rows.shape[0], weight_f.shape[1]), device=rows.device, dtype=torch.float32 + ) + weight_t = weight_f.t().contiguous() + for index in range(rows.shape[0]): + out_rows[index] = torch.mv(weight_t, rows[index]) + return out_rows.reshape(*grad_output.shape[:-1], weight_f.shape[1]) + + +def row_local_linear_dw_fp32(grad_output: torch.Tensor, hidden: torch.Tensor) -> torch.Tensor: + """dW = sum_t outer(grad[t], hidden[t]) in physical row order.""" + + grad_rows = grad_output.reshape(-1, grad_output.size(-1)).float() + hidden_rows = hidden.reshape(-1, hidden.size(-1)).float() + if grad_rows.shape[0] != hidden_rows.shape[0]: + raise ValueError(f"grad rows {grad_rows.shape[0]} != hidden rows {hidden_rows.shape[0]}") + dweight = torch.zeros( + (grad_rows.shape[1], hidden_rows.shape[1]), + device=grad_rows.device, + dtype=torch.float32, + ) + for index in range(grad_rows.shape[0]): + dweight.addmm_(grad_rows[index].unsqueeze(1), hidden_rows[index].unsqueeze(0)) + return dweight + + +def row_local_bias_fp32(grad_output: torch.Tensor) -> torch.Tensor: + rows = grad_output.reshape(-1, grad_output.size(-1)).float() + acc = torch.zeros((rows.shape[1],), device=rows.device, dtype=torch.float32) + for index in range(rows.shape[0]): + acc = acc + rows[index] + return acc + + +def rmsnorm_dweight_rows_fp32( + x: torch.Tensor, + grad_output: torch.Tensor, + *, + rstd: torch.Tensor | None = None, + eps: float = 1e-6, +) -> torch.Tensor: + """Per-row dweight contributions, shape [..., H].""" + + x32 = x.float() + grad32 = grad_output.float() + if rstd is None: + rstd = torch.rsqrt(x32.square().mean(dim=-1) + float(eps)) + else: + rstd = rstd.float() + return grad32 * x32 * rstd.unsqueeze(-1) + + +def reduce_rows_fp32(rows: torch.Tensor) -> torch.Tensor: + """Left-fold dim 0 in FP32. Deterministic for a fixed row order.""" + + flat = rows.reshape(rows.shape[0], -1).float() + acc = torch.zeros((flat.shape[1],), device=flat.device, dtype=torch.float32) + for index in range(flat.shape[0]): + acc = acc + flat[index] + return acc.reshape(rows.shape[1:]) + + +def reduce_keyed_rows_fp32( + contributions: Mapping[tuple[str, int], torch.Tensor], +) -> torch.Tensor: + if not contributions: + raise RuntimeError("no logical-token contributions to reduce") + keys = sorted(contributions) + acc = contributions[keys[0]].float().clone() + for key in keys[1:]: + acc = acc + contributions[key].float() + return acc + + +def reduce_keyed_outers_fp32( + rows_g: Mapping[tuple[str, int], torch.Tensor], + rows_x: Mapping[tuple[str, int], torch.Tensor], +) -> torch.Tensor: + keys = sorted(set(rows_g) | set(rows_x)) + if not keys or set(rows_g) != set(rows_x): + raise RuntimeError("logical-token sets for outer-product VJP do not match") + first_g = rows_g[keys[0]].float() + first_x = rows_x[keys[0]].float() + acc = torch.outer(first_g, first_x) + for key in keys[1:]: + acc = acc + torch.outer(rows_g[key].float(), rows_x[key].float()) + return acc + + +def merge_keyed( + target: dict[tuple[str, int], torch.Tensor], + source: Mapping[tuple[str, int], torch.Tensor], +) -> None: + overlap = set(target) & set(source) + if overlap: + raise RuntimeError(f"logical token collision: {sorted(overlap)[:4]}") + target.update(source) + + +__all__ = [ + "BACKWARD_IMPL", + "merge_keyed", + "reduce_keyed_outers_fp32", + "reduce_keyed_rows_fp32", + "reduce_rows_fp32", + "rmsnorm_dweight_rows_fp32", + "row_local_bias_fp32", + "row_local_linear_dw_fp32", + "row_local_linear_dx_fp32", +] From 2e5cc49ff200e5a3d1ac306d5bf7ade307b01c0d Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Tue, 18 Aug 2026 10:47:50 +0000 Subject: [PATCH 14/17] feat(attention): add unified ablation matrix wrapper --- .../kernels/ops/pytorch/attention/__init__.py | 13 +- .../kernels/ops/pytorch/attention/ablation.py | 517 ++++++++++++++++++ rl_engine/kernels/registry.py | 38 ++ tests/test_attention_ablation.py | 144 +++++ 4 files changed, 711 insertions(+), 1 deletion(-) create mode 100644 rl_engine/kernels/ops/pytorch/attention/ablation.py create mode 100644 tests/test_attention_ablation.py diff --git a/rl_engine/kernels/ops/pytorch/attention/__init__.py b/rl_engine/kernels/ops/pytorch/attention/__init__.py index d2454e67..977ab14c 100644 --- a/rl_engine/kernels/ops/pytorch/attention/__init__.py +++ b/rl_engine/kernels/ops/pytorch/attention/__init__.py @@ -4,6 +4,12 @@ import torch import torch.nn.functional as F +from rl_engine.kernels.ops.pytorch.attention.ablation import ( + AttentionAblationConfig, + AttentionAblationOp, + AttentionAblationResult, +) + class NativeAttentionOp: """PyTorch SDPA fallback for FlashAttention-layout tensors.""" @@ -46,4 +52,9 @@ def __call__( return out.transpose(1, 2) -__all__ = ["NativeAttentionOp"] +__all__ = [ + "AttentionAblationConfig", + "AttentionAblationOp", + "AttentionAblationResult", + "NativeAttentionOp", +] diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py new file mode 100644 index 00000000..7d0d6dad --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -0,0 +1,517 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Unified Attention entry point for the PR230 cross-configuration matrix. + +The matrix needs one stable callable shape even though training and rollout may +materialize different Attention backends. This adapter owns the common +contract checks and provenance only; numerical work remains in the existing +deterministic Attention implementations. +""" + +from __future__ import annotations + +import hashlib +import inspect +import json +import math +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import Any, Callable, Mapping + +import torch +from torch import Tensor + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + AttentionContract, + AttentionContractError, + AttentionDType, + SplitKVMode, +) + +BACKEND_ID = "rlkernel.attention.deterministic.v1" +REFERENCE_BACKEND_ID = "rlkernel.attention.reference.v1" + +_TORCH_DTYPES = { + AttentionDType.BF16: torch.bfloat16, + AttentionDType.FP16: torch.float16, + AttentionDType.FP32: torch.float32, +} + + +@dataclass(frozen=True) +class AttentionAblationConfig: + """Per-invocation settings materialized by the ablation runner.""" + + backend: str = "auto" + deterministic: bool = True + communication_backend: str = "none" + return_lse: bool = True + return_gradients: bool = False + strict_core_id: str = STRICT_ATTENTION_CORE_ID + validate: bool = True + + def __post_init__(self) -> None: + if not isinstance(self.backend, str) or not self.backend.strip(): + raise AttentionContractError("Attention backend must be a non-empty string") + if not isinstance(self.deterministic, bool): + raise AttentionContractError("deterministic must be a bool") + if not isinstance(self.communication_backend, str) or not self.communication_backend.strip(): + raise AttentionContractError("communication_backend must be a non-empty string") + for name in ("return_lse", "return_gradients", "validate"): + if not isinstance(getattr(self, name), bool): + raise AttentionContractError(f"{name} must be a bool") + if not isinstance(self.strict_core_id, str) or not self.strict_core_id.strip(): + raise AttentionContractError("strict_core_id must be a non-empty string") + object.__setattr__(self, "backend", self.backend.strip().lower()) + object.__setattr__(self, "communication_backend", self.communication_backend.strip()) + + +@dataclass(frozen=True) +class AttentionAblationResult: + """Standardized Attention result consumed by cross-config artifacts.""" + + out: Tensor + lse: Tensor | None + dq: Tensor | None = None + dk: Tensor | None = None + dv: Tensor | None = None + backend_id: str = BACKEND_ID + deterministic: bool = True + provenance: Mapping[str, Any] = field(default_factory=dict) + + def __post_init__(self) -> None: + if not isinstance(self.out, Tensor): + raise TypeError("Attention result out must be a torch.Tensor") + if self.lse is not None and not isinstance(self.lse, Tensor): + raise TypeError("Attention result lse must be a torch.Tensor or None") + for name in ("dq", "dk", "dv"): + value = getattr(self, name) + if value is not None and not isinstance(value, Tensor): + raise TypeError(f"Attention result {name} must be a torch.Tensor or None") + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise ValueError("Attention result backend_id must be non-empty") + if not isinstance(self.deterministic, bool): + raise TypeError("Attention result deterministic must be a bool") + if not isinstance(self.provenance, Mapping): + raise TypeError("Attention result provenance must be a mapping") + object.__setattr__(self, "provenance", MappingProxyType(dict(self.provenance))) + + @property + def out_lse(self) -> tuple[Tensor, Tensor | None]: + """Compatibility tuple for callers that consume ``(out, lse)``.""" + + return self.out, self.lse + + def readback(self) -> dict[str, Any]: + """Return JSON-compatible execution evidence for PR230 artifacts.""" + + return { + "backend_id": self.backend_id, + "deterministic": self.deterministic, + "out_shape": list(self.out.shape), + "out_dtype": str(self.out.dtype).replace("torch.", ""), + "lse_shape": None if self.lse is None else list(self.lse.shape), + "lse_dtype": None + if self.lse is None + else str(self.lse.dtype).replace("torch.", ""), + "gradients": { + "dq": self.dq is not None, + "dk": self.dk is not None, + "dv": self.dv is not None, + }, + "provenance": dict(self.provenance), + } + + +class AttentionAblationOp: + """PR230/PR314-style unified Attention wrapper. + + ``core`` and ``reference`` are injectable so the wrapper is usable by the + semantic operator session without importing CUDA at construction time. + ``core`` should expose ``forward_with_lse``; ``reference`` is the existing + pure-PyTorch CP reference with the same method. + """ + + op_class = "attention" + is_batch_invariant = True + backend_id = BACKEND_ID + + def __init__( + self, + *, + core: Any | None = None, + reference: Any | None = None, + native: Any | None = None, + communication_backend: str = "none", + ) -> None: + if not isinstance(communication_backend, str) or not communication_backend.strip(): + raise AttentionContractError("communication_backend must be a non-empty string") + self.core = core + self.reference = reference + self.native = native + self.communication_backend = communication_backend.strip() + + def __call__( + self, + q: Tensor, + k: Tensor, + v: Tensor, + *, + contract: AttentionContract, + config: AttentionAblationConfig | Mapping[str, Any] | None = None, + backend: str | Callable[..., Any] | None = None, + deterministic: bool | None = None, + return_lse: bool | None = None, + return_gradients: bool | None = None, + dout: Tensor | None = None, + communication_backend: str | None = None, + validate: bool | None = None, + **kwargs: Any, + ) -> AttentionAblationResult: + return self.apply( + q, + k, + v, + contract=contract, + config=config, + backend=backend, + deterministic=deterministic, + return_lse=return_lse, + return_gradients=return_gradients, + dout=dout, + communication_backend=communication_backend, + validate=validate, + **kwargs, + ) + + def apply( + self, + q: Tensor, + k: Tensor, + v: Tensor, + *, + contract: AttentionContract, + config: AttentionAblationConfig | Mapping[str, Any] | None = None, + backend: str | Callable[..., Any] | None = None, + deterministic: bool | None = None, + return_lse: bool | None = None, + return_gradients: bool | None = None, + dout: Tensor | None = None, + communication_backend: str | None = None, + validate: bool | None = None, + **kwargs: Any, + ) -> AttentionAblationResult: + if not isinstance(contract, AttentionContract): + raise AttentionContractError("contract must be an AttentionContract") + backend_request = ( + backend + if callable(backend) or hasattr(backend, "forward_with_lse") or hasattr(backend, "apply") + else None + ) + cfg = _resolve_config( + config, + backend=backend if isinstance(backend, str) else None, + deterministic=deterministic, + return_lse=return_lse, + return_gradients=return_gradients, + communication_backend=( + communication_backend + if communication_backend is not None + else self.communication_backend + ), + validate=validate, + ) + if cfg.validate: + self._validate_inputs(q, k, v, contract) + if cfg.deterministic and contract.split_kv.mode is SplitKVMode.AUTO: + raise AttentionContractError( + "deterministic Attention cannot use runtime-dependent Split-KV=auto" + ) + if cfg.return_gradients and dout is None: + raise AttentionContractError("dout is required when return_gradients=True") + if dout is not None and dout.shape != q.shape: + raise AttentionContractError("dout must have the same shape as q") + + requested = backend_request if backend_request is not None else cfg.backend + selected, selected_id = self._select_backend(requested, q, contract) + if cfg.deterministic and selected_id == "native": + raise AttentionContractError( + "deterministic=True cannot execute an unverified native Attention backend" + ) + + call_kwargs = dict(kwargs) + call_kwargs.setdefault("causal", contract.causal) + call_kwargs.setdefault("scale", 1.0 / math.sqrt(contract.head_dim)) + call_kwargs.setdefault("cp_world_size", contract.sharding.cp_world_size) + if contract.split_kv.mode is SplitKVMode.FIXED: + call_kwargs.setdefault("kv_chunk_size", contract.split_kv.fixed_split_size) + out, lse = self._invoke(selected, q, k, v, call_kwargs, contract) + if cfg.validate: + self._validate_outputs(out, lse, q, contract) + + dq = dk = dv = None + if cfg.return_gradients: + dq, dk, dv = self._backward(selected, q, k, v, out, dout, call_kwargs) + + provenance = { + "schema_version": "rlkernel.attention.ablation_result.v1", + "semantic_operator": "attention", + "backend_id": selected_id, + "deterministic": cfg.deterministic, + "strict_core_id": cfg.strict_core_id if cfg.deterministic else None, + "communication_backend": cfg.communication_backend, + "split_kv": contract.split_kv.to_dict(), + "actual_split_kv": _actual_split_provenance(contract), + "reduction": _reduction_provenance(contract), + "contract_fingerprint": _contract_fingerprint(contract), + "return_lse": cfg.return_lse, + "return_gradients": cfg.return_gradients, + } + return AttentionAblationResult( + out=out, + lse=lse if cfg.return_lse else None, + dq=dq, + dk=dk, + dv=dv, + backend_id=selected_id, + deterministic=cfg.deterministic, + provenance=provenance, + ) + + def apply_fp32(self, *args: Any, **kwargs: Any) -> AttentionAblationResult: + """Stable fingerprint entry point used by ``OperatorSession``.""" + + return self.apply(*args, **kwargs) + + def _select_backend( + self, + requested: str | Callable[..., Any], + q: Tensor, + contract: AttentionContract, + ) -> tuple[Any, str]: + if callable(requested) or hasattr(requested, "forward_with_lse") or hasattr(requested, "apply"): + return requested, _callable_backend_id(requested) + normalized = str(requested).strip().lower() + if normalized in {"native", "te", "flashinfer"}: + if self.native is None: + raise AttentionContractError( + "native Attention backend was requested but no native callable was injected" + ) + return self.native, "native" + if normalized in {"reference", "pytorch_reference"}: + return self._reference_backend(), REFERENCE_BACKEND_ID + if normalized not in {"auto", "deterministic", "rlkernel"}: + raise AttentionContractError(f"unsupported Attention backend {requested!r}") + if ( + contract.sharding.cp_world_size > 1 + or q.device.type != "cuda" + or torch.version.hip is not None + ): + return self._reference_backend(), REFERENCE_BACKEND_ID + if self.core is None: + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionOp, + ) + + self.core = DeterministicAttentionOp() + return self.core, BACKEND_ID + + def _reference_backend(self) -> Any: + if self.reference is None: + from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + DeterministicCPAttentionReferenceOp, + ) + + self.reference = DeterministicCPAttentionReferenceOp() + return self.reference + + @staticmethod + def _validate_inputs(q: Tensor, k: Tensor, v: Tensor, contract: AttentionContract) -> None: + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise AttentionContractError("q, k, and v must use [B, H, S, D] layout") + expected_dtype = _TORCH_DTYPES[contract.dtype] + if q.dtype is not expected_dtype or k.dtype is not expected_dtype or v.dtype is not expected_dtype: + raise AttentionContractError( + f"q, k, and v must match contract dtype {contract.dtype.value}" + ) + if q.device != k.device or q.device != v.device: + raise AttentionContractError("q, k, and v must be on the same device") + batch, q_heads, q_seq, dim = q.shape + if batch != contract.batch_size: + raise AttentionContractError( + f"q batch={batch} does not match contract batch_size={contract.batch_size}" + ) + if q_seq != contract.query_sequence_length: + raise AttentionContractError( + "q sequence length does not match AttentionContract query_sequence_length" + ) + sharding = contract.sharding + if q_heads != sharding.local_q_heads or k.shape[1] != sharding.local_kv_heads: + raise AttentionContractError("q/k head counts do not match TP sharding in contract") + if k.shape[0] != batch or v.shape[:3] != k.shape[:3] or k.shape[-1] != dim or v.shape[-1] != dim: + raise AttentionContractError("q, k, and v shapes are inconsistent") + if dim != contract.head_dim: + raise AttentionContractError("tensor head_dim does not match AttentionContract") + + @staticmethod + def _validate_outputs(out: Tensor, lse: Tensor, q: Tensor, contract: AttentionContract) -> None: + if out.shape != q.shape: + raise AttentionContractError( + f"Attention output shape {tuple(out.shape)} does not match q {tuple(q.shape)}" + ) + expected_lse = (q.shape[0], q.shape[1], q.shape[2]) + if lse.shape != expected_lse: + raise AttentionContractError( + f"attention-domain LSE shape {tuple(lse.shape)} does not match {expected_lse}" + ) + if lse.dtype is not torch.float32: + raise AttentionContractError("attention-domain LSE must remain fp32") + expected_dtype = _TORCH_DTYPES[contract.dtype] + if out.dtype is not expected_dtype: + raise AttentionContractError( + f"Attention output must be written in {contract.dtype.value}, got {out.dtype}" + ) + + @staticmethod + def _invoke( + backend: Any, + q: Tensor, + k: Tensor, + v: Tensor, + kwargs: Mapping[str, Any], + contract: AttentionContract, + ) -> tuple[Tensor, Tensor]: + method = getattr(backend, "forward_with_lse", None) + if not callable(method): + method = getattr(backend, "apply", None) + if not callable(method): + method = backend if callable(backend) else None + if method is None: + raise AttentionContractError( + "Attention backend must expose forward_with_lse, apply, or __call__" + ) + accepted = _accepted_kwargs(method, kwargs) + if contract.sharding.cp_world_size > 1 and "cp_world_size" not in accepted: + raise AttentionContractError( + "CP>1 requires an Attention backend that explicitly accepts cp_world_size" + ) + result = method(q, k, v, **accepted) + if isinstance(result, AttentionAblationResult): + out, lse = result.out, result.lse + elif isinstance(result, tuple) and len(result) == 2: + out, lse = result + else: + raise AttentionContractError( + "Attention backend must return (out, lse) or AttentionAblationResult" + ) + if not isinstance(out, Tensor) or not isinstance(lse, Tensor): + raise AttentionContractError("Attention backend returned non-tensor output or LSE") + return out, lse + + @staticmethod + def _backward( + backend: Any, + q: Tensor, + k: Tensor, + v: Tensor, + out: Tensor, + dout: Tensor | None, + kwargs: Mapping[str, Any], + ) -> tuple[Tensor, Tensor, Tensor]: + backward = getattr(backend, "backward_reference", None) + if callable(backward): + result = backward(q, k, v, dout, **_accepted_kwargs(backward, kwargs)) + gradients = getattr(result, "gradients", None) + if gradients is not None: + return gradients.dq, gradients.dk, gradients.dv + if not out.requires_grad: + raise AttentionContractError( + "Attention backend did not retain an autograd graph for gradients" + ) + return torch.autograd.grad( + out, + (q, k, v), + grad_outputs=dout.to(dtype=out.dtype), + allow_unused=False, + retain_graph=True, + ) + + +def _resolve_config( + config: AttentionAblationConfig | Mapping[str, Any] | None, + **overrides: Any, +) -> AttentionAblationConfig: + if config is None: + values: dict[str, Any] = {} + elif isinstance(config, AttentionAblationConfig): + values = { + name: getattr(config, name) + for name in ( + "backend", + "deterministic", + "communication_backend", + "return_lse", + "return_gradients", + "strict_core_id", + "validate", + ) + } + elif isinstance(config, Mapping): + values = dict(config) + else: + raise AttentionContractError("config must be AttentionAblationConfig, mapping, or None") + values.update({name: value for name, value in overrides.items() if value is not None}) + return AttentionAblationConfig(**values) + + +def _accepted_kwargs(method: Callable[..., Any], kwargs: Mapping[str, Any]) -> dict[str, Any]: + try: + signature = inspect.signature(method) + except (TypeError, ValueError): + return dict(kwargs) + parameters = signature.parameters.values() + if any(parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in parameters): + return dict(kwargs) + return {name: value for name, value in kwargs.items() if name in signature.parameters} + + +def _callable_backend_id(value: Any) -> str: + explicit = getattr(value, "backend_id", None) or getattr(value, "__attention_backend_id__", None) + if isinstance(explicit, str) and explicit.strip(): + return explicit.strip() + return f"injected.{type(value).__module__}.{type(value).__qualname__}" + + +def _actual_split_provenance(contract: AttentionContract) -> dict[str, Any]: + plan = contract.split_kv.resolve( + contract.sharding.global_sequence_length, + backend=BACKEND_ID, + ) + return plan.to_dict() + + +def _reduction_provenance(contract: AttentionContract) -> dict[str, Any]: + reduction = contract.reduction + return { + "merge": reduction.merge.value, + "acc_dtype": reduction.acc_dtype.value, + "order": reduction.order.value, + "downcast_at": reduction.downcast_at.value, + "engine": reduction.engine.value, + } + + +def _contract_fingerprint(contract: AttentionContract) -> str: + payload = json.dumps(contract.to_dict(), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +__all__ = [ + "AttentionAblationConfig", + "AttentionAblationOp", + "AttentionAblationResult", + "BACKEND_ID", + "REFERENCE_BACKEND_ID", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 060234be..d9b4541c 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -185,6 +185,44 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, version_or_build_fingerprint="runtime-native-unresolved-v1", ), + OperatorBackendDescriptor( + semantic_op="attention", + backend_id="rlkernel.attention.deterministic.v1", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu", "cuda", "rocm"}), + supported_dtypes=frozenset({"float32", "bfloat16", "float16"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "algorithm": "standard_softmax_attention", + "batch_invariant": True, + "deterministic": True, + "split_kv": "contract_bound", + "reduction_order": "global_block_index", + "strict_observable": True, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=( + "rl_engine.kernels.ops.pytorch.attention.ablation.AttentionAblationOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="AttentionAblationOp-v1", + ), + OperatorBackendDescriptor( + semantic_op="attention", + backend_id="native", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"*"}), + supported_dtypes=frozenset({"*"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={ + "selection": "runtime_native", + "strict_observable": False, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=None, + fallback_policy=OperatorFallbackPolicy.RUNTIME_MANAGED, + version_or_build_fingerprint="runtime-native-attention-unresolved-v1", + ), ) diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py new file mode 100644 index 00000000..6d26f15e --- /dev/null +++ b/tests/test_attention_ablation.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.attention_contract import ( + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMode, + AttentionRole, + ReductionSpec, + ShardingSpec, + SplitKVSpec, +) +from rl_engine.kernels.ops.pytorch.attention.ablation import ( + BACKEND_ID, + REFERENCE_BACKEND_ID, + AttentionAblationOp, +) + + +def _contract(*, split_kv: SplitKVSpec | None = None) -> AttentionContract: + sharding = ShardingSpec( + tp_rank=0, + tp_world_size=1, + cp_rank=0, + cp_world_size=1, + global_q_heads=2, + global_kv_heads=1, + local_q_head_start=0, + local_q_heads=2, + local_kv_head_start=0, + local_kv_heads=1, + global_sequence_length=4, + local_sequence_length=4, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 4), + ) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=4, + head_dim=4, + causal=True, + causal_offsets=(0,), + sharding=sharding, + reduction=ReductionSpec(), + split_kv=split_kv or SplitKVSpec.disabled(), + ) + + +def _qkv(): + torch.manual_seed(0) + return ( + torch.randn(1, 2, 4, 4, dtype=torch.bfloat16), + torch.randn(1, 1, 4, 4, dtype=torch.bfloat16), + torch.randn(1, 1, 4, 4, dtype=torch.bfloat16), + ) + + +def test_attention_wrapper_has_unified_result_and_provenance(): + q, k, v = _qkv() + result = AttentionAblationOp()(q, k, v, contract=_contract()) + + assert result.backend_id == REFERENCE_BACKEND_ID + assert result.deterministic + assert result.out.shape == q.shape + assert result.lse is not None + assert result.lse.dtype is torch.float32 + assert result.provenance["semantic_operator"] == "attention" + assert result.provenance["split_kv"]["mode"] == "disabled" + assert result.readback()["out_shape"] == list(q.shape) + + +def test_attention_wrapper_supports_explicit_injected_backend(): + q, k, v = _qkv() + + class FakeBackend: + backend_id = "test.attention.backend" + + def forward_with_lse(self, q, k, v, *, causal, scale): + del k, v, causal, scale + return q.clone(), torch.zeros(q.shape[:3], dtype=torch.float32) + + result = AttentionAblationOp()(q, k, v, contract=_contract(), backend=FakeBackend()) + assert result.backend_id == "test.attention.backend" + assert torch.equal(result.out, q) + + +def test_deterministic_attention_rejects_runtime_split_kv_auto(): + q, k, v = _qkv() + contract = _contract(split_kv=SplitKVSpec.auto(strict_consistency=False)) + with pytest.raises(AttentionContractError, match="Split-KV=auto"): + AttentionAblationOp()(q, k, v, contract=contract) + + +def test_deterministic_native_backend_requires_explicit_native_callable(): + q, k, v = _qkv() + with pytest.raises(AttentionContractError, match="native Attention backend"): + AttentionAblationOp()(q, k, v, contract=_contract(), backend="native") + + +def test_attention_wrapper_can_return_dq_dk_dv_from_reference_backend(): + q, k, v = _qkv() + result = AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + return_gradients=True, + dout=torch.ones_like(q), + ) + + assert result.dq is not None and result.dq.shape == q.shape + assert result.dk is not None and result.dk.shape == k.shape + assert result.dv is not None and result.dv.shape == v.shape + + +def test_attention_backend_is_registered_for_pr230_semantic_resolution(): + from rl_engine.kernels.registry import kernel_registry + from rl_engine.kernels.semantic_registry import OperatorRequirements + + session = kernel_registry.semantic.session() + resolution = session.resolve( + semantic_op="attention", + requested_backend=BACKEND_ID, + target="training", + requirements=OperatorRequirements( + device="cpu", + dtype="bfloat16", + topology={"world_size": 1, "tensor_parallel_size": 1, "context_parallel_size": 1}, + alignment_properties={"deterministic": True}, + ), + ) + instance = session.instantiate(resolution) + assert isinstance(instance, AttentionAblationOp) + provenance = session.instance_provenance(resolution, instance) + assert provenance.backend_id == BACKEND_ID From a3c3d1268f47a9c2accf4793124a2b9ce8d91603 Mon Sep 17 00:00:00 2001 From: Codex H100 Integration Date: Tue, 18 Aug 2026 10:52:29 +0000 Subject: [PATCH 15/17] fix(attention): report selected core provenance --- .../kernels/ops/pytorch/attention/ablation.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index 7d0d6dad..0b0b5106 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -260,10 +260,16 @@ def apply( "semantic_operator": "attention", "backend_id": selected_id, "deterministic": cfg.deterministic, - "strict_core_id": cfg.strict_core_id if cfg.deterministic else None, + "strict_core_id": ( + cfg.strict_core_id + if cfg.deterministic and selected_id == BACKEND_ID + else None + ), + "core_id": selected_id, + "backend_deterministic": selected_id in {BACKEND_ID, REFERENCE_BACKEND_ID}, "communication_backend": cfg.communication_backend, "split_kv": contract.split_kv.to_dict(), - "actual_split_kv": _actual_split_provenance(contract), + "actual_split_kv": _actual_split_provenance(contract, backend=selected_id), "reduction": _reduction_provenance(contract), "contract_fingerprint": _contract_fingerprint(contract), "return_lse": cfg.return_lse, @@ -484,10 +490,14 @@ def _callable_backend_id(value: Any) -> str: return f"injected.{type(value).__module__}.{type(value).__qualname__}" -def _actual_split_provenance(contract: AttentionContract) -> dict[str, Any]: +def _actual_split_provenance( + contract: AttentionContract, + *, + backend: str = BACKEND_ID, +) -> dict[str, Any]: plan = contract.split_kv.resolve( contract.sharding.global_sequence_length, - backend=BACKEND_ID, + backend=backend, ) return plan.to_dict() From 43e67c5a37df7e2d4991f717d9b282f255275bc5 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 18 Aug 2026 21:47:09 +0800 Subject: [PATCH 16/17] feat(attention): enforce canonical bitwise ablation core --- .../cross_config/attention_binding.py | 33 +++++ rl_engine/kernels/attention_contract.py | 2 + .../kernels/ops/pytorch/attention/ablation.py | 46 ++++++- .../ops/pytorch/attention/cp_attention.py | 127 ++++++++++++++++-- rl_engine/kernels/registry.py | 3 +- tests/test_attention_ablation.py | 41 +++++- tests/test_attention_cross_config_binding.py | 6 + tests/test_cp_attention.py | 26 ++++ 8 files changed, 264 insertions(+), 20 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index d5f53df4..57a54205 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -45,6 +45,7 @@ from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, AttentionContract, AttentionContractError, AttentionDType, @@ -125,6 +126,7 @@ class BindingErrorCode(str, Enum): ATTENTION_PROJECTION_MISMATCH = "ATTENTION_PROJECTION_MISMATCH" ATTENTION_CORE_MISSING = "ATTENTION_CORE_MISSING" ATTENTION_CORE_MISMATCH = "ATTENTION_CORE_MISMATCH" + ATTENTION_CORE_SCHEDULE = "ATTENTION_CORE_SCHEDULE" ATTENTION_NATIVE_ARITHMETIC = "ATTENTION_NATIVE_ARITHMETIC" ATTENTION_CORE_SPLIT_K = "ATTENTION_CORE_SPLIT_K" @@ -146,6 +148,7 @@ class AttentionRuntimeReadback: projection_plans: Mapping[str, Mapping[str, Any]] = field(default_factory=dict) strict_mode: bool = False strict_core_id: str | None = None + strict_schedule: str | None = None native_attention_arithmetic: bool = True strict_split_kv_policy: str | None = None @@ -185,6 +188,10 @@ def __post_init__(self) -> None: not isinstance(self.strict_core_id, str) or not self.strict_core_id.strip() ): raise ValueError("strict_core_id must be a non-empty string when provided") + if self.strict_schedule is not None and ( + not isinstance(self.strict_schedule, str) or not self.strict_schedule.strip() + ): + raise ValueError("strict_schedule must be a non-empty string when provided") if not isinstance(self.native_attention_arithmetic, bool): raise TypeError("native_attention_arithmetic must be a bool") if self.strict_split_kv_policy is not None and self.strict_split_kv_policy not in { @@ -237,6 +244,7 @@ def to_dict(self) -> dict[str, Any]: "strict_attention": { "enabled": self.strict_mode, "core_id": self.strict_core_id, + "schedule": self.strict_schedule, "native_attention_arithmetic": self.native_attention_arithmetic, "split_kv_policy": self.strict_split_kv_policy, }, @@ -1020,6 +1028,7 @@ def bind_attention_runtime_readbacks( "preprocess.policy_id": rollout.preprocess_policy_id, "strict.enabled": rollout.strict_mode, "strict.core_id": rollout.strict_core_id, + "strict.schedule": rollout.strict_schedule, "strict.native_attention_arithmetic": rollout.native_attention_arithmetic, "strict.split_kv_policy": rollout.strict_split_kv_policy, **{ @@ -1038,6 +1047,7 @@ def bind_attention_runtime_readbacks( "preprocess.policy_id": training.preprocess_policy_id, "strict.enabled": training.strict_mode, "strict.core_id": training.strict_core_id, + "strict.schedule": training.strict_schedule, "strict.native_attention_arithmetic": training.native_attention_arithmetic, "strict.split_kv_policy": training.strict_split_kv_policy, **{ @@ -1067,6 +1077,18 @@ def _strict_attention_core_issues( ), ) ) + if readback.strict_schedule != STRICT_ATTENTION_SCHEDULE_ID: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_SCHEDULE, + tier=BindingTier.SEMANTIC, + field=f"{side}.strict.schedule", + message=( + f"{side} strict Attention did not execute the canonical schedule " + f"{STRICT_ATTENTION_SCHEDULE_ID!r}" + ), + ) + ) if readback.native_attention_arithmetic: issues.append( BindingIssue( @@ -1120,6 +1142,17 @@ def _strict_attention_core_pair_issues( message="training and rollout executed different Attention cores", ) ] + if rollout.strict_mode and rollout.strict_schedule != training.strict_schedule: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_SCHEDULE, + tier=BindingTier.SEMANTIC, + field="strict.schedule", + rollout=rollout.strict_schedule, + training=training.strict_schedule, + message="training and rollout executed different strict Attention schedules", + ) + ] return [] diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index b4d85f7a..8b5f588b 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -19,6 +19,7 @@ STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" +STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks" class AttentionContractError(ValueError): @@ -1525,6 +1526,7 @@ class AttentionDispatchResult: "SplitKVRuntimePlanSet", "SplitKVSpec", "STRICT_ATTENTION_CORE_ID", + "STRICT_ATTENTION_SCHEDULE_ID", "validate_split_kv_alignment", "validate_split_kv_plan_set_alignment", ] diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index 0b0b5106..de1a8d5c 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -24,6 +24,7 @@ from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, AttentionContract, AttentionContractError, AttentionDType, @@ -50,6 +51,7 @@ class AttentionAblationConfig: return_lse: bool = True return_gradients: bool = False strict_core_id: str = STRICT_ATTENTION_CORE_ID + strict_schedule: str = STRICT_ATTENTION_SCHEDULE_ID validate: bool = True def __post_init__(self) -> None: @@ -64,6 +66,8 @@ def __post_init__(self) -> None: raise AttentionContractError(f"{name} must be a bool") if not isinstance(self.strict_core_id, str) or not self.strict_core_id.strip(): raise AttentionContractError("strict_core_id must be a non-empty string") + if not isinstance(self.strict_schedule, str) or not self.strict_schedule.strip(): + raise AttentionContractError("strict_schedule must be a non-empty string") object.__setattr__(self, "backend", self.backend.strip().lower()) object.__setattr__(self, "communication_backend", self.communication_backend.strip()) @@ -229,6 +233,15 @@ def apply( raise AttentionContractError( "deterministic Attention cannot use runtime-dependent Split-KV=auto" ) + if cfg.deterministic and ( + cfg.strict_core_id != STRICT_ATTENTION_CORE_ID + or cfg.strict_schedule != STRICT_ATTENTION_SCHEDULE_ID + ): + raise AttentionContractError( + "strict deterministic Attention requires the canonical core and schedule" + ) + if cfg.deterministic and not cfg.return_lse: + raise AttentionContractError("strict deterministic Attention must return LSE") if cfg.return_gradients and dout is None: raise AttentionContractError("dout is required when return_gradients=True") if dout is not None and dout.shape != q.shape: @@ -240,6 +253,16 @@ def apply( raise AttentionContractError( "deterministic=True cannot execute an unverified native Attention backend" ) + selected_core_id = getattr(selected, "core_id", None) + selected_schedule = getattr(selected, "strict_schedule", None) + if cfg.deterministic and selected_id not in {BACKEND_ID, REFERENCE_BACKEND_ID}: + if ( + selected_core_id != cfg.strict_core_id + or selected_schedule != cfg.strict_schedule + ): + raise AttentionContractError( + "deterministic Attention requires the shared strict core and schedule" + ) call_kwargs = dict(kwargs) call_kwargs.setdefault("causal", contract.causal) @@ -262,14 +285,23 @@ def apply( "deterministic": cfg.deterministic, "strict_core_id": ( cfg.strict_core_id - if cfg.deterministic and selected_id == BACKEND_ID + if cfg.deterministic else None ), - "core_id": selected_id, - "backend_deterministic": selected_id in {BACKEND_ID, REFERENCE_BACKEND_ID}, + "strict_schedule": cfg.strict_schedule if cfg.deterministic else None, + "core_id": cfg.strict_core_id if cfg.deterministic else selected_id, + "backend_deterministic": cfg.deterministic, + "native_attention_arithmetic": False if cfg.deterministic else selected_id == "native", "communication_backend": cfg.communication_backend, + "communication_executed": bool( + getattr(selected, "communication_executed", False) + ), "split_kv": contract.split_kv.to_dict(), - "actual_split_kv": _actual_split_provenance(contract, backend=selected_id), + "actual_split_kv": _actual_split_provenance( + contract, + total_kv_tokens=k.size(2), + backend=selected_id, + ), "reduction": _reduction_provenance(contract), "contract_fingerprint": _contract_fingerprint(contract), "return_lse": cfg.return_lse, @@ -330,7 +362,7 @@ def _reference_backend(self) -> Any: DeterministicCPAttentionReferenceOp, ) - self.reference = DeterministicCPAttentionReferenceOp() + self.reference = DeterministicCPAttentionReferenceOp(strict_bitwise=True) return self.reference @staticmethod @@ -461,6 +493,7 @@ def _resolve_config( "return_lse", "return_gradients", "strict_core_id", + "strict_schedule", "validate", ) } @@ -493,10 +526,11 @@ def _callable_backend_id(value: Any) -> str: def _actual_split_provenance( contract: AttentionContract, *, + total_kv_tokens: int, backend: str = BACKEND_ID, ) -> dict[str, Any]: plan = contract.split_kv.resolve( - contract.sharding.global_sequence_length, + total_kv_tokens, backend=backend, ) return plan.to_dict() diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 0062dac4..194c80d5 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -208,6 +208,11 @@ class DeterministicCPAttentionReferenceOp: op_class = "attention" + def __init__(self, *, strict_bitwise: bool = False) -> None: + if not isinstance(strict_bitwise, bool): + raise TypeError("strict_bitwise must be a bool") + self.strict_bitwise = strict_bitwise + @staticmethod def split_kv_execution_plans( total_kv_tokens: int, @@ -336,18 +341,31 @@ def forward_with_lse( resolved_output_dtype = q.dtype if output_dtype is None else output_dtype _validate_output_dtype(resolved_output_dtype) - out, lse = self._forward_impl( - q, - k, - v, - causal=causal, - scale=scale, - key_padding_mask=key_padding_mask, - query_position_offsets=query_position_offsets, - key_position_offsets=key_position_offsets, - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - ) + if self.strict_bitwise: + out, lse = self._forward_strict_bitwise( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + kv_chunk_size=kv_chunk_size, + ) + else: + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) out = out.to(resolved_output_dtype) return out, lse @@ -381,6 +399,91 @@ def forward_fp32_with_lse( output_dtype=torch.float32, ) + def _forward_strict_bitwise( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Execute one batch/CP-independent arithmetic schedule.""" + + _validate_qkv(q, k, v) + _validate_scale(scale) + batch, hq, sq, dim = q.shape + skv = k.size(2) + if key_padding_mask is not None: + if key_padding_mask.shape != (batch, skv): + raise ValueError("key_padding_mask must have shape [B, Skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=skv - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + kv_bounds = _kv_block_bounds(skv, 1, kv_chunk_size) + out_rows: list[torch.Tensor] = [] + lse_rows: list[torch.Tensor] = [] + for batch_index in range(batch): + q_batch = q[batch_index : batch_index + 1].contiguous() + k_batch = k[batch_index : batch_index + 1].contiguous() + v_batch = v[batch_index : batch_index + 1].contiguous() + pad_batch = ( + None + if key_padding_mask is None + else key_padding_mask[batch_index : batch_index + 1].contiguous() + ) + query_offset = query_offsets[batch_index : batch_index + 1] + key_offset = key_offsets[batch_index : batch_index + 1] + query_rows: list[torch.Tensor] = [] + lse_query_rows: list[torch.Tensor] = [] + for query_index in range(sq): + q_row = q_batch[:, :, query_index : query_index + 1, :].contiguous() + states = [ + self.local_partial_state( + q_row, + k_batch[:, :, key_start:key_end, :].contiguous(), + v_batch[:, :, key_start:key_end, :].contiguous(), + q_start=query_index, + k_start=key_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None + if pad_batch is None + else pad_batch[:, key_start:key_end].contiguous() + ), + query_position_offsets=query_offset, + key_position_offsets=key_offset, + ) + for key_start, key_end in kv_bounds + if key_start != key_end + ] + merged = merge_attention_partial_states(states) + query_rows.append(merged.out) + lse_query_rows.append(merged.lse) + out_rows.append(torch.cat(query_rows, dim=2)) + lse_rows.append(torch.cat(lse_query_rows, dim=2)) + return torch.cat(out_rows, dim=0), torch.cat(lse_rows, dim=0) + def backward_reference( self, q: torch.Tensor, diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index d9b4541c..edd1c535 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -198,6 +198,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "deterministic": True, "split_kv": "contract_bound", "reduction_order": "global_block_index", + "strict_schedule": "single_batch_single_query_global_kv_blocks", "strict_observable": True, }, lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, @@ -205,7 +206,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "rl_engine.kernels.ops.pytorch.attention.ablation.AttentionAblationOp" ), fallback_policy=OperatorFallbackPolicy.ERROR, - version_or_build_fingerprint="AttentionAblationOp-v1", + version_or_build_fingerprint="AttentionAblationOp-bitwise-v2", ), OperatorBackendDescriptor( semantic_op="attention", diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py index 6d26f15e..7023c6a5 100644 --- a/tests/test_attention_ablation.py +++ b/tests/test_attention_ablation.py @@ -6,6 +6,8 @@ import torch from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, AttentionContract, AttentionContractError, AttentionDType, @@ -75,6 +77,8 @@ def test_attention_wrapper_has_unified_result_and_provenance(): assert result.lse.dtype is torch.float32 assert result.provenance["semantic_operator"] == "attention" assert result.provenance["split_kv"]["mode"] == "disabled" + assert result.provenance["strict_core_id"] == STRICT_ATTENTION_CORE_ID + assert result.provenance["strict_schedule"] == STRICT_ATTENTION_SCHEDULE_ID assert result.readback()["out_shape"] == list(q.shape) @@ -88,7 +92,14 @@ def forward_with_lse(self, q, k, v, *, causal, scale): del k, v, causal, scale return q.clone(), torch.zeros(q.shape[:3], dtype=torch.float32) - result = AttentionAblationOp()(q, k, v, contract=_contract(), backend=FakeBackend()) + result = AttentionAblationOp()( + q, + k, + v, + contract=_contract(), + backend=FakeBackend(), + deterministic=False, + ) assert result.backend_id == "test.attention.backend" assert torch.equal(result.out, q) @@ -142,3 +153,31 @@ def test_attention_backend_is_registered_for_pr230_semantic_resolution(): assert isinstance(instance, AttentionAblationOp) provenance = session.instance_provenance(resolution, instance) assert provenance.backend_id == BACKEND_ID + + +def test_strict_wrapper_is_bitwise_invariant_to_batch_shape(): + q, k, v = _qkv() + noise_q, noise_k, noise_v = _qkv() + contract = _contract() + batch_contract = AttentionContract( + role=contract.role, + mode=contract.mode, + dtype=contract.dtype, + batch_size=2, + query_sequence_length=contract.query_sequence_length, + head_dim=contract.head_dim, + causal=contract.causal, + causal_offsets=(0, 0), + sharding=contract.sharding, + reduction=contract.reduction, + split_kv=contract.split_kv, + ) + single = AttentionAblationOp()(q, k, v, contract=contract) + batched = AttentionAblationOp()( + torch.cat((q, noise_q), dim=0), + torch.cat((k, noise_k), dim=0), + torch.cat((v, noise_v), dim=0), + contract=batch_contract, + ) + assert torch.equal(single.out[0], batched.out[0]) + assert torch.equal(single.lse[0], batched.lse[0]) diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index c5c7f2d4..282e5f39 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -46,6 +46,7 @@ from rl_engine.alignment.cross_config.schema import MaterializationStatus from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, AttentionContractError, AttentionMode, AttentionRole, @@ -665,6 +666,7 @@ def _strict_readback(materializer, flat, *, source): split_kv_plan_set=_plan_set(contract, backend=source), strict_mode=True, strict_core_id=STRICT_ATTENTION_CORE_ID, + strict_schedule=STRICT_ATTENTION_SCHEDULE_ID, native_attention_arithmetic=False, strict_split_kv_policy="disabled", ) @@ -819,6 +821,10 @@ def test_strict_runtime_readback_entrypoint_binds_executed_evidence(): {"strict_core_id": "different.core"}, BindingErrorCode.ATTENTION_CORE_MISSING, ), + ( + {"strict_schedule": "different_schedule"}, + BindingErrorCode.ATTENTION_CORE_SCHEDULE, + ), ( {"strict_split_kv_policy": "fixed"}, BindingErrorCode.ATTENTION_CORE_SPLIT_K, diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index cc93874b..4e58c7ba 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -688,3 +688,29 @@ def test_gapped_partial_ranges_raise(): def test_registry_dispatches_cp_attention_reference(): assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) + + +def test_strict_reference_is_bitwise_across_batch_cp_and_backward(): + q, k, v = _qkv(2, 4, 8, seed=41, heads=4, kv_heads=2, dim=8) + dout = torch.randn_like(q) + op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + + cp1_out, cp1_lse = op.forward_with_lse( + q, k, v, cp_world_size=1, kv_chunk_size=3 + ) + cp2_out, cp2_lse = op.forward_with_lse( + q, k, v, cp_world_size=2, kv_chunk_size=3 + ) + single_out, single_lse = op.forward_with_lse( + q[:1], k[:1], v[:1], cp_world_size=1, kv_chunk_size=3 + ) + assert torch.equal(cp1_out, cp2_out) + assert torch.equal(cp1_lse, cp2_lse) + assert torch.equal(cp1_out[:1], single_out) + assert torch.equal(cp1_lse[:1], single_lse) + + cp1 = op.backward_reference(q, k, v, dout, cp_world_size=1, kv_chunk_size=3) + cp2 = op.backward_reference(q, k, v, dout, cp_world_size=2, kv_chunk_size=3) + assert torch.equal(cp1.gradients.dq, cp2.gradients.dq) + assert torch.equal(cp1.gradients.dk, cp2.gradients.dk) + assert torch.equal(cp1.gradients.dv, cp2.gradients.dv) From fb58a1c2bd62d69db8e69c5b4777278c37b97f48 Mon Sep 17 00:00:00 2001 From: lamentropetion <3051000145@qq.com> Date: Tue, 18 Aug 2026 23:25:30 +0800 Subject: [PATCH 17/17] fix(attention): fail closed without production CP backend --- .../cross_config/attention_binding.py | 85 ++++++++++++++ .../kernels/ops/pytorch/attention/ablation.py | 108 ++++++++++++++++-- tests/test_attention_ablation.py | 79 +++++++++++++ tests/test_attention_cross_config_binding.py | 51 +++++++++ 4 files changed, 311 insertions(+), 12 deletions(-) diff --git a/rl_engine/alignment/cross_config/attention_binding.py b/rl_engine/alignment/cross_config/attention_binding.py index 57a54205..458a386e 100644 --- a/rl_engine/alignment/cross_config/attention_binding.py +++ b/rl_engine/alignment/cross_config/attention_binding.py @@ -129,6 +129,8 @@ class BindingErrorCode(str, Enum): ATTENTION_CORE_SCHEDULE = "ATTENTION_CORE_SCHEDULE" ATTENTION_NATIVE_ARITHMETIC = "ATTENTION_NATIVE_ARITHMETIC" ATTENTION_CORE_SPLIT_K = "ATTENTION_CORE_SPLIT_K" + ATTENTION_BACKEND_MISSING = "ATTENTION_BACKEND_MISSING" + ATTENTION_NOT_PRODUCTION_READY = "ATTENTION_NOT_PRODUCTION_READY" @dataclass(frozen=True) @@ -151,6 +153,9 @@ class AttentionRuntimeReadback: strict_schedule: str | None = None native_attention_arithmetic: bool = True strict_split_kv_policy: str | None = None + actual_backend: str | None = None + communication_backend: str | None = None + production_ready: bool = False def __post_init__(self) -> None: if not isinstance(self.contract, AttentionContract): @@ -200,6 +205,14 @@ def __post_init__(self) -> None: "auto", }: raise ValueError("strict_split_kv_policy must be disabled, fixed, or auto") + for name, value in ( + ("actual_backend", self.actual_backend), + ("communication_backend", self.communication_backend), + ): + if value is not None and (not isinstance(value, str) or not value.strip()): + raise ValueError(f"{name} must be a non-empty string when provided") + if not isinstance(self.production_ready, bool): + raise TypeError("production_ready must be a bool") normalized_projection_plans: dict[str, Mapping[str, Any]] = {} for name, plan in self.projection_plans.items(): if not isinstance(name, str) or not isinstance(plan, Mapping): @@ -248,6 +261,11 @@ def to_dict(self) -> dict[str, Any]: "native_attention_arithmetic": self.native_attention_arithmetic, "split_kv_policy": self.strict_split_kv_policy, }, + "runtime_backend": { + "actual_backend": self.actual_backend, + "communication_backend": self.communication_backend, + "production_ready": self.production_ready, + }, "split_kv_runtime_plan_set": self.split_kv_plan_set.to_dict(), } @@ -1031,6 +1049,9 @@ def bind_attention_runtime_readbacks( "strict.schedule": rollout.strict_schedule, "strict.native_attention_arithmetic": rollout.native_attention_arithmetic, "strict.split_kv_policy": rollout.strict_split_kv_policy, + "runtime.actual_backend": rollout.actual_backend, + "runtime.communication_backend": rollout.communication_backend, + "runtime.production_ready": rollout.production_ready, **{ f"projection.{projection}": dict(plan) for projection, plan in rollout.projection_plans.items() @@ -1050,6 +1071,9 @@ def bind_attention_runtime_readbacks( "strict.schedule": training.strict_schedule, "strict.native_attention_arithmetic": training.native_attention_arithmetic, "strict.split_kv_policy": training.strict_split_kv_policy, + "runtime.actual_backend": training.actual_backend, + "runtime.communication_backend": training.communication_backend, + "runtime.production_ready": training.production_ready, **{ f"projection.{projection}": dict(plan) for projection, plan in training.projection_plans.items() @@ -1065,6 +1089,45 @@ def _strict_attention_core_issues( if not readback.strict_mode: return [] issues = [] + if readback.actual_backend != "rlkernel.cuda.deterministic_attention": + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_BACKEND_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.runtime.actual_backend", + rollout=readback.actual_backend if side == "rollout" else None, + training=readback.actual_backend if side == "training" else None, + message=( + f"{side} strict Attention did not execute the CUDA deterministic core" + ), + ) + ) + if readback.communication_backend != "self_owned_cuda_ag_rs": + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_BACKEND_MISSING, + tier=BindingTier.SEMANTIC, + field=f"{side}.runtime.communication_backend", + rollout=readback.communication_backend if side == "rollout" else None, + training=readback.communication_backend if side == "training" else None, + message=( + f"{side} strict CP Attention did not execute the self-owned CUDA AG/RS path" + ), + ) + ) + if not readback.production_ready: + issues.append( + BindingIssue( + code=BindingErrorCode.ATTENTION_NOT_PRODUCTION_READY, + tier=BindingTier.SEMANTIC, + field=f"{side}.runtime.production_ready", + rollout=False if side == "rollout" else None, + training=False if side == "training" else None, + message=( + f"{side} evidence is reference-only and cannot close the production gate" + ), + ) + ) if readback.strict_core_id != STRICT_ATTENTION_CORE_ID: issues.append( BindingIssue( @@ -1153,6 +1216,28 @@ def _strict_attention_core_pair_issues( message="training and rollout executed different strict Attention schedules", ) ] + if rollout.strict_mode and rollout.actual_backend != training.actual_backend: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="runtime.actual_backend", + rollout=rollout.actual_backend, + training=training.actual_backend, + message="training and rollout executed different Attention backends", + ) + ] + if rollout.strict_mode and rollout.communication_backend != training.communication_backend: + return [ + BindingIssue( + code=BindingErrorCode.ATTENTION_CORE_MISMATCH, + tier=BindingTier.SEMANTIC, + field="runtime.communication_backend", + rollout=rollout.communication_backend, + training=training.communication_backend, + message="training and rollout executed different Attention communication backends", + ) + ] return [] diff --git a/rl_engine/kernels/ops/pytorch/attention/ablation.py b/rl_engine/kernels/ops/pytorch/attention/ablation.py index de1a8d5c..e1719085 100644 --- a/rl_engine/kernels/ops/pytorch/attention/ablation.py +++ b/rl_engine/kernels/ops/pytorch/attention/ablation.py @@ -148,6 +148,7 @@ def __init__( core: Any | None = None, reference: Any | None = None, native: Any | None = None, + cp_backend: Any | None = None, communication_backend: str = "none", ) -> None: if not isinstance(communication_backend, str) or not communication_backend.strip(): @@ -155,6 +156,10 @@ def __init__( self.core = core self.reference = reference self.native = native + # CP production execution is injected by the runtime adapter. Keeping + # it separate from the single-device core prevents an accidental + # fallback to the PyTorch reference when AG/RS is required. + self.cp_backend = cp_backend self.communication_backend = communication_backend.strip() def __call__( @@ -248,7 +253,12 @@ def apply( raise AttentionContractError("dout must have the same shape as q") requested = backend_request if backend_request is not None else cfg.backend - selected, selected_id = self._select_backend(requested, q, contract) + selected, selected_id = self._select_backend( + requested, + q, + contract, + communication_backend=cfg.communication_backend, + ) if cfg.deterministic and selected_id == "native": raise AttentionContractError( "deterministic=True cannot execute an unverified native Attention backend" @@ -270,9 +280,17 @@ def apply( call_kwargs.setdefault("cp_world_size", contract.sharding.cp_world_size) if contract.split_kv.mode is SplitKVMode.FIXED: call_kwargs.setdefault("kv_chunk_size", contract.split_kv.fixed_split_size) - out, lse = self._invoke(selected, q, k, v, call_kwargs, contract) + out, lse, backend_provenance = self._invoke( + selected, q, k, v, call_kwargs, contract + ) if cfg.validate: self._validate_outputs(out, lse, q, contract) + _validate_runtime_provenance( + selected, + selected_id, + backend_provenance, + cfg, + ) dq = dk = dv = None if cfg.return_gradients: @@ -307,6 +325,10 @@ def apply( "return_lse": cfg.return_lse, "return_gradients": cfg.return_gradients, } + provenance.update(backend_provenance) + provenance.setdefault("actual_backend", selected_id) + provenance.setdefault("communication_backend", cfg.communication_backend) + provenance.setdefault("production_ready", False) return AttentionAblationResult( out=out, lse=lse if cfg.return_lse else None, @@ -328,6 +350,8 @@ def _select_backend( requested: str | Callable[..., Any], q: Tensor, contract: AttentionContract, + *, + communication_backend: str, ) -> tuple[Any, str]: if callable(requested) or hasattr(requested, "forward_with_lse") or hasattr(requested, "apply"): return requested, _callable_backend_id(requested) @@ -342,11 +366,15 @@ def _select_backend( return self._reference_backend(), REFERENCE_BACKEND_ID if normalized not in {"auto", "deterministic", "rlkernel"}: raise AttentionContractError(f"unsupported Attention backend {requested!r}") - if ( - contract.sharding.cp_world_size > 1 - or q.device.type != "cuda" - or torch.version.hip is not None - ): + if contract.sharding.cp_world_size > 1: + if communication_backend == "self_owned_cuda_ag_rs": + if self.cp_backend is None: + raise AttentionContractError( + "CP production Attention requires an injected AG/RS backend" + ) + return self.cp_backend, _callable_backend_id(self.cp_backend) + return self._reference_backend(), REFERENCE_BACKEND_ID + if q.device.type != "cuda" or torch.version.hip is not None: return self._reference_backend(), REFERENCE_BACKEND_ID if self.core is None: from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( @@ -420,7 +448,7 @@ def _invoke( v: Tensor, kwargs: Mapping[str, Any], contract: AttentionContract, - ) -> tuple[Tensor, Tensor]: + ) -> tuple[Tensor, Tensor, dict[str, Any]]: method = getattr(backend, "forward_with_lse", None) if not callable(method): method = getattr(backend, "apply", None) @@ -436,17 +464,27 @@ def _invoke( "CP>1 requires an Attention backend that explicitly accepts cp_world_size" ) result = method(q, k, v, **accepted) + backend_provenance: dict[str, Any] = {} if isinstance(result, AttentionAblationResult): out, lse = result.out, result.lse elif isinstance(result, tuple) and len(result) == 2: out, lse = result else: - raise AttentionContractError( - "Attention backend must return (out, lse) or AttentionAblationResult" - ) + out = getattr(result, "out", None) + lse = getattr(result, "lse", None) + raw_provenance = getattr(result, "provenance", {}) + if isinstance(raw_provenance, Mapping): + backend_provenance = dict(raw_provenance) + if out is None or lse is None: + raise AttentionContractError( + "Attention backend must return (out, lse), AttentionAblationResult, " + "or an object with out/lse/provenance" + ) + if isinstance(result, AttentionAblationResult): + backend_provenance = dict(result.provenance) if not isinstance(out, Tensor) or not isinstance(lse, Tensor): raise AttentionContractError("Attention backend returned non-tensor output or LSE") - return out, lse + return out, lse, backend_provenance @staticmethod def _backward( @@ -523,6 +561,52 @@ def _callable_backend_id(value: Any) -> str: return f"injected.{type(value).__module__}.{type(value).__qualname__}" +def _validate_runtime_provenance( + selected: Any, + selected_id: str, + runtime: Mapping[str, Any], + config: AttentionAblationConfig, +) -> None: + """Fail closed when a production strict backend did not prove its identity.""" + + if not config.deterministic: + return + if selected_id == "native": + raise AttentionContractError( + "deterministic Attention cannot execute native Attention arithmetic" + ) + + actual_core = runtime.get("strict_core_id", getattr(selected, "core_id", None)) + actual_schedule = runtime.get("strict_schedule", getattr(selected, "strict_schedule", None)) + if selected_id != REFERENCE_BACKEND_ID and ( + actual_core != config.strict_core_id or actual_schedule != config.strict_schedule + ): + raise AttentionContractError( + "deterministic Attention backend did not prove the shared strict core and schedule" + ) + + if config.communication_backend != "self_owned_cuda_ag_rs": + return + + expected = { + "strict_core_id": config.strict_core_id, + "strict_schedule": config.strict_schedule, + "actual_backend": "rlkernel.cuda.deterministic_attention", + "communication_backend": "self_owned_cuda_ag_rs", + "production_ready": True, + "native_attention_arithmetic": False, + "fallback": False, + } + mismatches = [ + name for name, value in expected.items() if runtime.get(name) != value + ] + if mismatches: + raise AttentionContractError( + "CP production Attention runtime provenance is incomplete or mismatched: " + + ", ".join(mismatches) + ) + + def _actual_split_provenance( contract: AttentionContract, *, diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py index 7023c6a5..70dbd2ca 100644 --- a/tests/test_attention_ablation.py +++ b/tests/test_attention_ablation.py @@ -2,6 +2,8 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch @@ -104,6 +106,83 @@ def forward_with_lse(self, q, k, v, *, causal, scale): assert torch.equal(result.out, q) +def test_cp_production_configuration_fails_closed_without_ag_rs_backend(): + q, k, v = _qkv() + cp_sharding = ShardingSpec( + tp_rank=0, + tp_world_size=1, + cp_rank=0, + cp_world_size=2, + global_q_heads=2, + global_kv_heads=1, + local_q_head_start=0, + local_q_heads=2, + local_kv_head_start=0, + local_kv_heads=1, + global_sequence_length=4, + local_sequence_length=2, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 2), + ) + contract = AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=2, + head_dim=4, + causal=True, + causal_offsets=(0,), + sharding=cp_sharding, + reduction=ReductionSpec(), + split_kv=SplitKVSpec.disabled(), + ) + with pytest.raises(AttentionContractError, match="injected AG/RS backend"): + AttentionAblationOp(communication_backend="self_owned_cuda_ag_rs")( + q[:, :, :2], k[:, :, :2], v[:, :, :2], contract=contract + ) + + +def test_cp_production_wrapper_preserves_runtime_backend_provenance(): + q, k, v = _qkv() + + class StrictCPBackend: + backend_id = "injected.strict_cp_backend" + core_id = STRICT_ATTENTION_CORE_ID + strict_schedule = STRICT_ATTENTION_SCHEDULE_ID + + def __call__(self, q, k, v, *, causal, scale): + del k, v, causal, scale + return SimpleNamespace( + out=q.clone(), + lse=torch.zeros(q.shape[:3], dtype=torch.float32), + provenance={ + "strict_core_id": STRICT_ATTENTION_CORE_ID, + "strict_schedule": STRICT_ATTENTION_SCHEDULE_ID, + "actual_backend": "rlkernel.cuda.deterministic_attention", + "communication_backend": "self_owned_cuda_ag_rs", + "production_ready": True, + "native_attention_arithmetic": False, + "fallback": False, + }, + ) + + result = AttentionAblationOp( + cp_backend=StrictCPBackend(), + communication_backend="self_owned_cuda_ag_rs", + )( + q, + k, + v, + contract=_contract(), + backend=StrictCPBackend(), + ) + assert result.provenance["actual_backend"] == "rlkernel.cuda.deterministic_attention" + assert result.provenance["communication_backend"] == "self_owned_cuda_ag_rs" + assert result.provenance["production_ready"] is True + + def test_deterministic_attention_rejects_runtime_split_kv_auto(): q, k, v = _qkv() contract = _contract(split_kv=SplitKVSpec.auto(strict_consistency=False)) diff --git a/tests/test_attention_cross_config_binding.py b/tests/test_attention_cross_config_binding.py index 282e5f39..c1c8d35d 100644 --- a/tests/test_attention_cross_config_binding.py +++ b/tests/test_attention_cross_config_binding.py @@ -651,6 +651,9 @@ def _readback(materializer, flat, *, source): preprocess_backends=MANDATED_ATTENTION_PREPROCESS_BACKENDS, preprocess_fallback=False, projection_plans=projection_plans, + actual_backend=f"reference.{source}", + communication_backend="none", + production_ready=False, ) @@ -669,6 +672,9 @@ def _strict_readback(materializer, flat, *, source): strict_schedule=STRICT_ATTENTION_SCHEDULE_ID, native_attention_arithmetic=False, strict_split_kv_policy="disabled", + actual_backend="rlkernel.cuda.deterministic_attention", + communication_backend="self_owned_cuda_ag_rs", + production_ready=True, ) @@ -876,6 +882,51 @@ def test_strict_runtime_readback_accepts_shared_no_split_k_core(): assert result.passed assert result.provenance["rollout"]["recorded"]["strict.core_id"] == (STRICT_ATTENTION_CORE_ID) + assert result.provenance["rollout"]["recorded"]["runtime.actual_backend"] == ( + "rlkernel.cuda.deterministic_attention" + ) + assert result.provenance["rollout"]["recorded"]["runtime.communication_backend"] == ( + "self_owned_cuda_ag_rs" + ) + assert result.provenance["rollout"]["recorded"]["runtime.production_ready"] is True + + +@pytest.mark.parametrize( + ("field", "value", "error_code"), + [ + ( + "actual_backend", + "rlkernel.attention.cp_reference", + BindingErrorCode.ATTENTION_BACKEND_MISSING, + ), + ( + "communication_backend", + "p2p_nccl_reference", + BindingErrorCode.ATTENTION_BACKEND_MISSING, + ), + ("production_ready", False, BindingErrorCode.ATTENTION_NOT_PRODUCTION_READY), + ], +) +def test_strict_runtime_readback_rejects_reference_only_evidence(field, value, error_code): + rollout = replace( + _strict_readback(VllmRolloutMaterializer(), ROLLOUT_KNOBS, source="vllm.runtime_readback"), + **{field: value}, + ) + training = _strict_readback( + MegatronAttentionMaterializer(), + TRAINING_KNOBS, + source="megatron.runtime_readback", + ) + result = bind_attention_runtime_readbacks( + rollout=rollout, + training=training, + rollout_identity=_identity(), + training_identity=_identity(), + rollout_backend_id="rlkernel.cuda.deterministic_attention", + training_backend_id="rlkernel.cuda.deterministic_attention", + ) + assert not result.passed + assert result.issues_by_code(error_code) def test_strict_runtime_readback_accepts_common_deterministic_preprocess_fallback():