diff --git a/docs/operators/attention.md b/docs/operators/attention.md index e3ba7196..32b8767f 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 @@ -95,6 +101,34 @@ Existing WS1 implementations do not yet export attention-domain LSE or implement CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md). +### WS2 deterministic CP reference + +### WS2 CP-aware dispatch + +WS2 distributed callers use a separate contract-aware entry point, +`kernel_registry.get_attention_op(contract)`. It validates explicit TP/CP ownership, fixed +`(out, lse)` merge semantics, causal or packed-sequence offsets, and decode KV-cache identity +before selecting a backend. Legacy `get_op("attention")` behavior remains unchanged. + +Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic +CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as +a silent fallback. See [WS2 CP-aware Attention contract](../design/ws2-cp-attention-contract.md). + +Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow +`disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block +boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason. +Runtime-selected `auto` plans are diagnostic only unless both training and rollout export and +validate the same actual plan. + +The rank-aware drift benchmark can emit a CPU smoke artifact or a torchrun-friendly GPU report: + +```bash +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 \ + --cp-world-sizes 2 --kv-chunk-sizes none,1 --include-backward \ + --output artifacts/ws2-cp-attention-drift.json +``` + Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow `disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason. diff --git a/rl_engine/kernels/gtest/tolerance.py b/rl_engine/kernels/gtest/tolerance.py index 4fb5bcbf..83aa49a0 100644 --- a/rl_engine/kernels/gtest/tolerance.py +++ b/rl_engine/kernels/gtest/tolerance.py @@ -202,6 +202,62 @@ 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.""" @@ -1056,6 +1112,8 @@ 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/tests/test_attention_ablation.py b/tests/test_attention_ablation.py index e07201da..a77a97eb 100644 --- a/tests/test_attention_ablation.py +++ b/tests/test_attention_ablation.py @@ -139,6 +139,44 @@ 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() @@ -279,6 +317,13 @@ def __call__(self, q, k, v, *, causal, scale): 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)) + with pytest.raises(AttentionContractError, match="Split-KV"): + AttentionAblationOp()(q, k, v, contract=contract) + + @pytest.mark.parametrize( ("missing_field", "replacement"), [ diff --git a/tests/test_attention_preprocess.py b/tests/test_attention_preprocess.py index 9b800dca..1abc4f2e 100644 --- a/tests/test_attention_preprocess.py +++ b/tests/test_attention_preprocess.py @@ -152,6 +152,7 @@ def _inputs(): @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) result = H100AttentionPreprocessor(reuse_transformer_engine_qk_norm=False)( q, k, q_weight, k_weight, positions ) @@ -179,6 +180,7 @@ def test_h100_preprocessor_executes_cuda_qk_norm_and_zigzag_rope(): @requires_h100_preprocess def test_h100_preprocessor_is_bitwise_batch_invariant_for_2d_positions(): q, k, q_weight, k_weight, positions = _inputs() + op = H100AttentionPreprocessor() op = H100AttentionPreprocessor(reuse_transformer_engine_qk_norm=False) full = op(q, k, q_weight, k_weight, positions) diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py index 44cc81bd..5e548999 100644 --- a/tests/test_cp_attention.py +++ b/tests/test_cp_attention.py @@ -745,18 +745,15 @@ def test_strict_core_is_bitwise_invariant_to_batch_and_cp_schedule(dtype): op = DeterministicCPAttentionReferenceOp(strict_bitwise=True) cp1_out, cp1_lse = op.forward_with_lse( - q, - k, - v, - cp_world_size=1, - kv_chunk_size=None, + 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, + 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_rms_norm.py b/tests/test_rms_norm.py index 43a9cf85..6572603e 100644 --- a/tests/test_rms_norm.py +++ b/tests/test_rms_norm.py @@ -5,14 +5,17 @@ import torch import torch.nn.functional as F -from rl_engine.kernels.ops.cuda.norm.rmsnorm import rmsnorm_cuda +from rl_engine.kernels.ops.cuda.norm.rmsnorm import RMSNormCudaOp, rmsnorm_cuda from rl_engine.kernels.ops.pytorch.norm.rms_norm import NativeRMSNormOp from rl_engine.kernels.ops.triton.rmsnorm_triton import rmsnorm_triton try: from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE - _HAS_CUDA_RMSNORM = _EXT_AVAILABLE and hasattr(_C, "rmsnorm_forward") + _HAS_CUDA_RMSNORM = _EXT_AVAILABLE and all( + hasattr(_C, name) + for name in ("rmsnorm_forward", "rmsnorm_backward_dx", "rmsnorm_backward_dw") + ) except ImportError: # pragma: no cover - import can fail when the extension is not built. _HAS_CUDA_RMSNORM = False @@ -236,8 +239,12 @@ def test_registry_dispatches_rms_norm(): from rl_engine.kernels.registry import kernel_registry op = kernel_registry.get_op("rms_norm") - assert isinstance(op, NativeRMSNormOp) - assert hasattr(op, "forward") and hasattr(op, "forward_fp32") + if torch.cuda.is_available() and _HAS_CUDA_RMSNORM: + assert isinstance(op, RMSNormCudaOp) + assert hasattr(op, "forward") + else: + assert isinstance(op, NativeRMSNormOp) + assert hasattr(op, "forward") and hasattr(op, "forward_fp32") @requires_cuda diff --git a/tests/test_tolerance_contract.py b/tests/test_tolerance_contract.py index 11af399a..778fab24 100644 --- a/tests/test_tolerance_contract.py +++ b/tests/test_tolerance_contract.py @@ -6,11 +6,15 @@ 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, @@ -26,8 +30,10 @@ 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, )