diff --git a/benchmarks/benchmark_ws2_cp_attention_drift.py b/benchmarks/benchmark_ws2_cp_attention_drift.py new file mode 100644 index 00000000..d64dc7d8 --- /dev/null +++ b/benchmarks/benchmark_ws2_cp_attention_drift.py @@ -0,0 +1,1233 @@ +# 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 ( # noqa: E402 + 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.pytorch.rotary_embedding.rope import NativeRoPEOp # noqa: E402 +from rl_engine.testing.reference_ops import selected_logprobs_reference # noqa: E402 + +SCHEMA_VERSION = "ws2_cp_attention_drift/v2" +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( + "--include-dlogp", + action="store_true", + help=( + "Project attention outputs through a deterministic synthetic lm_head " + "and report active-token dlogp drift." + ), + ) + 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": "requested" if args.include_dlogp else "not_requested", + "reason": None if args.include_dlogp else "selected-logprob chain was not requested", + "source": "synthetic_fp32_lm_head_projection", + }, + "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, + include_dlogp=args.include_dlogp, + ) + ) + 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, + include_dlogp: bool, +) -> 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", + ) + dlogp_report = _dlogp_report( + candidate_dtype_out, + reference_out.to(dtype), + batch=args.batch, + seq_len=seq_len, + local_hidden=local_hq * QWEN3_8B_HEAD_DIM, + seed=case_seed + 101, + device=device, + enabled=include_dlogp, + ) + 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"}, + "dlogp": dlogp_report, + } + 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, + } + + try: + from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionCPMergedState, + AttentionCPPartialState, + AttentionParallelSpec, + P2PNCCLAttentionCPCommunication, + ) + except ModuleNotFoundError as exc: + if exc.name != "rl_engine.kernels.ops.cuda.attention.cp_comm": + raise + return { + "status": "unavailable", + "reason": "PR7 CP communication module is not present in this checkout", + "required_dependency": "#279", + } + + 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 _dlogp_report( + candidate_out: torch.Tensor, + reference_out: torch.Tensor, + *, + batch: int, + seq_len: int, + local_hidden: int, + seed: int, + device: torch.device, + enabled: bool, +) -> dict[str, object]: + """Run the selected-token log-probability leg on the attention outputs. + + The benchmark intentionally uses a small deterministic synthetic projection, + but performs both logits and log-softmax in FP32. The reference is first cast + to the candidate's final-write dtype, so this gate compares the same BF16 cast + boundary and does not mislabel the expected FP32-to-BF16 write as CP drift. + """ + + if not enabled: + return { + "status": "not_requested", + "reason": "use --include-dlogp to exercise the selected-logprob chain", + } + if candidate_out.shape != reference_out.shape: + raise ValueError("candidate and reference attention outputs must have matching shapes") + generator = torch.Generator(device="cpu").manual_seed(seed) + vocab_size = 17 + weight = torch.randn( + vocab_size, + local_hidden, + generator=generator, + dtype=torch.float32, + ).to(device=device) + target_ids = ( + torch.arange(batch * seq_len, device=device, dtype=torch.long).reshape(batch, seq_len) + % vocab_size + ) + active_mask = torch.ones((batch, seq_len), device=device, dtype=torch.bool) + if seq_len > 1: + active_mask[:, 0] = False + + def project(out: torch.Tensor) -> torch.Tensor: + hidden = out.float().transpose(1, 2).reshape(batch, seq_len, local_hidden) + logits = torch.matmul(hidden, weight.transpose(0, 1)) + return selected_logprobs_reference( + logits, + target_ids, + mask=active_mask, + output_dtype=torch.float32, + ) + + candidate_logp = project(candidate_out) + reference_logp = project(reference_out) + return { + "status": "available", + "projection": "synthetic_fp32_lm_head_projection", + "reference_cast_dtype": str(reference_out.dtype).removeprefix("torch."), + "vocab_size": vocab_size, + "active_token_count": int(active_mask.sum().item()), + "drift": _drift_stats( + candidate_logp[active_mask], + reference_logp[active_mask], + ), + } + + +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..9b3d09c6 --- /dev/null +++ b/docs/design/ws2-attention-pr5-distributed-drift-benchmark.md @@ -0,0 +1,159 @@ +# 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. | + +With `--include-dlogp`, PR5 projects both Attention outputs through the same +deterministic synthetic FP32 lm_head and reports active-token selected-logprob +drift. This closes the operator-attribution leg without pretending to replace +PR4's full Qwen3 model/runtime integration. Without that flag, dlogp is recorded +as `not_requested` rather than silently omitted. + +## 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 \ + --include-dlogp \ + --output artifacts/ws2-cp-attention-drift.json +``` + +Strict GPU acceptance manifest (expected to fail until every required GPU/NCCL +case and the self-owned CUDA AG/RS operators are executable): + +```bash +python scripts/ws2_attention_gpu_acceptance.py \ + --mode manifest \ + --output artifacts/ws2-attention-acceptance-manifest.json +``` + +Strict GPU run after stacking the issue #235 implementation PRs in one checkout: + +```bash +python scripts/ws2_attention_gpu_acceptance.py \ + --mode run \ + --output artifacts/ws2-attention-gpu-acceptance.json +``` + +The orchestrator requires the Qwen3-8B `TP=2, CP=2, BF16` matrix, full and +chunked prefill, FlashInfer paged prefill/decode with disabled and fixed +Split-K, attention-domain `out/lse`, active-token `dlogp`, PR8 `dq/dk/dv`, +batch/page-layout invariance, the P2P NCCL reference, and the self-owned CUDA +AG/RS path. Missing scripts, dry-runs, requested-only Split-K provenance, +skipped collectives, or unavailable metrics fail closed. + +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`. + +The PR5 report schema is `ws2_cp_attention_drift/v2`. The strict aggregate +report schema is `ws2_attention_gpu_acceptance/v1`. diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 32b8767f..3f2c63e8 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -171,6 +171,34 @@ python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 --output artifacts/ws2-cp-attention-drift.json ``` +### WS2 CP-aware dispatch + +WS2 distributed callers use a separate contract-aware entry point, +`kernel_registry.get_attention_op(contract)`. It validates explicit TP/CP ownership, fixed +`(out, lse)` merge semantics, causal or packed-sequence offsets, and decode KV-cache identity +before selecting a backend. Legacy `get_op("attention")` behavior remains unchanged. + +Existing WS1 implementations do not yet export attention-domain LSE or implement deterministic +CP merge, so they are declared incompatible with strict WS2 requests instead of being selected as +a silent fallback. See the [WS2 CP-aware Attention contract PR][attention-contract-pr]. + +[attention-contract-pr]: https://github.com/RL-Align/RL-Kernel/pull/236 + +Split-KV is part of that contract rather than a recorded backend extra. Strict runs allow +`disabled` or a fixed logical KV chunk size, and must export the actual per-CP-owner block +boundaries, FP32 `(out, lse)` merge order, final downcast point, backend, and fallback reason. +Runtime-selected `auto` plans are diagnostic only unless both training and rollout export and +validate the same actual plan. + +The rank-aware drift benchmark can emit a CPU smoke artifact or a torchrun-friendly GPU report: + +```bash +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --json +python benchmarks/benchmark_ws2_cp_attention_drift.py --smoke --tp-world-sizes 2 \ + --cp-world-sizes 2 --kv-chunk-sizes none,1 --include-backward \ + --output artifacts/ws2-cp-attention-drift.json +``` + ## Accuracy Reference semantics (`forward_fp32`, fp32 accumulation, TF32/autocast disabled): diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index bed70208..7d79d713 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -1,60 +1,60 @@ -# 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 ( - STRICT_ATTENTION_CORE_ID, - STRICT_ATTENTION_SCHEDULE_ID, - SplitKVExecutionPlan, - SplitKVMode, - SplitKVRuntimeCoordinate, - SplitKVRuntimePlanEntry, - SplitKVRuntimePlanSet, - SplitKVSpec, -) -from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp - - -@dataclass(frozen=True) +# 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 ( + STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_SCHEDULE_ID, + SplitKVExecutionPlan, + SplitKVMode, + SplitKVRuntimeCoordinate, + SplitKVRuntimePlanEntry, + SplitKVRuntimePlanSet, + SplitKVSpec, +) +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.out.device != self.lse.device: - raise ValueError("partial attention out/lse must be on the same device") - if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: - raise ValueError("partial attention out/lse must remain FP32 before merge") - if self.block_start < 0: - raise ValueError("block_start must be non-negative") + """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.out.device != self.lse.device: + raise ValueError("partial attention out/lse must be on the same device") + if self.out.dtype is not torch.float32 or self.lse.dtype is not torch.float32: + raise ValueError("partial attention out/lse must remain FP32 before merge") + 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") @@ -130,158 +130,158 @@ def forward_with_lse( "strict_schedule": STRICT_ATTENTION_SCHEDULE_ID, }, ) - - -@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, - ) - - + + +@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. - """ - + """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" def __init__(self, *, strict_bitwise: bool = False) -> None: @@ -296,20 +296,20 @@ def __init__(self, *, strict_bitwise: bool = False) -> None: if not isinstance(strict_bitwise, bool): raise TypeError("strict_bitwise must be a bool") self.strict_bitwise = strict_bitwise - - @staticmethod + + @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.""" - + 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, + total_kv_tokens, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, backend="deterministic_cp_reference", ) @@ -328,135 +328,44 @@ def execution_provenance( kv_chunk_size=kv_chunk_size, backend="deterministic_cp_reference", ) - return { - "execution_scope": "logical_single_process_cp_reference", - "runtime_verified": False, - "input_boundary": "projected_post_qk_norm_post_rope_qkv", - "query_scope": "logical_global_query_reference", - "kv_scope": "logical_owner_local_cp_shards", - "production_cp_protocol": "ag_query_local_kv_rs_out_lse", - "communication_executed": "none", - "partial_state": "fp32_out_attention_lse", - "merge_order": "global_block_index", - "accum_dtype": "fp32", - "downcast_at": "final_write", - "requested_split_kv_policy": "disabled" if kv_chunk_size is None else "fixed", - "requested_split_kv_size": kv_chunk_size, - "actual_split_kv_plans": plans, - } - - def __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. - """ - - resolved_output_dtype = q.dtype if output_dtype is None else output_dtype - _validate_output_dtype(resolved_output_dtype) + return out + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, + query_position_offsets: Optional[torch.Tensor] = None, + key_position_offsets: Optional[torch.Tensor] = None, + cp_world_size: int = 1, + kv_chunk_size: Optional[int] = None, + output_dtype: Optional[torch.dtype] = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return ``(out, lse)`` for the CP reference path. + + ``lse`` is always fp32 and in the attention domain. ``out`` is fp32 + until the final write, then downcast to ``output_dtype``. When omitted, + ``output_dtype`` defaults to the input dtype. + """ + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) + out, lse = self._forward_impl( + q, + k, + v, + causal=causal, + scale=scale, + key_padding_mask=key_padding_mask, + query_position_offsets=query_position_offsets, + key_position_offsets=key_position_offsets, + cp_world_size=cp_world_size, + kv_chunk_size=kv_chunk_size, + ) if self.strict_bitwise: out, lse = self._forward_strict_bitwise( q, @@ -482,36 +391,36 @@ def forward_with_lse( cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size, ) - out = out.to(resolved_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.""" - + out = out.to(resolved_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, + 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, ) @@ -601,694 +510,695 @@ def _forward_strict_bitwise( out_rows.append(torch.cat(query_rows, dim=2)) lse_rows.append(torch.cat(lse_query_rows, dim=2)) return torch.cat(out_rows, dim=0), torch.cat(lse_rows, dim=0) - - def backward_reference( - self, - q: torch.Tensor, - 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") - if dout.device != q.device: - raise ValueError("dout must be on the same device as q, k, and v") - if dout.dtype != q.dtype: - raise ValueError("dout must have the same dtype as q") - q_leaf = q.detach().clone().requires_grad_(True) - k_leaf = k.detach().clone().requires_grad_(True) - v_leaf = v.detach().clone().requires_grad_(True) - - resolved_output_dtype = q.dtype if output_dtype is None else output_dtype - _validate_output_dtype(resolved_output_dtype) - out, lse = self.forward_with_lse( - q_leaf, - k_leaf, - v_leaf, - 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(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) - _validate_scale(scale) - 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) - _validate_scale(scale) - if ( - isinstance(cp_world_size, bool) - or not isinstance(cp_world_size, int) - or cp_world_size < 1 - ): - raise ValueError("cp_world_size must be >= 1") - if kv_chunk_size is not None and ( - isinstance(kv_chunk_size, bool) - or not isinstance(kv_chunk_size, int) - or kv_chunk_size < 1 - ): - raise ValueError("kv_chunk_size must be >= 1 when provided") - - 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) < 1 or k.size(1) < 1 or q.size(3) < 1: - raise ValueError("q, k, and v must have positive head counts and head dim") - if not all(torch.is_floating_point(tensor) for tensor in (q, k, v)) or any( - torch.is_complex(tensor) for tensor in (q, k, v) - ): - raise ValueError("q, k, and v must be real floating-point tensors") - if q.dtype != k.dtype or q.dtype != v.dtype: - raise ValueError("q, k, and v must have the same dtype") - if q.device != k.device or q.device != v.device: - raise ValueError("q, k, and v must be on the same device") - 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 _validate_scale(scale: Optional[float]) -> None: - if scale is None: - return - if isinstance(scale, bool) or not isinstance(scale, (int, float)): - raise ValueError("scale must be a positive finite number") - if not math.isfinite(float(scale)) or float(scale) <= 0: - raise ValueError("scale must be a positive finite number") - - -def _validate_output_dtype(output_dtype: torch.dtype) -> None: - if not isinstance(output_dtype, torch.dtype): - raise ValueError("output_dtype must be a real floating-point torch dtype") - probe = torch.empty((), dtype=output_dtype) - if not torch.is_floating_point(probe) or torch.is_complex(probe): - raise ValueError("output_dtype must be a real floating-point torch dtype") - - -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]] = [] + + 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") + if dout.device != q.device: + raise ValueError("dout must be on the same device as q, k, and v") + if dout.dtype != q.dtype: + raise ValueError("dout must have the same dtype as q") + q_leaf = q.detach().clone().requires_grad_(True) + k_leaf = k.detach().clone().requires_grad_(True) + v_leaf = v.detach().clone().requires_grad_(True) + + resolved_output_dtype = q.dtype if output_dtype is None else output_dtype + _validate_output_dtype(resolved_output_dtype) + out, lse = self.forward_with_lse( + q_leaf, + k_leaf, + v_leaf, + 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(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) + _validate_scale(scale) + 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) + _validate_scale(scale) + if ( + isinstance(cp_world_size, bool) + or not isinstance(cp_world_size, int) + or cp_world_size < 1 + ): + raise ValueError("cp_world_size must be >= 1") + if kv_chunk_size is not None and ( + isinstance(kv_chunk_size, bool) + or not isinstance(kv_chunk_size, int) + or kv_chunk_size < 1 + ): + raise ValueError("kv_chunk_size must be >= 1 when provided") + + 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) < 1 or k.size(1) < 1 or q.size(3) < 1: + raise ValueError("q, k, and v must have positive head counts and head dim") + if not all(torch.is_floating_point(tensor) for tensor in (q, k, v)) or any( + torch.is_complex(tensor) for tensor in (q, k, v) + ): + raise ValueError("q, k, and v must be real floating-point tensors") + if q.dtype != k.dtype or q.dtype != v.dtype: + raise ValueError("q, k, and v must have the same dtype") + if q.device != k.device or q.device != v.device: + raise ValueError("q, k, and v must be on the same device") + 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 _validate_scale(scale: Optional[float]) -> None: + if scale is None: + return + if isinstance(scale, bool) or not isinstance(scale, (int, float)): + raise ValueError("scale must be a positive finite number") + if not math.isfinite(float(scale)) or float(scale) <= 0: + raise ValueError("scale must be a positive finite number") + + +def _validate_output_dtype(output_dtype: torch.dtype) -> None: + if not isinstance(output_dtype, torch.dtype): + raise ValueError("output_dtype must be a real floating-point torch dtype") + probe = torch.empty((), dtype=output_dtype) + if not torch.is_floating_point(probe) or torch.is_complex(probe): + raise ValueError("output_dtype must be a real floating-point torch dtype") + + +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]] = [] + 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 boundaries: tuple[tuple[int, int], ...] 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): + 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] = [] + boundaries: tuple[tuple[int, int], ...] + 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): boundaries: tuple[tuple[int, int], ...] 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", + 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", @@ -1298,7 +1208,7 @@ def build_reference_split_kv_runtime_plan_set( "GradientDriftStats", "STRICT_ATTENTION_CORE_ID", "STRICT_ATTENTION_SCHEDULE_ID", - "compare_cp_attention_backward", - "merge_attention_partial_states", - "split_kv_execution_plan_provenance", -] + "compare_cp_attention_backward", + "merge_attention_partial_states", + "split_kv_execution_plan_provenance", +] diff --git a/scripts/ws2_attention_gpu_acceptance.py b/scripts/ws2_attention_gpu_acceptance.py new file mode 100644 index 00000000..18043b47 --- /dev/null +++ b/scripts/ws2_attention_gpu_acceptance.py @@ -0,0 +1,1314 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Strict issue #235 WS2 Attention GPU acceptance orchestrator. + +This runner combines reports produced by the existing PR branches. A required +case that is missing, skipped, dry-run only, or lacks actual runtime provenance +fails closed. It therefore separates a useful local report from a GPU/NCCL +acceptance artifact that is eligible to close issue #235. +""" + +from __future__ import annotations + +import argparse +import datetime as _datetime +import json +import math +import os +import platform +import shlex +import subprocess +import sys +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from typing import Any, Callable, Mapping, Sequence, cast + +REPO_ROOT = Path(__file__).resolve().parents[1] +SCHEMA_VERSION = "ws2_attention_gpu_acceptance/v1" +STRICT_ATTENTION_CORE_ID = "rlkernel.attention.deterministic_core.v1" +STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks" +DEFAULT_IMAGE = "ghcr.io/rl-align/rl-kernel/rl-kernel-ci:cuda" +DEFAULT_PR7_OUT_ATOL = 1.0e-2 +DEFAULT_PR7_LSE_ATOL = 2.0e-3 +DEFAULT_PR7_DLOGP_ATOL = 2.0e-3 + + +@dataclass(frozen=True) +class AcceptanceCase: + name: str + command: tuple[str, ...] | None + required: bool = True + report_path: Path | None = None + validator: Callable[[Mapping[str, Any]], list[str]] | None = None + unavailable_reason: str | None = None + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--mode", + choices=["manifest", "run"], + default="manifest", + help="manifest records the matrix without executing GPU commands", + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--torchrun", default="torchrun") + parser.add_argument("--image", default=DEFAULT_IMAGE) + parser.add_argument("--head-sha", default=os.environ.get("GITHUB_SHA")) + parser.add_argument( + "--megatron-te-script", + type=Path, + help="Megatron Bridge teacher script used for the native TE CP comparison", + ) + parser.add_argument("--megatron-model", type=Path) + parser.add_argument("--megatron-token-artifact", type=Path) + parser.add_argument("--megatron-python", default=sys.executable) + parser.add_argument("--timeout-seconds", type=int, default=1800) + parser.add_argument( + "--collective-world-size", + type=int, + choices=(2, 4, 8), + default=8, + help="rank count for the self-owned AG/RS/AllReduce probe", + ) + parser.add_argument("--out-atol", type=float, default=2.0e-4) + parser.add_argument("--lse-atol", type=float, default=2.0e-4) + parser.add_argument("--pr7-out-atol", type=float, default=DEFAULT_PR7_OUT_ATOL) + parser.add_argument("--pr7-lse-atol", type=float, default=DEFAULT_PR7_LSE_ATOL) + parser.add_argument("--pr7-dlogp-atol", type=float, default=DEFAULT_PR7_DLOGP_ATOL) + # The synthetic dlogp leg consumes the final BF16 Attention write. Use + # the shared WS1 logprob/BF16 tolerance instead of an FP32-only threshold. + parser.add_argument("--dlogp-atol", type=float, default=5.0e-2) + parser.add_argument("--grad-atol", type=float, default=5.0e-2) + return parser.parse_args(argv) + + +def build_acceptance_cases(args: argparse.Namespace) -> tuple[AcceptanceCase, ...]: + artifact_dir = args.output.resolve().parent + pr5_report = artifact_dir / "ws2-pr5-forward-backward.json" + pr7_reports = { + name: artifact_dir / f"ws2-pr7-{name}.json" + for name in ( + "decode-disabled", + "decode-fixed", + "prefill-disabled", + "prefill-fixed", + ) + } + python = str(args.python) + torchrun = str(args.torchrun) + pr7_script = REPO_ROOT / "scripts" / "ws2_pr7_flashinfer_attention_check.py" + pr7_available = pr7_script.is_file() + pr7_unavailable = None if pr7_available else "PR7 validation script is absent; integrate #279" + p2p_script = REPO_ROOT / "scripts" / "ws2_p2p_nccl_attention_reference_check.py" + p2p_available = p2p_script.is_file() + p2p_unavailable = ( + None + if p2p_available + else "three-stage Attention communication check is absent; integrate #279" + ) + collective_script = REPO_ROOT / "scripts" / "ws2_deterministic_collective_attention_check.py" + collective_report = artifact_dir / "ws2-self-owned-attention-collectives.json" + collective_available = collective_script.is_file() + collective_unavailable = ( + None if collective_available else "self-owned deterministic collective check is absent" + ) + te_compare_script = REPO_ROOT / "scripts" / "ws2_megatron_te_cp_compare.py" + te_inputs = ( + args.megatron_te_script, + args.megatron_model, + args.megatron_token_artifact, + ) + te_available = te_compare_script.is_file() and all( + path is not None and path.exists() for path in te_inputs + ) + te_unavailable = ( + None + if te_available + else "native Megatron/TE comparison requires --megatron-te-script, " + "--megatron-model, and --megatron-token-artifact" + ) + te_report = artifact_dir / "ws2-megatron-te-cp-compare.json" + + cases: list[AcceptanceCase] = [ + AcceptanceCase( + name="pr5_cp_forward_backward_dlogp", + command=( + python, + str(REPO_ROOT / "benchmarks" / "benchmark_ws2_cp_attention_drift.py"), + "--device", + "cuda", + "--tp-world-sizes", + "2", + "--cp-world-sizes", + "2", + "--kv-chunk-sizes", + "none,4", + "--include-backward", + "--include-dlogp", + "--output", + str(pr5_report), + ), + report_path=pr5_report, + validator=lambda report: validate_pr5_report(report, args), + ), + AcceptanceCase( + name="native_te_kv_ring_cp_compare", + # TE's native KV ring is a diagnostic/performance baseline. It + # currently exposes CP-dependent drift and must not gate the + # self-owned AG/RS acceptance path. + required=False, + command=( + ( + python, + str(te_compare_script), + "--teacher-script", + str(args.megatron_te_script), + "--model", + str(args.megatron_model), + "--token-artifact", + str(args.megatron_token_artifact), + "--output-dir", + str(artifact_dir / "megatron-te-cp-runs"), + "--output", + str(te_report), + "--python", + str(args.megatron_python), + "--cp-comm-type", + "p2p", + ) + if te_available + else None + ), + report_path=te_report, + validator=lambda report: validate_native_te_report(report, args), + unavailable_reason=te_unavailable, + ), + ] + for transport, prefix in ( + ("p2p_nccl_reference", "p2p_nccl_reference"), + ("cuda_ag_rs", "custom_cuda_ag_rs"), + ): + for world_size, suffix in ((2, ""), (4, "_tp2_cp2"), (8, "_tp2_cp2_replica2")): + name = f"{prefix}{suffix}" + report_path = artifact_dir / f"ws2-{name}.json" + strict_core_expected = transport == "cuda_ag_rs" + command = [ + torchrun, + "--standalone", + f"--nproc-per-node={world_size}", + str(p2p_script), + "--transport", + transport, + "--repeats", + "3", + "--atol", + str(args.out_atol), + "--final-write-atol", + str(max(args.out_atol * 100.0, 2.0e-2)), + "--output", + str(report_path), + ] + if strict_core_expected: + command.append("--strict-shared-core") + cases.append( + AcceptanceCase( + name=name, + command=tuple(command) if p2p_available else None, + report_path=report_path, + validator=partial( + validate_p2p_report, + expected_transport=transport, + expected_world_size=world_size, + expected_strict_core=strict_core_expected, + ), + unavailable_reason=p2p_unavailable, + ) + ) + for name, mode, query_len, policy, fixed_size in ( + ("decode-disabled", "decode", 1, "disabled", None), + ("decode-fixed", "decode", 1, "fixed", 4), + ("prefill-disabled", "prefill", 4, "disabled", None), + ("prefill-fixed", "prefill", 4, "fixed", 4), + ): + command = [ + python, + str(pr7_script), + "--no-dry-run", + "--device", + "cuda", + "--mode", + mode, + "--query-len", + str(query_len), + "--split-kv-policy", + policy, + "--output", + str(pr7_reports[name]), + ] + strict_expected = policy == "disabled" + if strict_expected: + command.append("--strict") + if fixed_size is not None: + command.extend(("--fixed-split-size", str(fixed_size))) + + def pr7_validator( + report: Mapping[str, Any], + expected_policy: str = policy, + strict: bool = strict_expected, + ) -> list[str]: + return validate_pr7_report( + report, + args, + expected_policy=expected_policy, + strict_expected=strict, + ) + + cases.append( + AcceptanceCase( + name=f"pr7_flashinfer_{name.replace('-', '_')}", + command=tuple(command) if pr7_available else None, + required=strict_expected, + report_path=pr7_reports[name], + validator=pr7_validator, + unavailable_reason=pr7_unavailable, + ) + ) + cases.append( + AcceptanceCase( + name="custom_cuda_allreduce", + required=False, + command=( + ( + torchrun, + "--standalone", + f"--nproc-per-node={args.collective_world_size}", + str(collective_script), + "--output", + str(collective_report), + ) + if collective_available + else None + ), + report_path=collective_report, + validator=lambda report: validate_collective_report( + report, args, operation="allreduce" + ), + unavailable_reason=collective_unavailable, + ) + ) + return tuple(cases) + + +def run_acceptance( + args: argparse.Namespace, + *, + runner: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run, +) -> dict[str, Any]: + if args.timeout_seconds < 1: + raise ValueError("timeout_seconds must be positive") + rows: list[dict[str, Any]] = [] + for case in build_acceptance_cases(args): + rows.append(_run_case(case, args, runner=runner)) + failed_required = [row["name"] for row in rows if row["required"] and not row["passed"]] + return { + "schema_version": SCHEMA_VERSION, + "issue": 235, + "created_at_utc": _datetime.datetime.now(_datetime.UTC).isoformat(), + "mode": args.mode, + "status": "passed" if not failed_required else "failed", + "passed": not failed_required, + "failed_required_cases": failed_required, + "runtime": { + "python": sys.version.split()[0], + "platform": platform.platform(), + "image": args.image, + "head_sha": args.head_sha, + "command": " ".join(shlex.quote(item) for item in sys.argv), + }, + "thresholds": { + "out_max_abs": args.out_atol, + "lse_max_abs": args.lse_atol, + "dlogp_max_abs": args.dlogp_atol, + "gradient_max_abs": args.grad_atol, + "flashinfer_out_max_abs": args.pr7_out_atol, + "flashinfer_lse_max_abs": args.pr7_lse_atol, + "flashinfer_dlogp_max_abs": args.pr7_dlogp_atol, + }, + "required_matrix": { + "topology": "Qwen3-8B TP=2 CP=2 BF16", + "attention_modes": ["prefill", "chunked_prefill", "paged_prefill", "decode"], + "split_kv": ["disabled", "fixed", "auto_diagnostic_only"], + "outputs": ["out", "attention_lse", "active_token_dlogp", "dq", "dk", "dv"], + "invariance": [ + "batch_composition", + "query_position", + "physical_page_order", + "prefix_cache_identity", + "global_block_merge_order", + ], + "strict_schedule": STRICT_ATTENTION_SCHEDULE_ID, + "communication": [ + "p2p_nccl_reference", + "self_owned_cuda_ag_rs", + "self_owned_cuda_allreduce", + ], + }, + "cases": rows, + } + + +def _run_case( + case: AcceptanceCase, + args: argparse.Namespace, + *, + runner: Callable[..., subprocess.CompletedProcess[str]], +) -> dict[str, Any]: + row: dict[str, Any] = { + "name": case.name, + "required": case.required, + "command": None if case.command is None else list(case.command), + "report_path": None if case.report_path is None else str(case.report_path), + "status": "pending", + "passed": False, + "errors": [], + } + if case.command is None: + row.update(status="unavailable") + row["errors"] = [case.unavailable_reason or "no executable implementation"] + return row + if args.mode == "manifest": + row.update(status="not_run") + row["errors"] = ["manifest mode does not execute GPU validation"] + return row + if case.report_path is not None: + case.report_path.parent.mkdir(parents=True, exist_ok=True) + case.report_path.unlink(missing_ok=True) + try: + completed = runner( + list(case.command), + cwd=REPO_ROOT, + capture_output=True, + text=True, + timeout=args.timeout_seconds, + check=False, + ) + except (OSError, subprocess.TimeoutExpired) as exc: + row.update(status="execution_error") + row["errors"] = [str(exc)] + return row + row["returncode"] = completed.returncode + row["stdout_tail"] = completed.stdout[-4000:] + row["stderr_tail"] = completed.stderr[-4000:] + if completed.returncode != 0: + if case.report_path is not None: + try: + unavailable_report = json.loads(case.report_path.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError, ValueError): + unavailable_report = None + if ( + isinstance(unavailable_report, dict) + and unavailable_report.get("status") == "not_available" + ): + row.update(status="not_available") + row["errors"] = list(unavailable_report.get("errors") or []) or [ + f"command exited with {completed.returncode}" + ] + row["report_summary"] = _report_summary(unavailable_report) + return row + row.update(status="failed") + row["errors"] = [f"command exited with {completed.returncode}"] + return row + try: + if case.report_path is not None: + report = json.loads(case.report_path.read_text(encoding="utf-8")) + else: + report = _last_json_document(completed.stdout) + except (OSError, UnicodeError, json.JSONDecodeError, ValueError) as exc: + row.update(status="invalid_report") + row["errors"] = [str(exc)] + return row + errors = [] if case.validator is None else case.validator(report) + row["errors"] = errors + row["status"] = "passed" if not errors else "failed" + row["passed"] = not errors + row["report_summary"] = _report_summary(report) + return row + + +def validate_native_te_report(report: Mapping[str, Any], args: argparse.Namespace) -> list[str]: + """Validate the delegated native Megatron/TE CP=1 vs CP=2 run.""" + + errors: list[str] = [] + if report.get("schema_version") != "ws2_megatron_te_cp_compare/v1": + errors.append("native TE report schema is invalid") + if report.get("transport") != "native_te_kv_ring": + errors.append("native TE report did not use cp_comm_type=p2p") + if report.get("status") != "passed" or report.get("passed") is not True: + report_errors = report.get("errors") + if isinstance(report_errors, list): + errors.extend(str(error) for error in report_errors) + else: + errors.append("native TE report did not provide structured errors") + requested = report.get("requested") + if not isinstance(requested, Mapping): + errors.append("native TE request metadata is missing") + requested = {} + if requested.get("cp_comm_type") != "p2p": + errors.append("native TE request did not use cp_comm_type=p2p") + if requested.get("context_parallel_sizes") != [1, 2]: + errors.append("native TE comparison must cover CP=1 and CP=2") + comparison = report.get("comparison") + if not isinstance(comparison, Mapping): + errors.append("native TE report is missing CP comparison") + comparison_hash = None + else: + if comparison.get("pass") is not True: + errors.append("native TE CP comparison did not pass") + if comparison.get("left_cp_size") != 1 or comparison.get("right_cp_size") != 2: + errors.append("native TE comparison order is not CP=1 then CP=2") + errors.extend( + _scalar_threshold_errors( + comparison.get("max_abs"), + args.dlogp_atol, + "native TE CP logprob drift", + ) + ) + comparison_hash = comparison.get("token_ids_sha256") + if not ( + isinstance(comparison_hash, str) + and len(comparison_hash) == 64 + and all(character in "0123456789abcdef" for character in comparison_hash) + ): + errors.append("native TE comparison token hash is invalid") + runs = report.get("runs") + if not isinstance(runs, list) or len(runs) != 2: + errors.append("native TE report must contain exactly two runs") + return errors + seen_cp_sizes: set[int] = set() + for index, run in enumerate(runs): + if not isinstance(run, Mapping): + errors.append(f"native TE run {index} is invalid") + continue + cp_size = run.get("cp_size") + if isinstance(cp_size, int) and not isinstance(cp_size, bool): + seen_cp_sizes.add(cp_size) + if run.get("status") != "passed": + errors.append(f"native TE CP={cp_size} run did not pass") + active_token_count = run.get("active_token_count") + if ( + isinstance(active_token_count, bool) + or not isinstance(active_token_count, int) + or active_token_count < 1 + ): + errors.append(f"native TE CP={cp_size} active-token evidence is missing") + if run.get("token_ids_sha256") != comparison_hash: + errors.append(f"native TE CP={cp_size} token hash differs from the comparison") + actual = run.get("actual") + provider = actual.get("provider") if isinstance(actual, Mapping) else None + if not isinstance(provider, Mapping): + provider = {} + if provider.get("transformer_impl") != "transformer_engine": + errors.append("native TE run did not record transformer_engine") + if provider.get("cp_comm_type") != "p2p": + errors.append("native TE run did not record cp_comm_type=p2p") + if seen_cp_sizes != {1, 2}: + errors.append("native TE runs must cover CP=1 and CP=2 exactly") + return errors + + +def validate_pr5_report(report: Mapping[str, Any], args: argparse.Namespace) -> list[str]: + errors: list[str] = [] + if report.get("schema_version") != "ws2_cp_attention_drift/v2": + errors.append("PR5 report schema is not ws2_cp_attention_drift/v2") + if report.get("issue") != 235 or report.get("pr") != 5: + errors.append("PR5 report identity is not issue #235 PR5") + runtime = report.get("runtime") + if not isinstance(runtime, dict) or not str(runtime.get("device", "")).startswith("cuda"): + errors.append("PR5 report was not produced on CUDA") + target = report.get("target") + if not isinstance(target, dict): + errors.append("PR5 target metadata is missing") + else: + if target.get("model") != "qwen3-8b" or target.get("dtype") != "bf16": + errors.append("PR5 target must be Qwen3-8B BF16") + if target.get("global_num_query_heads") != 32: + errors.append("PR5 target query-head count must be 32") + if target.get("global_num_kv_heads") != 8 or target.get("head_dim") != 128: + errors.append("PR5 target KV-head/head-dim metadata is invalid") + cases = report.get("cases") + if not isinstance(cases, list) or not cases: + errors.append("PR5 report has no cases") + return errors + expected_modes = {"prefill", "chunked_prefill"} + actual_modes = {case.get("attention_mode") for case in cases if isinstance(case, dict)} + if not expected_modes.issubset(actual_modes): + errors.append("PR5 report must contain prefill and chunked_prefill") + actual_policies = { + case.get("provenance", {}).get("requested_split_kv_policy") + for case in cases + if isinstance(case, dict) and isinstance(case.get("provenance"), dict) + } + if not {"disabled", "fixed"}.issubset(actual_policies): + errors.append("PR5 report must contain disabled and fixed Split-KV") + for case in cases: + if not isinstance(case, dict): + errors.append("PR5 case must be an object") + continue + topology = case.get("topology", {}) + if topology.get("tp_world_size") != 2 or topology.get("cp_world_size") != 2: + errors.append(f"{case.get('case_name')}: topology is not TP=2 CP=2") + provenance = case.get("provenance", {}) + if provenance.get("rope", {}).get("rope_state") != "post_rope": + errors.append(f"{case.get('case_name')}: RoPE was not composed before Attention") + requested_policy = provenance.get("requested_split_kv_policy") + requested_size = provenance.get("requested_split_kv_size") + if requested_policy == "disabled" and requested_size is not None: + errors.append(f"{case.get('case_name')}: disabled Split-KV has a split size") + if requested_policy == "fixed" and not isinstance(requested_size, int): + errors.append(f"{case.get('case_name')}: fixed Split-KV lacks an integer size") + plan_set = provenance.get("actual_split_kv_plan_set") + errors.extend( + _validate_runtime_plan_set( + plan_set, + expected_batch=_report_positive_int(target, "batch"), + expected_tp=2, + expected_cp=2, + expected_policy=requested_policy, + label=f"{case.get('case_name')}.actual_split_kv_plan_set", + ) + ) + drift = case.get("drift", {}).get("cp_merge_fp32", {}) + errors.extend( + _threshold_errors( + drift.get("out"), + args.out_atol, + f"{case.get('case_name')}.out", + ) + ) + errors.extend( + _threshold_errors( + drift.get("lse"), + args.lse_atol, + f"{case.get('case_name')}.lse", + ) + ) + dlogp = case.get("dlogp", {}) + if dlogp.get("status") != "available": + errors.append(f"{case.get('case_name')}: active-token dlogp is unavailable") + else: + errors.extend( + _threshold_errors( + dlogp.get("drift"), + args.dlogp_atol, + f"{case.get('case_name')}.dlogp", + ) + ) + backward = case.get("backward", {}) + if backward.get("status") != "available": + errors.append(f"{case.get('case_name')}: backward drift is unavailable") + else: + backward_drifts = backward.get("report", {}).get("drifts") + if not isinstance(backward_drifts, list) or not backward_drifts: + errors.append(f"{case.get('case_name')}: backward drift rows are missing") + continue + for item in backward_drifts: + if not isinstance(item, dict): + errors.append(f"{case.get('case_name')}: backward drift row is invalid") + continue + for name in ("dq", "dk", "dv"): + errors.extend( + _threshold_errors( + item.get(name), + args.grad_atol, + f"{case.get('case_name')}.{name}", + ) + ) + return errors + + +def validate_pr7_report( + report: Mapping[str, Any], + args: argparse.Namespace, + *, + expected_policy: str, + strict_expected: bool = False, +) -> list[str]: + errors: list[str] = [] + if report.get("status") != "passed" or report.get("passed") is not True: + errors.append("PR7 report is not an executed pass") + provenance = report.get("candidate_provenance") + if not isinstance(provenance, dict): + errors.append("PR7 report lacks candidate runtime provenance") + return errors + if provenance.get("arithmetic_semantics_verified") is not True: + errors.append("PR7 arithmetic semantics are not runtime-verified") + if strict_expected: + if provenance.get("strict_mode") is not True: + errors.append("PR7 strict mode was not executed") + if provenance.get("strict_core_id") != STRICT_ATTENTION_CORE_ID: + errors.append("PR7 strict core identity is invalid") + if provenance.get("strict_schedule") != STRICT_ATTENTION_SCHEDULE_ID: + errors.append("PR7 strict arithmetic schedule is invalid") + if provenance.get("native_attention_arithmetic") is not False: + errors.append("PR7 strict path entered native FlashInfer Attention arithmetic") + if provenance.get("fallback") is not False: + errors.append("PR7 strict path used a fallback") + plans = provenance.get("strict_core_row_plans") + else: + plans = provenance.get("actual_split_kv_plans") + if not isinstance(plans, list) or not plans: + errors.append("PR7 actual Split-K plans are missing") + else: + for plan in plans: + if plan.get("actual_split_kv_policy") != expected_policy: + errors.append("PR7 actual Split-K policy differs from the requested policy") + if not plan.get("actual_split_boundaries"): + errors.append("PR7 actual Split-K boundaries are missing") + if not strict_expected: + plan_set = provenance.get("actual_split_kv_plan_set") + shape = report.get("shape", {}) + errors.extend( + _validate_runtime_plan_set( + plan_set, + expected_batch=_report_positive_int(shape, "batch_size"), + expected_tp=2, + expected_cp=2, + expected_policy=expected_policy, + label="PR7 actual Split-KV plan set", + ) + ) + drift = report.get("drift", {}) + errors.extend( + _threshold_errors( + drift.get("out"), + 0.0 if strict_expected else args.pr7_out_atol, + "PR7.out", + ) + ) + errors.extend( + _threshold_errors( + drift.get("lse"), + 0.0 if strict_expected else args.pr7_lse_atol, + "PR7.lse", + ) + ) + errors.extend( + _threshold_errors( + drift.get("dlogp"), + 0.0 if strict_expected else args.pr7_dlogp_atol, + "PR7.dlogp", + ) + ) + for key in ("batch_invariant_sweep", "page_layout_invariant_sweep"): + sweep = report.get(key) + if not isinstance(sweep, dict) or sweep.get("passed") is not True: + errors.append(f"PR7 {key} did not pass") + elif strict_expected: + errors.extend(_validate_strict_invariance_sweep(sweep, label=f"PR7 {key}")) + return errors + + +def _validate_strict_invariance_sweep( + sweep: Mapping[str, Any], + *, + label: str, +) -> list[str]: + """Require explicit zero-drift evidence from strict invariance sweeps.""" + + errors: list[str] = [] + scalar_fields = ("out_max_abs", "lse_max_abs") + nested_fields = ("out", "lse") + observed = False + for field in scalar_fields: + if field in sweep: + observed = True + if sweep.get(field) != 0.0: + errors.append(f"{label} {field} is not exactly zero") + for field in nested_fields: + stats = sweep.get(field) + if isinstance(stats, Mapping): + observed = True + if stats.get("max_abs") != 0.0: + errors.append(f"{label} {field}.max_abs is not exactly zero") + if not observed and sweep.get("status") != "not_applicable": + errors.append(f"{label} lacks explicit zero-drift evidence") + return errors + + +def validate_p2p_report( + report: Mapping[str, Any], + *, + expected_transport: str = "p2p_nccl_reference", + expected_world_size: int | None = None, + expected_strict_core: bool = False, +) -> list[str]: + """Validate CP-only, TP2/CP2, and replicated TP2/CP2 communication runs.""" + + errors: list[str] = [] + if expected_transport not in {"p2p_nccl_reference", "cuda_ag_rs"}: + return [f"unsupported P2P transport expectation: {expected_transport}"] + expected_schema = ( + "ws2_p2p_nccl_attention_reference/v1" + if expected_transport == "p2p_nccl_reference" + else "ws2_cuda_ag_rs_attention/v1" + ) + if report.get("schema_version") != expected_schema: + errors.append("P2P report schema is invalid") + if report.get("transport") != expected_transport: + errors.append(f"P2P report did not use {expected_transport}") + if "nccl" not in str(report.get("backend", "")).lower(): + errors.append("P2P report backend is not NCCL") + world_size = report.get("world_size") + if not isinstance(world_size, int) or world_size not in {2, 4, 8}: + errors.append("P2P report world size must be 2, 4, or 8") + return errors + if expected_world_size is not None and world_size != expected_world_size: + errors.append(f"P2P report world size is not {expected_world_size}") + expected_tp_world_size = 1 if world_size == 2 else 2 + expected_replica_count = 2 if world_size == 8 else 1 + if report.get("tp_world_size") != expected_tp_world_size: + errors.append("P2P report TP world size is inconsistent with the rank topology") + if report.get("cp_world_size") != 2: + errors.append("P2P report CP world size is not 2") + if report.get("replica_count") != expected_replica_count: + errors.append("P2P report replica count is inconsistent with the rank topology") + if report.get("global_failure_count") != 0: + errors.append("P2P report has global rank failures") + ranks = report.get("ranks") + if not isinstance(ranks, list) or len(ranks) != world_size: + errors.append(f"P2P report must contain exactly {world_size} rank reports") + return errors + + seen_ranks: set[int] = set() + seen_coords: set[tuple[int, int, int]] = set() + query_ranges_by_group: dict[tuple[int, int], dict[int, list[int]]] = {} + manifests_by_group: dict[tuple[int, int], list[list[Any]]] = {} + for index, row in enumerate(ranks): + if not isinstance(row, dict): + errors.append(f"P2P rank {index} report is invalid") + continue + rank = row.get("rank") + tp_rank = row.get("tp_rank") + cp_rank = row.get("cp_rank") + replica_index = row.get("replica_index") + if not isinstance(rank, int): + errors.append(f"P2P row {index} lacks an integer rank") + continue + seen_ranks.add(rank) + if rank < 0 or rank >= world_size: + errors.append(f"P2P rank {index} is outside the world") + replica_rank = rank % 4 if world_size == 8 else rank + expected_replica = rank // 4 if world_size == 8 else 0 + expected_tp = 0 if world_size == 2 else replica_rank // 2 + expected_cp = replica_rank % 2 + if (replica_index, tp_rank, cp_rank) != ( + expected_replica, + expected_tp, + expected_cp, + ): + errors.append( + f"P2P rank {index} replica/TP/CP coordinates are inconsistent with rank order" + ) + if all(isinstance(value, int) for value in (replica_index, tp_rank, cp_rank)): + seen_coords.add((replica_index, tp_rank, cp_rank)) + if row.get("global_world_size") != world_size: + errors.append(f"P2P rank {index} global world size is inconsistent") + if row.get("cp_world_size") != 2 or row.get("replica_count") != expected_replica_count: + errors.append(f"P2P rank {index} CP/replica topology is inconsistent") + if row.get("global_failure_count") != 0: + errors.append(f"P2P rank {index} observed global rank failures") + if row.get("passed") is not True: + errors.append(f"P2P rank {index} did not pass") + if row.get("transport") != expected_transport: + errors.append(f"P2P rank {index} did not use {expected_transport}") + if row.get("query_ag") != expected_transport: + errors.append(f"P2P rank {index} did not execute the expected Q AllGather") + if row.get("protocol") != "ag_query_local_kv_rs_out_lse": + errors.append(f"P2P rank {index} did not execute the three-stage protocol") + strict_report = row.get("strict_shared_core") + if expected_strict_core: + errors.extend(_validate_strict_shared_core_report(strict_report, rank=index)) + if row.get("strict_protocol") != "ag_qkv_positions_shared_core_rs_out_lse": + errors.append(f"P2P rank {index} strict protocol is invalid") + elif isinstance(strict_report, Mapping) and strict_report.get("executed") is not False: + errors.append(f"P2P rank {index} unexpectedly claimed strict shared-core execution") + elif strict_report is not None and not isinstance(strict_report, Mapping): + errors.append(f"P2P rank {index} strict shared-core report is invalid") + if row.get("query_ag_max_abs") != 0.0: + errors.append(f"P2P rank {index} Q AllGather was not bitwise exact") + if row.get("dtype") != "bf16" or row.get("accum_dtype") != "fp32": + errors.append(f"P2P rank {index} arithmetic provenance is invalid") + if row.get("downcast_at") != "final_write": + errors.append(f"P2P rank {index} downcast provenance is invalid") + if row.get("final_output_dtype") != "bfloat16": + errors.append(f"P2P rank {index} final output dtype is not BF16") + if not str(row.get("device", "")).startswith("cuda"): + errors.append(f"P2P rank {index} was not executed on CUDA") + if not isinstance(row.get("repeat_count"), int) or row["repeat_count"] < 2: + errors.append(f"P2P rank {index} repeat count is insufficient") + for repeat_name in ( + "repeat_query_bitwise", + "repeat_out_bitwise", + "repeat_lse_bitwise", + "repeat_manifest_bitwise", + ): + if row.get(repeat_name) is not True: + errors.append(f"P2P rank {index} {repeat_name} did not pass") + + query_range = row.get("query_range") + if not ( + isinstance(query_range, list) + and len(query_range) == 2 + and all(isinstance(value, int) and not isinstance(value, bool) for value in query_range) + and query_range[0] < query_range[1] + ): + errors.append(f"P2P rank {index} query ownership is invalid") + elif all(isinstance(value, int) for value in (replica_index, tp_rank, cp_rank)): + query_ranges_by_group.setdefault((replica_index, tp_rank), {})[cp_rank] = query_range + + gathered_indices = row.get("gathered_block_indices") + block_manifest = row.get("expected_block_manifest") + manifest_errors, manifest_indices = _validate_p2p_block_manifest( + block_manifest, + expected_tp_rank=tp_rank if isinstance(tp_rank, int) else None, + expected_tp_world_size=expected_tp_world_size, + ) + errors.extend(f"P2P rank {index}: {error}" for error in manifest_errors) + if not ( + isinstance(gathered_indices, list) + and gathered_indices + and gathered_indices == list(range(len(gathered_indices))) + and manifest_indices == gathered_indices + ): + errors.append(f"P2P rank {index} gathered block order/coverage is invalid") + if ( + isinstance(tp_rank, int) + and isinstance(cp_rank, int) + and isinstance(manifest_indices, list) + and isinstance(block_manifest, list) + ): + local_indices = row.get("local_block_indices") + expected_local = [ + block_index + for block_index, block in enumerate(block_manifest or []) + if isinstance(block, Mapping) and block.get("owner_cp_rank") == cp_rank + ] + if local_indices != expected_local: + errors.append(f"P2P rank {index} local block ownership is invalid") + if isinstance(replica_index, int): + manifests_by_group.setdefault((replica_index, tp_rank), []).append(block_manifest) + for name in ("out_max_abs", "lse_max_abs"): + errors.extend( + _scalar_threshold_errors(row.get(name), row.get("atol"), f"P2P rank {index}.{name}") + ) + errors.extend( + _scalar_threshold_errors( + row.get("final_out_max_abs"), + row.get("final_write_atol"), + f"P2P rank {index}.final_out_max_abs", + ) + ) + + expected_ranks = set(range(world_size)) + if seen_ranks != expected_ranks: + errors.append(f"P2P report must cover ranks 0 through {world_size - 1} exactly") + expected_coords = { + (replica_index, tp_rank, cp_rank) + for replica_index in range(expected_replica_count) + for tp_rank in range(expected_tp_world_size) + for cp_rank in range(2) + } + if seen_coords != expected_coords: + errors.append("P2P report must cover the canonical replica/TP/CP coordinate grid") + + reference_ranges: dict[int, list[int]] | None = None + for replica_index in range(expected_replica_count): + for tp_rank in range(expected_tp_world_size): + group = (replica_index, tp_rank) + ranges = query_ranges_by_group.get(group, {}) + first_range = ranges.get(0) + second_range = ranges.get(1) + if ( + first_range is None + or second_range is None + or first_range[0] != 0 + or first_range[1] != second_range[0] + or second_range[1] <= second_range[0] + ): + errors.append( + "P2P group " + f"replica={replica_index}, tp={tp_rank} query ownership is not canonical" + ) + if reference_ranges is None: + reference_ranges = ranges + elif ranges != reference_ranges: + errors.append("P2P replica/TP groups have different query ownership ranges") + manifests = manifests_by_group.get(group, []) + if len(manifests) != 2: + errors.append( + f"P2P group replica={replica_index}, tp={tp_rank} lacks both CP manifests" + ) + elif manifests[0] != manifests[1]: + errors.append( + f"P2P group replica={replica_index}, tp={tp_rank} gathered different manifests" + ) + return errors + + +def _validate_strict_shared_core_report(report: Any, *, rank: int) -> list[str]: + label = f"P2P rank {rank} strict shared core" + if not isinstance(report, Mapping): + return [f"{label} report is missing"] + errors: list[str] = [] + expected = { + "executed": True, + "passed": True, + "strict_core_id": STRICT_ATTENTION_CORE_ID, + "strict_schedule": STRICT_ATTENTION_SCHEDULE_ID, + "actual_backend": "rlkernel.cuda.deterministic_attention", + "communication_backend": "self_owned_cuda_ag_rs", + "production_ready": True, + "strict_mode": True, + "native_attention_arithmetic": False, + "fallback": False, + "split_kv_policy": "disabled", + "communication_autograd": True, + "repeat_out_bitwise": True, + "repeat_lse_bitwise": True, + } + for field, value in expected.items(): + if report.get(field) != value: + errors.append(f"{label} has invalid {field}") + bitwise = report.get("bitwise") + if not isinstance(bitwise, Mapping) or any( + bitwise.get(name) is not True for name in ("out", "lse", "dq", "dk", "dv") + ): + errors.append(f"{label} Out/LSE/gradient bitwise evidence is incomplete") + max_abs = report.get("max_abs") + if not isinstance(max_abs, Mapping) or any( + max_abs.get(name) != 0.0 for name in ("out", "lse", "dq", "dk", "dv") + ): + errors.append(f"{label} Out/LSE/gradient drift is not exactly zero") + return errors + + +def validate_collective_report( + report: Mapping[str, Any], + args: argparse.Namespace, + *, + operation: str, +) -> list[str]: + """Validate executed PR310/311/312 evidence, never configured-only claims.""" + + errors: list[str] = [] + if report.get("schema_version") != "ws2_deterministic_attention_collectives/v1": + errors.append("self-owned collective report schema is invalid") + if report.get("world_size") != args.collective_world_size: + errors.append("self-owned collective report world size differs from the requested size") + if report.get("transport") != "self_owned_cuda_ag_rs": + errors.append("self-owned report did not execute the CUDA AG/RS backend") + if report.get("allreduce_transport") != "self_owned_cuda_allreduce": + errors.append("self-owned report did not execute the CUDA AllReduce backend") + if report.get("global_failure_count") != 0 or report.get("passed") is not True: + errors.append("self-owned collective report contains rank failures") + ranks = report.get("ranks") + if not isinstance(ranks, list) or len(ranks) != args.collective_world_size: + errors.append("self-owned collective report must contain every rank") + return errors + required = { + "ag_rs": ("all_gather_q", "reduce_scatter_out_lse"), + "allreduce": ("all_reduce_o_proj",), + }[operation] + for index, row in enumerate(ranks): + if not isinstance(row, Mapping): + errors.append(f"self-owned rank {index} report is invalid") + continue + if row.get("passed") is not True: + errors.append(f"self-owned rank {index} did not pass") + operations = row.get("operations") + if not isinstance(operations, Mapping): + errors.append(f"self-owned rank {index} operation evidence is missing") + continue + for name in required: + evidence = operations.get(name) + if not isinstance(evidence, Mapping) or evidence.get("passed") is not True: + errors.append(f"self-owned rank {index} {name} evidence did not pass") + return errors + + +def _validate_p2p_block_manifest( + manifest: Any, + *, + expected_tp_rank: int | None, + expected_tp_world_size: int = 1, +) -> tuple[list[str], list[int] | None]: + if not isinstance(manifest, list) or not manifest: + return ["expected block manifest is missing"], None + required = { + "global_block_index", + "kv_block_start", + "kv_block_end", + "owner_cp_rank", + "owner_tp_rank", + } + errors: list[str] = [] + indices: list[int] = [] + cursor = 0 + owners: set[int] = set() + for index, block in enumerate(manifest): + if not isinstance(block, dict) or not required.issubset(block): + errors.append(f"manifest block {index} is missing required metadata") + continue + values = {name: block[name] for name in required} + if not all( + isinstance(value, int) and not isinstance(value, bool) for value in values.values() + ): + errors.append(f"manifest block {index} metadata must contain integers") + continue + global_index = values["global_block_index"] + start = values["kv_block_start"] + end = values["kv_block_end"] + owner_cp_rank = values["owner_cp_rank"] + owner_tp_rank = values["owner_tp_rank"] + indices.append(global_index) + owners.add(owner_cp_rank) + if global_index != index: + errors.append(f"manifest block {index} has a non-canonical global index") + if start != cursor or end <= start: + errors.append(f"manifest block {index} does not preserve gap-free KV coverage") + cursor = end + if owner_cp_rank not in {0, 1} or not 0 <= owner_tp_rank < expected_tp_world_size: + errors.append(f"manifest block {index} owner is outside the TP-local CP=2 group") + if expected_tp_rank is not None and owner_tp_rank != expected_tp_rank: + errors.append(f"manifest block {index} owner TP rank does not match the report") + if owners != {0, 1}: + errors.append("manifest does not assign KV blocks to both CP ranks") + return errors, indices + + +def _threshold_errors(stats: Any, threshold: float, label: str) -> list[str]: + if not isinstance(stats, dict) or "max_abs" not in stats: + return [f"{label} drift is missing"] + try: + value = float(stats["max_abs"]) + except (TypeError, ValueError): + return [f"{label} max_abs is not numeric"] + if not math.isfinite(value) or value < 0: + return [f"{label} max_abs must be finite and non-negative"] + return [] if value <= threshold else [f"{label} max_abs={value} exceeds {threshold}"] + + +def _scalar_threshold_errors(value: Any, threshold: Any, label: str) -> list[str]: + try: + numeric_value = float(value) + numeric_threshold = float(threshold) + except (TypeError, ValueError): + return [f"{label} or its threshold is not numeric"] + if not math.isfinite(numeric_value) or numeric_value < 0: + return [f"{label} must be finite and non-negative"] + if not math.isfinite(numeric_threshold) or numeric_threshold < 0: + return [f"{label} threshold must be finite and non-negative"] + if numeric_value > numeric_threshold: + return [f"{label}={numeric_value} exceeds {numeric_threshold}"] + return [] + + +def _validate_runtime_plan_set( + plan_set: Any, + *, + expected_batch: int, + expected_tp: int, + expected_cp: int, + expected_policy: Any, + label: str, +) -> list[str]: + if expected_batch < 1: + return [f"{label} expected batch size is invalid"] + if not isinstance(plan_set, dict): + return [f"{label} is missing"] + errors: list[str] = [] + if plan_set.get("coverage") != "complete_batch_tp_cp_owner_cartesian_product": + errors.append(f"{label} coverage marker is invalid") + topology = ( + plan_set.get("batch_size"), + plan_set.get("tp_world_size"), + plan_set.get("cp_world_size"), + ) + expected_topology = (expected_batch, expected_tp, expected_cp) + if topology != expected_topology: + errors.append(f"{label} topology {topology} does not match {expected_topology}") + totals = plan_set.get("total_kv_tokens") + if not ( + isinstance(totals, list) + and len(totals) == expected_batch + and all( + isinstance(total, int) and not isinstance(total, bool) and total > 0 for total in totals + ) + ): + errors.append(f"{label} total_kv_tokens is invalid") + return errors + entries = plan_set.get("entries") + expected_coordinates = { + (batch_index, tp_rank, cp_rank, owner_cp_rank) + for batch_index in range(expected_batch) + for tp_rank in range(expected_tp) + for cp_rank in range(expected_cp) + for owner_cp_rank in range(expected_cp) + } + if not isinstance(entries, list): + errors.append(f"{label} entries are missing") + return errors + coordinates: list[tuple[Any, Any, Any, Any]] = [] + owner_ranges: dict[tuple[int, int, int], tuple[int, int]] = {} + for index, entry in enumerate(entries): + entry_label = f"{label}.entries[{index}]" + if not isinstance(entry, dict): + errors.append(f"{entry_label} is not an object") + continue + coordinate_values = tuple( + entry.get(key) for key in ("batch_index", "tp_rank", "cp_rank", "owner_cp_rank") + ) + if not all( + isinstance(value, int) and not isinstance(value, bool) for value in coordinate_values + ): + errors.append(f"{entry_label} coordinate must contain integers") + continue + coordinate: tuple[int, int, int, int] = ( + cast(int, coordinate_values[0]), + cast(int, coordinate_values[1]), + cast(int, coordinate_values[2]), + cast(int, coordinate_values[3]), + ) + coordinates.append(coordinate) + if coordinate not in expected_coordinates: + errors.append(f"{entry_label} coordinate is out of range") + continue + batch_index, tp_rank, _, owner_cp_rank = coordinate + expected_range = entry.get("expected_kv_range") + if not ( + isinstance(expected_range, list) + and len(expected_range) == 2 + and all( + isinstance(value, int) and not isinstance(value, bool) for value in expected_range + ) + and 0 <= expected_range[0] < expected_range[1] <= totals[batch_index] + ): + errors.append(f"{entry_label} expected_kv_range is invalid") + continue + range_key = (batch_index, tp_rank, owner_cp_rank) + range_tuple = (expected_range[0], expected_range[1]) + previous_range = owner_ranges.setdefault(range_key, range_tuple) + if previous_range != range_tuple: + errors.append(f"{entry_label} owner range differs across CP consumers") + if entry.get("requested_split_kv_policy") != expected_policy: + errors.append(f"{entry_label} requested Split-KV policy is wrong") + if entry.get("actual_split_kv_policy") != expected_policy: + errors.append(f"{entry_label} actual Split-KV policy is wrong") + if entry.get("split_kv_merge_order") != "global_block_index": + errors.append(f"{entry_label} merge order is not global_block_index") + if entry.get("split_kv_accum_dtype") != "fp32": + errors.append(f"{entry_label} accumulation dtype is not fp32") + if entry.get("split_kv_downcast_at") != "final_write": + errors.append(f"{entry_label} downcast point is not final_write") + if entry.get("split_kv_fallback") is not False: + errors.append(f"{entry_label} used a fallback") + if not isinstance(entry.get("split_kv_plan_source"), str): + errors.append(f"{entry_label} runtime plan source is missing") + boundaries = entry.get("actual_split_boundaries") + if not isinstance(boundaries, list) or not boundaries: + errors.append(f"{entry_label} actual split boundaries are missing") + continue + cursor = expected_range[0] + valid_boundaries = True + for boundary in boundaries: + if not ( + isinstance(boundary, list) + and len(boundary) == 2 + and all( + isinstance(value, int) and not isinstance(value, bool) for value in boundary + ) + and boundary[0] == cursor + and boundary[0] < boundary[1] <= expected_range[1] + ): + valid_boundaries = False + break + cursor = boundary[1] + if not valid_boundaries or cursor != expected_range[1]: + errors.append(f"{entry_label} boundaries do not cover the owner range exactly") + if entry.get("actual_split_kv_count") != len(boundaries): + errors.append(f"{entry_label} actual split count is inconsistent") + if len(coordinates) != len(set(coordinates)): + errors.append(f"{label} contains duplicate coordinates") + actual_coordinates = set(coordinates) + if actual_coordinates != expected_coordinates: + errors.append(f"{label} coordinate coverage is incomplete") + for batch_index in range(expected_batch): + for tp_rank in range(expected_tp): + cursor = 0 + for owner_cp_rank in range(expected_cp): + owner_range = owner_ranges.get((batch_index, tp_rank, owner_cp_rank)) + if owner_range is None or owner_range[0] != cursor: + errors.append( + f"{label} owner ranges are not contiguous for " + f"batch={batch_index}, tp={tp_rank}" + ) + break + cursor = owner_range[1] + if cursor != totals[batch_index]: + errors.append( + f"{label} owner ranges do not cover total KV for " + f"batch={batch_index}, tp={tp_rank}" + ) + return errors + + +def _report_positive_int(container: Any, key: str) -> int: + if not isinstance(container, dict): + return 0 + value = container.get(key) + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + return 0 + return value + + +def _last_json_document(stdout: str) -> Mapping[str, Any]: + decoder = json.JSONDecoder() + for index, character in enumerate(stdout): + if character != "{": + continue + try: + value, end = decoder.raw_decode(stdout[index:]) + except json.JSONDecodeError: + continue + if stdout[index + end :].strip() or not isinstance(value, dict): + continue + return value + raise ValueError("command stdout does not end with a JSON object") + + +def _report_summary(report: Mapping[str, Any]) -> dict[str, Any]: + return { + key: report.get(key) + for key in ("schema_version", "status", "passed", "issue", "pr", "mode") + if key in report + } + + +def write_report(report: Mapping[str, Any], 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_acceptance(args) + write_report(report, args.output) + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 if report["passed"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/ws2_deterministic_collective_attention_check.py b/scripts/ws2_deterministic_collective_attention_check.py new file mode 100644 index 00000000..ad48284f --- /dev/null +++ b/scripts/ws2_deterministic_collective_attention_check.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Exercise the self-owned deterministic AG/RS/AllReduce on an Attention shape. + +This is intentionally a small transport probe, not a replacement for the CP +attention reference. It checks the exact communication primitives required by +the table: AG for Q/SP, FP32 RS for `(Out, LSE)`, and AllReduce for the o_proj +partial sum. Run under ``torchrun`` on one host with 2, 4, or 8 ranks. +""" + +from __future__ import annotations + +import argparse +import json +import os +from pathlib import Path +from typing import Sequence + +import torch +import torch.distributed as dist + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--rows", type=int, default=4) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + if args.rows < 1: + raise ValueError("rows must be positive") + if not torch.cuda.is_available(): + raise RuntimeError("self-owned CUDA collective check requires CUDA") + dist.init_process_group("nccl", init_method="env://") + rank = dist.get_rank() + world_size = dist.get_world_size() + if world_size not in (2, 4, 8): + raise RuntimeError(f"expected 2, 4, or 8 ranks, got {world_size}") + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + try: + from rl_engine.distributed import DeterministicCollective + + with DeterministicCollective(device=device, max_size_bytes=16 * 1024 * 1024) as collective: + operations = { + "all_gather_q": _check_all_gather(collective, rank, world_size, args.rows, device), + "reduce_scatter_out_lse": _check_reduce_scatter( + collective, rank, world_size, args.rows, device + ), + "all_reduce_o_proj": _check_all_reduce( + collective, rank, world_size, args.rows, device + ), + } + passed = all(bool(item["passed"]) for item in operations.values()) + result = { + "rank": rank, + "device": str(device), + "world_size": world_size, + "transport": "self_owned_cuda_ag_rs", + "allreduce_transport": "self_owned_cuda_allreduce", + "accumulation_dtype": "fp32", + "downcast_at": "final_write", + "operations": operations, + "passed": passed, + } + failures = torch.tensor([0 if passed else 1], dtype=torch.int32, device=device) + dist.all_reduce(failures, op=dist.ReduceOp.SUM) + reports: list[dict[str, object] | None] = [None] * world_size + dist.all_gather_object(reports, result) + if rank == 0: + payload = { + "schema_version": "ws2_deterministic_attention_collectives/v1", + "world_size": world_size, + "transport": "self_owned_cuda_ag_rs", + "allreduce_transport": "self_owned_cuda_allreduce", + "global_failure_count": int(failures.item()), + "ranks": reports, + "passed": int(failures.item()) == 0, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8") + print(json.dumps(payload, indent=2, sort_keys=True)) + return 0 if int(failures.item()) == 0 else 1 + finally: + dist.destroy_process_group() + + +def _check_all_gather(collective, rank: int, world_size: int, rows: int, device: torch.device): + local = ( + torch.arange(rows * 8, device=device, dtype=torch.bfloat16).reshape(rows, 8) + rank * 100 + ) + expected = torch.cat( + [ + torch.arange(rows * 8, device=device, dtype=torch.bfloat16).reshape(rows, 8) + + peer_rank * 100 + for peer_rank in range(world_size) + ], + dim=0, + ) + out = collective.all_gather(local) + repeat = collective.all_gather(local) + return { + "dtype": "bf16", + "passed": bool(torch.equal(out, expected) and torch.equal(out, repeat)), + } + + +def _check_reduce_scatter(collective, rank: int, world_size: int, rows: int, device: torch.device): + local = torch.full((rows * world_size, 9), float(rank + 1), device=device, dtype=torch.float32) + expected = torch.full( + (rows, 9), + float(sum(range(1, world_size + 1))), + device=device, + dtype=torch.float32, + ) + out = collective.reduce_scatter(local) + repeat = collective.reduce_scatter(local) + return { + "dtype": "fp32", + "passed": bool(torch.equal(out, expected) and torch.equal(out, repeat)), + } + + +def _check_all_reduce(collective, rank: int, world_size: int, rows: int, device: torch.device): + local = torch.full((rows, 11), float(rank + 1), device=device, dtype=torch.float32) + expected = torch.full( + (rows, 11), + float(sum(range(1, world_size + 1))), + device=device, + dtype=torch.float32, + ) + out = collective.all_reduce(local) + repeat = collective.all_reduce(local) + return { + "dtype": "fp32", + "passed": bool(torch.equal(out, expected) and torch.equal(out, repeat)), + } + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/ws2_megatron_te_cp_compare.py b/scripts/ws2_megatron_te_cp_compare.py new file mode 100644 index 00000000..f3dd38ae --- /dev/null +++ b/scripts/ws2_megatron_te_cp_compare.py @@ -0,0 +1,334 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Run the native Megatron/Transformer Engine CP KV-ring comparison. + +This runner deliberately delegates model execution to Megatron Bridge. RL-Kernel +does not reimplement TE's KV ring here. The delegated teacher script must set +``transformer_impl=transformer_engine`` and receives ``cp_comm_type`` from this +runner. A non-zero child exit or a missing runtime provenance is a failure. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +import os +import subprocess +import sys +from pathlib import Path +from typing import Any, Mapping, Sequence + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--teacher-script", type=Path, required=True) + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--token-artifact", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--python", default=sys.executable) + parser.add_argument("--cp-comm-type", choices=("p2p", "all_gather"), default="p2p") + parser.add_argument("--tensor-parallel-size", type=int, default=2) + parser.add_argument("--cp-sizes", default="1,2") + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--atol", type=float, default=5.0e-2) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> int: + args = parse_args(argv) + cp_sizes = _parse_cp_sizes(args.cp_sizes) + _validate_args(args, cp_sizes) + args.output_dir.mkdir(parents=True, exist_ok=True) + runs: list[dict[str, Any]] = [] + for cp_size in cp_sizes: + result = _run_teacher(args, cp_size=cp_size) + runs.append(result) + + comparison = None + errors: list[str] = [] + if all(run.get("status") == "passed" for run in runs): + try: + comparison = _compare_runs(runs, atol=args.atol) + except (RuntimeError, ValueError) as exc: + comparison = {"pass": False, "identity_error": str(exc), "atol": args.atol} + errors.append(f"CP token identity validation failed: {exc}") + else: + if not comparison["pass"]: + errors.append( + f"CP=1 vs CP=2 native TE drift exceeds atol={args.atol}: " + f"max_abs={comparison['max_abs']}" + ) + else: + errors.extend( + f"CP={run['cp_size']} teacher run failed: {run.get('error', 'unknown error')}" + for run in runs + if run.get("status") != "passed" + ) + + report = { + "schema_version": "ws2_megatron_te_cp_compare/v1", + "status": "passed" if not errors else "failed", + "passed": not errors, + "transport": ( + "native_te_kv_ring" if args.cp_comm_type == "p2p" else "native_te_kv_all_gather" + ), + "requested": { + "cp_comm_type": args.cp_comm_type, + "tensor_parallel_size": args.tensor_parallel_size, + "context_parallel_sizes": cp_sizes, + "dtype": "bfloat16", + "seed": args.seed, + }, + "runs": runs, + "comparison": comparison, + "errors": errors, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 if not errors else 1 + + +def _run_teacher(args: argparse.Namespace, *, cp_size: int) -> dict[str, Any]: + stem = f"megatron_tp{args.tensor_parallel_size}_cp{cp_size}_{args.cp_comm_type}" + output = args.output_dir / f"{stem}.json" + log = args.output_dir / f"{stem}.log" + command = [ + args.python, + "-m", + "torch.distributed.run", + "--standalone", + f"--nproc-per-node={args.tensor_parallel_size * cp_size}", + str(args.teacher_script), + "--model", + str(args.model), + "--token-artifact", + str(args.token_artifact), + "--output", + str(output), + "--tensor-parallel-size", + str(args.tensor_parallel_size), + "--context-parallel-size", + str(cp_size), + "--cp-comm-type", + args.cp_comm_type, + "--seed", + str(args.seed), + ] + env = os.environ.copy() + env.setdefault("TOKENIZERS_PARALLELISM", "false") + env.setdefault("OMP_NUM_THREADS", "1") + with log.open("w", encoding="utf-8") as stream: + completed = subprocess.run( + command, + cwd=args.teacher_script.resolve().parents[2], + env=env, + stdout=stream, + stderr=subprocess.STDOUT, + check=False, + ) + row: dict[str, Any] = { + "cp_size": cp_size, + "world_size": args.tensor_parallel_size * cp_size, + "command": command, + "output": str(output), + "log": str(log), + "returncode": completed.returncode, + } + if completed.returncode != 0: + row.update({"status": "failed", "error": f"returncode={completed.returncode}"}) + return row + if not output.is_file(): + row.update({"status": "failed", "error": "teacher output JSON is missing"}) + return row + try: + artifact = json.loads(output.read_text(encoding="utf-8")) + except (OSError, UnicodeError, json.JSONDecodeError, ValueError) as exc: + row.update({"status": "failed", "error": f"invalid teacher output JSON: {exc}"}) + return row + if not isinstance(artifact, dict): + row.update({"status": "failed", "error": "teacher output JSON is not an object"}) + return row + if artifact.get("schema") != "ws2.megatron_teacher_logprobs.v1": + row.update({"status": "failed", "error": "teacher output schema is invalid"}) + return row + actual = artifact.get("actual") + provider = actual.get("provider") if isinstance(actual, Mapping) else None + if not isinstance(provider, Mapping): + provider = {} + active_token_logprobs = artifact.get("active_token_logprobs") + required = { + "context_parallel_size": cp_size, + "tensor_model_parallel_size": args.tensor_parallel_size, + "cp_comm_type": args.cp_comm_type, + "transformer_impl": "transformer_engine", + } + mismatches = { + key: {"expected": expected, "actual": provider.get(key)} + for key, expected in required.items() + if provider.get(key) != expected + } + row.update( + { + "status": "passed" if not mismatches else "failed", + "actual": actual, + "mismatches": mismatches, + "active_token_count": ( + len(active_token_logprobs) if isinstance(active_token_logprobs, list) else 0 + ), + } + ) + try: + _, _, _, token_ids_sha256 = _token_identity(artifact, label=f"CP={cp_size}") + except (RuntimeError, ValueError) as exc: + row.update(status="failed", error=str(exc)) + else: + row["token_ids_sha256"] = token_ids_sha256 + if mismatches: + row["error"] = "runtime provider provenance does not match native TE request" + return row + + +def _compare_runs(runs: list[dict[str, Any]], *, atol: float) -> dict[str, Any]: + if len(runs) != 2: + raise ValueError("CP comparison requires exactly two runs") + if [run.get("cp_size") for run in runs] != [1, 2]: + raise ValueError("CP comparison requires runs ordered as CP=1 then CP=2") + artifacts = [json.loads(Path(run["output"]).read_text(encoding="utf-8")) for run in runs] + left = artifacts[0]["active_token_logprobs"] + right = artifacts[1]["active_token_logprobs"] + left_positions, left_token_ids, left_all_token_ids, left_hash = _token_identity( + artifacts[0], label="left CP run" + ) + right_positions, right_token_ids, right_all_token_ids, right_hash = _token_identity( + artifacts[1], label="right CP run" + ) + if len(left) != len(right): + raise RuntimeError( + f"CP runs produced different active-token counts: left={len(left)}, right={len(right)}" + ) + if left_positions != right_positions: + raise RuntimeError("CP runs produced different active-token positions") + if left_all_token_ids != right_all_token_ids: + raise RuntimeError("CP runs used different complete token ID sequences") + if left_token_ids != right_token_ids: + mismatch = next( + index + for index, (left_id, right_id) in enumerate( + zip(left_token_ids, right_token_ids, strict=True) + ) + if left_id != right_id + ) + raise RuntimeError( + "CP runs produced different active-token IDs at index " + f"{mismatch}: left={left_token_ids[mismatch]}, right={right_token_ids[mismatch]}" + ) + if left_hash != right_hash: + raise RuntimeError("CP runs produced different token_ids_sha256 values") + diffs = [ + abs(float(left_row["logprob"]) - float(right_row["logprob"])) + for left_row, right_row in zip(left, right, strict=True) + ] + worst_index = max(range(len(diffs)), key=diffs.__getitem__) + return { + "left_cp_size": runs[0]["cp_size"], + "right_cp_size": runs[1]["cp_size"], + "active_token_count": len(diffs), + "token_ids_sha256": left_hash, + "max_abs": max(diffs, default=0.0), + "mean_abs": sum(diffs) / len(diffs) if diffs else 0.0, + "worst": { + "position": left[worst_index]["position"] if diffs else None, + "token_id": left[worst_index]["token_id"] if diffs else None, + "abs_diff": diffs[worst_index] if diffs else 0.0, + }, + "atol": atol, + "pass": max(diffs, default=0.0) <= atol, + } + + +def _token_identity( + artifact: dict[str, Any], + *, + label: str, +) -> tuple[list[int], list[int], list[int], str]: + entries = artifact.get("active_token_logprobs") + if not isinstance(entries, list) or not entries: + raise RuntimeError(f"{label} active_token_logprobs is missing or empty") + positions: list[int] = [] + token_ids: list[int] = [] + for index, entry in enumerate(entries): + if not isinstance(entry, dict): + raise RuntimeError(f"{label} token entry {index} is not an object") + position = entry.get("position") + token_id = entry.get("token_id") + if isinstance(position, bool) or not isinstance(position, int): + raise RuntimeError(f"{label} token entry {index} has an invalid position") + if isinstance(token_id, bool) or not isinstance(token_id, int): + raise RuntimeError(f"{label} token entry {index} has an invalid token_id") + logprob = entry.get("logprob") + if ( + isinstance(logprob, bool) + or not isinstance(logprob, (int, float)) + or not math.isfinite(logprob) + ): + raise RuntimeError(f"{label} token entry {index} has a non-finite logprob") from None + positions.append(position) + token_ids.append(token_id) + all_token_ids = artifact.get("token_ids") + if not ( + isinstance(all_token_ids, list) + and all_token_ids + and all( + isinstance(token_id, int) and not isinstance(token_id, bool) + for token_id in all_token_ids + ) + ): + raise RuntimeError(f"{label} complete token_ids is missing or invalid") + expected_positions = list(range(1, len(entries) + 1)) + if positions != expected_positions: + raise RuntimeError(f"{label} active-token positions are not canonical") + if len(all_token_ids) != len(entries) + 1: + raise RuntimeError(f"{label} complete token_ids length is inconsistent") + if any( + all_token_ids[position] != token_id + for position, token_id in zip(positions, token_ids, strict=True) + ): + raise RuntimeError(f"{label} active token IDs do not match complete token_ids") + digest = hashlib.sha256( + json.dumps(all_token_ids, separators=(",", ":")).encode("ascii") + ).hexdigest() + declared_digest = artifact.get("token_ids_sha256") + if declared_digest != digest: + raise RuntimeError(f"{label} token_ids_sha256 does not match complete token_ids") + return positions, token_ids, all_token_ids, digest + + +def _parse_cp_sizes(raw: str) -> tuple[int, ...]: + try: + values = tuple(int(item.strip()) for item in raw.split(",") if item.strip()) + except ValueError as exc: + raise ValueError("cp-sizes must be a comma-separated list of positive integers") from exc + if values != (1, 2): + raise ValueError("cp-sizes must be exactly 1,2 for the CP consistency comparison") + return values + + +def _validate_args(args: argparse.Namespace, cp_sizes: tuple[int, ...]) -> None: + if not args.teacher_script.is_file(): + raise FileNotFoundError(f"teacher script does not exist: {args.teacher_script}") + if not args.model.exists(): + raise FileNotFoundError(f"model path does not exist: {args.model}") + if not args.token_artifact.is_file(): + raise FileNotFoundError(f"token artifact does not exist: {args.token_artifact}") + if args.tensor_parallel_size < 1: + raise ValueError("parallel sizes must be positive") + if not math.isfinite(args.atol) or args.atol < 0: + raise ValueError("atol must be finite and non-negative") + + +if __name__ == "__main__": + raise SystemExit(main()) 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..5eed0ec5 --- /dev/null +++ b/scripts/ws2_p2p_nccl_attention_reference_check.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Two-GPU P2P NCCL correctness reference 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 math +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 ( # noqa: E402 + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionCPMergedState, + AttentionCPPartialState, + AttentionParallelSpec, + P2PNCCLAttentionCPCommunication, +) +from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( # noqa: E402 + 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) + parser.add_argument("--final-write-atol", type=float, default=2.0e-2) + 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( + { + "schema_version": "ws2_p2p_nccl_attention_reference/v1", + "backend": str(dist.get_backend()), + "world_size": world_size, + "global_failure_count": int(failures.item()), + "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.batch < 1: + raise ValueError("batch must be positive") + 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 != 16 or args.kv_heads != 4 or args.head_dim != 128: + raise ValueError("TP=2 Qwen3-8B local heads must be Hq=16, Hkv=4, D=128") + for name in ("atol", "final_write_atol"): + value = float(getattr(args, name)) + if not math.isfinite(value) or value < 0: + raise ValueError(f"{name} must be finite and non-negative") + + 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)) + 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=owner_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 = owner_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()) + final_out = local.out.to(q.dtype) + expected_final_out = full_out[:, :, start:end, :].to(q.dtype) + final_out_max_abs = float((final_out.float() - expected_final_out.float()).abs().max().item()) + gathered_indices = [state.block.global_block_index for state in gathered] + passed = ( + gathered_indices == list(range(len(blocks))) + and out_max_abs <= args.atol + and lse_max_abs <= args.atol + and final_out.dtype == q.dtype + and final_out_max_abs <= args.final_write_atol + ) + return { + "rank": rank, + "device": str(device), + "dtype": "bf16", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "final_output_dtype": str(final_out.dtype).removeprefix("torch."), + "transport": "p2p_nccl_reference", + "query_range": [start, end], + "world_size": 2, + "expected_block_manifest": [block.provenance() for block in blocks], + "gathered_block_indices": gathered_indices, + "out_max_abs": out_max_abs, + "lse_max_abs": lse_max_abs, + "final_out_max_abs": final_out_max_abs, + "atol": args.atol, + "final_write_atol": args.final_write_atol, + "passed": passed, + } + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 19ca52df..7c0d9841 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): @@ -27,6 +32,7 @@ def _args(**overrides): "n_dim": 32, "theta": 1.0e6, "eps": 1.0e-6, + "arch_key": None, } values.update(overrides) return argparse.Namespace(**values) @@ -86,6 +92,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_attention_gpu_acceptance.py b/tests/test_ws2_attention_gpu_acceptance.py new file mode 100644 index 00000000..d75e617c --- /dev/null +++ b/tests/test_ws2_attention_gpu_acceptance.py @@ -0,0 +1,620 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""CPU-safe validation for the strict WS2 Attention GPU acceptance runner.""" + +from __future__ import annotations + +import json +import subprocess + +from scripts.ws2_attention_gpu_acceptance import ( + AcceptanceCase, + _run_case, + build_acceptance_cases, + parse_args, + run_acceptance, + validate_native_te_report, + validate_p2p_report, + validate_pr5_report, + validate_pr7_report, +) + + +def test_manifest_fails_closed_for_every_unexecuted_required_case(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + report = run_acceptance(args) + + assert report["status"] == "failed" + assert report["passed"] is False + assert "custom_cuda_ag_rs" in report["failed_required_cases"] + assert all(not case["passed"] for case in report["cases"]) + + +def test_matrix_contains_required_modes_splitk_and_communication(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + cases = build_acceptance_cases(args) + names = {case.name for case in cases} + + assert "pr5_cp_forward_backward_dlogp" in names + assert "p2p_nccl_reference" in names + assert "native_te_kv_ring_cp_compare" in names + assert "pr7_flashinfer_decode_disabled" in names + assert "pr7_flashinfer_decode_fixed" in names + assert "pr7_flashinfer_prefill_disabled" in names + assert "pr7_flashinfer_prefill_fixed" in names + assert "custom_cuda_ag_rs" in names + assert "p2p_nccl_reference_tp2_cp2" in names + assert "p2p_nccl_reference_tp2_cp2_replica2" in names + assert "custom_cuda_ag_rs_tp2_cp2" in names + assert "custom_cuda_ag_rs_tp2_cp2_replica2" in names + communication_cases = [ + case for case in cases if case.name.startswith(("p2p_nccl_reference", "custom_cuda_ag_rs")) + ] + assert len(communication_cases) == 6 + assert all( + case.command is not None and "--transport" in case.command for case in communication_cases + ) + by_name = {case.name: case for case in cases} + assert by_name["pr7_flashinfer_decode_disabled"].required is True + strict_command = by_name["pr7_flashinfer_decode_disabled"].command + if strict_command is not None: + assert "--strict" in strict_command + assert by_name["pr7_flashinfer_decode_fixed"].required is False + diagnostic_command = by_name["pr7_flashinfer_decode_fixed"].command + if diagnostic_command is not None: + assert "--strict" not in diagnostic_command + + +def test_formal_communication_cases_cover_four_and_eight_rank_entrypoints(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + cases = {case.name: case for case in build_acceptance_cases(args)} + + p2p = cases["p2p_nccl_reference_tp2_cp2"] + assert p2p.command is not None + assert "--nproc-per-node=4" in p2p.command + assert "--transport" in p2p.command + assert "p2p_nccl_reference" in p2p.command + assert "--repeats" in p2p.command + assert p2p.report_path is not None + + custom = cases["custom_cuda_ag_rs_tp2_cp2"] + assert custom.command is not None + assert "--nproc-per-node=4" in custom.command + assert "cuda_ag_rs" in custom.command + assert "--strict-shared-core" in custom.command + assert custom.report_path is not None + + replicated = cases["custom_cuda_ag_rs_tp2_cp2_replica2"] + assert replicated.command is not None + assert "--nproc-per-node=8" in replicated.command + assert "--strict-shared-core" in replicated.command + + p2p_replica = cases["p2p_nccl_reference_tp2_cp2_replica2"] + assert p2p_replica.command is not None + assert "--strict-shared-core" not in p2p_replica.command + + +def test_native_te_kv_ring_is_optional_diagnostic(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + case = next( + case for case in build_acceptance_cases(args) if case.name == "native_te_kv_ring_cp_compare" + ) + + assert case.required is False + + +def test_allreduce_is_optional_for_attention_gate(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + case = next( + case for case in build_acceptance_cases(args) if case.name == "custom_cuda_allreduce" + ) + + assert case.required is False + + +def test_dlogp_default_uses_shared_bf16_logprob_tolerance(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + + assert args.dlogp_atol == 5.0e-2 + assert args.pr7_out_atol == 1.0e-2 + assert args.pr7_lse_atol == 2.0e-3 + assert args.pr7_dlogp_atol == 2.0e-3 + + +def test_native_te_validator_requires_native_kv_ring_and_cp_compare(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + report = { + "schema_version": "ws2_megatron_te_cp_compare/v1", + "status": "passed", + "passed": True, + "transport": "native_te_kv_ring", + "requested": {"cp_comm_type": "p2p", "context_parallel_sizes": [1, 2]}, + "comparison": { + "pass": True, + "left_cp_size": 1, + "right_cp_size": 2, + "max_abs": 0.0, + "token_ids_sha256": "a" * 64, + }, + "runs": [ + { + "cp_size": cp_size, + "status": "passed", + "active_token_count": 2, + "token_ids_sha256": "a" * 64, + "actual": { + "provider": { + "transformer_impl": "transformer_engine", + "cp_comm_type": "p2p", + } + }, + } + for cp_size in (1, 2) + ], + "errors": [], + } + + assert validate_native_te_report(report, args) == [] + report["comparison"]["max_abs"] = "not-a-number" + assert any("not numeric" in error for error in validate_native_te_report(report, args)) + report["comparison"]["max_abs"] = 0.0 + report["transport"] = "native_te_kv_all_gather" + assert validate_native_te_report(report, args) + report["transport"] = "native_te_kv_ring" + report["runs"] = report["runs"][:1] + assert "native TE report must contain exactly two runs" in validate_native_te_report( + report, args + ) + + +def test_run_mode_preserves_structured_not_available_reports(tmp_path): + report_path = tmp_path / "not-available.json" + args = parse_args(["--mode", "run", "--output", str(tmp_path / "acceptance.json")]) + case = AcceptanceCase( + name="optional", + command=("fake",), + report_path=report_path, + ) + + def fake_runner(command, **kwargs): + report_path.write_text( + json.dumps( + { + "status": "not_available", + "errors": ["FlashInfer unavailable: missing wheel"], + } + ), + encoding="utf-8", + ) + return subprocess.CompletedProcess(command, 1, stdout="", stderr="") + + row = _run_case(case, args, runner=fake_runner) + assert row["status"] == "not_available" + assert row["passed"] is False + assert row["errors"] == ["FlashInfer unavailable: missing wheel"] + + +def test_run_mode_does_not_pass_when_reports_are_missing(tmp_path): + args = parse_args( + [ + "--mode", + "run", + "--output", + str(tmp_path / "acceptance.json"), + ] + ) + + def fake_runner(command, **kwargs): + return subprocess.CompletedProcess(command, 0, stdout="{}", stderr="") + + report = run_acceptance(args, runner=fake_runner) + + assert report["passed"] is False + assert "custom_cuda_ag_rs" in report["failed_required_cases"] + assert any(case["status"] == "invalid_report" for case in report["cases"]) + + +def test_pr7_strict_validation_rejects_requested_only_split_plan(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + report = { + "status": "passed", + "passed": True, + "candidate_provenance": { + "arithmetic_semantics_verified": True, + "actual_split_kv_plans": [ + { + "actual_split_kv_policy": None, + "actual_split_boundaries": [], + } + ], + "actual_split_kv_plan_set": None, + }, + "drift": { + "out": {"max_abs": 5.0e-3}, + "lse": {"max_abs": 1.0e-3}, + "dlogp": {"max_abs": 1.0e-3}, + }, + "batch_invariant_sweep": {"passed": True}, + "page_layout_invariant_sweep": {"passed": True}, + } + + errors = validate_pr7_report(report, args, expected_policy="fixed") + + assert any("actual Split-K policy" in error for error in errors) + assert any("boundaries" in error for error in errors) + assert any("plan set" in error for error in errors) + assert not any(error.startswith("PR7.") for error in errors) + + report["drift"]["out"]["max_abs"] = 2.0e-2 + assert any( + error.startswith("PR7.out") + for error in validate_pr7_report(report, args, expected_policy="fixed") + ) + + +def test_pr7_strict_validation_accepts_shared_no_split_core(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + report = { + "status": "passed", + "passed": True, + "candidate_provenance": { + "arithmetic_semantics_verified": True, + "strict_mode": True, + "strict_core_id": "rlkernel.attention.deterministic_core.v1", + "strict_schedule": "single_batch_single_query_global_kv_blocks", + "native_attention_arithmetic": False, + "fallback": False, + "strict_core_row_plans": [ + { + "actual_split_kv_policy": "disabled", + "actual_split_boundaries": [[0, 8]], + } + ], + }, + "drift": { + "out": {"max_abs": 0.0}, + "lse": {"max_abs": 0.0}, + "dlogp": {"max_abs": 0.0}, + }, + "batch_invariant_sweep": { + "passed": True, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + }, + "page_layout_invariant_sweep": { + "passed": True, + "out": {"max_abs": 0.0}, + "lse": {"max_abs": 0.0}, + }, + } + + assert ( + validate_pr7_report( + report, + args, + expected_policy="disabled", + strict_expected=True, + ) + == [] + ) + + report["candidate_provenance"]["strict_schedule"] = "different_schedule" + assert any( + "strict arithmetic schedule" in error + for error in validate_pr7_report( + report, + args, + expected_policy="disabled", + strict_expected=True, + ) + ) + + +def test_acceptance_report_is_json_serializable(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + json.dumps(run_acceptance(args)) + + +def _valid_pr5_report(): + def case(mode, policy, split_size): + stats = {"max_abs": 0.0} + entries = [] + for tp_rank in range(2): + for cp_rank in range(2): + for owner_cp_rank, owner_range in enumerate(([0, 2], [2, 4])): + entries.append( + { + "batch_index": 0, + "tp_rank": tp_rank, + "cp_rank": cp_rank, + "owner_cp_rank": owner_cp_rank, + "expected_kv_range": owner_range, + "requested_split_kv_policy": policy, + "actual_split_kv_policy": policy, + "actual_split_kv_size": split_size, + "actual_split_kv_count": 1 if split_size is None else 2, + "actual_split_boundaries": ( + [owner_range] + if split_size is None + else [ + [owner_range[0], owner_range[0] + 1], + [owner_range[0] + 1, owner_range[1]], + ] + ), + "split_kv_merge_order": "global_block_index", + "split_kv_accum_dtype": "fp32", + "split_kv_downcast_at": "final_write", + "split_kv_plan_source": "test_runtime", + "split_kv_fallback": False, + "split_kv_fallback_reason": None, + } + ) + return { + "case_name": f"{mode}-{policy}", + "attention_mode": mode, + "topology": {"tp_world_size": 2, "cp_world_size": 2}, + "provenance": { + "requested_split_kv_policy": policy, + "requested_split_kv_size": split_size, + "rope": {"rope_state": "post_rope"}, + "actual_split_kv_plan_set": { + "batch_size": 1, + "tp_world_size": 2, + "cp_world_size": 2, + "total_kv_tokens": [4], + "entries": entries, + "coverage": "complete_batch_tp_cp_owner_cartesian_product", + }, + }, + "drift": {"cp_merge_fp32": {"out": stats, "lse": stats}}, + "dlogp": {"status": "available", "drift": stats}, + "backward": { + "status": "available", + "report": {"drifts": [{"dq": stats, "dk": stats, "dv": stats}]}, + }, + } + + return { + "schema_version": "ws2_cp_attention_drift/v2", + "issue": 235, + "pr": 5, + "runtime": {"device": "cuda:0"}, + "target": { + "model": "qwen3-8b", + "dtype": "bf16", + "global_num_query_heads": 32, + "global_num_kv_heads": 8, + "head_dim": 128, + "batch": 1, + }, + "cases": [case("prefill", "disabled", None), case("chunked_prefill", "fixed", 4)], + } + + +def test_pr5_validation_binds_gpu_identity_and_nonempty_backward(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + report = _valid_pr5_report() + assert validate_pr5_report(report, args) == [] + + report["runtime"]["device"] = "cpu" + report["cases"][0]["backward"]["report"]["drifts"] = [] + errors = validate_pr5_report(report, args) + assert any("not produced on CUDA" in error for error in errors) + assert any("backward drift rows" in error for error in errors) + + +def test_pr5_validation_rejects_nonfinite_or_negative_drift(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + report = _valid_pr5_report() + report["cases"][0]["drift"]["cp_merge_fp32"]["out"] = {"max_abs": float("nan")} + report["cases"][1]["dlogp"]["drift"] = {"max_abs": -1.0} + + errors = validate_pr5_report(report, args) + assert sum("finite and non-negative" in error for error in errors) == 2 + + +def _valid_p2p_report(world_size=4, transport="p2p_nccl_reference"): + tp_world_size = 1 if world_size == 2 else 2 + replica_count = 2 if world_size == 8 else 1 + manifest_by_tp = {} + rows = [] + for rank in range(world_size): + replica_rank = rank % 4 if world_size == 8 else rank + replica_index = rank // 4 if world_size == 8 else 0 + tp_rank = 0 if world_size == 2 else replica_rank // 2 + cp_rank = replica_rank % 2 + manifest = manifest_by_tp.setdefault( + tp_rank, + [ + { + "global_block_index": block, + "kv_block_start": block * 4, + "kv_block_end": block * 4 + 4, + "owner_cp_rank": 0 if block < 2 else 1, + "owner_tp_rank": tp_rank, + } + for block in range(4) + ], + ) + rows.append( + { + "rank": rank, + "global_world_size": world_size, + "tp_rank": tp_rank, + "tp_world_size": 2, + "cp_rank": cp_rank, + "cp_world_size": 2, + "replica_index": replica_index, + "replica_count": replica_count, + "passed": True, + "global_failure_count": 0, + "transport": transport, + "query_ag": transport, + "protocol": "ag_query_local_kv_rs_out_lse", + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "query_ag_max_abs": 0.0, + "device": f"cuda:{rank}", + "dtype": "bf16", + "accum_dtype": "fp32", + "downcast_at": "final_write", + "final_output_dtype": "bfloat16", + "query_range": [0, 8] if cp_rank == 0 else [8, 16], + "expected_block_manifest": manifest, + "local_block_indices": [0, 1] if cp_rank == 0 else [2, 3], + "gathered_block_indices": [0, 1, 2, 3], + "repeat_count": 3, + "repeat_query_bitwise": True, + "repeat_out_bitwise": True, + "repeat_lse_bitwise": True, + "repeat_manifest_bitwise": True, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "final_out_max_abs": 0.0, + "atol": 2.0e-4, + "final_write_atol": 2.0e-2, + "strict_shared_core": ( + { + "executed": True, + "passed": True, + "strict_core_id": "rlkernel.attention.deterministic_core.v1", + "strict_schedule": "single_batch_single_query_global_kv_blocks", + "actual_backend": "rlkernel.cuda.deterministic_attention", + "communication_backend": "self_owned_cuda_ag_rs", + "production_ready": True, + "strict_mode": True, + "native_attention_arithmetic": False, + "fallback": False, + "split_kv_policy": "disabled", + "communication_autograd": True, + "bitwise": { + "out": True, + "lse": True, + "dq": True, + "dk": True, + "dv": True, + }, + "max_abs": { + "out": 0.0, + "lse": 0.0, + "dq": 0.0, + "dk": 0.0, + "dv": 0.0, + }, + "repeat_out_bitwise": True, + "repeat_lse_bitwise": True, + } + if transport == "cuda_ag_rs" + else {"executed": False, "passed": False} + ), + } + ) + return { + "schema_version": ( + "ws2_p2p_nccl_attention_reference/v1" + if transport == "p2p_nccl_reference" + else "ws2_cuda_ag_rs_attention/v1" + ), + "backend": "nccl", + "transport": transport, + "world_size": world_size, + "tp_world_size": tp_world_size, + "cp_world_size": 2, + "replica_count": replica_count, + "global_failure_count": 0, + "ranks": rows, + } + + +def test_p2p_validation_binds_nccl_rank_and_arithmetic_provenance(): + report = _valid_p2p_report() + assert validate_p2p_report(report, expected_world_size=4) == [] + + report["ranks"][1]["transport"] = "gloo" + report["ranks"][1]["rank"] = 0 + errors = validate_p2p_report(report, expected_world_size=4) + assert any("p2p_nccl_reference" in error for error in errors) + assert any("ranks 0 through 3" in error for error in errors) + + +def test_p2p_validation_accepts_legacy_two_rank_artifact(): + report = _valid_p2p_report(world_size=2) + assert validate_p2p_report(report) == [] + assert validate_p2p_report(report, expected_world_size=4) + + +def test_cuda_ag_rs_validation_accepts_two_tp2_cp2_replicas(): + report = _valid_p2p_report(world_size=8, transport="cuda_ag_rs") + + assert ( + validate_p2p_report( + report, + expected_transport="cuda_ag_rs", + expected_world_size=8, + expected_strict_core=True, + ) + == [] + ) + + +def test_cuda_ag_rs_validation_rejects_missing_strict_gradient_bitwise_evidence(): + report = _valid_p2p_report(world_size=4, transport="cuda_ag_rs") + report["ranks"][0]["strict_shared_core"]["bitwise"]["dk"] = False + + errors = validate_p2p_report( + report, + expected_transport="cuda_ag_rs", + expected_world_size=4, + expected_strict_core=True, + ) + assert any("gradient bitwise evidence" in error for error in errors) + + +def test_p2p_validation_rejects_claimed_downcast_without_final_output_evidence(): + report = _valid_p2p_report() + for row in report["ranks"]: + row.pop("final_output_dtype") + row.pop("expected_block_manifest") + row.pop("final_out_max_abs") + row.pop("final_write_atol") + + errors = validate_p2p_report(report, expected_world_size=4) + assert any("final output dtype" in error for error in errors) + assert any("gathered block order/coverage" in error for error in errors) + assert any("final_out_max_abs" in error for error in errors) + + +def test_p2p_validation_rejects_forged_manifest_and_rank_query_mapping(): + report = _valid_p2p_report() + report["ranks"][0]["expected_block_manifest"][1]["kv_block_start"] = 5 + report["ranks"][0]["expected_block_manifest"][1]["owner_cp_rank"] = 3 + report["ranks"][0]["query_range"] = [8, 16] + report["ranks"][1]["query_range"] = [0, 8] + + errors = validate_p2p_report(report, expected_world_size=4) + assert any("gap-free KV coverage" in error for error in errors) + assert any("outside the TP-local CP=2 group" in error for error in errors) + assert any("both CP ranks" in error for error in errors) + assert any("query ownership ranges" in error for error in errors) + + +def test_pr5_validation_rejects_forged_plan_set_coverage(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + report = _valid_pr5_report() + plan_set = report["cases"][0]["provenance"]["actual_split_kv_plan_set"] + plan_set["entries"] = plan_set["entries"][:-1] + plan_set["entries"][0]["split_kv_accum_dtype"] = "bf16" + + errors = validate_pr5_report(report, args) + assert any("coordinate coverage is incomplete" in error for error in errors) + assert any("accumulation dtype is not fp32" in error for error in errors) + + +def test_pr5_validation_reports_malformed_coordinates_without_crashing(tmp_path): + args = parse_args(["--output", str(tmp_path / "acceptance.json")]) + report = _valid_pr5_report() + plan_set = report["cases"][0]["provenance"]["actual_split_kv_plan_set"] + plan_set["entries"][0]["batch_index"] = [] + + errors = validate_pr5_report(report, args) + assert any("coordinate must contain integers" in error for error in errors) + assert any("coordinate coverage is incomplete" in error for error in errors) 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..b5e8e913 --- /dev/null +++ b/tests/test_ws2_cp_attention_drift_benchmark.py @@ -0,0 +1,170 @@ +# 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_requested" + 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 + assert chunked["dlogp"]["status"] == "not_requested" + + +def test_include_dlogp_reports_active_selected_token_drift(): + report = run_benchmark( + parse_args( + [ + "--smoke", + "--include-dlogp", + "--tp-world-sizes", + "2", + "--cp-world-sizes", + "2", + "--kv-chunk-sizes", + "none,1", + ] + ) + ) + + assert report["dlogp"]["status"] == "requested" + for case in report["cases"]: + dlogp = case["dlogp"] + assert dlogp["status"] == "available" + assert dlogp["projection"] == "synthetic_fp32_lm_head_projection" + assert dlogp["active_token_count"] == 3 + assert dlogp["drift"]["active_count"] == 3 + assert dlogp["drift"]["max_abs"] >= 0.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", + ] + ) + ) diff --git a/tests/test_ws2_megatron_te_cp_compare.py b/tests/test_ws2_megatron_te_cp_compare.py new file mode 100644 index 00000000..bed578de --- /dev/null +++ b/tests/test_ws2_megatron_te_cp_compare.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import hashlib +import json +from argparse import Namespace +from pathlib import Path + +import pytest + +from scripts.ws2_megatron_te_cp_compare import _compare_runs, _parse_cp_sizes, _validate_args + + +def _artifact(path: Path, *, offset: float) -> Path: + token_ids = [42, 7, 8] + path.write_text( + json.dumps( + { + "token_ids": token_ids, + "token_ids_sha256": hashlib.sha256( + json.dumps(token_ids, separators=(",", ":")).encode("ascii") + ).hexdigest(), + "active_token_logprobs": [ + {"position": 1, "token_id": 7, "logprob": -1.0 + offset}, + {"position": 2, "token_id": 8, "logprob": -2.0 + offset}, + ], + } + ) + ) + return path + + +def test_native_te_comparison_is_ordered_and_records_worst_token(tmp_path): + left = _artifact(tmp_path / "cp1.json", offset=0.0) + right = _artifact(tmp_path / "cp2.json", offset=0.02) + + report = _compare_runs( + [ + {"cp_size": 1, "output": str(left)}, + {"cp_size": 2, "output": str(right)}, + ], + atol=0.05, + ) + + assert report["pass"] is True + assert report["max_abs"] == pytest.approx(0.02) + assert report["worst"]["position"] == 1 + assert report["token_ids_sha256"] + + +def test_native_te_comparison_rejects_token_identity_mismatch(tmp_path): + left = _artifact(tmp_path / "cp1.json", offset=0.0) + right = _artifact(tmp_path / "cp2.json", offset=0.02) + payload = json.loads(right.read_text()) + payload["active_token_logprobs"][1]["token_id"] = 99 + right.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match="token IDs"): + _compare_runs( + [ + {"cp_size": 1, "output": str(left)}, + {"cp_size": 2, "output": str(right)}, + ], + atol=0.05, + ) + + +def test_native_te_comparison_rejects_position_mismatch(tmp_path): + left = _artifact(tmp_path / "cp1.json", offset=0.0) + right = _artifact(tmp_path / "cp2.json", offset=0.02) + payload = json.loads(right.read_text()) + payload["active_token_logprobs"][1]["position"] = 3 + right.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match="positions"): + _compare_runs( + [ + {"cp_size": 1, "output": str(left)}, + {"cp_size": 2, "output": str(right)}, + ], + atol=0.05, + ) + + +def test_native_te_comparison_rejects_invalid_token_hash(tmp_path): + left = _artifact(tmp_path / "cp1.json", offset=0.0) + right = _artifact(tmp_path / "cp2.json", offset=0.02) + payload = json.loads(right.read_text()) + payload["token_ids_sha256"] = "0" * 64 + right.write_text(json.dumps(payload)) + + with pytest.raises(RuntimeError, match="token_ids_sha256"): + _compare_runs( + [ + {"cp_size": 1, "output": str(left)}, + {"cp_size": 2, "output": str(right)}, + ], + atol=0.05, + ) + + +def test_cp_sizes_require_exactly_cp1_and_cp2(): + assert _parse_cp_sizes("1,2") == (1, 2) + + for value in ("2", "1,2,4", "0,2", "2,1", "a,2"): + try: + _parse_cp_sizes(value) + except ValueError: + continue + raise AssertionError(f"expected invalid cp-sizes to fail: {value}") + + +def test_native_te_rejects_nonfinite_tolerance(tmp_path): + args = Namespace( + teacher_script=tmp_path / "teacher.py", + model=tmp_path / "model", + token_artifact=tmp_path / "tokens.json", + tensor_parallel_size=2, + atol=float("nan"), + ) + args.teacher_script.touch() + args.model.mkdir() + args.token_artifact.touch() + + with pytest.raises(ValueError, match="finite and non-negative"): + _validate_args(args, (1, 2))