diff --git a/benchmarks/benchmark_ws2_cp_attention_drift.py b/benchmarks/benchmark_ws2_cp_attention_drift.py new file mode 100644 index 00000000..48cf9be0 --- /dev/null +++ b/benchmarks/benchmark_ws2_cp_attention_drift.py @@ -0,0 +1,1133 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""WS2 CP attention drift benchmark and report artifact generator. + +This is the PR5 artifact path for issue #235. It is intentionally rank-aware +and torchrun-friendly, but the correctness surface remains the deterministic +PyTorch CP reference. The benchmark can run as a CPU smoke test on one process +or under torchrun; rank 0 writes the shared JSON report. +""" + +from __future__ import annotations + +import argparse +import contextlib +import datetime as _datetime +import hashlib +import importlib +import importlib.metadata +import json +import os +import platform +import shlex +import sys +from pathlib import Path +from typing import Any, Iterator, Sequence + +import torch + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + build_reference_split_kv_runtime_plan_set, + compare_cp_attention_backward, + merge_attention_partial_states, + split_kv_execution_plan_provenance, +) +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionCPMergedState, + AttentionCPPartialState, + AttentionParallelSpec, + P2PNCCLAttentionCPCommunication, +) +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp + +SCHEMA_VERSION = "ws2_cp_attention_drift/v1" +ISSUE = 235 +PR = 5 +DEFAULT_SEQ_LEN = 16 +QWEN3_8B_HEADS = 32 +QWEN3_8B_KV_HEADS = 8 +QWEN3_8B_HEAD_DIM = 128 +QWEN3_8B_ROPE_THETA = 1_000_000.0 +TE_CONTEXT_PARALLEL_MODULE = ( + "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" +) +TE_SYMBOLS = ( + "flash_attn_fwd_softmax_lse_correction", + "flash_attn_fwd_out_correction_init", + "flash_attn_fwd_out_correction", +) + + +class TEContextParallelMergeAdapter: + """Optional Transformer Engine CP merge oracle used only by PR5 reports.""" + + def __init__(self, module: Any, *, version: str) -> None: + self._module = module + self.version = version + + @classmethod + def probe(cls) -> tuple["TEContextParallelMergeAdapter | None", dict[str, object]]: + status: dict[str, object] = { + "te_available": False, + "te_version": None, + "te_module": TE_CONTEXT_PARALLEL_MODULE, + "te_symbols": list(TE_SYMBOLS), + "te_capability_probe": "unavailable", + "te_signature_checked": False, + "te_numeric_selftest": "not_run", + "fallback": True, + "fallback_reason": None, + } + try: + module = importlib.import_module(TE_CONTEXT_PARALLEL_MODULE) + version = _transformer_engine_version() + missing = [name for name in TE_SYMBOLS if not hasattr(module, name)] + if missing: + status.update( + { + "te_version": version, + "te_capability_probe": "missing_symbols", + "fallback_reason": f"missing symbols: {', '.join(missing)}", + } + ) + return None, status + adapter = cls(module, version=version) + status.update( + { + "te_available": True, + "te_version": version, + "te_signature_checked": True, + } + ) + adapter._numeric_selftest() + except ( + ImportError, + OSError, + RuntimeError, + AttributeError, + TypeError, + AssertionError, + ) as exc: + status.update( + { + "te_capability_probe": "failed", + "te_numeric_selftest": "failed", + "fallback_reason": str(exc), + } + ) + return None, status + + status.update( + { + "te_capability_probe": "passed", + "te_numeric_selftest": "passed", + "fallback": False, + "fallback_reason": None, + } + ) + return adapter, status + + def merge(self, states: Sequence[AttentionPartialState]) -> AttentionPartialState: + if not states: + raise ValueError("at least one partial state is required") + ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) + if len(ordered) == 1: + state = ordered[0] + return AttentionPartialState( + out=state.out.float().clone(), + lse=state.lse.float().clone(), + block_start=state.block_start, + block_end=state.block_end, + ) + + merged_lse = ordered[0].lse.float().clone() + merged_out = ordered[0].out.float().clone() + for state in ordered[1:]: + next_lse = state.lse.float() + previous_lse = merged_lse.clone() + self._module.flash_attn_fwd_softmax_lse_correction(merged_lse, next_lse) + merged_out = self._module.flash_attn_fwd_out_correction_init( + merged_out, + merged_lse, + previous_lse, + seq_dim=2, + ) + self._module.flash_attn_fwd_out_correction( + merged_out, + state.out.float(), + merged_lse, + next_lse, + seq_dim=2, + ) + + return AttentionPartialState( + out=merged_out, + lse=merged_lse, + block_start=ordered[0].block_start, + block_end=ordered[-1].block_end, + ) + + def _numeric_selftest(self) -> None: + gen = torch.Generator().manual_seed(5) + states = [ + AttentionPartialState( + out=torch.randn(1, 2, 3, 4, generator=gen), + lse=torch.randn(1, 2, 3, generator=gen), + block_start=0, + block_end=2, + ), + AttentionPartialState( + out=torch.randn(1, 2, 3, 4, generator=gen), + lse=torch.randn(1, 2, 3, generator=gen), + block_start=2, + block_end=5, + ), + ] + ours = merge_attention_partial_states(states) + te = self.merge(states) + torch.testing.assert_close(te.lse, ours.lse, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(te.out, ours.out, atol=1.0e-6, rtol=0.0) + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run the WS2 CP attention drift benchmark for issue #235 PR5." + ) + parser.add_argument("--model", default="qwen3-8b", choices=["qwen3-8b"]) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seq-len", type=int, default=DEFAULT_SEQ_LEN) + parser.add_argument("--seed", type=int, default=2355) + parser.add_argument("--device", choices=["auto", "cpu", "cuda"], default="cpu") + parser.add_argument("--dtype", choices=["bf16", "fp32"], default="bf16") + parser.add_argument("--tp-world-sizes", default="1,2") + parser.add_argument("--cp-world-sizes", default="1,2") + parser.add_argument( + "--kv-chunk-sizes", + default="none,4", + help="Comma list such as 'none,4'. 'none' means full prefill.", + ) + parser.add_argument("--smoke", action="store_true", help="Use a tiny CPU-friendly shape.") + parser.add_argument( + "--include-backward", + action="store_true", + help="Include optional PR8 dq/dk/dv drift fields.", + ) + parser.add_argument( + "--no-rope", + action="store_false", + dest="compose_rope", + help="Disable the pre-attention RoPE composition step.", + ) + parser.set_defaults(compose_rope=True) + parser.add_argument("--num-threads", type=int, default=1) + parser.add_argument( + "--init-process-group", + action="store_true", + help="Initialize torch.distributed from torchrun env vars before benchmarking.", + ) + parser.add_argument("--output", type=Path, help="Optional JSON artifact path.") + parser.add_argument("--json", action="store_true", help="Print the JSON report on rank 0.") + return parser.parse_args(argv) + + +def run_benchmark(args: argparse.Namespace) -> dict[str, object]: + rank_env = _rank_env() + device = _resolve_device(args.device, rank_env) + _validate_args(args) + distributed = _maybe_init_process_group(args, device, rank_env) + te_adapter, te_status = TEContextParallelMergeAdapter.probe() + seq_len = 4 if args.smoke and args.seq_len == DEFAULT_SEQ_LEN else args.seq_len + kv_chunk_sizes = _parse_kv_chunk_sizes(args.kv_chunk_sizes) + if args.smoke and args.kv_chunk_sizes == "none,4": + kv_chunk_sizes = (None, 1) + + cases: list[dict[str, object]] = [] + report: dict[str, object] = { + "schema_version": SCHEMA_VERSION, + "report_family": "ws2_cross_config_drift_report", + "tolerance_source": "#108", + "issue": ISSUE, + "pr": PR, + "created_at_utc": _datetime.datetime.now(_datetime.UTC).isoformat(), + "launch": _launch_metadata(rank_env), + "runtime": _runtime_metadata(device, distributed, rank_env), + "target": { + "model": args.model, + "global_num_query_heads": QWEN3_8B_HEADS, + "global_num_kv_heads": QWEN3_8B_KV_HEADS, + "head_dim": QWEN3_8B_HEAD_DIM, + "dtype": args.dtype, + "batch": args.batch, + "seq_len": seq_len, + "causal": True, + }, + "te_context_parallel_merge": te_status, + "dlogp": { + "status": "not_available", + "reason": "selected-logprob chain integration is outside PR5 benchmark scope", + }, + "cases": cases, + } + + try: + with _thread_limit(args.num_threads): + for tp_world_size in _parse_int_csv(args.tp_world_sizes, name="tp_world_sizes"): + for cp_world_size in _parse_int_csv(args.cp_world_sizes, name="cp_world_sizes"): + for kv_chunk_size in kv_chunk_sizes: + cases.append( + _run_case( + args, + device=device, + seq_len=seq_len, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + te_adapter=te_adapter, + ) + ) + finally: + if distributed["initialized"]: + import torch.distributed as dist + + if sys.exc_info()[0] is None: + dist.barrier() + dist.destroy_process_group() + return report + + +def write_report(report: dict[str, object], output: Path) -> None: + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + report = run_benchmark(args) + rank = int(report["launch"]["rank"]) + if rank == 0: + if args.output is not None: + write_report(report, args.output) + if args.json: + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 + + +def _run_case( + args: argparse.Namespace, + *, + device: torch.device, + seq_len: int, + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: int | None, + te_adapter: TEContextParallelMergeAdapter | None, +) -> dict[str, object]: + _validate_topology(tp_world_size, cp_world_size) + dtype = _dtype_from_name(args.dtype) + local_hq = QWEN3_8B_HEADS // tp_world_size + local_hkv = QWEN3_8B_KV_HEADS // tp_world_size + case_seed = _case_seed(args.seed, tp_world_size, cp_world_size, kv_chunk_size) + q, k, v, rope_report = _make_qkv( + batch=args.batch, + local_hq=local_hq, + local_hkv=local_hkv, + seq_len=seq_len, + dtype=dtype, + device=device, + seed=case_seed, + compose_rope=args.compose_rope, + ) + dout = _make_dout( + batch=args.batch, + local_hq=local_hq, + seq_len=seq_len, + dtype=dtype, + device=device, + seed=case_seed + 17, + ) + attention = DeterministicCPAttentionReferenceOp() + reference_out, reference_lse = attention.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=1, + kv_chunk_size=None, + ) + candidate_fp32_out, candidate_fp32_lse = attention.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + candidate_dtype_out, candidate_dtype_lse = attention.forward_with_lse( + q, + k, + v, + causal=True, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + output_dtype=dtype, + ) + + split_kv_policy = "disabled" if kv_chunk_size is None else "fixed" + attention_mode = "prefill" if kv_chunk_size is None else "chunked_prefill" + q_bounds = _split_bounds(seq_len, cp_world_size) + kv_bounds = _kv_block_bounds(seq_len, cp_world_size, kv_chunk_size) + runtime_plan_set = build_reference_split_kv_runtime_plan_set( + (seq_len,) * args.batch, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ) + case: dict[str, object] = { + "case_name": _case_name(tp_world_size, cp_world_size, kv_chunk_size, args.dtype), + "attention_mode": attention_mode, + "model": args.model, + "topology": { + "tp_world_size": tp_world_size, + "cp_world_size": cp_world_size, + "tp_rank": 0, + "logical_cp_ranks": cp_world_size, + "local_num_query_heads": local_hq, + "local_num_kv_heads": local_hkv, + "local_query_head_range": [0, local_hq], + "local_kv_head_range": [0, local_hkv], + "head_dim": QWEN3_8B_HEAD_DIM, + "q_sequence_bounds": [list(item) for item in q_bounds], + "kv_block_bounds": [list(item) for item in kv_bounds], + }, + "provenance": { + "backend": "deterministic_cp_reference", + "reference_backend": "cp1_fp32_prefill", + "candidate_backend": "cp_reference", + "dtype": args.dtype, + "accum_dtype": "fp32", + "downcast_at": "final_write", + "lse_domain": "attention", + "merge_order": "global_block_index", + "split_kv_policy": split_kv_policy, + "requested_split_kv_policy": split_kv_policy, + "requested_split_kv_size": kv_chunk_size, + "actual_split_kv_plans": split_kv_execution_plan_provenance( + seq_len, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ), + "actual_split_kv_plan_set": runtime_plan_set.to_dict(), + "kv_chunk_size": kv_chunk_size, + "block_metadata_hash": _block_metadata_hash(kv_bounds), + "scale_placement": "scores_after_qk_matmul", + "mask_application_order": ["scale", "causal_mask", "key_padding_mask"], + "dropout_policy": "disabled", + "deterministic_controls": { + "reference": "strict_fp32_math_inside_cp_attention", + "num_threads": args.num_threads, + }, + "rope": rope_report["provenance"], + }, + "drift": { + "cp_merge_fp32": { + "out": _drift_stats(candidate_fp32_out, reference_out), + "lse": _drift_stats(candidate_fp32_lse, reference_lse), + "source_class": "reduction_and_collective_drift", + }, + "dtype_path_vs_fp32": { + "out": _drift_stats(candidate_dtype_out, reference_out), + "lse": _drift_stats(candidate_dtype_lse, reference_lse), + "source_class": "arithmetic_schedule_drift", + }, + "rope": rope_report["drift"], + }, + "merge_order_probe": _merge_order_probe( + attention, + q, + k, + v, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ), + "te_merge_oracle": _te_merge_oracle_probe( + attention, + q, + k, + v, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + te_adapter=te_adapter, + ), + "per_rank": _per_rank_forward_drifts( + candidate_fp32_out, + candidate_fp32_lse, + reference_out, + reference_lse, + cp_world_size, + ), + "backward": {"status": "not_requested"}, + } + distributed_reference = _run_distributed_p2p_reference( + q, + k, + v, + reference_out, + reference_lse, + device=device, + seq_len=seq_len, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) + case["distributed_p2p_reference"] = distributed_reference + if args.include_backward: + backward = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=cp_world_size, + candidate_kv_chunk_size=kv_chunk_size, + output_dtype=dtype, + ) + case["backward"] = { + "status": "available", + "report": backward.to_dict(), + } + return case + + +def _run_distributed_p2p_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + reference_out: torch.Tensor, + reference_lse: torch.Tensor, + *, + device: torch.device, + seq_len: int, + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: int | None, +) -> dict[str, object]: + """Exercise the actual P2P reference when launched as a matching NCCL job.""" + + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + return {"status": "not_requested", "reason": "process_group_not_initialized"} + backend = str(dist.get_backend()).lower() + world_size = int(dist.get_world_size()) + if device.type != "cuda": + return { + "status": "skipped", + "reason": "P2P NCCL reference requires CUDA", + "backend": backend, + } + if "nccl" not in backend: + return { + "status": "skipped", + "reason": "P2P reference requires NCCL", + "backend": backend, + } + if world_size != cp_world_size: + return { + "status": "skipped", + "reason": "WORLD_SIZE must equal cp_world_size for the CP reference", + "world_size": world_size, + "cp_world_size": cp_world_size, + } + + rank = int(dist.get_rank()) + owner_ranges = _split_bounds(seq_len, cp_world_size) + block_bounds = _kv_block_bounds(seq_len, cp_world_size, kv_chunk_size) + blocks: list[AttentionCPBlockMetadata] = [] + owner_block_counts = [0] * cp_world_size + for block_index, (start, end) in enumerate(block_bounds): + owner = next( + owner_rank + for owner_rank, (owner_start, owner_end) in enumerate(owner_ranges) + if owner_start <= start < owner_end + ) + blocks.append( + AttentionCPBlockMetadata( + global_block_index=block_index, + kv_block_start=start, + kv_block_end=end, + owner_cp_rank=owner, + owner_tp_rank=0, + ) + ) + owner_block_counts[owner] += 1 + + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec( + tp_world_size=tp_world_size, + tp_rank=0, + cp_world_size=cp_world_size, + cp_rank=rank, + ), + backend="p2p_nccl_reference", + status="implemented", + expected_blocks=tuple(blocks), + expected_kv_token_range=(0, seq_len), + query_token_ranges=tuple(_split_bounds(q.size(2), cp_world_size)), + ) + attention = DeterministicCPAttentionReferenceOp() + local_states: list[AttentionCPPartialState] = [] + for block in blocks: + if block.owner_cp_rank != rank: + continue + state = attention.local_partial_state( + q, + k[:, :, block.kv_block_start : block.kv_block_end, :], + v[:, :, block.kv_block_start : block.kv_block_end, :], + q_start=0, + k_start=block.kv_block_start, + total_kv_len=seq_len, + total_query_len=q.size(2), + causal=True, + ) + local_states.append( + AttentionCPPartialState( + out=state.out, + lse=state.lse, + block=block, + ) + ) + communication = P2PNCCLAttentionCPCommunication() + gathered = communication.all_gather_partial_states(tuple(local_states), plan) + merged = merge_attention_partial_states( + [ + AttentionPartialState( + out=state.out, + lse=state.lse, + block_start=state.block.kv_block_start, + block_end=state.block.kv_block_end, + ) + for state in gathered + ] + ) + local = communication.reduce_scatter_merged_state( + AttentionCPMergedState(out=merged.out, lse=merged.lse), + plan, + ) + q_start, q_end = plan.query_token_ranges[rank] + reference_local_out = reference_out[:, :, q_start:q_end, :] + reference_local_lse = reference_lse[:, :, q_start:q_end] + return { + "status": "available", + "backend": backend, + "rank": rank, + "world_size": world_size, + "transport": "p2p_nccl_reference", + "manifest_block_count": len(blocks), + "owner_block_counts": owner_block_counts, + "gathered_block_indices": [state.block.global_block_index for state in gathered], + "query_range": [q_start, q_end], + "out": _drift_stats(local.out, reference_local_out), + "lse": _drift_stats(local.lse, reference_local_lse), + } + + +def _make_qkv( + *, + batch: int, + local_hq: int, + local_hkv: int, + seq_len: int, + dtype: torch.dtype, + device: torch.device, + seed: int, + compose_rope: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict[str, object]]: + gen = torch.Generator(device="cpu").manual_seed(seed) + q_pre = torch.randn(batch, local_hq, seq_len, QWEN3_8B_HEAD_DIM, generator=gen) + k_pre = torch.randn(batch, local_hkv, seq_len, QWEN3_8B_HEAD_DIM, generator=gen) + v = torch.randn(batch, local_hkv, seq_len, QWEN3_8B_HEAD_DIM, generator=gen) + q_pre = q_pre.to(device=device, dtype=dtype) + k_pre = k_pre.to(device=device, dtype=dtype) + v = v.to(device=device, dtype=dtype) + + positions = ( + torch.arange(seq_len, dtype=torch.long, device=device) + .unsqueeze(0) + .expand( + batch, + -1, + ) + ) + if not compose_rope: + return q_pre, k_pre, v, _rope_report_disabled() + + rope = NativeRoPEOp() + q_rope_dtype = rope.forward(q_pre, positions, theta=QWEN3_8B_ROPE_THETA) + k_rope_dtype = rope.forward(k_pre, positions, theta=QWEN3_8B_ROPE_THETA) + q_rope_fp32 = rope.forward_fp32(q_pre, positions, theta=QWEN3_8B_ROPE_THETA) + k_rope_fp32 = rope.forward_fp32(k_pre, positions, theta=QWEN3_8B_ROPE_THETA) + return ( + q_rope_dtype, + k_rope_dtype, + v, + { + "provenance": { + "rope_state": "post_rope", + "rope_theta": QWEN3_8B_ROPE_THETA, + "rope_scaling": None, + "rotary_dim": QWEN3_8B_HEAD_DIM, + "position_ids": "arange(seq_len)", + "cache_position": "same_as_position_ids", + "query_position_offsets": [0 for _ in range(batch)], + "key_position_offsets": [0 for _ in range(batch)], + "k_cache_rope_state": "post_rope", + "rope_cast_at": "rope_output", + "rope_output_dtype": _dtype_name(dtype), + "fusion_boundary": "unfused_rope_attention_reference", + }, + "drift": { + "status": "available", + "q": _drift_stats(q_rope_dtype, q_rope_fp32), + "k": _drift_stats(k_rope_dtype, k_rope_fp32), + }, + }, + ) + + +def _make_dout( + *, + batch: int, + local_hq: int, + seq_len: int, + dtype: torch.dtype, + device: torch.device, + seed: int, +) -> torch.Tensor: + gen = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn( + batch, + local_hq, + seq_len, + QWEN3_8B_HEAD_DIM, + generator=gen, + dtype=dtype, + ).to(device=device) + + +def _rope_report_disabled() -> dict[str, object]: + return { + "provenance": { + "rope_state": "not_composed", + "fusion_boundary": "attention_only", + }, + "drift": { + "status": "not_composed", + "q": None, + "k": None, + }, + } + + +def _merge_order_probe( + attention: DeterministicCPAttentionReferenceOp, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cp_world_size: int, + kv_chunk_size: int | None, +) -> dict[str, object]: + reversed_out: list[torch.Tensor] = [] + reversed_lse: list[torch.Tensor] = [] + sorted_out: list[torch.Tensor] = [] + sorted_lse: list[torch.Tensor] = [] + kv_bounds = _kv_block_bounds(k.size(2), cp_world_size, kv_chunk_size) + for q_start, q_end in _split_bounds(q.size(2), cp_world_size): + if q_start == q_end: + continue + states = _partial_states_for_query_block( + attention, + q, + k, + v, + q_start=q_start, + q_end=q_end, + kv_bounds=kv_bounds, + ) + sorted_merge = merge_attention_partial_states(states) + reversed_merge = merge_attention_partial_states(list(reversed(states))) + sorted_out.append(sorted_merge.out) + sorted_lse.append(sorted_merge.lse) + reversed_out.append(reversed_merge.out) + reversed_lse.append(reversed_merge.lse) + if not sorted_out: + return { + "status": "empty_query", + "arrival_order_policy": "ignored_then_sorted_by_global_block_index", + } + return { + "status": "available", + "arrival_order_policy": "ignored_then_sorted_by_global_block_index", + "out": _drift_stats(torch.cat(reversed_out, dim=2), torch.cat(sorted_out, dim=2)), + "lse": _drift_stats(torch.cat(reversed_lse, dim=2), torch.cat(sorted_lse, dim=2)), + } + + +def _te_merge_oracle_probe( + attention: DeterministicCPAttentionReferenceOp, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cp_world_size: int, + kv_chunk_size: int | None, + te_adapter: TEContextParallelMergeAdapter | None, +) -> dict[str, object]: + if te_adapter is None: + return { + "status": "unavailable", + "te_module": TE_CONTEXT_PARALLEL_MODULE, + "fallback": "deterministic_cp_reference", + } + + ours_out: list[torch.Tensor] = [] + ours_lse: list[torch.Tensor] = [] + te_out: list[torch.Tensor] = [] + te_lse: list[torch.Tensor] = [] + kv_bounds = _kv_block_bounds(k.size(2), cp_world_size, kv_chunk_size) + for q_start, q_end in _split_bounds(q.size(2), cp_world_size): + if q_start == q_end: + continue + states = _partial_states_for_query_block( + attention, + q, + k, + v, + q_start=q_start, + q_end=q_end, + kv_bounds=kv_bounds, + ) + ours = merge_attention_partial_states(states) + te = te_adapter.merge(states) + ours_out.append(ours.out) + ours_lse.append(ours.lse) + te_out.append(te.out) + te_lse.append(te.lse) + if not ours_out: + return {"status": "empty_query", "te_version": te_adapter.version} + return { + "status": "available", + "te_version": te_adapter.version, + "te_module": TE_CONTEXT_PARALLEL_MODULE, + "te_symbols": list(TE_SYMBOLS), + "out": _drift_stats(torch.cat(te_out, dim=2), torch.cat(ours_out, dim=2)), + "lse": _drift_stats(torch.cat(te_lse, dim=2), torch.cat(ours_lse, dim=2)), + } + + +def _partial_states_for_query_block( + attention: DeterministicCPAttentionReferenceOp, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + q_start: int, + q_end: int, + kv_bounds: Sequence[tuple[int, int]], +) -> list[AttentionPartialState]: + return [ + attention.local_partial_state( + q[:, :, q_start:q_end, :], + k[:, :, k_start:k_end, :], + v[:, :, k_start:k_end, :], + q_start=q_start, + k_start=k_start, + total_kv_len=k.size(2), + total_query_len=q.size(2), + causal=True, + ) + for k_start, k_end in kv_bounds + if k_start != k_end + ] + + +def _per_rank_forward_drifts( + candidate_out: torch.Tensor, + candidate_lse: torch.Tensor, + reference_out: torch.Tensor, + reference_lse: torch.Tensor, + cp_world_size: int, +) -> list[dict[str, object]]: + per_rank = [] + for rank, (q_start, q_end) in enumerate(_split_bounds(candidate_out.size(2), cp_world_size)): + per_rank.append( + { + "rank": rank, + "query_start": q_start, + "query_end": q_end, + "out": _drift_stats( + candidate_out[:, :, q_start:q_end, :], + reference_out[:, :, q_start:q_end, :], + ), + "lse": _drift_stats( + candidate_lse[:, :, q_start:q_end], + reference_lse[:, :, q_start:q_end], + ), + } + ) + return per_rank + + +def _drift_stats(candidate: torch.Tensor, reference: torch.Tensor) -> dict[str, object]: + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match " + f"reference shape {tuple(reference.shape)}" + ) + diff = (candidate.float() - reference.float()).abs().reshape(-1) + active_count = int(diff.numel()) + if active_count == 0: + return { + "max_abs": 0.0, + "mean_abs": 0.0, + "p95_abs": 0.0, + "p99_abs": 0.0, + "active_count": 0, + } + return { + "max_abs": float(diff.max().item()), + "mean_abs": float(diff.mean().item()), + "p95_abs": float(torch.quantile(diff, 0.95).item()), + "p99_abs": float(torch.quantile(diff, 0.99).item()), + "active_count": active_count, + } + + +def _rank_env() -> dict[str, int | bool]: + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + return { + "rank": rank, + "world_size": world_size, + "local_rank": local_rank, + "torchrun": "RANK" in os.environ or "WORLD_SIZE" in os.environ, + } + + +def _resolve_device(device_arg: str, rank_env: dict[str, int | bool]) -> torch.device: + if device_arg == "auto": + device_arg = "cuda" if torch.cuda.is_available() else "cpu" + if device_arg == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA was requested but is not available") + local_rank = int(rank_env["local_rank"]) + if torch.cuda.device_count() > 0: + torch.cuda.set_device(local_rank % torch.cuda.device_count()) + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") + + +def _maybe_init_process_group( + args: argparse.Namespace, + device: torch.device, + rank_env: dict[str, int | bool], +) -> dict[str, object]: + initialized = False + backend = None + if args.init_process_group and int(rank_env["world_size"]) > 1: + import torch.distributed as dist + + backend = "nccl" if device.type == "cuda" else "gloo" + dist.init_process_group(backend=backend, init_method="env://") + initialized = True + return { + "initialized": initialized, + "backend": backend, + "transport": "torchrun_env_rank_aware", + } + + +def _runtime_metadata( + device: torch.device, + distributed: dict[str, object], + rank_env: dict[str, int | bool], +) -> dict[str, object]: + return { + "python": sys.version.split()[0], + "platform": platform.platform(), + "torch_version": torch.__version__, + "cuda_available": torch.cuda.is_available(), + "device": str(device), + "distributed": distributed, + "rank_env": rank_env, + } + + +def _launch_metadata(rank_env: dict[str, int | bool]) -> dict[str, object]: + return { + "command": _shell_join(sys.argv), + "rank": int(rank_env["rank"]), + "world_size": int(rank_env["world_size"]), + "local_rank": int(rank_env["local_rank"]), + "torchrun": bool(rank_env["torchrun"]), + } + + +def _validate_args(args: argparse.Namespace) -> None: + if args.batch < 1: + raise ValueError("batch must be >= 1") + if args.seq_len < 1: + raise ValueError("seq_len must be >= 1") + if args.num_threads < 1: + raise ValueError("num_threads must be >= 1") + for tp_world_size in _parse_int_csv(args.tp_world_sizes, name="tp_world_sizes"): + _validate_topology(tp_world_size, 1) + for cp_world_size in _parse_int_csv(args.cp_world_sizes, name="cp_world_sizes"): + _validate_topology(1, cp_world_size) + _parse_kv_chunk_sizes(args.kv_chunk_sizes) + + +def _validate_topology(tp_world_size: int, cp_world_size: int) -> None: + if tp_world_size < 1 or cp_world_size < 1: + raise ValueError("tp_world_size and cp_world_size must be >= 1") + if QWEN3_8B_HEADS % tp_world_size != 0: + raise ValueError("Qwen3 query heads must be divisible by tp_world_size") + if QWEN3_8B_KV_HEADS % tp_world_size != 0: + raise ValueError("Qwen3 KV heads must be divisible by tp_world_size") + + +def _parse_int_csv(value: str, *, name: str) -> tuple[int, ...]: + parsed: list[int] = [] + for raw in value.split(","): + item = raw.strip() + if not item: + continue + try: + parsed.append(int(item)) + except ValueError as exc: + raise ValueError(f"{name} must be a comma-separated integer list") from exc + if not parsed: + raise ValueError(f"{name} must contain at least one integer") + return tuple(parsed) + + +def _parse_kv_chunk_sizes(value: str) -> tuple[int | None, ...]: + parsed: list[int | None] = [] + for raw in value.split(","): + item = raw.strip().lower() + if not item: + continue + if item in {"none", "full", "no_split"}: + parsed.append(None) + continue + try: + size = int(item) + except ValueError as exc: + raise ValueError("kv_chunk_sizes must contain integers or 'none'") from exc + if size < 1: + raise ValueError("kv chunk sizes must be >= 1") + parsed.append(size) + if not parsed: + raise ValueError("kv_chunk_sizes must contain at least one entry") + return tuple(parsed) + + +def _dtype_from_name(name: str) -> torch.dtype: + if name == "bf16": + return torch.bfloat16 + if name == "fp32": + return torch.float32 + raise ValueError(f"unsupported dtype: {name}") + + +def _dtype_name(dtype: torch.dtype) -> str: + return str(dtype).replace("torch.", "") + + +def _split_bounds(length: int, parts: int) -> list[tuple[int, int]]: + base, extra = divmod(length, parts) + bounds: list[tuple[int, int]] = [] + start = 0 + for index in range(parts): + width = base + (1 if index < extra else 0) + end = start + width + bounds.append((start, end)) + start = end + return bounds + + +def _kv_block_bounds( + length: int, + cp_world_size: int, + kv_chunk_size: int | None, +) -> list[tuple[int, int]]: + bounds: list[tuple[int, int]] = [] + for start, end in _split_bounds(length, cp_world_size): + if kv_chunk_size is None: + bounds.append((start, end)) + continue + cursor = start + while cursor < end: + chunk_end = min(cursor + kv_chunk_size, end) + bounds.append((cursor, chunk_end)) + cursor = chunk_end + return bounds + + +def _block_metadata_hash(bounds: Sequence[tuple[int, int]]) -> str: + payload = json.dumps([list(item) for item in bounds], separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest()[:16] + + +def _case_seed(seed: int, tp_world_size: int, cp_world_size: int, kv_chunk_size: int | None) -> int: + return seed + tp_world_size * 101 + cp_world_size * 17 + (kv_chunk_size or 0) + + +def _case_name( + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: int | None, + dtype: str, +) -> str: + mode = "prefill" if kv_chunk_size is None else f"chunk{kv_chunk_size}" + return f"qwen3_8b_tp{tp_world_size}_cp{cp_world_size}_{mode}_{dtype}" + + +def _transformer_engine_version() -> str: + for package in ("transformer-engine", "transformer_engine"): + try: + return importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + continue + return "unknown" + + +def _shell_join(argv: Sequence[str]) -> str: + if os.name == "nt": + return " ".join(argv) + return shlex.join(argv) + + +@contextlib.contextmanager +def _thread_limit(num_threads: int) -> Iterator[None]: + previous = torch.get_num_threads() + torch.set_num_threads(num_threads) + try: + yield + finally: + torch.set_num_threads(previous) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md b/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md new file mode 100644 index 00000000..88fbd199 --- /dev/null +++ b/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md @@ -0,0 +1,129 @@ +# WS2 Attention PR5 Drift Benchmark + +PR5 adds the report artifact path for issue #235. It does not introduce a +production communication kernel. The benchmark is a rank-aware, torchrun-style +driver around the deterministic CP attention reference. Under a matching +two-rank CUDA/NCCL launch it executes the P2P reference transport; CPU/Gloo +remains a report-generation smoke path. + +## Scope + +The benchmark covers the Qwen3-8B Attention target: + +- global heads: `Hq=32`, `Hkv=8`, `D=128` +- TP sweep: `TP=1/2`; TP only changes the local head shard shape +- CP sweep: `CP=1/2` +- modes: full prefill and chunked-prefill replay +- dtype path: BF16 candidate path compared with FP32 reference +- optional backward: `dq`, `dk`, `dv` drift from the PR8 reference +- optional RoPE composition before Attention, while CP Attention still consumes + post-RoPE Q/K + +The report separates two drift classes: + +| Field | Meaning | +| --- | --- | +| `drift.cp_merge_fp32` | CP/chunked candidate with FP32 output vs CP=1 FP32 prefill. This isolates CP merge and split-KV order. | +| `drift.dtype_path_vs_fp32` | BF16 candidate path vs FP32 reference. This exposes arithmetic/final-write drift. | +| `merge_order_probe` | Reversed-arrival partial states vs canonical sorted merge. This verifies that arrival order is ignored. | +| `te_merge_oracle` | Optional Transformer Engine merge-oracle drift when TE is installed and passes capability probes. | +| `backward` | Optional PR8 `dq/dk/dv` drift report when `--include-backward` is used. | +| `distributed_p2p_reference` | Real NCCL P2P partial-state gather, FP32 merge, and query scatter drift. | + +Selected-logprob `dlogp` remains `not_available` here because the full logprob +chain integration is outside PR5. PR4/WS2 runtime integration should fill that +field once Attention is wired into the chain. + +## Commands + +Local smoke: + +```bash +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json +``` + +Qwen3 TP=2 / CP=2 with backward drift and a JSON artifact: + +```bash +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 +``` + +Two-GPU NCCL transport check: + +```bash +torchrun --standalone --nproc-per-node=2 \ + scripts/ws2_p2p_nccl_attention_reference_check.py +``` + +Two-GPU benchmark report with real P2P transport: + +```bash +torchrun --standalone --nproc-per-node=2 \ + benchmarks/benchmark_ws2_cp_attention_drift.py \ + --smoke \ + --device cuda \ + --init-process-group \ + --tp-world-sizes 2 \ + --cp-world-sizes 2 \ + --json +``` + +Rank 0 prints or writes the shared report. Other ranks can run the same +rank-aware benchmark without changing the numerical reducer. The recommended +container is the repository CUDA image built from `docker/Dockerfile.cuda` +(`ghcr.io/rl-align/rl-kernel/rl-kernel-ci:cuda` when using the repository image +workflow). It is based on PyTorch 2.4 / CUDA 12.4 and includes NCCL support. + +## Transformer Engine Reuse + +PR5 reuses Transformer Engine only as an optional merge oracle, not as the +source of truth. The adapter imports: + +```text +transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +``` + +and uses these APIs when available: + +```text +flash_attn_fwd_softmax_lse_correction +flash_attn_fwd_out_correction_init +flash_attn_fwd_out_correction +``` + +The benchmark first builds RL-Kernel partial states: + +```text +state_i = (out_i, lse_i, global_block_index_i) +``` + +then sorts them by `global_block_index`. TE is allowed to perform only the +online-softmax correction arithmetic for those already-sorted states. If TE is +missing, incompatible, or fails the numeric self-test, the report records a +provenance fallback and continues with the deterministic RL-Kernel merge. + +## Report Contract + +The JSON root contains: + +```text +schema_version +issue / pr +launch.command +runtime.rank_env +target +te_context_parallel_merge +dlogp +cases[] +``` + +Each case records topology, RoPE/cache provenance, split-KV policy, block +metadata hash, drift summaries, per-logical-CP-rank metrics, and optional +backward drift. The merge order is always `global_block_index`, and +`downcast_at` is always `final_write`. diff --git a/docs/operators/attention.md b/docs/operators/attention.md index ebff9a58..bc7f4a38 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -85,6 +85,32 @@ Calling it (`__call__` -> `forward(...)`) computes in the input dtype; `forward_ 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 Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): @@ -137,6 +163,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 +221,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 new file mode 100644 index 00000000..eb4994b3 --- /dev/null +++ b/rl_engine/kernels/attention_contract.py @@ -0,0 +1,1472 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for context-parallel standard softmax attention. + +The objects in this module describe a distributed attention invocation. They +do not shard tensors, launch collectives, or implement the ``(out, lse)`` +merge. Keeping description and materialization separate lets dispatch reject +an incompatible backend before any numerically different path is launched. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Iterable, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + + +class AttentionContractError(ValueError): + """Raised when attention metadata does not describe a valid invocation.""" + + +class AttentionRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class AttentionMode(str, Enum): + PREFILL = "prefill" + CHUNKED_PREFILL = "chunked_prefill" + DECODE = "decode" + + +class AttentionDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class AttentionMerge(str, Enum): + ONLINE_SOFTMAX_LSE = "online_softmax_lse" + + +class ReductionOrder(str, Enum): + GLOBAL_BLOCK_INDEX = "global_block_index" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +class SplitKVMode(str, Enum): + DISABLED = "disabled" + FIXED = "fixed" + AUTO = "auto" + + +class RoPEState(str, Enum): + PRE_ROPE = "pre_rope" + POST_ROPE = "post_rope" + + +class RoPECastPoint(str, Enum): + NONE = "none" + BEFORE_ROPE = "before_rope" + AFTER_ROPE = "after_rope" + + +class RoPEFusionBoundary(str, Enum): + UNFUSED_ROPE_ATTENTION = "unfused_rope_attention" + FUSED_ROPE_ATTENTION = "fused_rope_attention" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise AttentionContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise AttentionContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise AttentionContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _integer_tuple(values: Iterable[int], field: str) -> tuple[int, ...]: + try: + result = tuple(values) + except TypeError as exc: + raise AttentionContractError(f"{field} must be an iterable of integers") from exc + for index, value in enumerate(result): + if isinstance(value, bool) or not isinstance(value, int): + raise AttentionContractError(f"{field}[{index}] must be an integer; got {value!r}") + return result + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical TP/CP ownership for one attention invocation. + + TP head shards are currently required to be equal and contiguous. CP + sequence ownership may be uneven, but every local block must carry a stable + logical global index so a later implementation can merge by logical order + instead of collective arrival order. + """ + + tp_rank: int + tp_world_size: int + cp_rank: int + cp_world_size: int + global_q_heads: int + global_kv_heads: int + local_q_head_start: int + local_q_heads: int + local_kv_head_start: int + local_kv_heads: int + global_sequence_length: int + local_sequence_length: int + global_block_indices: tuple[int, ...] + global_block_token_starts: tuple[int, ...] + local_block_offsets: tuple[int, ...] + packed_sequence_offsets: tuple[int, ...] | None = None + + 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") + 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: + raise AttentionContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + if cp_rank >= cp_world_size: + raise AttentionContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_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") + if global_q_heads % global_kv_heads != 0: + raise AttentionContractError( + f"global_q_heads={global_q_heads} must be divisible by " + f"global_kv_heads={global_kv_heads} for GQA" + ) + if global_q_heads % tp_world_size != 0 or global_kv_heads % tp_world_size != 0: + raise AttentionContractError( + "global Q/KV heads must be evenly divisible by tp_world_size; " + f"got Hq={global_q_heads}, Hkv={global_kv_heads}, TP={tp_world_size}" + ) + + local_q_heads = _positive_int(self.local_q_heads, "local_q_heads") + local_kv_heads = _positive_int(self.local_kv_heads, "local_kv_heads") + expected_q_heads = global_q_heads // tp_world_size + expected_kv_heads = global_kv_heads // tp_world_size + if local_q_heads != expected_q_heads or local_kv_heads != expected_kv_heads: + raise AttentionContractError( + "local TP head counts do not preserve the global Q/KV mapping; " + f"expected ({expected_q_heads}, {expected_kv_heads}), got " + f"({local_q_heads}, {local_kv_heads})" + ) + + local_q_head_start = _non_negative_int(self.local_q_head_start, "local_q_head_start") + local_kv_head_start = _non_negative_int(self.local_kv_head_start, "local_kv_head_start") + expected_q_start = tp_rank * expected_q_heads + expected_kv_start = tp_rank * expected_kv_heads + if local_q_head_start != expected_q_start or local_kv_head_start != expected_kv_start: + raise AttentionContractError( + "local TP head starts do not match contiguous rank ownership; " + f"expected ({expected_q_start}, {expected_kv_start}), got " + f"({local_q_head_start}, {local_kv_head_start})" + ) + + global_sequence_length = _positive_int( + self.global_sequence_length, "global_sequence_length" + ) + local_sequence_length = _positive_int(self.local_sequence_length, "local_sequence_length") + + block_indices = _integer_tuple(self.global_block_indices, "global_block_indices") + if not block_indices: + raise AttentionContractError("global_block_indices must not be empty") + if any(index < 0 for index in block_indices): + raise AttentionContractError("global_block_indices must be non-negative") + if any( + left >= right for left, right in zip(block_indices, block_indices[1:], strict=False) + ): + raise AttentionContractError( + "global_block_indices must be unique and strictly increasing" + ) + object.__setattr__(self, "global_block_indices", block_indices) + + block_token_starts = _integer_tuple( + self.global_block_token_starts, "global_block_token_starts" + ) + local_block_offsets = _integer_tuple(self.local_block_offsets, "local_block_offsets") + if len(block_token_starts) != len(block_indices): + raise AttentionContractError( + "global_block_token_starts must contain one entry per global_block_indices entry" + ) + if any(start < 0 for start in block_token_starts): + raise AttentionContractError("global_block_token_starts must be non-negative") + if len(local_block_offsets) != len(block_indices) + 1: + raise AttentionContractError( + "local_block_offsets must contain one boundary more than global_block_indices" + ) + if local_block_offsets[0] != 0 or local_block_offsets[-1] != local_sequence_length: + raise AttentionContractError( + "local_block_offsets must start at 0 and end at local_sequence_length" + ) + if any( + left >= right + for left, right in zip(local_block_offsets, local_block_offsets[1:], strict=False) + ): + raise AttentionContractError("local_block_offsets must be strictly increasing") + + previous_global_end = 0 + for index, (global_start, local_start, local_end) in enumerate( + zip( + block_token_starts, + local_block_offsets[:-1], + local_block_offsets[1:], + strict=True, + ) + ): + global_end = global_start + (local_end - local_start) + if global_end > global_sequence_length: + raise AttentionContractError( + f"global block {block_indices[index]} exceeds global_sequence_length" + ) + if index > 0 and global_start < previous_global_end: + raise AttentionContractError( + "global block token ranges must be non-overlapping and ordered" + ) + previous_global_end = global_end + object.__setattr__(self, "global_block_token_starts", block_token_starts) + object.__setattr__(self, "local_block_offsets", local_block_offsets) + + if self.packed_sequence_offsets is not None: + offsets = _integer_tuple(self.packed_sequence_offsets, "packed_sequence_offsets") + if len(offsets) < 2 or offsets[0] != 0: + raise AttentionContractError( + "packed_sequence_offsets must start at 0 and contain an end offset" + ) + if any(left >= right for left, right in zip(offsets, offsets[1:], strict=False)): + raise AttentionContractError("packed_sequence_offsets must be strictly increasing") + if offsets[-1] != local_sequence_length: + raise AttentionContractError( + "the final packed_sequence_offsets value must equal local_sequence_length; " + f"got {offsets[-1]} and {local_sequence_length}" + ) + object.__setattr__(self, "packed_sequence_offsets", offsets) + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic CP ``(out, lse)`` merge semantics.""" + + merge: AttentionMerge = AttentionMerge.ONLINE_SOFTMAX_LSE + acc_dtype: AttentionDType = AttentionDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + + def __post_init__(self) -> None: + object.__setattr__(self, "merge", _enum_value(AttentionMerge, self.merge, "merge")) + object.__setattr__( + self, "acc_dtype", _enum_value(AttentionDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError( + f"CP attention accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class SplitKVExecutionPlan: + """Actual backend-local Split-KV schedule emitted by a runtime. + + CP ownership is intentionally separate from this plan. ``boundaries`` + describe the canonical logical KV ranges reduced by one backend invocation; + CP may transport those partial states between ranks but must not change the + FP32 merge contract recorded here. + """ + + requested_mode: SplitKVMode + requested_split_size: int | None + actual_mode: SplitKVMode | None + actual_split_size: int | None + boundaries: tuple[tuple[int, int], ...] + merge_order: ReductionOrder = ReductionOrder.GLOBAL_BLOCK_INDEX + acc_dtype: AttentionDType = AttentionDType.FP32 + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + backend: str = "reference" + source: str = "contract_exact" + fallback: bool = False + fallback_reason: str | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "requested_mode", + _enum_value(SplitKVMode, self.requested_mode, "requested_mode"), + ) + if self.actual_mode is not None: + object.__setattr__( + self, + "actual_mode", + _enum_value(SplitKVMode, self.actual_mode, "actual_mode"), + ) + for field_name in ("requested_split_size", "actual_split_size"): + value = getattr(self, field_name) + if value is not None: + _positive_int(value, field_name) + if self.requested_mode is SplitKVMode.FIXED: + if self.requested_split_size is None: + raise AttentionContractError( + "requested fixed Split-KV mode requires requested_split_size" + ) + elif self.requested_split_size is not None: + raise AttentionContractError( + "requested_split_size is only valid for requested fixed Split-KV mode" + ) + try: + boundaries = tuple(tuple(boundary) for boundary in self.boundaries) + except TypeError as exc: + raise AttentionContractError( + "Split-KV boundaries must be an iterable of (start, end) pairs" + ) from exc + previous_end = 0 + for index, boundary in enumerate(boundaries): + if len(boundary) != 2: + raise AttentionContractError( + f"Split-KV boundary {index} must contain exactly start and end" + ) + start, end = boundary + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + ): + raise AttentionContractError( + "Split-KV boundaries must satisfy 0 <= start < end" + ) + if index > 0 and start != previous_end: + raise AttentionContractError( + "Split-KV boundaries must be contiguous and in logical KV order" + ) + previous_end = end + if self.actual_mode is None and boundaries: + raise AttentionContractError("unknown actual Split-KV plan cannot declare boundaries") + if self.actual_mode is not None and not boundaries: + raise AttentionContractError("actual Split-KV plan must declare logical boundaries") + if self.actual_mode is SplitKVMode.FIXED: + if self.actual_split_size is None: + raise AttentionContractError( + "actual fixed Split-KV mode requires actual_split_size" + ) + widths = tuple(end - start for start, end in boundaries) + if any(width != self.actual_split_size for width in widths[:-1]) or ( + widths and widths[-1] > self.actual_split_size + ): + raise AttentionContractError( + "fixed Split-KV boundaries must use actual_split_size except " + "for a shorter final split" + ) + elif self.actual_split_size is not None: + raise AttentionContractError( + "actual_split_size is only valid for actual fixed Split-KV mode" + ) + if self.actual_mode is SplitKVMode.DISABLED and len(boundaries) != 1: + raise AttentionContractError( + "disabled Split-KV execution must contain exactly one boundary" + ) + object.__setattr__(self, "boundaries", boundaries) + object.__setattr__( + self, + "merge_order", + _enum_value(ReductionOrder, self.merge_order, "merge_order"), + ) + object.__setattr__( + self, + "acc_dtype", + _enum_value(AttentionDType, self.acc_dtype, "acc_dtype"), + ) + object.__setattr__( + self, + "downcast_at", + _enum_value(DowncastPoint, self.downcast_at, "downcast_at"), + ) + if self.acc_dtype is not AttentionDType.FP32: + raise AttentionContractError("Split-KV partial states must be merged in fp32") + if not isinstance(self.backend, str) or not self.backend.strip(): + raise AttentionContractError("Split-KV backend must be a non-empty string") + if not isinstance(self.source, str) or not self.source.strip(): + raise AttentionContractError("Split-KV plan source must be a non-empty string") + if not isinstance(self.fallback, bool): + raise AttentionContractError("Split-KV fallback must be a bool") + if self.fallback and not self.fallback_reason: + raise AttentionContractError("Split-KV fallback_reason is required for a fallback") + if not self.fallback and self.fallback_reason is not None: + raise AttentionContractError( + "Split-KV fallback_reason must be None when fallback=False" + ) + if not self.fallback and self.actual_mode is not None: + if self.actual_mode is not self.requested_mode: + raise AttentionContractError( + "actual Split-KV mode may differ from requested mode only for a fallback" + ) + if self.actual_split_size != self.requested_split_size: + raise AttentionContractError( + "actual Split-KV size may differ from requested size only for a fallback" + ) + + @property + def actual_split_count(self) -> int | None: + return None if self.actual_mode is None else len(self.boundaries) + + def to_dict(self) -> dict[str, Any]: + return { + "requested_split_kv_policy": self.requested_mode.value, + "requested_split_kv_size": self.requested_split_size, + "actual_split_kv_policy": ( + None if self.actual_mode is None else self.actual_mode.value + ), + "actual_split_kv_size": self.actual_split_size, + "actual_split_kv_count": self.actual_split_count, + "actual_split_boundaries": [list(boundary) for boundary in self.boundaries], + "split_kv_merge_order": self.merge_order.value, + "split_kv_accum_dtype": self.acc_dtype.value, + "split_kv_downcast_at": self.downcast_at.value, + "split_kv_backend": self.backend, + "split_kv_plan_source": self.source, + "split_kv_fallback": self.fallback, + "split_kv_fallback_reason": self.fallback_reason, + } + + +@dataclass(frozen=True) +class SplitKVSpec: + """Requested Split-KV policy shared by training and rollout paths.""" + + mode: SplitKVMode = SplitKVMode.DISABLED + fixed_split_size: int | None = None + strict_consistency: bool = True + + @classmethod + def disabled(cls, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls(mode=SplitKVMode.DISABLED, strict_consistency=strict_consistency) + + @classmethod + def fixed(cls, split_size: int, *, strict_consistency: bool = True) -> "SplitKVSpec": + return cls( + mode=SplitKVMode.FIXED, + fixed_split_size=split_size, + strict_consistency=strict_consistency, + ) + + @classmethod + def auto(cls, *, strict_consistency: bool = False) -> "SplitKVSpec": + return cls(mode=SplitKVMode.AUTO, strict_consistency=strict_consistency) + + def __post_init__(self) -> None: + object.__setattr__(self, "mode", _enum_value(SplitKVMode, self.mode, "split_kv.mode")) + if not isinstance(self.strict_consistency, bool): + raise AttentionContractError("split_kv.strict_consistency must be a bool") + if self.mode is SplitKVMode.FIXED: + if self.fixed_split_size is None: + raise AttentionContractError( + "fixed Split-KV policy requires fixed_split_size" + ) + _positive_int(self.fixed_split_size, "split_kv.fixed_split_size") + elif self.fixed_split_size is not None: + raise AttentionContractError( + "fixed_split_size is only valid for fixed Split-KV policy" + ) + if self.strict_consistency and self.mode is SplitKVMode.AUTO: + raise AttentionContractError( + "auto Split-KV is runtime-shape dependent and is not allowed in strict consistency" + ) + + def resolve(self, total_kv_tokens: int, *, backend: str) -> SplitKVExecutionPlan: + """Resolve policies whose logical schedule is fully known by contract.""" + + total_kv_tokens = _positive_int(total_kv_tokens, "total_kv_tokens") + if self.mode is SplitKVMode.AUTO: + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=None, + actual_mode=None, + actual_split_size=None, + boundaries=(), + backend=backend, + source="runtime_required", + ) + split_size = total_kv_tokens if self.mode is SplitKVMode.DISABLED else self.fixed_split_size + assert split_size is not None + boundaries = tuple( + (start, min(start + split_size, total_kv_tokens)) + for start in range(0, total_kv_tokens, split_size) + ) + return SplitKVExecutionPlan( + requested_mode=self.mode, + requested_split_size=self.fixed_split_size, + actual_mode=self.mode, + actual_split_size=self.fixed_split_size, + boundaries=boundaries, + backend=backend, + source="contract_exact", + ) + + def to_dict(self) -> dict[str, Any]: + return { + "mode": self.mode.value, + "fixed_split_size": self.fixed_split_size, + "strict_consistency": self.strict_consistency, + } + + +def validate_split_kv_alignment( + training: SplitKVExecutionPlan, + rollout: SplitKVExecutionPlan, +) -> None: + """Fail closed unless train and rollout executed the same logical plan.""" + + if training.actual_mode is None or rollout.actual_mode is None: + raise AttentionContractError( + "strict Split-KV alignment requires actual runtime plans from both sides" + ) + fields = ( + "requested_mode", + "requested_split_size", + "actual_mode", + "actual_split_size", + "boundaries", + "merge_order", + "acc_dtype", + "downcast_at", + "fallback", + ) + mismatches = [ + field_name + for field_name in fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if mismatches: + raise AttentionContractError( + "training/rollout Split-KV execution plans differ: " + ", ".join(mismatches) + ) + + +@dataclass(frozen=True, order=True) +class SplitKVRuntimeCoordinate: + """Identity of one batch/rank/owner Split-KV runtime plan.""" + + batch_index: int + tp_rank: int + cp_rank: int + owner_cp_rank: int + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + ) -> None: + batch_index = _non_negative_int(self.batch_index, "batch_index") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + owner_cp_rank = _non_negative_int(self.owner_cp_rank, "owner_cp_rank") + if batch_index >= batch_size: + raise AttentionContractError("Split-KV batch_index is outside batch_size") + if tp_rank >= tp_world_size: + raise AttentionContractError("Split-KV tp_rank is outside tp_world_size") + if cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV cp_rank is outside cp_world_size") + if owner_cp_rank >= cp_world_size: + raise AttentionContractError("Split-KV owner_cp_rank is outside cp_world_size") + + def to_dict(self) -> dict[str, int]: + return { + "batch_index": self.batch_index, + "tp_rank": self.tp_rank, + "cp_rank": self.cp_rank, + "owner_cp_rank": self.owner_cp_rank, + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanEntry: + """Actual plan for one batch/TP/CP consumer and logical KV owner.""" + + coordinate: SplitKVRuntimeCoordinate + expected_kv_range: tuple[int, int] + execution: SplitKVExecutionPlan + + def validate( + self, + *, + batch_size: int, + tp_world_size: int, + cp_world_size: int, + total_kv_tokens: int, + ) -> None: + self.coordinate.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + ) + try: + start, end = self.expected_kv_range + except (TypeError, ValueError) as exc: + raise AttentionContractError( + "expected_kv_range must contain exactly (start, end)" + ) from exc + if ( + isinstance(start, bool) + or isinstance(end, bool) + or not isinstance(start, int) + or not isinstance(end, int) + or start < 0 + or end <= start + or end > total_kv_tokens + ): + raise AttentionContractError( + "expected_kv_range must satisfy 0 <= start < end <= total_kv_tokens" + ) + if self.execution.actual_mode is None: + raise AttentionContractError( + "complete Split-KV plan sets require actual runtime plans" + ) + if ( + self.execution.boundaries[0][0] != start + or self.execution.boundaries[-1][1] != end + ): + raise AttentionContractError( + "Split-KV execution boundaries must exactly cover expected_kv_range" + ) + if any( + boundary_start < start or boundary_end > end + for boundary_start, boundary_end in self.execution.boundaries + ): + raise AttentionContractError( + "Split-KV execution boundary escapes expected_kv_range" + ) + + def to_dict(self) -> dict[str, Any]: + return { + **self.coordinate.to_dict(), + "expected_kv_range": list(self.expected_kv_range), + **self.execution.to_dict(), + } + + +@dataclass(frozen=True) +class SplitKVRuntimePlanSet: + """Complete actual Split-KV plans across batch, TP, CP, and KV owners. + + Every CP consumer must report the plan used for every CP-owned KV range. + This duplicates owner plans across consumers intentionally: it detects one + rank silently choosing a different Split-K schedule or merge policy. + """ + + batch_size: int + tp_world_size: int + cp_world_size: int + total_kv_tokens: tuple[int, ...] + entries: tuple[SplitKVRuntimePlanEntry, ...] + + def __post_init__(self) -> None: + batch_size = _positive_int(self.batch_size, "batch_size") + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + totals = _integer_tuple(self.total_kv_tokens, "total_kv_tokens") + if len(totals) != batch_size or any(total <= 0 for total in totals): + raise AttentionContractError( + "total_kv_tokens must contain one positive length per batch item" + ) + object.__setattr__(self, "total_kv_tokens", totals) + entries = tuple(self.entries) + object.__setattr__(self, "entries", entries) + expected_coordinates = { + SplitKVRuntimeCoordinate(batch_index, tp_rank, cp_rank, owner_cp_rank) + for batch_index in range(batch_size) + for tp_rank in range(tp_world_size) + for cp_rank in range(cp_world_size) + for owner_cp_rank in range(cp_world_size) + } + actual_coordinates = [entry.coordinate for entry in entries] + if len(set(actual_coordinates)) != len(actual_coordinates): + raise AttentionContractError( + "Split-KV runtime plan set contains duplicate coordinates" + ) + missing = expected_coordinates.difference(actual_coordinates) + extra = set(actual_coordinates).difference(expected_coordinates) + if missing or extra: + raise AttentionContractError( + "Split-KV runtime plan set coordinate coverage is incomplete; " + f"missing={_format_split_kv_coordinates(missing)}, " + f"extra={_format_split_kv_coordinates(extra)}" + ) + for entry in entries: + entry.validate( + batch_size=batch_size, + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals[entry.coordinate.batch_index], + ) + self._validate_owner_coverage() + self._validate_rank_invariance() + + def _validate_owner_coverage(self) -> None: + for batch_index, total in enumerate(self.total_kv_tokens): + for tp_rank in range(self.tp_world_size): + ranges = [] + for owner_cp_rank in range(self.cp_world_size): + matches = [ + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.tp_rank == tp_rank + and entry.coordinate.cp_rank == 0 + and entry.coordinate.owner_cp_rank == owner_cp_rank + ] + ranges.append(matches[0].expected_kv_range) + previous_end = 0 + for start, end in ranges: + if start != previous_end: + raise AttentionContractError( + "Split-KV owner ranges must be gap-free in CP owner order" + ) + previous_end = end + if previous_end != total: + raise AttentionContractError( + "Split-KV owner ranges do not cover total_kv_tokens" + ) + + def _validate_rank_invariance(self) -> None: + for batch_index in range(self.batch_size): + for owner_cp_rank in range(self.cp_world_size): + entries = sorted( + ( + entry + for entry in self.entries + if entry.coordinate.batch_index == batch_index + and entry.coordinate.owner_cp_rank == owner_cp_rank + ), + key=lambda entry: ( + entry.coordinate.tp_rank, + entry.coordinate.cp_rank, + ), + ) + reference = entries[0] + for entry in entries[1:]: + if entry.expected_kv_range != reference.expected_kv_range: + raise AttentionContractError( + "Split-KV owner range differs across TP/CP consumers" + ) + try: + validate_split_kv_alignment( + reference.execution, + entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + "Split-KV plan differs across TP/CP consumers for " + f"batch={batch_index}, owner_cp={owner_cp_rank}: {exc}" + ) from exc + + def to_dict(self) -> dict[str, Any]: + return { + "batch_size": self.batch_size, + "tp_world_size": self.tp_world_size, + "cp_world_size": self.cp_world_size, + "total_kv_tokens": list(self.total_kv_tokens), + "entries": [ + entry.to_dict() + for entry in sorted(self.entries, key=lambda entry: entry.coordinate) + ], + "coverage": "complete_batch_tp_cp_owner_cartesian_product", + } + + +def validate_split_kv_plan_set_alignment( + training: SplitKVRuntimePlanSet, + rollout: SplitKVRuntimePlanSet, +) -> None: + """Fail closed unless complete train/rollout runtime plan sets align.""" + + topology_fields = ( + "batch_size", + "tp_world_size", + "cp_world_size", + "total_kv_tokens", + ) + topology_mismatches = [ + field_name + for field_name in topology_fields + if getattr(training, field_name) != getattr(rollout, field_name) + ] + if topology_mismatches: + raise AttentionContractError( + "training/rollout Split-KV plan-set topology differs: " + + ", ".join(topology_mismatches) + ) + training_by_coordinate = { + entry.coordinate: entry for entry in training.entries + } + rollout_by_coordinate = {entry.coordinate: entry for entry in rollout.entries} + if training_by_coordinate.keys() != rollout_by_coordinate.keys(): + raise AttentionContractError( + "training/rollout Split-KV plan-set coordinates differ" + ) + for coordinate in sorted(training_by_coordinate): + train_entry = training_by_coordinate[coordinate] + rollout_entry = rollout_by_coordinate[coordinate] + if train_entry.expected_kv_range != rollout_entry.expected_kv_range: + raise AttentionContractError( + f"training/rollout expected KV range differs at {coordinate}" + ) + try: + validate_split_kv_alignment( + train_entry.execution, + rollout_entry.execution, + ) + except AttentionContractError as exc: + raise AttentionContractError( + f"training/rollout Split-KV plan differs at {coordinate}: {exc}" + ) from exc + + +def _format_split_kv_coordinates( + coordinates: Iterable[SplitKVRuntimeCoordinate], +) -> list[dict[str, int]]: + return [coordinate.to_dict() for coordinate in sorted(coordinates)] + + +@dataclass(frozen=True) +class KVCacheSpec: + """Logical identity of the paged/block KV cache used for replay.""" + + cache_positions: tuple[int, ...] + kv_seq_lens: tuple[int, ...] + block_table: tuple[tuple[int, ...], ...] + global_token_positions: tuple[int, ...] + page_size: int + prefix_cache_enabled: bool = False + prefix_cache_key: str | None = None + shared_prefix_page_count: int = 0 + + def __post_init__(self) -> None: + cache_positions = _integer_tuple(self.cache_positions, "cache_positions") + kv_seq_lens = _integer_tuple(self.kv_seq_lens, "kv_seq_lens") + page_size = _positive_int(self.page_size, "page_size") + global_token_positions = _integer_tuple( + self.global_token_positions, "global_token_positions" + ) + if not cache_positions or any(position < 0 for position in cache_positions): + raise AttentionContractError("cache_positions must contain non-negative positions") + if not kv_seq_lens or any(length <= 0 for length in kv_seq_lens): + raise AttentionContractError("kv_seq_lens must contain positive sequence lengths") + if len(cache_positions) != len(kv_seq_lens): + raise AttentionContractError( + "cache_positions must contain one entry per kv_seq_lens entry" + ) + if not global_token_positions or any(position < 0 for position in global_token_positions): + raise AttentionContractError( + "global_token_positions must contain non-negative positions" + ) + if len(global_token_positions) != sum(kv_seq_lens): + raise AttentionContractError( + "global_token_positions must describe every logical cached token; " + f"expected {sum(kv_seq_lens)}, got {len(global_token_positions)}" + ) + token_offset = 0 + sequence_position_rows: list[tuple[int, ...]] = [] + for sequence_index, sequence_length in enumerate(kv_seq_lens): + sequence_positions = global_token_positions[ + token_offset : token_offset + sequence_length + ] + if any( + left >= right + for left, right in zip(sequence_positions, sequence_positions[1:], strict=False) + ): + raise AttentionContractError( + "global_token_positions must be strictly increasing within each sequence; " + f"sequence {sequence_index} is invalid" + ) + sequence_position_rows.append(sequence_positions) + token_offset += sequence_length + + for sequence_index, (cache_position, sequence_positions) in enumerate( + zip(cache_positions, sequence_position_rows, strict=True) + ): + terminal_position = sequence_positions[-1] + if cache_position != terminal_position: + raise AttentionContractError( + "cache_positions must equal the terminal global token position for each " + f"sequence; sequence {sequence_index} expected {terminal_position}, " + f"got {cache_position}" + ) + + try: + block_table = tuple(tuple(row) for row in self.block_table) + except TypeError as exc: + raise AttentionContractError( + "block_table must be a two-dimensional integer table" + ) from exc + if len(block_table) != len(kv_seq_lens) or any(not row for row in block_table): + raise AttentionContractError( + "block_table must contain one non-empty row per kv_seq_lens entry" + ) + active_block_rows: list[tuple[int, ...]] = [] + for row_index, (row, sequence_length) in enumerate( + zip(block_table, kv_seq_lens, strict=True) + ): + row_active_blocks: list[int] = [] + saw_padding = False + for column_index, block in enumerate(row): + if isinstance(block, bool) or not isinstance(block, int) or block < -1: + raise AttentionContractError( + "block_table entries must be integer block ids or -1 padding; " + f"got block_table[{row_index}][{column_index}]={block!r}" + ) + if block == -1: + saw_padding = True + continue + if saw_padding: + raise AttentionContractError( + "block_table -1 padding must be trailing; " + f"row {row_index} contains an active block after padding" + ) + row_active_blocks.append(block) + + expected_blocks = (sequence_length + page_size - 1) // page_size + if len(row_active_blocks) != expected_blocks: + raise AttentionContractError( + "block_table active page count must match kv_seq_lens and page_size; " + f"row {row_index} expected {expected_blocks}, got {len(row_active_blocks)}" + ) + if len(set(row_active_blocks)) != len(row_active_blocks): + raise AttentionContractError( + f"block_table row {row_index} contains duplicate active page ids" + ) + active_block_rows.append(tuple(row_active_blocks)) + + if not isinstance(self.prefix_cache_enabled, bool): + raise AttentionContractError("prefix_cache_enabled must be a bool") + shared_prefix_page_count = _non_negative_int( + self.shared_prefix_page_count, "shared_prefix_page_count" + ) + if self.prefix_cache_enabled and not self.prefix_cache_key: + raise AttentionContractError( + "prefix_cache_key is required when prefix_cache_enabled=True" + ) + if not self.prefix_cache_enabled and self.prefix_cache_key is not None: + raise AttentionContractError( + "prefix_cache_key must be None when prefix_cache_enabled=False" + ) + if not self.prefix_cache_enabled and shared_prefix_page_count != 0: + raise AttentionContractError( + "shared_prefix_page_count must be 0 when prefix_cache_enabled=False" + ) + + shared_prefix_pages: tuple[int, ...] = () + if shared_prefix_page_count > 0: + if any(len(row_blocks) < shared_prefix_page_count for row_blocks in active_block_rows): + raise AttentionContractError( + "shared_prefix_page_count exceeds an active block-table row" + ) + shared_prefix_token_count = shared_prefix_page_count * page_size + if any(length < shared_prefix_token_count for length in kv_seq_lens): + raise AttentionContractError( + "shared prefix pages must be fully populated and read-only" + ) + + shared_prefix_pages = active_block_rows[0][:shared_prefix_page_count] + shared_prefix_positions = sequence_position_rows[0][:shared_prefix_token_count] + for sequence_index, (row_blocks, positions) in enumerate( + zip(active_block_rows[1:], sequence_position_rows[1:], strict=True), start=1 + ): + if row_blocks[:shared_prefix_page_count] != shared_prefix_pages: + raise AttentionContractError( + "shared prefix page ids must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + if positions[:shared_prefix_token_count] != shared_prefix_positions: + raise AttentionContractError( + "shared prefix token positions must match across every sequence; " + f"sequence {sequence_index} is inconsistent" + ) + + exclusive_page_owners: dict[int, int] = {} + shared_prefix_page_ids = set(shared_prefix_pages) + for sequence_index, row_blocks in enumerate(active_block_rows): + for page_id in row_blocks[shared_prefix_page_count:]: + if page_id in shared_prefix_page_ids: + raise AttentionContractError( + "a writable suffix page cannot alias a read-only shared prefix page" + ) + previous_owner = exclusive_page_owners.get(page_id) + if previous_owner is not None: + raise AttentionContractError( + "active pages may be shared across sequences only when declared as " + "read-only prefix pages; " + f"page {page_id} is used by sequences {previous_owner} and " + f"{sequence_index}" + ) + exclusive_page_owners[page_id] = sequence_index + + object.__setattr__(self, "cache_positions", cache_positions) + object.__setattr__(self, "kv_seq_lens", kv_seq_lens) + object.__setattr__(self, "block_table", block_table) + object.__setattr__(self, "global_token_positions", global_token_positions) + + +@dataclass(frozen=True) +class RoPESpec: + """Qwen3 RoPE identity for attention inputs and KV-cache replay. + + The CP attention reference consumes attention inputs. This object records + whether those inputs are pre- or post-RoPE and pins the position/cache + metadata needed to compare fused ``RoPE+Attention`` and unfused + ``RoPE -> Attention`` materializations. + """ + + q_state: RoPEState = RoPEState.POST_ROPE + k_state: RoPEState = RoPEState.POST_ROPE + k_cache_state: RoPEState = RoPEState.POST_ROPE + theta: float = 1.0e6 + rotary_dim: int = 128 + rope_scaling: str | None = None + position_ids: tuple[int, ...] | None = None + query_position_offsets: tuple[int, ...] | None = None + key_position_offsets: tuple[int, ...] | None = None + cast_at: RoPECastPoint = RoPECastPoint.AFTER_ROPE + output_dtype: AttentionDType = AttentionDType.BF16 + fusion_boundary: RoPEFusionBoundary = RoPEFusionBoundary.UNFUSED_ROPE_ATTENTION + + def __post_init__(self) -> None: + object.__setattr__(self, "q_state", _enum_value(RoPEState, self.q_state, "q_state")) + object.__setattr__(self, "k_state", _enum_value(RoPEState, self.k_state, "k_state")) + object.__setattr__( + self, "k_cache_state", _enum_value(RoPEState, self.k_cache_state, "k_cache_state") + ) + if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + theta = float(self.theta) + if theta <= 0.0: + raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") + object.__setattr__(self, "theta", theta) + _positive_int(self.rotary_dim, "rotary_dim") + if self.rope_scaling is not None and ( + not isinstance(self.rope_scaling, str) or not self.rope_scaling.strip() + ): + raise AttentionContractError("rope_scaling must be a non-empty string when provided") + for field in ("position_ids", "query_position_offsets", "key_position_offsets"): + values = getattr(self, field) + if values is None: + continue + normalized = _integer_tuple(values, field) + if not normalized or any(value < 0 for value in normalized): + raise AttentionContractError(f"{field} must contain non-negative positions") + object.__setattr__(self, field, normalized) + object.__setattr__(self, "cast_at", _enum_value(RoPECastPoint, self.cast_at, "cast_at")) + object.__setattr__( + self, "output_dtype", _enum_value(AttentionDType, self.output_dtype, "output_dtype") + ) + object.__setattr__( + self, + "fusion_boundary", + _enum_value(RoPEFusionBoundary, self.fusion_boundary, "fusion_boundary"), + ) + + +@dataclass(frozen=True) +class AttentionContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: AttentionRole + mode: AttentionMode + dtype: AttentionDType + batch_size: int + query_sequence_length: int + head_dim: int + causal: bool + causal_offsets: tuple[int, ...] | None + sharding: ShardingSpec + reduction: ReductionSpec + split_kv: SplitKVSpec = field(default_factory=SplitKVSpec.disabled) + kv_cache: KVCacheSpec | None = None + rope: RoPESpec | None = None + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(AttentionRole, self.role, "role")) + object.__setattr__(self, "mode", _enum_value(AttentionMode, self.mode, "mode")) + object.__setattr__(self, "dtype", _enum_value(AttentionDType, self.dtype, "dtype")) + batch_size = _positive_int(self.batch_size, "batch_size") + query_sequence_length = _positive_int(self.query_sequence_length, "query_sequence_length") + _positive_int(self.head_dim, "head_dim") + if not isinstance(self.sharding, ShardingSpec): + raise AttentionContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise AttentionContractError("reduction must be a ReductionSpec") + if not isinstance(self.split_kv, SplitKVSpec): + raise AttentionContractError("split_kv must be a SplitKVSpec") + if ( + self.mode is AttentionMode.PREFILL + and query_sequence_length != self.sharding.local_sequence_length + ): + raise AttentionContractError( + "prefill query_sequence_length must equal sharding.local_sequence_length; " + f"got {query_sequence_length} and {self.sharding.local_sequence_length}" + ) + if self.sharding.packed_sequence_offsets is not None: + packed_sequence_count = len(self.sharding.packed_sequence_offsets) - 1 + if packed_sequence_count != batch_size: + raise AttentionContractError( + "packed sequence count must equal logical batch_size; " + f"got {packed_sequence_count} packed sequences and batch_size={batch_size}" + ) + if not isinstance(self.causal, bool): + raise AttentionContractError("causal must be a bool") + if self.causal: + if self.causal_offsets is None: + raise AttentionContractError("causal_offsets are required for causal attention") + if self.causal_offsets is not None: + causal_offsets = _integer_tuple(self.causal_offsets, "causal_offsets") + if not causal_offsets or any(offset < 0 for offset in causal_offsets): + raise AttentionContractError("causal_offsets must contain non-negative offsets") + if self.sharding.packed_sequence_offsets is not None: + expected_causal_offsets = batch_size + offset_owner = "packed sequence" + else: + expected_causal_offsets = batch_size + offset_owner = "batch entry" + if len(causal_offsets) != expected_causal_offsets: + raise AttentionContractError( + f"causal_offsets must contain one entry per {offset_owner}" + ) + object.__setattr__(self, "causal_offsets", causal_offsets) + + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise AttentionContractError( + "export_lse must be True for the WS2 attention-domain LSE contract" + ) + + if self.mode is AttentionMode.DECODE and self.kv_cache is None: + raise AttentionContractError("kv_cache metadata is required for decode attention") + if self.kv_cache is not None and not isinstance(self.kv_cache, KVCacheSpec): + raise AttentionContractError("kv_cache must be a KVCacheSpec when provided") + if self.rope is not None: + if not isinstance(self.rope, RoPESpec): + raise AttentionContractError("rope must be a RoPESpec when provided") + if self.rope.rotary_dim > self.head_dim: + raise AttentionContractError( + f"rotary_dim={self.rope.rotary_dim} must not exceed head_dim={self.head_dim}" + ) + if self.rope.position_ids is not None and len(self.rope.position_ids) not in { + query_sequence_length, + self.sharding.local_sequence_length, + }: + raise AttentionContractError( + "position_ids must describe the local query sequence or full local " + "sequence length" + ) + for field in ("query_position_offsets", "key_position_offsets"): + offsets = getattr(self.rope, field) + if offsets is not None and len(offsets) != batch_size: + raise AttentionContractError( + f"{field} must contain one entry per logical batch entry" + ) + if self.mode is AttentionMode.DECODE and self.kv_cache is not None: + if len(self.kv_cache.kv_seq_lens) != batch_size: + raise AttentionContractError( + "decode kv_seq_lens must contain one entry per batch entry" + ) + if len(self.kv_cache.cache_positions) != batch_size: + raise AttentionContractError( + "decode cache_positions must contain one entry per batch entry" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_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, + "local_q_heads": self.sharding.local_q_heads, + "local_kv_head_start": self.sharding.local_kv_head_start, + "local_kv_heads": self.sharding.local_kv_heads, + "global_sequence_length": self.sharding.global_sequence_length, + "local_sequence_length": self.sharding.local_sequence_length, + "global_block_indices": list(self.sharding.global_block_indices), + "global_block_token_starts": list(self.sharding.global_block_token_starts), + "local_block_offsets": list(self.sharding.local_block_offsets), + "packed_sequence_offsets": ( + list(self.sharding.packed_sequence_offsets) + if self.sharding.packed_sequence_offsets is not None + else None + ), + } + reduction = { + "merge": self.reduction.merge.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + } + kv_cache = None + if self.kv_cache is not None: + kv_cache = { + "cache_positions": list(self.kv_cache.cache_positions), + "kv_seq_lens": list(self.kv_cache.kv_seq_lens), + "block_table": [list(row) for row in self.kv_cache.block_table], + "global_token_positions": list(self.kv_cache.global_token_positions), + "page_size": self.kv_cache.page_size, + "prefix_cache_enabled": self.kv_cache.prefix_cache_enabled, + "prefix_cache_key": self.kv_cache.prefix_cache_key, + "shared_prefix_page_count": self.kv_cache.shared_prefix_page_count, + } + rope = None + if self.rope is not None: + rope = { + "q_state": self.rope.q_state.value, + "k_state": self.rope.k_state.value, + "k_cache_state": self.rope.k_cache_state.value, + "theta": self.rope.theta, + "rotary_dim": self.rope.rotary_dim, + "rope_scaling": self.rope.rope_scaling, + "position_ids": ( + list(self.rope.position_ids) if self.rope.position_ids is not None else None + ), + "query_position_offsets": ( + list(self.rope.query_position_offsets) + if self.rope.query_position_offsets is not None + else None + ), + "key_position_offsets": ( + list(self.rope.key_position_offsets) + if self.rope.key_position_offsets is not None + else None + ), + "cast_at": self.rope.cast_at.value, + "output_dtype": self.rope.output_dtype.value, + "fusion_boundary": self.rope.fusion_boundary.value, + } + return { + "semantic_operator": "standard_softmax_attention", + "role": self.role.value, + "mode": self.mode.value, + "dtype": self.dtype.value, + "batch_size": self.batch_size, + "query_sequence_length": self.query_sequence_length, + "head_dim": self.head_dim, + "causal": self.causal, + "causal_offsets": ( + list(self.causal_offsets) if self.causal_offsets is not None else None + ), + "export_lse": self.export_lse, + "lse_domain": "attention", + "sharding": sharding, + "reduction": reduction, + "split_kv": self.split_kv.to_dict(), + "kv_cache": kv_cache, + "rope": rope, + } + + +@dataclass(frozen=True) +class AttentionBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[AttentionRole] + modes: frozenset[AttentionMode] + dtypes: frozenset[AttentionDType] + cp_world_sizes: tuple[int, ...] + tp_world_sizes: tuple[int, ...] | None = None + exports_attention_lse: bool = False + deterministic_cp_merge: bool = False + supports_packed_varlen: bool = False + supports_kv_cache: bool = False + supports_rope_metadata: bool = False + supports_fused_rope_attention: bool = False + supports_split_kv_disabled: bool = True + supports_split_kv_fixed: bool = False + supports_split_kv_auto: bool = False + reports_actual_split_kv_plan: bool = False + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise AttentionContractError("backend_id must be a non-empty string") + roles = frozenset(_enum_value(AttentionRole, value, "roles") for value in self.roles) + modes = frozenset(_enum_value(AttentionMode, value, "modes") for value in self.modes) + dtypes = frozenset(_enum_value(AttentionDType, value, "dtypes") for value in self.dtypes) + if not roles or not modes or not dtypes: + raise AttentionContractError("backend roles, modes, and dtypes must not be empty") + cp_world_sizes = _integer_tuple(self.cp_world_sizes, "cp_world_sizes") + if not cp_world_sizes or any(size <= 0 for size in cp_world_sizes): + raise AttentionContractError("cp_world_sizes must contain positive values") + if len(set(cp_world_sizes)) != len(cp_world_sizes): + raise AttentionContractError("cp_world_sizes must not contain duplicates") + tp_world_sizes = None + if self.tp_world_sizes is not None: + tp_world_sizes = _integer_tuple(self.tp_world_sizes, "tp_world_sizes") + if not tp_world_sizes or any(size <= 0 for size in tp_world_sizes): + raise AttentionContractError("tp_world_sizes must contain positive values") + if len(set(tp_world_sizes)) != len(tp_world_sizes): + raise AttentionContractError("tp_world_sizes must not contain duplicates") + for field in ( + "exports_attention_lse", + "deterministic_cp_merge", + "supports_packed_varlen", + "supports_kv_cache", + "supports_rope_metadata", + "supports_fused_rope_attention", + "supports_split_kv_disabled", + "supports_split_kv_fixed", + "supports_split_kv_auto", + "reports_actual_split_kv_plan", + ): + if not isinstance(getattr(self, field), bool): + raise AttentionContractError(f"{field} must be a bool") + if self.implementation_kind not in {"production", "reference", "deterministic"}: + raise AttentionContractError( + "implementation_kind must be production, reference, or deterministic" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "modes", modes) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + + def incompatibilities(self, contract: AttentionContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.mode not in self.modes: + reasons.append(f"mode={contract.mode.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if contract.export_lse and not self.exports_attention_lse: + reasons.append("attention-domain LSE export is unsupported") + if cp_size > 1 and not self.deterministic_cp_merge: + reasons.append("deterministic CP (out, lse) merge is unsupported") + if ( + contract.sharding.packed_sequence_offsets is not None + and not self.supports_packed_varlen + ): + reasons.append("packed varlen layout is unsupported") + if contract.kv_cache is not None and not self.supports_kv_cache: + reasons.append("KV-cache identity materialization is unsupported") + if contract.rope is not None and not self.supports_rope_metadata: + reasons.append("RoPE/position metadata is unsupported") + if ( + contract.rope is not None + and contract.rope.fusion_boundary is RoPEFusionBoundary.FUSED_ROPE_ATTENTION + and not self.supports_fused_rope_attention + ): + reasons.append("fused RoPE+Attention boundary is unsupported") + split_support = { + SplitKVMode.DISABLED: self.supports_split_kv_disabled, + SplitKVMode.FIXED: self.supports_split_kv_fixed, + SplitKVMode.AUTO: self.supports_split_kv_auto, + } + if not split_support[contract.split_kv.mode]: + reasons.append( + f"Split-KV policy={contract.split_kv.mode.value} is unsupported" + ) + if contract.split_kv.strict_consistency and not self.reports_actual_split_kv_plan: + reasons.append("actual Split-KV execution-plan provenance is unsupported") + return tuple(reasons) + + def supports(self, contract: AttentionContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "modes": sorted(mode.value for mode in self.modes), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes), + "exports_attention_lse": self.exports_attention_lse, + "deterministic_cp_merge": self.deterministic_cp_merge, + "supports_packed_varlen": self.supports_packed_varlen, + "supports_kv_cache": self.supports_kv_cache, + "supports_rope_metadata": self.supports_rope_metadata, + "supports_fused_rope_attention": self.supports_fused_rope_attention, + "supports_split_kv_disabled": self.supports_split_kv_disabled, + "supports_split_kv_fixed": self.supports_split_kv_fixed, + "supports_split_kv_auto": self.supports_split_kv_auto, + "reports_actual_split_kv_plan": self.reports_actual_split_kv_plan, + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class AttentionDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: AttentionBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "AttentionContract", + "AttentionContractError", + "AttentionBackendCapability", + "AttentionDispatchResult", + "AttentionDType", + "AttentionMerge", + "AttentionMode", + "AttentionRole", + "DowncastPoint", + "KVCacheSpec", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "RoPECastPoint", + "RoPEFusionBoundary", + "RoPESpec", + "RoPEState", + "ShardingSpec", + "SplitKVExecutionPlan", + "SplitKVMode", + "SplitKVRuntimeCoordinate", + "SplitKVRuntimePlanEntry", + "SplitKVRuntimePlanSet", + "SplitKVSpec", + "validate_split_kv_alignment", + "validate_split_kv_plan_set_alignment", +] diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 75a2ab7f..c9777cf7 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -29,6 +29,7 @@ def make_operator_inputs( "matmul": _make_matmul_inputs, "det_gemm": _make_det_gemm_inputs, "attention": _make_attention_inputs, + "cp_attention": _make_cp_attention_inputs, "logp": _make_logp_inputs, "linear_logp": _make_linear_logp_inputs, "batch_invariant_logp": _make_batch_invariant_logp_inputs, @@ -53,6 +54,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}", @@ -139,6 +141,26 @@ def _make_attention_inputs( return inputs +def _make_cp_attention_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + return { + "q": _floating_tensor( + (batch, DEFAULT_N_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 0 + ), + "k": _floating_tensor( + (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 1 + ), + "v": _floating_tensor( + (batch, DEFAULT_N_KV_HEADS, seq, DEFAULT_HEAD_DIM), args, dtype, device, 2 + ), + "causal": True, + "cp_world_size": 2, + "kv_chunk_size": max(1, seq // 2), + } + + def _make_logp_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index c3d27848..76ba3edc 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -63,6 +63,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/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py new file mode 100644 index 00000000..c334f49a --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -0,0 +1,1035 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic context-parallel attention reference. + +This module is the correctness-first WS2 reference for CP-aware standard +softmax attention. It intentionally stays in PyTorch and uses fp32 partial +states so fused CUDA/Triton backends can validate their CP/LSE merge semantics +against a small, inspectable implementation. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional, Sequence + +import torch + +from rl_engine.kernels.attention_contract import ( + SplitKVExecutionPlan, + SplitKVMode, + SplitKVRuntimeCoordinate, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + +@dataclass(frozen=True) +class AttentionPartialState: + """One KV block's attention state before deterministic LSE merge. + + ``out`` is already normalized within the local KV block and has shape + ``[B, Hq, Sq, D]``. ``lse`` is the local attention-domain log-sum-exp with + shape ``[B, Hq, Sq]``. ``block_start`` / ``block_end`` are logical global KV + positions and define the canonical merge order. + """ + + out: torch.Tensor + lse: torch.Tensor + block_start: int + block_end: int + + def __post_init__(self) -> None: + if self.out.ndim != 4: + raise ValueError("partial attention out must have shape [B, Hq, Sq, D]") + if self.lse.shape != self.out.shape[:3]: + raise ValueError("partial attention lse must have shape [B, Hq, Sq]") + if self.block_start < 0: + raise ValueError("block_start must be non-negative") + if self.block_end < self.block_start: + raise ValueError("block_end must be >= block_start") + + +@dataclass(frozen=True) +class AttentionBackwardGradients: + """Training-side gradients emitted by the CP attention backward reference.""" + + dq: torch.Tensor + dk: torch.Tensor + dv: torch.Tensor + + +@dataclass(frozen=True) +class AttentionBackwardPathResult: + """One materialized CP attention backward path.""" + + name: str + out: torch.Tensor + lse: torch.Tensor + gradients: AttentionBackwardGradients + provenance: dict[str, object] + + +@dataclass(frozen=True) +class GradientDriftStats: + """Shape-aware absolute drift summary for backward validation reports.""" + + max_abs: float + mean_abs: float + p95_abs: float + p99_abs: float + active_count: int + + def to_dict(self) -> dict[str, object]: + return { + "max_abs": self.max_abs, + "mean_abs": self.mean_abs, + "p95_abs": self.p95_abs, + "p99_abs": self.p99_abs, + "active_count": self.active_count, + } + + +@dataclass(frozen=True) +class AttentionBackwardRankDrift: + """Backward drift for one logical CP rank's sequence ownership.""" + + rank: int + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + + def to_dict(self) -> dict[str, object]: + return { + "rank": self.rank, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + } + + +@dataclass(frozen=True) +class AttentionBackwardPathDrift: + """Candidate-vs-reference backward drift for one CP path.""" + + candidate_name: str + dq: GradientDriftStats + dk: GradientDriftStats + dv: GradientDriftStats + out: GradientDriftStats + lse: GradientDriftStats + per_rank: tuple[AttentionBackwardRankDrift, ...] + provenance: dict[str, object] + + def to_dict(self) -> dict[str, object]: + return { + "candidate_name": self.candidate_name, + "dq": self.dq.to_dict(), + "dk": self.dk.to_dict(), + "dv": self.dv.to_dict(), + "out": self.out.to_dict(), + "lse": self.lse.to_dict(), + "per_rank": [item.to_dict() for item in self.per_rank], + "provenance": self.provenance, + } + + +@dataclass(frozen=True) +class AttentionBackwardComparisonReport: + """Structured PR8 report for CP attention gradient drift validation.""" + + reference_name: str + drifts: tuple[AttentionBackwardPathDrift, ...] + + def to_dict(self) -> dict[str, object]: + return { + "reference_name": self.reference_name, + "drifts": [drift.to_dict() for drift in self.drifts], + } + + +def merge_attention_partial_states( + states: Sequence[AttentionPartialState], +) -> AttentionPartialState: + """Merge CP/chunk partial states in logical block order. + + The merge is the online-softmax/LSE merge used by attention, not a plain + sum. The input order is deliberately ignored: states are sorted by logical + ``block_start`` so the result depends on global block indices rather than + arrival order. + """ + + if not states: + raise ValueError("at least one attention partial state is required") + + ordered = sorted(states, key=lambda item: (item.block_start, item.block_end)) + _validate_merge_shapes_and_ranges(ordered) + + merged = ordered[0] + merged_out = merged.out.float() + merged_lse = merged.lse.float() + for state in ordered[1:]: + merged_out, merged_lse = _merge_two_states( + merged_out, + merged_lse, + state.out.float(), + state.lse.float(), + ) + + return AttentionPartialState( + out=merged_out, + lse=merged_lse, + block_start=ordered[0].block_start, + block_end=ordered[-1].block_end, + ) + + +class DeterministicCPAttentionReferenceOp: + """Correctness-first CP attention reference for prefill and chunked prefill. + + The reference consumes attention-ready Q/K. For Qwen3 WS2 this means Q/K + have already passed QK-Norm and RoPE unless an outer contract explicitly + marks them as pre-RoPE. RoPE is intentionally kept outside this CP merge + implementation so fused and unfused ``RoPE+Attention`` paths can compare the + same post-RoPE Q/K boundary before validating CP communication. + + The op emulates CP by splitting query and KV sequence dimensions into + logical CP shards. Each query shard computes one partial attention state per + KV block, then merges those states in fixed global-block order using fp32 + LSE arithmetic. ``forward`` returns the input dtype after the final write; + ``forward_fp32`` keeps the fp32 merged output. + """ + + op_class = "attention" + + @staticmethod + def split_kv_execution_plans( + total_kv_tokens: int, + *, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + ) -> list[dict[str, object]]: + """Export the actual logical Split-KV plan before execution.""" + + return split_kv_execution_plan_provenance( + total_kv_tokens, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + backend="deterministic_cp_reference", + ) + + def __call__( + 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: + 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, + ) + + 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.""" + + out, _ = self.forward_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, + 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. + """ + + 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(q.dtype if output_dtype is None else output_dtype) + return out, lse + + def forward_fp32_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, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return fp32 ``(out, lse)`` for the CP reference path.""" + + return self.forward_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, + output_dtype=torch.float32, + ) + + def backward_reference( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: 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] = torch.float32, + name: Optional[str] = None, + ) -> AttentionBackwardPathResult: + """Run the deterministic training-side backward validation path. + + The semantic backward input is ``dout`` plus the forward attention state + produced from the same Q/K/V, masks, position offsets, CP world, and KV + block order. The reference keeps the softmax/merge math in fp32 and + records the final-write dtype in provenance; decode backward is + intentionally out of scope for PR8. + """ + + _validate_qkv(q, k, v) + if dout.shape != q.shape: + raise ValueError("dout must have shape [B, Hq, Sq, D], matching q") + if not torch.is_floating_point(dout) or torch.is_complex(dout): + raise ValueError("dout must be a real floating-point tensor") + 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 + out, lse = self.forward_with_lse( + q_leaf, + k_leaf, + v_leaf, + 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, + output_dtype=resolved_output_dtype, + ) + torch.autograd.backward(out, dout.to(device=out.device, 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") + + return AttentionBackwardPathResult( + name=name + or _backward_path_name(cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size), + out=out.detach(), + lse=lse.detach(), + gradients=AttentionBackwardGradients( + dq=q_leaf.grad.detach(), + dk=k_leaf.grad.detach(), + dv=v_leaf.grad.detach(), + ), + provenance={ + "attention_mode": "prefill" if kv_chunk_size is None else "chunked_prefill", + "gradient_mode": "training_backward", + "gradient_inputs": ["q", "k", "v"], + "gradient_outputs": ["out"], + "saved_forward_state": [ + "out", + "attention_lse", + "causal_mask", + "key_padding_mask", + "query_position_offsets", + "key_position_offsets", + "global_block_index", + ], + "cp_world_size": cp_world_size, + "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", + ), + "merge_order": "global_block_index", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "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.", ""), + "te_backward_oracle": "not_used", + "decode_backward": "not_supported", + }, + ) + + def local_partial_state( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + q_start: int, + k_start: int, + total_kv_len: int, + total_query_len: Optional[int] = None, + 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, + ) -> AttentionPartialState: + """Compute one query shard against one logical KV block. + + ``query_position_offsets`` and ``key_position_offsets`` are optional + per-batch-row base positions. They let the reference express varlen or + packed metadata while retaining the dense [B, H, S, D] tensor layout. + For post-RoPE Q/K, these offsets must describe the same absolute token + positions used when RoPE was applied. + """ + + _validate_qkv(q, k, v) + if q_start < 0 or k_start < 0: + raise ValueError("q_start and k_start must be non-negative") + if total_kv_len < k_start + k.size(2): + raise ValueError("total_kv_len must cover the local KV block") + if total_query_len is None: + total_query_len = q.size(2) + if total_query_len < q_start + q.size(2): + raise ValueError("total_query_len must cover the local query block") + if key_padding_mask is not None: + if key_padding_mask.shape != (q.size(0), k.size(2)): + raise ValueError("local key_padding_mask must have shape [B, local_skv]") + if key_padding_mask.dtype != torch.bool: + raise ValueError("local key_padding_mask must be bool") + query_offsets = _normalize_position_offsets( + query_position_offsets, + q.size(0), + q.device, + default=total_kv_len - total_query_len, + name="query_position_offsets", + ) + key_offsets = _normalize_position_offsets( + key_position_offsets, + q.size(0), + q.device, + default=0, + name="key_position_offsets", + ) + + ctx = NativeAttentionOp._strict_fp32_math(q.device.type) + with ctx: + qf = q.float() + kf = k.float() + vf = v.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hkv != hq: + repeat = hq // hkv + kf = kf.repeat_interleave(repeat, dim=1) + vf = vf.repeat_interleave(repeat, dim=1) + + if skv == 0: + zero_dep = _zero_dependency(qf, kf, vf) + return AttentionPartialState( + out=torch.zeros(q.size(0), hq, sq, dim, device=q.device, dtype=torch.float32) + + zero_dep, + lse=torch.full( + (q.size(0), hq, sq), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep, + block_start=k_start, + block_end=k_start, + ) + + scale_value = scale if scale is not None else (1.0 / math.sqrt(dim)) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * scale_value + if causal: + query_base = query_offsets[:, None] + q_start + key_base = key_offsets[:, None] + k_start + q_pos = torch.arange(sq, device=q.device, dtype=torch.long) + query_base + k_pos = torch.arange(skv, device=q.device, dtype=torch.long) + key_base + causal_mask = k_pos[:, None, :] > q_pos[:, :, None] + scores = scores.masked_fill(causal_mask[:, None, :, :], float("-inf")) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + + lse = torch.logsumexp(scores, dim=-1) + finite_lse = torch.isfinite(lse) + weights = torch.exp(scores - lse.unsqueeze(-1)) + weights = torch.where(finite_lse.unsqueeze(-1), weights, torch.zeros_like(weights)) + out = torch.matmul(weights, vf) + return AttentionPartialState( + out=out, + lse=lse, + block_start=k_start, + block_end=k_start + skv, + ) + + def _forward_impl( + 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], + cp_world_size: int, + kv_chunk_size: Optional[int], + ) -> tuple[torch.Tensor, torch.Tensor]: + _validate_qkv(q, k, v) + if cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + + 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", + ) + + q_bounds = _split_bounds(sq, cp_world_size) + kv_bounds = _kv_block_bounds(skv, cp_world_size, kv_chunk_size) + out_chunks: list[torch.Tensor] = [] + lse_chunks: list[torch.Tensor] = [] + for q_start, q_end in q_bounds: + if q_start == q_end: + continue + q_block = q[:, :, q_start:q_end, :] + states = [ + self.local_partial_state( + q_block, + k[:, :, k_start:k_end, :], + v[:, :, k_start:k_end, :], + q_start=q_start, + k_start=k_start, + total_kv_len=skv, + total_query_len=sq, + causal=causal, + scale=scale, + key_padding_mask=( + None if key_padding_mask is None else key_padding_mask[:, k_start:k_end] + ), + query_position_offsets=query_offsets, + key_position_offsets=key_offsets, + ) + for k_start, k_end in kv_bounds + if k_start != k_end + ] + if states: + merged = merge_attention_partial_states(states) + out_chunks.append(merged.out) + lse_chunks.append(merged.lse) + else: + zero_dep = _zero_dependency(q_block.float(), k.float(), v.float()) + out_chunks.append( + torch.zeros(batch, hq, q_end - q_start, dim, device=q.device) + zero_dep + ) + lse_chunks.append( + torch.full( + (batch, hq, q_end - q_start), + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + + zero_dep + ) + + if not out_chunks: + 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, + ) + return torch.cat(out_chunks, dim=2), torch.cat(lse_chunks, dim=2) + + +def compare_cp_attention_backward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dout: 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, + candidate_cp_world_size: int = 2, + candidate_kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = torch.float32, +) -> AttentionBackwardComparisonReport: + """Compare CP=1 backward with a CP/chunked-prefill candidate. + + The report includes whole-tensor ``dq/dk/dv`` drift and per-logical-CP-rank + slices. It is a validation/reporting helper, not a separate production + backward kernel. + """ + + op = DeterministicCPAttentionReferenceOp() + reference = op.backward_reference( + q, + k, + v, + dout, + 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=1, + kv_chunk_size=None, + output_dtype=output_dtype, + name="cp1_backward_reference", + ) + candidate = op.backward_reference( + q, + k, + v, + dout, + 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=candidate_cp_world_size, + kv_chunk_size=candidate_kv_chunk_size, + output_dtype=output_dtype, + ) + return AttentionBackwardComparisonReport( + reference_name=reference.name, + drifts=(_compare_backward_path(candidate, reference),), + ) + + +def _compare_backward_path( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, +) -> AttentionBackwardPathDrift: + cp_world_size = _provenance_int(candidate.provenance, "cp_world_size") + return AttentionBackwardPathDrift( + candidate_name=candidate.name, + dq=_drift_stats(candidate.gradients.dq, reference.gradients.dq), + dk=_drift_stats(candidate.gradients.dk, reference.gradients.dk), + dv=_drift_stats(candidate.gradients.dv, reference.gradients.dv), + out=_drift_stats(candidate.out, reference.out), + lse=_drift_stats(candidate.lse, reference.lse), + per_rank=_per_rank_backward_drifts(candidate, reference, cp_world_size), + provenance=candidate.provenance, + ) + + +def _per_rank_backward_drifts( + candidate: AttentionBackwardPathResult, + reference: AttentionBackwardPathResult, + cp_world_size: int, +) -> tuple[AttentionBackwardRankDrift, ...]: + q_bounds = _split_bounds(candidate.gradients.dq.size(2), cp_world_size) + kv_bounds = _split_bounds(candidate.gradients.dk.size(2), cp_world_size) + per_rank = [] + for rank, ((q_start, q_end), (kv_start, kv_end)) in enumerate(zip(q_bounds, kv_bounds)): + per_rank.append( + AttentionBackwardRankDrift( + rank=rank, + dq=_drift_stats( + candidate.gradients.dq[:, :, q_start:q_end, :], + reference.gradients.dq[:, :, q_start:q_end, :], + ), + dk=_drift_stats( + candidate.gradients.dk[:, :, kv_start:kv_end, :], + reference.gradients.dk[:, :, kv_start:kv_end, :], + ), + dv=_drift_stats( + candidate.gradients.dv[:, :, kv_start:kv_end, :], + reference.gradients.dv[:, :, kv_start:kv_end, :], + ), + ) + ) + return tuple(per_rank) + + +def _drift_stats(candidate: torch.Tensor, reference: torch.Tensor) -> GradientDriftStats: + if candidate.shape != reference.shape: + raise ValueError( + f"candidate shape {tuple(candidate.shape)} must match " + f"reference shape {tuple(reference.shape)}" + ) + diff = (candidate.float() - reference.float()).abs().reshape(-1) + active_count = int(diff.numel()) + if active_count == 0: + return GradientDriftStats(0.0, 0.0, 0.0, 0.0, 0) + return GradientDriftStats( + max_abs=float(diff.max().item()), + mean_abs=float(diff.mean().item()), + p95_abs=float(torch.quantile(diff, 0.95).item()), + p99_abs=float(torch.quantile(diff, 0.99).item()), + active_count=active_count, + ) + + +def _backward_path_name(*, cp_world_size: int, kv_chunk_size: Optional[int]) -> str: + prefix = f"cp{cp_world_size}" + if kv_chunk_size is None: + return f"{prefix}_backward" + return f"{prefix}_chunked_backward" + + +def _provenance_int(provenance: dict[str, object], key: str) -> int: + value = provenance[key] + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"provenance field {key!r} must be an int") + return value + + +def _merge_two_states( + out_a: torch.Tensor, + lse_a: torch.Tensor, + out_b: torch.Tensor, + lse_b: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + merged_lse = torch.logaddexp(lse_a, lse_b) + finite = torch.isfinite(merged_lse) + weight_a = torch.where(finite, torch.exp(lse_a - merged_lse), torch.zeros_like(merged_lse)) + weight_b = torch.where(finite, torch.exp(lse_b - merged_lse), torch.zeros_like(merged_lse)) + merged_out = weight_a.unsqueeze(-1) * out_a + weight_b.unsqueeze(-1) * out_b + return merged_out, merged_lse + + +def _validate_merge_shapes_and_ranges(states: Sequence[AttentionPartialState]) -> None: + first = states[0] + previous_end = first.block_end + for state in states[1:]: + if state.out.shape != first.out.shape or state.lse.shape != first.lse.shape: + raise ValueError("all partial states must have matching out/lse shapes") + if state.block_start != previous_end: + raise ValueError("partial state block ranges must be gap-free and non-overlapping") + previous_end = state.block_end + + +def _validate_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q, k, and v must have shape [B, H, S, D]") + if k.shape != v.shape: + raise ValueError("k and v must have the same shape") + if q.size(0) != k.size(0) or q.size(3) != k.size(3): + raise ValueError("q, k, and v must share batch size and head dim") + if q.size(1) % k.size(1) != 0: + raise ValueError(f"Hq={q.size(1)} not divisible by Hkv={k.size(1)} (GQA group)") + + +def _zero_dependency(*tensors: torch.Tensor) -> torch.Tensor: + total = torch.tensor(0.0, device=tensors[0].device) + for tensor in tensors: + total = total + tensor.sum() + return total * 0.0 + + +def _normalize_position_offsets( + offsets: Optional[torch.Tensor], + batch: int, + device: torch.device, + *, + default: int, + name: str, +) -> torch.Tensor: + if offsets is None: + return torch.full((batch,), default, dtype=torch.long, device=device) + if offsets.ndim != 1 or offsets.numel() != batch: + raise ValueError(f"{name} must have shape [B]") + if torch.is_floating_point(offsets) or torch.is_complex(offsets) or offsets.dtype == torch.bool: + raise ValueError(f"{name} must contain integer positions") + return offsets.to(device=device, dtype=torch.long) + + +def _split_bounds(length: int, parts: int) -> list[tuple[int, int]]: + base, extra = divmod(length, parts) + bounds: list[tuple[int, int]] = [] + start = 0 + for index in range(parts): + width = base + (1 if index < extra else 0) + end = start + width + bounds.append((start, end)) + start = end + return bounds + + +def _kv_block_bounds( + length: int, + cp_world_size: int, + kv_chunk_size: Optional[int], +) -> list[tuple[int, int]]: + bounds: list[tuple[int, int]] = [] + for start, end in _split_bounds(length, cp_world_size): + if kv_chunk_size is None: + bounds.append((start, end)) + continue + cursor = start + while cursor < end: + chunk_end = min(cursor + kv_chunk_size, end) + bounds.append((cursor, chunk_end)) + cursor = chunk_end + return bounds + + +def split_kv_execution_plan_provenance( + length: int, + *, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str, +) -> list[dict[str, object]]: + """Return the actual backend-local Split-KV plan for every CP owner.""" + + if length < 1: + raise ValueError("Split-KV sequence length must be >= 1") + if cp_world_size < 1: + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + result: list[dict[str, object]] = [] + for owner_cp_rank, (rank_start, rank_end) in enumerate( + _split_bounds(length, cp_world_size) + ): + if rank_start == rank_end: + continue + if kv_chunk_size is None: + boundaries = ((rank_start, rank_end),) + mode = SplitKVMode.DISABLED + else: + boundaries = tuple( + (start, min(start + kv_chunk_size, rank_end)) + for start in range(rank_start, rank_end, kv_chunk_size) + ) + mode = SplitKVMode.FIXED + plan = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + result.append({"owner_cp_rank": owner_cp_rank, **plan.to_dict()}) + return result + + +def build_reference_split_kv_runtime_plan_set( + total_kv_tokens: Sequence[int], + *, + tp_world_size: int, + cp_world_size: int, + kv_chunk_size: Optional[int], + backend: str = "deterministic_cp_reference", +) -> SplitKVRuntimePlanSet: + """Build complete per-batch/TP/CP/owner plans for the reference path.""" + + totals = tuple(total_kv_tokens) + if not totals or any(total < cp_world_size for total in totals): + raise ValueError( + "reference runtime plan sets require at least one KV token per CP owner" + ) + if tp_world_size < 1 or cp_world_size < 1: + raise ValueError("TP and CP world sizes must be >= 1") + if kv_chunk_size is not None and kv_chunk_size < 1: + raise ValueError("kv_chunk_size must be >= 1 when provided") + + entries: list[SplitKVRuntimePlanEntry] = [] + for batch_index, total in enumerate(totals): + owner_ranges = _split_bounds(total, cp_world_size) + 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): + if kv_chunk_size is None: + mode = SplitKVMode.DISABLED + boundaries = ((owner_start, owner_end),) + else: + mode = SplitKVMode.FIXED + boundaries = tuple( + (start, min(start + kv_chunk_size, owner_end)) + for start in range(owner_start, owner_end, kv_chunk_size) + ) + execution = SplitKVExecutionPlan( + requested_mode=mode, + requested_split_size=kv_chunk_size, + actual_mode=mode, + actual_split_size=kv_chunk_size, + boundaries=boundaries, + backend=backend, + source="reference_execution", + ) + entries.append( + SplitKVRuntimePlanEntry( + coordinate=SplitKVRuntimeCoordinate( + batch_index=batch_index, + tp_rank=tp_rank, + cp_rank=cp_rank, + owner_cp_rank=owner_cp_rank, + ), + expected_kv_range=(owner_start, owner_end), + execution=execution, + ) + ) + return SplitKVRuntimePlanSet( + batch_size=len(totals), + tp_world_size=tp_world_size, + cp_world_size=cp_world_size, + total_kv_tokens=totals, + entries=tuple(entries), + ) + + +CPAttentionReferenceOp = DeterministicCPAttentionReferenceOp + +__all__ = [ + "AttentionBackwardComparisonReport", + "AttentionBackwardGradients", + "AttentionBackwardPathDrift", + "AttentionBackwardPathResult", + "AttentionBackwardRankDrift", + "AttentionPartialState", + "build_reference_split_kv_runtime_plan_set", + "CPAttentionReferenceOp", + "DeterministicCPAttentionReferenceOp", + "GradientDriftStats", + "compare_cp_attention_backward", + "merge_attention_partial_states", + "split_kv_execution_plan_provenance", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 8d9ff739..49f89e74 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -97,6 +97,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_KV_CACHE_ATTN = ( "rl_engine.kernels.ops.pytorch.attention.kv_cache.NativeKVCacheAttnOp" ) + # WS2 correctness-first context-parallel attention reference. It emulates + # CP prefill/chunked-prefill with fp32 attention-domain LSE merges. + PYTORCH_CP_ATTENTION = ( + "rl_engine.kernels.ops.pytorch.attention.cp_attention." + "DeterministicCPAttentionReferenceOp" + ) # WS1 pure-PyTorch ground-truth linear ops PYTORCH_NATIVE_LM_HEAD = "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp" # WS1 pure-PyTorch ground-truth embedding ops @@ -193,6 +199,7 @@ def __init__(self): OpBackend.CUDA_DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION, ], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], @@ -226,6 +233,7 @@ def __init__(self): OpBackend.TRITON_GENERIC, ], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], @@ -250,6 +258,7 @@ def __init__(self): "logp_deterministic_indexed": [OpBackend.PYTORCH_NATIVE], "attn": [OpBackend.PYTORCH_ATTN], "attention": [OpBackend.PYTORCH_NATIVE_ATTENTION], + "cp_attention": [OpBackend.PYTORCH_CP_ATTENTION], "kv_cache_attention": [OpBackend.PYTORCH_NATIVE_KV_CACHE_ATTN], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], diff --git a/scripts/ws2_p2p_nccl_attention_reference_check.py b/scripts/ws2_p2p_nccl_attention_reference_check.py new file mode 100644 index 00000000..8d95c99d --- /dev/null +++ b/scripts/ws2_p2p_nccl_attention_reference_check.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Two-GPU P2P NCCL reference check for issue #235. + +Run with: + + torchrun --standalone --nproc-per-node=2 \ + scripts/ws2_p2p_nccl_attention_reference_check.py +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Sequence + +import torch +import torch.distributed as dist + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionCPMergedState, + AttentionCPPartialState, + AttentionParallelSpec, + P2PNCCLAttentionCPCommunication, +) +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + merge_attention_partial_states, +) + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batch", type=int, default=2) + parser.add_argument("--seq-len", type=int, default=16) + parser.add_argument("--q-heads", type=int, default=16) + parser.add_argument("--kv-heads", type=int, default=4) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--chunk-size", type=int, default=4) + parser.add_argument("--seed", type=int, default=2357) + parser.add_argument("--atol", type=float, default=2.0e-4) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + raise RuntimeError("this check requires at least two visible CUDA devices") + dist.init_process_group("nccl", init_method="env://") + try: + world_size = dist.get_world_size() + rank = dist.get_rank() + if world_size != 2: + raise RuntimeError("this reference check requires exactly two NCCL ranks") + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + result = run_check(args, rank=rank, device=device) + failures = torch.tensor( + [0 if result["passed"] else 1], + dtype=torch.int32, + device=device, + ) + dist.all_reduce(failures, op=dist.ReduceOp.SUM) + result["global_failure_count"] = int(failures.item()) + reports: list[dict[str, object] | None] = [None] * world_size + dist.all_gather_object(reports, result) + if rank == 0: + print(json.dumps({"ranks": reports}, indent=2, sort_keys=True)) + return 0 if int(failures.item()) == 0 else 1 + finally: + dist.destroy_process_group() + + +def run_check( + args: argparse.Namespace, + *, + rank: int, + device: torch.device, +) -> dict[str, object]: + if args.seq_len < 2 or args.seq_len % 2 != 0: + raise ValueError("seq_len must be positive and divisible by CP=2") + if args.chunk_size < 1: + raise ValueError("chunk_size must be positive") + if args.q_heads % args.kv_heads != 0: + raise ValueError("q_heads must be divisible by kv_heads") + + generator = torch.Generator(device="cpu").manual_seed(args.seed) + shape_q = (args.batch, args.q_heads, args.seq_len, args.head_dim) + shape_kv = (args.batch, args.kv_heads, args.seq_len, args.head_dim) + q = torch.randn(shape_q, generator=generator, dtype=torch.bfloat16).to(device) + k = torch.randn(shape_kv, generator=generator, dtype=torch.bfloat16).to(device) + v = torch.randn(shape_kv, generator=generator, dtype=torch.bfloat16).to(device) + owner_ranges = ((0, args.seq_len // 2), (args.seq_len // 2, args.seq_len)) + query_ranges = owner_ranges + blocks: list[AttentionCPBlockMetadata] = [] + for owner, (owner_start, owner_end) in enumerate(owner_ranges): + for start in range(owner_start, owner_end, args.chunk_size): + blocks.append( + AttentionCPBlockMetadata( + global_block_index=len(blocks), + kv_block_start=start, + kv_block_end=min(start + args.chunk_size, owner_end), + owner_cp_rank=owner, + owner_tp_rank=0, + ) + ) + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec( + tp_world_size=2, + tp_rank=0, + cp_world_size=2, + cp_rank=rank, + ), + backend="p2p_nccl_reference", + status="implemented", + expected_blocks=tuple(blocks), + expected_kv_token_range=(0, args.seq_len), + query_token_ranges=query_ranges, + ) + reference = DeterministicCPAttentionReferenceOp() + local_states: list[AttentionCPPartialState] = [] + for block in reversed(blocks): + if block.owner_cp_rank != rank: + continue + state = reference.local_partial_state( + q, + k[:, :, block.kv_block_start : block.kv_block_end, :], + v[:, :, block.kv_block_start : block.kv_block_end, :], + q_start=0, + k_start=block.kv_block_start, + total_kv_len=args.seq_len, + total_query_len=args.seq_len, + causal=True, + ) + local_states.append( + AttentionCPPartialState(state.out, state.lse, block) + ) + + communication = P2PNCCLAttentionCPCommunication() + gathered = communication.all_gather_partial_states(tuple(local_states), plan) + merged = merge_attention_partial_states( + [ + AttentionPartialState( + state.out, + state.lse, + state.block.kv_block_start, + state.block.kv_block_end, + ) + for state in gathered + ] + ) + local = communication.reduce_scatter_merged_state( + AttentionCPMergedState(merged.out, merged.lse), + plan, + ) + full_out, full_lse = reference.forward_fp32_with_lse(q, k, v, causal=True) + start, end = query_ranges[rank] + out_max_abs = float((local.out - full_out[:, :, start:end, :]).abs().max().item()) + lse_max_abs = float((local.lse - full_lse[:, :, start:end]).abs().max().item()) + expected_indices = list(range(len(blocks))) + gathered_indices = [state.block.global_block_index for state in gathered] + passed = ( + gathered_indices == expected_indices + and out_max_abs <= args.atol + and lse_max_abs <= args.atol + ) + return { + "rank": rank, + "device": str(device), + "dtype": "bf16", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "transport": "p2p_nccl_reference", + "query_range": [start, end], + "gathered_block_indices": gathered_indices, + "out_max_abs": out_max_abs, + "lse_max_abs": lse_max_abs, + "atol": args.atol, + "passed": passed, + } + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_cp_attention.py b/tests/test_cp_attention.py new file mode 100644 index 00000000..196cfcaa --- /dev/null +++ b/tests/test_cp_attention.py @@ -0,0 +1,636 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the WS2 deterministic CP attention reference. + +The implementation is a correctness-first prefill/chunked-prefill reference: +local KV blocks produce ``(out, lse)`` partial states and CP merges those states +with fp32 online-softmax arithmetic in logical global-block order. +""" + +import contextlib +import json +import math + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + DeterministicCPAttentionReferenceOp, + compare_cp_attention_backward, + merge_attention_partial_states, + split_kv_execution_plan_provenance, +) +from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp +from rl_engine.kernels.ops.pytorch.rotary_embedding.rope import NativeRoPEOp +from rl_engine.kernels.registry import kernel_registry + +_N_HEADS = 32 +_N_KV = 8 +_HEAD_DIM = 128 +_ATOL = 3.0e-6 +_GRAD_ATOL = 1.0e-5 + + +@contextlib.contextmanager +def _single_thread(): + prev = torch.get_num_threads() + torch.set_num_threads(1) + try: + yield + finally: + torch.set_num_threads(prev) + + +def _qkv( + batch, + sq, + skv, + *, + seed, + dtype=torch.float32, + heads=_N_HEADS, + kv_heads=_N_KV, + dim=_HEAD_DIM, +): + gen = torch.Generator().manual_seed(seed) + q = torch.randn(batch, heads, sq, dim, generator=gen, dtype=dtype) + k = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + v = torch.randn(batch, kv_heads, skv, dim, generator=gen, dtype=dtype) + return q, k, v + + +def _full_lse(q, k, *, causal, scale=None, key_padding_mask=None): + qf, kf = q.float(), k.float() + hq, sq, dim = qf.shape[1], qf.shape[2], qf.shape[3] + hkv, skv = kf.shape[1], kf.shape[2] + if hq % hkv != 0: + raise ValueError("invalid GQA shape") + if hq != hkv: + kf = kf.repeat_interleave(hq // hkv, dim=1) + scores = torch.matmul(qf, kf.transpose(-1, -2)) * ( + scale if scale is not None else 1.0 / math.sqrt(dim) + ) + if causal: + query_pos = torch.arange(skv - sq, skv) + key_pos = torch.arange(skv) + scores = scores.masked_fill( + (key_pos[None, :] > query_pos[:, None])[None, None, :, :], + float("-inf"), + ) + if key_padding_mask is not None: + scores = scores.masked_fill(~key_padding_mask[:, None, None, :], float("-inf")) + return torch.logsumexp(scores, dim=-1) + + +def test_cp1_matches_native_attention_and_exports_lse(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 8, 8, seed=1) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + want = native.forward_fp32(q, k, v, causal=True) + want_lse = _full_lse(q, k, causal=True) + + torch.testing.assert_close(out, want, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse, want_lse, atol=_ATOL, rtol=0.0) + assert lse.dtype == torch.float32 + assert lse.shape == q.shape[:3] + + +def test_cp2_prefill_matches_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 9, 9, seed=2) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=1) + out2, lse2 = op.forward_fp32_with_lse(q, k, v, causal=True, cp_world_size=2) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + +def test_cp2_consumes_post_rope_qk_with_shared_global_position_metadata(): + op = DeterministicCPAttentionReferenceOp() + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 7, 7, seed=14, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([17, 103], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + + assert not torch.equal(q, pre_rope_q.float()) + assert not torch.equal(k, pre_rope_k.float()) + + with _single_thread(): + out1, lse1 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=1, + ) + out2, lse2 = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + cp_world_size=2, + kv_chunk_size=2, + ) + + torch.testing.assert_close(out2, out1, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(lse2, lse1, atol=_ATOL, rtol=0.0) + + +def test_chunked_prefill_replay_matches_unchunked_cp2(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 10, 10, seed=3) + + with _single_thread(): + unchunked_out, unchunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + ) + chunked_out, chunked_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=3, + ) + + torch.testing.assert_close(chunked_out, unchunked_out, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(chunked_lse, unchunked_lse, atol=_ATOL, rtol=0.0) + + +def test_causal_mask_uses_global_positions_across_cp_boundary(): + op = DeterministicCPAttentionReferenceOp() + batch, heads, kv_heads, seq, dim = 1, 2, 1, 5, 3 + q = torch.zeros(batch, heads, seq, dim) + k = torch.zeros(batch, kv_heads, seq, dim) + v = torch.arange(seq * dim, dtype=torch.float32).reshape(1, 1, seq, dim) + out = op.forward_fp32(q, k, v, causal=True, cp_world_size=2) + + expected = torch.stack([v[0, 0, : index + 1].mean(dim=0) for index in range(seq)]) + expected = expected.reshape(1, 1, seq, dim).repeat(1, heads, 1, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + + +def test_position_offsets_apply_varlen_causal_metadata_per_batch_row(): + op = DeterministicCPAttentionReferenceOp() + q = torch.zeros(2, 2, 2, 1) + k = torch.zeros(2, 1, 4, 1) + v = torch.arange(8, dtype=torch.float32).reshape(2, 1, 4, 1) + + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + query_position_offsets=torch.tensor([0, 11]), + key_position_offsets=torch.tensor([0, 10]), + cp_world_size=2, + kv_chunk_size=1, + ) + + expected = torch.tensor([0.0, 0.5, 4.5, 5.0]).reshape(2, 1, 2, 1).repeat(1, 2, 1, 1) + expected_lse = torch.log(torch.tensor([1.0, 2.0, 2.0, 3.0])).reshape(2, 1, 2) + expected_lse = expected_lse.repeat(1, 2, 1) + torch.testing.assert_close(out, expected, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(lse, expected_lse, atol=1.0e-6, rtol=0.0) + + +def test_merge_order_uses_global_block_index_not_arrival_order(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 6, 6, seed=4) + first = op.local_partial_state( + q, + k[:, :, :3], + v[:, :, :3], + q_start=0, + k_start=0, + total_kv_len=6, + causal=True, + ) + second = op.local_partial_state( + q, + k[:, :, 3:], + v[:, :, 3:], + q_start=0, + k_start=3, + total_kv_len=6, + causal=True, + ) + + forward = merge_attention_partial_states([first, second]) + reversed_arrival = merge_attention_partial_states([second, first]) + assert torch.equal(forward.out, reversed_arrival.out) + assert torch.equal(forward.lse, reversed_arrival.lse) + + +def test_key_padding_mask_and_all_masked_rows_are_stable(): + op = DeterministicCPAttentionReferenceOp() + native = NativeAttentionOp() + q, k, v = _qkv(2, 6, 6, seed=5) + mask = torch.tensor( + [ + [True, True, True, False, False, False], + [False, False, False, False, False, False], + ], + dtype=torch.bool, + ) + + with _single_thread(): + out, lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=False, + key_padding_mask=mask, + cp_world_size=2, + kv_chunk_size=2, + ) + want = native.forward_fp32(q, k, v, causal=False, key_padding_mask=mask) + + torch.testing.assert_close(out[:1], want[:1], atol=_ATOL, rtol=0.0) + assert torch.equal(out[1], torch.zeros_like(out[1])) + assert torch.isneginf(lse[1]).all() + assert torch.isfinite(out).all() + + +def test_empty_query_and_empty_kv_edges_are_stable(): + op = DeterministicCPAttentionReferenceOp() + q_empty = torch.randn(1, 2, 0, 4, requires_grad=True) + k_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + v_empty = torch.randn(1, 1, 0, 4, requires_grad=True) + out, lse = op.forward_fp32_with_lse(q_empty, k_empty, v_empty, cp_world_size=2) + assert out.shape == (1, 2, 0, 4) + assert lse.shape == (1, 2, 0) + assert out.requires_grad + out.sum().backward() + assert torch.equal(q_empty.grad, torch.zeros_like(q_empty)) + assert torch.equal(k_empty.grad, torch.zeros_like(k_empty)) + assert torch.equal(v_empty.grad, torch.zeros_like(v_empty)) + + q = torch.randn(1, 2, 3, 4) + out, lse = op.forward_fp32_with_lse(q, k_empty, v_empty, causal=False, cp_world_size=4) + assert torch.equal(out, torch.zeros_like(out)) + assert torch.isneginf(lse).all() + + +def test_empty_kv_backward_returns_zero_grads(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 2, 3, 4, requires_grad=True) + k = torch.randn(1, 1, 0, 4, requires_grad=True) + v = torch.randn(1, 1, 0, 4, requires_grad=True) + + out = op.forward_fp32(q, k, v, causal=False, cp_world_size=4) + assert out.requires_grad + out.sum().backward() + + assert torch.equal(q.grad, torch.zeros_like(q)) + assert torch.equal(k.grad, torch.zeros_like(k)) + assert torch.equal(v.grad, torch.zeros_like(v)) + + +def test_bf16_forward_uses_fp32_merge_then_final_write(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 8, 8, seed=6, dtype=torch.bfloat16) + + out, lse = op.forward_with_lse(q, k, v, causal=True, cp_world_size=2, kv_chunk_size=2) + fp32_out, fp32_lse = op.forward_fp32_with_lse( + q, + k, + v, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + + assert out.dtype == torch.bfloat16 + assert lse.dtype == torch.float32 + assert torch.equal(out, fp32_out.to(torch.bfloat16)) + assert torch.equal(lse, fp32_lse) + + +def test_cp2_chunked_gradients_match_cp1_reference(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 5, 5, seed=12, heads=4, kv_heads=2, dim=8) + q_ref = q.detach().clone().requires_grad_(True) + k_ref = k.detach().clone().requires_grad_(True) + v_ref = v.detach().clone().requires_grad_(True) + q_cp = q.detach().clone().requires_grad_(True) + k_cp = k.detach().clone().requires_grad_(True) + v_cp = v.detach().clone().requires_grad_(True) + gen = torch.Generator().manual_seed(13) + dy = torch.randn(1, 4, 5, 8, generator=gen) + + with _single_thread(): + out_ref = op.forward_fp32(q_ref, k_ref, v_ref, causal=True, cp_world_size=1) + out_cp = op.forward_fp32( + q_cp, + k_cp, + v_cp, + causal=True, + cp_world_size=2, + kv_chunk_size=2, + ) + out_ref.backward(dy) + out_cp.backward(dy) + + torch.testing.assert_close(out_cp, out_ref, atol=_ATOL, rtol=0.0) + torch.testing.assert_close(q_cp.grad, q_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(k_cp.grad, k_ref.grad, atol=1.0e-5, rtol=0.0) + torch.testing.assert_close(v_cp.grad, v_ref.grad, atol=1.0e-5, rtol=0.0) + + +def test_backward_report_cp2_prefill_matches_cp1_reference(): + q, k, v = _qkv(1, 5, 5, seed=15, heads=4, kv_heads=2, dim=8) + dout = torch.randn(1, 4, 5, 8, generator=torch.Generator().manual_seed(16)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + output_dtype=torch.float32, + ) + + assert report.reference_name == "cp1_backward_reference" + drift = report.drifts[0] + assert drift.candidate_name == "cp2_backward" + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + assert drift.out.max_abs <= _ATOL + assert drift.lse.max_abs <= _ATOL + assert len(drift.per_rank) == 2 + assert drift.per_rank[0].dq.active_count > 0 + assert drift.per_rank[1].dk.active_count > 0 + assert drift.provenance["saved_forward_state"][0] == "out" + assert drift.provenance["merge_order"] == "global_block_index" + assert drift.provenance["te_backward_oracle"] == "not_used" + assert drift.provenance["decode_backward"] == "not_supported" + json.dumps(report.to_dict()) + + +def test_backward_report_cp2_chunked_prefill_matches_cp1_reference(): + q, k, v = _qkv(1, 6, 6, seed=17, heads=4, kv_heads=2, dim=8) + dout = torch.randn(1, 4, 6, 8, generator=torch.Generator().manual_seed(18)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.float32, + ) + + drift = report.drifts[0] + assert drift.candidate_name == "cp2_chunked_backward" + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + assert drift.provenance["attention_mode"] == "chunked_prefill" + assert drift.provenance["kv_chunk_size"] == 2 + assert drift.provenance["requested_split_kv_policy"] == "fixed" + assert drift.provenance["actual_split_kv_plans"] == [ + { + "owner_cp_rank": 0, + "requested_split_kv_policy": "fixed", + "requested_split_kv_size": 2, + "actual_split_kv_policy": "fixed", + "actual_split_kv_size": 2, + "actual_split_kv_count": 2, + "actual_split_boundaries": [[0, 2], [2, 3]], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "deterministic_cp_backward_reference", + "split_kv_plan_source": "reference_execution", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + }, + { + "owner_cp_rank": 1, + "requested_split_kv_policy": "fixed", + "requested_split_kv_size": 2, + "actual_split_kv_policy": "fixed", + "actual_split_kv_size": 2, + "actual_split_kv_count": 2, + "actual_split_boundaries": [[3, 5], [5, 6]], + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_backend": "deterministic_cp_backward_reference", + "split_kv_plan_source": "reference_execution", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + }, + ] + + +def test_split_kv_plan_never_crosses_cp_owner_boundaries(): + plans = split_kv_execution_plan_provenance( + 10, + cp_world_size=3, + kv_chunk_size=3, + backend="test-reference", + ) + + assert [plan["actual_split_boundaries"] for plan in plans] == [ + [[0, 3], [3, 4]], + [[4, 7]], + [[7, 10]], + ] + assert [plan["owner_cp_rank"] for plan in plans] == [0, 1, 2] + + +def test_backward_report_preserves_post_rope_position_metadata(): + rope = NativeRoPEOp() + pre_rope_q, pre_rope_k, v = _qkv(2, 5, 5, seed=19, heads=4, kv_heads=2, dim=8) + position_offsets = torch.tensor([23, 101], dtype=torch.long) + positions = position_offsets[:, None] + torch.arange(pre_rope_q.size(2), dtype=torch.long) + q = rope.forward_fp32(pre_rope_q, positions, theta=1_000_000.0) + k = rope.forward_fp32(pre_rope_k, positions, theta=1_000_000.0) + dout = torch.randn(2, 4, 5, 8, generator=torch.Generator().manual_seed(20)) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + query_position_offsets=position_offsets, + key_position_offsets=position_offsets, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.float32, + ) + + drift = report.drifts[0] + assert drift.dq.max_abs <= _GRAD_ATOL + assert drift.dk.max_abs <= _GRAD_ATOL + assert drift.dv.max_abs <= _GRAD_ATOL + + +def test_qwen3_8b_local_tp2_cp2_bf16_backward_report_smoke(): + # Qwen3-8B global Hq/Hkv is 32/8. A TP=2 local shard owns 16/4 heads. + q, k, v = _qkv( + 1, + 4, + 4, + seed=21, + dtype=torch.bfloat16, + heads=16, + kv_heads=4, + dim=_HEAD_DIM, + ) + dout = torch.randn( + 1, + 16, + 4, + _HEAD_DIM, + generator=torch.Generator().manual_seed(22), + dtype=torch.bfloat16, + ) + + with _single_thread(): + report = compare_cp_attention_backward( + q, + k, + v, + dout, + causal=True, + candidate_cp_world_size=2, + candidate_kv_chunk_size=2, + output_dtype=torch.bfloat16, + ) + + drift = report.drifts[0] + assert drift.provenance["q_dtype"] == "bfloat16" + assert drift.provenance["output_dtype"] == "bfloat16" + assert drift.provenance["downcast_at"] == "final_write" + assert drift.dq.max_abs <= 5.0e-2 + assert drift.dk.max_abs <= 5.0e-2 + assert drift.dv.max_abs <= 5.0e-2 + + +def test_backward_report_validates_dout_shape_and_dtype(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=23, heads=4, kv_heads=2, dim=8) + + with pytest.raises(ValueError, match="dout must have shape"): + op.backward_reference(q, k, v, torch.randn(1, 4, 3, 8), cp_world_size=2) + + with pytest.raises(ValueError, match="dout must be a real floating-point tensor"): + op.backward_reference( + q, + k, + v, + torch.ones(1, 4, 4, 8, dtype=torch.long), + cp_world_size=2, + ) + + +def test_inputs_are_not_mutated(): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(2, 6, 6, seed=7) + mask = torch.ones(2, 6, dtype=torch.bool) + qc, kc, vc, mc = q.clone(), k.clone(), v.clone(), mask.clone() + + op.forward_fp32_with_lse(q, k, v, causal=True, key_padding_mask=mask, cp_world_size=2) + + assert torch.equal(q, qc) + assert torch.equal(k, kc) + assert torch.equal(v, vc) + assert torch.equal(mask, mc) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"cp_world_size": 0}, "cp_world_size"), + ({"cp_world_size": 2, "kv_chunk_size": 0}, "kv_chunk_size"), + ], +) +def test_invalid_parallelism_arguments_raise(kwargs, message): + op = DeterministicCPAttentionReferenceOp() + q, k, v = _qkv(1, 4, 4, seed=8) + with pytest.raises(ValueError, match=message): + op.forward_fp32_with_lse(q, k, v, causal=True, **kwargs) + + +def test_invalid_gqa_and_mask_shapes_raise(): + op = DeterministicCPAttentionReferenceOp() + q = torch.randn(1, 6, 4, _HEAD_DIM) + k = torch.randn(1, 4, 4, _HEAD_DIM) + v = torch.randn(1, 4, 4, _HEAD_DIM) + with pytest.raises(ValueError, match="not divisible"): + op.forward_fp32_with_lse(q, k, v, causal=True) + + q, k, v = _qkv(1, 4, 4, seed=9) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 3, dtype=torch.bool)) + with pytest.raises(ValueError, match="key_padding_mask"): + op.forward_fp32_with_lse(q, k, v, key_padding_mask=torch.ones(1, 4)) + with pytest.raises(ValueError, match="query_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + query_position_offsets=torch.ones(2, dtype=torch.long), + ) + with pytest.raises(ValueError, match="key_position_offsets"): + op.forward_fp32_with_lse( + q, + k, + v, + key_position_offsets=torch.ones(1, dtype=torch.float32), + ) + + +def test_overlapping_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="overlap"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=3), + AttentionPartialState(out=out, lse=lse, block_start=2, block_end=4), + ] + ) + + +def test_gapped_partial_ranges_raise(): + out = torch.zeros(1, 1, 1, 1) + lse = torch.zeros(1, 1, 1) + with pytest.raises(ValueError, match="gap-free"): + merge_attention_partial_states( + [ + AttentionPartialState(out=out, lse=lse, block_start=0, block_end=2), + AttentionPartialState(out=out, lse=lse, block_start=3, block_end=4), + ] + ) + + +def test_registry_dispatches_cp_attention_reference(): + assert isinstance(kernel_registry.get_op("cp_attention"), DeterministicCPAttentionReferenceOp) diff --git a/tests/test_cp_attention_transformer_engine.py b/tests/test_cp_attention_transformer_engine.py new file mode 100644 index 00000000..d98e31a1 --- /dev/null +++ b/tests/test_cp_attention_transformer_engine.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Optional Transformer Engine oracle tests for CP attention merging.""" + +from __future__ import annotations + +import importlib + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( + AttentionPartialState, + merge_attention_partial_states, +) + + +def _te_context_parallel_module(): + try: + return importlib.import_module( + "transformer_engine.pytorch.attention.dot_product_attention.context_parallel" + ) + except (ImportError, OSError, RuntimeError) as exc: + pytest.skip(f"Transformer Engine context-parallel attention is unavailable: {exc}") + + +def test_cp_attention_merge_matches_transformer_engine_corrections(): + te_cp = _te_context_parallel_module() + gen = torch.Generator().manual_seed(238) + out_a = torch.randn(2, 3, 5, 4, generator=gen) + out_b = torch.randn(2, 3, 5, 4, generator=gen) + lse_a = torch.randn(2, 3, 5, generator=gen) + lse_b = torch.randn(2, 3, 5, generator=gen) + + ours = merge_attention_partial_states( + [ + AttentionPartialState(out=out_b, lse=lse_b, block_start=5, block_end=9), + AttentionPartialState(out=out_a, lse=lse_a, block_start=0, block_end=5), + ] + ) + + te_lse = lse_a.clone() + te_cp.flash_attn_fwd_softmax_lse_correction(te_lse, lse_b) + te_out = te_cp.flash_attn_fwd_out_correction_init(out_a.clone(), te_lse, lse_a, seq_dim=2) + te_cp.flash_attn_fwd_out_correction(te_out, out_b, te_lse, lse_b, seq_dim=2) + + torch.testing.assert_close(ours.lse, te_lse, atol=1.0e-6, rtol=0.0) + torch.testing.assert_close(ours.out, te_out, atol=1.0e-6, rtol=0.0) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 3f2b4863..fee1cf94 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -9,6 +9,11 @@ import torch from rl_engine.kernels.gtest.operator_inputs import make_operator_inputs, operator_shape_name +from rl_engine.kernels.gtest.operator_specs import ( + make_candidate, + make_operator_case, + operator_names, +) def _args(**overrides): @@ -25,6 +30,7 @@ def _args(**overrides): "n_dim": 32, "theta": 1.0e6, "eps": 1.0e-6, + "arch_key": None, } values.update(overrides) return argparse.Namespace(**values) @@ -36,6 +42,7 @@ def _args(**overrides): "rms_norm", "matmul", "attention", + "cp_attention", "logp", "linear_logp", "batch_invariant_logp", @@ -80,6 +87,18 @@ def test_random_logp_inputs_are_seeded(): assert torch.equal(first["token_ids"], second["token_ids"]) +def test_cp_attention_operator_spec_registers_backward_grad_inputs(): + args = _args(op="cp_attention", input_mode="constant", batch=1, seq=2) + + assert "cp_attention" in operator_names() + case = make_operator_case(args, torch.float32, torch.device("cpu")) + candidate = make_candidate(argparse.Namespace(**{**vars(args), "candidate": "pytorch"})) + + assert case.op_class == "attention" + assert case.grad_input_names == ("q", "k", "v") + assert candidate.name == "pytorch-cp_attention" + + def test_constant_linear_logp_inputs_match_operator_contract(): args = _args(input_mode="constant", constant_value=0.5, token_value=3) inputs = make_operator_inputs("linear_logp", args, torch.float32, torch.device("cpu")) diff --git a/tests/test_ws2_cp_attention_drift_benchmark.py b/tests/test_ws2_cp_attention_drift_benchmark.py new file mode 100644 index 00000000..0f348505 --- /dev/null +++ b/tests/test_ws2_cp_attention_drift_benchmark.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Tests for the WS2 CP attention PR5 drift benchmark artifact.""" + +from __future__ import annotations + +import json + +import pytest + +from benchmarks.benchmark_ws2_cp_attention_drift import ( + SCHEMA_VERSION, + parse_args, + run_benchmark, + write_report, +) + + +def test_smoke_report_has_pr5_schema_and_qwen3_tp2_cp2_case(): + report = run_benchmark( + parse_args( + [ + "--smoke", + "--tp-world-sizes", + "2", + "--cp-world-sizes", + "2", + "--kv-chunk-sizes", + "none,1", + ] + ) + ) + + assert report["schema_version"] == SCHEMA_VERSION + assert report["report_family"] == "ws2_cross_config_drift_report" + assert report["tolerance_source"] == "#108" + assert report["issue"] == 235 + assert report["pr"] == 5 + assert report["target"]["model"] == "qwen3-8b" + assert report["target"]["global_num_query_heads"] == 32 + assert report["target"]["global_num_kv_heads"] == 8 + assert report["te_context_parallel_merge"]["te_module"].endswith("context_parallel") + assert report["dlogp"]["status"] == "not_available" + assert len(report["cases"]) == 2 + + names = {case["case_name"] for case in report["cases"]} + assert "qwen3_8b_tp2_cp2_prefill_bf16" in names + assert "qwen3_8b_tp2_cp2_chunk1_bf16" in names + + chunked = next(case for case in report["cases"] if case["attention_mode"] == "chunked_prefill") + assert chunked["topology"]["local_num_query_heads"] == 16 + assert chunked["topology"]["local_num_kv_heads"] == 4 + assert chunked["topology"]["local_query_head_range"] == [0, 16] + assert chunked["topology"]["local_kv_head_range"] == [0, 4] + assert chunked["provenance"]["merge_order"] == "global_block_index" + assert chunked["provenance"]["split_kv_policy"] == "fixed" + assert chunked["provenance"]["requested_split_kv_size"] == 1 + assert chunked["provenance"]["actual_split_kv_plans"][0][ + "actual_split_boundaries" + ] + assert chunked["provenance"]["actual_split_kv_plans"][0][ + "split_kv_accum_dtype" + ] == "fp32" + assert chunked["provenance"]["actual_split_kv_plans"][0][ + "split_kv_downcast_at" + ] == "final_write" + plan_set = chunked["provenance"]["actual_split_kv_plan_set"] + assert plan_set["coverage"] == "complete_batch_tp_cp_owner_cartesian_product" + assert len(plan_set["entries"]) == 8 + assert chunked["distributed_p2p_reference"]["status"] == "not_requested" + assert chunked["provenance"]["block_metadata_hash"] + assert chunked["provenance"]["rope"]["rope_state"] == "post_rope" + assert chunked["drift"]["rope"]["status"] == "available" + assert chunked["drift"]["cp_merge_fp32"]["out"]["max_abs"] <= 1.0e-5 + assert chunked["drift"]["cp_merge_fp32"]["lse"]["max_abs"] <= 1.0e-5 + assert chunked["merge_order_probe"]["out"]["max_abs"] == 0.0 + assert len(chunked["per_rank"]) == 2 + assert chunked["per_rank"][0]["out"]["active_count"] > 0 + + +def test_report_writes_reproducible_json_artifact(tmp_path): + output = tmp_path / "ws2-cp-attention-drift.json" + report = run_benchmark( + parse_args( + [ + "--smoke", + "--no-rope", + "--tp-world-sizes", + "1", + "--cp-world-sizes", + "1", + "--kv-chunk-sizes", + "none", + ] + ) + ) + + write_report(report, output) + loaded = json.loads(output.read_text(encoding="utf-8")) + + assert loaded["schema_version"] == SCHEMA_VERSION + assert loaded["cases"][0]["provenance"]["rope"]["rope_state"] == "not_composed" + assert loaded["cases"][0]["attention_mode"] == "prefill" + + +def test_include_backward_adds_pr8_gradient_drift_report(): + report = run_benchmark( + parse_args( + [ + "--smoke", + "--include-backward", + "--tp-world-sizes", + "2", + "--cp-world-sizes", + "2", + "--kv-chunk-sizes", + "1", + ] + ) + ) + + backward = report["cases"][0]["backward"] + assert backward["status"] == "available" + drift = backward["report"]["drifts"][0] + assert drift["candidate_name"] == "cp2_chunked_backward" + assert drift["provenance"]["attention_mode"] == "chunked_prefill" + assert drift["provenance"]["downcast_at"] == "final_write" + assert drift["dq"]["max_abs"] <= 5.0e-2 + assert drift["dk"]["max_abs"] <= 5.0e-2 + assert drift["dv"]["max_abs"] <= 5.0e-2 + assert len(drift["per_rank"]) == 2 + + +def test_invalid_qwen3_tp_topology_is_rejected(): + with pytest.raises(ValueError, match="query heads"): + run_benchmark( + parse_args( + [ + "--tp-world-sizes", + "3", + "--cp-world-sizes", + "1", + "--kv-chunk-sizes", + "none", + ] + ) + )