diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 3f2c63e8..e1aff500 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -13,12 +13,6 @@ 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 @@ -90,18 +84,6 @@ the inputs' device. Calling it (`__call__` -> `forward(...)`) computes in the input dtype; `forward_fp32(...)` is the explicit fp32 golden path (NativeAttentionOp only). The production `"attn"` op_type (SDPA-based `PYTORCH_ATTN`, FlashAttention, etc.) is a separate dispatch chain and is unaffected. -### WS2 CP-aware dispatch - -WS2 distributed callers use a separate contract-aware entry point, -`kernel_registry.get_attention_op(contract)`. It validates explicit TP/CP ownership, fixed -`(out, lse)` merge semantics, causal or packed-sequence offsets, and decode KV-cache identity -before selecting a backend. Legacy `get_op("attention")` behavior remains unchanged. - -Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic -CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as -a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md). - -### WS2 deterministic CP reference ### WS2 CP-aware dispatch @@ -129,76 +111,6 @@ python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 --output artifacts/ws2-cp-attention-drift.json ``` -Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow -`disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block -boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason. -Runtime-selected `auto` plans are diagnostic only unless both training and rollout export and -validate the same actual plan. - -The rank-aware drift benchmark can emit a CPU smoke artifact or a torchrun-friendly GPU report: - -```bash -python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json -python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 \ - --cp-world-sizes 2 --kv-chunk-sizes none,1 --include-backward \ - --output artifacts/ws2-cp-attention-drift.json -``` - -### WS2 CP-aware dispatch - -WS2 distributed callers use a separate contract-aware entry point, -`kernel_registry.get_attention_op(contract)`. It validates explicit TP/CP ownership, fixed -`(out, lse)` merge semantics, causal or packed-sequence offsets, and decode KV-cache identity -before selecting a backend. Legacy `get_op("attention")` behavior remains unchanged. - -Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic -CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as -a silent fallback. See the -[WS2 CP-aware Attention contract in PR #236](https://github.com/RL-Align/RL-Kernel/blob/feat/issue-235-attention-cp-contract/docs/design/ws2-cp-attention-contract.md). - -Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow -`disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block -boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason. -Runtime-selected `auto` plans are diagnostic only unless both training and rollout export and -validate the same actual plan. - -The rank-aware drift benchmark can emit a CPU smoke artifact or a torchrun-friendly GPU report: - -```bash -python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json -python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 \ - --cp-world-sizes 2 --kv-chunk-sizes none,1 --include-backward \ - --output artifacts/ws2-cp-attention-drift.json -``` - -### WS2 CP-aware dispatch - -WS2 distributed callers use a separate contract-aware entry point, -`kernel_registry.get_attention_op(contract)`. It validates explicit TP/CP ownership, fixed -`(out, lse)` merge semantics, causal or packed-sequence offsets, and decode KV-cache identity -before selecting a backend. Legacy `get_op("attention")` behavior remains unchanged. - -Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic -CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as -a silent fallback. See the [WS2 CP-aware Attention contract PR][attention-contract-pr]. - -[attention-contract-pr]: https://github.com/RL-Align/RL-Kernel/pull/236 - -Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow -`disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block -boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason. -Runtime-selected `auto` plans are diagnostic only unless both training and rollout export and -validate the same actual plan. - -The rank-aware drift benchmark can emit a CPU smoke artifact or a torchrun-friendly GPU report: - -```bash -python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json -python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 \ - --cp-world-sizes 2 --kv-chunk-sizes none,1 --include-backward \ - --output artifacts/ws2-cp-attention-drift.json -``` - ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -309,8 +221,12 @@ Hooks: - `forward(q, k, v, ...)` — main path (registry, #108 harness). Differentiable. - `forward_with_lse(q, k, v, ...)` — returns `(out, lse)` for LSE verification, debugging, and future KV-cache / training integration. -- `backward_reference(q, k, v, dout, ...)` — runs the deterministic training backward - validation path and returns `dq`, `dk`, `dv`, `out`, `lse`, and provenance. +- `save_forward_state(q, k, v, ...)` captures FP32 `out/lse`, masks, position metadata, + topology, Split-KV boundaries, and content fingerprints for the exact forward invocation. +- `backward_reference(q, k, v, dout, ..., saved_forward_state=state)` consumes that saved + state in canonical global KV-block order. It fails closed if Q/K/V, saved tensors, masks, + offsets, topology, or Split-KV metadata changed, and returns `dq`, `dk`, `dv`, `out`, `lse`, + and provenance. - `compare_cp_attention_backward(q, k, v, dout, ...)` — compares CP=1 backward against CP/chunked-prefill backward and emits whole-tensor plus per-logical-rank drift stats. diff --git a/rl_engine/kernels/gtest/tolerance.py b/rl_engine/kernels/gtest/tolerance.py index 83aa49a0..4fb5bcbf 100644 --- a/rl_engine/kernels/gtest/tolerance.py +++ b/rl_engine/kernels/gtest/tolerance.py @@ -202,62 +202,6 @@ def load_contract( return contract -def resolve_logprob_threshold(dtype: Any) -> float: - """Return the fixed WS1 selected-logprob absolute-difference threshold. - - The contract path is intentionally not configurable through this accessor. - Cross-configuration experiment definitions may select a dtype, but they cannot - inject or override a numerical threshold. - """ - - dtype_name = _normalize_dtype_name(dtype) - contract = load_contract() - try: - values = contract["accuracy"]["default"]["logprob"][dtype_name] - raw_threshold = values["atol"] - except (KeyError, TypeError) as exc: - raise ValueError(f"WS1 has no logprob threshold for dtype {dtype_name!r}") from exc - if isinstance(raw_threshold, bool) or not isinstance(raw_threshold, (int, float)): - raise ValueError(f"invalid WS1 logprob threshold for dtype {dtype_name!r}") - threshold = float(raw_threshold) - if not math.isfinite(threshold) or threshold < 0.0: - raise ValueError(f"invalid WS1 logprob threshold for dtype {dtype_name!r}") - return threshold - - -def tolerance_contract_fingerprint() -> str: - """Return a deterministic fingerprint of the current WS1 contract contents.""" - - canonical = json.dumps( - load_contract(), - ensure_ascii=True, - separators=(",", ":"), - sort_keys=True, - ).encode("utf-8") - return hashlib.sha256(canonical).hexdigest() - - -def _normalize_dtype_name(dtype: Any) -> str: - normalized = str(dtype).strip().lower().replace("torch.", "").replace("-", "") - aliases = { - "bf16": "bfloat16", - "bfloat16": "bfloat16", - "fp16": "float16", - "float16": "float16", - "half": "float16", - "fp32": "float32", - "float32": "float32", - "float": "float32", - } - try: - return aliases[normalized] - except KeyError as exc: - valid = ", ".join(sorted(set(aliases.values()))) - raise ValueError( - f"unsupported WS1 logprob dtype {dtype!r}; expected one of: {valid}" - ) from exc - - def validate_contract_schema(contract: Mapping[str, Any]) -> None: """Validate four-judgment schema, dtype policy, roles, and aggregates.""" @@ -1112,8 +1056,6 @@ def tolerance_contract_fingerprint() -> str: "tolerance_contract_fingerprint", "resolve_tolerance", "resolve_tolerance_support", - "resolve_logprob_threshold", - "tolerance_contract_fingerprint", "validate_backend_provenance", "validate_contract_schema", ] diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 7d79d713..d68dbb18 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -10,6 +10,7 @@ from __future__ import annotations +import hashlib import math from dataclasses import dataclass from typing import Optional, Sequence @@ -141,6 +142,50 @@ class AttentionBackwardGradients: dv: torch.Tensor +@dataclass(frozen=True) +class AttentionSavedForwardState: + """Exact FP32 forward state consumed by the PR8 backward reference.""" + + out: torch.Tensor + lse: torch.Tensor + causal: bool + scale: float + key_padding_mask: Optional[torch.Tensor] + query_position_offsets: torch.Tensor + key_position_offsets: torch.Tensor + cp_world_size: int + kv_chunk_size: Optional[int] + query_bounds: tuple[tuple[int, int], ...] + kv_block_bounds: tuple[tuple[int, int], ...] + q_shape: tuple[int, ...] + k_shape: tuple[int, ...] + v_shape: tuple[int, ...] + q_dtype: torch.dtype + k_dtype: torch.dtype + v_dtype: torch.dtype + q_fingerprint: str + k_fingerprint: str + v_fingerprint: str + out_fingerprint: str + lse_fingerprint: str + key_padding_mask_fingerprint: Optional[str] + query_position_offsets_fingerprint: str + key_position_offsets_fingerprint: str + strict_bitwise: bool + strict_schedule: Optional[str] + + def __post_init__(self) -> None: + if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: + raise ValueError("saved attention out/lse must be FP32") + if self.out.ndim != 4 or self.lse.shape != self.out.shape[:3]: + raise ValueError("saved attention out/lse shapes are invalid") + if not math.isfinite(self.scale) or self.scale <= 0: + raise ValueError("saved attention scale must be positive and finite") + expected_schedule = STRICT_ATTENTION_SCHEDULE_ID if self.strict_bitwise else None + if self.strict_schedule != expected_schedule: + raise ValueError("saved attention strict schedule does not match strict_bitwise") + + @dataclass(frozen=True) class AttentionBackwardPathResult: """One materialized CP attention backward path.""" @@ -149,6 +194,7 @@ class AttentionBackwardPathResult: out: torch.Tensor lse: torch.Tensor gradients: AttentionBackwardGradients + saved_forward_state: AttentionSavedForwardState provenance: dict[str, object] @@ -328,9 +374,24 @@ def execution_provenance( kv_chunk_size=kv_chunk_size, backend="deterministic_cp_reference", ) - return out + return { + "execution_scope": "logical_single_process_cp_reference", + "runtime_verified": False, + "input_boundary": "projected_post_qk_norm_post_rope_qkv", + "query_scope": "logical_global_query_reference", + "kv_scope": "logical_owner_local_cp_shards", + "production_cp_protocol": "ag_query_local_kv_rs_out_lse", + "communication_executed": "none", + "partial_state": "fp32_out_attention_lse", + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "requested_split_kv_policy": "disabled" if kv_chunk_size is None else "fixed", + "requested_split_kv_size": kv_chunk_size, + "actual_split_kv_plans": plans, + } - def forward_with_lse( + def __call__( self, q: torch.Tensor, k: torch.Tensor, @@ -343,18 +404,37 @@ def forward_with_lse( key_position_offsets: Optional[torch.Tensor] = None, cp_world_size: int = 1, kv_chunk_size: Optional[int] = None, - output_dtype: Optional[torch.dtype] = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Return ``(out, lse)`` for the CP reference path. + ) -> torch.Tensor: + return self.forward( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) - ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 - until the final write, then downcast to ``output_dtype``. When omitted, - ``output_dtype`` defaults to the input dtype. - """ + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and final input-dtype write.""" - resolved_output_dtype = q.dtype if output_dtype is None else output_dtype - _validate_output_dtype(resolved_output_dtype) - out, lse = self._forward_impl( + out, _ = self.forward_with_lse( q, k, v, @@ -365,7 +445,64 @@ def forward_with_lse( key_position_offsets=key_position_offsets, cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size, + output_dtype=q.dtype, ) + return out + + def forward_fp32( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> torch.Tensor: + """Compute CP attention with fp32 accumulation and fp32 output.""" + + out, _ = self.forward_fp32_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(out, lse)`` for the CP reference path. + + ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 + until the final write, then downcast to ``output_dtype``. When omitted, + ``output_dtype`` defaults to the input dtype. + """ + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) if self.strict_bitwise: out, lse = self._forward_strict_bitwise( q, @@ -376,6 +513,7 @@ def forward_with_lse( 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, ) else: @@ -435,6 +573,7 @@ def _forward_strict_bitwise( key_padding_mask: Optional[torch.Tensor], query_position_offsets: Optional[torch.Tensor], key_position_offsets: Optional[torch.Tensor], + cp_world_size: int, kv_chunk_size: Optional[int], ) -> tuple[torch.Tensor, torch.Tensor]: """Run one canonical arithmetic schedule for every caller. @@ -450,6 +589,7 @@ def _forward_strict_bitwise( _validate_qkv(q, k, v) _validate_scale(scale) + _validate_partition_args(cp_world_size, kv_chunk_size) batch, hq, sq, dim = q.shape skv = k.size(2) if key_padding_mask is not None: @@ -472,6 +612,13 @@ def _forward_strict_bitwise( default=0, name="key_position_offsets", ) + if sq == 0: + zero_dep = _zero_dependency(q.float(), k.float(), v.float()) + return ( + torch.empty(batch, hq, 0, dim, device=q.device, dtype=torch.float32) + zero_dep, + torch.empty(batch, hq, 0, device=q.device, dtype=torch.float32) + zero_dep, + ) + # The canonical schedule is global. CP ownership and arrival order # must not change which partial states are generated or merged. out_rows: list[torch.Tensor] = [] @@ -527,6 +674,7 @@ def backward_reference( kv_chunk_size: Optional[int] = None, output_dtype: Optional[torch.dtype] = torch.float32, name: Optional[str] = None, + saved_forward_state: Optional[AttentionSavedForwardState] = None, ) -> AttentionBackwardPathResult: """Run the deterministic training-side backward validation path. @@ -546,16 +694,12 @@ def backward_reference( raise ValueError("dout must be on the same device as q, k, and v") if dout.dtype != q.dtype: raise ValueError("dout must have the same dtype as q") - q_leaf = q.detach().clone().requires_grad_(True) - k_leaf = k.detach().clone().requires_grad_(True) - v_leaf = v.detach().clone().requires_grad_(True) - resolved_output_dtype = q.dtype if output_dtype is None else output_dtype _validate_output_dtype(resolved_output_dtype) - out, lse = self.forward_with_lse( - q_leaf, - k_leaf, - v_leaf, + state = saved_forward_state or self.save_forward_state( + q, + k, + v, causal=causal, scale=scale, key_padding_mask=key_padding_mask, @@ -563,11 +707,31 @@ def backward_reference( key_position_offsets=key_position_offsets, cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size, - output_dtype=resolved_output_dtype, ) - torch.autograd.backward(out, dout.to(dtype=out.dtype)) - if q_leaf.grad is None or k_leaf.grad is None or v_leaf.grad is None: - raise RuntimeError("CP attention backward did not produce dq/dk/dv") + _validate_saved_forward_state( + state, + 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, + strict_bitwise=self.strict_bitwise, + ) + gradients = _backward_from_saved_state( + q, + k, + v, + dout, + state, + strict_bitwise=self.strict_bitwise, + ) + out = state.out.to(resolved_output_dtype) + lse = state.lse return AttentionBackwardPathResult( name=name @@ -575,10 +739,11 @@ def backward_reference( out=out.detach(), lse=lse.detach(), gradients=AttentionBackwardGradients( - dq=q_leaf.grad.detach(), - dk=k_leaf.grad.detach(), - dv=v_leaf.grad.detach(), + dq=gradients.dq, + dk=gradients.dk, + dv=gradients.dv, ), + saved_forward_state=state, provenance={ "attention_mode": "prefill" if kv_chunk_size is None else "chunked_prefill", "gradient_mode": "training_backward", @@ -597,25 +762,130 @@ def backward_reference( "kv_chunk_size": kv_chunk_size, "requested_split_kv_policy": ("disabled" if kv_chunk_size is None else "fixed"), "requested_split_kv_size": kv_chunk_size, - "actual_split_kv_plans": split_kv_execution_plan_provenance( - k.size(2), - cp_world_size=cp_world_size, - kv_chunk_size=kv_chunk_size, - backend="deterministic_cp_backward_reference", + "actual_split_kv_plans": ( + _strict_no_split_plan_provenance( + k.size(2), + cp_world_size=cp_world_size, + backend="deterministic_cp_backward_strict_reference", + ) + if self.strict_bitwise + else split_kv_execution_plan_provenance( + k.size(2), + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_backward_reference", + ) ), "merge_order": "global_block_index", "accum_dtype": "fp32", "downcast_at": "final_write", + "strict_bitwise": self.strict_bitwise, + "strict_core_id": (STRICT_ATTENTION_CORE_ID if self.strict_bitwise else None), + "strict_schedule": (STRICT_ATTENTION_SCHEDULE_ID if self.strict_bitwise else None), + "actual_split_kv_policy": ( + "disabled" + if self.strict_bitwise + else ("disabled" if kv_chunk_size is None else "fixed") + ), "output_dtype": str(resolved_output_dtype).replace("torch.", ""), "q_dtype": str(q.dtype).replace("torch.", ""), "k_dtype": str(k.dtype).replace("torch.", ""), "v_dtype": str(v.dtype).replace("torch.", ""), "dout_dtype": str(dout.dtype).replace("torch.", ""), + "saved_forward_state_source": ( + "caller" if saved_forward_state is not None else "captured_reference" + ), + "backward_algorithm": ( + "saved_out_lse_canonical_row_reference" + if self.strict_bitwise + else "saved_out_lse_block_order_reference" + ), "te_backward_oracle": "not_used", "decode_backward": "not_supported", + "projection_scope": "attention_core_only", + "qkv_projection_backward_dgrad_collective": "all_reduce", + "qkv_projection_sp_backward_collective": "reduce_scatter", + "o_proj_backward_dgrad_collective": "none", + "o_proj_sp_backward_collective": "all_gather", + "projection_collectives_executed": False, + "projection_collectives_source": "attention_contract_runtime_adapter", }, ) + def save_forward_state( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> AttentionSavedForwardState: + """Capture the exact state a production training backward must consume.""" + + out, lse = self.forward_fp32_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + batch, _, sq, dim = q.shape + query_offsets = _normalize_position_offsets( + query_position_offsets, + batch, + q.device, + default=k.size(2) - sq, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + batch, + q.device, + default=0, + name="key_position_offsets", + ) + mask = None if key_padding_mask is None else key_padding_mask.detach().clone() + return AttentionSavedForwardState( + out=out.detach().clone(), + lse=lse.detach().clone(), + causal=causal, + scale=float(scale if scale is not None else 1.0 / math.sqrt(dim)), + key_padding_mask=mask, + query_position_offsets=query_offsets.detach().clone(), + key_position_offsets=key_offsets.detach().clone(), + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + query_bounds=tuple(_split_bounds(sq, cp_world_size)), + kv_block_bounds=tuple(_kv_block_bounds(k.size(2), cp_world_size, kv_chunk_size)), + q_shape=tuple(q.shape), + k_shape=tuple(k.shape), + v_shape=tuple(v.shape), + q_dtype=q.dtype, + k_dtype=k.dtype, + v_dtype=v.dtype, + q_fingerprint=_tensor_fingerprint(q), + k_fingerprint=_tensor_fingerprint(k), + v_fingerprint=_tensor_fingerprint(v), + out_fingerprint=_tensor_fingerprint(out), + lse_fingerprint=_tensor_fingerprint(lse), + key_padding_mask_fingerprint=(None if mask is None else _tensor_fingerprint(mask)), + query_position_offsets_fingerprint=_tensor_fingerprint(query_offsets), + key_position_offsets_fingerprint=_tensor_fingerprint(key_offsets), + strict_bitwise=self.strict_bitwise, + strict_schedule=(STRICT_ATTENTION_SCHEDULE_ID if self.strict_bitwise else None), + ) + def local_partial_state( self, q: torch.Tensor, @@ -888,6 +1158,305 @@ def compare_cp_attention_backward( ) +def _backward_from_saved_state( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + state: AttentionSavedForwardState, + *, + strict_bitwise: bool, +) -> AttentionBackwardGradients: + """Apply standard-softmax backward in canonical global KV-block order.""" + + if strict_bitwise: + return _backward_strict_from_saved_state(q, k, v, dout, state) + + batch, hq, _, dim = q.shape + hkv = k.size(1) + group_size = hq // hkv + with NativeAttentionOp._strict_fp32_math(q.device.type): + qf = q.float() + kf = k.float() + vf = v.float() + doutf = dout.float() + k_expanded = kf.repeat_interleave(group_size, dim=1) + v_expanded = vf.repeat_interleave(group_size, dim=1) + dq = torch.zeros_like(qf) + dk_expanded = torch.zeros( + batch, + hq, + k.size(2), + dim, + dtype=torch.float32, + device=q.device, + ) + dv_expanded = torch.zeros_like(dk_expanded) + + for q_start, q_end in state.query_bounds: + if q_start == q_end: + continue + q_block = qf[:, :, q_start:q_end, :] + dout_block = doutf[:, :, q_start:q_end, :] + out_block = state.out[:, :, q_start:q_end, :] + lse_block = state.lse[:, :, q_start:q_end] + dq_block = torch.zeros_like(q_block) + for k_start, k_end in state.kv_block_bounds: + if k_start == k_end: + continue + k_block = k_expanded[:, :, k_start:k_end, :] + v_block = v_expanded[:, :, k_start:k_end, :] + scores = torch.matmul(q_block, k_block.transpose(-1, -2)) * state.scale + if state.causal: + query_base = state.query_position_offsets[:, None] + q_start + key_base = state.key_position_offsets[:, None] + k_start + q_pos = ( + torch.arange( + q_end - q_start, + device=q.device, + dtype=torch.long, + ) + + query_base + ) + k_pos = ( + torch.arange( + k_end - k_start, + device=q.device, + dtype=torch.long, + ) + + key_base + ) + scores = scores.masked_fill( + (k_pos[:, None, :] > q_pos[:, :, None])[:, None, :, :], + float("-inf"), + ) + if state.key_padding_mask is not None: + scores = scores.masked_fill( + ~state.key_padding_mask[:, None, None, k_start:k_end], + float("-inf"), + ) + probability = torch.exp(scores - lse_block.unsqueeze(-1)) + probability = torch.where( + torch.isfinite(lse_block).unsqueeze(-1), + probability, + torch.zeros_like(probability), + ) + dv_expanded[:, :, k_start:k_end, :] += torch.matmul( + probability.transpose(-1, -2), + dout_block, + ) + dp = torch.matmul(dout_block, v_block.transpose(-1, -2)) + # The global softmax dot term is dout dot the saved global output. + ds = probability * (dp - (dout_block * out_block).sum(dim=-1, keepdim=True)) + dq_block += torch.matmul(ds, k_block) * state.scale + dk_expanded[:, :, k_start:k_end, :] += ( + torch.matmul(ds.transpose(-1, -2), q_block) * state.scale + ) + dq[:, :, q_start:q_end, :] = dq_block + + dk = dk_expanded.reshape( + batch, + hkv, + group_size, + k.size(2), + dim, + ).sum(dim=2) + dv = dv_expanded.reshape( + batch, + hkv, + group_size, + v.size(2), + dim, + ).sum(dim=2) + return AttentionBackwardGradients( + dq=dq.to(q.dtype), + dk=dk.to(k.dtype), + dv=dv.to(v.dtype), + ) + + +def _backward_strict_from_saved_state( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: torch.Tensor, + state: AttentionSavedForwardState, +) -> AttentionBackwardGradients: + """Run one batch row and the complete logical Q/KV domain at a time.""" + + batch, hq, sq, dim = q.shape + hkv = k.size(1) + skv = k.size(2) + group_size = hq // hkv + if sq == 0: + return AttentionBackwardGradients( + dq=torch.zeros_like(q), + dk=torch.zeros_like(k), + dv=torch.zeros_like(v), + ) + + dq_rows: list[torch.Tensor] = [] + dk_rows: list[torch.Tensor] = [] + dv_rows: list[torch.Tensor] = [] + with NativeAttentionOp._strict_fp32_math(q.device.type): + for batch_index in range(batch): + qf = q[batch_index : batch_index + 1].float().contiguous() + kf = k[batch_index : batch_index + 1].float().contiguous() + vf = v[batch_index : batch_index + 1].float().contiguous() + doutf = dout[batch_index : batch_index + 1].float().contiguous() + k_expanded = kf.repeat_interleave(group_size, dim=1) + v_expanded = vf.repeat_interleave(group_size, dim=1) + scores = torch.matmul(qf, k_expanded.transpose(-1, -2)) * state.scale + if state.causal: + q_pos = state.query_position_offsets[batch_index : batch_index + 1, None] + q_pos = q_pos + torch.arange(sq, device=q.device, dtype=torch.long) + k_pos = state.key_position_offsets[batch_index : batch_index + 1, None] + k_pos = k_pos + torch.arange(skv, device=q.device, dtype=torch.long) + scores = scores.masked_fill( + (k_pos[:, None, :] > q_pos[:, :, None])[:, None, :, :], + float("-inf"), + ) + if state.key_padding_mask is not None: + scores = scores.masked_fill( + ~state.key_padding_mask[batch_index : batch_index + 1, None, None, :], + float("-inf"), + ) + lse = state.lse[batch_index : batch_index + 1] + probability = torch.exp(scores - lse.unsqueeze(-1)) + probability = torch.where( + torch.isfinite(lse).unsqueeze(-1), + probability, + torch.zeros_like(probability), + ) + dp = torch.matmul(doutf, v_expanded.transpose(-1, -2)) + out = state.out[batch_index : batch_index + 1] + delta = (doutf * out).sum(dim=-1, keepdim=True) + ds = probability * (dp - delta) + dq_rows.append(torch.matmul(ds, k_expanded) * state.scale) + dk_expanded = torch.matmul(ds.transpose(-1, -2), qf) * state.scale + dv_expanded = torch.matmul(probability.transpose(-1, -2), doutf) + dk_rows.append(dk_expanded.reshape(1, hkv, group_size, skv, dim).sum(dim=2)) + dv_rows.append(dv_expanded.reshape(1, hkv, group_size, skv, dim).sum(dim=2)) + return AttentionBackwardGradients( + dq=torch.cat(dq_rows, dim=0).to(q.dtype), + dk=torch.cat(dk_rows, dim=0).to(k.dtype), + dv=torch.cat(dv_rows, dim=0).to(v.dtype), + ) + + +def _validate_saved_forward_state( + state: AttentionSavedForwardState, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: Optional[float], + key_padding_mask: Optional[torch.Tensor], + query_position_offsets: Optional[torch.Tensor], + key_position_offsets: Optional[torch.Tensor], + cp_world_size: int, + kv_chunk_size: Optional[int], + strict_bitwise: bool, +) -> None: + if not isinstance(state, AttentionSavedForwardState): + raise ValueError("saved_forward_state must be an AttentionSavedForwardState") + expected_scale = float(scale if scale is not None else 1.0 / math.sqrt(q.size(-1))) + expected_query_offsets = _normalize_position_offsets( + query_position_offsets, + q.size(0), + q.device, + default=k.size(2) - q.size(2), + name="query_position_offsets", + ) + expected_key_offsets = _normalize_position_offsets( + key_position_offsets, + q.size(0), + q.device, + default=0, + name="key_position_offsets", + ) + checks = { + "out_shape": (tuple(state.out.shape), tuple(q.shape)), + "lse_shape": (tuple(state.lse.shape), tuple(q.shape[:3])), + "out_device": (state.out.device, q.device), + "lse_device": (state.lse.device, q.device), + "q_shape": (state.q_shape, tuple(q.shape)), + "k_shape": (state.k_shape, tuple(k.shape)), + "v_shape": (state.v_shape, tuple(v.shape)), + "q_dtype": (state.q_dtype, q.dtype), + "k_dtype": (state.k_dtype, k.dtype), + "v_dtype": (state.v_dtype, v.dtype), + "causal": (state.causal, causal), + "scale": (state.scale, expected_scale), + "cp_world_size": (state.cp_world_size, cp_world_size), + "kv_chunk_size": (state.kv_chunk_size, kv_chunk_size), + "strict_bitwise": (state.strict_bitwise, strict_bitwise), + "strict_schedule": ( + state.strict_schedule, + STRICT_ATTENTION_SCHEDULE_ID if strict_bitwise else None, + ), + "query_bounds": (state.query_bounds, tuple(_split_bounds(q.size(2), cp_world_size))), + "kv_block_bounds": ( + state.kv_block_bounds, + tuple(_kv_block_bounds(k.size(2), cp_world_size, kv_chunk_size)), + ), + "q_fingerprint": (state.q_fingerprint, _tensor_fingerprint(q)), + "k_fingerprint": (state.k_fingerprint, _tensor_fingerprint(k)), + "v_fingerprint": (state.v_fingerprint, _tensor_fingerprint(v)), + "out_fingerprint": (state.out_fingerprint, _tensor_fingerprint(state.out)), + "lse_fingerprint": (state.lse_fingerprint, _tensor_fingerprint(state.lse)), + "query_position_offsets_fingerprint": ( + state.query_position_offsets_fingerprint, + _tensor_fingerprint(state.query_position_offsets), + ), + "key_position_offsets_fingerprint": ( + state.key_position_offsets_fingerprint, + _tensor_fingerprint(state.key_position_offsets), + ), + "query_position_offsets_device": ( + state.query_position_offsets.device, + q.device, + ), + "key_position_offsets_device": ( + state.key_position_offsets.device, + q.device, + ), + } + mismatches = [name for name, (actual, expected) in checks.items() if actual != expected] + if not torch.equal(state.query_position_offsets, expected_query_offsets): + mismatches.append("query_position_offsets") + if not torch.equal(state.key_position_offsets, expected_key_offsets): + mismatches.append("key_position_offsets") + masks_match = ( + state.key_padding_mask is None + and key_padding_mask is None + or state.key_padding_mask is not None + and key_padding_mask is not None + and torch.equal(state.key_padding_mask, key_padding_mask) + ) + if not masks_match: + mismatches.append("key_padding_mask") + actual_mask_fingerprint = ( + None if state.key_padding_mask is None else _tensor_fingerprint(state.key_padding_mask) + ) + if state.key_padding_mask_fingerprint != actual_mask_fingerprint: + mismatches.append("key_padding_mask_fingerprint") + if mismatches: + raise ValueError( + "saved_forward_state does not match the backward invocation: " + ", ".join(mismatches) + ) + + +def _tensor_fingerprint(tensor: torch.Tensor) -> str: + digest = hashlib.sha256() + digest.update(str(tuple(tensor.shape)).encode()) + digest.update(str(tensor.dtype).encode()) + digest.update(str(tensor.device).encode()) + digest.update(tensor.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()) + return digest.hexdigest() + + def _compare_backward_path( candidate: AttentionBackwardPathResult, reference: AttentionBackwardPathResult, @@ -1030,6 +1599,18 @@ def _validate_output_dtype(output_dtype: torch.dtype) -> None: raise ValueError("output_dtype must be a real floating-point torch dtype") +def _validate_partition_args( + cp_world_size: int, + kv_chunk_size: Optional[int], +) -> None: + if isinstance(cp_world_size, bool) or not isinstance(cp_world_size, int) or cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and ( + isinstance(kv_chunk_size, bool) or not isinstance(kv_chunk_size, int) or kv_chunk_size < 1 + ): + raise ValueError("kv_chunk_size must be >= 1 when provided") + + def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: total = torch.tensor(0.0, device=tensors[0].device) for tensor in tensors: @@ -1126,6 +1707,29 @@ def split_kv_execution_plan_provenance( return result +def _strict_no_split_plan_provenance( + length: int, + *, + cp_world_size: int, + backend: str, +) -> list[dict[str, object]]: + """Describe the full logical KV row consumed by each strict CP executor.""" + + if length < 1: + raise ValueError("strict no-Split-KV sequence length must be >= 1") + _validate_partition_args(cp_world_size, None) + plan = SplitKVExecutionPlan( + requested_mode=SplitKVMode.DISABLED, + requested_split_size=None, + actual_mode=SplitKVMode.DISABLED, + actual_split_size=None, + boundaries=((0, length),), + backend=backend, + source="canonical_strict_execution", + ).to_dict() + return [{"owner_cp_rank": cp_rank, **plan} for cp_rank in range(cp_world_size)] + + def build_reference_split_kv_runtime_plan_set( total_kv_tokens: Sequence[int], *, @@ -1151,7 +1755,6 @@ 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),) @@ -1200,6 +1803,7 @@ def build_reference_split_kv_runtime_plan_set( "AttentionBackwardPathResult", "AttentionBackwardRankDrift", "AttentionPartialState", + "AttentionSavedForwardState", "build_reference_split_kv_runtime_plan_set", "CPAttentionReferenceOp", "DeterministicAttentionCore", diff --git a/rl_engine/testing/__init__.py b/rl_engine/testing/__init__.py index 602c875b..6e1af25b 100644 --- a/rl_engine/testing/__init__.py +++ b/rl_engine/testing/__init__.py @@ -16,8 +16,8 @@ compare_single_gpu_attention, compare_single_gpu_rope_attention, decode_kv_cache_fingerprint, - decode_rope_identity_fingerprint, decode_prefix_cache_fingerprint, + decode_rope_identity_fingerprint, run_chunked_query_attention, run_decode_full_prefill_reference, run_decode_kv_replay, diff --git a/rl_engine/testing/attention_comparison.py b/rl_engine/testing/attention_comparison.py index def7a449..10cbf029 100644 --- a/rl_engine/testing/attention_comparison.py +++ b/rl_engine/testing/attention_comparison.py @@ -270,9 +270,7 @@ def compare_decode_kv_replay( _validate_decode_inputs(inputs) if strict_bitwise: reference = _run_decode_strict_shared_core(inputs, materialization="logical_prefill") - candidates = [ - _run_decode_strict_shared_core(inputs, materialization="paged_kv_layout") - ] + candidates = [_run_decode_strict_shared_core(inputs, materialization="paged_kv_layout")] else: reference = _run_decode_full_prefill_reference(inputs) candidates = [_run_decode_kv_replay(inputs, merge_backend="rl_kernel")] diff --git a/tests/test_attention_ablation.py b/tests/test_attention_ablation.py index a77a97eb..3edb0a43 100644 --- a/tests/test_attention_ablation.py +++ b/tests/test_attention_ablation.py @@ -139,44 +139,6 @@ 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_wrapper_owned_deterministic_core_does_not_require_external_provenance(): q, k, v = _qkv() diff --git a/tests/test_attention_comparison.py b/tests/test_attention_comparison.py index f26f5f66..a8c16238 100644 --- a/tests/test_attention_comparison.py +++ b/tests/test_attention_comparison.py @@ -25,8 +25,8 @@ compare_single_gpu_attention, compare_single_gpu_rope_attention, decode_kv_cache_fingerprint, - decode_rope_identity_fingerprint, decode_prefix_cache_fingerprint, + decode_rope_identity_fingerprint, run_decode_kv_replay, run_paged_kv_attention, ) @@ -308,12 +308,8 @@ def test_strict_decode_replay_is_bitwise_across_logical_and_paged_materializatio assert drift.candidate_name == "strict_shared_core_paged_kv_layout" assert drift.out.max_abs == 0.0 assert drift.lse.max_abs == 0.0 - assert drift.provenance["strict_core_id"] == ( - "rlkernel.attention.deterministic_core.v1" - ) - assert drift.provenance["strict_schedule"] == ( - "single_batch_single_query_global_kv_blocks" - ) + assert drift.provenance["strict_core_id"] == ("rlkernel.attention.deterministic_core.v1") + assert drift.provenance["strict_schedule"] == ("single_batch_single_query_global_kv_blocks") assert drift.provenance["split_kv_policy"] == "disabled" diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 5e548999..c3aec008 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -14,11 +14,14 @@ import pytest import torch -from rl_engine.kernels.attention_contract import SplitKVSpec -from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( +from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_CORE_ID, STRICT_ATTENTION_SCHEDULE_ID, + SplitKVSpec, +) +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( AttentionPartialState, + AttentionSavedForwardState, DeterministicAttentionCore, DeterministicCPAttentionReferenceOp, compare_cp_attention_backward, @@ -410,6 +413,12 @@ def test_backward_report_cp2_prefill_matches_cp1_reference(): assert drift.provenance["merge_order"] == "global_block_index" assert drift.provenance["te_backward_oracle"] == "not_used" assert drift.provenance["decode_backward"] == "not_supported" + assert drift.provenance["projection_scope"] == "attention_core_only" + assert drift.provenance["qkv_projection_backward_dgrad_collective"] == "all_reduce" + assert drift.provenance["qkv_projection_sp_backward_collective"] == "reduce_scatter" + assert drift.provenance["o_proj_backward_dgrad_collective"] == "none" + assert drift.provenance["o_proj_sp_backward_collective"] == "all_gather" + assert drift.provenance["projection_collectives_executed"] is False json.dumps(report.to_dict()) @@ -473,6 +482,183 @@ def test_backward_report_cp2_chunked_prefill_matches_cp1_reference(): ] +def test_saved_forward_backward_matches_independent_dense_autograd(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 5, 7, seed=31, heads=4, kv_heads=2, dim=8) + mask = torch.tensor( + [[True, True, True, True, True, True, False], [True] * 7], + dtype=torch.bool, + ) + query_offsets = torch.tensor([11, 23], dtype=torch.long) + key_offsets = torch.tensor([9, 21], dtype=torch.long) + dout = torch.randn(q.shape, generator=torch.Generator().manual_seed(32)) + + state = op.save_forward_state( + q, + k, + v, + causal=True, + scale=0.37, + key_padding_mask=mask, + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + result = op.backward_reference( + q, + k, + v, + dout, + causal=True, + scale=0.37, + key_padding_mask=mask, + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + cp_world_size=2, + kv_chunk_size=2, + saved_forward_state=state, + ) + + q_ref = q.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + v_ref = v.detach().clone().requires_grad_(True) + k_expanded = k_ref.repeat_interleave(2, dim=1) + v_expanded = v_ref.repeat_interleave(2, dim=1) + scores = torch.matmul(q_ref, k_expanded.transpose(-1, -2)) * 0.37 + q_pos = query_offsets[:, None] + torch.arange(q.size(2)) + k_pos = key_offsets[:, None] + torch.arange(k.size(2)) + scores = scores.masked_fill( + (k_pos[:, None, :] > q_pos[:, :, None])[:, None, :, :], + float("-inf"), + ) + scores = scores.masked_fill(~mask[:, None, None, :], float("-inf")) + out_ref = torch.matmul(torch.softmax(scores, dim=-1), v_expanded) + out_ref.backward(dout) + + assert isinstance(result.saved_forward_state, AttentionSavedForwardState) + assert result.saved_forward_state is state + assert result.provenance["saved_forward_state_source"] == "caller" + torch.testing.assert_close(result.out, out_ref, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(result.gradients.dq, q_ref.grad, atol=_GRAD_ATOL, rtol=0.0) + torch.testing.assert_close(result.gradients.dk, k_ref.grad, atol=_GRAD_ATOL, rtol=0.0) + torch.testing.assert_close(result.gradients.dv, v_ref.grad, atol=_GRAD_ATOL, rtol=0.0) + + +@pytest.mark.parametrize( + ("tensor_name", "message"), + [("q", "q_fingerprint"), ("k", "k_fingerprint"), ("v", "v_fingerprint")], +) +def test_saved_forward_state_rejects_stale_qkv(tensor_name, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=33, heads=4, kv_heads=2, dim=8) + state = op.save_forward_state(q, k, v, cp_world_size=2, kv_chunk_size=2) + inputs = {"q": q.clone(), "k": k.clone(), "v": v.clone()} + inputs[tensor_name].flatten()[0] += 1.0 + + with pytest.raises(ValueError, match=message): + op.backward_reference( + inputs["q"], + inputs["k"], + inputs["v"], + torch.ones_like(q), + cp_world_size=2, + kv_chunk_size=2, + saved_forward_state=state, + ) + + +@pytest.mark.parametrize( + ("tensor_name", "message"), + [ + ("out", "out_fingerprint"), + ("lse", "lse_fingerprint"), + ("key_padding_mask", "key_padding_mask_fingerprint"), + ("query_position_offsets", "query_position_offsets_fingerprint"), + ("key_position_offsets", "key_position_offsets_fingerprint"), + ], +) +def test_saved_forward_state_rejects_mutated_saved_tensors(tensor_name, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=34, heads=4, kv_heads=2, dim=8) + mask = torch.ones(1, 4, dtype=torch.bool) + offsets = torch.tensor([7], dtype=torch.long) + state = op.save_forward_state( + q, + k, + v, + key_padding_mask=mask, + query_position_offsets=offsets, + key_position_offsets=offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + tensor = getattr(state, tensor_name) + if tensor.dtype == torch.bool: + tensor.flatten()[0].logical_not_() + else: + tensor.flatten()[0].add_(1) + + with pytest.raises(ValueError, match=message): + op.backward_reference( + q, + k, + v, + torch.ones_like(q), + key_padding_mask=mask, + query_position_offsets=offsets, + key_position_offsets=offsets, + cp_world_size=2, + kv_chunk_size=2, + saved_forward_state=state, + ) + + +@pytest.mark.parametrize( + ("override", "message"), + [ + ({"causal": False}, "causal"), + ({"scale": 0.5}, "scale"), + ({"cp_world_size": 1}, "cp_world_size"), + ({"kv_chunk_size": None}, "kv_chunk_size"), + ({"query_position_offsets": torch.tensor([8])}, "query_position_offsets"), + ({"key_position_offsets": torch.tensor([8])}, "key_position_offsets"), + ], +) +def test_saved_forward_state_rejects_execution_metadata_mismatch(override, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=35, heads=4, kv_heads=2, dim=8) + offsets = torch.tensor([7], dtype=torch.long) + state = op.save_forward_state( + q, + k, + v, + query_position_offsets=offsets, + key_position_offsets=offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + kwargs = { + "causal": True, + "scale": None, + "query_position_offsets": offsets, + "key_position_offsets": offsets, + "cp_world_size": 2, + "kv_chunk_size": 2, + } + kwargs.update(override) + + with pytest.raises(ValueError, match=message): + op.backward_reference( + q, + k, + v, + torch.ones_like(q), + saved_forward_state=state, + **kwargs, + ) + + def test_split_kv_plan_never_crosses_cp_owner_boundaries(): plans = split_kv_execution_plan_provenance( 10, @@ -737,6 +923,90 @@ def test_registry_dispatches_cp_attention_reference(): assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) +def test_strict_forward_is_bitwise_invariant_to_batch_cp_and_chunk(dtype): + q, k, v = _qkv(2, 5, 9, seed=41, dtype=dtype, heads=4, kv_heads=2, dim=8) + op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + + full_out, full_lse = op.forward_with_lse(q, k, v, cp_world_size=1) + chunked_out, chunked_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=4, + kv_chunk_size=1, + ) + assert torch.equal(full_out, chunked_out) + assert torch.equal(full_lse, chunked_lse) + assert torch.equal(full_out[:1], single_out) + assert torch.equal(full_lse[:1], single_lse) + + +def test_strict_backward_is_bitwise_invariant_to_batch_cp_and_chunk(): + q, k, v = _qkv(2, 5, 9, seed=42, heads=4, kv_heads=2, dim=8) + dout = torch.randn(q.shape, generator=torch.Generator().manual_seed(43)) + op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) + + full = op.backward_reference(q, k, v, dout, cp_world_size=1) + chunked = op.backward_reference( + q, + k, + v, + dout, + cp_world_size=2, + kv_chunk_size=3, + ) + single = op.backward_reference( + q[:1], + k[:1], + v[:1], + dout[:1], + cp_world_size=4, + kv_chunk_size=1, + ) + + for full_tensor, chunked_tensor, single_tensor in ( + (full.out, chunked.out, single.out), + (full.lse, chunked.lse, single.lse), + (full.gradients.dq, chunked.gradients.dq, single.gradients.dq), + (full.gradients.dk, chunked.gradients.dk, single.gradients.dk), + (full.gradients.dv, chunked.gradients.dv, single.gradients.dv), + ): + assert torch.equal(full_tensor, chunked_tensor) + assert torch.equal(full_tensor[:1], single_tensor) + + assert chunked.provenance["strict_core_id"] == STRICT_ATTENTION_CORE_ID + assert chunked.provenance["strict_schedule"] == STRICT_ATTENTION_SCHEDULE_ID + assert chunked.provenance["actual_split_kv_policy"] == "disabled" + assert chunked.provenance["backward_algorithm"] == ("saved_out_lse_canonical_row_reference") + assert all( + plan["actual_split_kv_policy"] == "disabled" + and plan["actual_split_boundaries"] == [[0, k.size(2)]] + for plan in chunked.provenance["actual_split_kv_plans"] + ) + + +def test_strict_backward_rejects_non_strict_saved_forward_state(): + q, k, v = _qkv(1, 4, 6, seed=44, heads=4, kv_heads=2, dim=8) + state = DeterministicCPAttentionReferenceOp().save_forward_state(q, k, v) + + with pytest.raises(ValueError, match="strict_bitwise"): + DeterministicCPAttentionReferenceOp(strict_bitwise=True).backward_reference( + q, + k, + v, + torch.ones_like(q), + saved_forward_state=state, + ) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float16]) def test_strict_core_is_bitwise_invariant_to_batch_and_cp_schedule(dtype): """The strict candidate must not change arithmetic with batch/CP shape.""" @@ -744,17 +1014,8 @@ def test_strict_core_is_bitwise_invariant_to_batch_and_cp_schedule(dtype): q, k, v = _qkv(2, 5, 9, seed=41, dtype=dtype, heads=4, kv_heads=2, dim=8) 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 - ) 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) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 7c0d9841..3b92af3b 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -47,7 +47,6 @@ def _args(**overrides): "matmul", "det_gemm", "attention", - "cp_attention", "logp", "linear_logp", "batch_invariant_logp", @@ -119,6 +118,8 @@ def test_constant_embedding_inputs_match_operator_contract(): inputs = make_operator_inputs("embedding", args, torch.float32, torch.device("cpu")) assert torch.equal(inputs["token_ids"], torch.full((1, 2), 3, dtype=torch.long)) + assert inputs["weight"].shape == (17, 128) + assert inputs["weight"].dtype is torch.float32 assert torch.equal(inputs["weight"], torch.full((17, 128), 0.5)) assert operator_shape_name("embedding", args) == "1x2x17x128" @@ -127,6 +128,10 @@ def test_constant_lm_head_inputs_match_operator_contract(): args = _args(input_mode="constant", constant_value=0.5) inputs = make_operator_inputs("lm_head", args, torch.float32, torch.device("cpu")) + assert inputs["hidden"].shape == (1, 2, 128) + assert inputs["weight"].shape == (17, 128) + assert inputs["hidden"].dtype is torch.float32 + assert inputs["weight"].dtype is torch.float32 assert torch.equal(inputs["hidden"], torch.full((1, 2, 128), 0.5)) assert torch.equal(inputs["weight"], torch.full((17, 128), 0.51)) assert inputs["bias"] is None diff --git a/tests/test_tolerance_contract.py b/tests/test_tolerance_contract.py index 778fab24..11af399a 100644 --- a/tests/test_tolerance_contract.py +++ b/tests/test_tolerance_contract.py @@ -6,15 +6,11 @@ from __future__ import annotations import copy -import hashlib -import inspect -import json import math import pytest import torch -from rl_engine.kernels.gtest import tolerance as tolerance_module from rl_engine.kernels.gtest.tolerance import ( CHAIN_AGGREGATE_METRICS, JUDGMENTS, @@ -30,10 +26,8 @@ resolve_chain_aggregate_thresholds, resolve_comparison_roles, resolve_dtype_policy, - resolve_logprob_threshold, resolve_tolerance, resolve_tolerance_support, - tolerance_contract_fingerprint, validate_backend_provenance, validate_contract_schema, )