diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7ada3eb6..5006574c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -96,6 +96,10 @@ jobs: run: | python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu" +<<<<<<< HEAD + - name: Run WS2 Attention Contract Tests (CPU-safe) + run: python -m pytest tests/test_attention_contract.py -v +======= - name: Run WS2 Logprob Contract Tests (CPU-safe) run: python -m pytest tests/test_logprob_contract.py -v @@ -107,6 +111,7 @@ jobs: python -m pytest -q \ tests/test_alignment_wrapper_interfaces.py \ tests/test_qwen_ffn.py +>>>>>>> origin/test docs: runs-on: ubuntu-latest diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index 23c1586a..c8bf3d14 100644 --- a/docs/design/runtime-dispatch.md +++ b/docs/design/runtime-dispatch.md @@ -11,6 +11,11 @@ logical type, and the registry selects the first available backend for the curre 4. Cache successfully constructed operator instances. 5. Skip backends that already failed in the current process. +WS2 Attention uses the stricter `KernelRegistry.get_attention_op(contract)` path. In addition to +platform priority, this path requires a backend capability descriptor and checks the requested +role, mode, dtype, TP/CP layout, LSE export, deterministic merge, packed varlen, and KV-cache +semantics. Incompatible candidates produce explicit rejection reasons and are never used as an +undeclared fallback. See [WS2 CP-aware Attention contract](ws2-cp-attention-contract.md). WS2 TP-aware logprob uses the stricter `KernelRegistry.get_logprob_op(contract)` path. In addition to platform priority, this path requires a backend capability descriptor and checks the requested role, dtype, TP/CP layout, padded-vs-real vocab masking, inactive-token diff --git a/docs/operators/attention.md b/docs/operators/attention.md index ebff9a58..2b9b4817 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -84,6 +84,31 @@ 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). + +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 @@ -137,6 +162,7 @@ memory. ```bash python -m pytest tests/test_attention.py -v +python -m pytest tests/test_cp_attention.py -v ``` Covers: `forward_fp32` vs an independent fp32 reference (bitwise), strict-fp32 under hostile @@ -194,6 +220,10 @@ 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. +- `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. ## Tolerance diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 8b5f588b..29faedda 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -18,6 +18,8 @@ _EnumT = TypeVar("_EnumT", bound=Enum) +# Stable identity shared by the training and rollout deterministic Attention +# core. Backend adapters may differ, but strict mode must report this ID. STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks" @@ -65,6 +67,13 @@ class SplitKVMode(str, Enum): AUTO = "auto" +class ProjectionCollective(str, Enum): + NONE = "none" + ALL_REDUCE = "all_reduce" + ALL_GATHER = "all_gather" + REDUCE_SCATTER = "reduce_scatter" + + class RoPEState(str, Enum): PRE_ROPE = "pre_rope" POST_ROPE = "post_rope" @@ -138,10 +147,16 @@ class ShardingSpec: global_block_token_starts: tuple[int, ...] local_block_offsets: tuple[int, ...] packed_sequence_offsets: tuple[int, ...] | None = None + sp_rank: int = 0 + sp_world_size: int = 1 def __post_init__(self) -> None: tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + sp_world_size = _positive_int(self.sp_world_size, "sp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + sp_rank = _non_negative_int(self.sp_rank, "sp_rank") tp_rank = _non_negative_int(self.tp_rank, "tp_rank") cp_rank = _non_negative_int(self.cp_rank, "cp_rank") if tp_rank >= tp_world_size: @@ -152,6 +167,10 @@ def __post_init__(self) -> None: raise AttentionContractError( f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" ) + if sp_rank >= sp_world_size: + raise AttentionContractError( + f"sp_rank={sp_rank} must be smaller than sp_world_size={sp_world_size}" + ) global_q_heads = _positive_int(self.global_q_heads, "global_q_heads") global_kv_heads = _positive_int(self.global_kv_heads, "global_kv_heads") @@ -1153,6 +1172,98 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True) +class AttentionProjectionSpec: + """Deterministic QKV or output-projection execution contract.""" + + name: str + input_dtype: AttentionDType = AttentionDType.BF16 + output_dtype: AttentionDType = AttentionDType.BF16 + acc_dtype: AttentionDType = AttentionDType.FP32 + split_kv: SplitKVMode = SplitKVMode.DISABLED + k_order: str = "ascending" + backend_policy: str = "native_verified_then_common_deterministic" + deterministic_backend: str = "rlkernel.cuda.det_gemm" + tp_forward_collective: ProjectionCollective = ProjectionCollective.NONE + tp_backward_dgrad_collective: ProjectionCollective = ProjectionCollective.NONE + sp_forward_collective: ProjectionCollective = ProjectionCollective.NONE + sp_backward_collective: ProjectionCollective = ProjectionCollective.NONE + qkv_split_order: tuple[str, ...] = () + require_runtime_readback: bool = True + + @classmethod + def qkv(cls) -> "AttentionProjectionSpec": + return cls( + name="qkv", + tp_backward_dgrad_collective=ProjectionCollective.ALL_REDUCE, + sp_forward_collective=ProjectionCollective.ALL_GATHER, + sp_backward_collective=ProjectionCollective.REDUCE_SCATTER, + qkv_split_order=("q", "k", "v"), + ) + + @classmethod + def output(cls) -> "AttentionProjectionSpec": + return cls( + name="o_proj", + tp_forward_collective=ProjectionCollective.ALL_REDUCE, + sp_forward_collective=ProjectionCollective.REDUCE_SCATTER, + sp_backward_collective=ProjectionCollective.ALL_GATHER, + ) + + def __post_init__(self) -> None: + if self.name not in {"qkv", "o_proj"}: + raise AttentionContractError("projection name must be qkv or o_proj") + for field_name in ("input_dtype", "output_dtype", "acc_dtype"): + object.__setattr__( + self, + field_name, + _enum_value(AttentionDType, getattr(self, field_name), field_name), + ) + object.__setattr__( + self, + "split_kv", + _enum_value(SplitKVMode, self.split_kv, "projection.split_kv"), + ) + for field_name in ( + "tp_forward_collective", + "tp_backward_dgrad_collective", + "sp_forward_collective", + "sp_backward_collective", + ): + object.__setattr__( + self, + field_name, + _enum_value( + ProjectionCollective, + getattr(self, field_name), + field_name, + ), + ) + if self.input_dtype is not AttentionDType.BF16: + raise AttentionContractError("Attention projections require BF16 input") + if self.output_dtype is not AttentionDType.BF16: + raise AttentionContractError("Attention projections require BF16 output") + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError("Attention projections require FP32 accumulation") + if self.split_kv is not SplitKVMode.DISABLED: + raise AttentionContractError("Attention projection GEMMs must disable Split-K") + if self.k_order != "ascending": + raise AttentionContractError("Attention projection GEMMs require ascending K order") + for field_name in ("backend_policy", "deterministic_backend"): + value = getattr(self, field_name) + if not isinstance(value, str) or not value.strip(): + raise AttentionContractError(f"{field_name} must be a non-empty string") + if not isinstance(self.require_runtime_readback, bool) or not self.require_runtime_readback: + raise AttentionContractError("projection runtime readback must be required") + split_order = tuple(self.qkv_split_order) + if self.name == "qkv": + if split_order != ("q", "k", "v"): + raise AttentionContractError("QKV projection must split in Q, K, V order") + elif split_order: + raise AttentionContractError("o_proj must not declare a QKV split order") + object.__setattr__(self, "qkv_split_order", split_order) + + @dataclass(frozen=True) class AttentionContract: """Complete semantic request consumed by contract-aware dispatch.""" @@ -1171,6 +1282,10 @@ class AttentionContract: kv_cache: KVCacheSpec | None = None rope: RoPESpec | None = None export_lse: bool = True + qkv_projection: AttentionProjectionSpec = field(default_factory=AttentionProjectionSpec.qkv) + output_projection: AttentionProjectionSpec = field( + default_factory=AttentionProjectionSpec.output + ) def __post_init__(self) -> None: object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role")) @@ -1185,6 +1300,14 @@ def __post_init__(self) -> None: raise AttentionContractError("reduction must be a ReductionSpec") if not isinstance(self.split_kv, SplitKVSpec): raise AttentionContractError("split_kv must be a SplitKVSpec") + if not isinstance(self.qkv_projection, AttentionProjectionSpec): + raise AttentionContractError("qkv_projection must be an AttentionProjectionSpec") + if self.qkv_projection.name != "qkv": + raise AttentionContractError("qkv_projection must use name='qkv'") + if not isinstance(self.output_projection, AttentionProjectionSpec): + raise AttentionContractError("output_projection must be an AttentionProjectionSpec") + if self.output_projection.name != "o_proj": + raise AttentionContractError("output_projection must use name='o_proj'") if ( self.mode is AttentionMode.PREFILL and query_sequence_length != self.sharding.local_sequence_length @@ -1269,6 +1392,8 @@ def to_dict(self) -> dict[str, Any]: "tp_world_size": self.sharding.tp_world_size, "cp_rank": self.sharding.cp_rank, "cp_world_size": self.sharding.cp_world_size, + "sp_rank": self.sharding.sp_rank, + "sp_world_size": self.sharding.sp_world_size, "global_q_heads": self.sharding.global_q_heads, "global_kv_heads": self.sharding.global_kv_heads, "local_q_head_start": self.sharding.local_q_head_start, @@ -1331,6 +1456,24 @@ def to_dict(self) -> dict[str, Any]: "output_dtype": self.rope.output_dtype.value, "fusion_boundary": self.rope.fusion_boundary.value, } + projections = { + spec.name: { + "input_dtype": spec.input_dtype.value, + "output_dtype": spec.output_dtype.value, + "acc_dtype": spec.acc_dtype.value, + "split_kv": spec.split_kv.value, + "k_order": spec.k_order, + "backend_policy": spec.backend_policy, + "deterministic_backend": spec.deterministic_backend, + "tp_forward_collective": spec.tp_forward_collective.value, + "tp_backward_dgrad_collective": spec.tp_backward_dgrad_collective.value, + "sp_forward_collective": spec.sp_forward_collective.value, + "sp_backward_collective": spec.sp_backward_collective.value, + "qkv_split_order": list(spec.qkv_split_order), + "require_runtime_readback": spec.require_runtime_readback, + } + for spec in (self.qkv_projection, self.output_projection) + } return { "semantic_operator": "standard_softmax_attention", "role": self.role.value, @@ -1350,6 +1493,7 @@ def to_dict(self) -> dict[str, Any]: "split_kv": self.split_kv.to_dict(), "kv_cache": kv_cache, "rope": rope, + "projections": projections, } @@ -1505,6 +1649,7 @@ class AttentionDispatchResult: "AttentionBackendCapability", "AttentionDispatchResult", "AttentionDType", + "AttentionProjectionSpec", "AttentionMerge", "AttentionMode", "AttentionRole", @@ -1513,6 +1658,7 @@ class AttentionDispatchResult: "ReductionEngine", "ReductionOrder", "ReductionSpec", + "ProjectionCollective", "RoPECastPoint", "RoPEFusionBoundary", "RoPESpec", diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index f37fce36..07d6d448 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -31,6 +31,7 @@ def make_operator_inputs( "matmul": _make_matmul_inputs, "det_gemm": _make_det_gemm_inputs, "attention": _make_attention_inputs, + "cp_attention": _make_attention_inputs, "logp": _make_logp_inputs, "linear_logp": _make_linear_logp_inputs, "batch_invariant_logp": _make_batch_invariant_logp_inputs, @@ -58,6 +59,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "matmul": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "det_gemm": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", + "cp_attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}xcp2", "logp": f"{batch}x{seq}x{vocab}", "linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "batch_invariant_logp": f"{batch}x{seq}x{vocab}", diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index bde87edb..ca4a462e 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -76,6 +76,22 @@ def _load_object(path: str) -> Any: }, grad_input_names=("q", "k", "v"), ), + "cp_attention": OperatorSpec( + name="cp_attention", + op_class="attention", + gold_path=( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ), + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ), + }, + grad_input_names=("q", "k", "v"), + ), "logp": OperatorSpec( name="logp", op_class="logprob", diff --git a/rl_engine/kernels/ops/cuda/attention/__init__.py b/rl_engine/kernels/ops/cuda/attention/__init__.py index 09775c8e..8ca8df7b 100644 --- a/rl_engine/kernels/ops/cuda/attention/__init__.py +++ b/rl_engine/kernels/ops/cuda/attention/__init__.py @@ -1,11 +1,62 @@ -# File: rl_engine/kernels/ops/cuda/attention/__init__.py - from .deterministic_attn import DeterministicAttentionOp from .flash_attn import FlashAttentionOp from .prefix_shared_attn import PrefixSharedAttentionOp -__all__ = [ - "DeterministicAttentionOp", - "FlashAttentionOp", - "PrefixSharedAttentionOp", -] +__all__ = ["DeterministicAttentionOp", "FlashAttentionOp", "PrefixSharedAttentionOp"] + +# CP communication and FlashInfer are optional layers owned by later WS2 PRs. +# Keep the base Attention package importable while those PRs are developed or +# tested independently, then expose their symbols automatically when present. +try: + from .cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunication, + AttentionCPCommunicationPlan, + AttentionCPCommunicationUnavailable, + AttentionCPMergedState, + AttentionCPPartialState, + AttentionParallelSpec, + CPCommunicationBackend, + CPCommunicationStatus, + CUDAAGRSAttentionCPCommunication, + P2PNCCLAttentionCPCommunication, + sort_attention_cp_partial_states, + ) +except ModuleNotFoundError as exc: + if exc.name != f"{__package__}.cp_comm": + raise +else: + __all__ += [ + "AttentionCPBlockMetadata", + "AttentionCPCommunication", + "AttentionCPCommunicationPlan", + "AttentionCPCommunicationUnavailable", + "AttentionCPMergedState", + "AttentionCPPartialState", + "AttentionParallelSpec", + "CPCommunicationBackend", + "CPCommunicationStatus", + "CUDAAGRSAttentionCPCommunication", + "P2PNCCLAttentionCPCommunication", + "sort_attention_cp_partial_states", + ] + +try: + from .flashinfer_paged_attention import ( + FlashInferPagedAttentionConfig, + FlashInferQwen3PagedAttentionOp, + FlashInferRoPEFusionConfig, + FlashInferSplitKVPolicy, + FlashInferUnavailable, + ) +except ModuleNotFoundError as exc: + if exc.name != f"{__package__}.flashinfer_paged_attention": + raise +else: + __all__ += [ + "FlashInferPagedAttentionConfig", + "FlashInferQwen3PagedAttentionOp", + "FlashInferRoPEFusionConfig", + "FlashInferSplitKVPolicy", + "FlashInferUnavailable", + ] diff --git a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py index 01c51c99..1e3d01e2 100644 --- a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py @@ -17,6 +17,10 @@ from torch.autograd import Function from torch.autograd.function import once_differentiable +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, +) from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE from rl_engine.utils.logger import logger @@ -82,6 +86,10 @@ class DeterministicAttentionOp: so #108 harness can call forward(**inputs) with key_padding_mask. """ + core_id = STRICT_ATTENTION_CORE_ID + strict_schedule = STRICT_ATTENTION_SCHEDULE_ID + backend_id = "rlkernel.cuda.deterministic_attention" + def __init__(self) -> None: if not _EXT_AVAILABLE or not hasattr(_C, "deterministic_attention_forward"): raise RuntimeError( diff --git a/tests/test_attention.py b/tests/test_attention.py index 469c6d30..40ee6fd5 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -443,6 +443,18 @@ def test_registry_dispatches_native_attention_op(): assert isinstance(op, (NativeAttentionOp, DeterministicAttentionOp)) +def test_deterministic_attention_op_exposes_shared_strict_identity(): + from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + ) + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionOp + + assert DeterministicAttentionOp.core_id == STRICT_ATTENTION_CORE_ID + assert DeterministicAttentionOp.strict_schedule == STRICT_ATTENTION_SCHEDULE_ID + assert DeterministicAttentionOp.backend_id == "rlkernel.cuda.deterministic_attention" + + # --------------------------------------------------------------------------- # # Qwen3-8B LARGE real-scale GPU smoke test # --------------------------------------------------------------------------- # diff --git a/tests/test_attention_contract.py b/tests/test_attention_contract.py index 350545a1..eabddea9 100644 --- a/tests/test_attention_contract.py +++ b/tests/test_attention_contract.py @@ -11,11 +11,17 @@ import pytest from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, AttentionBackendCapability, AttentionContract, AttentionContractError, AttentionDType, AttentionMode, + AttentionProjectionSpec, + AttentionRole, + KVCacheSpec, + ProjectionCollective, AttentionRole, KVCacheSpec, ReductionSpec, @@ -32,6 +38,11 @@ from rl_engine.kernels.registry import KernelRegistry, OpBackend +def test_strict_attention_identity_pins_core_and_arithmetic_schedule(): + assert STRICT_ATTENTION_CORE_ID == "rlkernel.attention.deterministic_core.v1" + assert STRICT_ATTENTION_SCHEDULE_ID == "single_batch_single_query_global_kv_blocks" + + def _sharding( *, tp_rank: int = 0, @@ -44,6 +55,8 @@ def _sharding( global_block_token_starts: tuple[int, ...] = (0,), local_block_offsets: tuple[int, ...] = (0, 2048), packed_sequence_offsets: tuple[int, ...] | None = None, + sp_rank: int = 0, + sp_world_size: int = 1, ) -> ShardingSpec: local_q_heads = 32 // tp_world_size local_kv_heads = 8 // tp_world_size @@ -64,6 +77,8 @@ def _sharding( global_block_token_starts=global_block_token_starts, local_block_offsets=local_block_offsets, packed_sequence_offsets=packed_sequence_offsets, + sp_rank=sp_rank, + sp_world_size=sp_world_size, ) @@ -681,3 +696,48 @@ def test_packed_sequence_count_must_match_logical_batch_size(): contract = _contract(sharding=sharding, causal_offsets=(0, 0), batch_size=2) assert contract.batch_size == 2 + + +def test_projection_contract_pins_gemm_and_tp_sp_collectives(): + contract = _contract(sharding=_sharding(sp_rank=1, sp_world_size=2)) + provenance = contract.to_dict() + + assert provenance["sharding"]["sp_rank"] == 1 + assert provenance["sharding"]["sp_world_size"] == 2 + assert provenance["projections"]["qkv"] == { + "input_dtype": "bf16", + "output_dtype": "bf16", + "acc_dtype": "fp32", + "split_kv": "disabled", + "k_order": "ascending", + "backend_policy": "native_verified_then_common_deterministic", + "deterministic_backend": "rlkernel.cuda.det_gemm", + "tp_forward_collective": "none", + "tp_backward_dgrad_collective": "all_reduce", + "sp_forward_collective": "all_gather", + "sp_backward_collective": "reduce_scatter", + "qkv_split_order": ["q", "k", "v"], + "require_runtime_readback": True, + } + assert provenance["projections"]["o_proj"]["tp_forward_collective"] == "all_reduce" + assert provenance["projections"]["o_proj"]["sp_forward_collective"] == "reduce_scatter" + assert provenance["projections"]["o_proj"]["sp_backward_collective"] == "all_gather" + + +def test_projection_contract_rejects_split_k_and_wrong_collective_identity(): + with pytest.raises(AttentionContractError, match="disable Split-K"): + replace(AttentionProjectionSpec.qkv(), split_kv="fixed") + + with pytest.raises(AttentionContractError, match="qkv_projection"): + replace( + _contract(), + qkv_projection=replace( + AttentionProjectionSpec.output(), + tp_forward_collective=ProjectionCollective.NONE, + ), + ) + + +def test_sharding_rejects_out_of_range_sp_rank(): + with pytest.raises(AttentionContractError, match="sp_rank=2"): + _sharding(sp_rank=2, sp_world_size=2)