diff --git a/benchmarks/benchmark_rocm_attention.py b/benchmarks/benchmark_rocm_attention.py new file mode 100644 index 00000000..e3c250e2 --- /dev/null +++ b/benchmarks/benchmark_rocm_attention.py @@ -0,0 +1,469 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Operator-only benchmark for the strict ROCm attention path. + +Seeded Q/K/V only: no checkpoint, tokenizer, or serving engine. Defaults to the +Qwen3-8B dense head layout (``Hq=32``, ``Hkv=8``, ``D=128``), BF16, causal +prefill. + +Three backends are compared: + +``native`` + ``torch.nn.functional.scaled_dot_product_attention`` with the KV heads + expanded to the Q head count. Not batch-invariant; present as the + throughput reference every ROCm deployment already has. +``triton`` + ``flash_attn`` with the ROCm Triton backend enabled. +``strict`` + ``aiter.rocm.ck_dense_mha`` through the WS2 contract dispatch, i.e. the path + this PR adds. + +Three extra sections quantify what the strict contract actually costs and +where it stops holding. All three are properties of the vendor kernel rather +than of the integration: + +* ``determinism_cost`` — AITER ``mha_bwd`` with ``deterministic`` on vs off. +* ``batch_composition`` — whether raw AITER returns the same bits for a batch + and for the same rows submitted one at a time, swept over shapes because the + answer varies with shape. +* ``tp_head_count_sensitivity`` — whether a head shard depends on how many + heads shared its launch, i.e. whether the path is TP-degree invariant. +""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +from pathlib import Path +from typing import Any, Callable + +import torch + +DEFAULT_Q_HEADS = 32 +DEFAULT_KV_HEADS = 8 +DEFAULT_HEAD_DIM = 128 + + +class _ContextParallel: + rank = 0 + world_size = 1 + layout = "single" + + +class _Request: + """Structural request understood by the Vime attention provider.""" + + def __init__(self, query, key, value, metadata): + self.query = query + self.key = key + self.value = value + self.metadata = metadata + self.context_parallel = _ContextParallel() + self.tensor_parallel_group = None + self.key_padding_mask = None + + +def _metadata(q_heads: int, kv_heads: int) -> dict[str, Any]: + return { + "global_q_heads": q_heads, + "global_kv_heads": kv_heads, + "tp_rank": 0, + "tp_world_size": 1, + "attention_mode": "prefill", + "role": "train", + "causal": True, + } + + +def _tensors(batch, q_heads, kv_heads, seq_len, head_dim, dtype, seed=0): + generator = torch.Generator(device="cuda").manual_seed(seed) + query = torch.randn( + batch, q_heads, seq_len, head_dim, generator=generator, device="cuda", dtype=dtype + ) + key = torch.randn( + batch, kv_heads, seq_len, head_dim, generator=generator, device="cuda", dtype=dtype + ) + value = torch.randn( + batch, kv_heads, seq_len, head_dim, generator=generator, device="cuda", dtype=dtype + ) + return query, key, value + + +def _measure(run: Callable[[bool], None], *, backward: bool, warmup: int, iters: int): + torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + for _ in range(warmup): + run(backward) + torch.cuda.synchronize() + samples: list[float] = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run(backward) + end.record() + torch.cuda.synchronize() + samples.append(start.elapsed_time(end)) + samples.sort() + return { + "median_ms": round(statistics.median(samples), 4), + "p95_ms": round(samples[int(0.95 * (len(samples) - 1))], 4), + "peak_mib": round(torch.cuda.max_memory_allocated() / 2**20, 1), + } + + +def _native_runner(query, key, value, q_heads, kv_heads): + repeats = q_heads // kv_heads + + def run(backward: bool) -> None: + q = query.detach().requires_grad_(backward) + k = key.detach().requires_grad_(backward) + v = value.detach().requires_grad_(backward) + out = torch.nn.functional.scaled_dot_product_attention( + q, + k.repeat_interleave(repeats, dim=1), + v.repeat_interleave(repeats, dim=1), + is_causal=True, + ) + if backward: + out.backward(torch.ones_like(out)) + + return run + + +def _triton_runner(query, key, value, head_dim): + import os + + os.environ["FLASH_ATTENTION_TRITON_AMD_ENABLE"] = "TRUE" + from flash_attn import flash_attn_func + + scale = 1.0 / math.sqrt(head_dim) + + def run(backward: bool) -> None: + q = query.detach().transpose(1, 2).contiguous().requires_grad_(backward) + k = key.detach().transpose(1, 2).contiguous().requires_grad_(backward) + v = value.detach().transpose(1, 2).contiguous().requires_grad_(backward) + out = flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=scale, causal=True) + if backward: + out.backward(torch.ones_like(out)) + + return run + + +def _strict_runner(query, key, value, metadata): + from rl_engine.integrations.vime import attention_provider + + def run(backward: bool) -> None: + q = query.detach().requires_grad_(backward) + k = key.detach().requires_grad_(backward) + v = value.detach().requires_grad_(backward) + result = attention_provider(_Request(q, k, v, metadata)) + if backward: + result.out.backward(torch.ones_like(result.out)) + + return run + + +def _determinism_cost(seq_lens, q_heads, kv_heads, head_dim, dtype, warmup, iters): + """AITER deterministic backward vs the non-deterministic one.""" + + from aiter.ops.mha import mha_bwd, mha_fwd + + scale = 1.0 / math.sqrt(head_dim) + rows = [] + for seq_len in seq_lens: + query, key, value = _tensors(1, q_heads, kv_heads, seq_len, head_dim, dtype) + # AITER consumes [B, S, H, D]. + q = query.transpose(1, 2).contiguous() + k = key.transpose(1, 2).contiguous() + v = value.transpose(1, 2).contiguous() + out, lse, _mask, rng_state = mha_fwd(q, k, v, 0.0, scale, True, -1, -1, 0, True, False) + grad_out = torch.ones_like(out) + entry: dict[str, Any] = {"seq_len": seq_len} + for deterministic in (True, False): + + # Bind the tensors as defaults: they are released below, and the + # closure must not depend on the enclosing names still existing. + def run( + _backward: bool, + deterministic=deterministic, + grad_out=grad_out, + q=q, + k=k, + v=v, + out=out, + lse=lse, + rng_state=rng_state, + ) -> None: + mha_bwd( + grad_out, + q, + k, + v, + out, + lse, + 0.0, + scale, + True, + -1, + -1, + deterministic, + rng_state=rng_state, + ) + + key_name = "deterministic" if deterministic else "non_deterministic" + entry[key_name] = _measure(run, backward=False, warmup=warmup, iters=iters) + entry["time_ratio"] = round( + entry["deterministic"]["median_ms"] / entry["non_deterministic"]["median_ms"], 2 + ) + entry["memory_ratio"] = round( + entry["deterministic"]["peak_mib"] / entry["non_deterministic"]["peak_mib"], 1 + ) + rows.append(entry) + del query, key, value, q, k, v, out, lse, grad_out + torch.cuda.empty_cache() + return rows + + +def _batch_composition(q_heads, kv_heads, head_dim, dtype, shapes): + """Whether raw AITER is batch-composition invariant, swept over shapes. + + The strict core sidesteps this by executing one logical row at a time; this + section records what the vendor kernel does without that constraint. The + sweep matters: invariance holds for most shapes and breaks for a few, so a + single-shape probe would report whichever answer it happened to land on. + """ + + from aiter.ops.mha import mha_fwd + + scale = 1.0 / math.sqrt(head_dim) + + def forward(query, key, value): + out, lse, _mask, _rng = mha_fwd( + query.transpose(1, 2).contiguous(), + key.transpose(1, 2).contiguous(), + value.transpose(1, 2).contiguous(), + 0.0, + scale, + True, + -1, + -1, + 0, + True, + False, + ) + return out.transpose(1, 2).contiguous(), lse + + rows = [] + for batch, seq_len in shapes: + if batch < 2: + continue + query, key, value = _tensors(batch, q_heads, kv_heads, seq_len, head_dim, dtype, seed=11) + batched_out, batched_lse = forward(query, key, value) + worst_out = worst_lse = 0.0 + for row in range(batch): + row_out, row_lse = forward( + query[row : row + 1], key[row : row + 1], value[row : row + 1] + ) + worst_out = max(worst_out, (batched_out[row : row + 1] - row_out).abs().max().item()) + worst_lse = max(worst_lse, (batched_lse[row : row + 1] - row_lse).abs().max().item()) + rows.append( + { + "batch": batch, + "seq_len": seq_len, + "raw_aiter_out_max_abs": worst_out, + "raw_aiter_lse_max_abs": worst_lse, + "raw_aiter_is_batch_invariant": worst_out == 0.0 and worst_lse == 0.0, + } + ) + print( + f"batch-composition B={batch} S={seq_len:5d} " + f"out {worst_out:.6e} lse {worst_lse:.6e} " + f"{'invariant' if rows[-1]['raw_aiter_is_batch_invariant'] else 'NOT INVARIANT'}", + flush=True, + ) + del query, key, value, batched_out, batched_lse + torch.cuda.empty_cache() + return rows + + +def _tp_head_count_sensitivity(q_heads, kv_heads, head_dim, dtype, seq_lens, tp_degrees): + """Does a head shard depend on how many heads shared its launch? + + TP shards attention by head and performs no cross-rank reduction, so a rank + computing its own head slice ought to match the corresponding slice of an + unsharded run. Where it does not, the strict path is not TP-degree + invariant and training/rollout must be pinned to one degree. + """ + + from aiter.ops.mha import mha_fwd + + scale = 1.0 / math.sqrt(head_dim) + + def forward(query, key, value): + out, lse, _mask, _rng = mha_fwd( + query.transpose(1, 2).contiguous(), + key.transpose(1, 2).contiguous(), + value.transpose(1, 2).contiguous(), + 0.0, + scale, + True, + -1, + -1, + 0, + True, + False, + ) + return out.transpose(1, 2).contiguous(), lse + + rows = [] + for seq_len in seq_lens: + query, key, value = _tensors(1, q_heads, kv_heads, seq_len, head_dim, dtype, seed=3) + full_out, full_lse = forward(query, key, value) + for tp in tp_degrees: + if q_heads % tp or kv_heads % tp: + continue + local_q, local_kv = q_heads // tp, kv_heads // tp + worst_out = worst_lse = 0.0 + for rank in range(tp): + shard_out, shard_lse = forward( + query[:, rank * local_q : (rank + 1) * local_q], + key[:, rank * local_kv : (rank + 1) * local_kv], + value[:, rank * local_kv : (rank + 1) * local_kv], + ) + ref_out = full_out[:, rank * local_q : (rank + 1) * local_q] + ref_lse = full_lse[:, rank * local_q : (rank + 1) * local_q] + worst_out = max(worst_out, (shard_out - ref_out).abs().max().item()) + worst_lse = max(worst_lse, (shard_lse - ref_lse).abs().max().item()) + invariant = worst_out == 0.0 and worst_lse == 0.0 + rows.append( + { + "seq_len": seq_len, + "tp": tp, + "local_q_heads": local_q, + "local_kv_heads": local_kv, + "out_max_abs": worst_out, + "lse_max_abs": worst_lse, + "tp_degree_invariant": invariant, + } + ) + print( + f"tp-sensitivity S={seq_len:5d} TP={tp} (Hq={local_q},Hkv={local_kv}) " + f"out {worst_out:.6e} lse {worst_lse:.6e} " + f"{'invariant' if invariant else 'NOT INVARIANT'}", + flush=True, + ) + del query, key, value, full_out, full_lse + torch.cuda.empty_cache() + return rows + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, default=Path("attention_results.json")) + parser.add_argument("--q-heads", type=int, default=DEFAULT_Q_HEADS) + parser.add_argument("--kv-heads", type=int, default=DEFAULT_KV_HEADS) + parser.add_argument("--head-dim", type=int, default=DEFAULT_HEAD_DIM) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=50) + parser.add_argument( + "--shapes", + default="1x1024,1x2048,1x4096,2x2048,4x2048", + help="comma-separated BATCHxSEQ pairs", + ) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("this benchmark requires a ROCm (or CUDA) device") + + dtype = torch.bfloat16 + metadata = _metadata(args.q_heads, args.kv_heads) + shapes = [tuple(int(part) for part in pair.split("x")) for pair in args.shapes.split(",")] + + rows = [] + for batch, seq_len in shapes: + query, key, value = _tensors( + batch, args.q_heads, args.kv_heads, seq_len, args.head_dim, dtype + ) + factories = { + "native": lambda q=query, k=key, v=value: _native_runner( + q, k, v, args.q_heads, args.kv_heads + ), + "triton": lambda q=query, k=key, v=value: _triton_runner(q, k, v, args.head_dim), + "strict": lambda q=query, k=key, v=value: _strict_runner(q, k, v, metadata), + } + for name, factory in factories.items(): + try: + runner = factory() + forward = _measure(runner, backward=False, warmup=args.warmup, iters=args.iters) + combined = _measure(runner, backward=True, warmup=args.warmup, iters=args.iters) + except Exception as exc: # noqa: BLE001 - report, do not abort the sweep + print(f"B={batch} S={seq_len} {name}: FAILED {type(exc).__name__}: {exc}") + continue + rows.append( + { + "batch": batch, + "seq_len": seq_len, + "backend": name, + "forward": forward, + "forward_backward": combined, + } + ) + print( + f"B={batch} S={seq_len:5d} {name:8s} " + f"fwd {forward['median_ms']:8.4f} p95 {forward['p95_ms']:8.4f} " + f"peak {forward['peak_mib']:9.1f} | " + f"fwd+bwd {combined['median_ms']:8.4f} peak {combined['peak_mib']:9.1f}", + flush=True, + ) + del query, key, value + torch.cuda.empty_cache() + + seq_lens = sorted({seq for _batch, seq in shapes}) + determinism = _determinism_cost( + seq_lens, args.q_heads, args.kv_heads, args.head_dim, dtype, args.warmup, args.iters + ) + composition_shapes = [ + (batch, seq_len) for batch in (2, 4) for seq_len in sorted({128, 256, 512, *seq_lens}) + ] + composition = _batch_composition( + args.q_heads, args.kv_heads, args.head_dim, dtype, composition_shapes + ) + tp_sensitivity = _tp_head_count_sensitivity( + args.q_heads, + args.kv_heads, + args.head_dim, + dtype, + sorted({512, *seq_lens}), + (2, 4, 8), + ) + + properties = torch.cuda.get_device_properties(0) + payload = { + "environment": { + "gpu": properties.name, + "arch": getattr(properties, "gcnArchName", "unknown"), + "device_count": torch.cuda.device_count(), + "torch": torch.__version__, + "hip": torch.version.hip, + "dtype": "bf16", + "q_heads": args.q_heads, + "kv_heads": args.kv_heads, + "head_dim": args.head_dim, + }, + "latency": rows, + "determinism_cost": determinism, + "batch_composition": composition, + "tp_head_count_sensitivity": tp_sensitivity, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2)) + print("wrote", args.output) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_rocm_collectives.py b/benchmarks/benchmark_rocm_collectives.py new file mode 100644 index 00000000..6202e2e5 --- /dev/null +++ b/benchmarks/benchmark_rocm_collectives.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark deterministic ROCm collectives against native RCCL. + +Example: + + torchrun --standalone --nproc-per-node=8 \ + benchmarks/benchmark_rocm_collectives.py \ + --size-bytes 4096 65536 1048576 16777216 \ + --output benchmarks/results/rocm_collectives_mi300x.json + +The native RCCL rows are performance references only. They are not used as a +bitwise correctness oracle because their floating-point reduction order is not +part of the strict deterministic contract. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import time +from pathlib import Path +from typing import Callable, Sequence + +import torch +import torch.distributed as dist + +from rl_engine.distributed import RCCLDeterministicCollective + +_DTYPES = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, +} +_OPERATIONS = ("all_reduce", "all_gather", "reduce_scatter") + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--size-bytes", + type=int, + nargs="+", + default=[4 * 1024, 64 * 1024, 1024 * 1024, 16 * 1024 * 1024], + ) + parser.add_argument("--dtype", choices=tuple(_DTYPES), default="bf16") + parser.add_argument("--operations", nargs="+", choices=_OPERATIONS, default=_OPERATIONS) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iterations", type=int, default=50) + parser.add_argument("--samples", type=int, default=5) + parser.add_argument("--output", type=Path) + return parser.parse_args(argv) + + +def _validate_args(args: argparse.Namespace) -> None: + if any(size <= 0 for size in args.size_bytes): + raise ValueError("every --size-bytes value must be positive") + if args.warmup < 0: + raise ValueError("--warmup must be non-negative") + if args.iterations <= 0 or args.samples <= 0: + raise ValueError("--iterations and --samples must be positive") + + +def _timed_sample(operation: Callable[[], None], *, warmup: int, iterations: int) -> float: + for _ in range(warmup): + operation() + torch.cuda.synchronize() + dist.barrier() + start = time.perf_counter() + for _ in range(iterations): + operation() + torch.cuda.synchronize() + elapsed = (time.perf_counter() - start) / iterations + + # Report the slowest rank, which is the end-to-end collective latency. + elapsed_tensor = torch.tensor([elapsed], dtype=torch.float64, device="cuda") + dist.all_reduce(elapsed_tensor, op=dist.ReduceOp.MAX) + return float(elapsed_tensor.item()) + + +def _benchmark( + operation: Callable[[], None], + *, + warmup: int, + iterations: int, + samples: int, +) -> dict[str, object]: + timings = [ + _timed_sample(operation, warmup=warmup if index == 0 else 0, iterations=iterations) + for index in range(samples) + ] + median = statistics.median(timings) + return { + "median_us": median * 1.0e6, + "min_us": min(timings) * 1.0e6, + "max_us": max(timings) * 1.0e6, + "samples_us": [value * 1.0e6 for value in timings], + } + + +def _make_inputs( + *, + size_bytes: int, + dtype: torch.dtype, + world_size: int, + rank: int, + device: torch.device, +) -> tuple[torch.Tensor, int]: + element_size = torch.empty((), dtype=dtype).element_size() + elements = max(world_size, size_bytes // element_size) + elements -= elements % world_size + generator = torch.Generator(device="cpu").manual_seed(942 + rank) + tensor = torch.randn(elements, generator=generator, dtype=torch.float32).to( + device=device, + dtype=dtype, + ) + return tensor.contiguous(), elements * element_size + + +def _operation_pair( + name: str, + input_tensor: torch.Tensor, + collective: RCCLDeterministicCollective, + world_size: int, +) -> tuple[Callable[[], None], Callable[[], None], torch.Tensor, torch.Tensor]: + if name == "all_reduce": + deterministic_out = torch.empty_like(input_tensor) + native_out = torch.empty_like(input_tensor) + + def deterministic() -> None: + collective.all_reduce(input_tensor, out=deterministic_out) + + def native() -> None: + native_out.copy_(input_tensor) + dist.all_reduce(native_out) + + elif name == "all_gather": + output_shape = (input_tensor.numel() * world_size,) + deterministic_out = torch.empty(output_shape, dtype=input_tensor.dtype, device="cuda") + native_out = torch.empty_like(deterministic_out) + + def deterministic() -> None: + collective.all_gather(input_tensor, out=deterministic_out) + + def native() -> None: + dist.all_gather_into_tensor(native_out, input_tensor) + + elif name == "reduce_scatter": + output_shape = (input_tensor.numel() // world_size,) + deterministic_out = torch.empty(output_shape, dtype=input_tensor.dtype, device="cuda") + native_out = torch.empty_like(deterministic_out) + + def deterministic() -> None: + collective.reduce_scatter(input_tensor, out=deterministic_out) + + def native() -> None: + dist.reduce_scatter_tensor(native_out, input_tensor) + + else: # pragma: no cover - argparse constrains this value + raise ValueError(f"unsupported operation: {name}") + + return deterministic, native, deterministic_out, native_out + + +def run(args: argparse.Namespace) -> dict[str, object] | None: + _validate_args(args) + if torch.version.hip is None or not torch.cuda.is_available(): + raise RuntimeError("the ROCm collective benchmark requires an available AMD GPU") + + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + torch.cuda.set_device(local_rank) + 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"the benchmark requires 2, 4, or 8 ranks, got {world_size}") + device = torch.device("cuda", local_rank) + dtype = _DTYPES[args.dtype] + max_size_bytes = max(args.size_bytes) + dtype.itemsize * world_size + + rows: list[dict[str, object]] = [] + try: + with RCCLDeterministicCollective( + device=device, + max_size_bytes=max_size_bytes, + ) as collective: + for requested_size in args.size_bytes: + input_tensor, actual_size = _make_inputs( + size_bytes=requested_size, + dtype=dtype, + world_size=world_size, + rank=rank, + device=device, + ) + for name in args.operations: + deterministic, native, deterministic_out, native_out = _operation_pair( + name, + input_tensor, + collective, + world_size, + ) + deterministic() + deterministic_repeat = deterministic_out.clone() + deterministic() + repeat_bitwise = bool(torch.equal(deterministic_out, deterministic_repeat)) + native() + max_abs_vs_native = float( + (deterministic_out.float() - native_out.float()).abs().max().item() + ) + + torch.cuda.reset_peak_memory_stats(device) + deterministic_timing = _benchmark( + deterministic, + warmup=args.warmup, + iterations=args.iterations, + samples=args.samples, + ) + deterministic_peak = int(torch.cuda.max_memory_allocated(device)) + torch.cuda.reset_peak_memory_stats(device) + native_timing = _benchmark( + native, + warmup=args.warmup, + iterations=args.iterations, + samples=args.samples, + ) + native_peak = int(torch.cuda.max_memory_allocated(device)) + deterministic_us = float(deterministic_timing["median_us"]) + native_us = float(native_timing["median_us"]) + rows.append( + { + "operation": name, + "requested_size_bytes": requested_size, + "actual_input_bytes": actual_size, + "dtype": args.dtype, + "deterministic": deterministic_timing, + "native_rccl": native_timing, + "latency_ratio_vs_native": deterministic_us / native_us, + "deterministic_input_gbps": actual_size / (deterministic_us * 1.0e3), + "native_input_gbps": actual_size / (native_us * 1.0e3), + "repeat_bitwise": repeat_bitwise, + "max_abs_vs_native": max_abs_vs_native, + "deterministic_workspace_bytes": collective.workspace_size_bytes, + "deterministic_peak_allocated_bytes": deterministic_peak, + "native_peak_allocated_bytes": native_peak, + } + ) + + reports: list[list[dict[str, object]] | None] = [None] * world_size + dist.all_gather_object(reports, rows) + if rank != 0: + return None + payload = { + "schema_version": "rlkernel.rocm_collective_benchmark.v1", + "world_size": world_size, + "device": torch.cuda.get_device_name(device), + "hip_version": torch.version.hip, + "collective_backend": RCCLDeterministicCollective.backend_id, + "reduction_order": RCCLDeterministicCollective.reduction_order, + "supports_compute_communication_fusion": False, + "warmup": args.warmup, + "iterations": args.iterations, + "samples": args.samples, + "rows": rows, + "all_rank_repeat_bitwise": all( + bool(row["repeat_bitwise"]) + for rank_rows in reports + if rank_rows is not None + for row in rank_rows + ), + } + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2), encoding="utf-8") + return payload + finally: + dist.destroy_process_group() + + +def main(argv: Sequence[str] | None = None) -> int: + payload = run(parse_args(argv)) + if payload is not None: + print(json.dumps(payload, indent=2)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/benchmark_ws2_rocm_attention.py b/benchmarks/benchmark_ws2_rocm_attention.py new file mode 100644 index 00000000..e8a0c7b2 --- /dev/null +++ b/benchmarks/benchmark_ws2_rocm_attention.py @@ -0,0 +1,2093 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 strict ROCm Attention performance and bitwise-parity benchmark. + +Operator-only. No model checkpoint or serving engine is loaded; the shapes are +Qwen3-8B's attention shapes (``Hq=32``, ``Hkv=8``, ``D=128``). + +Measurement matrix and presentation follow PR #325 (`benchmark_rocm_ffn.py`) and +PR #328 (`benchmark_rocm_logp.py`): the timing/accuracy helpers, the spawned +distributed world, and the figure style are taken from those scripts so the three +reports can be read side by side. + +Paths measured: + +- ``sdpa`` PyTorch ``scaled_dot_product_attention``. Speed baseline only, + exactly as PR #325 uses upstream ``Qwen3MLP`` at TP=1: no + accuracy claim is mixed into the speed comparison. +- ``strict-aiter`` ``StrictRocmAiterCKAttentionCore`` — the ROCm production core + (AITER CK dense MHA, non-split API). +- ``reference-native`` ``_C.deterministic_attention_forward/backward`` — the materializing + FP32 reference core, hipified from the shared ``.cu``. +- ``triton-bitwise`` ``TritonDeterministicAttentionOp`` — the Triton port whose + contract is bit-identity with ``reference-native``. + +The headline column is ``triton-bitwise`` versus ``reference-native``: acceptance is +0 mismatched elements on out, lse, dQ, dK and dV. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import math +import multiprocessing as mp +import os +import platform +import queue +import statistics +import tempfile +import threading +import time +import traceback +from datetime import timedelta +from pathlib import Path +from typing import Any, Callable + +import torch +import torch.distributed as dist + +QWEN3_8B_Q_HEADS = 32 +QWEN3_8B_KV_HEADS = 8 +QWEN3_8B_HEAD_DIM = 128 + +DEFAULT_SEQ_LENS = (512, 1024, 2048, 4096) +DEFAULT_TP_DEGREES = (2, 4, 8) +# (label, tp_world_size, cp_world_size, replicas) -- world_size = tp * cp * replicas. +# Replicas run independent CP groups side by side, which is how PR #319 exercised +# 8 ranks at TP=2/CP=2. +DISTRIBUTED_TOPOLOGIES = ( + ("tp1_cp2", 1, 2, 1), + ("tp2_cp2", 2, 2, 1), + ("tp1_cp4", 1, 4, 1), + ("tp2_cp2_x2", 2, 2, 2), + ("tp2_cp4", 2, 4, 1), + ("tp1_cp8", 1, 8, 1), +) + + +# --------------------------------------------------------------------------- +# Measurement helpers (PR #328 benchmark_rocm_logp.py) +# --------------------------------------------------------------------------- + + +def _percentile(values: list[float], percentile: float) -> float: + ordered = sorted(values) + if not ordered: + return float("nan") + position = (len(ordered) - 1) * percentile + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return ordered[lower] + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def _summary_ms(values: list[float]) -> dict[str, float]: + return { + "median_ms": statistics.median(values), + "p95_ms": _percentile(values, 0.95), + "min_ms": min(values), + "max_ms": max(values), + } + + +def _relative_l2(actual: torch.Tensor, expected: torch.Tensor) -> float: + actual_float = actual.detach().double() + expected_float = expected.detach().double() + denominator = torch.linalg.vector_norm(expected_float) + if denominator.item() == 0.0: + return float(torch.linalg.vector_norm(actual_float - expected_float).item()) + return float((torch.linalg.vector_norm(actual_float - expected_float) / denominator).item()) + + +def _accuracy(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: + difference = actual.detach().double() - expected.detach().double() + return { + "max_abs": float(difference.abs().max().item()) if difference.numel() else 0.0, + "relative_l2": _relative_l2(actual, expected), + } + + +def _bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: + return a.shape == b.shape and a.dtype == b.dtype and bool(torch.equal(a, b)) + + +def _mismatch_count(a: torch.Tensor, b: torch.Tensor) -> int: + if a.shape != b.shape: + return -1 + return int((a != b).sum().item()) + + +def _gpu_event_samples( + function: Callable[[], Any], *, warmup: int, samples: int, deadline: float = 0.0 +) -> list[float]: + for _ in range(warmup): + function() + if deadline and time.perf_counter() > deadline: + break + torch.cuda.synchronize() + events = [] + for _ in range(samples): + if deadline and events and time.perf_counter() > deadline: + break + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + function() + end.record() + events.append((start, end)) + torch.cuda.synchronize() + return [float(start.elapsed_time(end)) for start, end in events] + + +def _host_wall_samples( + function: Callable[[], Any], *, warmup: int, samples: int, deadline: float = 0.0 +) -> list[float]: + """Wall-clock timing for host execution, where CUDA events do not apply.""" + for _ in range(warmup): + function() + if deadline and time.perf_counter() > deadline: + break + timings = [] + for _ in range(samples): + if deadline and timings and time.perf_counter() > deadline: + break + start = time.perf_counter() + function() + timings.append((time.perf_counter() - start) * 1000.0) + return timings + + +def _timed_samples( + function: Callable[[], Any], + *, + warmup: int, + samples: int, + device: torch.device, + deadline: float = 0.0, +) -> list[float]: + """Sample, stopping early once ``deadline`` (a perf_counter value) passes. + + The pre-flight projection can under-estimate badly: at S=4096 the materialized + score matrix leaves cache and the compute model stops holding. So the budget is + also enforced here as a hard wall clock, not only as an estimate. At least one + sample is always taken, so a row is never empty. + """ + if device.type == "cuda": + return _gpu_event_samples(function, warmup=warmup, samples=samples, deadline=deadline) + return _host_wall_samples(function, warmup=warmup, samples=samples, deadline=deadline) + + +def _rss_mib() -> float: + with open("/proc/self/statm", "r", encoding="ascii") as handle: + resident_pages = int(handle.read().split()[1]) + return resident_pages * os.sysconf("SC_PAGE_SIZE") / (1024.0 * 1024.0) + + +def _host_peak_rss_mib(function: Callable[[], Any]) -> float: + """Peak resident-set increase during one host call, sampled from /proc. + + The closest host analogue of ``torch.cuda.max_memory_allocated``, but an RSS + high-water delta rather than an allocator statistic: it includes caching-allocator + reuse and page granularity, so it is an approximation and not directly comparable + to the device figures. A call served from already-resident pages can report ~0. + """ + gc.collect() + baseline = _rss_mib() + peak = baseline + stop = threading.Event() + + def sampler() -> None: + nonlocal peak + while not stop.is_set(): + peak = max(peak, _rss_mib()) + stop.wait(0.001) + + thread = threading.Thread(target=sampler, daemon=True) + thread.start() + try: + function() + finally: + stop.set() + thread.join() + return float(max(peak, _rss_mib()) - baseline) + + +def _peak_memory_mib(function: Callable[[], Any], device: torch.device) -> float: + """Peak memory used by one call, above what was live before it.""" + if device.type != "cuda": + return _host_peak_rss_mib(function) + torch.cuda.synchronize() + _empty_cache(device) + torch.cuda.reset_peak_memory_stats() + baseline = torch.cuda.memory_allocated() + function() + torch.cuda.synchronize() + return float((torch.cuda.max_memory_allocated() - baseline) / (1024.0 * 1024.0)) + + +def _device_sync(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.synchronize() + + +def _empty_cache(device: torch.device) -> None: + if device.type == "cuda": + torch.cuda.empty_cache() + else: + gc.collect() + + +# --------------------------------------------------------------------------- +# Attention paths +# --------------------------------------------------------------------------- + + +def _seeded_qkv( + batch: int, + q_heads: int, + kv_heads: int, + seq_len: int, + head_dim: int, + dtype: torch.dtype, + device: torch.device, + *, + seed: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + generator = torch.Generator(device=device).manual_seed(seed) + q = torch.randn( + batch, q_heads, seq_len, head_dim, device=device, dtype=dtype, generator=generator + ) + k = torch.randn( + batch, kv_heads, seq_len, head_dim, device=device, dtype=dtype, generator=generator + ) + v = torch.randn( + batch, kv_heads, seq_len, head_dim, device=device, dtype=dtype, generator=generator + ) + return q, k, v + + +def _positions(batch: int, seq_len: int, device: torch.device) -> torch.Tensor: + return torch.arange(seq_len, device=device, dtype=torch.int32).unsqueeze(0).expand(batch, -1) + + +def _fp64_oracle( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, causal: bool, scale: float +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference (out, lse) in FP64 from the BF16/FP16-rounded inputs.""" + q64 = q.double() + k64 = k.double() + v64 = v.double() + group = q.size(1) // k.size(1) + k64 = k64.repeat_interleave(group, dim=1) + v64 = v64.repeat_interleave(group, dim=1) + scores = torch.matmul(q64, k64.transpose(-1, -2)) * scale + if causal: + sq, skv = q.size(2), k.size(2) + offset = skv - sq + mask = torch.arange(skv, device=q.device)[None, :] > ( + torch.arange(sq, device=q.device)[:, None] + offset + ) + scores = scores.masked_fill(mask, float("-inf")) + lse = torch.logsumexp(scores, dim=-1) + probs = torch.softmax(scores, dim=-1) + return torch.matmul(probs, v64), lse + + +def _sdpa_forward(q, k, v, *, causal, scale): + group = q.size(1) // k.size(1) + return torch.nn.functional.scaled_dot_product_attention( + q, + k.repeat_interleave(group, dim=1), + v.repeat_interleave(group, dim=1), + is_causal=causal, + scale=scale, + ) + + +class _Paths: + """Lazily constructed attention paths, so a missing backend skips one row.""" + + def __init__(self, device: torch.device) -> None: + self.device = device + self.errors: dict[str, str] = {} + # The PyTorch reference is the only non-SDPA path that also runs on the host. + self.native = self._try("pytorch-native", self._make_native) + self.is_rocm = torch.version.hip is not None + if device.type == "cuda": + if self.is_rocm: + self.strict = self._try("strict-aiter", self._make_strict) + self.strict_fa4 = None + self.errors["strict-fa4"] = "CUDA-only path; this run is ROCm" + else: + self.strict = None + self.errors["strict-aiter"] = "ROCm-only path; this run is CUDA" + self.strict_fa4 = self._try("strict-fa4", self._make_strict_fa4) + self.reference = self._try("reference-native", self._make_reference) + self.triton = self._try("triton-bitwise", self._make_triton) + else: + self.strict = self.strict_fa4 = self.reference = self.triton = None + for name in ("strict-aiter", "strict-fa4", "reference-native", "triton-bitwise"): + self.errors[name] = "GPU-only path; not available on the host" + + def _try(self, name: str, factory: Callable[[], Any]) -> Any: + try: + return factory() + except Exception as exc: # noqa: BLE001 - a missing backend is a reported row + self.errors[name] = f"{type(exc).__name__}: {exc}" + return None + + @staticmethod + def _make_native(): + from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp + + return NativeAttentionOp() + + @staticmethod + def _make_strict(): + from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore + + return StrictRocmAiterCKAttentionCore() + + @staticmethod + def _make_strict_fa4(): + from rl_engine.kernels.ops.cuda.attention.flash_attn import StrictFlashAttention4Core + + return StrictFlashAttention4Core() + + @staticmethod + def _make_reference(): + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionOp, + ) + + return DeterministicAttentionOp() + + def _make_triton(self): + from rl_engine.kernels.ops.triton.attention.deterministic_attn import ( + BITWISE_LIBM_PARITY, + TritonDeterministicAttentionOp, + ) + + # The bitwise expf/logf sequence is only ported for HIP, so on CUDA the op + # refuses by default. Measure it anyway, but the report must not call it bitwise. + return TritonDeterministicAttentionOp(require_bitwise_libm=BITWISE_LIBM_PARITY) + + def runner(self, name: str, q, k, v, *, causal: bool, scale: float, positions): + """Return ``() -> (out, lse|None)`` for one path, or None when unavailable.""" + if name == "sdpa": + return lambda: (_sdpa_forward(q, k, v, causal=causal, scale=scale), None) + if name == "pytorch-native" and self.native is not None: + # NativeAttentionOp is the repo's ground-truth reference; it returns out only. + return lambda: ( + self.native.forward(q, k, v, causal=causal, scale=scale), + None, + ) + if name == "strict-aiter" and self.strict is not None: + + def run_strict(): + result = self.strict.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + query_position_ids=positions if causal else None, + key_position_ids=positions if causal else None, + ) + return result.out, result.lse + + return run_strict + if name == "strict-fa4" and self.strict_fa4 is not None: + + def run_fa4(): + result = self.strict_fa4.forward_with_lse( + q, + k, + v, + causal=causal, + scale=scale, + query_position_ids=positions if causal else None, + key_position_ids=positions if causal else None, + ) + return result.out, result.lse + + return run_fa4 + if name == "reference-native" and self.reference is not None: + return lambda: self.reference.forward_with_lse(q, k, v, causal=causal, scale=scale) + if name == "triton-bitwise" and self.triton is not None: + return lambda: self.triton.forward_with_lse(q, k, v, causal=causal, scale=scale) + return None + + +PATH_NAMES = ( + "sdpa", + "pytorch-native", + "strict-aiter", + "strict-fa4", + "reference-native", + "triton-bitwise", +) + + +def _single_gpu_benchmarks( + *, + paths: _Paths, + seq_lens: tuple[int, ...], + dtypes: tuple[torch.dtype, ...], + q_heads: int, + kv_heads: int, + head_dim: int, + batch: int, + warmup: int, + samples: int, + training_samples: int, + device: torch.device, + budget_seconds: float = 0.0, +) -> list[dict[str, Any]]: + cases: list[dict[str, Any]] = [] + scale = 1.0 / math.sqrt(head_dim) + + for dtype in dtypes: + dtype_name = str(dtype).replace("torch.", "").replace("float", "fp").replace("bfp", "bf") + for seq_len in seq_lens: + q, k, v = _seeded_qkv( + batch, q_heads, kv_heads, seq_len, head_dim, dtype, device, seed=1234 + ) + positions = _positions(batch, seq_len, device) + oracle_out, oracle_lse = _fp64_oracle(q, k, v, causal=True, scale=scale) + + row: dict[str, Any] = { + "dtype": dtype_name, + "seq_len": seq_len, + "batch": batch, + "q_heads": q_heads, + "kv_heads": kv_heads, + "head_dim": head_dim, + "paths": {}, + } + captured: dict[str, tuple[torch.Tensor, torch.Tensor | None]] = {} + + for name in PATH_NAMES: + runner = paths.runner(name, q, k, v, causal=True, scale=scale, positions=positions) + if runner is None: + continue + + # One untimed call both captures the outputs and prices the path. A cell + # whose sampling would blow the budget is skipped and says so, rather + # than silently costing an hour. + probe_start = time.perf_counter() + out, lse = runner() + _device_sync(device) + probe_seconds = time.perf_counter() - probe_start + captured[name] = ( + out.detach().clone(), + None if lse is None else lse.detach().clone(), + ) + + forward_calls = warmup + samples + 2 + training_calls = max(1, warmup // 2) + training_samples + 1 + # Backward is empirically 2-4x the forward on these paths; 3x is the + # midpoint and only decides whether to run, never a reported number. + estimated = probe_seconds * (forward_calls + 3 * training_calls) + if budget_seconds > 0 and estimated > budget_seconds: + row["paths"][name] = { + "skipped": ( + f"one call took {probe_seconds:.1f}s; sampling would need about " + f"{estimated / 60:.0f} min, over the " + f"{budget_seconds / 60:.0f} min per-path budget" + ), + "probe_seconds": probe_seconds, + "estimated_seconds": estimated, + "out_vs_fp64": _accuracy(out.double(), oracle_out), + } + del out, lse + _empty_cache(device) + continue + + deadline = time.perf_counter() + budget_seconds if budget_seconds else 0.0 + forward_ms = _timed_samples( + runner, warmup=warmup, samples=samples, device=device, deadline=deadline + ) + forward = _summary_ms(forward_ms) + forward_peak = _peak_memory_mib(runner, device) + + entry: dict[str, Any] = { + "forward": forward, + "forward_samples": len(forward_ms), + "forward_truncated": len(forward_ms) < samples, + "forward_peak_mib": forward_peak, + "out_vs_fp64": _accuracy(out.double(), oracle_out), + } + if lse is not None: + entry["lse_vs_fp64"] = _accuracy(lse.double(), oracle_lse) + + # Repeat determinism: two identical calls must be bitwise equal. + repeat_out, repeat_lse = runner() + entry["repeat_bitwise"] = _bitwise_equal(out, repeat_out) and ( + lse is None or _bitwise_equal(lse, repeat_lse) + ) + + training = _training_runner( + paths, name, q, k, v, causal=True, scale=scale, positions=positions + ) + if training is not None: + train_deadline = time.perf_counter() + budget_seconds if budget_seconds else 0.0 + train_ms = _timed_samples( + training, + warmup=max(1, warmup // 2), + samples=training_samples, + device=device, + deadline=train_deadline, + ) + entry["train_fwd_bwd"] = _summary_ms(train_ms) + entry["train_samples"] = len(train_ms) + entry["train_truncated"] = len(train_ms) < training_samples + entry["train_peak_mib"] = _peak_memory_mib(training, device) + + row["paths"][name] = entry + del out, lse, repeat_out, repeat_lse + _empty_cache(device) + + # Headline: Triton must be bit-identical to the native reference core. + if "triton-bitwise" in captured and "reference-native" in captured: + t_out, t_lse = captured["triton-bitwise"] + r_out, r_lse = captured["reference-native"] + row["triton_vs_reference"] = { + "out_mismatched": _mismatch_count(t_out, r_out), + "lse_mismatched": _mismatch_count(t_lse, r_lse), + "out_relative_l2": _relative_l2(t_out, r_out), + "bitwise": _bitwise_equal(t_out, r_out) and _bitwise_equal(t_lse, r_lse), + } + # The production core is a different vendor kernel; report the gap, do + # not claim parity with it. + production = "strict-aiter" if "strict-aiter" in captured else "strict-fa4" + if production in captured and "reference-native" in captured: + s_out, s_lse = captured[production] + r_out, r_lse = captured["reference-native"] + row["strict_vs_reference"] = { + "production_path": production, + "out": _accuracy(s_out, r_out), + "lse": _accuracy(s_lse, r_lse), + "out_mismatched": _mismatch_count(s_out, r_out), + } + + cases.append(row) + del q, k, v, oracle_out, oracle_lse, captured + _empty_cache(device) + return cases + + +def _training_runner(paths, name, q, k, v, *, causal, scale, positions): + """Return ``() -> None`` running one forward+backward, or None.""" + if name == "sdpa": + + def train_sdpa() -> None: + qr = q.detach().requires_grad_(True) + kr = k.detach().requires_grad_(True) + vr = v.detach().requires_grad_(True) + out = _sdpa_forward(qr, kr, vr, causal=causal, scale=scale) + out.sum().backward() + + return train_sdpa + + if name == "pytorch-native": + if paths.native is None: + return None + + def train_native() -> None: + qr = q.detach().requires_grad_(True) + kr = k.detach().requires_grad_(True) + vr = v.detach().requires_grad_(True) + out = paths.native.forward(qr, kr, vr, causal=causal, scale=scale) + out.sum().backward() + + return train_native + + op = { + "strict-aiter": paths.strict, + "strict-fa4": paths.strict_fa4, + "reference-native": paths.reference, + "triton-bitwise": paths.triton, + }.get(name) + if op is None: + return None + + def train_op() -> None: + qr = q.detach().requires_grad_(True) + kr = k.detach().requires_grad_(True) + vr = v.detach().requires_grad_(True) + kwargs: dict[str, Any] = {"causal": causal, "scale": scale} + if name in ("strict-aiter", "strict-fa4") and causal: + kwargs["query_position_ids"] = positions + kwargs["key_position_ids"] = positions + result = op.forward_with_lse(qr, kr, vr, **kwargs) + out = result.out if hasattr(result, "out") else result[0] + out.sum().backward() + + return train_op + + +def _backward_parity( + *, + paths: _Paths, + seq_lens: tuple[int, ...], + q_heads: int, + kv_heads: int, + head_dim: int, + batch: int, + device: torch.device, +) -> list[dict[str, Any]]: + """dQ/dK/dV bitwise parity, Triton port versus the native reference core.""" + if paths.reference is None or paths.triton is None: + return [] + rows = [] + scale = 1.0 / math.sqrt(head_dim) + for seq_len in seq_lens: + q, k, v = _seeded_qkv( + batch, q_heads, kv_heads, seq_len, head_dim, torch.bfloat16, device, seed=99 + ) + grad_out = torch.randn( + batch, + q_heads, + seq_len, + head_dim, + device=device, + dtype=torch.bfloat16, + generator=torch.Generator(device=device).manual_seed(100), + ) + grads = {} + for name, op in (("reference-native", paths.reference), ("triton-bitwise", paths.triton)): + qr = q.detach().requires_grad_(True) + kr = k.detach().requires_grad_(True) + vr = v.detach().requires_grad_(True) + out, _lse = op.forward_with_lse(qr, kr, vr, causal=True, scale=scale) + out.backward(grad_out) + grads[name] = (qr.grad.clone(), kr.grad.clone(), vr.grad.clone()) + reference = grads["reference-native"] + triton_grads = grads["triton-bitwise"] + rows.append( + { + "seq_len": seq_len, + "dq_mismatched": _mismatch_count(triton_grads[0], reference[0]), + "dk_mismatched": _mismatch_count(triton_grads[1], reference[1]), + "dv_mismatched": _mismatch_count(triton_grads[2], reference[2]), + "bitwise": all(_bitwise_equal(t, r) for t, r in zip(triton_grads, reference)), + } + ) + del q, k, v, grad_out, grads + _empty_cache(device) + return rows + + +def _batch_composition( + *, + paths: _Paths, + seq_lens: tuple[int, ...], + q_heads: int, + kv_heads: int, + head_dim: int, + device: torch.device, +) -> list[dict[str, Any]]: + """A row computed alone must be bitwise equal to the same row inside a batch. + + The strict ROCm core refuses ``B > 1`` outright (``_validate_inputs``: "executes + one logical batch row at a time"), so for that path the property is structural + rather than measured, and the row records that instead of a comparison. + """ + rows = [] + scale = 1.0 / math.sqrt(head_dim) + for seq_len in seq_lens: + q, k, v = _seeded_qkv( + 4, q_heads, kv_heads, seq_len, head_dim, torch.bfloat16, device, seed=7 + ) + positions_batch = _positions(4, seq_len, device) + positions_one = _positions(1, seq_len, device) + row: dict[str, Any] = {"seq_len": seq_len, "paths": {}} + for name in PATH_NAMES: + if name == "strict-aiter": + if paths.strict is not None: + row["paths"][name] = { + "batch_gt1_rejected": True, + "out_bitwise": True, + "out_mismatched": 0, + "note": "core executes one logical batch row per launch", + } + continue + batched = paths.runner( + name, q, k, v, causal=True, scale=scale, positions=positions_batch + ) + single = paths.runner( + name, + q[2:3].contiguous(), + k[2:3].contiguous(), + v[2:3].contiguous(), + causal=True, + scale=scale, + positions=positions_one, + ) + if batched is None or single is None: + continue + batch_out, batch_lse = batched() + single_out, single_lse = single() + row["paths"][name] = { + "batch_gt1_rejected": False, + "out_bitwise": _bitwise_equal(single_out[0], batch_out[2].contiguous()), + "out_mismatched": _mismatch_count(single_out[0], batch_out[2].contiguous()), + "out_max_abs": _accuracy(single_out[0], batch_out[2])["max_abs"], + "lse_bitwise": ( + None + if batch_lse is None + else _bitwise_equal(single_lse[0], batch_lse[2].contiguous()) + ), + } + del batch_out, batch_lse, single_out, single_lse + _empty_cache(device) + rows.append(row) + del q, k, v + _empty_cache(device) + return rows + + +def _tp_head_sensitivity( + *, + paths: _Paths, + seq_lens: tuple[int, ...], + tp_degrees: tuple[int, ...], + q_heads: int, + kv_heads: int, + head_dim: int, + device: torch.device, +) -> list[dict[str, Any]]: + """Is a head shard under TP=N bitwise equal to the same slice of an unsharded run? + + TP performs no cross-rank reduction in attention, so any nonzero value here + means the kernel's arithmetic depends on how many heads shared the launch. + Measured both on the raw production core and through the per-KV-group launch + schedule that the Vime provider uses. + """ + rows: list[dict[str, Any]] = [] + scale = 1.0 / math.sqrt(head_dim) + for seq_len in seq_lens: + q, k, v = _seeded_qkv( + 1, q_heads, kv_heads, seq_len, head_dim, torch.bfloat16, device, seed=21 + ) + positions = _positions(1, seq_len, device) + + for schedule in ("raw_launch", "one_kv_group_per_launch"): + full = _tp_schedule_forward(paths, q, k, v, scale, positions, schedule) + if full is None: + continue + full_out, full_lse = full + for tp in tp_degrees: + if q_heads % tp or kv_heads % tp: + continue + local_q = q_heads // tp + local_kv = kv_heads // tp + shard = _tp_schedule_forward( + paths, + q[:, :local_q], + k[:, :local_kv], + v[:, :local_kv], + scale, + positions, + schedule, + ) + shard_out, shard_lse = shard + rows.append( + { + "seq_len": seq_len, + "schedule": schedule, + "tp": tp, + "local_q_heads": local_q, + "local_kv_heads": local_kv, + "out_max_abs": _accuracy(shard_out, full_out[:, :local_q])["max_abs"], + "lse_max_abs": _accuracy(shard_lse, full_lse[:, :local_q])["max_abs"], + "invariant": _bitwise_equal(shard_out, full_out[:, :local_q].contiguous()) + and _bitwise_equal(shard_lse, full_lse[:, :local_q].contiguous()), + } + ) + del shard_out, shard_lse + _empty_cache(device) + del full_out, full_lse + _empty_cache(device) + del q, k, v + _empty_cache(device) + return rows + + +def _tp_schedule_cost( + *, + paths: _Paths, + seq_lens: tuple[int, ...], + q_heads: int, + kv_heads: int, + head_dim: int, + device: torch.device, + warmup: int, + samples: int, +) -> list[dict[str, Any]]: + """What the TP-degree invariance costs. + + ``raw_launch`` is one launch for all heads and is NOT the production schedule; + ``one_kv_group_per_launch`` is what the provider runs (``Hkv`` launches per row) + and is what makes the result independent of the TP degree. + """ + if paths.strict is None and paths.strict_fa4 is None: + return [] + rows: list[dict[str, Any]] = [] + scale = 1.0 / math.sqrt(head_dim) + for seq_len in seq_lens: + q, k, v = _seeded_qkv( + 1, q_heads, kv_heads, seq_len, head_dim, torch.bfloat16, device, seed=21 + ) + positions = _positions(1, seq_len, device) + entry: dict[str, Any] = {"seq_len": seq_len, "launches": kv_heads} + for schedule in ("raw_launch", "one_kv_group_per_launch"): + + # Bind the tensors as defaults: they are deleted at the end of each + # iteration, so a late-binding closure would reference a dead name. + def run(chosen=schedule, q=q, k=k, v=v, positions=positions): + return _tp_schedule_forward(paths, q, k, v, scale, positions, chosen) + + entry[schedule] = _summary_ms( + _timed_samples(run, warmup=warmup, samples=samples, device=device) + ) + entry[f"{schedule}_peak_mib"] = _peak_memory_mib(run, device) + + def run_sdpa(q=q, k=k, v=v): + return _sdpa_forward(q, k, v, causal=True, scale=scale) + + entry["sdpa"] = _summary_ms( + _timed_samples(run_sdpa, warmup=warmup, samples=samples, device=device) + ) + rows.append(entry) + del q, k, v + _empty_cache(device) + return rows + + +def _tp_schedule_forward(paths, q, k, v, scale, positions, schedule): + """Run the strict core either in one launch or one launch per KV group.""" + core = paths.strict if paths.strict is not None else paths.strict_fa4 + if core is None: + return None + if schedule == "raw_launch": + result = core.forward_with_lse( + q, + k, + v, + causal=True, + scale=scale, + query_position_ids=positions, + key_position_ids=positions, + ) + return result.out.contiguous(), result.lse.contiguous() + + group = q.size(1) // k.size(1) + outs, lses = [], [] + for kv_index in range(k.size(1)): + lo, hi = kv_index * group, (kv_index + 1) * group + result = core.forward_with_lse( + q[:, lo:hi], + k[:, kv_index : kv_index + 1], + v[:, kv_index : kv_index + 1], + causal=True, + scale=scale, + query_position_ids=positions, + key_position_ids=positions, + ) + outs.append(result.out) + lses.append(result.lse) + return torch.cat(outs, dim=1).contiguous(), torch.cat(lses, dim=1).contiguous() + + +# --------------------------------------------------------------------------- +# Distributed CP (spawned world; harness shape from PR #325) +# --------------------------------------------------------------------------- + + +def _distributed_cp_worker( + rank: int, + world_size: int, + topology: tuple[str, int, int, int], + init_method: str, + result_queue: Any, + warmup: int, + samples: int, + seq_len: int, + q_heads: int, + kv_heads: int, + head_dim: int, +) -> None: + try: + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + device_id=torch.device("cuda", rank), + timeout=timedelta(minutes=20), + ) + payload = _distributed_cp_case( + rank, + world_size, + topology, + warmup=warmup, + samples=samples, + seq_len=seq_len, + q_heads=q_heads, + kv_heads=kv_heads, + head_dim=head_dim, + ) + if rank == 0: + result_queue.put({"ok": True, "topology": topology[0], "payload": payload}) + except Exception: + result_queue.put( + { + "ok": False, + "rank": rank, + "topology": topology[0], + "traceback": traceback.format_exc(), + } + ) + raise + finally: + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _distributed_cp_case( + rank: int, + world_size: int, + topology: tuple[str, int, int], + *, + warmup: int, + samples: int, + seq_len: int, + q_heads: int, + kv_heads: int, + head_dim: int, +) -> dict[str, Any]: + """One CP topology running the real AG/RS schedule over the platform's transport. + + Schedule (same one ``scripts/ws2_p2p_nccl_attention_reference_check.py`` accepts): + all-gather Q/K/V and the position ids over the CP group, run the strict core + once on the full sequence, then reduce-scatter the ``(out, lse)`` result back + to this rank's query range. Bitwise acceptance is against a CP=1 run of the + same core on the same full-sequence inputs. + """ + from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionParallelSpec, + CUDAAGRSAttentionCPCommunication, + RCCLAGRSAttentionCPCommunication, + ) + + # ROCm runs AITER over the RCCL transport; CUDA runs FA4 over the CUDA IPC transport. + is_rocm = torch.version.hip is not None + if is_rocm: + from rl_engine.kernels.ops.rocm.attention.flash_attn import ( + StrictRocmAiterCKAttentionCore as _StrictCore, + ) + + _Transport = RCCLAGRSAttentionCPCommunication + transport_id = "rccl_ag_rs" + else: + from rl_engine.kernels.ops.cuda.attention.flash_attn import ( + StrictFlashAttention4Core as _StrictCore, + ) + + _Transport = CUDAAGRSAttentionCPCommunication + transport_id = "cuda_ag_rs" + + label, tp_world, cp_world, replicas = topology + device = torch.device("cuda", rank) + scale = 1.0 / math.sqrt(head_dim) + chunk_size = seq_len // (cp_world * 2) + + # Ranks that share a TP index form one CP group. Every rank must call + # new_group for every group, in the same order. + group_index = rank // cp_world + tp_index = group_index % tp_world + replica_index = group_index // tp_world + cp_rank = rank % cp_world + cp_group = None + for slice_index in range(world_size // cp_world): + ranks = list(range(slice_index * cp_world, (slice_index + 1) * cp_world)) + group = dist.new_group(ranks=ranks) + if slice_index == group_index: + cp_group = group + + local_q_heads = q_heads // tp_world + local_kv_heads = kv_heads // tp_world + + generator = torch.Generator(device="cpu").manual_seed(2357 + tp_index + 100 * replica_index) + q = torch.randn( + 1, local_q_heads, seq_len, head_dim, generator=generator, dtype=torch.bfloat16 + ).to(device) + k = torch.randn( + 1, local_kv_heads, seq_len, head_dim, generator=generator, dtype=torch.bfloat16 + ).to(device) + v = torch.randn( + 1, local_kv_heads, seq_len, head_dim, generator=generator, dtype=torch.bfloat16 + ).to(device) + positions = _positions(1, seq_len, device).to(torch.int32) + + span = seq_len // cp_world + owner_ranges = tuple((i * span, (i + 1) * span) for i in range(cp_world)) + blocks: list[AttentionCPBlockMetadata] = [] + for owner, (owner_start, owner_end) in enumerate(owner_ranges): + for start in range(owner_start, owner_end, chunk_size): + blocks.append( + AttentionCPBlockMetadata( + global_block_index=len(blocks), + kv_block_start=start, + kv_block_end=min(start + chunk_size, owner_end), + owner_cp_rank=owner, + owner_tp_rank=tp_index, + ) + ) + plan = AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec( + tp_world_size=tp_world, + tp_rank=tp_index, + cp_world_size=cp_world, + cp_rank=cp_rank, + ), + backend=transport_id, + status="implemented", + expected_blocks=tuple(blocks), + expected_kv_token_range=(0, seq_len), + query_token_ranges=owner_ranges, + ) + + core = _StrictCore() + communication = _Transport(process_group=cp_group) + + query_start, query_end = owner_ranges[cp_rank] + q_local = q[:, :, query_start:query_end, :].contiguous() + k_local = k[:, :, query_start:query_end, :].contiguous() + v_local = v[:, :, query_start:query_end, :].contiguous() + positions_local = positions[:, query_start:query_end].contiguous() + + def cp_forward(): + q_full = communication.all_gather_query(q_local, plan) + k_full, v_full = communication.all_gather_kv(k_local, v_local, plan) + query_positions, key_positions = communication.all_gather_position_ids( + positions_local, positions_local, plan + ) + result = core.forward_with_lse( + q_full, + k_full, + v_full, + causal=True, + scale=scale, + query_position_ids=query_positions, + key_position_ids=key_positions, + ) + return communication.reduce_scatter_strict_result(result.out, result.lse, plan) + + shard = cp_forward() + forward_ms = _summary_ms( + _gpu_event_samples(lambda: cp_forward(), warmup=warmup, samples=samples) + ) + peak = _peak_memory_mib(lambda: cp_forward(), device) + + # CP=1 acceptance: the same core on the same full-sequence inputs, then take + # this rank's query range out of it. + def cp1_forward(): + return core.forward_with_lse( + q, + k, + v, + causal=True, + scale=scale, + query_position_ids=positions, + key_position_ids=positions, + ) + + single = cp1_forward() + cp1_ms = _summary_ms(_gpu_event_samples(cp1_forward, warmup=warmup, samples=samples)) + expected_out = single.out[:, :, query_start:query_end, :].contiguous() + expected_lse = single.lse[:, :, query_start:query_end].contiguous() + + out_bitwise = _bitwise_equal(shard.out.contiguous(), expected_out) + lse_bitwise = _bitwise_equal(shard.lse.contiguous(), expected_lse) + repeat = cp_forward() + repeat_bitwise = _bitwise_equal(shard.out.contiguous(), repeat.out.contiguous()) + + flags = torch.tensor( + [ + 1.0 if out_bitwise else 0.0, + 1.0 if lse_bitwise else 0.0, + 1.0 if repeat_bitwise else 0.0, + ], + device=device, + ) + dist.all_reduce(flags, op=dist.ReduceOp.MIN) + mismatches = torch.tensor( + [ + float(_mismatch_count(shard.out.contiguous(), expected_out)), + float(_mismatch_count(shard.lse.contiguous(), expected_lse)), + ], + device=device, + ) + dist.all_reduce(mismatches, op=dist.ReduceOp.SUM) + + return { + "topology": label, + "world_size": world_size, + "tp_world_size": tp_world, + "cp_world_size": cp_world, + "replicas": replicas, + "seq_len": seq_len, + "local_q_heads": local_q_heads, + "local_kv_heads": local_kv_heads, + "transport": transport_id, + "forward": forward_ms, + "cp1_baseline": cp1_ms, + "peak_mib_per_rank": peak, + "out_bitwise_vs_cp1": bool(flags[0].item() == 1.0), + "lse_bitwise_vs_cp1": bool(flags[1].item() == 1.0), + "repeat_bitwise": bool(flags[2].item() == 1.0), + "out_mismatched_all_ranks": int(mismatches[0].item()), + "lse_mismatched_all_ranks": int(mismatches[1].item()), + } + + +def _run_distributed_topology( + topology: tuple[str, int, int, int], + *, + warmup: int, + samples: int, + seq_len: int, + q_heads: int, + kv_heads: int, + head_dim: int, +) -> dict[str, Any]: + _label, tp_world, cp_world, replicas = topology + world_size = tp_world * cp_world * replicas + context = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as temporary_directory: + init_method = (Path(temporary_directory) / "rccl_init").as_uri() + result_queue = context.Queue() + processes = [ + context.Process( + target=_distributed_cp_worker, + args=( + rank, + world_size, + topology, + init_method, + result_queue, + warmup, + samples, + seq_len, + q_heads, + kv_heads, + head_dim, + ), + ) + for rank in range(world_size) + ] + for process in processes: + process.start() + result = None + try: + result = result_queue.get(timeout=1800) + except queue.Empty as exc: + for process in processes: + if process.is_alive(): + process.terminate() + raise RuntimeError(f"timed out waiting for {topology[0]}") from exc + finally: + for process in processes: + process.join(timeout=90) + if process.is_alive(): + process.terminate() + process.join(timeout=30) + result_queue.close() + result_queue.join_thread() + if result is None or not result["ok"]: + raise RuntimeError((result or {}).get("traceback", f"{topology[0]} returned no result")) + return result["payload"] + + +# --------------------------------------------------------------------------- +# Environment, report and figures +# --------------------------------------------------------------------------- + + +def _default_platform_label(device: torch.device) -> str: + if device.type != "cuda": + return "cpu" + name = torch.cuda.get_device_name(0).lower() + if "mi300" in name: + return "mi300x" + if "h100" in name: + return "h100" + return name.replace(" ", "-")[:24] + + +def _load_comparisons(specs: list[str]) -> list[tuple[str, dict[str, Any]]]: + """Parse ``--compare-with LABEL=PATH`` into (label, payload) pairs.""" + loaded: list[tuple[str, dict[str, Any]]] = [] + for spec in specs: + if "=" not in spec: + raise SystemExit(f"--compare-with expects LABEL=PATH, got {spec!r}") + label, _, path = spec.partition("=") + loaded.append((label, json.loads(Path(path).read_text()))) + return loaded + + +def _environment(device: torch.device | None = None) -> dict[str, Any]: + on_gpu = device is None or device.type == "cuda" + properties = ( + torch.cuda.get_device_properties(0) if on_gpu and torch.cuda.is_available() else None + ) + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + symbols = sorted(name for name in dir(_C) if "attention" in name) if _EXT_AVAILABLE else [] + except Exception: # noqa: BLE001 + symbols = [] + try: + import triton + + triton_version = triton.__version__ + except Exception: # noqa: BLE001 + triton_version = "unavailable" + # Device facts must not leak into a host run's column: a CPU row reporting + # gpu_count=8 and an RCCL collective would misdescribe what was measured. + return { + "cpu_count": os.cpu_count(), + "torch_threads": torch.get_num_threads(), + "gpu": properties.name if properties else "n/a (host execution)", + "architecture": getattr(properties, "gcnArchName", "unknown") if properties else "n/a", + "gpu_count": (torch.cuda.device_count() if on_gpu and torch.cuda.is_available() else 0), + "hip": torch.version.hip if on_gpu else None, + "cuda": torch.version.cuda if on_gpu else None, + "torch": torch.__version__, + "triton": triton_version if on_gpu else "n/a (host execution)", + "python": platform.python_version(), + "extension_attention_symbols": symbols if on_gpu else [], + "native_collective": ( + "torch.distributed ProcessGroupNCCL (RCCL on ROCm)" + if on_gpu + else "n/a (single-process host run)" + ), + } + + +def _case_for(payload: dict[str, Any], dtype: str, seq_len: int) -> dict[str, Any] | None: + for case in payload.get("single_gpu", {}).get("cases", []): + if case["dtype"] == dtype and case["seq_len"] == seq_len: + return case + return None + + +def _fmt(value: Any, spec: str = ".4f") -> str: + if value is None: + return "n/a" + if isinstance(value, bool): + return "yes" if value else "**no**" + if isinstance(value, (int,)) and not isinstance(value, bool): + return str(value) + try: + return format(float(value), spec) + except (TypeError, ValueError): + return str(value) + + +def _write_report( + payload: dict[str, Any], + output_directory: Path, + comparisons: list[tuple[str, dict[str, Any]]] | None = None, +) -> None: + platforms: list[tuple[str, dict[str, Any]]] = [ + (payload.get("platform_label", "this run"), payload) + ] + list(comparisons or []) + configuration = payload["configuration"] + lines: list[str] = [] + add = lines.append + + add("# WS2 strict ROCm Attention — bitwise parity and performance") + add("") + add("> Operator-only benchmark. No model checkpoint or serving engine was used;") + add("> the shapes are Qwen3-8B's attention shapes.") + add("") + add("## Environment") + add("") + keys = sorted({k for _, pl in platforms for k in pl["environment"]}) + add("| Item | " + " | ".join(label for label, _ in platforms) + " |") + add("|---|" + "---|" * len(platforms)) + for key in keys: + cells = [] + for _, pl in platforms: + value = pl["environment"].get(key, "n/a") + if isinstance(value, list): + value = ", ".join(value) or "none" + cells.append(str(value)) + add(f"| {key} | " + " | ".join(cells) + " |") + add("") + if len(platforms) > 1: + add( + "A missing row below means the backend cannot exist on that platform, not that it " + "failed: `strict-aiter` is ROCm-only, `reference-native` and `triton-bitwise` need a " + "GPU, and only `sdpa` and `pytorch-native` also run on the host." + ) + add("") + + add("## Methodology") + add("") + add( + f"- Operator shape: `Hq={configuration['q_heads']}`, `Hkv={configuration['kv_heads']}`, " + f"`D={configuration['head_dim']}`, `B={configuration['batch']}`, causal; sequence sweep " + + ", ".join(str(s) for s in configuration["seq_lens"]) + + "." + ) + add("- Measured paths:") + add( + " - `sdpa`: `torch.nn.functional.scaled_dot_product_attention`. **Speed baseline only** — " + "as in PR #325, no accuracy comparison is mixed into the speed table." + ) + add( + " - `strict-aiter`: `StrictRocmAiterCKAttentionCore` called **once for all heads**. " + "This is the core, not the production schedule: the Vime provider launches it once " + "per (batch row, KV group). See the per-KV-group schedule table for that cost." + ) + add( + " - `reference-native`: `_C.deterministic_attention_forward/backward`, the materializing " + "FP32 reference core hipified from the shared `.cu`." + ) + add( + " - `triton-bitwise`: `TritonDeterministicAttentionOp`, whose contract is bit-identity " + "with `reference-native`." + ) + add( + "- Timing: CUDA events, median and p95. Peak memory is the per-call increase in " + "`torch.cuda.max_memory_allocated` above what was live before the call." + ) + add( + "- Accuracy is against an FP64 oracle over the same BF16/FP16-rounded inputs. " + "Repeat = two identical calls are bitwise equal; batch-invariant = a row computed " + "alone is bitwise equal to the same row inside a batch." + ) + add( + f"- {configuration['warmup']} warmups, {configuration['samples']} measured forward " + f"samples, {configuration['training_samples']} measured forward+backward samples. " + "Raw medians, p95, min and max are in `results.json`." + ) + add("") + add("Reproduce from the repository root:") + add("") + add("```bash") + add("python benchmarks/benchmark_ws2_rocm_attention.py \\") + add(f" --seq-lens {','.join(str(s) for s in configuration['seq_lens'])} \\") + add(f" --dtypes {','.join(configuration['dtypes'])} \\") + add(f" --warmup {configuration['warmup']} --samples {configuration['samples']} \\") + add(f" --training-samples {configuration['training_samples']} \\") + add(" --output-dir benchmarks/results/ws2_rocm_mi300x") + add("```") + add("") + + if payload.get("unavailable_paths"): + add("### Unavailable paths") + add("") + for name, reason in payload["unavailable_paths"].items(): + add(f"- `{name}`: {reason}") + add("") + + # ---- headline + add("## Bitwise parity: Triton port vs the native reference core") + add("") + add("Acceptance is 0 mismatched elements. This is the contract the Triton core exists to hold.") + add("") + add("| dtype | S | out mismatched | lse mismatched | dQ | dK | dV | bitwise |") + add("|---|---:|---:|---:|---:|---:|---:|:---:|") + backward = {row["seq_len"]: row for row in payload.get("backward_parity", [])} + for case in payload["single_gpu"]["cases"]: + parity = case.get("triton_vs_reference") + if not parity: + continue + grads = backward.get(case["seq_len"], {}) if case["dtype"] == "bf16" else {} + add( + f"| {case['dtype']} | {case['seq_len']} | {parity['out_mismatched']} | " + f"{parity['lse_mismatched']} | {_fmt(grads.get('dq_mismatched'))} | " + f"{_fmt(grads.get('dk_mismatched'))} | {_fmt(grads.get('dv_mismatched'))} | " + f"{_fmt(parity['bitwise'])} |" + ) + add("") + add("`dQ/dK/dV` are measured on the BF16 sweep only; `n/a` marks the FP16 rows.") + add("") + + skipped_rows = [ + (label, case["dtype"], case["seq_len"], name, entry["skipped"]) + for label, pl in platforms + for case in pl.get("single_gpu", {}).get("cases", []) + for name, entry in case["paths"].items() + if isinstance(entry, dict) and "skipped" in entry + ] + if skipped_rows: + add("## Skipped cells") + add("") + add( + "A cell whose single call implies more sampling time than the per-path budget " + "is not measured. The observed single-call cost is reported instead, which is " + "the useful part: it is what made the cell unaffordable." + ) + add("") + add("| Platform | dtype | S | Path | Why |") + add("|---|---|---:|---|---|") + for label, dtype, seq_len, name, why in skipped_rows: + add(f"| {label} | {dtype} | {seq_len} | {name} | {why} |") + add("") + + # ---- single GPU speed + multi = len(platforms) > 1 + platform_column = "Platform | " if multi else "" + platform_rule = "---|" if multi else "" + for dtype in configuration["dtypes"]: + seq_lens = sorted( + { + c["seq_len"] + for _, pl in platforms + for c in pl.get("single_gpu", {}).get("cases", []) + if c["dtype"] == dtype + } + ) + if not seq_lens: + continue + add(f"## Single-device Attention ({dtype})") + add("") + add("### Forward") + add("") + add( + f"| S | {platform_column}Path | Median (ms) | p95 (ms) | vs sdpa | Peak MiB | " + "out max-abs vs FP64 | lse max-abs vs FP64 | Repeat |" + ) + add(f"|---:|{platform_rule}---|---:|---:|---:|---:|---:|---:|:---:|") + for seq_len in seq_lens: + for label, pl in platforms: + case = _case_for(pl, dtype, seq_len) + if case is None: + continue + baseline = case["paths"].get("sdpa", {}).get("forward", {}).get("median_ms") + for name in PATH_NAMES: + entry = case["paths"].get(name) + if not entry: + continue + if "skipped" in entry: + prefix = f"{label} | " if multi else "" + add( + f"| {seq_len} | {prefix}{name} | skipped | — | — | — | " + f"{entry['out_vs_fp64']['max_abs']:.3e} | — | — |" + ) + continue + median = entry["forward"]["median_ms"] + ratio = f"{median / baseline:.2f}x" if baseline else "n/a" + prefix = f"{label} | " if multi else "" + add( + f"| {seq_len} | {prefix}{name} | {median:.4f} | " + f"{entry['forward']['p95_ms']:.4f} | {ratio} | " + f"{entry['forward_peak_mib']:.1f} | " + f"{entry['out_vs_fp64']['max_abs']:.3e} | " + f"{_fmt(entry.get('lse_vs_fp64', {}).get('max_abs'), '.3e')} | " + f"{_fmt(entry['repeat_bitwise'])} |" + ) + add("") + add("### Forward+backward") + add("") + add(f"| S | {platform_column}Path | Median (ms) | p95 (ms) | vs sdpa | Peak MiB |") + add(f"|---:|{platform_rule}---|---:|---:|---:|---:|") + for seq_len in seq_lens: + for label, pl in platforms: + case = _case_for(pl, dtype, seq_len) + if case is None: + continue + baseline = case["paths"].get("sdpa", {}).get("train_fwd_bwd", {}).get("median_ms") + for name in PATH_NAMES: + entry = case["paths"].get(name) + if not entry or "train_fwd_bwd" not in entry: + continue + median = entry["train_fwd_bwd"]["median_ms"] + ratio = f"{median / baseline:.2f}x" if baseline else "n/a" + prefix = f"{label} | " if multi else "" + add( + f"| {seq_len} | {prefix}{name} | {median:.4f} | " + f"{entry['train_fwd_bwd']['p95_ms']:.4f} | {ratio} | " + f"{entry['train_peak_mib']:.1f} |" + ) + add("") + if multi: + add( + "Host peak memory is an RSS high-water delta sampled from `/proc`, not an " + "allocator statistic, so the `cpu` rows approximate and are not directly " + "comparable to the device figures." + ) + add("") + + # ---- production core vs reference + add("## Production core versus the reference core") + add("") + add( + "These are two different kernels, so this is a tolerance comparison, not a parity " + "claim. It is here to size the gap, not to assert equality." + ) + add("") + add("| dtype | S | out max-abs | out relative-L2 | lse max-abs |") + add("|---|---:|---:|---:|---:|") + for case in payload["single_gpu"]["cases"]: + gap = case.get("strict_vs_reference") + if not gap: + continue + add( + f"| {case['dtype']} | {case['seq_len']} | {gap['out']['max_abs']:.3e} | " + f"{gap['out']['relative_l2']:.3e} | {gap['lse']['max_abs']:.3e} |" + ) + add("") + + # ---- batch composition + add("## Batch-composition invariance") + add("") + add( + "A row computed alone must be bitwise equal to the same row inside a batch. " + "The strict ROCm core rejects `B > 1` outright, so for that path the property is " + "structural rather than measured." + ) + add("") + add("| S | Path | Bitwise | Mismatched | Note |") + add("|---:|---|:---:|---:|---|") + for row in payload.get("batch_composition", []): + for name, entry in row["paths"].items(): + note = entry.get("note", "measured") + add( + f"| {row['seq_len']} | {name} | {_fmt(entry['out_bitwise'])} | " + f"{entry['out_mismatched']} | {note} |" + ) + add("") + + # ---- TP degree + add("## TP-degree invariance of the strict ROCm core") + add("") + add( + "A head shard computed under TP=N versus the same slice of an unsharded run. TP performs " + "no cross-rank reduction in attention, so any nonzero value means the kernel's result " + "depends on how many heads shared the launch. `raw_launch` is one launch for all heads; " + "`one_kv_group_per_launch` is the schedule the Vime provider actually uses." + ) + add("") + add("| S | Schedule | TP | Local Hq | Local Hkv | out max-abs | lse max-abs | Invariant |") + add("|---:|---|---:|---:|---:|---:|---:|:---:|") + for row in payload.get("tp_head_sensitivity", []): + add( + f"| {row['seq_len']} | {row['schedule']} | {row['tp']} | {row['local_q_heads']} | " + f"{row['local_kv_heads']} | {row['out_max_abs']:.6e} | {row['lse_max_abs']:.6e} | " + f"{_fmt(row['invariant'])} |" + ) + add("") + + # ---- schedule cost + schedule = payload.get("tp_schedule_cost") or [] + if schedule: + add("## Cost of the per-KV-group launch schedule") + add("") + add( + "§ TP-degree invariance is bought by launching the core once per " + "`(batch row, KV group)` instead of once for all heads. This table is that " + "bill. `raw_launch` is one launch for all heads and is **not** the production " + "schedule; `per_kv_group` is what the Vime provider actually runs " + "(`Hkv` launches per row)." + ) + add("") + add( + "| S | Launches | sdpa (ms) | raw_launch (ms) | per_kv_group (ms) | " + "vs raw | vs sdpa |" + ) + add("|---:|---:|---:|---:|---:|---:|---:|") + for row in schedule: + raw = row["raw_launch"]["median_ms"] + group = row["one_kv_group_per_launch"]["median_ms"] + sdpa = row["sdpa"]["median_ms"] + add( + f"| {row['seq_len']} | {row['launches']} | {sdpa:.4f} | {raw:.4f} | " + f"{group:.4f} | {group / raw:.2f}x | {group / sdpa:.2f}x |" + ) + add("") + + # ---- distributed + distributed = payload.get("distributed") or [] + if distributed: + add("## Distributed CP (RCCL AG/RS transport)") + add("") + add( + "Schedule: all-gather Q/K/V and the position ids over the CP group, run the strict " + "core once on the full sequence, reduce-scatter `(out, lse)` back to this rank's " + "query range. Acceptance is bitwise against a CP=1 run of the same core." + ) + add("") + add( + "| Topology | World | TP | CP | Replicas | S | Median (ms) | p95 (ms) | " + "Peak MiB/rank | out bitwise | lse bitwise | Repeat |" + ) + add("|---|---:|---:|---:|---:|---:|---:|---:|---:|:---:|:---:|:---:|") + for row in distributed: + if "error" in row: + add( + f"| {row['topology']} | — | — | — | — | — | — | — | — | " + "error | error | error |" + ) + continue + add( + f"| {row['topology']} | {row['world_size']} | {row['tp_world_size']} | " + f"{row['cp_world_size']} | {row.get('replicas', 1)} | {row['seq_len']} | " + f"{row['forward']['median_ms']:.4f} | " + f"{row['forward']['p95_ms']:.4f} | {row['peak_mib_per_rank']:.1f} | " + f"{_fmt(row['out_bitwise_vs_cp1'])} | {_fmt(row['lse_bitwise_vs_cp1'])} | " + f"{_fmt(row['repeat_bitwise'])} |" + ) + add("") + errors = [row for row in distributed if "error" in row] + for row in errors: + add(f"- `{row['topology']}` failed: {row['error']}") + if errors: + add("") + + add("## Figures") + add("") + add( + "`reference-native` and `triton-bitwise` allocate exactly the same buffers, so their " + "memory curves coincide and the later-drawn series hides the earlier one." + ) + add("") + add("![Single-device latency and memory grid](single_gpu_grid.png)") + add("") + add("![Single-device latency](single_gpu_latency.png)") + add("") + add("![Single-device peak memory](single_gpu_memory.png)") + add("") + add("![Bitwise exactness matrix](exactness_matrix.png)") + add("") + add("![TP-degree invariance](tp_degree_invariance.png)") + add("") + if distributed: + add("![Distributed CP latency](distributed_cp_latency.png)") + add("") + + (output_directory / "report.md").write_text("\n".join(lines) + "\n") + + +def _write_figures(payload: dict[str, Any], output_directory: Path) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + plt.style.use("seaborn-v0_8-whitegrid") + plt.rcParams.update({"font.size": 10, "axes.titlesize": 11, "legend.fontsize": 8}) + + style = { + "sdpa": {"marker": "o", "color": "#888888", "linestyle": "--"}, + "strict-aiter": {"marker": "s", "color": "#d62728"}, + "reference-native": {"marker": "^", "color": "#1f77b4"}, + "triton-bitwise": {"marker": "D", "color": "#2ca02c"}, + } + + cases = [c for c in payload["single_gpu"]["cases"] if c["dtype"] == "bf16"] + if not cases: + return + seq_lens = sorted({c["seq_len"] for c in cases}) + present = [n for n in PATH_NAMES if any(n in c["paths"] for c in cases)] + + panels = ( + ("forward", "median_ms", "Forward latency", "median ms", True), + ("train_fwd_bwd", "median_ms", "Forward+backward latency", "median ms", True), + ("forward_peak_mib", None, "Forward peak memory", "peak MiB above live", True), + ("train_peak_mib", None, "Forward+backward peak memory", "peak MiB above live", True), + ) + + def value(case, name, key, sub): + entry = case["paths"].get(name) + if entry is None or key not in entry: + return float("nan") + return entry[key][sub] if sub else entry[key] + + def draw(axis, key, sub, title, ylabel, log_y): + for index, name in enumerate(present): + ys = [ + value(next(c for c in cases if c["seq_len"] == s), name, key, sub) for s in seq_lens + ] + axis.plot( + seq_lens, + ys, + label=name, + linewidth=3.0 - 0.35 * index, + markersize=7 - 0.5 * index, + zorder=3 + index, + **style.get(name, {}), + ) + axis.set_xscale("log", base=2) + if log_y: + axis.set_yscale("log") + axis.set_xlabel("sequence length") + axis.set_ylabel(ylabel) + axis.set_title(f"BF16: {title}") + axis.grid(True, which="both", alpha=0.3) + axis.legend() + + for filename, chosen in ( + ("single_gpu_latency.png", panels[:2]), + ("single_gpu_memory.png", panels[2:]), + ): + figure, axes = plt.subplots(1, 2, figsize=(12, 4.5)) + for axis, (key, sub, title, ylabel, log_y) in zip(axes, chosen): + draw(axis, key, sub, title, ylabel, log_y) + figure.tight_layout() + figure.savefig(output_directory / filename, dpi=180) + plt.close(figure) + + figure, axes = plt.subplots(2, 2, figsize=(13, 9)) + for axis, (key, sub, title, ylabel, log_y) in zip(axes.flat, panels): + draw(axis, key, sub, title, ylabel, log_y) + figure.suptitle( + "Single-device strict Attention, BF16, Qwen3-8B shape " + f"(Hq={QWEN3_8B_Q_HEADS}, Hkv={QWEN3_8B_KV_HEADS}, D={QWEN3_8B_HEAD_DIM})", + fontsize=12, + ) + figure.tight_layout(rect=(0, 0, 1, 0.96)) + figure.savefig(output_directory / "single_gpu_grid.png", dpi=180) + plt.close(figure) + + # TP head-count sensitivity: raw launch versus one KV group per launch. + tp_rows = payload.get("tp_head_sensitivity") or [] + if tp_rows: + figure, axis = plt.subplots(figsize=(11, 5)) + for schedule, marker in (("raw_launch", "o"), ("one_kv_group_per_launch", "s")): + rows = [r for r in tp_rows if r["schedule"] == schedule] + if not rows: + continue + labels = [f"S={r['seq_len']}\nTP={r['tp']}" for r in rows] + axis.plot( + range(len(rows)), + [max(r["out_max_abs"], 1e-12) for r in rows], + marker=marker, + label=schedule, + linewidth=2.4, + ) + axis.set_xticks(range(len([r for r in tp_rows if r["schedule"] == "raw_launch"]))) + axis.set_xticklabels( + [f"S={r['seq_len']}\nTP={r['tp']}" for r in tp_rows if r["schedule"] == "raw_launch"], + fontsize=8, + ) + axis.set_yscale("log") + axis.set_ylabel("out max-abs vs unsharded slice (1e-12 == bitwise)") + axis.set_title("TP-degree invariance of the strict ROCm core, BF16") + axis.grid(True, which="both", alpha=0.3) + axis.legend() + figure.tight_layout() + figure.savefig(output_directory / "tp_degree_invariance.png", dpi=180) + plt.close(figure) + + distributed = [r for r in (payload.get("distributed") or []) if "error" not in r] + if distributed: + figure, axis = plt.subplots(figsize=(max(9, 1.7 * len(distributed)), 5.0)) + labels = [f"{r['topology']}\nS={r['seq_len']}" for r in distributed] + xs = list(range(len(distributed))) + width = 0.38 + baseline = [r.get("cp1_baseline", {}).get("median_ms", float("nan")) for r in distributed] + measured = [r["forward"]["median_ms"] for r in distributed] + bars_a = axis.bar( + [x - width / 2 for x in xs], baseline, width, label="CP=1 baseline", color="#888888" + ) + bars_b = axis.bar( + [x + width / 2 for x in xs], measured, width, label="CP AG/RS", color="#2ca02c" + ) + for bars in (bars_a, bars_b): + axis.bar_label(bars, fmt="%.2f", fontsize=8, padding=2) + axis.set_xticks(xs) + axis.set_xticklabels(labels, fontsize=9) + axis.set_ylabel("median ms") + axis.set_title( + "Strict ROCm Attention: CP=1 baseline vs RCCL AG/RS CP transport, BF16 S=4096" + ) + axis.grid(True, axis="y", alpha=0.3) + axis.legend() + figure.tight_layout() + figure.savefig(output_directory / "distributed_cp_latency.png", dpi=180) + plt.close(figure) + + # ---- exactness matrix, in the shape PR #325 uses for its topology mismatch heatmap + single_rows, single_labels = [], [] + for case in payload["single_gpu"]["cases"]: + parity = case.get("triton_vs_reference") + if not parity: + continue + grads = next( + (r for r in payload.get("backward_parity", []) if r["seq_len"] == case["seq_len"]), + {}, + ) + use_grads = case["dtype"] == "bf16" and grads + # Gradients are only measured on the BF16 sweep; NaN renders as "not measured" + # rather than a zero that would read as "measured and equal". + missing = float("nan") + single_rows.append( + [ + parity["out_mismatched"], + parity["lse_mismatched"], + grads.get("dq_mismatched", missing) if use_grads else missing, + grads.get("dk_mismatched", missing) if use_grads else missing, + grads.get("dv_mismatched", missing) if use_grads else missing, + ] + ) + single_labels.append(f"{case['dtype']}, S={case['seq_len']}") + + dist_rows = [ + [r["out_mismatched_all_ranks"], r["lse_mismatched_all_ranks"]] for r in distributed + ] + dist_labels = [ + f"{r['topology']} (TP={r['tp_world_size']}, CP={r['cp_world_size']})" for r in distributed + ] + + if single_rows or dist_rows: + panels = [] + if single_rows: + panels.append( + ( + single_rows, + single_labels, + ["out", "lse", "dQ", "dK", "dV"], + "Triton core vs native reference core", + ) + ) + if dist_rows: + panels.append((dist_rows, dist_labels, ["out", "lse"], "CP topology vs CP=1")) + figure, axes = plt.subplots(1, len(panels), figsize=(7.5 * len(panels), 5.6), squeeze=False) + for axis, (rows, row_labels, col_labels, title) in zip(axes.flat, panels): + matrix = [[float(value) for value in row] for row in rows] + colormap = plt.get_cmap("RdYlGn_r").copy() + colormap.set_bad("#d9d9d9") + image = axis.imshow(matrix, aspect="auto", cmap=colormap, vmin=0, vmax=1.0) + axis.grid(False) + for y in range(len(matrix)): + for x in range(len(matrix[y])): + cell = matrix[y][x] + label = "n/m" if cell != cell else str(int(cell)) + axis.text( + x, + y, + label, + ha="center", + va="center", + fontsize=15, + fontweight="bold", + color="#222222", + ) + axis.set_xticks(range(len(col_labels)), col_labels) + axis.set_yticks(range(len(row_labels)), row_labels) + axis.set_xlabel("Compared tensor") + axis.set_title(title) + figure.colorbar(image, ax=axis, fraction=0.03, pad=0.03).set_label( + "Mismatched elements" + ) + figure.suptitle( + "Bitwise exactness — every measured cell must be 0 (n/m = not measured)", + fontsize=13, + ) + figure.tight_layout(rect=(0, 0, 1, 0.95)) + figure.savefig(output_directory / "exactness_matrix.png", dpi=180) + plt.close(figure) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", type=Path, default=Path("benchmarks/results/ws2_rocm_mi300x") + ) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--samples", type=int, default=20) + parser.add_argument("--training-samples", type=int, default=10) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--q-heads", type=int, default=QWEN3_8B_Q_HEADS) + parser.add_argument("--kv-heads", type=int, default=QWEN3_8B_KV_HEADS) + parser.add_argument("--head-dim", type=int, default=QWEN3_8B_HEAD_DIM) + parser.add_argument( + "--seq-lens", + type=lambda s: tuple(int(x) for x in s.split(",")), + default=DEFAULT_SEQ_LENS, + ) + parser.add_argument("--dtypes", default="bf16,fp16") + parser.add_argument( + "--path-budget-seconds", + type=float, + default=300.0, + help=( + "Skip a (path, case) whose measured single call implies more than this many " + "seconds of sampling. The row records the observed cost instead, which is " + "itself the finding. 0 disables the budget." + ), + ) + parser.add_argument( + "--device", + default="auto", + help="auto | cuda | cpu. On cpu only the sdpa and pytorch-native paths exist.", + ) + parser.add_argument( + "--platform-label", + default=None, + help="Name for this run in merged reports (default: mi300x / h100 / cpu, auto-detected).", + ) + parser.add_argument( + "--compare-with", + action="append", + default=[], + metavar="LABEL=PATH", + help="Merge another platform's results.json into the report. Repeatable.", + ) + parser.add_argument("--skip-distributed", action="store_true") + parser.add_argument("--skip-figures", action="store_true") + parser.add_argument( + "--distributed-only", + action="store_true", + help="Re-run only the distributed topologies and merge into an existing results.json.", + ) + parser.add_argument( + "--report-only", + action="store_true", + help="Re-render report.md and the figures from an existing results.json.", + ) + arguments = parser.parse_args() + + if arguments.report_only: + payload = json.loads((arguments.output_dir / "results.json").read_text()) + _write_report(payload, arguments.output_dir, _load_comparisons(arguments.compare_with)) + if not arguments.skip_figures: + _write_figures(payload, arguments.output_dir) + print(json.dumps({"output_dir": str(arguments.output_dir)}, indent=2)) + return + + if arguments.device == "auto": + device_type = "cuda" if torch.cuda.is_available() else "cpu" + else: + device_type = arguments.device + if device_type == "cuda" and not torch.cuda.is_available(): + raise SystemExit("--device cuda requested but no CUDA/ROCm device is visible") + + device = torch.device(device_type, 0) if device_type == "cuda" else torch.device("cpu") + if device.type == "cuda": + torch.cuda.set_device(device) + dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} + dtypes = tuple(dtype_map[name] for name in arguments.dtypes.split(",")) + + paths = _Paths(device) + payload: dict[str, Any] = { + "platform_label": arguments.platform_label or _default_platform_label(device), + "environment": _environment(device), + "configuration": { + "batch": arguments.batch, + "q_heads": arguments.q_heads, + "kv_heads": arguments.kv_heads, + "head_dim": arguments.head_dim, + "seq_lens": list(arguments.seq_lens), + "dtypes": arguments.dtypes.split(","), + "warmup": arguments.warmup, + "samples": arguments.samples, + "training_samples": arguments.training_samples, + }, + "unavailable_paths": paths.errors, + } + + existing: dict[str, Any] = {} + if arguments.distributed_only: + existing = json.loads((arguments.output_dir / "results.json").read_text()) + payload = dict(existing) + payload["environment"] = _environment() + + if not arguments.distributed_only: + payload["single_gpu"] = { + "cases": _single_gpu_benchmarks( + paths=paths, + seq_lens=arguments.seq_lens, + dtypes=dtypes, + q_heads=arguments.q_heads, + kv_heads=arguments.kv_heads, + head_dim=arguments.head_dim, + batch=arguments.batch, + warmup=arguments.warmup, + samples=arguments.samples, + training_samples=arguments.training_samples, + device=device, + ) + } + payload["backward_parity"] = _backward_parity( + paths=paths, + seq_lens=arguments.seq_lens, + q_heads=arguments.q_heads, + kv_heads=arguments.kv_heads, + head_dim=arguments.head_dim, + batch=arguments.batch, + device=device, + ) + payload["batch_composition"] = _batch_composition( + paths=paths, + seq_lens=arguments.seq_lens, + q_heads=arguments.q_heads, + kv_heads=arguments.kv_heads, + head_dim=arguments.head_dim, + device=device, + ) + payload["tp_head_sensitivity"] = ( + [] + if device.type != "cuda" + else _tp_head_sensitivity( + paths=paths, + seq_lens=arguments.seq_lens, + tp_degrees=DEFAULT_TP_DEGREES, + q_heads=arguments.q_heads, + kv_heads=arguments.kv_heads, + head_dim=arguments.head_dim, + device=device, + ) + ) + + distributed: list[dict[str, Any]] = [] + if not arguments.skip_distributed and device.type == "cuda": + available = torch.cuda.device_count() + for topology in DISTRIBUTED_TOPOLOGIES: + if topology[1] * topology[2] * topology[3] > available: + continue + try: + distributed.append( + _run_distributed_topology( + topology, + warmup=arguments.warmup, + samples=arguments.samples, + seq_len=arguments.seq_lens[-1] if arguments.seq_lens else 2048, + q_heads=arguments.q_heads, + kv_heads=arguments.kv_heads, + head_dim=arguments.head_dim, + ) + ) + except Exception as exc: # noqa: BLE001 - a failed topology is a reported row + distributed.append( + {"topology": topology[0], "error": f"{type(exc).__name__}: {exc}"} + ) + payload["distributed"] = distributed + + arguments.output_dir.mkdir(parents=True, exist_ok=True) + (arguments.output_dir / "results.json").write_text(json.dumps(payload, indent=2) + "\n") + _write_report(payload, arguments.output_dir, _load_comparisons(arguments.compare_with)) + if not arguments.skip_figures: + _write_figures(payload, arguments.output_dir) + print(json.dumps({"output_dir": str(arguments.output_dir)}, indent=2)) + + +if __name__ == "__main__": + os.environ.setdefault("NCCL_IB_DISABLE", "1") + main() diff --git a/benchmarks/results/pr319_rocm_mi300x/distributed/report.md b/benchmarks/results/pr319_rocm_mi300x/distributed/report.md new file mode 100644 index 00000000..e46222af --- /dev/null +++ b/benchmarks/results/pr319_rocm_mi300x/distributed/report.md @@ -0,0 +1,16 @@ +# PR #319 — strict ROCm CP attention, multi-rank acceptance (MI300X) + +`torchrun --standalone --nproc-per-node=N scripts/ws2_p2p_nccl_attention_reference_check.py \ + --transport rccl_ag_rs --strict-shared-core` + +Requires the native extension built for the active GPU platform +(`PYTORCH_ROCM_ARCH=gfx942 python setup.py build_ext --inplace`); without it the strict path +fails closed on the ROCm deterministic RoPE operator rather than running a different one. + +| world | TP | CP | replicas | transport | ranks passed | out | lse | dQ | dK | dV | +|---:|---:|---:|---:|---|---|---|---|---|---|---| +| 2 | 1 | 2 | 1 | rccl_ag_rs | 2/2 | bitwise | bitwise | bitwise | bitwise | bitwise | +| 4 | 2 | 2 | 1 | rccl_ag_rs | 4/4 | bitwise | bitwise | bitwise | bitwise | bitwise | +| 8 | 2 | 2 | 2 | rccl_ag_rs | 8/8 | bitwise | bitwise | bitwise | bitwise | bitwise | + +Every rank reports `strict_shared_core.executed=true` and `passed=true`. diff --git a/benchmarks/results/pr319_rocm_mi300x/distributed/strict_rccl_ag_rs_w2.json b/benchmarks/results/pr319_rocm_mi300x/distributed/strict_rccl_ag_rs_w2.json new file mode 100644 index 00000000..b7f1a00d --- /dev/null +++ b/benchmarks/results/pr319_rocm_mi300x/distributed/strict_rccl_ag_rs_w2.json @@ -0,0 +1,582 @@ +{ + "backend": "nccl", + "collective_version": [ + 2, + 28, + 9 + ], + "cp_world_size": 2, + "device_name": "AMD Instinct MI300X", + "git_commit": "3247e3ef390ac2f55bf93864ffa70b0a1350ec4c", + "global_failure_count": 0, + "platform": "rocm", + "ranks": [ + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 0, + "cp_world_size": 2, + "device": "cuda:0", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "final_out_max_abs": 0.00390625, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 2, + "local_block_indices": [ + 0, + 1 + ], + "lse_max_abs": 2.384185791015625e-07, + "out_max_abs": 7.152557373046875e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 0, + 8 + ], + "rank": 0, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 1, + "replica_index": 0, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 0, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 0, + 8 + ], + "strict_local_query_range": [ + 0, + 8 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 0, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 0, + "tp_world_size": 1, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 1, + "cp_world_size": 2, + "device": "cuda:1", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "final_out_max_abs": 0.00390625, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 2, + "local_block_indices": [ + 2, + 3 + ], + "lse_max_abs": 4.76837158203125e-07, + "out_max_abs": 5.960464477539062e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 8, + 16 + ], + "rank": 1, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 1, + "replica_index": 0, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 1, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 8, + 16 + ], + "strict_local_query_range": [ + 8, + 16 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 0, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 0, + "tp_world_size": 1, + "transport": "rccl_ag_rs" + } + ], + "replica_count": 1, + "runtime_version": "7.14.60850", + "schema_version": "ws2_rccl_ag_rs_attention/v2", + "torch_version": "2.12.0+rocm7.14.0a20260608", + "tp_world_size": 1, + "transport": "rccl_ag_rs", + "world_size": 2 +} diff --git a/benchmarks/results/pr319_rocm_mi300x/distributed/strict_rccl_ag_rs_w4.json b/benchmarks/results/pr319_rocm_mi300x/distributed/strict_rccl_ag_rs_w4.json new file mode 100644 index 00000000..72c267f1 --- /dev/null +++ b/benchmarks/results/pr319_rocm_mi300x/distributed/strict_rccl_ag_rs_w4.json @@ -0,0 +1,1142 @@ +{ + "backend": "nccl", + "collective_version": [ + 2, + 28, + 9 + ], + "cp_world_size": 2, + "device_name": "AMD Instinct MI300X", + "git_commit": "3247e3ef390ac2f55bf93864ffa70b0a1350ec4c", + "global_failure_count": 0, + "platform": "rocm", + "ranks": [ + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 0, + "cp_world_size": 2, + "device": "cuda:0", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "final_out_max_abs": 0.00390625, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 4, + "local_block_indices": [ + 0, + 1 + ], + "lse_max_abs": 2.384185791015625e-07, + "out_max_abs": 7.152557373046875e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 0, + 8 + ], + "rank": 0, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 1, + "replica_index": 0, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 0, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 0, + 8 + ], + "strict_local_query_range": [ + 0, + 8 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 0, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 0, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 1, + "cp_world_size": 2, + "device": "cuda:1", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "final_out_max_abs": 0.00390625, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 4, + "local_block_indices": [ + 2, + 3 + ], + "lse_max_abs": 4.76837158203125e-07, + "out_max_abs": 5.960464477539062e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 8, + 16 + ], + "rank": 1, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 1, + "replica_index": 0, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 1, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 8, + 16 + ], + "strict_local_query_range": [ + 8, + 16 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 0, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 0, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 0, + "cp_world_size": 2, + "device": "cuda:2", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "final_out_max_abs": 0.001953125, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 4, + "local_block_indices": [ + 0, + 1 + ], + "lse_max_abs": 2.384185791015625e-07, + "out_max_abs": 7.152557373046875e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 0, + 8 + ], + "rank": 2, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 1, + "replica_index": 0, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 0, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 0, + 8 + ], + "strict_local_query_range": [ + 0, + 8 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 1, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 1, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 1, + "cp_world_size": 2, + "device": "cuda:3", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "final_out_max_abs": 3.814697265625e-06, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 4, + "local_block_indices": [ + 2, + 3 + ], + "lse_max_abs": 4.76837158203125e-07, + "out_max_abs": 7.152557373046875e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 8, + 16 + ], + "rank": 3, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 1, + "replica_index": 0, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 1, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 8, + 16 + ], + "strict_local_query_range": [ + 8, + 16 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 1, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 1, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + } + ], + "replica_count": 1, + "runtime_version": "7.14.60850", + "schema_version": "ws2_rccl_ag_rs_attention/v2", + "torch_version": "2.12.0+rocm7.14.0a20260608", + "tp_world_size": 2, + "transport": "rccl_ag_rs", + "world_size": 4 +} diff --git a/benchmarks/results/pr319_rocm_mi300x/distributed/strict_rccl_ag_rs_w8.json b/benchmarks/results/pr319_rocm_mi300x/distributed/strict_rccl_ag_rs_w8.json new file mode 100644 index 00000000..0c9a7d1a --- /dev/null +++ b/benchmarks/results/pr319_rocm_mi300x/distributed/strict_rccl_ag_rs_w8.json @@ -0,0 +1,2262 @@ +{ + "backend": "nccl", + "collective_version": [ + 2, + 28, + 9 + ], + "cp_world_size": 2, + "device_name": "AMD Instinct MI300X", + "git_commit": "3247e3ef390ac2f55bf93864ffa70b0a1350ec4c", + "global_failure_count": 0, + "platform": "rocm", + "ranks": [ + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 0, + "cp_world_size": 2, + "device": "cuda:0", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "final_out_max_abs": 0.00390625, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 8, + "local_block_indices": [ + 0, + 1 + ], + "lse_max_abs": 2.384185791015625e-07, + "out_max_abs": 7.152557373046875e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 0, + 8 + ], + "rank": 0, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 2, + "replica_index": 0, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 0, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 0, + 8 + ], + "strict_local_query_range": [ + 0, + 8 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 0, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 0, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 1, + "cp_world_size": 2, + "device": "cuda:1", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "final_out_max_abs": 0.00390625, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 8, + "local_block_indices": [ + 2, + 3 + ], + "lse_max_abs": 4.76837158203125e-07, + "out_max_abs": 5.960464477539062e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 8, + 16 + ], + "rank": 1, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 2, + "replica_index": 0, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 1, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 8, + 16 + ], + "strict_local_query_range": [ + 8, + 16 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 0, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 0, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 0, + "cp_world_size": 2, + "device": "cuda:2", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "final_out_max_abs": 0.001953125, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 8, + "local_block_indices": [ + 0, + 1 + ], + "lse_max_abs": 2.384185791015625e-07, + "out_max_abs": 7.152557373046875e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 0, + 8 + ], + "rank": 2, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 2, + "replica_index": 0, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 0, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 0, + 8 + ], + "strict_local_query_range": [ + 0, + 8 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 1, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 1, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 1, + "cp_world_size": 2, + "device": "cuda:3", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "final_out_max_abs": 3.814697265625e-06, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 8, + "local_block_indices": [ + 2, + 3 + ], + "lse_max_abs": 4.76837158203125e-07, + "out_max_abs": 7.152557373046875e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 8, + 16 + ], + "rank": 3, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 2, + "replica_index": 0, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 1, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 8, + 16 + ], + "strict_local_query_range": [ + 8, + 16 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 1, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 1, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 0, + "cp_world_size": 2, + "device": "cuda:4", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "final_out_max_abs": 1.52587890625e-05, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 8, + "local_block_indices": [ + 0, + 1 + ], + "lse_max_abs": 2.384185791015625e-07, + "out_max_abs": 7.152557373046875e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 0, + 8 + ], + "rank": 4, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 2, + "replica_index": 1, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 0, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 0, + 8 + ], + "strict_local_query_range": [ + 0, + 8 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 0, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 0, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 1, + "cp_world_size": 2, + "device": "cuda:5", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "final_out_max_abs": 0.00048828125, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 8, + "local_block_indices": [ + 2, + 3 + ], + "lse_max_abs": 4.76837158203125e-07, + "out_max_abs": 9.5367431640625e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 8, + 16 + ], + "rank": 5, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 2, + "replica_index": 1, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 0 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 0 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 1, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 8, + 16 + ], + "strict_local_query_range": [ + 8, + 16 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 0, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 0, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 0, + "cp_world_size": 2, + "device": "cuda:6", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "final_out_max_abs": 0.000244140625, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 8, + "local_block_indices": [ + 0, + 1 + ], + "lse_max_abs": 2.384185791015625e-07, + "out_max_abs": 7.152557373046875e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 0, + 8 + ], + "rank": 6, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 2, + "replica_index": 1, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 0, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 0, + 8 + ], + "strict_local_query_range": [ + 0, + 8 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 1, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 1, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + }, + { + "accum_dtype": "fp32", + "atol": 0.0002, + "cp_rank": 1, + "cp_world_size": 2, + "device": "cuda:7", + "downcast_at": "final_write", + "dtype": "bf16", + "expected_block_manifest": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "final_out_max_abs": 0.001953125, + "final_output_dtype": "bfloat16", + "final_write_atol": 0.02, + "gathered_block_indices": [ + 0, + 1, + 2, + 3 + ], + "global_failure_count": 0, + "global_world_size": 8, + "local_block_indices": [ + 2, + 3 + ], + "lse_max_abs": 4.76837158203125e-07, + "out_max_abs": 7.152557373046875e-07, + "passed": true, + "protocol": "ag_query_local_kv_rs_out_lse", + "query_ag": "rccl_ag_rs", + "query_ag_max_abs": 0.0, + "query_range": [ + 8, + 16 + ], + "rank": 7, + "repeat_count": 3, + "repeat_lse_bitwise": true, + "repeat_manifest_bitwise": true, + "repeat_out_bitwise": true, + "repeat_query_bitwise": true, + "replica_count": 2, + "replica_index": 1, + "strict_protocol": "ag_qkv_positions_shared_core_rs_out_lse", + "strict_shared_core": { + "actual_backend": "aiter.rocm.ck_dense_mha", + "bitwise": { + "dk": true, + "dq": true, + "dv": true, + "lse": true, + "out": true + }, + "communication_autograd": true, + "communication_backend": "rccl_ag_rs", + "executed": true, + "fallback": false, + "identity_errors": [], + "max_abs": { + "dk": 0.0, + "dq": 0.0, + "dv": 0.0, + "lse": 0.0, + "out": 0.0 + }, + "native_attention_arithmetic": true, + "passed": true, + "production_ready": true, + "repeat_lse_bitwise": true, + "repeat_out_bitwise": true, + "split_kv_policy": "disabled", + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_mode": true, + "strict_provenance": { + "accum_dtype": "fp32", + "actual_backend": "aiter.rocm.ck_dense_mha", + "adapter_backend": "flashinfer", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "db61d7ce62a907c5079830622723f0abcc7fd35816d88006e381140cb6566ff4", + "arithmetic_plan_source": "aiter.ops.mha", + "arithmetic_semantics_verified": true, + "attention_backend": "aiter.rocm.ck_dense_mha", + "attention_mode": "prefill", + "batch_invariant_claim": "strict_runtime_verified", + "causal": true, + "communication_backend": "rccl_ag_rs", + "communication_overlap": "disabled", + "compute_communication": "decoupled", + "compute_order": [ + 0, + 1 + ], + "compute_schedule": "rlkernel.attention.strict_ring_state.v1", + "cp_comm_accum_dtype": "fp32", + "cp_comm_attention_numeric_reduction": false, + "cp_comm_backend": "rccl_ag_rs", + "cp_comm_compute_communication": "decoupled", + "cp_comm_contract": "partial_out_lse_global_block_index", + "cp_comm_expected_blocks": [ + { + "global_block_index": 0, + "kv_block_end": 4, + "kv_block_start": 0, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 1, + "kv_block_end": 8, + "kv_block_start": 4, + "owner_cp_rank": 0, + "owner_tp_rank": 1 + }, + { + "global_block_index": 2, + "kv_block_end": 12, + "kv_block_start": 8, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + }, + { + "global_block_index": 3, + "kv_block_end": 16, + "kv_block_start": 12, + "owner_cp_rank": 1, + "owner_tp_rank": 1 + } + ], + "cp_comm_expected_kv_token_range": [ + 0, + 16 + ], + "cp_comm_merge_order": "global_block_index", + "cp_comm_merge_root_cp_rank": 0, + "cp_comm_pattern": "ag_rs", + "cp_comm_query_token_ranges": [ + [ + 0, + 8 + ], + [ + 8, + 16 + ] + ], + "cp_comm_required": true, + "cp_comm_return_lse": true, + "cp_comm_runtime": "rccl", + "cp_comm_status": "implemented", + "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_strict_contract": "ag_qkv_positions_shared_core_rs_out_lse", + "cp_comm_strict_kv_communication": "all_gather", + "cp_comm_strict_position_communication": "all_gather", + "cp_rank": 1, + "cp_world_size": 2, + "deterministic_backward": true, + "downcast_at": "final_write", + "fa_api_source": null, + "fa_package_version": null, + "fallback": false, + "fallback_reason": null, + "k_cache_rope_state": "post_rope", + "lse_domain": "attention", + "lse_dtype": "fp32", + "lse_exported": true, + "materialization": "ag_qkv_positions_shared_core_rs", + "merge_order_indices": [ + 0, + 1 + ], + "native_attention_arithmetic": true, + "num_splits": 1, + "platform": "rocm", + "production_ready": true, + "q_rope_state": "post_rope", + "reference_only": false, + "requested_backend": "flashinfer_layout_adapter", + "ring_partial_arithmetic": false, + "ring_schedule_default": true, + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_fusion": false, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1000000.0, + "rotary_dim": null, + "softmax_scale": null, + "split_kv_control": "dense_non_split_api", + "strict_comm_autograd": true, + "strict_core_id": "rlkernel.attention.rocm.aiter_ck_dense_mha.v1", + "strict_core_row_plans": [ + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + }, + { + "actual_split_boundaries": [ + [ + 0, + 16 + ] + ], + "actual_split_kv_count": 1, + "actual_split_kv_policy": "disabled", + "actual_split_kv_size": null, + "requested_split_kv_policy": "disabled", + "requested_split_kv_size": null, + "split_kv_accum_dtype": "fp32", + "split_kv_backend": "aiter.rocm.ck_dense_mha", + "split_kv_downcast_at": "final_write", + "split_kv_fallback": false, + "split_kv_fallback_reason": null, + "split_kv_merge_order": "global_block_index", + "split_kv_plan_source": "contract_exact" + } + ], + "strict_full_qkv_all_gather": true, + "strict_local_kv_range": [ + 8, + 16 + ], + "strict_local_query_range": [ + 8, + 16 + ], + "strict_mode": true, + "strict_position_ids_all_gather": true, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv", + "strict_split_kv": "disabled", + "tp_rank": 1, + "tp_world_size": 2 + }, + "strict_schedule": "single_batch_aiter_ck_dense_mha_no_splitkv" + }, + "tp_rank": 1, + "tp_world_size": 2, + "transport": "rccl_ag_rs" + } + ], + "replica_count": 2, + "runtime_version": "7.14.60850", + "schema_version": "ws2_rccl_ag_rs_attention/v2", + "torch_version": "2.12.0+rocm7.14.0a20260608", + "tp_world_size": 2, + "transport": "rccl_ag_rs", + "world_size": 8 +} diff --git a/benchmarks/results/ws2_cpu/report.md b/benchmarks/results/ws2_cpu/report.md new file mode 100644 index 00000000..e5c5b510 --- /dev/null +++ b/benchmarks/results/ws2_cpu/report.md @@ -0,0 +1,150 @@ +# WS2 strict ROCm Attention — bitwise parity and performance + +> Operator-only benchmark. No model checkpoint or serving engine was used; +> the shapes are Qwen3-8B's attention shapes. + +## Environment + +| Item | cpu | +|---|---| +| architecture | n/a | +| cpu_count | 192 | +| cuda | None | +| extension_attention_symbols | none | +| gpu | n/a (host execution) | +| gpu_count | 0 | +| hip | None | +| native_collective | n/a (single-process host run) | +| python | 3.12.3 | +| torch | 2.12.0+rocm7.14.0a20260608 | +| torch_threads | 192 | +| triton | n/a (host execution) | + +## Methodology + +- Operator shape: `Hq=32`, `Hkv=8`, `D=128`, `B=1`, causal; sequence sweep 512, 1024, 2048. +- Measured paths: + - `sdpa`: `torch.nn.functional.scaled_dot_product_attention`. **Speed baseline only** — as in PR #325, no accuracy comparison is mixed into the speed table. + - `strict-aiter`: `StrictRocmAiterCKAttentionCore` called **once for all heads**. This is the core, not the production schedule: the Vime provider launches it once per (batch row, KV group). See the per-KV-group schedule table for that cost. + - `reference-native`: `_C.deterministic_attention_forward/backward`, the materializing FP32 reference core hipified from the shared `.cu`. + - `triton-bitwise`: `TritonDeterministicAttentionOp`, whose contract is bit-identity with `reference-native`. +- Timing: CUDA events, median and p95. Peak memory is the per-call increase in `torch.cuda.max_memory_allocated` above what was live before the call. +- Accuracy is against an FP64 oracle over the same BF16/FP16-rounded inputs. Repeat = two identical calls are bitwise equal; batch-invariant = a row computed alone is bitwise equal to the same row inside a batch. +- 2 warmups, 5 measured forward samples, 3 measured forward+backward samples. Raw medians, p95, min and max are in `results.json`. + +Reproduce from the repository root: + +```bash +python benchmarks/benchmark_ws2_rocm_attention.py \ + --seq-lens 512,1024,2048 \ + --dtypes bf16,fp16 \ + --warmup 2 --samples 5 \ + --training-samples 3 \ + --output-dir benchmarks/results/ws2_rocm_mi300x +``` + +### Unavailable paths + +- `strict-aiter`: GPU-only path; not available on the host +- `strict-fa4`: GPU-only path; not available on the host +- `reference-native`: GPU-only path; not available on the host +- `triton-bitwise`: GPU-only path; not available on the host + +## Bitwise parity: Triton port vs the native reference core + +Acceptance is 0 mismatched elements. This is the contract the Triton core exists to hold. + +| dtype | S | out mismatched | lse mismatched | dQ | dK | dV | bitwise | +|---|---:|---:|---:|---:|---:|---:|:---:| + +`dQ/dK/dV` are measured on the BF16 sweep only; `n/a` marks the FP16 rows. + +## Single-device Attention (bf16) + +### Forward + +| S | Path | Median (ms) | p95 (ms) | vs sdpa | Peak MiB | out max-abs vs FP64 | lse max-abs vs FP64 | Repeat | +|---:|---|---:|---:|---:|---:|---:|---:|:---:| +| 512 | sdpa | 38.1068 | 43.2369 | 1.00x | 0.1 | 8.606e-03 | n/a | yes | +| 512 | pytorch-native | 7.0910 | 9.0480 | 0.19x | 0.0 | 1.947e-02 | n/a | yes | +| 1024 | sdpa | 174.7517 | 252.6524 | 1.00x | 109.8 | 8.191e-03 | n/a | yes | +| 1024 | pytorch-native | 58.5536 | 87.5402 | 0.34x | 130.2 | 1.610e-02 | n/a | yes | +| 2048 | sdpa | 192.7010 | 232.4717 | 1.00x | 111.9 | 9.314e-03 | n/a | yes | +| 2048 | pytorch-native | 227.7564 | 307.3500 | 1.18x | 548.6 | 1.790e-02 | n/a | yes | + +### Forward+backward + +| S | Path | Median (ms) | p95 (ms) | vs sdpa | Peak MiB | +|---:|---|---:|---:|---:|---:| +| 512 | sdpa | 170.4205 | 175.4842 | 1.00x | 7.5 | +| 512 | pytorch-native | 48.0205 | 156.2514 | 0.28x | 64.8 | +| 1024 | sdpa | 297.3551 | 358.7597 | 1.00x | 135.3 | +| 1024 | pytorch-native | 267.1181 | 313.6076 | 0.90x | 190.8 | +| 2048 | sdpa | 555.5966 | 569.2160 | 1.00x | 112.9 | +| 2048 | pytorch-native | 489.2399 | 550.9883 | 0.88x | 815.6 | + +## Single-device Attention (fp16) + +### Forward + +| S | Path | Median (ms) | p95 (ms) | vs sdpa | Peak MiB | out max-abs vs FP64 | lse max-abs vs FP64 | Repeat | +|---:|---|---:|---:|---:|---:|---:|---:|:---:| +| 512 | sdpa | 31.7511 | 33.4260 | 1.00x | 0.0 | 1.045e-03 | n/a | yes | +| 512 | pytorch-native | 697.2643 | 707.5163 | 21.96x | 0.0 | 2.811e-03 | n/a | yes | +| 1024 | sdpa | 96.1593 | 113.6039 | 1.00x | 79.9 | 1.075e-03 | n/a | yes | +| 1024 | pytorch-native | 2726.1935 | 2804.8509 | 28.35x | 65.7 | 2.117e-03 | n/a | yes | +| 2048 | sdpa | 250.5798 | 286.6925 | 1.00x | 81.5 | 1.065e-03 | n/a | yes | +| 2048 | pytorch-native | 21112.2733 | 21152.1648 | 84.25x | 515.2 | 2.143e-03 | n/a | yes | + +### Forward+backward + +| S | Path | Median (ms) | p95 (ms) | vs sdpa | Peak MiB | +|---:|---|---:|---:|---:|---:| +| 512 | sdpa | 204.9800 | 242.0397 | 1.00x | 0.0 | +| 512 | pytorch-native | 2154.7387 | 2248.4103 | 10.51x | 0.0 | +| 1024 | sdpa | 386.7709 | 387.0739 | 1.00x | 79.7 | +| 1024 | pytorch-native | 8579.2534 | 12587.9295 | 22.18x | 189.2 | +| 2048 | sdpa | 1122.4169 | 1212.3130 | 1.00x | 79.7 | +| 2048 | pytorch-native | 53718.0483 | 54692.0301 | 47.86x | 764.6 | + +## Production core versus the reference core + +These are two different kernels, so this is a tolerance comparison, not a parity claim. It is here to size the gap, not to assert equality. + +| dtype | S | out max-abs | out relative-L2 | lse max-abs | +|---|---:|---:|---:|---:| + +## Batch-composition invariance + +A row computed alone must be bitwise equal to the same row inside a batch. The strict ROCm core rejects `B > 1` outright, so for that path the property is structural rather than measured. + +| S | Path | Bitwise | Mismatched | Note | +|---:|---|:---:|---:|---| +| 512 | sdpa | yes | 0 | measured | +| 512 | pytorch-native | yes | 0 | measured | +| 1024 | sdpa | yes | 0 | measured | +| 1024 | pytorch-native | yes | 0 | measured | +| 2048 | sdpa | yes | 0 | measured | +| 2048 | pytorch-native | yes | 0 | measured | + +## TP-degree invariance of the strict ROCm core + +A head shard computed under TP=N versus the same slice of an unsharded run. TP performs no cross-rank reduction in attention, so any nonzero value means the kernel's result depends on how many heads shared the launch. `raw_launch` is one launch for all heads; `one_kv_group_per_launch` is the schedule the Vime provider actually uses. + +| S | Schedule | TP | Local Hq | Local Hkv | out max-abs | lse max-abs | Invariant | +|---:|---|---:|---:|---:|---:|---:|:---:| + +## Figures + +`reference-native` and `triton-bitwise` allocate exactly the same buffers, so their memory curves coincide and the later-drawn series hides the earlier one. + +![Single-device latency and memory grid](single_gpu_grid.png) + +![Single-device latency](single_gpu_latency.png) + +![Single-device peak memory](single_gpu_memory.png) + +![Bitwise exactness matrix](exactness_matrix.png) + +![TP-degree invariance](tp_degree_invariance.png) + diff --git a/benchmarks/results/ws2_cpu/results.json b/benchmarks/results/ws2_cpu/results.json new file mode 100644 index 00000000..d6bc0783 --- /dev/null +++ b/benchmarks/results/ws2_cpu/results.json @@ -0,0 +1,467 @@ +{ + "platform_label": "cpu", + "environment": { + "cpu_count": 192, + "torch_threads": 192, + "gpu": "n/a (host execution)", + "architecture": "n/a", + "gpu_count": 0, + "hip": null, + "cuda": null, + "torch": "2.12.0+rocm7.14.0a20260608", + "triton": "n/a (host execution)", + "python": "3.12.3", + "extension_attention_symbols": [], + "native_collective": "n/a (single-process host run)" + }, + "configuration": { + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "seq_lens": [ + 512, + 1024, + 2048 + ], + "dtypes": [ + "bf16", + "fp16" + ], + "warmup": 2, + "samples": 5, + "training_samples": 3 + }, + "unavailable_paths": { + "strict-aiter": "GPU-only path; not available on the host", + "strict-fa4": "GPU-only path; not available on the host", + "reference-native": "GPU-only path; not available on the host", + "triton-bitwise": "GPU-only path; not available on the host" + }, + "single_gpu": { + "cases": [ + { + "dtype": "bf16", + "seq_len": 512, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 38.106778636574745, + "p95_ms": 43.23689788579941, + "min_ms": 28.195404447615147, + "max_ms": 44.51697692275047 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 0.0859375, + "out_vs_fp64": { + "max_abs": 0.00860644332381133, + "relative_l2": 0.001973059990724691 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 170.42051907628775, + "p95_ms": 175.48417346552014, + "min_ms": 132.9573979601264, + "max_ms": 176.0468017309904 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 7.4765625 + }, + "pytorch-native": { + "forward": { + "median_ms": 7.091020233929157, + "p95_ms": 9.047956205904484, + "min_ms": 4.163389094173908, + "max_ms": 9.431971237063408 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 0.0, + "out_vs_fp64": { + "max_abs": 0.01947146352388973, + "relative_l2": 0.004330172876834844 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 48.0204951018095, + "p95_ms": 156.25138729810712, + "min_ms": 43.106830678880215, + "max_ms": 168.27704198658466 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 64.8203125 + } + } + }, + { + "dtype": "bf16", + "seq_len": 1024, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 174.75166637450457, + "p95_ms": 252.65235546976325, + "min_ms": 164.2555631697178, + "max_ms": 271.48198056966066 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 109.82421875, + "out_vs_fp64": { + "max_abs": 0.008190526417996224, + "relative_l2": 0.002021095362635071 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 297.35514242202044, + "p95_ms": 358.75966083258385, + "min_ms": 296.0860254243016, + "max_ms": 365.5823851004243 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 135.33203125 + }, + "pytorch-native": { + "forward": { + "median_ms": 58.55359323322773, + "p95_ms": 87.5402009114623, + "min_ms": 54.07743901014328, + "max_ms": 92.4938665702939 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 130.1875, + "out_vs_fp64": { + "max_abs": 0.016100370761021, + "relative_l2": 0.004516966595870689 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 267.11810380220413, + "p95_ms": 313.60761895775795, + "min_ms": 249.75966848433018, + "max_ms": 318.7731206417084 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 190.8046875 + } + } + }, + { + "dtype": "bf16", + "seq_len": 2048, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 192.7009578794241, + "p95_ms": 232.47170187532902, + "min_ms": 168.04722882807255, + "max_ms": 234.8661571741104 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 111.88671875, + "out_vs_fp64": { + "max_abs": 0.00931387262518557, + "relative_l2": 0.002037123315687043 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 555.5966263636947, + "p95_ms": 569.2159625701606, + "min_ms": 521.6728867962956, + "max_ms": 570.7292221486568 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 112.91796875 + }, + "pytorch-native": { + "forward": { + "median_ms": 227.75637917220592, + "p95_ms": 307.34998527914286, + "min_ms": 209.55913793295622, + "max_ms": 317.96263344585896 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 548.5625, + "out_vs_fp64": { + "max_abs": 0.0178950704113916, + "relative_l2": 0.004656276082238093 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 489.23985194414854, + "p95_ms": 550.9883309714496, + "min_ms": 478.3158637583256, + "max_ms": 557.8492730855942 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 815.6015625 + } + } + }, + { + "dtype": "fp16", + "seq_len": 512, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 31.751069240272045, + "p95_ms": 33.425997383892536, + "min_ms": 29.143651947379112, + "max_ms": 33.47691521048546 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 0.0, + "out_vs_fp64": { + "max_abs": 0.0010446820342941976, + "relative_l2": 0.00025098233651414496 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 204.98001947999, + "p95_ms": 242.03968066722155, + "min_ms": 192.74852704256773, + "max_ms": 246.15742079913616 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 0.0 + }, + "pytorch-native": { + "forward": { + "median_ms": 697.2642932087183, + "p95_ms": 707.5163403525949, + "min_ms": 684.5130370929837, + "max_ms": 709.9788626655936 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 0.0, + "out_vs_fp64": { + "max_abs": 0.002810583458336957, + "relative_l2": 0.0005388467180647695 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 2154.7386962920427, + "p95_ms": 2248.4102914109826, + "min_ms": 2068.0867824703455, + "max_ms": 2258.818246424198 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 0.0 + } + } + }, + { + "dtype": "fp16", + "seq_len": 1024, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 96.15929331630468, + "p95_ms": 113.60393781214952, + "min_ms": 95.77664453536272, + "max_ms": 116.9054713100195 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 79.91796875, + "out_vs_fp64": { + "max_abs": 0.0010750978941995726, + "relative_l2": 0.00025550971614830375 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 386.77085004746914, + "p95_ms": 387.07391703501344, + "min_ms": 327.51816138625145, + "max_ms": 387.1075911447406 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 79.66015625 + }, + "pytorch-native": { + "forward": { + "median_ms": 2726.1935137212276, + "p95_ms": 2804.850871488452, + "min_ms": 2686.2719180062413, + "max_ms": 2823.695234954357 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 65.7421875, + "out_vs_fp64": { + "max_abs": 0.002116843588507722, + "relative_l2": 0.0005576223691406817 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 8579.253423959017, + "p95_ms": 12587.929507251827, + "min_ms": 8301.5665281564, + "max_ms": 13033.33796095103 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 189.20703125 + } + } + }, + { + "dtype": "fp16", + "seq_len": 2048, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 250.57981442660093, + "p95_ms": 286.6925349459052, + "min_ms": 194.098518230021, + "max_ms": 291.58885311335325 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 81.46484375, + "out_vs_fp64": { + "max_abs": 0.0010650871035631226, + "relative_l2": 0.0002589060030057917 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 1122.4169470369816, + "p95_ms": 1212.3129985295236, + "min_ms": 1065.4740231111646, + "max_ms": 1222.3014486953616 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 79.66015625 + }, + "pytorch-native": { + "forward": { + "median_ms": 21112.273322418332, + "p95_ms": 21152.164766564965, + "min_ms": 20729.920755140483, + "max_ms": 21160.758836194873 + }, + "forward_samples": 5, + "forward_truncated": false, + "forward_peak_mib": 515.18359375, + "out_vs_fp64": { + "max_abs": 0.0021428477134313173, + "relative_l2": 0.0005852208464456109 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 53718.04826427251, + "p95_ms": 54692.0300597325, + "min_ms": 53648.340058512986, + "max_ms": 54800.25025922805 + }, + "train_samples": 3, + "train_truncated": false, + "train_peak_mib": 764.57421875 + } + } + } + ] + }, + "backward_parity": [], + "batch_composition": [ + { + "seq_len": 512, + "paths": { + "sdpa": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "pytorch-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + } + } + }, + { + "seq_len": 1024, + "paths": { + "sdpa": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "pytorch-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + } + } + }, + { + "seq_len": 2048, + "paths": { + "sdpa": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "pytorch-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + } + } + } + ], + "tp_head_sensitivity": [], + "distributed": [] +} diff --git a/benchmarks/results/ws2_rocm_mi300x/PR_DESCRIPTION.md b/benchmarks/results/ws2_rocm_mi300x/PR_DESCRIPTION.md new file mode 100644 index 00000000..b187f07b --- /dev/null +++ b/benchmarks/results/ws2_rocm_mi300x/PR_DESCRIPTION.md @@ -0,0 +1,292 @@ +# WS2: bitwise-exact Attention on ROCm + +Brings the strict ROCm Attention path to a stated, measured bitwise standard, and adds a +Triton core that is bit-identical to the native reference kernel so the arithmetic +contract is testable without the vendor kernel. + +Operator-only. No model checkpoint or serving engine is loaded anywhere in this PR; the +shapes are Qwen3-8B's (`Hq=32`, `Hkv=8`, `D=128`). Measured on 8×MI300X (`gfx942`), +ROCm 7.14.60850, torch 2.12.0, Triton 3.7.0. + +--- + +## 1. What "bitwise" means here, and what it does not + +Several different guarantees get conflated in attention work, so this PR states which one +it claims at each boundary: + +| Scope | Claim | Enforced by | +| --- | --- | --- | +| Varying batch composition | **bitwise** | `B > 1` rejected at the core (§2.2) | +| Varying padding | **bitwise** | `key_padding_mask` rejected (§2.3) | +| Varying TP degree | **bitwise** | one KV group per launch (§2.4) | +| Varying CP degree | **bitwise** | RCCL-as-transport AG/RS (§2.5) | +| Triton core vs native reference core | **bitwise** | §2.8, measured in §3.1 | +| CUDA production core vs ROCm production core | **not claimed** | different vendor kernels | +| ROCm production core vs reference core | **not claimed** | different kernels; gap sized in §3.4 | + +The last two rows are deliberate. CUDA runs FlashAttention 4 CuTe and ROCm runs AITER CK +dense MHA; the tile decomposition, the online-softmax rescale order, and MFMA versus MMA +accumulation all differ. Nothing in the tree claims cross-platform bit equality and this +PR does not add such a claim. What is shared across platforms is the *contract*, not the +bits. + +## 2. The algorithmic arrangements + +The strict ROCm core does not reimplement attention. It removes every source of +run-to-run and shape-to-shape arithmetic variation from the vendor kernel and records what +was removed, so a mismatch becomes a contract violation rather than a debugging session. + +**2.1 Split-KV is structurally impossible, not merely switched off.** +Split-KV partitions the KV axis and merges partial softmax states; the partition count +depends on shape and occupancy, so the reduction order moves with it. CUDA can pass +`num_splits=1` to FA4. AITER exposes no such knob, so the ROCm core binds to the dense, +non-split API entry point instead and records `split_kv_control = "dense_non_split_api"`. +`SplitKVSpec` must be `DISABLED`; a non-disabled spec raises in `__init__` rather than +being quietly honoured. + +**2.2 Batch composition cannot change the bits, because `B > 1` is rejected.** +`StrictRocmAiterCKAttentionCore._validate_inputs` refuses any input with `q.size(0) != 1`: +*"strict AITER CK core executes one logical batch row at a time"*. This is stronger than +testing for batch invariance — there is no batched launch whose arithmetic could differ +from the single-row launch, because the batched launch does not exist. Callers materialise +each logical row separately. + +**2.3 Padding never enters a reduction.** +`key_padding_mask` is rejected outright: the core *"materializes each unpadded logical +row"*. A padded and an unpadded run of the same logical row cannot differ, because the +padded run is not expressible. + +**2.4 One KV group per launch, to make the result independent of TP degree.** +This is the ROCm-specific problem. AITER/CK's reduction order depends on how many heads +shared the launch, and TP performs no cross-rank reduction in attention — it is pure head +sharding — so a head shard computed under TP=4 was *not* bit-identical to the same shard +under TP=8 at some shapes. The provider therefore launches the core once per +`(batch row, KV group)` and concatenates, so every launch sees exactly one KV group and its +Q heads regardless of the TP degree that produced the shard. §3.2 measures both schedules +side by side; the cost is real and is reported. + +**2.5 RCCL is a transport, never a reduction.** +`_RCCLRankOrderedTransport` uses RCCL only for `all_gather` and a root-owned `scatter`. Its +`reduce_scatter` first gathers every source shard and then evaluates a fixed balanced rank +tree locally, so the floating-point combine order is ours and does not depend on RCCL's +internal algorithm selection, which varies with message size and topology. + +**2.6 The vendor kernel is fingerprinted, not version-pinned.** +AITER dispatches in Python, so a package version does not pin behaviour. +`_load_aiter_ck_ops()` takes a **sha256 of the `aiter.ops.mha` source file** and exports it +as `aiter_source_sha256` in the provenance, so a silent upstream change to the dispatch +logic invalidates the recorded arithmetic identity. + +**2.7 Fail closed, everywhere.** +A missing AITER entry point, a missing native extension, or a dispatch that resolves to a +different backend raises. No path substitutes a different kernel to keep a run alive. + +**2.8 A reference core shared with CUDA, and a Triton port that matches it bitwise.** +`csrc/cuda/attention/deterministic_attention.cu` is hipified to +`csrc/hip/attention/deterministic_attention.hip`, so the *reference* core genuinely is the +same algorithm on both platforms. This PR adds +`rl_engine/kernels/ops/triton/attention/deterministic_attn.py`, a Triton port whose +contract is bit-identity with that reference. Three things had to be reproduced rather than +re-derived: + +- **Dot products stay sequential FMA chains.** The C++ kernel accumulates one element at a + time in a single thread, so the contraction index is the *loop* and the head dim is the + *vector*. The opposite, much faster arrangement would reassociate the sum. +- **The row softmax keeps the 256-lane partial layout.** Key `k` belongs to lane `k % 256`; + each lane sums ascending, then a stride-halving fold combines the partials. + `_tree_sum_256` reproduces that fold exactly. +- **`expf`/`logf` are re-emitted instruction for instruction.** Every Triton exp/log + intrinsic — `tl.exp`, `tl.math.exp`, `libdevice.exp` — lowers to a bare `v_exp_f32`, about + 1 ULP away from the `expf` the C++ kernel calls, which alone broke parity on ~14% of + elements. The helpers reproduce hipcc's two-term argument reduction around that same + hardware instruction, with an inline-asm barrier to stop LLVM refolding the reduction into + an FMA. Verified bitwise over 4M+ inputs including subnormals, ±inf and NaN. + +The nvcc `expf`/`logf` sequences are not ported, so on CUDA the op refuses to construct +unless the caller passes `require_bitwise_libm=False`, rather than silently returning +non-bitwise results. + +--- + +## 3. Results + +Full report, `results.json` and figures: `benchmarks/results/ws2_rocm_mi300x/`. +Reproduce with `python benchmarks/benchmark_ws2_rocm_attention.py`. + +Paths: `sdpa` (`torch.nn.functional.scaled_dot_product_attention`, **speed baseline only** — +as in PR #325, no accuracy comparison is mixed into the speed table), `strict-aiter` (the +ROCm production core), `reference-hip` (`_C.deterministic_attention_*`), `triton-bitwise` +(this PR). + +### 3.1 Headline: Triton port vs the native reference core + +Acceptance is 0 mismatched elements. This is the contract the Triton core exists to hold. + +| dtype | S | out | lse | dQ | dK | dV | bitwise | +|---|---:|---:|---:|---:|---:|---:|:---:| +| bf16 | 512 | 0 | 0 | 0 | 0 | 0 | yes | +| bf16 | 1024 | 0 | 0 | 0 | 0 | 0 | yes | +| bf16 | 2048 | 0 | 0 | 0 | 0 | 0 | yes | +| bf16 | 4096 | 0 | 0 | 0 | 0 | 0 | yes | +| fp16 | 512 | 0 | 0 | — | — | — | yes | +| fp16 | 1024 | 0 | 0 | — | — | — | yes | +| fp16 | 2048 | 0 | 0 | — | — | — | yes | +| fp16 | 4096 | 0 | 0 | — | — | — | yes | + +`dQ/dK/dV` are measured on the BF16 sweep only. + +### 3.2 TP-degree invariance of the strict ROCm core + +A head shard computed under TP=N versus the same slice of an unsharded run. TP performs no +cross-rank reduction in attention, so any nonzero value means the kernel's result depends on +how many heads shared the launch. + +| S | TP | Local Hq | `raw_launch` out max-abs | `one_kv_group_per_launch` out max-abs | +|---:|---:|---:|---:|---:| +| 512 | 2 | 16 | 0.000000e+00 | 0.000000e+00 | +| 512 | 4 | 8 | 0.000000e+00 | 0.000000e+00 | +| 512 | 8 | 4 | 0.000000e+00 | 0.000000e+00 | +| 1024 | 2 | 16 | **7.812500e-03** | 0.000000e+00 | +| 1024 | 4 | 8 | **7.812500e-03** | 0.000000e+00 | +| 1024 | 8 | 4 | **7.812500e-03** | 0.000000e+00 | +| 2048 | 2 | 16 | 0.000000e+00 | 0.000000e+00 | +| 2048 | 4 | 8 | **3.906250e-03** | 0.000000e+00 | +| 2048 | 8 | 4 | **1.953125e-03** | 0.000000e+00 | +| 4096 | 2 | 16 | 0.000000e+00 | 0.000000e+00 | +| 4096 | 4 | 8 | 0.000000e+00 | 0.000000e+00 | +| 4096 | 8 | 4 | **3.906250e-03** | 0.000000e+00 | + +Raw AITER is non-invariant at 5 of 12 points, and *which* points is shape-dependent — the +failure is invisible at S=512 and at S=4096/TP=2, which is exactly what makes it dangerous: +training at TP=4 and rolling out at TP=8 would compare fine on most shapes. The per-KV-group +schedule is bitwise at **12 of 12**. This reproduces PR #319's finding on independent inputs. + +### 3.3 Single-GPU latency and memory (BF16) + +| S | Path | Fwd median (ms) | vs sdpa | Fwd+bwd (ms) | vs sdpa | Fwd peak MiB | Fwd+bwd peak MiB | out max-abs vs FP64 | +|---:|---|---:|---:|---:|---:|---:|---:|---:| +| 512 | sdpa | 0.0782 | 1.00x | 0.3228 | 1.00x | 12.1 | 32.2 | 8.195e-03 | +| 512 | strict-aiter | 0.2428 | 3.11x | 0.6039 | 1.87x | 14.1 | 288.2 | 2.468e-02 | +| 512 | reference-hip | 0.9659 | 12.36x | 2.9714 | 9.21x | 36.1 | 78.1 | 7.741e-03 | +| 512 | triton-bitwise | 1.3816 | 17.68x | 4.8578 | 15.05x | 36.1 | 78.1 | 7.741e-03 | +| 1024 | sdpa | 0.1319 | 1.00x | 0.4319 | 1.00x | 24.1 | 64.3 | 1.027e-02 | +| 1024 | strict-aiter | 0.2468 | 1.87x | 0.9293 | 2.15x | 28.1 | 1088.4 | 2.609e-02 | +| 1024 | reference-hip | 3.1389 | 23.79x | 12.2009 | 28.25x | 136.1 | 284.3 | 7.810e-03 | +| 1024 | triton-bitwise | 4.8240 | 36.56x | 20.1771 | 46.71x | 136.1 | 284.3 | 7.810e-03 | +| 2048 | sdpa | 0.2887 | 1.00x | 1.0735 | 1.00x | 48.3 | 128.8 | 7.994e-03 | +| 2048 | strict-aiter | 0.2962 | 1.03x | 1.9019 | 1.77x | 56.3 | 4224.8 | 2.027e-02 | +| 2048 | reference-hip | 12.8467 | 44.50x | 47.8741 | 44.60x | 528.2 | 1080.5 | 7.804e-03 | +| 2048 | triton-bitwise | 19.4226 | 67.28x | 76.7414 | 71.49x | 528.3 | 1080.5 | 7.804e-03 | +| 4096 | sdpa | 0.6936 | 1.00x | 3.2464 | 1.00x | 96.5 | 257.5 | 9.604e-03 | +| 4096 | strict-aiter | **0.5644** | **0.81x** | 5.7304 | 1.77x | 112.5 | 16641.5 | 2.138e-02 | +| 4096 | reference-hip | 49.3624 | 71.17x | 173.3365 | 53.39x | 2080.5 | 4209.0 | 7.808e-03 | +| 4096 | triton-bitwise | 86.0655 | 124.08x | 304.1032 | 93.67x | 2080.5 | 4209.0 | 7.808e-03 | + +Three things worth reading off this table: + +- **The strict production core is not a tax at long sequence.** At S=4096 forward it is + *faster* than SDPA (0.81x), and its worst case across the sweep is 3.11x at S=512 where + absolute cost is 0.24 ms. The bitwise arrangements in §2 cost almost nothing in the + production path. +- **AITER's backward is memory-hungry.** `strict-aiter` fwd+bwd peaks at 16.6 GiB at S=4096 + versus 4.2 GiB for the materializing reference core — the reference core materializes an + FP32 `[B, Hq, Sq, Skv]` score matrix and is *still* 4x smaller. Worth knowing before + sizing a training run. +- **The deterministic cores are the most accurate of the four.** Against an FP64 oracle they + sit at 7.8e-03 versus 9.6e-03 for SDPA and 2.1e-02 for AITER. Determinism here is not + bought with accuracy. + +FP16 is in the full report; the shape of the result is the same. + +### 3.4 Production core versus reference core + +Two different kernels, so this is a tolerance comparison, not a parity claim. It is here to +size the gap. + +| S | out max-abs | out relative-L2 | lse max-abs | +|---:|---:|---:|---:| +| 512 | 3.125e-02 | 5.420e-03 | 9.537e-07 | +| 1024 | 3.125e-02 | 5.505e-03 | 1.431e-06 | +| 2048 | 1.562e-02 | 5.595e-03 | 1.907e-06 | +| 4096 | 1.562e-02 | 5.633e-03 | 3.815e-06 | + +### 3.5 Batch-composition invariance + +Bitwise at every sequence length for every path. For `strict-aiter` the property is +structural (§2.2) rather than measured: the batched launch does not exist. + +### 3.6 Distributed CP over the RCCL AG/RS transport + +Schedule: all-gather Q/K/V and the position ids over the CP group, run the strict core once +on the full sequence, reduce-scatter `(out, lse)` back to this rank's query range. Acceptance +is bitwise against a CP=1 run of the same core on the same inputs. S=4096, BF16. + +| Topology | World | TP | CP | Replicas | Local Hq/Hkv | Median (ms) | p95 (ms) | Peak MiB/rank | out bitwise | lse bitwise | Repeat | +|---|---:|---:|---:|---:|---|---:|---:|---:|:---:|:---:|:---:| +| `tp1_cp2` | 2 | 1 | 2 | 1 | 32/8 | 1.8352 | 1.8726 | 160.5 | yes | yes | yes | +| `tp2_cp2` | 4 | 2 | 2 | 1 | 16/4 | 1.2114 | 1.3138 | 80.3 | yes | yes | yes | +| `tp1_cp4` | 4 | 1 | 4 | 1 | 32/8 | 1.3973 | 1.4363 | 160.5 | yes | yes | yes | +| `tp2_cp2_x2` | 8 | 2 | 2 | 2 | 16/4 | 1.2297 | 1.2948 | 80.3 | yes | yes | yes | +| `tp2_cp4` | 8 | 2 | 4 | 1 | 16/4 | 1.2594 | 1.3294 | 80.3 | yes | yes | yes | +| `tp1_cp8` | 8 | 1 | 8 | 1 | 32/8 | 1.4112 | 2.2029 | 160.5 | yes | yes | yes | + +All six topologies are bitwise against CP=1 on both `out` and `lse`, with 0 mismatched +elements summed across every rank, and repeat-bitwise on every rank. `tp2_cp2_x2` is the +8-rank case PR #319 used: two independent CP groups running side by side at TP=2/CP=2. + +--- + +## 4. Figures + +![Single-device latency and memory grid](benchmarks/results/ws2_rocm_mi300x/single_gpu_grid.png) + +![TP-degree invariance](benchmarks/results/ws2_rocm_mi300x/tp_degree_invariance.png) + +![Distributed CP latency](benchmarks/results/ws2_rocm_mi300x/distributed_cp_latency.png) + +`reference-hip` and `triton-bitwise` allocate exactly the same buffers, so their memory +curves coincide and the later-drawn series hides the earlier one. + +--- + +## 5. Files + +| Path | What | +| --- | --- | +| `rl_engine/kernels/ops/triton/attention/deterministic_attn.py` | New. Triton core, bit-identical to `_C.deterministic_attention_*`. | +| `rl_engine/kernels/ops/triton/attention/__init__.py` | Exports the new op and `BITWISE_LIBM_PARITY`. | +| `tests/test_triton_deterministic_attention.py` | New. 71 parity / invariance tests. | +| `benchmarks/benchmark_ws2_rocm_attention.py` | New. Measurement matrix and figures, reusing PR #325 / #328 helpers. | +| `benchmarks/results/ws2_rocm_mi300x/` | Report, `results.json`, figures. | +| `csrc/ops.cpp` | Fix: the merge from `feat/rocm-deterministic-collectives` dropped an `#if !defined(USE_ROCM)` around the `prefix_shared_attention` registration but kept its `#endif`, leaving 13 `#endif` against 12 `#if`. The ROCm build failed with `#endif without #if`. | + +## 6. Test plan + +- `pytest tests/test_triton_deterministic_attention.py` — 71 passed. 15 shape/mask/scale + configs × {bf16, fp16} × {forward, backward}, plus end-to-end autograd through both ops, + the fully-masked-row case, batch-slice invariance, and a direct pin on the `expf`/`logf` + helpers against the vendor libm. +- `pytest tests/test_deterministic_attention_cuda.py` — 614 passed (the native core is + unaffected). +- `python benchmarks/benchmark_ws2_rocm_attention.py` — the report above. + +## 7. Known limitations + +- **CUDA is not covered by the Triton core's bitwise claim.** The nvcc `expf`/`logf` argument + reductions are not ported, and no CUDA device was available to derive or verify them. + `TritonDeterministicAttentionOp` raises on CUDA unless the caller passes + `require_bitwise_libm=False`; a test pins that behaviour on both platforms. +- **The Triton core is a parity core, not a FlashAttention replacement.** Like the native + reference it materialises the full FP32 `[B, Hq, Sq, Skv]` score matrix and runs + scalar-order reductions; §3.3 shows the cost. +- **`_C` does not register `deterministic_attention_forward_fp32`.** The `.cu` defines it but + the pybind registration is missing, so the *native* `DeterministicAttentionOp.forward_fp32` + raises `AttributeError`. Pre-existing, not touched here; the Triton `forward_fp32` works and + its test validates against the op's own downcast instead of the native path. +- **A stale comment contradicts the shipped TP policy.** + `rl_engine/integrations/vime/attention.py` still carries a comment saying RL-Kernel "binds + the degree rather than paying ~3x forward time", from before the merge that introduced the + per-KV-group launch loop. The provenance dict immediately below it correctly reports + `tp_degree_invariant: True`. §3.2 shows the code is right and the comment is wrong; flagged + here rather than silently rewritten. diff --git a/benchmarks/results/ws2_rocm_mi300x/distributed_cp_latency.png b/benchmarks/results/ws2_rocm_mi300x/distributed_cp_latency.png new file mode 100644 index 00000000..d2113dce Binary files /dev/null and b/benchmarks/results/ws2_rocm_mi300x/distributed_cp_latency.png differ diff --git a/benchmarks/results/ws2_rocm_mi300x/exactness_matrix.png b/benchmarks/results/ws2_rocm_mi300x/exactness_matrix.png new file mode 100644 index 00000000..dcc60896 Binary files /dev/null and b/benchmarks/results/ws2_rocm_mi300x/exactness_matrix.png differ diff --git a/benchmarks/results/ws2_rocm_mi300x/report.md b/benchmarks/results/ws2_rocm_mi300x/report.md new file mode 100644 index 00000000..36ce7166 --- /dev/null +++ b/benchmarks/results/ws2_rocm_mi300x/report.md @@ -0,0 +1,283 @@ +# WS2 strict ROCm Attention — bitwise parity and performance + +> Operator-only benchmark. No model checkpoint or serving engine was used; +> the shapes are Qwen3-8B's attention shapes. + +## Environment + +| Item | mi300x | +|---|---| +| architecture | gfx942:sramecc+:xnack- | +| cpu_count | 192 | +| cuda | None | +| extension_attention_symbols | deterministic_attention_backward, deterministic_attention_forward | +| gpu | AMD Instinct MI300X | +| gpu_count | 8 | +| hip | 7.14.60850 | +| native_collective | torch.distributed ProcessGroupNCCL (RCCL on ROCm) | +| python | 3.12.3 | +| torch | 2.12.0+rocm7.14.0a20260608 | +| torch_threads | 4 | +| triton | 3.7.0 | + +## Methodology + +- Operator shape: `Hq=32`, `Hkv=8`, `D=128`, `B=1`, causal; sequence sweep 512, 1024, 2048, 4096. +- Measured paths: + - `sdpa`: `torch.nn.functional.scaled_dot_product_attention`. **Speed baseline only** — as in PR #325, no accuracy comparison is mixed into the speed table. + - `strict-aiter`: `StrictRocmAiterCKAttentionCore` called **once for all heads**. This is the core, not the production schedule: the Vime provider launches it once per (batch row, KV group). See the per-KV-group schedule table for that cost. + - `reference-native`: `_C.deterministic_attention_forward/backward`, the materializing FP32 reference core hipified from the shared `.cu`. + - `triton-bitwise`: `TritonDeterministicAttentionOp`, whose contract is bit-identity with `reference-native`. +- Timing: CUDA events, median and p95. Peak memory is the per-call increase in `torch.cuda.max_memory_allocated` above what was live before the call. +- Accuracy is against an FP64 oracle over the same BF16/FP16-rounded inputs. Repeat = two identical calls are bitwise equal; batch-invariant = a row computed alone is bitwise equal to the same row inside a batch. +- 5 warmups, 20 measured forward samples, 10 measured forward+backward samples. Raw medians, p95, min and max are in `results.json`. + +Reproduce from the repository root: + +```bash +python benchmarks/benchmark_ws2_rocm_attention.py \ + --seq-lens 512,1024,2048,4096 \ + --dtypes bf16,fp16 \ + --warmup 5 --samples 20 \ + --training-samples 10 \ + --output-dir benchmarks/results/ws2_rocm_mi300x +``` + +### Unavailable paths + +- `strict-fa4`: CUDA-only path; this run is ROCm + +## Bitwise parity: Triton port vs the native reference core + +Acceptance is 0 mismatched elements. This is the contract the Triton core exists to hold. + +| dtype | S | out mismatched | lse mismatched | dQ | dK | dV | bitwise | +|---|---:|---:|---:|---:|---:|---:|:---:| +| bf16 | 512 | 0 | 0 | 0 | 0 | 0 | yes | +| bf16 | 1024 | 0 | 0 | 0 | 0 | 0 | yes | +| bf16 | 2048 | 0 | 0 | 0 | 0 | 0 | yes | +| bf16 | 4096 | 0 | 0 | 0 | 0 | 0 | yes | +| fp16 | 512 | 0 | 0 | n/a | n/a | n/a | yes | +| fp16 | 1024 | 0 | 0 | n/a | n/a | n/a | yes | +| fp16 | 2048 | 0 | 0 | n/a | n/a | n/a | yes | +| fp16 | 4096 | 0 | 0 | n/a | n/a | n/a | yes | + +`dQ/dK/dV` are measured on the BF16 sweep only; `n/a` marks the FP16 rows. + +## Single-device Attention (bf16) + +### Forward + +| S | Path | Median (ms) | p95 (ms) | vs sdpa | Peak MiB | out max-abs vs FP64 | lse max-abs vs FP64 | Repeat | +|---:|---|---:|---:|---:|---:|---:|---:|:---:| +| 512 | sdpa | 0.0785 | 0.0843 | 1.00x | 12.1 | 8.195e-03 | n/a | yes | +| 512 | pytorch-native | 0.1981 | 0.3121 | 2.52x | 44.2 | 1.391e-02 | n/a | yes | +| 512 | strict-aiter | 0.2475 | 0.2714 | 3.15x | 14.1 | 2.468e-02 | 8.359e-07 | yes | +| 512 | reference-native | 1.0180 | 1.0542 | 12.97x | 36.1 | 7.741e-03 | 8.111e-07 | yes | +| 512 | triton-bitwise | 1.3627 | 1.3866 | 17.36x | 36.1 | 7.741e-03 | 8.111e-07 | yes | +| 1024 | sdpa | 0.1327 | 0.1440 | 1.00x | 24.1 | 1.027e-02 | n/a | yes | +| 1024 | pytorch-native | 0.3096 | 0.3200 | 2.33x | 153.0 | 1.571e-02 | n/a | yes | +| 1024 | strict-aiter | 0.2451 | 0.2578 | 1.85x | 28.1 | 2.609e-02 | 1.213e-06 | yes | +| 1024 | reference-native | 3.1537 | 3.8402 | 23.77x | 136.1 | 7.810e-03 | 8.732e-07 | yes | +| 1024 | triton-bitwise | 4.8326 | 4.8614 | 36.42x | 136.1 | 7.810e-03 | 8.732e-07 | yes | +| 2048 | sdpa | 0.2875 | 0.3083 | 1.00x | 48.3 | 7.994e-03 | n/a | yes | +| 2048 | pytorch-native | 1.0848 | 1.1129 | 3.77x | 564.0 | 1.803e-02 | n/a | yes | +| 2048 | strict-aiter | 0.2938 | 0.3060 | 1.02x | 56.3 | 2.027e-02 | 2.288e-06 | yes | +| 2048 | reference-native | 12.7428 | 12.9149 | 44.32x | 528.2 | 7.804e-03 | 1.142e-06 | yes | +| 2048 | triton-bitwise | 19.4210 | 19.5038 | 67.55x | 528.3 | 7.804e-03 | 1.142e-06 | yes | +| 4096 | sdpa | 0.6965 | 0.7457 | 1.00x | 96.5 | 9.604e-03 | n/a | yes | +| 4096 | pytorch-native | 3.9513 | 4.2331 | 5.67x | 2160.0 | 1.398e-02 | n/a | yes | +| 4096 | strict-aiter | 0.5569 | 0.5848 | 0.80x | 112.5 | 2.138e-02 | 3.967e-06 | yes | +| 4096 | reference-native | 49.3536 | 49.4533 | 70.86x | 2080.5 | 7.808e-03 | 1.381e-06 | yes | +| 4096 | triton-bitwise | 105.4049 | 110.6211 | 151.34x | 2080.5 | 7.808e-03 | 1.381e-06 | yes | + +### Forward+backward + +| S | Path | Median (ms) | p95 (ms) | vs sdpa | Peak MiB | +|---:|---|---:|---:|---:|---:| +| 512 | sdpa | 0.2836 | 0.5891 | 1.00x | 32.2 | +| 512 | pytorch-native | 0.9766 | 1.0728 | 3.44x | 76.3 | +| 512 | strict-aiter | 0.6328 | 0.6556 | 2.23x | 288.2 | +| 512 | reference-native | 3.1217 | 3.1910 | 11.01x | 78.1 | +| 512 | triton-bitwise | 4.8656 | 4.9558 | 17.15x | 78.1 | +| 1024 | sdpa | 0.4391 | 0.4822 | 1.00x | 64.3 | +| 1024 | pytorch-native | 0.7501 | 0.8035 | 1.71x | 281.0 | +| 1024 | strict-aiter | 0.8448 | 0.9276 | 1.92x | 1088.4 | +| 1024 | reference-native | 12.0942 | 12.2326 | 27.54x | 284.3 | +| 1024 | triton-bitwise | 20.0742 | 20.2126 | 45.72x | 284.3 | +| 2048 | sdpa | 1.0837 | 1.1348 | 1.00x | 128.8 | +| 2048 | pytorch-native | 2.3871 | 2.4252 | 2.20x | 1076.0 | +| 2048 | strict-aiter | 1.9570 | 2.0058 | 1.81x | 4224.8 | +| 2048 | reference-native | 47.8157 | 48.0608 | 44.12x | 1080.5 | +| 2048 | triton-bitwise | 76.7789 | 77.9824 | 70.85x | 1080.5 | +| 4096 | sdpa | 3.2163 | 3.3038 | 1.00x | 257.5 | +| 4096 | pytorch-native | 9.3608 | 9.4943 | 2.91x | 4208.0 | +| 4096 | strict-aiter | 5.7263 | 5.8219 | 1.78x | 16641.5 | +| 4096 | reference-native | 173.3246 | 177.7302 | 53.89x | 4209.0 | +| 4096 | triton-bitwise | 306.5966 | 346.2128 | 95.33x | 4209.0 | + +## Single-device Attention (fp16) + +### Forward + +| S | Path | Median (ms) | p95 (ms) | vs sdpa | Peak MiB | out max-abs vs FP64 | lse max-abs vs FP64 | Repeat | +|---:|---|---:|---:|---:|---:|---:|---:|:---:| +| 512 | sdpa | 0.0703 | 0.1083 | 1.00x | 12.1 | 1.042e-03 | n/a | yes | +| 512 | pytorch-native | 0.1926 | 0.2989 | 2.74x | 44.2 | 1.811e-03 | n/a | yes | +| 512 | strict-aiter | 0.2729 | 0.2851 | 3.88x | 14.1 | 2.046e-03 | 8.756e-07 | yes | +| 512 | reference-native | 0.9919 | 1.0175 | 14.12x | 36.1 | 9.702e-04 | 8.340e-07 | yes | +| 512 | triton-bitwise | 1.3618 | 1.4008 | 19.38x | 36.1 | 9.702e-04 | 8.340e-07 | yes | +| 1024 | sdpa | 0.1264 | 0.1373 | 1.00x | 24.1 | 1.053e-03 | n/a | yes | +| 1024 | pytorch-native | 0.3178 | 0.3455 | 2.51x | 153.0 | 2.487e-03 | n/a | yes | +| 1024 | strict-aiter | 0.2891 | 0.3381 | 2.29x | 28.1 | 2.299e-03 | 1.250e-06 | yes | +| 1024 | reference-native | 3.0663 | 3.0975 | 24.26x | 136.1 | 9.757e-04 | 1.011e-06 | yes | +| 1024 | triton-bitwise | 4.7755 | 4.8063 | 37.78x | 136.1 | 9.757e-04 | 1.011e-06 | yes | +| 2048 | sdpa | 0.2897 | 0.2998 | 1.00x | 48.3 | 9.099e-04 | n/a | yes | +| 2048 | pytorch-native | 1.0828 | 1.1133 | 3.74x | 564.0 | 1.727e-03 | n/a | yes | +| 2048 | strict-aiter | 0.2991 | 0.3337 | 1.03x | 56.3 | 2.053e-03 | 2.540e-06 | yes | +| 2048 | reference-native | 12.5615 | 12.6840 | 43.36x | 528.2 | 9.099e-04 | 1.103e-06 | yes | +| 2048 | triton-bitwise | 19.2360 | 19.3075 | 66.40x | 528.3 | 9.099e-04 | 1.103e-06 | yes | +| 4096 | sdpa | 1.0774 | 1.1026 | 1.00x | 96.5 | 9.905e-04 | n/a | yes | +| 4096 | pytorch-native | 3.9607 | 4.2144 | 3.68x | 2160.0 | 2.339e-03 | n/a | yes | +| 4096 | strict-aiter | 0.5825 | 0.6701 | 0.54x | 112.5 | 1.953e-03 | 3.805e-06 | yes | +| 4096 | reference-native | 48.5714 | 48.6788 | 45.08x | 2080.5 | 9.681e-04 | 1.511e-06 | yes | +| 4096 | triton-bitwise | 84.3475 | 86.5872 | 78.29x | 2080.5 | 9.681e-04 | 1.511e-06 | yes | + +### Forward+backward + +| S | Path | Median (ms) | p95 (ms) | vs sdpa | Peak MiB | +|---:|---|---:|---:|---:|---:| +| 512 | sdpa | 0.3330 | 0.3590 | 1.00x | 32.2 | +| 512 | pytorch-native | 0.9847 | 1.0736 | 2.96x | 76.3 | +| 512 | strict-aiter | 1.1291 | 1.5507 | 3.39x | 288.2 | +| 512 | reference-native | 3.0893 | 3.1265 | 9.28x | 78.1 | +| 512 | triton-bitwise | 4.7979 | 4.8599 | 14.41x | 78.1 | +| 1024 | sdpa | 0.4161 | 0.4702 | 1.00x | 64.4 | +| 1024 | pytorch-native | 0.7710 | 0.8476 | 1.85x | 281.0 | +| 1024 | strict-aiter | 0.8597 | 0.9174 | 2.07x | 1088.4 | +| 1024 | reference-native | 11.8471 | 11.9334 | 28.47x | 284.3 | +| 1024 | triton-bitwise | 19.5915 | 19.7842 | 47.09x | 284.3 | +| 2048 | sdpa | 1.1873 | 1.5985 | 1.00x | 128.5 | +| 2048 | pytorch-native | 2.3897 | 2.9295 | 2.01x | 1076.0 | +| 2048 | strict-aiter | 1.7668 | 1.7958 | 1.49x | 4224.8 | +| 2048 | reference-native | 47.4088 | 47.5834 | 39.93x | 1080.5 | +| 2048 | triton-bitwise | 75.9789 | 76.2766 | 63.99x | 1080.5 | +| 4096 | sdpa | 3.9710 | 4.0939 | 1.00x | 257.0 | +| 4096 | pytorch-native | 9.4752 | 9.7886 | 2.39x | 4208.0 | +| 4096 | strict-aiter | 5.9360 | 5.9734 | 1.49x | 16641.5 | +| 4096 | reference-native | 171.5709 | 171.7084 | 43.21x | 4209.0 | +| 4096 | triton-bitwise | 301.2615 | 302.5984 | 75.87x | 4209.0 | + +## Production core versus the reference core + +These are two different kernels, so this is a tolerance comparison, not a parity claim. It is here to size the gap, not to assert equality. + +| dtype | S | out max-abs | out relative-L2 | lse max-abs | +|---|---:|---:|---:|---:| +| bf16 | 512 | 3.125e-02 | 5.420e-03 | 9.537e-07 | +| bf16 | 1024 | 3.125e-02 | 5.505e-03 | 1.431e-06 | +| bf16 | 2048 | 1.562e-02 | 5.595e-03 | 1.907e-06 | +| bf16 | 4096 | 1.562e-02 | 5.633e-03 | 3.815e-06 | +| fp16 | 512 | 1.953e-03 | 4.323e-04 | 9.537e-07 | +| fp16 | 1024 | 1.953e-03 | 4.466e-04 | 1.431e-06 | +| fp16 | 2048 | 1.953e-03 | 4.503e-04 | 2.861e-06 | +| fp16 | 4096 | 1.953e-03 | 4.580e-04 | 3.815e-06 | + +## Batch-composition invariance + +A row computed alone must be bitwise equal to the same row inside a batch. The strict ROCm core rejects `B > 1` outright, so for that path the property is structural rather than measured. + +| S | Path | Bitwise | Mismatched | Note | +|---:|---|:---:|---:|---| +| 512 | sdpa | yes | 0 | measured | +| 512 | pytorch-native | yes | 0 | measured | +| 512 | strict-aiter | yes | 0 | core executes one logical batch row per launch | +| 512 | reference-native | yes | 0 | measured | +| 512 | triton-bitwise | yes | 0 | measured | +| 1024 | sdpa | yes | 0 | measured | +| 1024 | pytorch-native | yes | 0 | measured | +| 1024 | strict-aiter | yes | 0 | core executes one logical batch row per launch | +| 1024 | reference-native | yes | 0 | measured | +| 1024 | triton-bitwise | yes | 0 | measured | +| 2048 | sdpa | yes | 0 | measured | +| 2048 | pytorch-native | yes | 0 | measured | +| 2048 | strict-aiter | yes | 0 | core executes one logical batch row per launch | +| 2048 | reference-native | yes | 0 | measured | +| 2048 | triton-bitwise | yes | 0 | measured | +| 4096 | sdpa | yes | 0 | measured | +| 4096 | pytorch-native | yes | 0 | measured | +| 4096 | strict-aiter | yes | 0 | core executes one logical batch row per launch | +| 4096 | reference-native | yes | 0 | measured | +| 4096 | triton-bitwise | yes | 0 | measured | + +## TP-degree invariance of the strict ROCm core + +A head shard computed under TP=N versus the same slice of an unsharded run. TP performs no cross-rank reduction in attention, so any nonzero value means the kernel's result depends on how many heads shared the launch. `raw_launch` is one launch for all heads; `one_kv_group_per_launch` is the schedule the Vime provider actually uses. + +| S | Schedule | TP | Local Hq | Local Hkv | out max-abs | lse max-abs | Invariant | +|---:|---|---:|---:|---:|---:|---:|:---:| +| 512 | raw_launch | 2 | 16 | 4 | 0.000000e+00 | 0.000000e+00 | yes | +| 512 | raw_launch | 4 | 8 | 2 | 0.000000e+00 | 0.000000e+00 | yes | +| 512 | raw_launch | 8 | 4 | 1 | 0.000000e+00 | 0.000000e+00 | yes | +| 512 | one_kv_group_per_launch | 2 | 16 | 4 | 0.000000e+00 | 0.000000e+00 | yes | +| 512 | one_kv_group_per_launch | 4 | 8 | 2 | 0.000000e+00 | 0.000000e+00 | yes | +| 512 | one_kv_group_per_launch | 8 | 4 | 1 | 0.000000e+00 | 0.000000e+00 | yes | +| 1024 | raw_launch | 2 | 16 | 4 | 7.812500e-03 | 1.907349e-06 | **no** | +| 1024 | raw_launch | 4 | 8 | 2 | 7.812500e-03 | 1.907349e-06 | **no** | +| 1024 | raw_launch | 8 | 4 | 1 | 7.812500e-03 | 1.907349e-06 | **no** | +| 1024 | one_kv_group_per_launch | 2 | 16 | 4 | 0.000000e+00 | 0.000000e+00 | yes | +| 1024 | one_kv_group_per_launch | 4 | 8 | 2 | 0.000000e+00 | 0.000000e+00 | yes | +| 1024 | one_kv_group_per_launch | 8 | 4 | 1 | 0.000000e+00 | 0.000000e+00 | yes | +| 2048 | raw_launch | 2 | 16 | 4 | 0.000000e+00 | 0.000000e+00 | yes | +| 2048 | raw_launch | 4 | 8 | 2 | 3.906250e-03 | 2.861023e-06 | **no** | +| 2048 | raw_launch | 8 | 4 | 1 | 1.953125e-03 | 2.861023e-06 | **no** | +| 2048 | one_kv_group_per_launch | 2 | 16 | 4 | 0.000000e+00 | 0.000000e+00 | yes | +| 2048 | one_kv_group_per_launch | 4 | 8 | 2 | 0.000000e+00 | 0.000000e+00 | yes | +| 2048 | one_kv_group_per_launch | 8 | 4 | 1 | 0.000000e+00 | 0.000000e+00 | yes | +| 4096 | raw_launch | 2 | 16 | 4 | 0.000000e+00 | 0.000000e+00 | yes | +| 4096 | raw_launch | 4 | 8 | 2 | 0.000000e+00 | 0.000000e+00 | yes | +| 4096 | raw_launch | 8 | 4 | 1 | 3.906250e-03 | 4.768372e-06 | **no** | +| 4096 | one_kv_group_per_launch | 2 | 16 | 4 | 0.000000e+00 | 0.000000e+00 | yes | +| 4096 | one_kv_group_per_launch | 4 | 8 | 2 | 0.000000e+00 | 0.000000e+00 | yes | +| 4096 | one_kv_group_per_launch | 8 | 4 | 1 | 0.000000e+00 | 0.000000e+00 | yes | + +## Cost of the per-KV-group launch schedule + +§ TP-degree invariance is bought by launching the core once per `(batch row, KV group)` instead of once for all heads. This table is that bill. `raw_launch` is one launch for all heads and is **not** the production schedule; `per_kv_group` is what the Vime provider actually runs (`Hkv` launches per row). + +| S | Launches | sdpa (ms) | raw_launch (ms) | per_kv_group (ms) | vs raw | vs sdpa | +|---:|---:|---:|---:|---:|---:|---:| +| 512 | 8 | 0.0712 | 0.2579 | 1.7759 | 6.89x | 24.95x | +| 1024 | 8 | 0.1302 | 0.2513 | 1.9985 | 7.95x | 15.35x | +| 2048 | 8 | 0.2802 | 0.2917 | 1.7507 | 6.00x | 6.25x | +| 4096 | 8 | 0.7046 | 0.5682 | 2.0472 | 3.60x | 2.91x | + +## Distributed CP (RCCL AG/RS transport) + +Schedule: all-gather Q/K/V and the position ids over the CP group, run the strict core once on the full sequence, reduce-scatter `(out, lse)` back to this rank's query range. Acceptance is bitwise against a CP=1 run of the same core. + +| Topology | World | TP | CP | Replicas | S | Median (ms) | p95 (ms) | Peak MiB/rank | out bitwise | lse bitwise | Repeat | +|---|---:|---:|---:|---:|---:|---:|---:|---:|:---:|:---:|:---:| +| tp1_cp2 | 2 | 1 | 2 | 1 | 4096 | 1.8379 | 1.9000 | 160.5 | yes | yes | yes | +| tp2_cp2 | 4 | 2 | 2 | 1 | 4096 | 1.2288 | 1.2961 | 80.3 | yes | yes | yes | +| tp1_cp4 | 4 | 1 | 4 | 1 | 4096 | 1.3988 | 1.4461 | 160.5 | yes | yes | yes | +| tp2_cp2_x2 | 8 | 2 | 2 | 2 | 4096 | 1.2809 | 3.3578 | 80.3 | yes | yes | yes | +| tp2_cp4 | 8 | 2 | 4 | 1 | 4096 | 2.3537 | 6.9351 | 80.3 | yes | yes | yes | +| tp1_cp8 | 8 | 1 | 8 | 1 | 4096 | 3.0636 | 35.8611 | 160.5 | yes | yes | yes | + +## Figures + +`reference-native` and `triton-bitwise` allocate exactly the same buffers, so their memory curves coincide and the later-drawn series hides the earlier one. + +![Single-device latency and memory grid](single_gpu_grid.png) + +![Single-device latency](single_gpu_latency.png) + +![Single-device peak memory](single_gpu_memory.png) + +![Bitwise exactness matrix](exactness_matrix.png) + +![TP-degree invariance](tp_degree_invariance.png) + +![Distributed CP latency](distributed_cp_latency.png) + diff --git a/benchmarks/results/ws2_rocm_mi300x/results.json b/benchmarks/results/ws2_rocm_mi300x/results.json new file mode 100644 index 00000000..0f5b2176 --- /dev/null +++ b/benchmarks/results/ws2_rocm_mi300x/results.json @@ -0,0 +1,1910 @@ +{ + "platform_label": "mi300x", + "environment": { + "cpu_count": 192, + "torch_threads": 4, + "gpu": "AMD Instinct MI300X", + "architecture": "gfx942:sramecc+:xnack-", + "gpu_count": 8, + "hip": "7.14.60850", + "cuda": null, + "torch": "2.12.0+rocm7.14.0a20260608", + "triton": "3.7.0", + "python": "3.12.3", + "extension_attention_symbols": [ + "deterministic_attention_backward", + "deterministic_attention_forward" + ], + "native_collective": "torch.distributed ProcessGroupNCCL (RCCL on ROCm)" + }, + "configuration": { + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "seq_lens": [ + 512, + 1024, + 2048, + 4096 + ], + "dtypes": [ + "bf16", + "fp16" + ], + "warmup": 5, + "samples": 20, + "training_samples": 10 + }, + "unavailable_paths": { + "strict-fa4": "CUDA-only path; this run is ROCm" + }, + "single_gpu": { + "cases": [ + { + "dtype": "bf16", + "seq_len": 512, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 0.07849700003862381, + "p95_ms": 0.08434915244579318, + "min_ms": 0.07539200037717819, + "max_ms": 0.11144600063562393 + }, + "forward_peak_mib": 12.06396484375, + "out_vs_fp64": { + "max_abs": 0.008195295214645348, + "relative_l2": 0.001967118270169404 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.2836415022611618, + "p95_ms": 0.5891397461295125, + "min_ms": 0.2573019862174988, + "max_ms": 0.7238360047340393 + }, + "train_peak_mib": 32.189453125 + }, + "pytorch-native": { + "forward": { + "median_ms": 0.19809399545192719, + "p95_ms": 0.312086047232151, + "min_ms": 0.18271200358867645, + "max_ms": 0.31326499581336975 + }, + "forward_peak_mib": 44.25, + "out_vs_fp64": { + "max_abs": 0.013913263434108813, + "relative_l2": 0.004276855892935304 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.9766320288181305, + "p95_ms": 1.0728291511535644, + "min_ms": 0.6414740085601807, + "max_ms": 1.0736769437789917 + }, + "train_peak_mib": 76.2509765625 + }, + "strict-aiter": { + "forward": { + "median_ms": 0.24754799902439117, + "p95_ms": 0.2713805958628655, + "min_ms": 0.2415190041065216, + "max_ms": 0.35468798875808716 + }, + "forward_peak_mib": 14.06298828125, + "out_vs_fp64": { + "max_abs": 0.024676734518057408, + "relative_l2": 0.00520034717487522 + }, + "lse_vs_fp64": { + "max_abs": 8.359452712269899e-07, + "relative_l2": 3.778278560802271e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.6328204870223999, + "p95_ms": 0.6555787414312363, + "min_ms": 0.6018149852752686, + "max_ms": 0.657056987285614 + }, + "train_peak_mib": 288.18896484375 + }, + "reference-native": { + "forward": { + "median_ms": 1.01797354221344, + "p95_ms": 1.054203498363495, + "min_ms": 0.9671170115470886, + "max_ms": 1.0990339517593384 + }, + "forward_peak_mib": 36.0625, + "out_vs_fp64": { + "max_abs": 0.0077411426297544494, + "relative_l2": 0.0015903199116563222 + }, + "lse_vs_fp64": { + "max_abs": 8.111189337967062e-07, + "relative_l2": 4.340800016528509e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 3.1216800212860107, + "p95_ms": 3.1909892201423644, + "min_ms": 2.9844770431518555, + "max_ms": 3.211052894592285 + }, + "train_peak_mib": 78.1259765625 + }, + "triton-bitwise": { + "forward": { + "median_ms": 1.3627254962921143, + "p95_ms": 1.386585181951523, + "min_ms": 1.3438379764556885, + "max_ms": 1.4528800249099731 + }, + "forward_peak_mib": 36.06298828125, + "out_vs_fp64": { + "max_abs": 0.0077411426297544494, + "relative_l2": 0.0015903199116563222 + }, + "lse_vs_fp64": { + "max_abs": 8.111189337967062e-07, + "relative_l2": 4.340800016528509e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 4.8655924797058105, + "p95_ms": 4.955807328224182, + "min_ms": 4.812994003295898, + "max_ms": 5.008446216583252 + }, + "train_peak_mib": 78.1259765625 + } + }, + "triton_vs_reference": { + "out_mismatched": 0, + "lse_mismatched": 0, + "out_relative_l2": 0.0, + "bitwise": true + }, + "strict_vs_reference": { + "production_path": "strict-aiter", + "out": { + "max_abs": 0.03125, + "relative_l2": 0.0054199441038222254 + }, + "lse": { + "max_abs": 9.5367431640625e-07, + "relative_l2": 4.7804698254096993e-08 + }, + "out_mismatched": 1710619 + } + }, + { + "dtype": "bf16", + "seq_len": 1024, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 0.13267749547958374, + "p95_ms": 0.14403650015592576, + "min_ms": 0.1257070004940033, + "max_ms": 0.16348299384117126 + }, + "forward_peak_mib": 24.12646484375, + "out_vs_fp64": { + "max_abs": 0.01027133353685894, + "relative_l2": 0.0019939102141256666 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.43907299637794495, + "p95_ms": 0.48221664726734154, + "min_ms": 0.4229089915752411, + "max_ms": 0.5076339840888977 + }, + "train_peak_mib": 64.251953125 + }, + "pytorch-native": { + "forward": { + "median_ms": 0.30955949425697327, + "p95_ms": 0.32002418935298926, + "min_ms": 0.2979629933834076, + "max_ms": 0.3738360106945038 + }, + "forward_peak_mib": 153.0, + "out_vs_fp64": { + "max_abs": 0.015709640340966446, + "relative_l2": 0.004502714586301753 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.7501150071620941, + "p95_ms": 0.8035147368907928, + "min_ms": 0.7432649731636047, + "max_ms": 0.822983980178833 + }, + "train_peak_mib": 281.0009765625 + }, + "strict-aiter": { + "forward": { + "median_ms": 0.24514450132846832, + "p95_ms": 0.25779101252555847, + "min_ms": 0.24143899977207184, + "max_ms": 0.2708820104598999 + }, + "forward_peak_mib": 28.12548828125, + "out_vs_fp64": { + "max_abs": 0.02608916227826974, + "relative_l2": 0.00528885768244778 + }, + "lse_vs_fp64": { + "max_abs": 1.2127858237676037e-06, + "relative_l2": 4.3688457448489715e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.8448359966278076, + "p95_ms": 0.9276321947574615, + "min_ms": 0.8153319954872131, + "max_ms": 0.9654340147972107 + }, + "train_peak_mib": 1088.37646484375 + }, + "reference-native": { + "forward": { + "median_ms": 3.1537084579467773, + "p95_ms": 3.8402310609817505, + "min_ms": 3.126007080078125, + "max_ms": 3.872627019882202 + }, + "forward_peak_mib": 136.125, + "out_vs_fp64": { + "max_abs": 0.007810300996808017, + "relative_l2": 0.0016039478773382583 + }, + "lse_vs_fp64": { + "max_abs": 8.732049643356277e-07, + "relative_l2": 4.1651456242074347e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 12.094237804412842, + "p95_ms": 12.232606029510498, + "min_ms": 12.002681732177734, + "max_ms": 12.233705520629883 + }, + "train_peak_mib": 284.2509765625 + }, + "triton-bitwise": { + "forward": { + "median_ms": 4.832623481750488, + "p95_ms": 4.861410903930664, + "min_ms": 4.802298069000244, + "max_ms": 5.01409387588501 + }, + "forward_peak_mib": 136.12548828125, + "out_vs_fp64": { + "max_abs": 0.007810300996808017, + "relative_l2": 0.0016039478773382583 + }, + "lse_vs_fp64": { + "max_abs": 8.732049643356277e-07, + "relative_l2": 4.1651456242074347e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 20.0741605758667, + "p95_ms": 20.212563514709473, + "min_ms": 20.023523330688477, + "max_ms": 20.296171188354492 + }, + "train_peak_mib": 284.2509765625 + } + }, + "triton_vs_reference": { + "out_mismatched": 0, + "lse_mismatched": 0, + "out_relative_l2": 0.0, + "bitwise": true + }, + "strict_vs_reference": { + "production_path": "strict-aiter", + "out": { + "max_abs": 0.03125, + "relative_l2": 0.005504811866196166 + }, + "lse": { + "max_abs": 1.430511474609375e-06, + "relative_l2": 5.160771963046382e-08 + }, + "out_mismatched": 3454931 + } + }, + { + "dtype": "bf16", + "seq_len": 2048, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 0.2875075042247772, + "p95_ms": 0.3082621052861214, + "min_ms": 0.2809379994869232, + "max_ms": 0.33497798442840576 + }, + "forward_peak_mib": 48.25146484375, + "out_vs_fp64": { + "max_abs": 0.007994353511369567, + "relative_l2": 0.00202698986197318 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 1.0837309956550598, + "p95_ms": 1.1348404943943022, + "min_ms": 1.0423489809036255, + "max_ms": 1.1623669862747192 + }, + "train_peak_mib": 128.751953125 + }, + "pytorch-native": { + "forward": { + "median_ms": 1.0847724676132202, + "p95_ms": 1.1129266023635864, + "min_ms": 1.068107008934021, + "max_ms": 1.1333249807357788 + }, + "forward_peak_mib": 564.0, + "out_vs_fp64": { + "max_abs": 0.018031305229536443, + "relative_l2": 0.004707755028054482 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 2.3871285915374756, + "p95_ms": 2.425160896778107, + "min_ms": 2.368360996246338, + "max_ms": 2.4312539100646973 + }, + "train_peak_mib": 1076.0009765625 + }, + "strict-aiter": { + "forward": { + "median_ms": 0.2938365042209625, + "p95_ms": 0.3060245126485825, + "min_ms": 0.29071199893951416, + "max_ms": 0.3374220132827759 + }, + "forward_peak_mib": 56.25048828125, + "out_vs_fp64": { + "max_abs": 0.020265153423261406, + "relative_l2": 0.005397474562875508 + }, + "lse_vs_fp64": { + "max_abs": 2.288034266939576e-06, + "relative_l2": 6.118480465147299e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 1.9570285081863403, + "p95_ms": 2.005762046575546, + "min_ms": 1.8095699548721313, + "max_ms": 2.021085023880005 + }, + "train_peak_mib": 4224.75146484375 + }, + "reference-native": { + "forward": { + "median_ms": 12.742782592773438, + "p95_ms": 12.914891624450682, + "min_ms": 12.67115592956543, + "max_ms": 13.001167297363281 + }, + "forward_peak_mib": 528.25, + "out_vs_fp64": { + "max_abs": 0.007803990877593758, + "relative_l2": 0.0016089924958069535 + }, + "lse_vs_fp64": { + "max_abs": 1.1420866119493667e-06, + "relative_l2": 4.112077326317473e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 47.81572151184082, + "p95_ms": 48.06082630157471, + "min_ms": 47.7118034362793, + "max_ms": 48.176414489746094 + }, + "train_peak_mib": 1080.5009765625 + }, + "triton-bitwise": { + "forward": { + "median_ms": 19.420988082885742, + "p95_ms": 19.503754711151124, + "min_ms": 19.371715545654297, + "max_ms": 19.649168014526367 + }, + "forward_peak_mib": 528.25048828125, + "out_vs_fp64": { + "max_abs": 0.007803990877593758, + "relative_l2": 0.0016089924958069535 + }, + "lse_vs_fp64": { + "max_abs": 1.1420866119493667e-06, + "relative_l2": 4.112077326317473e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 76.77885818481445, + "p95_ms": 77.98238182067871, + "min_ms": 76.71944427490234, + "max_ms": 78.67646026611328 + }, + "train_peak_mib": 1080.5009765625 + } + }, + "triton_vs_reference": { + "out_mismatched": 0, + "lse_mismatched": 0, + "out_relative_l2": 0.0, + "bitwise": true + }, + "strict_vs_reference": { + "production_path": "strict-aiter", + "out": { + "max_abs": 0.015625, + "relative_l2": 0.005595395404504914 + }, + "lse": { + "max_abs": 1.9073486328125e-06, + "relative_l2": 6.75179843473521e-08 + }, + "out_mismatched": 6955314 + } + }, + { + "dtype": "bf16", + "seq_len": 4096, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 0.6964554786682129, + "p95_ms": 0.745677000284195, + "min_ms": 0.6806110143661499, + "max_ms": 0.8356419801712036 + }, + "forward_peak_mib": 96.50146484375, + "out_vs_fp64": { + "max_abs": 0.009603632718454325, + "relative_l2": 0.002050744344877689 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 3.216281533241272, + "p95_ms": 3.3038426995277406, + "min_ms": 3.1893410682678223, + "max_ms": 3.316930055618286 + }, + "train_peak_mib": 257.501953125 + }, + "pytorch-native": { + "forward": { + "median_ms": 3.9512734413146973, + "p95_ms": 4.233126997947693, + "min_ms": 3.7510159015655518, + "max_ms": 4.261174201965332 + }, + "forward_peak_mib": 2160.0, + "out_vs_fp64": { + "max_abs": 0.013981263962416168, + "relative_l2": 0.004776104082102705 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 9.36076021194458, + "p95_ms": 9.494323635101319, + "min_ms": 8.979938507080078, + "max_ms": 9.520913124084473 + }, + "train_peak_mib": 4208.0009765625 + }, + "strict-aiter": { + "forward": { + "median_ms": 0.5569075047969818, + "p95_ms": 0.5847899734973907, + "min_ms": 0.5425670146942139, + "max_ms": 0.6106669902801514 + }, + "forward_peak_mib": 112.50048828125, + "out_vs_fp64": { + "max_abs": 0.021384551260978935, + "relative_l2": 0.005439155102039385 + }, + "lse_vs_fp64": { + "max_abs": 3.966678205458152e-06, + "relative_l2": 1.0105794722703805e-07 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 5.726272106170654, + "p95_ms": 5.8219287395477295, + "min_ms": 5.679933071136475, + "max_ms": 5.875733852386475 + }, + "train_peak_mib": 16641.50146484375 + }, + "reference-native": { + "forward": { + "median_ms": 49.35360336303711, + "p95_ms": 49.45325679779052, + "min_ms": 49.06020736694336, + "max_ms": 49.60898208618164 + }, + "forward_peak_mib": 2080.5, + "out_vs_fp64": { + "max_abs": 0.0078084380373617535, + "relative_l2": 0.0016189978158770984 + }, + "lse_vs_fp64": { + "max_abs": 1.3814724333371942e-06, + "relative_l2": 4.407682874169915e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 173.32455444335938, + "p95_ms": 177.7302001953125, + "min_ms": 173.03814697265625, + "max_ms": 177.82281494140625 + }, + "train_peak_mib": 4209.0009765625 + }, + "triton-bitwise": { + "forward": { + "median_ms": 105.40489196777344, + "p95_ms": 110.62110214233398, + "min_ms": 91.67607116699219, + "max_ms": 114.09678649902344 + }, + "forward_peak_mib": 2080.50048828125, + "out_vs_fp64": { + "max_abs": 0.0078084380373617535, + "relative_l2": 0.0016189978158770984 + }, + "lse_vs_fp64": { + "max_abs": 1.3814724333371942e-06, + "relative_l2": 4.407682874169915e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 306.59657287597656, + "p95_ms": 346.2127853393554, + "min_ms": 304.967041015625, + "max_ms": 376.7984619140625 + }, + "train_peak_mib": 4209.0009765625 + } + }, + "triton_vs_reference": { + "out_mismatched": 0, + "lse_mismatched": 0, + "out_relative_l2": 0.0, + "bitwise": true + }, + "strict_vs_reference": { + "production_path": "strict-aiter", + "out": { + "max_abs": 0.015625, + "relative_l2": 0.00563327785605726 + }, + "lse": { + "max_abs": 3.814697265625e-06, + "relative_l2": 1.0616847591543177e-07 + }, + "out_mismatched": 13983861 + } + }, + { + "dtype": "fp16", + "seq_len": 512, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 0.07026449963450432, + "p95_ms": 0.10828655026853087, + "min_ms": 0.06782100349664688, + "max_ms": 0.12818999588489532 + }, + "forward_peak_mib": 12.06396484375, + "out_vs_fp64": { + "max_abs": 0.0010422287323277324, + "relative_l2": 0.00024548880248913073 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.33297500014305115, + "p95_ms": 0.3590096488595009, + "min_ms": 0.29275500774383545, + "max_ms": 0.36005499958992004 + }, + "train_peak_mib": 32.189453125 + }, + "pytorch-native": { + "forward": { + "median_ms": 0.19264650344848633, + "p95_ms": 0.2988662883639336, + "min_ms": 0.13620199263095856, + "max_ms": 0.31526899337768555 + }, + "forward_peak_mib": 44.25, + "out_vs_fp64": { + "max_abs": 0.0018111154495057402, + "relative_l2": 0.0005362038982946119 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.9847039878368378, + "p95_ms": 1.073615401983261, + "min_ms": 0.9752489924430847, + "max_ms": 1.112733006477356 + }, + "train_peak_mib": 76.2509765625 + }, + "strict-aiter": { + "forward": { + "median_ms": 0.27286599576473236, + "p95_ms": 0.2851421862840653, + "min_ms": 0.26487401127815247, + "max_ms": 0.30184701085090637 + }, + "forward_peak_mib": 14.06298828125, + "out_vs_fp64": { + "max_abs": 0.002045633587364648, + "relative_l2": 0.0003824173102168583 + }, + "lse_vs_fp64": { + "max_abs": 8.756207581228637e-07, + "relative_l2": 3.826813085957934e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 1.129139006137848, + "p95_ms": 1.5507168650627137, + "min_ms": 0.7949420213699341, + "max_ms": 1.552988052368164 + }, + "train_peak_mib": 288.18896484375 + }, + "reference-native": { + "forward": { + "median_ms": 0.9919345080852509, + "p95_ms": 1.0174889862537384, + "min_ms": 0.9166420102119446, + "max_ms": 1.068789005279541 + }, + "forward_peak_mib": 36.0625, + "out_vs_fp64": { + "max_abs": 0.000970187372193454, + "relative_l2": 0.00019891305228570232 + }, + "lse_vs_fp64": { + "max_abs": 8.340042594312536e-07, + "relative_l2": 4.5004346410580595e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 3.089252471923828, + "p95_ms": 3.126495563983917, + "min_ms": 2.990485906600952, + "max_ms": 3.1272881031036377 + }, + "train_peak_mib": 78.1259765625 + }, + "triton-bitwise": { + "forward": { + "median_ms": 1.3617845177650452, + "p95_ms": 1.4007869601249694, + "min_ms": 1.3310589790344238, + "max_ms": 1.4392999410629272 + }, + "forward_peak_mib": 36.06298828125, + "out_vs_fp64": { + "max_abs": 0.000970187372193454, + "relative_l2": 0.00019891305228570232 + }, + "lse_vs_fp64": { + "max_abs": 8.340042594312536e-07, + "relative_l2": 4.5004346410580595e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 4.797931909561157, + "p95_ms": 4.859868359565735, + "min_ms": 4.747137069702148, + "max_ms": 4.888025760650635 + }, + "train_peak_mib": 78.1259765625 + } + }, + "triton_vs_reference": { + "out_mismatched": 0, + "lse_mismatched": 0, + "out_relative_l2": 0.0, + "bitwise": true + }, + "strict_vs_reference": { + "production_path": "strict-aiter", + "out": { + "max_abs": 0.001953125, + "relative_l2": 0.00043232633045456683 + }, + "lse": { + "max_abs": 9.5367431640625e-07, + "relative_l2": 4.962085268945435e-08 + }, + "out_mismatched": 1200744 + } + }, + { + "dtype": "fp16", + "seq_len": 1024, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 0.12638749927282333, + "p95_ms": 0.1373237483203411, + "min_ms": 0.11969800293445587, + "max_ms": 0.15939700603485107 + }, + "forward_peak_mib": 24.12646484375, + "out_vs_fp64": { + "max_abs": 0.0010534945195725953, + "relative_l2": 0.000249303688450972 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.41607849299907684, + "p95_ms": 0.47015325427055354, + "min_ms": 0.40576300024986267, + "max_ms": 0.5112000107765198 + }, + "train_peak_mib": 64.376953125 + }, + "pytorch-native": { + "forward": { + "median_ms": 0.31775249540805817, + "p95_ms": 0.3454900071024895, + "min_ms": 0.3053340017795563, + "max_ms": 0.3967899978160858 + }, + "forward_peak_mib": 153.0, + "out_vs_fp64": { + "max_abs": 0.002487331960753014, + "relative_l2": 0.000563456696230733 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.770966500043869, + "p95_ms": 0.847625720500946, + "min_ms": 0.7560039758682251, + "max_ms": 0.8734579682350159 + }, + "train_peak_mib": 281.0009765625 + }, + "strict-aiter": { + "forward": { + "median_ms": 0.28908900916576385, + "p95_ms": 0.3380810514092445, + "min_ms": 0.26848000288009644, + "max_ms": 0.3384239971637726 + }, + "forward_peak_mib": 28.12548828125, + "out_vs_fp64": { + "max_abs": 0.0022991352720684866, + "relative_l2": 0.0003967929347774515 + }, + "lse_vs_fp64": { + "max_abs": 1.2495849039950713e-06, + "relative_l2": 4.3279582438585845e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 0.8596780002117157, + "p95_ms": 0.917380455136299, + "min_ms": 0.8416510224342346, + "max_ms": 0.94132000207901 + }, + "train_peak_mib": 1088.37646484375 + }, + "reference-native": { + "forward": { + "median_ms": 3.066338539123535, + "p95_ms": 3.0974503636360167, + "min_ms": 3.038316011428833, + "max_ms": 3.140949010848999 + }, + "forward_peak_mib": 136.125, + "out_vs_fp64": { + "max_abs": 0.0009756507335656472, + "relative_l2": 0.0002006938787806314 + }, + "lse_vs_fp64": { + "max_abs": 1.0114761641588643e-06, + "relative_l2": 4.218427372109014e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 11.847110748291016, + "p95_ms": 11.933381175994873, + "min_ms": 11.733922004699707, + "max_ms": 11.938265800476074 + }, + "train_peak_mib": 284.2509765625 + }, + "triton-bitwise": { + "forward": { + "median_ms": 4.7754786014556885, + "p95_ms": 4.8063427925109865, + "min_ms": 4.759194850921631, + "max_ms": 4.921235084533691 + }, + "forward_peak_mib": 136.12548828125, + "out_vs_fp64": { + "max_abs": 0.0009756507335656472, + "relative_l2": 0.0002006938787806314 + }, + "lse_vs_fp64": { + "max_abs": 1.0114761641588643e-06, + "relative_l2": 4.218427372109014e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 19.591503143310547, + "p95_ms": 19.784181976318358, + "min_ms": 19.13556671142578, + "max_ms": 19.899662017822266 + }, + "train_peak_mib": 284.2509765625 + } + }, + "triton_vs_reference": { + "out_mismatched": 0, + "lse_mismatched": 0, + "out_relative_l2": 0.0, + "bitwise": true + }, + "strict_vs_reference": { + "production_path": "strict-aiter", + "out": { + "max_abs": 0.001953125, + "relative_l2": 0.00044664733527359724 + }, + "lse": { + "max_abs": 1.430511474609375e-06, + "relative_l2": 5.213412779092794e-08 + }, + "out_mismatched": 2434520 + } + }, + { + "dtype": "fp16", + "seq_len": 2048, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 0.2896910011768341, + "p95_ms": 0.2998181506991387, + "min_ms": 0.2804969847202301, + "max_ms": 0.3403860032558441 + }, + "forward_peak_mib": 48.25146484375, + "out_vs_fp64": { + "max_abs": 0.000909865866452364, + "relative_l2": 0.0002537678065402769 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 1.1873445510864258, + "p95_ms": 1.598525464534759, + "min_ms": 1.1201050281524658, + "max_ms": 1.8953360319137573 + }, + "train_peak_mib": 128.501953125 + }, + "pytorch-native": { + "forward": { + "median_ms": 1.0827745199203491, + "p95_ms": 1.113287901878357, + "min_ms": 1.071632981300354, + "max_ms": 1.1308410167694092 + }, + "forward_peak_mib": 564.0, + "out_vs_fp64": { + "max_abs": 0.0017265887526813906, + "relative_l2": 0.0005855639752067646 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 2.3896775245666504, + "p95_ms": 2.9295324802398675, + "min_ms": 2.3408498764038086, + "max_ms": 3.3291189670562744 + }, + "train_peak_mib": 1076.0009765625 + }, + "strict-aiter": { + "forward": { + "median_ms": 0.299125000834465, + "p95_ms": 0.3336741864681244, + "min_ms": 0.29556000232696533, + "max_ms": 0.3492389917373657 + }, + "forward_peak_mib": 56.25048828125, + "out_vs_fp64": { + "max_abs": 0.002052942593323337, + "relative_l2": 0.000397435773927899 + }, + "lse_vs_fp64": { + "max_abs": 2.539945519686171e-06, + "relative_l2": 6.093248379166392e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 1.7668060064315796, + "p95_ms": 1.795847511291504, + "min_ms": 1.7248040437698364, + "max_ms": 1.7987140417099 + }, + "train_peak_mib": 4224.75146484375 + }, + "reference-native": { + "forward": { + "median_ms": 12.56151294708252, + "p95_ms": 12.683987283706665, + "min_ms": 12.459760665893555, + "max_ms": 12.691065788269043 + }, + "forward_peak_mib": 528.25, + "out_vs_fp64": { + "max_abs": 0.000909865866452364, + "relative_l2": 0.00020150767732632786 + }, + "lse_vs_fp64": { + "max_abs": 1.102904860772469e-06, + "relative_l2": 4.146891383633236e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 47.40879440307617, + "p95_ms": 47.5834342956543, + "min_ms": 47.31221008300781, + "max_ms": 47.60921096801758 + }, + "train_peak_mib": 1080.5009765625 + }, + "triton-bitwise": { + "forward": { + "median_ms": 19.23597526550293, + "p95_ms": 19.307478046417238, + "min_ms": 19.19112777709961, + "max_ms": 19.76598358154297 + }, + "forward_peak_mib": 528.25048828125, + "out_vs_fp64": { + "max_abs": 0.000909865866452364, + "relative_l2": 0.00020150767732632786 + }, + "lse_vs_fp64": { + "max_abs": 1.102904860772469e-06, + "relative_l2": 4.146891383633236e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 75.97890853881836, + "p95_ms": 76.2766300201416, + "min_ms": 75.94613647460938, + "max_ms": 76.34769439697266 + }, + "train_peak_mib": 1080.5009765625 + } + }, + "triton_vs_reference": { + "out_mismatched": 0, + "lse_mismatched": 0, + "out_relative_l2": 0.0, + "bitwise": true + }, + "strict_vs_reference": { + "production_path": "strict-aiter", + "out": { + "max_abs": 0.001953125, + "relative_l2": 0.00045030041010727414 + }, + "lse": { + "max_abs": 2.86102294921875e-06, + "relative_l2": 6.734176413189059e-08 + }, + "out_mismatched": 4922569 + } + }, + { + "dtype": "fp16", + "seq_len": 4096, + "batch": 1, + "q_heads": 32, + "kv_heads": 8, + "head_dim": 128, + "paths": { + "sdpa": { + "forward": { + "median_ms": 1.0774019956588745, + "p95_ms": 1.102609133720398, + "min_ms": 1.0514429807662964, + "max_ms": 1.1050820350646973 + }, + "forward_peak_mib": 96.50146484375, + "out_vs_fp64": { + "max_abs": 0.0009904582683271101, + "relative_l2": 0.0002562693823501264 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 3.97100293636322, + "p95_ms": 4.093861031532287, + "min_ms": 3.9423000812530518, + "max_ms": 4.121045112609863 + }, + "train_peak_mib": 257.001953125 + }, + "pytorch-native": { + "forward": { + "median_ms": 3.9607179164886475, + "p95_ms": 4.2144136190414425, + "min_ms": 3.68410587310791, + "max_ms": 4.284419059753418 + }, + "forward_peak_mib": 2160.0, + "out_vs_fp64": { + "max_abs": 0.0023392191337636703, + "relative_l2": 0.0006005882917972107 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 9.475195407867432, + "p95_ms": 9.788628911972046, + "min_ms": 8.91443157196045, + "max_ms": 9.886946678161621 + }, + "train_peak_mib": 4208.0009765625 + }, + "strict-aiter": { + "forward": { + "median_ms": 0.5824664831161499, + "p95_ms": 0.6701211005449296, + "min_ms": 0.569337010383606, + "max_ms": 0.7219929695129395 + }, + "forward_peak_mib": 112.50048828125, + "out_vs_fp64": { + "max_abs": 0.001953125, + "relative_l2": 0.00040417579976339364 + }, + "lse_vs_fp64": { + "max_abs": 3.8052896318419016e-06, + "relative_l2": 1.012403669168641e-07 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 5.9359588623046875, + "p95_ms": 5.9734298467636116, + "min_ms": 5.836415767669678, + "max_ms": 5.981771945953369 + }, + "train_peak_mib": 16641.50146484375 + }, + "reference-native": { + "forward": { + "median_ms": 48.57136344909668, + "p95_ms": 48.678818702697754, + "min_ms": 48.476505279541016, + "max_ms": 48.689056396484375 + }, + "forward_peak_mib": 2080.5, + "out_vs_fp64": { + "max_abs": 0.0009680650127679158, + "relative_l2": 0.00020199318189363435 + }, + "lse_vs_fp64": { + "max_abs": 1.5106917023999245e-06, + "relative_l2": 4.425631521647521e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 171.57088470458984, + "p95_ms": 171.70842437744142, + "min_ms": 171.45767211914062, + "max_ms": 171.72274780273438 + }, + "train_peak_mib": 4209.0009765625 + }, + "triton-bitwise": { + "forward": { + "median_ms": 84.34748458862305, + "p95_ms": 86.58723831176758, + "min_ms": 83.29586029052734, + "max_ms": 87.01377868652344 + }, + "forward_peak_mib": 2080.50048828125, + "out_vs_fp64": { + "max_abs": 0.0009680650127679158, + "relative_l2": 0.00020199318189363435 + }, + "lse_vs_fp64": { + "max_abs": 1.5106917023999245e-06, + "relative_l2": 4.425631521647521e-08 + }, + "repeat_bitwise": true, + "train_fwd_bwd": { + "median_ms": 301.26148986816406, + "p95_ms": 302.59836730957034, + "min_ms": 288.032958984375, + "max_ms": 303.03350830078125 + }, + "train_peak_mib": 4209.0009765625 + } + }, + "triton_vs_reference": { + "out_mismatched": 0, + "lse_mismatched": 0, + "out_relative_l2": 0.0, + "bitwise": true + }, + "strict_vs_reference": { + "production_path": "strict-aiter", + "out": { + "max_abs": 0.001953125, + "relative_l2": 0.0004580324082697752 + }, + "lse": { + "max_abs": 3.814697265625e-06, + "relative_l2": 1.0641371875326284e-07 + }, + "out_mismatched": 9932656 + } + } + ] + }, + "backward_parity": [ + { + "seq_len": 512, + "dq_mismatched": 0, + "dk_mismatched": 0, + "dv_mismatched": 0, + "bitwise": true + }, + { + "seq_len": 1024, + "dq_mismatched": 0, + "dk_mismatched": 0, + "dv_mismatched": 0, + "bitwise": true + }, + { + "seq_len": 2048, + "dq_mismatched": 0, + "dk_mismatched": 0, + "dv_mismatched": 0, + "bitwise": true + }, + { + "seq_len": 4096, + "dq_mismatched": 0, + "dk_mismatched": 0, + "dv_mismatched": 0, + "bitwise": true + } + ], + "batch_composition": [ + { + "seq_len": 512, + "paths": { + "sdpa": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "pytorch-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "strict-aiter": { + "batch_gt1_rejected": true, + "out_bitwise": true, + "out_mismatched": 0, + "note": "core executes one logical batch row per launch" + }, + "reference-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": true + }, + "triton-bitwise": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": true + } + } + }, + { + "seq_len": 1024, + "paths": { + "sdpa": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "pytorch-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "strict-aiter": { + "batch_gt1_rejected": true, + "out_bitwise": true, + "out_mismatched": 0, + "note": "core executes one logical batch row per launch" + }, + "reference-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": true + }, + "triton-bitwise": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": true + } + } + }, + { + "seq_len": 2048, + "paths": { + "sdpa": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "pytorch-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "strict-aiter": { + "batch_gt1_rejected": true, + "out_bitwise": true, + "out_mismatched": 0, + "note": "core executes one logical batch row per launch" + }, + "reference-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": true + }, + "triton-bitwise": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": true + } + } + }, + { + "seq_len": 4096, + "paths": { + "sdpa": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "pytorch-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": null + }, + "strict-aiter": { + "batch_gt1_rejected": true, + "out_bitwise": true, + "out_mismatched": 0, + "note": "core executes one logical batch row per launch" + }, + "reference-native": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": true + }, + "triton-bitwise": { + "batch_gt1_rejected": false, + "out_bitwise": true, + "out_mismatched": 0, + "out_max_abs": 0.0, + "lse_bitwise": true + } + } + } + ], + "tp_head_sensitivity": [ + { + "seq_len": 512, + "schedule": "raw_launch", + "tp": 2, + "local_q_heads": 16, + "local_kv_heads": 4, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 512, + "schedule": "raw_launch", + "tp": 4, + "local_q_heads": 8, + "local_kv_heads": 2, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 512, + "schedule": "raw_launch", + "tp": 8, + "local_q_heads": 4, + "local_kv_heads": 1, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 512, + "schedule": "one_kv_group_per_launch", + "tp": 2, + "local_q_heads": 16, + "local_kv_heads": 4, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 512, + "schedule": "one_kv_group_per_launch", + "tp": 4, + "local_q_heads": 8, + "local_kv_heads": 2, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 512, + "schedule": "one_kv_group_per_launch", + "tp": 8, + "local_q_heads": 4, + "local_kv_heads": 1, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 1024, + "schedule": "raw_launch", + "tp": 2, + "local_q_heads": 16, + "local_kv_heads": 4, + "out_max_abs": 0.0078125, + "lse_max_abs": 1.9073486328125e-06, + "invariant": false + }, + { + "seq_len": 1024, + "schedule": "raw_launch", + "tp": 4, + "local_q_heads": 8, + "local_kv_heads": 2, + "out_max_abs": 0.0078125, + "lse_max_abs": 1.9073486328125e-06, + "invariant": false + }, + { + "seq_len": 1024, + "schedule": "raw_launch", + "tp": 8, + "local_q_heads": 4, + "local_kv_heads": 1, + "out_max_abs": 0.0078125, + "lse_max_abs": 1.9073486328125e-06, + "invariant": false + }, + { + "seq_len": 1024, + "schedule": "one_kv_group_per_launch", + "tp": 2, + "local_q_heads": 16, + "local_kv_heads": 4, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 1024, + "schedule": "one_kv_group_per_launch", + "tp": 4, + "local_q_heads": 8, + "local_kv_heads": 2, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 1024, + "schedule": "one_kv_group_per_launch", + "tp": 8, + "local_q_heads": 4, + "local_kv_heads": 1, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 2048, + "schedule": "raw_launch", + "tp": 2, + "local_q_heads": 16, + "local_kv_heads": 4, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 2048, + "schedule": "raw_launch", + "tp": 4, + "local_q_heads": 8, + "local_kv_heads": 2, + "out_max_abs": 0.00390625, + "lse_max_abs": 2.86102294921875e-06, + "invariant": false + }, + { + "seq_len": 2048, + "schedule": "raw_launch", + "tp": 8, + "local_q_heads": 4, + "local_kv_heads": 1, + "out_max_abs": 0.001953125, + "lse_max_abs": 2.86102294921875e-06, + "invariant": false + }, + { + "seq_len": 2048, + "schedule": "one_kv_group_per_launch", + "tp": 2, + "local_q_heads": 16, + "local_kv_heads": 4, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 2048, + "schedule": "one_kv_group_per_launch", + "tp": 4, + "local_q_heads": 8, + "local_kv_heads": 2, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 2048, + "schedule": "one_kv_group_per_launch", + "tp": 8, + "local_q_heads": 4, + "local_kv_heads": 1, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 4096, + "schedule": "raw_launch", + "tp": 2, + "local_q_heads": 16, + "local_kv_heads": 4, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 4096, + "schedule": "raw_launch", + "tp": 4, + "local_q_heads": 8, + "local_kv_heads": 2, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 4096, + "schedule": "raw_launch", + "tp": 8, + "local_q_heads": 4, + "local_kv_heads": 1, + "out_max_abs": 0.00390625, + "lse_max_abs": 4.76837158203125e-06, + "invariant": false + }, + { + "seq_len": 4096, + "schedule": "one_kv_group_per_launch", + "tp": 2, + "local_q_heads": 16, + "local_kv_heads": 4, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 4096, + "schedule": "one_kv_group_per_launch", + "tp": 4, + "local_q_heads": 8, + "local_kv_heads": 2, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + }, + { + "seq_len": 4096, + "schedule": "one_kv_group_per_launch", + "tp": 8, + "local_q_heads": 4, + "local_kv_heads": 1, + "out_max_abs": 0.0, + "lse_max_abs": 0.0, + "invariant": true + } + ], + "distributed": [ + { + "topology": "tp1_cp2", + "world_size": 2, + "tp_world_size": 1, + "cp_world_size": 2, + "replicas": 1, + "seq_len": 4096, + "local_q_heads": 32, + "local_kv_heads": 8, + "transport": "rccl_ag_rs", + "forward": { + "median_ms": 1.837871491909027, + "p95_ms": 1.9000428080558778, + "min_ms": 1.8279180526733398, + "max_ms": 1.91103994846344 + }, + "cp1_baseline": { + "median_ms": 0.566241979598999, + "p95_ms": 0.586542186141014, + "min_ms": 0.561115026473999, + "max_ms": 0.5893959999084473 + }, + "peak_mib_per_rank": 160.53173828125, + "out_bitwise_vs_cp1": true, + "lse_bitwise_vs_cp1": true, + "repeat_bitwise": true, + "out_mismatched_all_ranks": 0, + "lse_mismatched_all_ranks": 0 + }, + { + "topology": "tp2_cp2", + "world_size": 4, + "tp_world_size": 2, + "cp_world_size": 2, + "replicas": 1, + "seq_len": 4096, + "local_q_heads": 16, + "local_kv_heads": 4, + "transport": "rccl_ag_rs", + "forward": { + "median_ms": 1.2287670373916626, + "p95_ms": 1.2961051762104034, + "min_ms": 1.1803940534591675, + "max_ms": 1.337548017501831 + }, + "cp1_baseline": { + "median_ms": 0.39462698996067047, + "p95_ms": 0.5805545553565048, + "min_ms": 0.3872550129890442, + "max_ms": 3.5400619506835938 + }, + "peak_mib_per_rank": 80.28173828125, + "out_bitwise_vs_cp1": true, + "lse_bitwise_vs_cp1": true, + "repeat_bitwise": true, + "out_mismatched_all_ranks": 0, + "lse_mismatched_all_ranks": 0 + }, + { + "topology": "tp1_cp4", + "world_size": 4, + "tp_world_size": 1, + "cp_world_size": 4, + "replicas": 1, + "seq_len": 4096, + "local_q_heads": 32, + "local_kv_heads": 8, + "transport": "rccl_ag_rs", + "forward": { + "median_ms": 1.398799479007721, + "p95_ms": 1.4461318969726564, + "min_ms": 1.3782479763031006, + "max_ms": 1.5850759744644165 + }, + "cp1_baseline": { + "median_ms": 0.5599325001239777, + "p95_ms": 0.5996211320161821, + "min_ms": 0.5533829927444458, + "max_ms": 0.6393910050392151 + }, + "peak_mib_per_rank": 160.53173828125, + "out_bitwise_vs_cp1": true, + "lse_bitwise_vs_cp1": true, + "repeat_bitwise": true, + "out_mismatched_all_ranks": 0, + "lse_mismatched_all_ranks": 0 + }, + { + "topology": "tp2_cp2_x2", + "world_size": 8, + "tp_world_size": 2, + "cp_world_size": 2, + "replicas": 2, + "seq_len": 4096, + "local_q_heads": 16, + "local_kv_heads": 4, + "transport": "rccl_ag_rs", + "forward": { + "median_ms": 1.2808839678764343, + "p95_ms": 3.3578428864479166, + "min_ms": 1.1812349557876587, + "max_ms": 16.83999252319336 + }, + "cp1_baseline": { + "median_ms": 0.393885001540184, + "p95_ms": 0.4397124022245408, + "min_ms": 0.38765600323677063, + "max_ms": 0.5222560167312622 + }, + "peak_mib_per_rank": 80.28173828125, + "out_bitwise_vs_cp1": true, + "lse_bitwise_vs_cp1": true, + "repeat_bitwise": true, + "out_mismatched_all_ranks": 0, + "lse_mismatched_all_ranks": 0 + }, + { + "topology": "tp2_cp4", + "world_size": 8, + "tp_world_size": 2, + "cp_world_size": 4, + "replicas": 1, + "seq_len": 4096, + "local_q_heads": 16, + "local_kv_heads": 4, + "transport": "rccl_ag_rs", + "forward": { + "median_ms": 2.353678584098816, + "p95_ms": 6.935053753852847, + "min_ms": 1.2289060354232788, + "max_ms": 10.850227355957031 + }, + "cp1_baseline": { + "median_ms": 0.39695000648498535, + "p95_ms": 1.3903744220733647, + "min_ms": 0.390980988740921, + "max_ms": 2.1617438793182373 + }, + "peak_mib_per_rank": 80.28173828125, + "out_bitwise_vs_cp1": true, + "lse_bitwise_vs_cp1": true, + "repeat_bitwise": true, + "out_mismatched_all_ranks": 0, + "lse_mismatched_all_ranks": 0 + }, + { + "topology": "tp1_cp8", + "world_size": 8, + "tp_world_size": 1, + "cp_world_size": 8, + "replicas": 1, + "seq_len": 4096, + "local_q_heads": 32, + "local_kv_heads": 8, + "transport": "rccl_ag_rs", + "forward": { + "median_ms": 3.06355357170105, + "p95_ms": 35.86110134124757, + "min_ms": 1.493980050086975, + "max_ms": 43.65816879272461 + }, + "cp1_baseline": { + "median_ms": 0.5652805268764496, + "p95_ms": 1.8745501130819493, + "min_ms": 0.5424069762229919, + "max_ms": 24.681163787841797 + }, + "peak_mib_per_rank": 160.53173828125, + "out_bitwise_vs_cp1": true, + "lse_bitwise_vs_cp1": true, + "repeat_bitwise": true, + "out_mismatched_all_ranks": 0, + "lse_mismatched_all_ranks": 0 + } + ], + "tp_schedule_cost": [ + { + "seq_len": 512, + "launches": 8, + "raw_launch": { + "median_ms": 0.2578835040330887, + "p95_ms": 0.3862708032131196, + "min_ms": 0.2440830022096634, + "max_ms": 0.4992220103740692 + }, + "raw_launch_peak_mib": 14.06298828125, + "one_kv_group_per_launch": { + "median_ms": 1.775919497013092, + "p95_ms": 2.0744626879692083, + "min_ms": 1.7243629693984985, + "max_ms": 2.772800922393799 + }, + "one_kv_group_per_launch_peak_mib": 8.125, + "sdpa": { + "median_ms": 0.07118599861860275, + "p95_ms": 0.0796682476997376, + "min_ms": 0.069302998483181, + "max_ms": 0.15939700603485107 + } + }, + { + "seq_len": 1024, + "launches": 8, + "raw_launch": { + "median_ms": 0.25129400193691254, + "p95_ms": 0.2852045923471451, + "min_ms": 0.23883499205112457, + "max_ms": 0.3441919982433319 + }, + "raw_launch_peak_mib": 28.12548828125, + "one_kv_group_per_launch": { + "median_ms": 1.9984899759292603, + "p95_ms": 2.597748446464539, + "min_ms": 1.709501028060913, + "max_ms": 2.769155979156494 + }, + "one_kv_group_per_launch_peak_mib": 16.25, + "sdpa": { + "median_ms": 0.13017350435256958, + "p95_ms": 0.1343752935528755, + "min_ms": 0.12150000035762787, + "max_ms": 0.13772499561309814 + } + }, + { + "seq_len": 2048, + "launches": 8, + "raw_launch": { + "median_ms": 0.2916930019855499, + "p95_ms": 0.3531085982918741, + "min_ms": 0.2873469889163971, + "max_ms": 0.5141639709472656 + }, + "raw_launch_peak_mib": 56.25048828125, + "one_kv_group_per_launch": { + "median_ms": 1.7506614923477173, + "p95_ms": 1.8041546821594239, + "min_ms": 1.7337770462036133, + "max_ms": 1.838291049003601 + }, + "one_kv_group_per_launch_peak_mib": 32.5, + "sdpa": { + "median_ms": 0.2801560014486313, + "p95_ms": 0.2873334392905235, + "min_ms": 0.27601000666618347, + "max_ms": 0.3060950040817261 + } + }, + { + "seq_len": 4096, + "launches": 8, + "raw_launch": { + "median_ms": 0.568244993686676, + "p95_ms": 0.6004684329032898, + "min_ms": 0.5574290156364441, + "max_ms": 0.6061009764671326 + }, + "raw_launch_peak_mib": 112.50048828125, + "one_kv_group_per_launch": { + "median_ms": 2.047242522239685, + "p95_ms": 2.1095147371292113, + "min_ms": 2.00606107711792, + "max_ms": 2.118267059326172 + }, + "one_kv_group_per_launch_peak_mib": 65.0, + "sdpa": { + "median_ms": 0.7046480178833008, + "p95_ms": 0.722521898150444, + "min_ms": 0.6748430132865906, + "max_ms": 0.7470300197601318 + } + } + ] +} \ No newline at end of file diff --git a/benchmarks/results/ws2_rocm_mi300x/single_gpu_grid.png b/benchmarks/results/ws2_rocm_mi300x/single_gpu_grid.png new file mode 100644 index 00000000..04cbefcd Binary files /dev/null and b/benchmarks/results/ws2_rocm_mi300x/single_gpu_grid.png differ diff --git a/benchmarks/results/ws2_rocm_mi300x/single_gpu_latency.png b/benchmarks/results/ws2_rocm_mi300x/single_gpu_latency.png new file mode 100644 index 00000000..6cf462b0 Binary files /dev/null and b/benchmarks/results/ws2_rocm_mi300x/single_gpu_latency.png differ diff --git a/benchmarks/results/ws2_rocm_mi300x/single_gpu_memory.png b/benchmarks/results/ws2_rocm_mi300x/single_gpu_memory.png new file mode 100644 index 00000000..e14504e2 Binary files /dev/null and b/benchmarks/results/ws2_rocm_mi300x/single_gpu_memory.png differ diff --git a/benchmarks/results/ws2_rocm_mi300x/tp_degree_invariance.png b/benchmarks/results/ws2_rocm_mi300x/tp_degree_invariance.png new file mode 100644 index 00000000..4d28bbf7 Binary files /dev/null and b/benchmarks/results/ws2_rocm_mi300x/tp_degree_invariance.png differ diff --git a/csrc/cuda/attention/deterministic_attention.cu b/csrc/cuda/attention/deterministic_attention.cu index aaa70b42..2c56d6c7 100644 --- a/csrc/cuda/attention/deterministic_attention.cu +++ b/csrc/cuda/attention/deterministic_attention.cu @@ -685,3 +685,90 @@ std::vector deterministic_attention_backward( return {dQ, dK, dV}; } + +#if defined(USE_ROCM) +namespace { + +template +__global__ void deterministic_rope_kernel( + const scalar_t* __restrict__ x, + const float* __restrict__ cos, + const float* __restrict__ sin, + scalar_t* __restrict__ out, + int64_t n_rows, + int table_rows, + int half, + float sin_sign) { + const int64_t index = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; + const int64_t count = n_rows * static_cast(half); + if (index >= count) { + return; + } + + const int64_t row = index / half; + const int pair = static_cast(index % half); + const int table_row = static_cast(row % table_rows); + const float c = cos[table_row * half + pair]; + const float s = sin[table_row * half + pair] * sin_sign; + const int64_t base = row * (2LL * half); + const float low = static_cast(x[base + pair]); + const float high = static_cast(x[base + pair + half]); + + out[base + pair] = static_cast(low * c - high * s); + out[base + pair + half] = static_cast(high * c + low * s); +} + +} // namespace + +torch::Tensor deterministic_rope_apply_rocm( + torch::Tensor x, + torch::Tensor cos, + torch::Tensor sin, + double sin_sign) { + TORCH_CHECK(x.is_cuda(), "ROCm RoPE: x must be a GPU tensor"); + TORCH_CHECK(x.dim() == 2 && x.is_contiguous(), + "ROCm RoPE: x must be contiguous [rows, head_dim]"); + TORCH_CHECK(cos.is_cuda() && sin.is_cuda(), + "ROCm RoPE: cos and sin must be GPU tensors"); + TORCH_CHECK(cos.scalar_type() == torch::kFloat32 && + sin.scalar_type() == torch::kFloat32, + "ROCm RoPE: cos and sin must be FP32"); + TORCH_CHECK(cos.is_contiguous() && sin.is_contiguous(), + "ROCm RoPE: cos and sin must be contiguous"); + TORCH_CHECK(cos.dim() == 2 && sin.sizes() == cos.sizes(), + "ROCm RoPE: cos and sin must have shape [table_rows, head_dim/2]"); + TORCH_CHECK(x.size(1) % 2 == 0, "ROCm RoPE: head_dim must be even"); + TORCH_CHECK(cos.size(0) > 0 && cos.size(1) == x.size(1) / 2, + "ROCm RoPE: invalid cos/sin table shape"); + TORCH_CHECK(x.size(0) % cos.size(0) == 0, + "ROCm RoPE: row count must be divisible by the position table size"); + + const at::cuda::OptionalCUDAGuard guard(device_of(x)); + auto out = torch::empty_like(x); + const int64_t n_rows = x.size(0); + const int half = static_cast(x.size(1) / 2); + const int table_rows = static_cast(cos.size(0)); + const int64_t count = n_rows * static_cast(half); + constexpr int threads = 256; + const int64_t blocks = (count + threads - 1) / threads; + auto stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + x.scalar_type(), + "deterministic_rope_apply_rocm", + [&] { + deterministic_rope_kernel<<>>( + x.data_ptr(), + cos.data_ptr(), + sin.data_ptr(), + out.data_ptr(), + n_rows, + table_rows, + half, + static_cast(sin_sign)); + }); + return out; +} +#endif diff --git a/csrc/ops.cpp b/csrc/ops.cpp index aecc30ed..883af5ba 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -78,7 +78,7 @@ torch::Tensor lm_head_sm90_forward_fp32(torch::Tensor hidden, torch::optional bias); #endif -#if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) +#if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) || defined(KERNEL_ALIGN_WITH_ROCM) torch::Tensor fused_logp_forward_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor output); torch::Tensor fused_logp_forward_fp32(torch::Tensor logits, torch::Tensor token_ids); torch::Tensor fused_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); @@ -93,7 +93,9 @@ torch::Tensor deterministic_logp_forward_fp32(torch::Tensor logits, torch::Tenso torch::Tensor deterministic_logp_forward_indexed_out(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices, torch::Tensor output); torch::Tensor deterministic_logp_forward_indexed_fp32(torch::Tensor logits, torch::Tensor token_ids, torch::Tensor row_indices); -// Single-node TP=8 deterministic collectives. +#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) +// Single-node TP=8 deterministic CUDA IPC collectives. ROCm uses the +// rank-ordered RCCL transport in rl_engine.distributed.collectives. std::tuple, int64_t> deterministic_collective_ipc_meta( torch::Tensor& tensor); int64_t deterministic_collective_create( @@ -110,6 +112,55 @@ void deterministic_collective_reduce_scatter(int64_t handle, torch::Tensor& outp void deterministic_collective_all_gather(int64_t handle, torch::Tensor& output); void deterministic_collective_all_gather_fused( int64_t handle, torch::Tensor& input, torch::Tensor& output); +#endif + +#if defined(KERNEL_ALIGN_WITH_ROCM) +// ROCm keeps arithmetic in a fixed balanced tree while using either RCCL or +// HIP IPC for rank-ordered transport. These kernels expose the local and IPC +// reduction paths without changing the CUDA implementation. +void deterministic_collective_rocm_all_reduce( + torch::Tensor rank_inputs, + torch::Tensor output); +void deterministic_collective_rocm_reduce_scatter( + torch::Tensor rank_inputs, + torch::Tensor output); +torch::Tensor deterministic_collective_rocm_ipc_allocate(int64_t size_bytes); +std::tuple, int64_t> +deterministic_collective_rocm_ipc_meta(torch::Tensor tensor); +int64_t deterministic_collective_rocm_ipc_create( + torch::Tensor staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank); +void deterministic_collective_rocm_ipc_synchronize(int64_t handle); +void deterministic_collective_rocm_ipc_destroy(int64_t handle); +void deterministic_collective_rocm_ipc_stage(int64_t handle, torch::Tensor input); +void deterministic_collective_rocm_ipc_all_reduce( + int64_t handle, + torch::Tensor output); +void deterministic_collective_rocm_ipc_all_reduce_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output); +void deterministic_collective_rocm_ipc_reduce_scatter( + int64_t handle, + torch::Tensor output); +void deterministic_collective_rocm_ipc_reduce_scatter_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output); +void deterministic_collective_rocm_ipc_reduce_scatter_many( + int64_t handle, + const std::vector& inputs, + const std::vector& outputs); +void deterministic_collective_rocm_ipc_all_gather( + int64_t handle, + torch::Tensor output); +void deterministic_collective_rocm_ipc_all_gather_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output); +#endif // Batch-Invariant Deterministic GEMM Declarations bool det_gemm_sm90_compiled(); @@ -312,9 +363,17 @@ std::vector deterministic_attention_backward( double scale, torch::optional key_padding_mask); -// Prefix-Shared Attention Declarations & Wrappers +#if defined(KERNEL_ALIGN_WITH_ROCM) +torch::Tensor deterministic_rope_apply_rocm( + torch::Tensor x, + torch::Tensor cos, + torch::Tensor sin, + double sin_sign); +#endif #if !defined(USE_ROCM) +// Prefix-Shared Attention Declarations & Wrappers (NVIDIA PTX only). + void prefix_shared_attention_forward( const __nv_bfloat16 *Q, // [bs, G, len_q, DIM] const __nv_bfloat16 *K, // [bs, len_kv, DIM] @@ -410,7 +469,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Single-card SM90 batch-invariant LM-head forward with fp32 output"); #endif -#if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) +#if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) || defined(KERNEL_ALIGN_WITH_ROCM) m.def("fused_logp_forward_out", &fused_logp_forward_out, "Fused logp out"); m.def("fused_logp_forward_fp32", &fused_logp_forward_fp32, "Fused logp fp32"); m.def("fused_logp_forward_indexed_out", &fused_logp_forward_indexed_out, "Fused logp indexed out"); @@ -425,7 +484,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("deterministic_logp_forward_indexed_out", &deterministic_logp_forward_indexed_out, "Batch-invariant deterministic logp indexed out"); m.def("deterministic_logp_forward_indexed_fp32", &deterministic_logp_forward_indexed_fp32, "Batch-invariant deterministic logp indexed fp32"); - // Single-node TP=8 fixed-tree collectives. +#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) + // Single-node TP=8 fixed-tree CUDA IPC collectives. ROCm dispatches to + // the Python RCCL transport implementation instead. m.def( "deterministic_collective_ipc_meta", &deterministic_collective_ipc_meta, @@ -462,9 +523,62 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "deterministic_collective_all_gather_fused", &deterministic_collective_all_gather_fused, "Run a fused small-message deterministic rank-ordered all-gather"); +#endif + +#if defined(KERNEL_ALIGN_WITH_ROCM) + m.def( + "deterministic_collective_rocm_all_reduce", + &deterministic_collective_rocm_all_reduce, + "Run the ROCm fixed-tree all-reduce kernel"); + m.def( + "deterministic_collective_rocm_reduce_scatter", + &deterministic_collective_rocm_reduce_scatter, + "Run the ROCm fixed-tree reduce-scatter kernel"); + m.def("deterministic_collective_rocm_ipc_meta", + &deterministic_collective_rocm_ipc_meta, + "Export a ROCm allocation for IPC deterministic collectives"); + m.def("deterministic_collective_rocm_ipc_allocate", + &deterministic_collective_rocm_ipc_allocate, + "Allocate ROCm memory that supports IPC export"); + m.def("deterministic_collective_rocm_ipc_create", + &deterministic_collective_rocm_ipc_create, + "Create a ROCm IPC deterministic collective state"); + m.def("deterministic_collective_rocm_ipc_destroy", + &deterministic_collective_rocm_ipc_destroy, + "Destroy a ROCm IPC deterministic collective state"); + m.def("deterministic_collective_rocm_ipc_synchronize", + &deterministic_collective_rocm_ipc_synchronize, + "Wait until every rank finishes reading ROCm IPC staging"); + m.def("deterministic_collective_rocm_ipc_stage", + &deterministic_collective_rocm_ipc_stage, + "Stage an input for ROCm IPC deterministic collectives"); + m.def("deterministic_collective_rocm_ipc_all_reduce", + &deterministic_collective_rocm_ipc_all_reduce, + "Run a direct ROCm IPC fixed-tree all-reduce"); + m.def("deterministic_collective_rocm_ipc_all_reduce_input", + &deterministic_collective_rocm_ipc_all_reduce_input, + "Stage and run a direct ROCm IPC fixed-tree all-reduce"); + m.def("deterministic_collective_rocm_ipc_reduce_scatter", + &deterministic_collective_rocm_ipc_reduce_scatter, + "Run a direct ROCm IPC fixed-tree reduce-scatter"); + m.def("deterministic_collective_rocm_ipc_reduce_scatter_input", + &deterministic_collective_rocm_ipc_reduce_scatter_input, + "Stage and run a direct ROCm IPC fixed-tree reduce-scatter"); + m.def("deterministic_collective_rocm_ipc_reduce_scatter_many", + &deterministic_collective_rocm_ipc_reduce_scatter_many, + "Run multiple ROCm IPC fixed-tree reduce-scatters with one synchronization"); + m.def("deterministic_collective_rocm_ipc_all_gather", + &deterministic_collective_rocm_ipc_all_gather, + "Run a direct ROCm IPC rank-ordered all-gather"); + m.def("deterministic_collective_rocm_ipc_all_gather_input", + &deterministic_collective_rocm_ipc_all_gather_input, + "Stage and run a direct ROCm IPC rank-ordered all-gather"); +#endif - // Prefix-shared attention uses NVIDIA PTX and falls back to PyTorch SDPA on ROCm. #if !defined(USE_ROCM) + // Prefix-shared attention uses NVIDIA PTX; the declaration above carries the + // same guard, so the registration must repeat it or a ROCm build fails on an + // undeclared identifier. m.def("prefix_shared_attention", &prefix_shared_attention, "Prefix-Shared Fused Attention for GRPO"); #endif @@ -512,5 +626,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "deterministic_attention_backward", &deterministic_attention_backward, "Deterministic standard softmax attention backward (dQ, dK, dV)"); +#if defined(KERNEL_ALIGN_WITH_ROCM) + m.def( + "deterministic_rope_apply_rocm", + &deterministic_rope_apply_rocm, + "Deterministic GPT-NeoX RoPE apply for ROCm"); +#endif #endif } diff --git a/csrc/rocm/distributed/deterministic_collective.hip b/csrc/rocm/distributed/deterministic_collective.hip new file mode 100644 index 00000000..c8448300 --- /dev/null +++ b/csrc/rocm/distributed/deterministic_collective.hip @@ -0,0 +1,1099 @@ +// ROCm fixed-tree reduction kernels for the RCCL transport collective. +// +// RCCL is intentionally used only to transport rank-ordered tensors. The +// kernels below perform the arithmetic locally in the exact balanced tree used +// by the Python reference implementation. ReduceScatter receives a view that +// contains only the destination rank's shard, so it does not reduce unrelated +// rows. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kThreads = 256; +constexpr int kMaxBlocks = 4096; +constexpr int kMaxWorldSize = 8; +constexpr int64_t kIPCControlBytes = 256; +constexpr int64_t kIPCReadyOffset = 0; +constexpr int64_t kIPCDoneOffset = 64; +constexpr int64_t kIPCCloseOffset = 128; + +struct PeerPointers { + const void* values[kMaxWorldSize]; +}; + +struct PeerSignals { + uint64_t* ready[kMaxWorldSize]; + uint64_t* done[kMaxWorldSize]; + uint64_t* closed[kMaxWorldSize]; +}; + +template +__device__ __forceinline__ scalar_t ordered_add(scalar_t lower, scalar_t upper) { + // Keep every parent as a separate expression. ROCm builds do not enable + // fast-math, so this is the same dtype operation as torch.add_ for the + // supported floating-point dtypes. + return lower + upper; +} + +template +__device__ __forceinline__ scalar_t fixed_tree_reduce( + const scalar_t* values, + int64_t rank_stride, + int64_t index) { + static_assert( + WorldSize == 1 || WorldSize == 2 || WorldSize == 4 || WorldSize == 8, + "unsupported deterministic collective world size"); + if constexpr (WorldSize == 1) { + return values[index]; + } else { + const scalar_t sum01 = ordered_add( + values[index], + values[rank_stride + index]); + if constexpr (WorldSize == 2) { + return sum01; + } else { + const scalar_t sum23 = ordered_add( + values[2 * rank_stride + index], + values[3 * rank_stride + index]); + const scalar_t sum03 = ordered_add(sum01, sum23); + if constexpr (WorldSize == 4) { + return sum03; + } else { + const scalar_t sum45 = ordered_add( + values[4 * rank_stride + index], + values[5 * rank_stride + index]); + const scalar_t sum67 = ordered_add( + values[6 * rank_stride + index], + values[7 * rank_stride + index]); + const scalar_t sum47 = ordered_add(sum45, sum67); + return ordered_add(sum03, sum47); + } + } + } +} + +template +__global__ void fixed_tree_reduce_kernel( + const scalar_t* __restrict__ values, + scalar_t* __restrict__ output, + int64_t rank_stride, + int64_t element_count) { + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < element_count; index += stride) { + output[index] = fixed_tree_reduce(values, rank_stride, index); + } +} + +template +void launch_fixed_tree_reduce( + const scalar_t* values, + scalar_t* output, + int64_t rank_stride, + int64_t element_count, + int64_t world_size, + hipStream_t stream) { + const int blocks = static_cast(std::min( + kMaxBlocks, + (element_count + kThreads - 1) / kThreads)); + switch (world_size) { + case 1: + hipLaunchKernelGGL( + (fixed_tree_reduce_kernel), + dim3(blocks), + dim3(kThreads), + 0, + stream, + values, + output, + rank_stride, + element_count); + break; + case 2: + hipLaunchKernelGGL( + (fixed_tree_reduce_kernel), + dim3(blocks), + dim3(kThreads), + 0, + stream, + values, + output, + rank_stride, + element_count); + break; + case 4: + hipLaunchKernelGGL( + (fixed_tree_reduce_kernel), + dim3(blocks), + dim3(kThreads), + 0, + stream, + values, + output, + rank_stride, + element_count); + break; + case 8: + hipLaunchKernelGGL( + (fixed_tree_reduce_kernel), + dim3(blocks), + dim3(kThreads), + 0, + stream, + values, + output, + rank_stride, + element_count); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void validate_inputs( + const torch::Tensor& rank_inputs, + const torch::Tensor& output, + const char* name) { + TORCH_CHECK(rank_inputs.is_cuda(), name, ": rank_inputs must be a CUDA/ROCm tensor"); + TORCH_CHECK(output.is_cuda(), name, ": output must be a CUDA/ROCm tensor"); + TORCH_CHECK(rank_inputs.scalar_type() == output.scalar_type(), name, ": dtype mismatch"); + TORCH_CHECK(rank_inputs.dim() >= 1, name, ": rank_inputs must have a rank dimension"); + TORCH_CHECK(rank_inputs.size(0) == 1 || rank_inputs.size(0) == 2 || + rank_inputs.size(0) == 4 || rank_inputs.size(0) == 8, + name, ": unsupported rank dimension ", rank_inputs.size(0)); + TORCH_CHECK(rank_inputs.select(0, 0).is_contiguous(), + name, ": each rank slice must be contiguous"); + TORCH_CHECK(output.is_contiguous(), name, ": output must be contiguous"); + TORCH_CHECK(rank_inputs.device() == output.device(), name, ": device mismatch"); + TORCH_CHECK(rank_inputs.numel() == output.numel() * rank_inputs.size(0), + name, ": rank_inputs/output element count mismatch"); +} + +void launch_dispatch( + const torch::Tensor& rank_inputs, + const torch::Tensor& output, + const char* name) { + validate_inputs(rank_inputs, output, name); + const int64_t world_size = rank_inputs.size(0); + const int64_t element_count = output.numel(); + if (element_count == 0) { + return; + } + const int64_t rank_stride = rank_inputs.stride(0); + const auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + rank_inputs.scalar_type(), + "deterministic_collective_rocm_fixed_tree", + [&] { + launch_fixed_tree_reduce( + rank_inputs.data_ptr(), + output.data_ptr(), + rank_stride, + element_count, + world_size, + stream); + }); +} + +template +__device__ __forceinline__ scalar_t ipc_fixed_tree_reduce( + const PeerPointers& peers, + int64_t index) { + const auto* rank0 = static_cast(peers.values[0]); + if constexpr (WorldSize == 1) { + return rank0[index]; + } else { + const auto* rank1 = static_cast(peers.values[1]); + const scalar_t sum01 = ordered_add(rank0[index], rank1[index]); + if constexpr (WorldSize == 2) { + return sum01; + } else { + const auto* rank2 = static_cast(peers.values[2]); + const auto* rank3 = static_cast(peers.values[3]); + const scalar_t sum23 = ordered_add(rank2[index], rank3[index]); + const scalar_t sum03 = ordered_add(sum01, sum23); + if constexpr (WorldSize == 4) { + return sum03; + } else { + const auto* rank4 = static_cast(peers.values[4]); + const auto* rank5 = static_cast(peers.values[5]); + const auto* rank6 = static_cast(peers.values[6]); + const auto* rank7 = static_cast(peers.values[7]); + const scalar_t sum45 = ordered_add(rank4[index], rank5[index]); + const scalar_t sum67 = ordered_add(rank6[index], rank7[index]); + return ordered_add(sum03, ordered_add(sum45, sum67)); + } + } + } +} + +template +__device__ __forceinline__ packed_t ordered_add_packed( + packed_t lower, + packed_t upper); + +template <> +__device__ __forceinline__ __half2 ordered_add_packed( + __half2 lower, + __half2 upper) { + return __hadd2(lower, upper); +} + +template <> +__device__ __forceinline__ __hip_bfloat162 ordered_add_packed( + __hip_bfloat162 lower, + __hip_bfloat162 upper) { + return __hadd2(lower, upper); +} + +template +__device__ __forceinline__ packed_t ipc_fixed_tree_reduce_packed( + const PeerPointers& peers, + int64_t index) { + const auto* rank0 = static_cast(peers.values[0]); + if constexpr (WorldSize == 1) { + return rank0[index]; + } else { + const auto* rank1 = static_cast(peers.values[1]); + const packed_t sum01 = ordered_add_packed(rank0[index], rank1[index]); + if constexpr (WorldSize == 2) { + return sum01; + } else { + const auto* rank2 = static_cast(peers.values[2]); + const auto* rank3 = static_cast(peers.values[3]); + const packed_t sum23 = ordered_add_packed(rank2[index], rank3[index]); + const packed_t sum03 = ordered_add_packed(sum01, sum23); + if constexpr (WorldSize == 4) { + return sum03; + } else { + const auto* rank4 = static_cast(peers.values[4]); + const auto* rank5 = static_cast(peers.values[5]); + const auto* rank6 = static_cast(peers.values[6]); + const auto* rank7 = static_cast(peers.values[7]); + const packed_t sum45 = ordered_add_packed(rank4[index], rank5[index]); + const packed_t sum67 = ordered_add_packed(rank6[index], rank7[index]); + return ordered_add_packed(sum03, ordered_add_packed(sum45, sum67)); + } + } + } +} + +template +__global__ void ipc_fixed_tree_reduce_kernel( + PeerPointers peers, + scalar_t* __restrict__ output, + int64_t input_offset, + int64_t element_count) { + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < element_count; index += stride) { + output[index] = ipc_fixed_tree_reduce( + peers, + input_offset + index); + } +} + +template +__global__ void ipc_fixed_tree_reduce_packed_kernel( + PeerPointers peers, + packed_t* __restrict__ output, + int64_t input_offset, + int64_t element_count) { + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < element_count; index += stride) { + output[index] = ipc_fixed_tree_reduce_packed( + peers, + input_offset + index); + } +} + +template +void launch_ipc_fixed_tree_reduce( + const PeerPointers& peers, + scalar_t* output, + int64_t input_offset, + int64_t element_count, + int64_t world_size, + hipStream_t stream) { + const int blocks = static_cast(std::min( + kMaxBlocks, + (element_count + kThreads - 1) / kThreads)); + switch (world_size) { + case 1: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 2: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 4: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 8: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void launch_ipc_fixed_tree_reduce_packed( + const PeerPointers& peers, + packed_t* output, + int64_t input_offset, + int64_t element_count, + int64_t world_size, + hipStream_t stream) { + const int blocks = static_cast(std::min( + kMaxBlocks, + (element_count + kThreads - 1) / kThreads)); + switch (world_size) { + case 1: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_packed_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 2: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_packed_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 4: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_packed_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + case 8: + hipLaunchKernelGGL( + (ipc_fixed_tree_reduce_packed_kernel), + dim3(blocks), dim3(kThreads), 0, stream, + peers, output, input_offset, element_count); + break; + default: + TORCH_CHECK(false, "unsupported deterministic collective world size ", world_size); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +__global__ void ipc_wait_signal_kernel( + PeerSignals signals, + uint64_t sequence, + int64_t world_size, + bool wait_for_done) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + for (int peer = 0; peer < world_size; ++peer) { + uint64_t* signal = wait_for_done ? signals.done[peer] : signals.ready[peer]; + while (__hip_atomic_load( + signal, + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM) < sequence) { + __builtin_amdgcn_s_sleep(1); + } + } +} + +__global__ void ipc_mark_signal_kernel(uint64_t* signal, uint64_t sequence) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + __hip_atomic_store( + signal, + sequence, + __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_SYSTEM); + } +} + +__global__ void ipc_mark_ready_and_wait_kernel( + PeerSignals signals, + int64_t rank, + uint64_t sequence, + int64_t world_size) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + __hip_atomic_store( + signals.ready[rank], + sequence, + __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_SYSTEM); + for (int peer = 0; peer < world_size; ++peer) { + while (__hip_atomic_load( + signals.ready[peer], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM) < sequence) { + __builtin_amdgcn_s_sleep(1); + } + } +} + +__global__ void ipc_close_and_wait_kernel( + PeerSignals signals, + int64_t rank, + int64_t world_size) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + for (int peer = 0; peer < world_size; ++peer) { + __hip_atomic_fetch_add( + signals.closed[peer], + static_cast(1), + __ATOMIC_ACQ_REL, + __HIP_MEMORY_SCOPE_SYSTEM); + } + while (__hip_atomic_load( + signals.closed[rank], + __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM) < static_cast(world_size)) { + __builtin_amdgcn_s_sleep(1); + } +} + +__global__ void ipc_all_gather_uint4_kernel( + PeerPointers peers, + uint4* __restrict__ output, + int64_t vectors_per_rank, + int64_t world_size) { + const int64_t total_vectors = vectors_per_rank * world_size; + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < total_vectors; index += stride) { + const int peer = static_cast(index / vectors_per_rank); + const int64_t peer_index = index - static_cast(peer) * vectors_per_rank; + output[index] = static_cast(peers.values[peer])[peer_index]; + } +} + +__global__ void ipc_all_gather_bytes_kernel( + PeerPointers peers, + uint8_t* __restrict__ output, + int64_t bytes_per_rank, + int64_t world_size) { + const int64_t total_bytes = bytes_per_rank * world_size; + const int64_t thread_index = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t index = thread_index; index < total_bytes; index += stride) { + const int peer = static_cast(index / bytes_per_rank); + const int64_t peer_index = index - static_cast(peer) * bytes_per_rank; + output[index] = static_cast(peers.values[peer])[peer_index]; + } +} + +void launch_ipc_all_gather( + const PeerPointers& peers, + void* output, + int64_t bytes_per_rank, + int64_t world_size, + hipStream_t stream) { + if (bytes_per_rank == 0) { + return; + } + if (bytes_per_rank % static_cast(sizeof(uint4)) == 0 && + reinterpret_cast(output) % alignof(uint4) == 0) { + const int64_t vectors_per_rank = bytes_per_rank / sizeof(uint4); + const int64_t total_vectors = vectors_per_rank * world_size; + const int blocks = static_cast(std::min( + kMaxBlocks, + (total_vectors + kThreads - 1) / kThreads)); + hipLaunchKernelGGL( + ipc_all_gather_uint4_kernel, + dim3(blocks), + dim3(kThreads), + 0, + stream, + peers, + static_cast(output), + vectors_per_rank, + world_size); + } else { + const int64_t total_bytes = bytes_per_rank * world_size; + const int blocks = static_cast(std::min( + kMaxBlocks, + (total_bytes + kThreads - 1) / kThreads)); + hipLaunchKernelGGL( + ipc_all_gather_bytes_kernel, + dim3(blocks), + dim3(kThreads), + 0, + stream, + peers, + static_cast(output), + bytes_per_rank, + world_size); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_wait_signal( + const PeerSignals& signals, + uint64_t sequence, + int64_t world_size, + bool wait_for_done, + hipStream_t stream) { + hipLaunchKernelGGL( + ipc_wait_signal_kernel, + dim3(1), + dim3(1), + 0, + stream, + signals, + sequence, + world_size, + wait_for_done); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_mark_signal(uint64_t* signal, uint64_t sequence, hipStream_t stream) { + hipLaunchKernelGGL( + ipc_mark_signal_kernel, + dim3(1), + dim3(1), + 0, + stream, + signal, + sequence); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_mark_ready_and_wait( + const PeerSignals& signals, + int64_t rank, + uint64_t sequence, + int64_t world_size, + hipStream_t stream) { + hipLaunchKernelGGL( + ipc_mark_ready_and_wait_kernel, + dim3(1), + dim3(1), + 0, + stream, + signals, + rank, + sequence, + world_size); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void launch_close_and_wait( + const PeerSignals& signals, + int64_t rank, + int64_t world_size, + hipStream_t stream) { + hipLaunchKernelGGL( + ipc_close_and_wait_kernel, + dim3(1), + dim3(1), + 0, + stream, + signals, + rank, + world_size); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +class ROCmIPCCollectiveState { + public: + ROCmIPCCollectiveState( + torch::Tensor staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank) + : rank_(rank), + world_size_(static_cast(handles.size())), + device_index_(staging.get_device()), + capacity_bytes_(staging.numel() * staging.element_size() - kIPCControlBytes) { + TORCH_CHECK(staging.is_cuda(), "ROCm IPC staging buffer must be on device"); + TORCH_CHECK(staging.is_contiguous(), "ROCm IPC staging buffer must be contiguous"); + TORCH_CHECK(staging.scalar_type() == torch::kUInt8, + "ROCm IPC staging buffer must have dtype uint8"); + TORCH_CHECK(capacity_bytes_ > 0, "ROCm IPC staging capacity must be positive"); + TORCH_CHECK( + world_size_ == 1 || world_size_ == 2 || world_size_ == 4 || world_size_ == 8, + "ROCm IPC deterministic collectives require world size 1, 2, 4, or 8"); + TORCH_CHECK(offsets.size() == handles.size(), "one IPC offset is required per rank"); + TORCH_CHECK(rank_ >= 0 && rank_ < world_size_, "invalid ROCm IPC rank"); + + set_peer_pointers(rank_, staging.data_ptr()); + try { + for (int peer = 0; peer < world_size_; ++peer) { + if (peer == rank_) { + continue; + } + TORCH_CHECK(handles[peer].size() == sizeof(hipIpcMemHandle_t), + "invalid ROCm IPC handle size for rank ", peer); + TORCH_CHECK(offsets[peer] >= 0, "invalid negative ROCm IPC offset"); + hipIpcMemHandle_t handle{}; + auto* raw_handle = reinterpret_cast(&handle); + for (size_t byte = 0; byte < sizeof(handle); ++byte) { + TORCH_CHECK(handles[peer][byte] >= 0 && handles[peer][byte] <= 255, + "invalid ROCm IPC handle byte for rank ", peer); + raw_handle[byte] = static_cast(handles[peer][byte]); + } + void* base = nullptr; + C10_HIP_CHECK(hipIpcOpenMemHandle( + &base, + handle, + hipIpcMemLazyEnablePeerAccess)); + imported_bases_[peer] = base; + set_peer_pointers( + peer, + static_cast(base) + offsets[peer]); + } + } catch (...) { + close_imports(); + throw; + } + } + + ~ROCmIPCCollectiveState() { + int previous_device = -1; + if (hipGetDevice(&previous_device) == hipSuccess && previous_device != device_index_) { + if (hipSetDevice(device_index_) != hipSuccess) { + return; + } + } + close_imports(); + if (previous_device >= 0 && previous_device != device_index_) { + C10_CUDA_IGNORE_ERROR(hipSetDevice(previous_device)); + } + } + + int device_index() const { + return device_index_; + } + + void stage(torch::Tensor input, hipStream_t stream) { + check_tensor(input, "input"); + const int64_t input_bytes = input.numel() * input.element_size(); + TORCH_CHECK(input_bytes <= capacity_bytes_, + "input exceeds ROCm IPC staging capacity"); + ++sequence_; + launch_wait_signal( + signals_, + sequence_ - 1, + world_size_, + true, + stream); + if (input_bytes > 0) { + C10_HIP_CHECK(hipMemcpyAsync( + const_cast(peers_.values[rank_]), + input.data_ptr(), + input_bytes, + hipMemcpyDeviceToDevice, + stream)); + } + launch_mark_ready_and_wait( + signals_, + rank_, + sequence_, + world_size_, + stream); + staged_bytes_ = input_bytes; + staged_type_ = input.scalar_type(); + } + + void all_reduce(torch::Tensor output, hipStream_t stream) const { + check_reduction_output(output, staged_bytes_, "all_reduce"); + launch(output, 0, output.numel(), stream); + launch_mark_signal(signals_.done[rank_], sequence_, stream); + } + + void reduce_scatter(torch::Tensor output, hipStream_t stream) const { + check_reduction_output( + output, + staged_bytes_ / world_size_, + "reduce_scatter"); + launch( + output, + rank_ * output.numel(), + output.numel(), + stream); + launch_mark_signal(signals_.done[rank_], sequence_, stream); + } + + void reduce_scatter_many( + const std::vector& inputs, + const std::vector& outputs, + hipStream_t stream) { + TORCH_CHECK(!inputs.empty(), "reduce_scatter_many requires at least one input"); + TORCH_CHECK(inputs.size() == outputs.size(), + "reduce_scatter_many input/output count mismatch"); + + int64_t total_bytes = 0; + const auto scalar_type = inputs.front().scalar_type(); + for (size_t index = 0; index < inputs.size(); ++index) { + const auto& input = inputs[index]; + const auto& output = outputs[index]; + check_tensor(input, "reduce_scatter_many input"); + check_tensor(output, "reduce_scatter_many output"); + TORCH_CHECK(input.scalar_type() == scalar_type && output.scalar_type() == scalar_type, + "reduce_scatter_many dtype mismatch"); + const int64_t input_bytes = input.numel() * input.element_size(); + TORCH_CHECK(input_bytes == output.numel() * output.element_size() * world_size_, + "reduce_scatter_many output size mismatch"); + TORCH_CHECK(input_bytes <= capacity_bytes_ - total_bytes, + "reduce_scatter_many inputs exceed ROCm IPC staging capacity"); + total_bytes += input_bytes; + } + + ++sequence_; + launch_wait_signal(signals_, sequence_ - 1, world_size_, true, stream); + int64_t byte_offset = 0; + for (const auto& input : inputs) { + const int64_t input_bytes = input.numel() * input.element_size(); + if (input_bytes > 0) { + C10_HIP_CHECK(hipMemcpyAsync( + static_cast(const_cast(peers_.values[rank_])) + byte_offset, + input.data_ptr(), + input_bytes, + hipMemcpyDeviceToDevice, + stream)); + } + byte_offset += input_bytes; + } + launch_mark_ready_and_wait( + signals_, + rank_, + sequence_, + world_size_, + stream); + + int64_t element_offset = 0; + for (size_t index = 0; index < outputs.size(); ++index) { + const auto& input = inputs[index]; + const auto& output = outputs[index]; + launch( + output, + element_offset + rank_ * output.numel(), + output.numel(), + stream); + element_offset += input.numel(); + } + launch_mark_signal(signals_.done[rank_], sequence_, stream); + } + + void all_gather(torch::Tensor output, hipStream_t stream) const { + check_tensor(output, "all_gather"); + TORCH_CHECK(staged_type_ != at::ScalarType::Undefined, "stage must be called first"); + TORCH_CHECK(output.scalar_type() == staged_type_, "all_gather dtype mismatch"); + TORCH_CHECK( + output.numel() * output.element_size() == staged_bytes_ * world_size_, + "all_gather output size mismatch"); + launch_ipc_all_gather( + peers_, + output.data_ptr(), + staged_bytes_, + world_size_, + stream); + launch_mark_signal(signals_.done[rank_], sequence_, stream); + } + + void synchronize(hipStream_t stream) const { + launch_wait_signal(signals_, sequence_, world_size_, true, stream); + launch_close_and_wait(signals_, rank_, world_size_, stream); + } + + private: + void set_peer_pointers(int peer, void* allocation_base) { + auto* bytes = static_cast(allocation_base); + signals_.ready[peer] = reinterpret_cast(bytes + kIPCReadyOffset); + signals_.done[peer] = reinterpret_cast(bytes + kIPCDoneOffset); + signals_.closed[peer] = reinterpret_cast(bytes + kIPCCloseOffset); + peers_.values[peer] = bytes + kIPCControlBytes; + } + + void check_tensor(const torch::Tensor& tensor, const char* name) const { + TORCH_CHECK(tensor.is_cuda(), name, " must be a ROCm tensor"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(tensor.get_device() == device_index_, name, " device mismatch"); + } + + void check_reduction_output( + const torch::Tensor& output, + int64_t expected_bytes, + const char* name) const { + check_tensor(output, name); + TORCH_CHECK(staged_type_ != at::ScalarType::Undefined, "stage must be called first"); + TORCH_CHECK(output.scalar_type() == staged_type_, name, " dtype mismatch"); + TORCH_CHECK(output.numel() * output.element_size() == expected_bytes, + name, " output size mismatch"); + } + + void launch( + torch::Tensor output, + int64_t input_offset, + int64_t element_count, + hipStream_t stream) const { + if (element_count == 0) { + return; + } + if (element_count % 2 == 0 && input_offset % 2 == 0) { + if (output.scalar_type() == at::ScalarType::Half && + reinterpret_cast(output.data_ptr()) % alignof(__half2) == 0) { + launch_ipc_fixed_tree_reduce_packed<__half2>( + peers_, + reinterpret_cast<__half2*>(output.data_ptr()), + input_offset / 2, + element_count / 2, + world_size_, + stream); + return; + } + if (output.scalar_type() == at::ScalarType::BFloat16 && + reinterpret_cast(output.data_ptr()) % + alignof(__hip_bfloat162) == + 0) { + launch_ipc_fixed_tree_reduce_packed<__hip_bfloat162>( + peers_, + reinterpret_cast<__hip_bfloat162*>(output.data_ptr()), + input_offset / 2, + element_count / 2, + world_size_, + stream); + return; + } + } + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, + at::ScalarType::BFloat16, + output.scalar_type(), + "deterministic_collective_rocm_ipc_fixed_tree", + [&] { + launch_ipc_fixed_tree_reduce( + peers_, + output.data_ptr(), + input_offset, + element_count, + world_size_, + stream); + }); + } + + void close_imports() noexcept { + for (int peer = 0; peer < world_size_; ++peer) { + if (imported_bases_[peer] != nullptr) { + C10_CUDA_IGNORE_ERROR(hipIpcCloseMemHandle(imported_bases_[peer])); + imported_bases_[peer] = nullptr; + } + } + } + + int64_t rank_; + int64_t world_size_; + int device_index_; + int64_t capacity_bytes_; + int64_t staged_bytes_{0}; + at::ScalarType staged_type_{at::ScalarType::Undefined}; + uint64_t sequence_{0}; + PeerPointers peers_{}; + PeerSignals signals_{}; + std::array imported_bases_{}; +}; + +ROCmIPCCollectiveState* ipc_state(int64_t handle) { + TORCH_CHECK(handle != 0, "ROCm IPC collective handle is closed"); + return reinterpret_cast(handle); +} + +} // namespace + +void deterministic_collective_rocm_all_reduce( + torch::Tensor rank_inputs, + torch::Tensor output) { + launch_dispatch(rank_inputs, output, "deterministic_collective_rocm_all_reduce"); +} + +void deterministic_collective_rocm_reduce_scatter( + torch::Tensor rank_inputs, + torch::Tensor output) { + launch_dispatch(rank_inputs, output, "deterministic_collective_rocm_reduce_scatter"); +} + +torch::Tensor deterministic_collective_rocm_ipc_allocate(int64_t size_bytes) { + TORCH_CHECK(size_bytes > 0, "ROCm IPC allocation size must be positive"); + int device_index = -1; + C10_HIP_CHECK(hipGetDevice(&device_index)); + const int64_t allocation_bytes = size_bytes + kIPCControlBytes; + void* pointer = nullptr; + C10_HIP_CHECK(hipMalloc(&pointer, static_cast(allocation_bytes))); + C10_HIP_CHECK(hipMemset(pointer, 0, static_cast(kIPCControlBytes))); + const auto options = torch::TensorOptions() + .dtype(torch::kUInt8) + .device(torch::Device(torch::kCUDA, device_index)); + return torch::from_blob( + pointer, + {allocation_bytes}, + [device_index](void* allocation) { + int previous_device = -1; + if (hipGetDevice(&previous_device) != hipSuccess) { + return; + } + if (previous_device != device_index && hipSetDevice(device_index) != hipSuccess) { + return; + } + C10_CUDA_IGNORE_ERROR(hipFree(allocation)); + if (previous_device != device_index) { + C10_CUDA_IGNORE_ERROR(hipSetDevice(previous_device)); + } + }, + options); +} + +std::tuple, int64_t> +deterministic_collective_rocm_ipc_meta(torch::Tensor tensor) { + const c10::cuda::CUDAGuard device_guard(tensor.device()); + TORCH_CHECK(tensor.is_cuda(), "ROCm IPC tensor must be on device"); + TORCH_CHECK(tensor.is_contiguous(), "ROCm IPC tensor must be contiguous"); + TORCH_CHECK(tensor.numel() > 0, "ROCm IPC tensor must be non-empty"); + + hipIpcMemHandle_t handle{}; + const hipError_t export_error = hipIpcGetMemHandle(&handle, tensor.data_ptr()); + TORCH_CHECK( + export_error == hipSuccess, + "hipIpcGetMemHandle failed: ", + hipGetErrorString(export_error)); + const auto* raw_handle = reinterpret_cast(&handle); + std::vector bytes(sizeof(handle)); + for (size_t byte = 0; byte < sizeof(handle); ++byte) { + bytes[byte] = raw_handle[byte]; + } + return std::make_tuple(bytes, 0); +} + +int64_t deterministic_collective_rocm_ipc_create( + torch::Tensor staging, + const std::vector>& handles, + const std::vector& offsets, + int64_t rank) { + const c10::cuda::CUDAGuard device_guard(staging.device()); + auto state = std::make_unique( + staging, + handles, + offsets, + rank); + return reinterpret_cast(state.release()); +} + +void deterministic_collective_rocm_ipc_destroy(int64_t handle) { + delete ipc_state(handle); +} + +void deterministic_collective_rocm_ipc_synchronize(int64_t handle) { + auto* state = ipc_state(handle); + const c10::cuda::CUDAGuard device_guard( + torch::Device(torch::kCUDA, state->device_index())); + const auto stream = at::cuda::getCurrentCUDAStream(); + state->synchronize(stream); +} + +void deterministic_collective_rocm_ipc_stage(int64_t handle, torch::Tensor input) { + const c10::cuda::CUDAGuard device_guard(input.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->stage(input, stream); +} + +void deterministic_collective_rocm_ipc_all_reduce( + int64_t handle, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(output.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->all_reduce(output, stream); +} + +void deterministic_collective_rocm_ipc_all_reduce_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(input.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + auto* state = ipc_state(handle); + state->stage(input, stream); + state->all_reduce(output, stream); +} + +void deterministic_collective_rocm_ipc_reduce_scatter( + int64_t handle, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(output.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->reduce_scatter(output, stream); +} + +void deterministic_collective_rocm_ipc_reduce_scatter_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(input.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + auto* state = ipc_state(handle); + state->stage(input, stream); + state->reduce_scatter(output, stream); +} + +void deterministic_collective_rocm_ipc_reduce_scatter_many( + int64_t handle, + const std::vector& inputs, + const std::vector& outputs) { + TORCH_CHECK(!inputs.empty(), "reduce_scatter_many requires at least one input"); + const c10::cuda::CUDAGuard device_guard(inputs.front().device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->reduce_scatter_many(inputs, outputs, stream); +} + +void deterministic_collective_rocm_ipc_all_gather( + int64_t handle, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(output.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + ipc_state(handle)->all_gather(output, stream); +} + +void deterministic_collective_rocm_ipc_all_gather_input( + int64_t handle, + torch::Tensor input, + torch::Tensor output) { + const c10::cuda::CUDAGuard device_guard(input.device()); + const auto stream = at::cuda::getCurrentCUDAStream(); + auto* state = ipc_state(handle); + state->stage(input, stream); + state->all_gather(output, stream); +} diff --git a/docs/design/rocm-deterministic-collectives.md b/docs/design/rocm-deterministic-collectives.md new file mode 100644 index 00000000..c105a537 --- /dev/null +++ b/docs/design/rocm-deterministic-collectives.md @@ -0,0 +1,141 @@ +# ROCm deterministic collectives + +This document defines the ROCm communication boundary used by deterministic +TP/CP kernels. The implementation is intentionally separate from the CUDA IPC +collective in `csrc/cuda/distributed/deterministic_collective.cu`. + +## Contract + +`RCCLDeterministicCollective` implements `all_reduce`, `all_gather`, and +`reduce_scatter` with the same Python call shape as the CUDA collective. It +supports process-group sizes 1, 2, 4, and 8 and FP32/FP16/BF16 reductions. +Inputs and outputs must be contiguous and all ranks must call operations in the +same order with matching shapes, dtypes, and capacity. + +Single-node ROCm uses a dedicated `hipMalloc` staging allocation on every rank. +The handles are exchanged once during construction and imported with HIP IPC. +Each call copies its local input into staging, publishes a system-scope GPU +sequence flag, and waits for every peer's matching sequence. The HIP kernel +then reads rank-ordered peer memory and evaluates exactly +`((rank0 + rank1) + (rank2 + rank3)) + ...`. A second sequence flag prevents a +rank from overwriting staging until every peer has finished reading it. + +FP16 and BF16 use two-element vector loads and `hadd2`. This changes only the +number of elements carried by an instruction: every scalar element retains the +same dtype, rank order, and expression grouping. FP32 and unaligned tails use +the scalar kernel. The executable Python tree remains the fallback for +CPU/reference backends and extensions built without the optional HIP source. + +The measured MI300X routing policy is: + +- AllReduce up to 768 KiB uses the direct IPC fixed-tree kernel. +- AllReduce from 768 KiB to 2.125 MiB uses the rank-major RCCL transport fallback. +- AllReduce at 2.125 MiB and above performs IPC ReduceScatter followed by RCCL + AllGather of the already-reduced shards. +- AllGather up to 256 KiB uses IPC peer copies; larger messages use RCCL. +- ReduceScatter uses IPC for all supported sizes and reduces only the local + destination shard. + +The ready publication and peer wait share one GPU kernel. The store is a +system-scope release and every peer load is a system-scope acquire. Done flags +remain a separate generation barrier because they protect staging reuse. At +close, every rank atomically acknowledges every peer allocation and waits for +all acknowledgements before releasing its local staging memory. + +Sequence-parallel FFN backward has two independent ReduceScatter lanes (gate +and up input gradients). On IPC, `reduce_scatter_many` copies the lanes into +disjoint staging ranges under one ready/done generation and launches one fixed +tree per output lane. It does not concatenate the inputs or mix their trees. +The RCCL fallback retains a measured packed-payload crossover because a larger +rank-major AllGather can be slower than two smaller calls. + +RCCL's `all_reduce` and `reduce_scatter` are not used for strict reductions. +They guarantee a mathematical reduction but do not expose a stable +floating-point operand order. Delegating the arithmetic to them would weaken +the cross-TP bitwise contract. + +`create_deterministic_collective` is the platform boundary. It selects the +existing CUDA IPC implementation on NVIDIA and the RCCL transport +implementation when `torch.version.hip` is set. The FFN path calls this factory +instead of importing the CUDA class directly. + +The ROCm extension build also excludes the CUDA IPC source and does not link +`libcuda`; the Python transport has no CUDA-driver dependency. + +## Relationship to vLLM + +The backend split follows the useful parts of vLLM's device communicator +design: + +- PyTorch exposes RCCL through the `nccl` process-group API. +- ROCm AllGather uses `torch.distributed.all_gather_into_tensor` rather than a + manually allocated PyNccl path. +- Backend, topology, world-size, dtype, layout, and capacity checks fail closed + before entering an optimized path. + +vLLM QuickReduce, custom all-reduce, and AITER all-reduce are not used in the +strict path. Those are valuable performance implementations, but their +reduction order is not the fixed balanced tree required here. They can be +added later as an explicitly non-strict performance mode with separate +provenance and toleranced correctness tests. + +## Compute/communication fusion + +The current ROCm collective is stream ordered and reports +`supports_async_overlap = False` and +`supports_compute_communication_fusion = False`. FFN and Attention keep the +dependency boundaries explicit: + +```text +sequence-parallel FFN: AllGather(input) -> GEMM -> ReduceScatter(output) +strict CP Attention: AllGather(Q/K/V/positions) -> Attention -> Scatter(output/LSE) +``` + +This is neither a fused GEMM+collective kernel nor two-stream overlap. vLLM's +generic AsyncTP GEMM/communication fusion is currently a CUDA path; ROCm AITER +has narrower fusions such as all-reduce plus RMSNorm, which do not replace this +contract. + +Future overlap must preserve collective issue order, the local reduction tree, +and stage boundaries. It also needs repeat-bitwise tests before being marked +strict. A useful first candidate is backward work that is independent of a +pending reduction; the forward SP AllGather and final row-parallel reduction +are data dependencies and cannot simply be overlapped with their adjacent +GEMMs. + +The “fill rank slots as messages arrive and merge ready contiguous blocks” +scheme needs a lower-level P2P or HIP/XGMI transport. PyTorch/RCCL +`all_gather_into_tensor` exposes completion of the whole collective, not +per-rank arrival events. Such a pipeline may merge only canonical sibling +subtrees when both are ready; merging arbitrary contiguous arrivals would +change floating-point parenthesization. It is therefore a follow-up transport, +not an optimization silently hidden inside this baseline. + +## Performance acceptance + +The IPC path favors a fixed arithmetic tree over native-RCCL reduction speed. +Large AllReduce avoids reducing the full tensor on every rank, but still moves +more data than a native RCCL AllReduce. A ROCm GPU PR should therefore report, for +world sizes 2/4/8 and representative FFN tensors: + +- latency and effective bandwidth for all three collectives; +- peak temporary memory; +- comparison with RCCL and vLLM's available ROCm communicator; +- repeat-bitwise and cross-TP results; +- end-to-end TP/CP/SP FFN timing, not only isolated transport timing. + +Multi-node or unsupported IPC configurations fall back behind the same factory +after topology and symbol checks fail closed. + +Run the included native-RCCL comparison on a single node, for example: + +```bash +torchrun --standalone --nproc-per-node=8 \ + benchmarks/benchmark_rocm_collectives.py \ + --size-bytes 4096 65536 1048576 16777216 \ + --output benchmarks/results/rocm_collectives_mi300x.json +``` + +The benchmark records slowest-rank latency, temporary allocation, repeat +bitwise status, and the ratio to native RCCL. Native RCCL remains a performance +reference only, not the strict arithmetic reference. diff --git a/pyproject.toml b/pyproject.toml index ca3b0c5d..72b62153 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,55 +1,55 @@ -[build-system] -requires = ["setuptools>=64", "wheel"] -build-backend = "setuptools.build_meta" - -[project] -name = "RL-Kernel" -version = "0.1.0" -description = "High-performance RL training engine focused on kernel fusion and memory efficiency." -readme = "README.md" -requires-python = ">=3.10" -license = {text = "Apache-2.0"} -authors = [ - {name = "RL-Kernel Contributors"} -] -dependencies = [ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", -] - -[project.entry-points."vllm.general_plugins"] -rl_kernel = "rl_engine.integrations.vllm_runtime:register_vllm_plugin" - -[project.optional-dependencies] -cuda = ["flashinfer-python>=0.1.6", "nvidia-ml-py"] +[build-system] +requires = ["setuptools>=64", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "RL-Kernel" +version = "0.1.0" +description = "High-performance RL training engine focused on kernel fusion and memory efficiency." +readme = "README.md" +requires-python = ">=3.10" +license = {text = "Apache-2.0"} +authors = [ + {name = "RL-Kernel Contributors"} +] +dependencies = [ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", +] + +[project.entry-points."vllm.general_plugins"] +rl_kernel = "rl_engine.integrations.vllm_runtime:register_vllm_plugin" + +[project.optional-dependencies] +cuda = ["flashinfer-python>=0.6.0,<0.7", "nvidia-ml-py"] rocm = ["aiter"] vllm = ["vllm>=0.6.0"] drift-viewer = ["Pillow>=10", "PySide6>=6.6"] dev = ["pytest", "black", "isort", "ruff", "mypy", "pre-commit"] - -[tool.setuptools.packages.find] -where = ["."] -include = ["rl_engine*"] - -[tool.ruff] -line-length = 100 - -[tool.ruff.lint] -select = ["E", "F", "B"] -ignore = [] - -[tool.ruff.lint.per-file-ignores] -"__init__.py" = ["F401"] - -[tool.mypy] -ignore_missing_imports = true -follow_imports = "silent" - -[tool.pytest.ini_options] -markers = [ - "smoke_operator: temporary smoke-only operator plumbing tests", - "unit: CPU-safe unit tests", -] + +[tool.setuptools.packages.find] +where = ["."] +include = ["rl_engine*"] + +[tool.ruff] +line-length = 100 + +[tool.ruff.lint] +select = ["E", "F", "B"] +ignore = [] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["F401"] + +[tool.mypy] +ignore_missing_imports = true +follow_imports = "silent" + +[tool.pytest.ini_options] +markers = [ + "smoke_operator: temporary smoke-only operator plumbing tests", + "unit: CPU-safe unit tests", +] diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8e0e865a..b60169e1 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -22,6 +22,14 @@ def deterministic_collective_all_gather(handle: int, output: torch.Tensor) -> No def deterministic_collective_all_gather_fused( handle: int, input: torch.Tensor, output: torch.Tensor ) -> None: ... +def deterministic_collective_rocm_all_reduce( + rank_inputs: torch.Tensor, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_reduce_scatter( + rank_inputs: torch.Tensor, + output: torch.Tensor, +) -> None: ... def fused_logp(logits: torch.Tensor, token_ids: torch.Tensor) -> torch.Tensor: ... def fused_logp_sm90(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: ... def batch_invariant_logp_sm90( @@ -222,3 +230,46 @@ def rmsnorm_backward_dw( rstd: torch.Tensor, mask: torch.Tensor, ) -> torch.Tensor: ... +def deterministic_collective_rocm_ipc_allocate(size_bytes: int) -> torch.Tensor: ... +def deterministic_collective_rocm_ipc_meta(tensor: torch.Tensor) -> tuple[list[int], int]: ... +def deterministic_collective_rocm_ipc_create( + staging: torch.Tensor, + handles: list[list[int]], + offsets: list[int], + rank: int, +) -> int: ... +def deterministic_collective_rocm_ipc_synchronize(handle: int) -> None: ... +def deterministic_collective_rocm_ipc_destroy(handle: int) -> None: ... +def deterministic_collective_rocm_ipc_stage(handle: int, input: torch.Tensor) -> None: ... +def deterministic_collective_rocm_ipc_all_reduce( + handle: int, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_all_reduce_input( + handle: int, + input: torch.Tensor, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_reduce_scatter( + handle: int, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_reduce_scatter_input( + handle: int, + input: torch.Tensor, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_reduce_scatter_many( + handle: int, + inputs: tuple[torch.Tensor, ...], + outputs: tuple[torch.Tensor, ...], +) -> None: ... +def deterministic_collective_rocm_ipc_all_gather( + handle: int, + output: torch.Tensor, +) -> None: ... +def deterministic_collective_rocm_ipc_all_gather_input( + handle: int, + input: torch.Tensor, + output: torch.Tensor, +) -> None: ... diff --git a/rl_engine/distributed/__init__.py b/rl_engine/distributed/__init__.py index 37698f1a..9010a534 100644 --- a/rl_engine/distributed/__init__.py +++ b/rl_engine/distributed/__init__.py @@ -1,6 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from rl_engine.distributed.collectives import DeterministicCollective +from rl_engine.distributed.collectives import ( + DeterministicCollective, + RCCLDeterministicCollective, + TorchDistributedDeterministicCollective, + create_deterministic_collective, +) -__all__ = ["DeterministicCollective"] +__all__ = [ + "DeterministicCollective", + "RCCLDeterministicCollective", + "TorchDistributedDeterministicCollective", + "create_deterministic_collective", +] diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index 8df550af..b599ad76 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -1,5 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +"""Deterministic collectives for CUDA IPC and ROCm rank-ordered transport. + +ROCm uses HIP IPC where it wins and RCCL otherwise. Reduction arithmetic stays +outside RCCL and follows the same fixed balanced rank tree on every rank. +""" from __future__ import annotations @@ -13,10 +18,18 @@ _SUPPORTED_WORLD_SIZES = (1, 2, 4, 8) _DEFAULT_MAX_SIZE_BYTES = 64 * 1024 * 1024 +# Packing two independent lanes saves a collective launch for small tensors, +# but doubles the message size seen by RCCL. On MI300X, separate AllGather +# transports win once the packed payload reaches the multi-megabyte regime. +# Keep the crossover explicit and easy to retune with new RCCL releases. +_PACKED_REDUCE_SCATTER_MAX_BYTES = 8 * 1024 * 1024 +_ROCM_IPC_DIRECT_ALL_REDUCE_MAX_BYTES = 768 * 1024 +_ROCM_IPC_SHARDED_ALL_REDUCE_MIN_BYTES = 2176 * 1024 +_ROCM_IPC_ALL_GATHER_MAX_BYTES = 256 * 1024 _COLLECTIVE_STAGING_FRAMES = 3 _COLLECTIVE_FRAME_METADATA_BYTES = 3 * 8 _REDUCTION_DTYPES = (torch.float32, torch.float16, torch.bfloat16) -_COLLECTIVES: dict[tuple[int, int, int, int], DeterministicCollective] = {} +_COLLECTIVES: dict[tuple[int, int, int, int], Any] = {} DETERMINISTIC_ALL_REDUCE_OP = "rl_kernel::deterministic_all_reduce_" @@ -115,8 +128,10 @@ def __init__( "deterministic_collective_destroy", "deterministic_collective_stage", "deterministic_collective_all_reduce", + "deterministic_collective_all_reduce_fused", "deterministic_collective_reduce_scatter", "deterministic_collective_all_gather", + "deterministic_collective_all_gather_fused", ) missing = [name for name in required_symbols if not hasattr(_C, name)] if missing: @@ -268,17 +283,33 @@ def reduce_scatter( def reduce_scatter_many( self, - inputs: tuple[torch.Tensor, ...], + inputs: tuple[torch.Tensor, ...] | list[torch.Tensor], *, + outs: tuple[torch.Tensor, ...] | list[torch.Tensor] | None = None, validate_signature: bool = True, ) -> tuple[torch.Tensor, ...]: - """Reduce-scatter several tensors through the single-tensor ABI.""" + """Compatibility fallback for CUDA IPC collectives. - if not inputs: + The native CUDA IPC backend has no packed transport primitive yet, so + it preserves its established behavior by issuing the individual + fixed-tree calls. The ROCm transport subclass overrides this method + with a packed implementation. + """ + + values = tuple(inputs) + if not values: raise ValueError("reduce_scatter_many requires at least one input") - return tuple( - self.reduce_scatter(input, validate_signature=validate_signature) for input in inputs + if outs is not None and len(outs) != len(values): + raise ValueError("reduce_scatter_many outs must match the number of inputs") + results = tuple( + self.reduce_scatter( + value, + out=None if outs is None else outs[index], + validate_signature=validate_signature, + ) + for index, value in enumerate(values) ) + return results def close(self) -> None: """Release imported CUDA IPC mappings after the last collective call.""" @@ -416,13 +447,821 @@ def _synchronize_ranks(self) -> None: dist.barrier(group=self.group) +class TorchDistributedDeterministicCollective: + """Correctness-first collectives using AllGather as transport only. + + Rank inputs are gathered without arithmetic and reduced locally as the + balanced tree ``((rank0 + rank1) + (rank2 + rank3)) + ...``. Consequently, + all ranks execute the exact same floating-point expression. TP sizes 1, + 2, 4, and 8 are nested prefixes of that expression and match the existing + CUDA IPC collective's ordering. + + The generic class also supports a CPU/Gloo process group, which is useful + as an executable reference. Production ROCm callers should use + :class:`RCCLDeterministicCollective` or + :func:`create_deterministic_collective` so backend validation fails closed. + All ranks must call methods in the same order with matching input shapes + and dtypes, and construct the instance with the same ``max_size_bytes``. + """ + + backend_id = "torch_distributed_balanced_tree" + transport_only = True + reduction_order = "balanced_rank_tree" + supports_async_overlap = False + supports_compute_communication_fusion = False + + def __init__( + self, + group: dist.ProcessGroup | None = None, + device: torch.device | str | int | None = None, + *, + max_size_bytes: int = _DEFAULT_MAX_SIZE_BYTES, + ) -> None: + if not dist.is_available() or not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before collectives") + if max_size_bytes <= 0: + raise ValueError("max_size_bytes must be positive") + + self.group = group if group is not None else dist.group.WORLD + self.rank = int(dist.get_rank(group=self.group)) + self.world_size = int(dist.get_world_size(group=self.group)) + if self.world_size not in _SUPPORTED_WORLD_SIZES: + raise ValueError( + "deterministic collectives require world_size in " + f"{_SUPPORTED_WORLD_SIZES}, got {self.world_size}" + ) + + self.device = self._normalize_device(device) + self.max_size_bytes = int(max_size_bytes) + self._backend = str(dist.get_backend(self.group)).lower() + self._lock = threading.Lock() + self._closed = False + # Keep a lifecycle marker for callers that historically inspected the + # CUDA IPC collective's ``_handle`` while managing the cache. Concrete + # transports own any native resource through their own state. + self._handle = id(self) + # One dtype-agnostic byte workspace is grown on demand and reused by + # reduction collectives. AllGather writes directly into its output. + self._workspace: torch.Tensor | None = None + # A Python-object collective is useful for catching a mismatched new + # signature, but running one on every hot-path call dominates small + # message latency. Validate each local signature once and then rely on + # the standard collective contract that ranks call operations in the + # same order. + self._validated_signatures: set[tuple[Any, ...]] = set() + self._validate_matching_capacity() + + @staticmethod + def _normalize_device( + device: torch.device | str | int | None, + ) -> torch.device: + if device is None: + if torch.cuda.is_available(): + normalized = torch.device("cuda", torch.cuda.current_device()) + else: + normalized = torch.device("cpu") + elif isinstance(device, int): + normalized = torch.device("cuda", device) + else: + normalized = torch.device(device) + + if normalized.type == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError("a CUDA/ROCm device was requested but none is available") + current_device = torch.cuda.current_device() + if normalized.index is None: + normalized = torch.device("cuda", current_device) + if normalized.index != current_device: + raise ValueError( + "the collective device must be the current CUDA/ROCm device; call " + f"torch.cuda.set_device({normalized.index}) first" + ) + return normalized + + @property + def closed(self) -> bool: + """Whether this instance rejects further collective calls.""" + + return self._closed + + @property + def workspace_size_bytes(self) -> int: + """Currently retained reduction workspace size in bytes.""" + + workspace = self._workspace + return 0 if workspace is None else int(workspace.numel()) + + def all_reduce( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + validate_signature: bool = True, + ) -> torch.Tensor: + """Return the fixed balanced-tree sum on every rank.""" + + self._check_open() + self._validate_reduction_input(input) + if out is None: + out = torch.empty_like(input) + self._validate_output(out, input, tuple(input.shape)) + if self.world_size == 1: + out.copy_(input) + return out + + with self._lock: + self._check_open() + if validate_signature: + self._validate_matching_signature("all_reduce", input) + if self._direct_all_reduce(input, out): + return out + rank_inputs = self._all_gather_transport(input) + if not self._fused_reduction( + rank_inputs, + out, + operation="all_reduce", + ): + reduced = self._balanced_tree_sum(rank_inputs) + out.copy_(reduced) + return out + + def all_gather( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + validate_signature: bool = True, + ) -> torch.Tensor: + """Gather rank-ordered input bit patterns along dimension 0.""" + + self._check_open() + self._validate_gather_input(input) + output_shape = (input.size(0) * self.world_size, *input.shape[1:]) + if out is None: + out = torch.empty(output_shape, dtype=input.dtype, device=input.device) + self._validate_output(out, input, output_shape) + if self.world_size == 1: + out.copy_(input) + return out + + with self._lock: + self._check_open() + if validate_signature: + self._validate_matching_signature("all_gather", input) + if self._direct_all_gather(input, out): + return out + self._all_gather_transport(input, gathered_flat=out.view(-1)) + return out + + def all_gather_many( + self, + inputs: tuple[torch.Tensor, ...], + *, + validate_signature: bool = True, + ) -> tuple[torch.Tensor, ...]: + """Gather several tensors through the single-tensor transport ABI.""" + + if not inputs: + raise ValueError("all_gather_many requires at least one input") + return tuple( + self.all_gather(input, validate_signature=validate_signature) for input in inputs + ) + + def reduce_scatter( + self, + input: torch.Tensor, + *, + out: torch.Tensor | None = None, + validate_signature: bool = True, + ) -> torch.Tensor: + """Fixed-tree sum followed by rank-ordered dimension-0 slicing.""" + + self._check_open() + self._validate_reduction_input(input) + if input.dim() == 0: + raise ValueError("reduce_scatter input must have at least one dimension") + if input.size(0) % self.world_size != 0: + raise ValueError( + "reduce_scatter input.size(0) must be divisible by " + f"world_size={self.world_size}; got {input.size(0)}" + ) + rows_per_rank = input.size(0) // self.world_size + output_shape = (rows_per_rank, *input.shape[1:]) + if out is None: + out = torch.empty(output_shape, dtype=input.dtype, device=input.device) + self._validate_output(out, input, output_shape) + if self.world_size == 1: + out.copy_(input) + return out + + with self._lock: + self._check_open() + if validate_signature: + self._validate_matching_signature("reduce_scatter", input) + if self._direct_reduce_scatter(input, out): + return out + rank_inputs = self._all_gather_transport(input) + begin = self.rank * rows_per_rank + # Only this rank's output shard participates in the reduction. The + # previous implementation reduced every global row and sliced the + # result afterwards, doing world_size times more arithmetic than + # ReduceScatter needs. The fixed rank tree is unchanged. + reduced = rank_inputs[:, begin : begin + rows_per_rank] + if not self._fused_reduction(reduced, out, operation="reduce_scatter"): + reduced = self._balanced_tree_sum(reduced) + out.copy_(reduced) + return out + + def reduce_scatter_many( + self, + inputs: tuple[torch.Tensor, ...] | list[torch.Tensor], + *, + outs: tuple[torch.Tensor, ...] | list[torch.Tensor] | None = None, + validate_signature: bool = True, + ) -> tuple[torch.Tensor, ...]: + """Reduce-scatter independent tensors in one fixed-tree collective. + + The tensors are packed along their final dimension, so each tensor's + element still follows the same balanced rank tree as an individual + ``reduce_scatter`` call. This is useful for independent gradient lanes: + packing them together removes one RCCL launch without changing the + floating-point expression for either lane. Inputs must have matching + shape/device/dtype except for the final dimension. + """ + + self._check_open() + values = tuple(inputs) + if not values: + raise ValueError("reduce_scatter_many requires at least one input") + if outs is not None and len(outs) != len(values): + raise ValueError("reduce_scatter_many outs must match the number of inputs") + if len(values) == 1: + return ( + self.reduce_scatter( + values[0], + out=None if outs is None else outs[0], + validate_signature=validate_signature, + ), + ) + + first = values[0] + self._validate_reduction_input(first) + if first.dim() < 2: + raise ValueError( + "reduce_scatter_many inputs must have at least two dimensions " + "when packing independent lanes" + ) + if first.size(0) % self.world_size != 0: + raise ValueError("reduce_scatter_many inputs must have a divisible leading dimension") + for value in values[1:]: + self._validate_reduction_input(value) + if value.dim() != first.dim() or value.shape[:-1] != first.shape[:-1]: + raise ValueError( + "reduce_scatter_many inputs must match in rank and all dimensions " + "except the final dimension" + ) + if value.device != first.device or value.dtype != first.dtype: + raise ValueError("reduce_scatter_many inputs must share device and dtype") + lane_sizes = tuple(int(value.size(-1)) for value in values) + rows_per_rank = first.size(0) // self.world_size + output_shape = (rows_per_rank, *first.shape[1:-1]) + if outs is not None: + for lane_size, out in zip(lane_sizes, outs, strict=True): + self._validate_output( + out, + first, + (*output_shape, lane_size), + ) + + packed_bytes = sum(value.numel() * value.element_size() for value in values) + if self._can_direct_reduce_scatter_many(): + if packed_bytes > self.max_size_bytes: + raise ValueError( + "reduce_scatter_many packed input requires " + f"{packed_bytes} bytes but max_size_bytes={self.max_size_bytes}" + ) + direct_outputs = tuple( + ( + outs[index] + if outs is not None + else torch.empty( + (*output_shape, lane_size), + dtype=first.dtype, + device=first.device, + ) + ) + for index, lane_size in enumerate(lane_sizes) + ) + with self._lock: + self._check_open() + if validate_signature: + self._validate_matching_signature( + f"reduce_scatter_many:{lane_sizes}", + first, + ) + if self._direct_reduce_scatter_many(values, direct_outputs): + return direct_outputs + + if packed_bytes > _PACKED_REDUCE_SCATTER_MAX_BYTES: + # A single packed AllGather moves the same bytes as two separate + # calls but loses RCCL's smaller-message algorithm. Use the + # established per-lane path above the measured crossover; this + # keeps the convenience API from regressing large FFN gradients. + return tuple( + self.reduce_scatter( + value, + out=None if outs is None else outs[index], + validate_signature=validate_signature, + ) + for index, value in enumerate(values) + ) + + if packed_bytes > self.max_size_bytes: + raise ValueError( + "reduce_scatter_many packed input requires " + f"{packed_bytes} bytes but max_size_bytes={self.max_size_bytes}" + ) + packed = torch.cat(values, dim=-1) + packed_out = torch.empty( + (packed.size(0) // self.world_size, *packed.shape[1:]), + dtype=packed.dtype, + device=packed.device, + ) + with self._lock: + self._check_open() + # Include lane boundaries in the signature. Equal packed shapes + # alone do not guarantee that every rank will split the result the + # same way, which could silently associate gradients with the + # wrong lane. + if validate_signature: + self._validate_matching_signature( + f"reduce_scatter_many:{lane_sizes}", + packed, + ) + if self._direct_reduce_scatter(packed, packed_out): + pieces = tuple(packed_out.split(tuple(value.size(-1) for value in values), dim=-1)) + if outs is None: + return pieces + result: list[torch.Tensor] = [] + for piece, out in zip(pieces, outs, strict=True): + out.copy_(piece) + result.append(out) + return tuple(result) + rank_inputs = self._all_gather_transport(packed) + begin = self.rank * rows_per_rank + reduced = rank_inputs[:, begin : begin + rows_per_rank] + if not self._fused_reduction( + reduced, + packed_out, + operation="reduce_scatter", + ): + reduced = self._balanced_tree_sum(reduced) + packed_out.copy_(reduced) + + pieces = tuple(packed_out.split(tuple(value.size(-1) for value in values), dim=-1)) + if outs is None: + return pieces + result: list[torch.Tensor] = [] + for piece, out in zip(pieces, outs, strict=True): + out.copy_(piece) + result.append(out) + return tuple(result) + + def close(self) -> None: + """Close the instance. + + Closing releases the lazily allocated reduction workspace and marks + the lifecycle boundary. Collective calls are blocking at this API. + """ + + with self._lock: + self._workspace = None + self._validated_signatures.clear() + self._closed = True + self._handle = 0 + + def __enter__(self) -> TorchDistributedDeterministicCollective: + self._check_open() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.close() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + def _check_open(self) -> None: + if getattr(self, "_closed", True): + raise RuntimeError("deterministic collective is closed") + + def _validate_tensor(self, input: torch.Tensor) -> None: + if not isinstance(input, torch.Tensor): + raise TypeError(f"input must be a torch.Tensor, got {type(input)!r}") + if input.device != self.device: + raise ValueError(f"input must be on {self.device}, got {input.device}") + if not input.is_contiguous(): + raise ValueError("input must be contiguous") + input_bytes = input.numel() * input.element_size() + if input_bytes > self.max_size_bytes: + raise ValueError( + f"input requires {input_bytes} bytes but max_size_bytes={self.max_size_bytes}" + ) + + def _validate_reduction_input(self, input: torch.Tensor) -> None: + self._validate_tensor(input) + if input.dtype not in _REDUCTION_DTYPES: + raise TypeError( + "deterministic reductions support float32, float16, and bfloat16; " + f"got {input.dtype}" + ) + + def _validate_gather_input(self, input: torch.Tensor) -> None: + self._validate_tensor(input) + if input.dim() == 0: + raise ValueError("all_gather input must have at least one dimension") + + @staticmethod + def _validate_output( + output: torch.Tensor, + input: torch.Tensor, + output_shape: tuple[int, ...], + ) -> None: + if not isinstance(output, torch.Tensor): + raise TypeError(f"out must be a torch.Tensor, got {type(output)!r}") + if output.device != input.device: + raise ValueError("out must be on the same device as input") + if output.dtype != input.dtype: + raise TypeError("out must have the same dtype as input") + if tuple(output.shape) != output_shape: + raise ValueError(f"out must have shape {output_shape}, got {tuple(output.shape)}") + if not output.is_contiguous(): + raise ValueError("out must be contiguous") + + def _validate_matching_signature(self, op_name: str, input: torch.Tensor) -> None: + if self.world_size == 1: + return + signature = (op_name, tuple(input.shape), str(input.dtype), input.numel()) + if signature in self._validated_signatures: + return + signatures: list[tuple[Any, ...] | None] = [None] * self.world_size + dist.all_gather_object(signatures, signature, group=self.group) + if any(peer_signature != signature for peer_signature in signatures): + raise ValueError( + f"all ranks must call {op_name} with matching shapes and dtypes; got {signatures}" + ) + self._validated_signatures.add(signature) + + def _validate_matching_capacity(self) -> None: + if self.world_size == 1: + return + capacities: list[int | None] = [None] * self.world_size + dist.all_gather_object(capacities, self.max_size_bytes, group=self.group) + if any(peer_capacity != self.max_size_bytes for peer_capacity in capacities): + raise ValueError(f"all ranks must use the same max_size_bytes; got {capacities}") + + def _all_gather_transport( + self, + input: torch.Tensor, + *, + gathered_flat: torch.Tensor | None = None, + ) -> torch.Tensor: + # Flattening makes the output contract independent of whether a given + # ProcessGroup implements the concatenation or stacking form of AG. + if self.world_size == 1: + if gathered_flat is None: + gathered_flat = input.clone().view(-1) + else: + gathered_flat.copy_(input.view(-1)) + return gathered_flat.reshape((1, *input.shape)) + input_flat = input.view(-1) + required_elements = self.world_size * input_flat.numel() + if gathered_flat is None: + gathered_flat = self._workspace_for(input, required_elements) + elif ( + gathered_flat.numel() != required_elements + or gathered_flat.dtype != input.dtype + or gathered_flat.device != input.device + or not gathered_flat.is_contiguous() + ): + raise ValueError("gathered transport output has an invalid layout") + if "nccl" in self._backend: + # PyTorch exposes RCCL through the NCCL ProcessGroup API. Keep this + # as a tensor-only transport; reduction happens below. + dist.all_gather_into_tensor(gathered_flat, input_flat, group=self.group) + else: + # Some reference backends (notably Gloo versions without + # all_gather_into_tensor) only implement the list API. + gathered_chunks = list( + gathered_flat.reshape(self.world_size, input_flat.numel()).unbind(0) + ) + dist.all_gather(gathered_chunks, input_flat, group=self.group) + return gathered_flat.reshape((self.world_size, *input.shape)) + + def _workspace_for(self, input: torch.Tensor, required_elements: int) -> torch.Tensor: + required_bytes = required_elements * input.element_size() + workspace = self._workspace + if workspace is None or workspace.numel() < required_bytes: + workspace = torch.empty(required_bytes, dtype=torch.uint8, device=self.device) + self._workspace = workspace + # Tensor.view(dtype) reinterprets the aligned byte allocation without + # an allocation or copy. Restrict the view to the current operation. + return workspace[:required_bytes].view(input.dtype) + + def _direct_all_reduce(self, input: torch.Tensor, output: torch.Tensor) -> bool: + return False + + def _direct_reduce_scatter(self, input: torch.Tensor, output: torch.Tensor) -> bool: + return False + + def _can_direct_reduce_scatter_many(self) -> bool: + return False + + def _direct_reduce_scatter_many( + self, + inputs: tuple[torch.Tensor, ...], + outputs: tuple[torch.Tensor, ...], + ) -> bool: + return False + + def _direct_all_gather(self, input: torch.Tensor, output: torch.Tensor) -> bool: + return False + + @staticmethod + def _balanced_tree_sum(rank_inputs: torch.Tensor) -> torch.Tensor: + world_size = rank_inputs.size(0) + if world_size not in _SUPPORTED_WORLD_SIZES: + raise ValueError( + "balanced reduction requires rank inputs for world_size in " + f"{_SUPPORTED_WORLD_SIZES}, got {world_size}" + ) + # ``rank_inputs`` is the private transport workspace for reductions, + # so fixed-tree nodes can be accumulated in place. This preserves the + # exact pairings while avoiding one temporary allocation per tree node. + stride = 1 + while stride < world_size: + for index in range(0, world_size, 2 * stride): + rank_inputs[index].add_(rank_inputs[index + stride]) + stride *= 2 + return rank_inputs[0] + + @staticmethod + def _fused_reduction( + rank_inputs: torch.Tensor, + output: torch.Tensor, + *, + operation: str, + ) -> bool: + """Use the optional ROCm fused fixed-tree kernel when available. + + The extension is deliberately optional: CPU/Gloo reference collectives + and installations built without the ROCm kernel retain the executable + Python implementation above. + """ + + if getattr(torch.version, "hip", None) is None or not rank_inputs.is_cuda: + return False + try: + from rl_engine import _C + except ImportError: + return False + if operation == "all_reduce": + fn = getattr(_C, "deterministic_collective_rocm_all_reduce", None) + if fn is not None: + fn(rank_inputs, output) + return True + elif operation == "reduce_scatter": + fn = getattr(_C, "deterministic_collective_rocm_reduce_scatter", None) + if fn is not None: + fn(rank_inputs, output) + return True + return False + + +class RCCLDeterministicCollective(TorchDistributedDeterministicCollective): + """Single-node ROCm fixed-tree collective using HIP IPC and RCCL.""" + + backend_id = "rocm_ipc_fixed_tree" + + def __init__( + self, + group: dist.ProcessGroup | None = None, + device: torch.device | str | int | None = None, + *, + max_size_bytes: int = _DEFAULT_MAX_SIZE_BYTES, + ) -> None: + if getattr(torch.version, "hip", None) is None: + raise RuntimeError("RCCL deterministic collectives require a ROCm PyTorch build") + if not torch.cuda.is_available(): + raise RuntimeError("RCCL deterministic collectives require an available ROCm device") + if device is not None and not isinstance(device, int): + requested_device = torch.device(device) + if requested_device.type != "cuda": + raise ValueError( + f"RCCL deterministic collectives require a ROCm device, got {device!r}" + ) + super().__init__(group=group, device=device, max_size_bytes=max_size_bytes) + if self.device.type != "cuda": + raise ValueError( + f"RCCL deterministic collectives require a ROCm device, got {device!r}" + ) + if "nccl" not in self._backend: + raise RuntimeError( + "RCCL deterministic collectives require PyTorch's NCCL process-group API" + ) + self._ipc_handle = 0 + self._ipc_staging: torch.Tensor | None = None + self._initialize_ipc_transport() + + @property + def workspace_size_bytes(self) -> int: + staging = self._ipc_staging + staging_bytes = 0 if staging is None else int(staging.numel()) + return staging_bytes + super().workspace_size_bytes + + def _initialize_ipc_transport(self) -> None: + if self.world_size == 1: + return + try: + from rl_engine import _C + except ImportError: + return + required_symbols = ( + "deterministic_collective_rocm_ipc_allocate", + "deterministic_collective_rocm_ipc_meta", + "deterministic_collective_rocm_ipc_create", + "deterministic_collective_rocm_ipc_synchronize", + "deterministic_collective_rocm_ipc_destroy", + "deterministic_collective_rocm_ipc_stage", + "deterministic_collective_rocm_ipc_all_reduce", + "deterministic_collective_rocm_ipc_all_reduce_input", + "deterministic_collective_rocm_ipc_reduce_scatter", + "deterministic_collective_rocm_ipc_reduce_scatter_input", + "deterministic_collective_rocm_ipc_reduce_scatter_many", + "deterministic_collective_rocm_ipc_all_gather", + "deterministic_collective_rocm_ipc_all_gather_input", + ) + if any(not hasattr(_C, symbol) for symbol in required_symbols): + return + + staging = _C.deterministic_collective_rocm_ipc_allocate(self.max_size_bytes) + handle, offset = _C.deterministic_collective_rocm_ipc_meta(staging) + local_metadata = (socket.gethostname(), handle, int(offset)) + gathered_metadata: list[tuple[str, list[int], int] | None] = [None] * self.world_size + dist.all_gather_object(gathered_metadata, local_metadata, group=self.group) + if any(metadata is None for metadata in gathered_metadata): + raise RuntimeError("failed to exchange ROCm IPC metadata") + complete_metadata = [metadata for metadata in gathered_metadata if metadata is not None] + if len({metadata[0] for metadata in complete_metadata}) != 1: + return + self._ipc_handle = int( + _C.deterministic_collective_rocm_ipc_create( + staging, + [metadata[1] for metadata in complete_metadata], + [metadata[2] for metadata in complete_metadata], + self.rank, + ) + ) + self._ipc_staging = staging + + def _direct_all_reduce(self, input: torch.Tensor, output: torch.Tensor) -> bool: + handle = self._ipc_handle + if not handle: + return False + from rl_engine import _C + + input_bytes = input.numel() * input.element_size() + if ( + _ROCM_IPC_DIRECT_ALL_REDUCE_MAX_BYTES + < input_bytes + < _ROCM_IPC_SHARDED_ALL_REDUCE_MIN_BYTES + and input.numel() % self.world_size == 0 + ): + return False + + if ( + input_bytes <= _ROCM_IPC_DIRECT_ALL_REDUCE_MAX_BYTES + or input.numel() % self.world_size != 0 + ): + _C.deterministic_collective_rocm_ipc_all_reduce_input( + handle, + input, + output, + ) + return True + + shard = self._workspace_for(input, input.numel() // self.world_size) + _C.deterministic_collective_rocm_ipc_reduce_scatter_input( + handle, + input, + shard, + ) + dist.all_gather_into_tensor(output.view(-1), shard, group=self.group) + return True + + def _direct_reduce_scatter(self, input: torch.Tensor, output: torch.Tensor) -> bool: + handle = self._ipc_handle + if not handle: + return False + from rl_engine import _C + + _C.deterministic_collective_rocm_ipc_reduce_scatter_input( + handle, + input, + output, + ) + return True + + def _can_direct_reduce_scatter_many(self) -> bool: + return bool(self._ipc_handle) + + def _direct_reduce_scatter_many( + self, + inputs: tuple[torch.Tensor, ...], + outputs: tuple[torch.Tensor, ...], + ) -> bool: + handle = self._ipc_handle + if not handle: + return False + from rl_engine import _C + + _C.deterministic_collective_rocm_ipc_reduce_scatter_many( + handle, + inputs, + outputs, + ) + return True + + def _direct_all_gather(self, input: torch.Tensor, output: torch.Tensor) -> bool: + handle = self._ipc_handle + input_bytes = input.numel() * input.element_size() + if not handle or input_bytes > _ROCM_IPC_ALL_GATHER_MAX_BYTES: + return False + from rl_engine import _C + + _C.deterministic_collective_rocm_ipc_all_gather_input( + handle, + input, + output, + ) + return True + + def close(self) -> None: + handle = getattr(self, "_ipc_handle", 0) + if handle: + from rl_engine import _C + + _C.deterministic_collective_rocm_ipc_synchronize(handle) + torch.cuda.synchronize(self.device) + self._ipc_handle = 0 + _C.deterministic_collective_rocm_ipc_destroy(handle) + self._ipc_staging = None + super().close() + + +def create_deterministic_collective( + group: dist.ProcessGroup | None = None, + device: torch.device | str | int | None = None, + *, + max_size_bytes: int = _DEFAULT_MAX_SIZE_BYTES, +) -> Any: + """Create the platform-appropriate deterministic collective. + + CUDA uses the native ``DeterministicCollective`` implementation. ROCm uses + HIP IPC or RCCL for rank-ordered transport while preserving the fixed local + reduction tree. The returned object has independent ownership. Shared caches + may replace an entry without closing it immediately because active autograd + contexts can retain the previous instance until their work completes. + """ + + if getattr(torch.version, "hip", None) is not None: + return RCCLDeterministicCollective( + group=group, + device=device, + max_size_bytes=max_size_bytes, + ) + + return DeterministicCollective( + group=group, + device=device, + max_size_bytes=max_size_bytes, + ) + + def collective_for_group( group: dist.ProcessGroup | None, *, min_size_bytes: int = 0, minimum_capacity_bytes: int = _DEFAULT_MAX_SIZE_BYTES, device: torch.device | str | int | None = None, -) -> DeterministicCollective | None: +) -> Any | None: """Return the process-local RL-Kernel collective shared by hot-path ops.""" if group is None: @@ -453,10 +1292,21 @@ def collective_for_group( # entry. Replacing an undersized entry must not invalidate those live # references; normal Python ownership closes it after the last borrower. - collective = DeterministicCollective( + collective = create_deterministic_collective( group=group, device=device_index, max_size_bytes=max(minimum_capacity_bytes, min_size_bytes), ) _COLLECTIVES[key] = collective return collective + + +__all__ = [ + "DETERMINISTIC_ALL_REDUCE_OP", + "DeterministicCollective", + "RCCLDeterministicCollective", + "TorchDistributedDeterministicCollective", + "collective_for_group", + "create_deterministic_collective", + "deterministic_all_reduce_inplace", +] diff --git a/rl_engine/integrations/vime/__init__.py b/rl_engine/integrations/vime/__init__.py index 7eeb8232..8e294f7c 100644 --- a/rl_engine/integrations/vime/__init__.py +++ b/rl_engine/integrations/vime/__init__.py @@ -3,6 +3,14 @@ """Vime adapter entry points without a Vime runtime dependency.""" +from .attention import AttentionProviderResult, AttentionProviderUnavailable, attention_provider from .linear_logp_provider import LinearLogpProviderUnavailable, LinearLogpResult, provider -__all__ = ["LinearLogpProviderUnavailable", "LinearLogpResult", "provider"] +__all__ = [ + "AttentionProviderResult", + "AttentionProviderUnavailable", + "LinearLogpProviderUnavailable", + "LinearLogpResult", + "attention_provider", + "provider", +] diff --git a/rl_engine/integrations/vime/attention.py b/rl_engine/integrations/vime/attention.py new file mode 100644 index 00000000..d9c006bf --- /dev/null +++ b/rl_engine/integrations/vime/attention.py @@ -0,0 +1,477 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Runtime strict-attention provider for Vime's Megatron backend. + +The adapter intentionally accepts and returns structural objects: RL-Kernel +never imports Vime. Vime remains responsible for materializing post-RoPE Q/K/V +for its locally owned rows; this provider owns only the attention core +arithmetic and the ``(out, lse)`` export. + +Two boundaries are deliberate and are enforced rather than documented: + +* Every launch carries exactly one logical batch row and one KV group (that KV + head plus the Q heads that attend to it). The AITER/CK reduction order + depends on the shape of the launch, so both batching and TP head-sharding + would otherwise change the bits: measured on MI300X, raw AITER differs by up + to ``1.5625e-02`` between a batch and its rows submitted singly, and by up to + ``7.8125e-03`` between TP degrees. Pinning the launch shape makes the result + of a row/group independent of how many rows or heads its caller happened to + hold, which is what lets training and rollout compare bitwise across + different batch sizes and TP degrees. It costs roughly 3x forward time. +* CP merges through the transport, never here. The strict ROCm core owns + single-rank attention arithmetic only. At ``CP > 1`` this provider hands the + schedule to :class:`StrictRocmAttentionRuntime`, whose RCCL AG/RS transport + combines the cross-rank ``(out, lse)`` in a fixed balanced rank tree, so no + second merge order is ever defined here. The layout must be ``allgather``; + ``zigzag`` fails closed because the strict CP plan describes one contiguous + block per rank. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import torch + +from rl_engine.kernels.attention_contract import ( + CROSS_CONFIG_BOUND_FIELDS, + AttentionContract, + AttentionDType, + AttentionMode, + AttentionRole, + ReductionSpec, + ShardingSpec, + SplitKVSpec, +) +from rl_engine.kernels.ops.rocm.attention.flash_attn import BACKEND_ID +from rl_engine.kernels.ops.rocm.attention.strict_runtime import StrictRocmAttentionRuntime +from rl_engine.kernels.registry import kernel_registry + + +class AttentionProviderUnavailable(RuntimeError): + """Request Vime's native attention fallback in ``auto`` mode. + + Vime recognizes the marker instead of importing this class, which keeps the + dependency direction from Vime to RL-Kernel at runtime only. + """ + + attention_provider_unavailable = True + + +@dataclass(frozen=True) +class AttentionProviderResult: + """Structural result understood by the Vime attention boundary.""" + + out: torch.Tensor + lse: torch.Tensor + backend_id: str + contract_id: str + provenance: Mapping[str, Any] + + +_DTYPE_TO_CONTRACT = { + torch.bfloat16: AttentionDType.BF16, + torch.float16: AttentionDType.FP16, +} + +# Decode is deliberately absent: it requires KV-cache identity metadata that +# this core does not materialize. A decode request fails closed here rather +# than being silently served by the dense prefill core over a cache the +# provider never validated. +_MODE_BY_NAME = { + "prefill": AttentionMode.PREFILL, + "chunked_prefill": AttentionMode.CHUNKED_PREFILL, +} + +_ROLE_BY_NAME = { + "train": AttentionRole.TRAIN, + "infer": AttentionRole.INFER, +} + + +def _as_positive_int(value: Any, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise AttentionProviderUnavailable(f"{name} must be a positive integer; got {value!r}") + return value + + +def _metadata(request: Any) -> Mapping[str, Any]: + value = getattr(request, "metadata", None) + if not isinstance(value, Mapping): + raise AttentionProviderUnavailable("request.metadata must provide attention metadata") + return value + + +def _request_tensor(request: Any, name: str) -> torch.Tensor: + value = getattr(request, name, None) + if not isinstance(value, torch.Tensor): + raise AttentionProviderUnavailable(f"request.{name} must be a torch.Tensor") + return value + + +def _tp_coordinates(tp_group: Any) -> tuple[int, int]: + if tp_group is not None and hasattr(tp_group, "rank") and hasattr(tp_group, "size"): + return int(tp_group.rank()), int(tp_group.size()) + + import torch.distributed as dist + + if dist.is_available() and dist.is_initialized(): + return dist.get_rank(group=tp_group), dist.get_world_size(group=tp_group) + return 0, 1 + + +def _reject_unsupported_materializations(request: Any, metadata: Mapping[str, Any]) -> None: + """Fail closed on every knob that would change the numerical definition.""" + + if getattr(request, "key_padding_mask", None) is not None: + raise AttentionProviderUnavailable( + "strict ROCm attention materializes each unpadded logical row; " + "pass unpadded per-row Q/K/V instead of a key padding mask" + ) + dropout_p = metadata.get("dropout_p", 0.0) + if dropout_p: + raise AttentionProviderUnavailable( + f"strict attention requires dropout_p=0.0; got {dropout_p!r}" + ) + for unsupported in ( + "alibi_slopes", + "attention_bias", + "logit_soft_cap", + "sliding_window", + "sink_tokens", + ): + if metadata.get(unsupported) is not None: + raise AttentionProviderUnavailable( + f"strict ROCm attention does not materialize {unsupported}" + ) + window = metadata.get("window_size") + if window is not None and tuple(window) != (-1, -1): + raise AttentionProviderUnavailable( + f"strict ROCm attention requires a full causal/full window; got {window!r}" + ) + + +def _contract_for_request( + request: Any, +) -> tuple[AttentionContract, float, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Validate the request and derive the explicit WS2 attention contract.""" + + metadata = _metadata(request) + _reject_unsupported_materializations(request, metadata) + + query = _request_tensor(request, "query") + key = _request_tensor(request, "key") + value = _request_tensor(request, "value") + if query.ndim != 4 or key.ndim != 4 or value.ndim != 4: + raise AttentionProviderUnavailable("query/key/value must be 4-D [B, H, S, D] tensors") + if key.shape != value.shape: + raise AttentionProviderUnavailable("key and value must share one shape") + if query.shape[0] != key.shape[0]: + raise AttentionProviderUnavailable("query and key must share the logical batch size") + if query.shape[-1] != key.shape[-1]: + raise AttentionProviderUnavailable("query and key must share head_dim") + if query.dtype not in _DTYPE_TO_CONTRACT: + raise AttentionProviderUnavailable( + f"strict ROCm attention supports BF16/FP16 only; got {query.dtype}" + ) + if key.dtype != query.dtype or value.dtype != query.dtype: + raise AttentionProviderUnavailable("query/key/value must share one dtype") + if not (query.device == key.device == value.device): + raise AttentionProviderUnavailable("query/key/value must share one device") + + batch_size, local_q_heads, query_len, head_dim = query.shape + local_kv_heads = key.shape[1] + kv_len = key.shape[2] + if local_q_heads % local_kv_heads: + raise AttentionProviderUnavailable( + f"local Q heads={local_q_heads} must be divisible by local KV heads={local_kv_heads}" + ) + + cp = getattr(request, "context_parallel", None) + cp_world_size = _as_positive_int(getattr(cp, "world_size", None), "context_parallel.world_size") + cp_rank = getattr(cp, "rank", None) + if ( + isinstance(cp_rank, bool) + or not isinstance(cp_rank, int) + or not 0 <= cp_rank < cp_world_size + ): + raise AttentionProviderUnavailable( + f"context_parallel.rank={cp_rank!r} is invalid for CP={cp_world_size}" + ) + cp_layout = getattr(cp, "layout", None) + if cp_layout not in ({"single"} if cp_world_size == 1 else {"zigzag", "allgather"}): + raise AttentionProviderUnavailable( + "context_parallel layout does not describe local CP token ownership" + ) + if cp_world_size > 1 and cp_layout != "allgather": + # The strict CP plan describes one contiguous block per rank. A zigzag + # rank owns two discontiguous token runs, so accepting it here would + # silently disagree with the block manifest the transport validates. + raise AttentionProviderUnavailable( + f"CP={cp_world_size} requires the 'allgather' layout; got {cp_layout!r}, " + "whose block ownership the strict CP plan does not describe" + ) + + tp_rank, tp_world_size = _tp_coordinates(getattr(request, "tensor_parallel_group", None)) + declared_tp_rank = metadata.get("tp_rank") + declared_tp_world_size = metadata.get("tp_world_size") + if declared_tp_rank is not None and declared_tp_rank != tp_rank: + raise AttentionProviderUnavailable( + f"metadata tp_rank={declared_tp_rank} disagrees with TP group rank={tp_rank}" + ) + if declared_tp_world_size is not None and declared_tp_world_size != tp_world_size: + raise AttentionProviderUnavailable( + f"metadata tp_world_size={declared_tp_world_size} disagrees with " + f"TP group size={tp_world_size}" + ) + + global_q_heads = _as_positive_int(metadata.get("global_q_heads"), "global_q_heads") + global_kv_heads = _as_positive_int(metadata.get("global_kv_heads"), "global_kv_heads") + if local_q_heads * tp_world_size != global_q_heads: + raise AttentionProviderUnavailable( + "local Q heads and TP group do not cover global_q_heads exactly: " + f"{local_q_heads} * {tp_world_size} != {global_q_heads}" + ) + if local_kv_heads * tp_world_size != global_kv_heads: + raise AttentionProviderUnavailable( + "local KV heads and TP group do not cover global_kv_heads exactly: " + f"{local_kv_heads} * {tp_world_size} != {global_kv_heads}" + ) + + mode_name = str(metadata.get("attention_mode", "prefill")) + if mode_name == "decode": + raise AttentionProviderUnavailable( + "decode requires KV-cache identity metadata (cache_position, block table, " + "prefix-cache key) that the strict dense core does not materialize; use the " + "paged decode path" + ) + if mode_name not in _MODE_BY_NAME: + raise AttentionProviderUnavailable(f"unsupported attention_mode={mode_name!r}") + mode = _MODE_BY_NAME[mode_name] + role_name = str(metadata.get("role", "train")) + if role_name not in _ROLE_BY_NAME: + raise AttentionProviderUnavailable(f"unsupported role={role_name!r}") + role = _ROLE_BY_NAME[role_name] + + causal = metadata.get("causal", True) + if not isinstance(causal, bool): + raise AttentionProviderUnavailable(f"causal must be a bool; got {causal!r}") + if mode is AttentionMode.PREFILL and query_len != kv_len: + raise AttentionProviderUnavailable( + "prefill requires the query and KV lengths to describe one logical sequence; " + f"got Sq={query_len} and Skv={kv_len}" + ) + + if query_len > kv_len: + raise AttentionProviderUnavailable( + f"causal attention requires Sq <= Skv; got Sq={query_len} and Skv={kv_len}" + ) + + scale = metadata.get("softmax_scale") + resolved_scale = 1.0 / math.sqrt(head_dim) if scale is None else float(scale) + + # Each CP rank owns one contiguous block of the logical sequence; at CP=1 + # that block is the whole sequence. The allgather layout checked above is + # what makes the block contiguous. + sharding = ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + global_q_heads=global_q_heads, + global_kv_heads=global_kv_heads, + local_q_head_start=tp_rank * local_q_heads, + local_q_heads=local_q_heads, + local_kv_head_start=tp_rank * local_kv_heads, + local_kv_heads=local_kv_heads, + global_sequence_length=kv_len * cp_world_size, + local_sequence_length=kv_len, + global_block_indices=(cp_rank,), + global_block_token_starts=(cp_rank * kv_len,), + local_block_offsets=(0, kv_len), + ) + # Causal alignment: the query block is the tail of the logical sequence, so + # every batch entry carries the same offset between Q row 0 and KV token 0. + causal_offsets = (kv_len - query_len,) * batch_size if causal else None + + contract = AttentionContract( + role=role, + mode=mode, + dtype=_DTYPE_TO_CONTRACT[query.dtype], + batch_size=batch_size, + query_sequence_length=query_len if mode is not AttentionMode.PREFILL else kv_len, + head_dim=head_dim, + causal=causal, + causal_offsets=causal_offsets, + sharding=sharding, + reduction=ReductionSpec(), + split_kv=SplitKVSpec.disabled(), + kv_cache=None, + rope=None, + export_lse=True, + ) + key_position_ids = _key_position_ids(metadata, query.device, batch_size, kv_len) + return contract, resolved_scale, query, key, value, key_position_ids + + +def _key_position_ids( + metadata: Mapping[str, Any], + device: torch.device, + batch_size: int, + kv_len: int, +) -> torch.Tensor: + """Resolve the global KV token positions for this request. + + Position identity is part of the contract, not an implementation detail: it + is what makes a training-side full-sequence call and a rollout-side chunk + provably describe the same logical tokens. Vime may declare it; otherwise + the canonical contiguous ``[0, kv_len)`` block is used. + """ + + declared = metadata.get("key_position_ids") + if declared is None: + return ( + torch.arange(kv_len, device=device, dtype=torch.int64) + .unsqueeze(0) + .expand(batch_size, kv_len) + .contiguous() + ) + positions = declared if isinstance(declared, torch.Tensor) else torch.as_tensor(declared) + positions = positions.to(device=device, dtype=torch.int64) + if positions.ndim == 1: + positions = positions.unsqueeze(0).expand(batch_size, -1) + if tuple(positions.shape) != (batch_size, kv_len): + raise AttentionProviderUnavailable( + f"key_position_ids must have shape {(batch_size, kv_len)}; " + f"got {tuple(positions.shape)}" + ) + if kv_len > 1 and bool((positions[:, 1:] - positions[:, :-1] != 1).any()): + raise AttentionProviderUnavailable( + "key_position_ids must describe one contiguous increasing token block" + ) + return positions.contiguous() + + +def attention_provider(request: Any) -> AttentionProviderResult: + """Compute Vime attention on the explicit WS2 strict ROCm contract. + + Materializes each logical batch row independently so batch composition + cannot change the bits, then returns the stacked ``(out, lse)`` together + with the dispatch provenance Vime records alongside the result. + """ + + contract, scale, query, key, value, key_positions = _contract_for_request(request) + + dispatch = kernel_registry.get_attention_op(contract, requested_backend=BACKEND_ID) + if dispatch.provenance["actual_backend"] != BACKEND_ID or dispatch.provenance["fallback"]: + raise RuntimeError("explicit strict attention dispatch changed during materialization") + + query_len = query.shape[2] + cp_world_size = contract.sharding.cp_world_size + if cp_world_size > 1: + cp_group = getattr(request, "context_parallel_group", None) + if cp_group is None: + raise AttentionProviderUnavailable( + f"CP={cp_world_size} requires request.context_parallel_group so the strict " + "RCCL AG/RS transport can be built; CP=1 does not need one" + ) + else: + cp_group = None + + # One runtime owns the launch schedule for both CP degrees, so the + # per-(batch row, KV group) launch loop that makes the result TP-degree + # invariant cannot drift between the single-rank and CP paths. + runtime = StrictRocmAttentionRuntime(process_group=cp_group, core=dispatch.op) + runtime_result = runtime.forward_with_lse( + query, + key, + value, + contract=contract, + causal=contract.causal, + scale=scale, + cp_world_size=cp_world_size, + query_position_ids=key_positions[:, -query_len:], + key_position_ids=key_positions, + # At CP=1 this rank already holds the logical sequence in position + # order, so the reorder the CP path needs would be a no-op copy. + positions_are_sorted=cp_world_size == 1, + ) + + out = runtime_result.out + lse = runtime_result.lse + core_provenance = runtime_result.provenance["core"] + launches = runtime_result.provenance["core_launch_count"] + + provenance = dict(dispatch.provenance) + provenance["core"] = core_provenance + provenance["request"] = { + "query_shape": list(query.shape), + "key_shape": list(key.shape), + "dtype": str(query.dtype).replace("torch.", ""), + "causal": contract.causal, + "softmax_scale": scale, + "tp_rank": contract.sharding.tp_rank, + "tp_world_size": contract.sharding.tp_world_size, + "cp_rank": contract.sharding.cp_rank, + "cp_world_size": contract.sharding.cp_world_size, + } + provenance["execution"] = { + "role": "vime_attention", + "strict_backend": True, + "launch_granularity": "one_batch_row_one_kv_group", + "core_launches": launches, + "batch_rows_materialized_independently": True, + "kv_groups_materialized_independently": True, + "attention_mode": contract.mode.value, + } + provenance["cp_row_ownership"] = { + "cp_rank": contract.sharding.cp_rank, + "cp_world_size": contract.sharding.cp_world_size, + "layout": getattr(request.context_parallel, "layout"), + "local_token_rows": contract.sharding.local_sequence_length, + # The merge happens in the RCCL AG/RS transport's fixed rank tree, not + # in this provider; the flag records that CP was an axis at all. + "cp_is_merge_axis": cp_world_size > 1, + "cp_merge_owner": ( + runtime_result.provenance["communication_backend"] if cp_world_size > 1 else "none" + ), + } + provenance["lse_domain"] = "attention" + # The qualified ROCm core's reduction order depends on the launch head + # count, so raw AITER gives a head shard computed under TP=4 different bits + # from the same shard under TP=8 at some shapes. RL-Kernel removes the + # dependence instead of binding the degree: every launch carries exactly one + # KV group, which is bitwise TP-invariant at 12 of 12 measured points. + # ``contract_id`` still encodes TP/CP so the preflight can compare the two + # sides, but it is no longer what buys the invariance. + provenance["cross_config_binding"] = { + "bound_fields": list(CROSS_CONFIG_BOUND_FIELDS), + "tp_world_size": contract.sharding.tp_world_size, + "cp_world_size": contract.sharding.cp_world_size, + "binding_token": "contract_id", + "tp_degree_invariant": True, + "invariance_mechanism": "one_kv_group_per_launch", + "reason": ( + "AITER/CK dense MHA reduction order depends on the launch head count, so " + "every launch is pinned to one KV group and its Q heads; the result of a " + "head shard is then independent of the TP degree that produced it" + ), + } + return AttentionProviderResult( + out=out, + lse=lse, + backend_id=dispatch.capability.backend_id, + contract_id=contract.cross_rank_fingerprint(), + provenance=provenance, + ) + + +__all__ = [ + "AttentionProviderResult", + "AttentionProviderUnavailable", + "attention_provider", +] diff --git a/rl_engine/kernels/attention_contract.py b/rl_engine/kernels/attention_contract.py index 79c24823..139c6bc1 100644 --- a/rl_engine/kernels/attention_contract.py +++ b/rl_engine/kernels/attention_contract.py @@ -11,6 +11,8 @@ from __future__ import annotations +import hashlib +import json from dataclasses import dataclass, field from enum import Enum from typing import Any, Iterable, TypeVar @@ -19,14 +21,20 @@ # Stable identities for Attention arithmetic shared by training and rollout. -# The FA4 core is the strict production path. The materializing RL-Kernel core -# remains available as an explicit reference and capability-gap fallback. +# The FA4 core is the strict production path on CUDA; AITER/CK dense MHA is its +# ROCm counterpart. The materializing RL-Kernel core remains available as an +# explicit reference and capability-gap fallback. STRICT_ATTENTION_PRODUCTION_CORE_ID = "rlkernel.attention.flash_attention4.num_splits1.v1" +STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID = "rlkernel.attention.rocm.aiter_ck_dense_mha.v1" STRICT_ATTENTION_REFERENCE_CORE_ID = "rlkernel.attention.deterministic_core.v1" # Compatibility alias for callers that explicitly select the original core. STRICT_ATTENTION_CORE_ID = STRICT_ATTENTION_REFERENCE_CORE_ID STRICT_ATTENTION_FA4_SCHEDULE_ID = "single_batch_flash_attention4_num_splits1" +STRICT_ATTENTION_ROCM_SCHEDULE_ID = "single_batch_aiter_ck_dense_mha_no_splitkv" STRICT_ATTENTION_SCHEDULE_ID = "single_batch_single_query_global_kv_blocks" +# The distributed strict path executes the full-KV core above. This identifies +# the fixed, pre-overlap communication schedule around that core. +STRICT_ATTENTION_RING_SCHEDULE_ID = "rlkernel.attention.strict_ring_state.v1" class AttentionContractError(ValueError): @@ -591,6 +599,61 @@ def validate_split_kv_alignment( ) +# What a bitwise cross-config comparison actually requires. +# +# TP degree and batch size are deliberately absent. The qualified ROCm vendor +# core (AITER/CK dense MHA) has a launch-shape-dependent reduction order, so +# both would otherwise change the bits. Measured on MI300X (BF16, Hq=32/ +# Hkv=8/D=128, causal), raw AITER differs by up to 1.5625e-02 between a batch +# and its rows submitted singly, and by up to 7.8125e-03 between TP degrees. +# The Vime provider removes the dependence by pinning every launch to one batch +# row and one KV group, at roughly 3x forward time, so a head shard's result no +# longer depends on the batch size or TP degree that produced it. Requiring +# those to match would therefore reject comparisons that are in fact bitwise +# equal. What remains here is the set that genuinely changes the arithmetic. +CROSS_CONFIG_BOUND_FIELDS = ("dtype", "head_dim", "causal", "export_lse") + + +def validate_cross_config_alignment( + training: "AttentionContract", + rollout: "AttentionContract", +) -> None: + """Fail closed unless both sides describe one comparable attention invocation. + + Names the field that diverged, so a cross-config drift investigation does + not start from one opaque fingerprint mismatch. + """ + + if not isinstance(training, AttentionContract) or not isinstance(rollout, AttentionContract): + raise AttentionContractError("both sides must be AttentionContract instances") + + layout_mismatches = [ + name + for name in ("global_q_heads", "global_kv_heads") + if getattr(training.sharding, name) != getattr(rollout.sharding, name) + ] + if layout_mismatches: + raise AttentionContractError( + "training and rollout describe different global head layouts: " + + ", ".join(layout_mismatches) + ) + + scalar_mismatches = [ + name + for name in CROSS_CONFIG_BOUND_FIELDS + if getattr(training, name) != getattr(rollout, name) + ] + if scalar_mismatches: + raise AttentionContractError( + "training and rollout attention contracts differ: " + ", ".join(scalar_mismatches) + ) + + if training.split_kv != rollout.split_kv: + raise AttentionContractError("training and rollout Split-KV policies differ") + if training.reduction != rollout.reduction: + raise AttentionContractError("training and rollout reduction specs differ") + + @dataclass(frozen=True, order=True) class SplitKVRuntimeCoordinate: """Identity of one batch/rank/owner Split-KV runtime plan.""" @@ -1501,6 +1564,41 @@ def to_dict(self) -> dict[str, Any]: "projections": projections, } + def cross_rank_fingerprint(self) -> str: + """Rank-independent identity for preflight agreement across ranks. + + Excludes ``tp_rank``/``cp_rank`` and the local head/sequence bounds they + derive, so every rank of one logical attention invocation computes the + same value. All-gathering this fingerprint together with the resolved + backend id and aborting on mismatch is the documented preflight for + distributed dispatch; ``requested_backend="auto"`` is not + distributed-safe without it. + """ + + payload = self.to_dict() + payload["sharding"] = { + key: value + for key, value in payload["sharding"].items() + if key + not in { + "tp_rank", + "cp_rank", + "local_q_head_start", + "local_q_heads", + "local_kv_head_start", + "local_kv_heads", + "local_sequence_length", + "global_block_indices", + "global_block_token_starts", + "local_block_offsets", + } + } + # Note: Any future extensions to this payload MUST maintain strict JSON + # serialization determinism across environments to prevent cross-rank + # hashing mismatches. + encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(encoded.encode("utf-8")).hexdigest() + @dataclass(frozen=True) class AttentionBackendCapability: @@ -1680,7 +1778,12 @@ class AttentionDispatchResult: "STRICT_ATTENTION_FA4_SCHEDULE_ID", "STRICT_ATTENTION_PRODUCTION_CORE_ID", "STRICT_ATTENTION_REFERENCE_CORE_ID", + "STRICT_ATTENTION_RING_SCHEDULE_ID", + "STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID", + "STRICT_ATTENTION_ROCM_SCHEDULE_ID", "STRICT_ATTENTION_SCHEDULE_ID", + "CROSS_CONFIG_BOUND_FIELDS", + "validate_cross_config_alignment", "validate_split_kv_alignment", "validate_split_kv_plan_set_alignment", ] diff --git a/rl_engine/kernels/ops/cuda/attention/__init__.py b/rl_engine/kernels/ops/cuda/attention/__init__.py index 91fa6f99..1d9c5621 100644 --- a/rl_engine/kernels/ops/cuda/attention/__init__.py +++ b/rl_engine/kernels/ops/cuda/attention/__init__.py @@ -1,15 +1,24 @@ -from .deterministic_attn import DeterministicAttentionOp +# File: rl_engine/kernels/ops/cuda/attention/__init__.py + +from .deterministic_attn import ( + DeterministicAttentionCoreResult, + DeterministicAttentionOp, + RLKernelDeterministicAttentionCore, +) from .flash_attn import FlashAttentionOp, StrictFlashAttention4Core, StrictFlashAttentionUnavailable from .prefix_shared_attn import PrefixSharedAttentionOp __all__ = [ + "DeterministicAttentionCoreResult", "DeterministicAttentionOp", + "RLKernelDeterministicAttentionCore", "FlashAttentionOp", "PrefixSharedAttentionOp", "StrictFlashAttention4Core", "StrictFlashAttentionUnavailable", ] + # CP communication and FlashInfer are optional layers owned by later WS2 PRs. # Keep the base Attention package importable while those PRs are developed or # tested independently, then expose their symbols automatically when present. @@ -27,6 +36,7 @@ CPCommunicationStatus, CUDAAGRSAttentionCPCommunication, P2PNCCLAttentionCPCommunication, + RCCLAGRSAttentionCPCommunication, sort_attention_cp_partial_states, ) except ModuleNotFoundError as exc: @@ -46,6 +56,7 @@ "CPCommunicationStatus", "CUDAAGRSAttentionCPCommunication", "P2PNCCLAttentionCPCommunication", + "RCCLAGRSAttentionCPCommunication", "sort_attention_cp_partial_states", ] diff --git a/rl_engine/kernels/ops/cuda/attention/cp_comm.py b/rl_engine/kernels/ops/cuda/attention/cp_comm.py index 03a05d23..da871488 100644 --- a/rl_engine/kernels/ops/cuda/attention/cp_comm.py +++ b/rl_engine/kernels/ops/cuda/attention/cp_comm.py @@ -40,7 +40,12 @@ import torch -CPCommunicationBackend = Literal["cuda_ag_rs", "p2p_nccl_reference", "local_debug"] +CPCommunicationBackend = Literal[ + "cuda_ag_rs", + "rccl_ag_rs", + "p2p_nccl_reference", + "local_debug", +] CPCommunicationStatus = Literal["interface_only", "implemented"] @@ -198,12 +203,17 @@ class AttentionCPCommunicationPlan: def validate(self) -> None: self.parallel.validate() - if self.backend not in {"cuda_ag_rs", "p2p_nccl_reference", "local_debug"}: + if self.backend not in { + "cuda_ag_rs", + "rccl_ag_rs", + "p2p_nccl_reference", + "local_debug", + }: raise ValueError(f"unsupported CP communication backend: {self.backend}") if self.status not in {"interface_only", "implemented"}: raise ValueError(f"unsupported CP communication status: {self.status}") if self.pattern != "ag_rs": - raise ValueError("PR7 CP communication must use the custom CUDA AG/RS interface") + raise ValueError("PR7 CP communication must use the self-owned AG/RS interface") if self.compute_communication != "decoupled": raise ValueError("PR7 CP communication must keep compute and communication decoupled") if self.merge_order != "global_block_index": @@ -242,6 +252,15 @@ def provenance(self) -> dict[str, object]: "cp_comm_strict_kv_communication": "all_gather", "cp_comm_strict_position_communication": "all_gather", "cp_comm_strict_backward": "rs_out_backward_ag_then_ag_qkv_backward_rs", + "cp_comm_runtime": ( + "rccl" + if self.backend == "rccl_ag_rs" + else "nccl" if self.backend in {"cuda_ag_rs", "p2p_nccl_reference"} else "local" + ), + # The ROCm path intentionally transports tensors and performs the + # arithmetic in the deterministic core; only the CUDA IPC path + # owns a numeric collective reduction kernel. + "cp_comm_attention_numeric_reduction": self.backend == "cuda_ag_rs", "cp_comm_expected_kv_token_range": ( None if self.expected_kv_token_range is None else list(self.expected_kv_token_range) ), @@ -335,15 +354,23 @@ def forward( ctx.sequence_dim = int(sequence_dim) ctx.rank = int(rank) ctx.root = int(root) + ctx.world_size = int(getattr(collective, "world_size", 1)) packed = full.movedim(ctx.sequence_dim, 0).contiguous() + ctx.full_shape = tuple(packed.shape) if ctx.rank != ctx.root: packed = torch.zeros_like(packed) - local = collective.reduce_scatter(packed) + # The ROCm transport adapter exposes an explicit root-owned scatter; + # CUDA's IPC collective keeps the historical reduce_scatter entrypoint. + scatter = getattr(collective, "scatter", None) + local = scatter(packed) if callable(scatter) else collective.reduce_scatter(packed) return local.movedim(0, ctx.sequence_dim).contiguous() @staticmethod def backward(ctx, grad_local: torch.Tensor) -> tuple[torch.Tensor, None, None, None, None]: packed = grad_local.movedim(ctx.sequence_dim, 0).contiguous() + # The forward scatter has one authoritative full input on root. Its + # backward is the dual gather of every rank's local output gradient; + # non-root full inputs were zeroed in forward and receive no gradient. grad_full = ctx.collective.all_gather(packed).movedim(0, ctx.sequence_dim).contiguous() if ctx.rank != ctx.root: grad_full.zero_() @@ -354,6 +381,7 @@ class CUDAAGRSAttentionCPCommunication: """Deterministic CUDA AG/RS adapter backed by PR311/PR312.""" backend_id = "cuda_ag_rs" + collective_label = "self-owned CUDA AG/RS" supports_autograd = True def __init__(self, *, process_group: Any = None, collective: Any = None) -> None: @@ -367,7 +395,7 @@ def _get_collective(self, plan: AttentionCPCommunicationPlan): from rl_engine.distributed.collectives import collective_for_group except ImportError as exc: raise AttentionCPCommunicationUnavailable( - "self-owned CUDA AG/RS requires PR311/PR312 DeterministicCollective" + f"{self.collective_label} requires PR311/PR312 DeterministicCollective" ) from exc try: dist = self._dist() @@ -380,7 +408,7 @@ def _get_collective(self, plan: AttentionCPCommunicationPlan): raise RuntimeError("the CP process group is unavailable") except (RuntimeError, ValueError, TypeError) as exc: raise AttentionCPCommunicationUnavailable( - f"self-owned CUDA AG/RS is unavailable: {exc}" + f"{self.collective_label} is unavailable: {exc}" ) from exc if self._collective.world_size != plan.parallel.cp_world_size: raise AttentionCPCommunicationUnavailable( @@ -561,6 +589,43 @@ def _validate_cuda_plan(self, plan: AttentionCPCommunicationPlan) -> None: raise AttentionCPCommunicationUnavailable("self-owned CUDA AG/RS requires CUDA") +class RCCLAGRSAttentionCPCommunication(CUDAAGRSAttentionCPCommunication): + """ROCm AG/RS adapter using RCCL only as rank-ordered tensor transport. + + The transport itself is the shared :func:`collective_for_group` collective, + which resolves to ``RCCLDeterministicCollective`` on ROCm. CUDA and ROCm + therefore evaluate one balanced rank tree from one implementation rather + than two copies that can silently drift apart. + """ + + backend_id = "rccl_ag_rs" + collective_label = "self-owned RCCL AG/RS" + supports_autograd = True + transport_only = True + supports_async_overlap = False + supports_compute_communication_fusion = False + + def _dist(self): + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL AG/RS requires initialized torch.distributed" + ) + return dist + + def _validate_cuda_plan(self, plan: AttentionCPCommunicationPlan) -> None: + plan.validate() + if plan.backend != "rccl_ag_rs" or plan.status != "implemented": + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL AG/RS requires an implemented rccl_ag_rs plan" + ) + if torch.version.hip is None or not torch.cuda.is_available(): + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL AG/RS requires an available ROCm device" + ) + + class P2PNCCLAttentionCPCommunication: """Correctness-first P2P NCCL implementation of the CP protocol. @@ -1292,6 +1357,7 @@ def _rank_in_world(rank: int, world_size: int, name: str) -> None: "CPCommunicationBackend", "CPCommunicationStatus", "CUDAAGRSAttentionCPCommunication", + "RCCLAGRSAttentionCPCommunication", "P2PNCCLAttentionCPCommunication", "sort_attention_cp_partial_states", ] diff --git a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py index c3ef6aa3..1ed6d47c 100644 --- a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""CUDA deterministic standard-softmax attention (issue #147). +"""Deterministic standard-softmax attention for CUDA and ROCm (issue #147). Forward: QK → masked softmax+LSE → PV (all FP32 intermediate). Backward: dP → softmax_bwd → dQ/dK/dV with §4.1 fixed GQA order. @@ -28,6 +28,8 @@ from rl_engine.utils.logger import logger _HEAD_DIM = 128 +_IS_ROCM = torch.version.hip is not None +_GPU_PLATFORM = "ROCm" if _IS_ROCM else "CUDA" @dataclass(frozen=True) @@ -92,7 +94,7 @@ def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): class DeterministicAttentionOp: - """Batch-invariant standard softmax attention on CUDA. + """Batch-invariant standard softmax attention on a CUDA or ROCm GPU. Materializes full FP32 scores/P. Public surface matches NativeAttentionOp so #108 harness can call forward(**inputs) with key_padding_mask. @@ -105,13 +107,13 @@ class DeterministicAttentionOp: def __init__(self) -> None: if not _EXT_AVAILABLE or not hasattr(_C, "deterministic_attention_forward"): raise RuntimeError( - "Deterministic CUDA attention kernel is unavailable. " - "Rebuild the extension with `pip install -e .` on a CUDA build." + f"Deterministic {_GPU_PLATFORM} attention kernel is unavailable. " + "Rebuild the native extension for the active GPU platform." ) if not hasattr(_C, "deterministic_attention_backward"): raise RuntimeError( - "Deterministic CUDA attention backward kernel is unavailable. " - "Rebuild the extension with `pip install -e .` on a CUDA build." + f"Deterministic {_GPU_PLATFORM} attention backward kernel is unavailable. " + "Rebuild the native extension for the active GPU platform." ) logger.info("Successfully linked to _C.deterministic_attention_forward/backward.") @@ -208,7 +210,7 @@ def _validate_inputs( if k.dtype != q.dtype or v.dtype != q.dtype: raise ValueError("q, k, v must share the same dtype") if not (q.is_cuda and k.is_cuda and v.is_cuda): - raise ValueError("q, k, v must be CUDA tensors") + raise ValueError("q, k, v must be GPU tensors") if key_padding_mask is not None: if key_padding_mask.dtype != torch.bool: raise ValueError("key_padding_mask must be bool") @@ -222,15 +224,20 @@ def _validate_inputs( class RLKernelDeterministicAttentionCore: - """Materializing CUDA reference core shared by training and rollout. + """Materializing GPU reference core shared by training and rollout. This remains useful for correctness and capability-gap diagnosis. The - production default is the shared FA4 CuTe core with ``num_splits=1``. + production default is the shared FA4 CuTe core with ``num_splits=1`` on + CUDA, and AITER CK dense MHA on ROCm. """ core_id = STRICT_ATTENTION_CORE_ID strict_schedule = STRICT_ATTENTION_SCHEDULE_ID - backend_id = "rlkernel.cuda.deterministic_attention" + backend_id = ( + "rlkernel.rocm.deterministic_attention" + if _IS_ROCM + else "rlkernel.cuda.deterministic_attention" + ) merge_order = "global_block_index" accum_dtype = "fp32" downcast_at = "final_write" @@ -248,7 +255,7 @@ def __init__( if not isinstance(requested, SplitKVSpec): raise TypeError("split_kv must be a SplitKVSpec") if requested.mode is not SplitKVMode.DISABLED: - raise ValueError("the strict CUDA Attention core requires Split-KV to be disabled") + raise ValueError("the strict GPU Attention core requires Split-KV to be disabled") self.split_kv = requested self._op = DeterministicAttentionOp() @@ -324,7 +331,7 @@ def _validate_positions( return if query_position_ids is None or key_position_ids is None: raise ValueError( - "strict CUDA Attention requires query_position_ids and " "key_position_ids" + "strict GPU Attention requires query_position_ids and " "key_position_ids" ) expected_q_shape = (q.size(0), q.size(2)) expected_k_shape = (k.size(0), k.size(2)) @@ -349,6 +356,6 @@ def _validate_positions( raise ValueError("key_position_ids must be contiguous and increasing") if not torch.equal(query_position_ids, key_position_ids[:, -q.size(2) :]): raise ValueError( - "strict CUDA Attention requires queries to be the trailing " + "strict GPU Attention requires queries to be the trailing " "contiguous positions of the logical KV sequence" ) diff --git a/rl_engine/kernels/ops/cuda/attention/flashinfer_paged_attention.py b/rl_engine/kernels/ops/cuda/attention/flashinfer_paged_attention.py index 1735a19e..48785e6c 100644 --- a/rl_engine/kernels/ops/cuda/attention/flashinfer_paged_attention.py +++ b/rl_engine/kernels/ops/cuda/attention/flashinfer_paged_attention.py @@ -28,6 +28,8 @@ STRICT_ATTENTION_CORE_ID, STRICT_ATTENTION_FA4_SCHEDULE_ID, STRICT_ATTENTION_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, STRICT_ATTENTION_SCHEDULE_ID, AttentionContractError, SplitKVExecutionPlan, @@ -776,7 +778,8 @@ def _run_strict_cp( communication, "supports_autograd", False ): raise FlashInferUnavailable( - "strict training requires the autograd-capable self-owned CUDA AG/RS backend" + "strict training requires an autograd-capable self-owned CUDA AG/RS " + "or ROCm RCCL AG/RS backend" ) query_start, query_end = plan.query_token_ranges[plan.parallel.cp_rank] key_start, key_end = _cp_owner_ranges(plan)[plan.parallel.cp_rank] @@ -1464,13 +1467,12 @@ def _validate_strict_core(core: Any) -> None: raise ValueError("strict Attention core must implement forward_with_lse") expected_schedules = { STRICT_ATTENTION_PRODUCTION_CORE_ID: STRICT_ATTENTION_FA4_SCHEDULE_ID, + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID: STRICT_ATTENTION_ROCM_SCHEDULE_ID, STRICT_ATTENTION_CORE_ID: STRICT_ATTENTION_SCHEDULE_ID, } core_id = getattr(core, "core_id", None) if core_id not in expected_schedules: - raise ValueError( - "strict Attention core ID must identify the FA4 production core or explicit reference" - ) + raise ValueError("strict Attention core ID is not an exact supported identity") if getattr(core, "strict_schedule", None) != expected_schedules[core_id]: raise ValueError("strict Attention core schedule does not match its exact core identity") required = { @@ -1486,9 +1488,8 @@ def _validate_strict_core(core: Any) -> None: raise ValueError( "strict Attention core has incompatible arithmetic identity: " + ", ".join(mismatches) ) - if core_id == STRICT_ATTENTION_PRODUCTION_CORE_ID: + if core_id in {STRICT_ATTENTION_PRODUCTION_CORE_ID, STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID}: production_required = { - "backend_id": "flash_attention_4.cute", "native_attention_arithmetic": True, "num_splits": 1, "deterministic_backward": True, @@ -1502,14 +1503,28 @@ def _validate_strict_core(core: Any) -> None: ] if production_mismatches: raise ValueError( - "strict FA4 production core has incompatible controls: " + "strict production core has incompatible controls: " + ", ".join(production_mismatches) ) + if ( + core_id == STRICT_ATTENTION_PRODUCTION_CORE_ID + and getattr(core, "backend_id", None) != "flash_attention_4.cute" + ): + raise ValueError("strict CUDA production core must be FlashAttention-4 CuTe") + if core_id == STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID: + if getattr(core, "backend_id", None) != "aiter.rocm.ck_dense_mha": + raise ValueError("strict ROCm production core must be AITER CK dense MHA") + if getattr(core, "split_kv_control", None) != "dense_non_split_api": + raise ValueError("strict ROCm production core must use the non-Split-K CK API") def _resolve_strict_core(cfg: FlashInferPagedAttentionConfig) -> Any: if cfg.deterministic_core is not None: return cfg.deterministic_core + if torch.version.hip is not None: + from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore + + return StrictRocmAiterCKAttentionCore(split_kv=cfg.split_kv) return StrictFlashAttention4Core(split_kv=cfg.split_kv) @@ -1551,12 +1566,17 @@ def _resolve_strict_rope(cfg: FlashInferPagedAttentionConfig) -> Any: if cfg.strict_rope_op is not None: return cfg.strict_rope_op try: - from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import ( + RocmDeterministicRoPEOp, + RoPESM90Op, + ) + if torch.version.hip is not None: + return RocmDeterministicRoPEOp() return RoPESM90Op() except (ImportError, RuntimeError) as exc: raise FlashInferUnavailable( - "strict Attention requires the RL-Kernel WS1 RoPE CUDA operator" + "strict Attention requires the RL-Kernel deterministic RoPE operator" ) from exc @@ -1566,7 +1586,7 @@ def _apply_strict_rope( position_ids: torch.Tensor, theta: float, ) -> torch.Tensor: - """RoPESM90Op accepts shared 1-D positions, so execute one batch row at a time.""" + """Execute one batch row at a time to preserve the strict row schedule.""" if position_ids.shape != (x.size(0), x.size(2)): raise ValueError("strict RoPE position IDs must have shape [B,S]") @@ -1705,6 +1725,7 @@ def _strict_attention_provenance( communication_backend = ( "self_owned_cuda_ag_rs" if communication_id == "cuda_ag_rs" else communication_id ) + expected_communication = "rccl_ag_rs" if torch.version.hip is not None else "cuda_ag_rs" return { "attention_backend": core_provenance["attention_backend"], "requested_backend": "flashinfer_layout_adapter", @@ -1722,12 +1743,15 @@ def _strict_attention_provenance( "strict_schedule": core_provenance["strict_schedule"], "accum_dtype": core_provenance["accum_dtype"], "downcast_at": core_provenance["downcast_at"], - "arithmetic_plan_source": core_provenance.get("fa_api_source", "rlkernel_reference_core"), + "arithmetic_plan_source": core_provenance.get( + "fa_api_source", + core_provenance.get("aiter_api_source", "rlkernel_reference_core"), + ), "arithmetic_semantics_verified": True, "native_attention_arithmetic": core_provenance["native_attention_arithmetic"], "fallback": False, "fallback_reason": None, - "rope_backend": getattr(rope, "backend_id", "rlkernel.cuda.rope_sm90"), + "rope_backend": getattr(rope, "backend_id", "rlkernel.unknown.rope"), "rope_theta": float(cfg.rope.rope_theta), "rotary_dim": cfg.rope.rotary_dim, "rope_fusion": False, @@ -1736,17 +1760,19 @@ def _strict_attention_provenance( "k_cache_rope_state": "post_rope", "batch_invariant_claim": "strict_runtime_verified", "cp_comm_required": cp_required, - "communication_backend": (communication_backend if cp_required else "none"), + "communication_backend": communication_backend if cp_required else "none", + "platform": core_provenance.get("platform", "cuda"), "num_splits": core_provenance.get("num_splits"), + "split_kv_control": core_provenance.get("split_kv_control"), "deterministic_backward": core_provenance.get("deterministic_backward"), "fa_api_source": core_provenance.get("fa_api_source"), "fa_package_version": core_provenance.get("fa_package_version"), + "aiter_api_source": core_provenance.get("aiter_api_source"), + "aiter_source_sha256": core_provenance.get("aiter_source_sha256"), "reference_only": bool(core_provenance.get("reference_only", False)), "production_ready": bool( core_provenance.get("production_ready", False) - and ( - not cp_required or getattr(cfg.cp_communication, "backend_id", None) == "cuda_ag_rs" - ) + and (not cp_required or communication_id == expected_communication) ), } diff --git a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py index 9a764012..928c3033 100644 --- a/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py +++ b/rl_engine/kernels/ops/cuda/rotary_embedding/rope.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Custom CUDA RoPE op for SM90 (GPT-NeoX rotate-half), matching NativeRoPEOp. +"""Deterministic GPU RoPE ops (GPT-NeoX rotate-half), matching NativeRoPEOp. cos/sin are built in fp32 with the exact reference math and passed to a small CUDA kernel (``_C.rope_apply_sm90``) that does the per-position rotation. Backward @@ -77,7 +77,7 @@ def forward(ctx, x: Tensor, positions: Tensor, theta: float) -> Tensor: ctx.save_for_backward(cos, sin) ctx.x_shape = tuple(x.shape) ctx.pos_dim = positions.dim() - out_2d = _C.rope_apply_sm90(x_2d, cos, sin, 1.0) + out_2d = _rope_apply(x_2d, cos, sin, 1.0) return _restore_rope(out_2d, x, positions) @staticmethod @@ -87,7 +87,7 @@ def backward(ctx, grad_out: Tensor): if ctx.needs_input_grad[0]: if ctx.pos_dim == 2 and len(ctx.x_shape) == 4: g_2d = grad_out.permute(1, 0, 2, 3).contiguous().reshape(-1, ctx.x_shape[-1]) - out_2d = _C.rope_apply_sm90(g_2d, cos, sin, -1.0) + out_2d = _rope_apply(g_2d, cos, sin, -1.0) heads, batch, seq, dim = ( ctx.x_shape[1], ctx.x_shape[0], @@ -97,10 +97,16 @@ def backward(ctx, grad_out: Tensor): grad_x = out_2d.reshape(heads, batch, seq, dim).permute(1, 0, 2, 3).contiguous() else: g_2d = grad_out.contiguous().reshape(-1, grad_out.shape[-1]) - grad_x = _C.rope_apply_sm90(g_2d, cos, sin, -1.0).reshape(grad_out.shape) + grad_x = _rope_apply(g_2d, cos, sin, -1.0).reshape(grad_out.shape) return grad_x, None, None +def _rope_apply(x: Tensor, cos: Tensor, sin: Tensor, sin_sign: float) -> Tensor: + if torch.version.hip is not None: + return _C.deterministic_rope_apply_rocm(x, cos, sin, sin_sign) + return _C.rope_apply_sm90(x, cos, sin, sin_sign) + + def _is_hopper(device: torch.device) -> bool: try: return torch.cuda.get_device_capability(device)[0] == 9 @@ -138,3 +144,32 @@ def forward(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) - f"got compute capability {torch.cuda.get_device_capability(x.device)}" ) return _RoPEFunction.apply(x, positions, theta) + + +class RocmDeterministicRoPEOp: + """Precompiled HIP RoPE path shared by ROCm training and rollout.""" + + backend_id = "rlkernel.rocm.deterministic_rope" + op_class = "elementwise" + fallback = False + + def __init__(self) -> None: + if torch.version.hip is None: + raise RuntimeError("RocmDeterministicRoPEOp requires a ROCm PyTorch build") + if not _EXT_AVAILABLE or not hasattr(_C, "deterministic_rope_apply_rocm"): + raise RuntimeError( + "ROCm deterministic RoPE is unavailable; rebuild rl_engine._C for ROCm" + ) + + def __call__(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: + return self.forward(x, positions, theta=theta) + + def forward(self, x: Tensor, positions: Tensor, *, theta: float = 1_000_000.0) -> Tensor: + if not x.is_cuda: + raise RuntimeError("ROCm deterministic RoPE requires a GPU tensor") + if x.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("ROCm deterministic RoPE requires FP16 or BF16") + return _RoPEFunction.apply(x, positions, theta) + + +__all__ = ["RoPESM90Op", "RocmDeterministicRoPEOp"] diff --git a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py index 8f952c12..666743a5 100644 --- a/rl_engine/kernels/ops/pytorch/attention/cp_attention.py +++ b/rl_engine/kernels/ops/pytorch/attention/cp_attention.py @@ -19,6 +19,7 @@ from rl_engine.kernels.attention_contract import ( STRICT_ATTENTION_CORE_ID, + STRICT_ATTENTION_RING_SCHEDULE_ID, STRICT_ATTENTION_SCHEDULE_ID, SplitKVExecutionPlan, SplitKVMode, @@ -139,7 +140,7 @@ def build( left += 1 right -= 1 return cls( - schedule_id="rlkernel.attention.strict_ring_state.v1", + schedule_id=STRICT_ATTENTION_RING_SCHEDULE_ID, total_kv_tokens=total_kv_tokens, cp_world_size=cp_world_size, kv_chunk_size=kv_chunk_size, diff --git a/rl_engine/kernels/ops/pytorch/ffn/ffn.py b/rl_engine/kernels/ops/pytorch/ffn/ffn.py index 41509f13..453878f4 100644 --- a/rl_engine/kernels/ops/pytorch/ffn/ffn.py +++ b/rl_engine/kernels/ops/pytorch/ffn/ffn.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Bias-free gated FFN assembled from deterministic CUDA kernels.""" +"""Bias-free gated FFN assembled from deterministic GPU kernels.""" from __future__ import annotations @@ -126,9 +126,9 @@ def _require_ffn_kernels(*, disable_split_k: bool, packed_gate_up: bool = False) if not _EXT_AVAILABLE or _C is None or missing: suffix = f" Missing symbols: {', '.join(missing)}." if missing else "" needed = ( - "compiled deterministic GEMM and SwiGLU CUDA kernels" + "compiled deterministic GEMM and SwiGLU GPU kernels" if disable_split_k - else "compiled SwiGLU CUDA kernels" + else "compiled SwiGLU GPU kernels" ) raise RuntimeError(f"qwen3_ffn requires the {needed}.{suffix}") @@ -231,7 +231,8 @@ def _validate_ffn_inputs( if tensor.dtype != torch.bfloat16: raise TypeError(f"{name} must have dtype bfloat16, got {tensor.dtype}.") if not tensor.is_cuda: - raise RuntimeError(f"{name} must be on a CUDA device, got '{tensor.device}'.") + # PyTorch exposes AMD GPU tensors through the torch.cuda API too. + raise RuntimeError(f"{name} must be on a CUDA/ROCm GPU device, got '{tensor.device}'.") if tensor.device != rmsnorm_output.device: raise RuntimeError( f"all FFN inputs must be on {rmsnorm_output.device}, " @@ -295,8 +296,14 @@ def forward( tp_world = tp_dist.get_world_size(group=tp_group) if tp_dist is not None else 1 gemm_tokens = rmsnorm_output_2d.size(0) * (tp_world if sequence_parallel else 1) element_size = rmsnorm_output_2d.element_size() + token_hidden_bytes = gemm_tokens * rmsnorm_output_2d.size(1) * element_size + # Sequence-parallel backward reduces the gate and up input-gradient + # lanes together. ``reduce_scatter_many`` packs those lanes along the + # final dimension, so reserve capacity for both lanes in one transport + # call rather than growing the collective (or failing) mid-backward. + reduction_bytes = token_hidden_bytes * (2 if sequence_parallel else 1) min_size_bytes = max( - gemm_tokens * rmsnorm_output_2d.size(1) * element_size, + reduction_bytes, gemm_tokens * gate_weight.size(0) * element_size, gate_weight.numel() * element_size, up_weight.numel() * element_size, @@ -466,28 +473,25 @@ def backward(ctx, grad_output: Tensor): gate_weight, disable_split_k=disable_split_k, ) - if ctx.sequence_parallel: - grad_rmsnorm_from_gate = _reduce_scatter_tokens( - grad_rmsnorm_from_gate, - tp_collective, - ) - elif tp_collective is not None: - grad_rmsnorm_from_gate = _all_reduce_inplace( - grad_rmsnorm_from_gate, - tp_collective, - ) - grad_rmsnorm_from_up = _linear_da( grad_up, up_weight, disable_split_k=disable_split_k, ) if ctx.sequence_parallel: - grad_rmsnorm_from_up = _reduce_scatter_tokens( - grad_rmsnorm_from_up, - tp_collective, + # These are independent reduction lanes. Pack them into one + # ReduceScatter while keeping each lane's balanced rank tree + # separate; adding them before the collective would change the + # floating-point parenthesization and break cross-TP bitwise + # invariance. + grad_rmsnorm_from_gate, grad_rmsnorm_from_up = tp_collective.reduce_scatter_many( + (grad_rmsnorm_from_gate, grad_rmsnorm_from_up) ) elif tp_collective is not None: + grad_rmsnorm_from_gate = _all_reduce_inplace( + grad_rmsnorm_from_gate, + tp_collective, + ) grad_rmsnorm_from_up = _all_reduce_inplace( grad_rmsnorm_from_up, tp_collective, @@ -537,7 +541,8 @@ def qwen3_ffn( unchanged. tp_group: Optional tensor-parallel process group. Gate and Up are column-parallel; Down is row-parallel. Reductions use the - deterministic fixed-tree collectives rather than NCCL. + platform deterministic fixed-tree collectives. On ROCm, RCCL only + transports rank inputs and the reduction tree executes locally. cp_group: Optional context-parallel process group. Each rank owns different token rows and the same local weight shards. Weight gradients AllGather tokens along CP and run the full-token @@ -629,6 +634,11 @@ def prepare_packed_inference( dist = _require_parallel_group(tp_group, "tensor") if dist is None: return 0, 1 + if getattr(torch.version, "hip", None) is not None: + raise RuntimeError( + "packed TP inference requires the native CUDA IPC collective and " + "is not available with the ROCm/RCCL transport" + ) tp_world_size = int(dist.get_world_size(group=tp_group)) if fused_gate_up_weight.size(0) % 2: raise ValueError("fused gate/up weight must contain two equal shards") diff --git a/rl_engine/kernels/ops/rocm/attention/__init__.py b/rl_engine/kernels/ops/rocm/attention/__init__.py index 150c937f..01a10e4f 100644 --- a/rl_engine/kernels/ops/rocm/attention/__init__.py +++ b/rl_engine/kernels/ops/rocm/attention/__init__.py @@ -1,8 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -from .flash_attn import RocmFlashAttentionOp +from .flash_attn import ( + RocmFlashAttentionOp, + StrictRocmAiterCKAttentionCore, + StrictRocmAttentionUnavailable, +) +from .strict_runtime import StrictRocmAttentionResult, StrictRocmAttentionRuntime __all__ = [ "RocmFlashAttentionOp", + "StrictRocmAiterCKAttentionCore", + "StrictRocmAttentionRuntime", + "StrictRocmAttentionResult", + "StrictRocmAttentionUnavailable", ] diff --git a/rl_engine/kernels/ops/rocm/attention/flash_attn.py b/rl_engine/kernels/ops/rocm/attention/flash_attn.py index a9781cfb..65cd94aa 100644 --- a/rl_engine/kernels/ops/rocm/attention/flash_attn.py +++ b/rl_engine/kernels/ops/rocm/attention/flash_attn.py @@ -1,13 +1,374 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +from __future__ import annotations + +import hashlib +import importlib +import inspect +import math import os +from pathlib import Path +from typing import Any, Callable import torch +from torch.autograd import Function +from torch.autograd.function import once_differentiable +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, + SplitKVMode, + SplitKVSpec, +) +from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionCoreResult, + RLKernelDeterministicAttentionCore, +) from rl_engine.utils.logger import logger _MAX_TESTED_ROCM_TRITON_HEAD_DIM = 512 +_AITER_API_SOURCE = "aiter.ops.mha" +_AITER_OP_NAMESPACE = "aiter" + +# AITER wraps its kernels in a JIT loader whose Python signature is +# ``(*args, **kwargs)``, so ``inspect.signature`` cannot see the contract the +# way it can for the FA4 CuTe API. The names live in the registered Torch +# schema instead, and the calls below are positional, so the order is part of +# what has to hold: an upstream insertion would silently reinterpret every +# argument after it. These tuples are the exact positional prefix each call +# site assumes. +_AITER_FWD_POSITIONAL_CONTRACT = ( + "q", + "k", + "v", + "dropout_p", + "softmax_scale", + "is_causal", + "window_size_left", + "window_size_right", + "sink_size", + "return_softmax_lse", + "return_dropout_randval", +) +_AITER_BWD_POSITIONAL_CONTRACT = ( + "dout", + "q", + "k", + "v", + "out", + "softmax_lse", + "dropout_p", + "softmax_scale", + "is_causal", + "window_size_left", + "window_size_right", + "deterministic", +) +# Passed by keyword, so only presence matters. +_AITER_BWD_REQUIRED_KEYWORDS = frozenset({"rng_state"}) + +# Stable dispatch identity for the strict ROCm attention core. Kept at module +# scope so contract-aware dispatch and the Vime adapter name one constant +# instead of duplicating the string. +BACKEND_ID = "aiter.rocm.ck_dense_mha" + + +class StrictRocmAttentionUnavailable(RuntimeError): + """Raised when the exact AITER CK strict contract is unavailable.""" + + +def _aiter_schema_argument_names(op_name: str) -> tuple[str, ...]: + """Return the registered Torch schema argument names for one AITER op.""" + + namespace = getattr(torch.ops, _AITER_OP_NAMESPACE, None) + if namespace is None: + raise StrictRocmAttentionUnavailable( + f"the '{_AITER_OP_NAMESPACE}' Torch operator namespace is not registered" + ) + try: + overload = getattr(namespace, op_name).default + arguments = overload._schema.arguments + except (AttributeError, RuntimeError) as exc: + raise StrictRocmAttentionUnavailable( + f"cannot read the Torch schema for {_AITER_OP_NAMESPACE}::{op_name}" + ) from exc + return tuple(argument.name for argument in arguments) + + +def _validate_aiter_schema( + op_name: str, + positional_contract: tuple[str, ...], + *, + required_keywords: frozenset[str] = frozenset(), +) -> None: + """Fail closed unless AITER still accepts what the call sites pass. + + The strict calls are positional, so a renamed *or reordered* argument + changes their meaning without changing their shape. Checking the ordered + prefix catches both, which name-presence alone would not. + """ + + names = _aiter_schema_argument_names(op_name) + prefix = names[: len(positional_contract)] + if prefix != positional_contract: + raise StrictRocmAttentionUnavailable( + f"AITER {op_name} positional contract changed: strict ROCm Attention " + f"passes {positional_contract} but the schema declares {prefix}" + ) + missing = sorted(required_keywords.difference(names)) + if missing: + raise StrictRocmAttentionUnavailable( + f"AITER {op_name} is missing strict controls: " + ", ".join(missing) + ) + + +def _load_aiter_ck_ops() -> tuple[Callable[..., Any], Callable[..., Any], str]: + try: + module = importlib.import_module(_AITER_API_SOURCE) + mha_fwd = getattr(module, "mha_fwd") + mha_bwd = getattr(module, "mha_bwd") + except (AttributeError, ImportError, OSError, RuntimeError) as exc: + raise StrictRocmAttentionUnavailable( + "strict ROCm Attention requires aiter.ops.mha.mha_fwd and mha_bwd" + ) from exc + if not callable(mha_fwd) or not callable(mha_bwd): + raise StrictRocmAttentionUnavailable("AITER CK MHA entry points are not callable") + _validate_aiter_schema("mha_fwd", _AITER_FWD_POSITIONAL_CONTRACT) + _validate_aiter_schema( + "mha_bwd", + _AITER_BWD_POSITIONAL_CONTRACT, + required_keywords=_AITER_BWD_REQUIRED_KEYWORDS, + ) + module_file = inspect.getsourcefile(module) + if not module_file: + raise StrictRocmAttentionUnavailable("cannot fingerprint the AITER MHA source module") + source_sha256 = hashlib.sha256(Path(module_file).read_bytes()).hexdigest() + return mha_fwd, mha_bwd, source_sha256 + + +class _AiterCKAttentionFn(Function): + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool, + scale: float, + mha_fwd: Callable[..., Any], + mha_bwd: Callable[..., Any], + ) -> tuple[torch.Tensor, torch.Tensor]: + q_fa = q.transpose(1, 2).contiguous() + k_fa = k.transpose(1, 2).contiguous() + v_fa = v.transpose(1, 2).contiguous() + result = mha_fwd( + q_fa, + k_fa, + v_fa, + 0.0, + float(scale), + bool(causal), + -1, + -1, + 0, + True, + False, + ) + if not isinstance(result, (tuple, list)) or len(result) != 4: + raise StrictRocmAttentionUnavailable( + "AITER mha_fwd must return (out, lse, dropout_mask, rng_state)" + ) + out_fa, lse, _dropout_mask, rng_state = result + if not all(isinstance(item, torch.Tensor) for item in (out_fa, lse, rng_state)): + raise StrictRocmAttentionUnavailable("AITER mha_fwd returned non-tensor state") + ctx.save_for_backward(q_fa, k_fa, v_fa, out_fa, lse, rng_state) + ctx.causal = bool(causal) + ctx.scale = float(scale) + ctx.mha_bwd = mha_bwd + ctx.mark_non_differentiable(lse) + return out_fa.transpose(1, 2).contiguous(), lse.contiguous() + + @staticmethod + @once_differentiable + def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): + q_fa, k_fa, v_fa, out_fa, lse, rng_state = ctx.saved_tensors + grad_out_fa = grad_out.transpose(1, 2).contiguous() + result = ctx.mha_bwd( + grad_out_fa, + q_fa, + k_fa, + v_fa, + out_fa, + lse, + 0.0, + ctx.scale, + ctx.causal, + -1, + -1, + True, + rng_state=rng_state, + ) + if not isinstance(result, (tuple, list)) or len(result) < 3: + raise StrictRocmAttentionUnavailable("AITER mha_bwd must return dQ/dK/dV") + dq, dk, dv = result[:3] + return ( + dq.transpose(1, 2).contiguous(), + dk.transpose(1, 2).contiguous(), + dv.transpose(1, 2).contiguous(), + None, + None, + None, + None, + ) + + +class StrictRocmAiterCKAttentionCore: + """Shared ROCm production core using the non-Split-K AITER CK dense MHA.""" + + core_id = STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID + strict_schedule = STRICT_ATTENTION_ROCM_SCHEDULE_ID + backend_id = BACKEND_ID + api_source = _AITER_API_SOURCE + merge_order = "global_block_index" + accum_dtype = "fp32" + downcast_at = "final_write" + fallback = False + native_attention_arithmetic = True + production_ready = True + reference_only = False + num_splits = 1 + split_kv_control = "dense_non_split_api" + deterministic_backward = True + + def __init__( + self, + *, + split_kv: SplitKVSpec | None = None, + _mha_fwd: Callable[..., Any] | None = None, + _mha_bwd: Callable[..., Any] | None = None, + _source_sha256: str | None = None, + ) -> None: + requested = SplitKVSpec.disabled() if split_kv is None else split_kv + if not isinstance(requested, SplitKVSpec): + raise TypeError("split_kv must be a SplitKVSpec") + if requested.mode is not SplitKVMode.DISABLED: + raise ValueError("strict AITER CK Attention requires Split-KV to be disabled") + if (_mha_fwd is None) != (_mha_bwd is None): + raise ValueError("test injection requires both AITER forward and backward callables") + if _mha_fwd is None: + mha_fwd, mha_bwd, source_sha256 = _load_aiter_ck_ops() + else: + assert _mha_bwd is not None + mha_fwd = _mha_fwd + mha_bwd = _mha_bwd + source_sha256 = "test-double" if _source_sha256 is None else _source_sha256 + if not callable(mha_fwd) or not callable(mha_bwd): + raise StrictRocmAttentionUnavailable("AITER CK MHA entry points are not callable") + self.split_kv = requested + self.source_sha256 = source_sha256 + self._mha_fwd = mha_fwd + self._mha_bwd = mha_bwd + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: float | None = None, + key_padding_mask: torch.Tensor | None = None, + query_position_ids: torch.Tensor | None = None, + key_position_ids: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, + ) -> DeterministicAttentionCoreResult: + self._validate_inputs(q, k, v, key_padding_mask) + RLKernelDeterministicAttentionCore._validate_positions( + q, + k, + causal=causal, + query_position_ids=query_position_ids, + key_position_ids=key_position_ids, + ) + resolved_dtype = q.dtype if output_dtype is None else output_dtype + if resolved_dtype != q.dtype: + raise ValueError("strict Attention output_dtype must match the Q/K/V input dtype") + resolved_scale = 1.0 / math.sqrt(q.size(-1)) if scale is None else float(scale) + out, lse = _AiterCKAttentionFn.apply( + q, + k, + v, + bool(causal), + resolved_scale, + self._mha_fwd, + self._mha_bwd, + ) + expected_lse_shape = (q.size(0), q.size(1), q.size(2)) + if out.shape != q.shape or out.dtype != resolved_dtype: + raise StrictRocmAttentionUnavailable("AITER CK output shape/dtype changed") + if tuple(lse.shape) != expected_lse_shape or lse.dtype != torch.float32: + raise StrictRocmAttentionUnavailable("AITER CK must export [B,H,Sq] FP32 LSE") + device_properties = torch.cuda.get_device_properties(q.device) + return DeterministicAttentionCoreResult( + out=out, + lse=lse, + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "attention_backend": self.backend_id, + "platform": "rocm", + "torch_version": torch.__version__, + "rocm_version": torch.version.hip, + "gpu_name": device_properties.name, + "gpu_arch": getattr(device_properties, "gcnArchName", "unknown"), + "aiter_api_source": self.api_source, + "aiter_source_sha256": self.source_sha256, + "num_splits": self.num_splits, + "split_kv_control": self.split_kv_control, + "deterministic_backward": self.deterministic_backward, + "dropout_p": 0.0, + "split_kv": self.split_kv.resolve(k.size(2), backend=self.backend_id).to_dict(), + "merge_order": self.merge_order, + "accum_dtype": self.accum_dtype, + "downcast_at": self.downcast_at, + "fallback": self.fallback, + "fallback_reason": None, + "native_attention_arithmetic": self.native_attention_arithmetic, + "production_ready": self.production_ready, + "reference_only": self.reference_only, + }, + ) + + @staticmethod + def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + key_padding_mask: torch.Tensor | None, + ) -> None: + if torch.version.hip is None: + raise StrictRocmAttentionUnavailable("strict AITER CK core requires ROCm PyTorch") + if key_padding_mask is not None: + raise ValueError("strict AITER CK core materializes each unpadded logical row") + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError("q/k/v must be 4-D [B,H,S,D]") + if q.size(0) != 1 or k.size(0) != 1 or v.size(0) != 1: + raise ValueError("strict AITER CK core executes one logical batch row at a time") + if k.shape != v.shape or q.size(-1) != k.size(-1): + raise ValueError("k/v shapes and q/k/v head dimensions must match") + if q.size(1) % k.size(1) != 0: + raise ValueError("Q heads must be divisible by KV heads for GQA") + if q.size(-1) > 256 or q.size(-1) % 8: + raise ValueError("AITER CK requires head_dim <= 256 and divisible by 8") + if q.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict AITER CK core supports FP16/BF16 only") + if k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("q/k/v must share one dtype") + if not (q.is_cuda and k.is_cuda and v.is_cuda): + raise ValueError("strict AITER CK core requires ROCm GPU tensors") + if not (q.device == k.device == v.device): + raise ValueError("q/k/v must be on one ROCm device") def _select_flash_attn_backend() -> str: diff --git a/rl_engine/kernels/ops/rocm/attention/strict_runtime.py b/rl_engine/kernels/ops/rocm/attention/strict_runtime.py new file mode 100644 index 00000000..49a5b9a3 --- /dev/null +++ b/rl_engine/kernels/ops/rocm/attention/strict_runtime.py @@ -0,0 +1,518 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Framework-neutral strict ROCm Attention runtime. + +Composes the production AITER/CK core with the self-owned RCCL AG/RS +transport, mirroring :class:`StrictCUDAAttentionRuntime` so both platforms +present one runtime shape to framework integrations. Before this existed the +CP schedule had no home on ROCm: the core is single-rank arithmetic, the Vime +provider fails closed at ``CP > 1``, and the only working AG/core/RS sequence +lived in the benchmark script. + +Two things differ from the CUDA runtime and both are load-bearing: + +* The core is launched once per ``(batch row, KV group)`` rather than once per + sequence. AITER/CK's reduction order depends on how many heads shared the + launch, so a head shard computed under TP=N is otherwise not bit-identical + to the same shard under a different TP degree. The CUDA FA4 core has no such + dependence and runs one launch per sequence. +* RCCL moves tensors but never reduces them. The cross-rank ``(out, lse)`` + combine order comes from the fixed balanced rank tree in the shared + ``RCCLDeterministicCollective``, not from RCCL's own algorithm selection. + That is the collective the CUDA runtime also resolves through + ``collective_for_group``, so both platforms run one reduction order from + one implementation. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, + AttentionContract, +) +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionParallelSpec, + RCCLAGRSAttentionCPCommunication, +) +from rl_engine.kernels.ops.cuda.attention.strict_runtime import StrictCUDAAttentionRuntime +from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore + +# The sequence reorder and the position validation are platform-neutral tensor +# bookkeeping. They are bound from the CUDA runtime rather than reimplemented +# so the two runtimes cannot drift into two different global orderings. +_sort_by_position = StrictCUDAAttentionRuntime._sort_by_position +_gather_sequence = StrictCUDAAttentionRuntime._gather_sequence +_validate_local_positions = StrictCUDAAttentionRuntime._validate_local_positions +_validate_global_positions = StrictCUDAAttentionRuntime._validate_global_positions + + +@dataclass(frozen=True) +class StrictRocmAttentionResult: + out: torch.Tensor + lse: torch.Tensor + provenance: dict[str, Any] + + +class StrictRocmAttentionRuntime: + """Run one AITER/CK arithmetic identity at CP=1 or through RCCL AG/RS.""" + + backend_id = "rlkernel.rocm.attention.aiter_ck_ag_rs.v1" + core_id = STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID + strict_schedule = STRICT_ATTENTION_ROCM_SCHEDULE_ID + communication_backend_id = "rccl_ag_rs" + # Communication and compute are deliberately decoupled; see + # ``AttentionCPCommunicationPlan.validate``. + supports_async_overlap = False + supports_compute_communication_fusion = False + + def __init__( + self, + *, + process_group: Any = None, + core: Any | None = None, + communication: Any | None = None, + ) -> None: + self._core = StrictRocmAiterCKAttentionCore() if core is None else core + self._communication = ( + RCCLAGRSAttentionCPCommunication(process_group=process_group) + if communication is None + else communication + ) + if getattr(self._core, "core_id", None) != self.core_id: + raise RuntimeError( + "strict ROCm Attention runtime requires the AITER/CK production core" + ) + if getattr(self._core, "strict_schedule", None) != self.strict_schedule: + raise RuntimeError("strict ROCm Attention runtime requires the AITER/CK fixed schedule") + self.communication_executed = False + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + contract: AttentionContract, + causal: bool, + scale: float | None, + cp_world_size: int, + query_position_ids: torch.Tensor, + key_position_ids: torch.Tensor, + positions_are_sorted: bool = False, + ) -> StrictRocmAttentionResult: + self._require_rocm(q) + if cp_world_size != contract.sharding.cp_world_size: + raise RuntimeError("runtime CP world size does not match AttentionContract") + _validate_local_positions(q, k, query_position_ids, key_position_ids) + + plan = None + if cp_world_size == 1: + global_q, global_k, global_v = q, k, v + global_q_positions = query_position_ids + global_k_positions = key_position_ids + communication_backend = "none" + self.communication_executed = False + else: + plan = self._communication_plan(contract, q.size(2), k.size(2)) + global_q = self._communication.all_gather_query(q, plan) + global_k, global_v = self._communication.all_gather_kv(k, v, plan) + global_q_positions, global_k_positions = self._communication.all_gather_position_ids( + query_position_ids, + key_position_ids, + plan, + ) + communication_backend = self.communication_backend_id + self.communication_executed = True + + if positions_are_sorted: + if cp_world_size != 1: + raise RuntimeError("pre-sorted Attention positions are supported only at CP=1") + q_sorted, k_sorted, v_sorted = global_q, global_k, global_v + q_positions_sorted, k_positions_sorted = global_q_positions, global_k_positions + q_sort = None + else: + q_sorted, q_positions_sorted, q_sort = _sort_by_position(global_q, global_q_positions) + k_sorted, k_positions_sorted, k_sort = _sort_by_position(global_k, global_k_positions) + v_sorted = _gather_sequence(global_v, k_sort) + _validate_global_positions(q_positions_sorted, k_positions_sorted, causal) + + out_sorted, lse_sorted, core_provenance, launches = self._run_core( + q_sorted, + k_sorted, + v_sorted, + causal=causal, + scale=scale, + query_position_ids=q_positions_sorted, + key_position_ids=k_positions_sorted, + output_dtype=q.dtype, + ) + + if cp_world_size > 1: + if q_sort is None: + raise RuntimeError("CP Attention requires a framework position reorder") + inverse_q_sort = torch.argsort(q_sort, dim=1) + out_rank_packed = _gather_sequence(out_sorted, inverse_q_sort) + lse_rank_packed = _gather_sequence(lse_sorted, inverse_q_sort) + shard = self._communication.reduce_scatter_strict_result( + out_rank_packed, + lse_rank_packed, + plan, + ) + out, lse = shard.out, shard.lse + else: + out, lse = out_sorted, lse_sorted + + backend = ( + core_provenance.get("attention_backend") + or core_provenance.get("actual_backend") + or getattr(self._core, "backend_id", None) + ) + return StrictRocmAttentionResult( + out=out, + lse=lse, + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": self.backend_id, + "communication_backend": communication_backend, + "communication_executed": self.communication_executed, + "native_attention_arithmetic": True, + "production_ready": True, + "fallback": False, + "fallback_reason": None, + "reference_only": False, + "split_kv": "disabled", + "framework_position_reorder": True, + # Unlike the CUDA runtime's single full-sequence launch, the + # query schedule here is one launch per (batch row, KV group). + "query_schedule": "one_batch_row_one_kv_group", + "backward_schedule": "aiter_ck_deterministic_per_kv_group", + "launch_granularity": "one_batch_row_one_kv_group", + "tp_degree_invariant": True, + "invariance_mechanism": "one_kv_group_per_launch", + "core_row_count": q_sorted.size(0) * q_sorted.size(2), + "core_launch_count": launches, + "core_batch_size": q_sorted.size(0), + "core_query_length": q_sorted.size(2), + "core_actual_backends": [] if backend is None else [str(backend)], + "core": core_provenance, + }, + ) + + def forward_paged_with_lse( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + *, + page_table: torch.Tensor, + seqused_k: torch.Tensor, + max_seqlen_k: int, + scale: float | None, + out: torch.Tensor | None = None, + ) -> StrictRocmAttentionResult: + """Run strict decode Attention over a paged KV cache. + + AITER exposes no paged entry point that this contract can use. Every + ``paged_attention_*`` kernel partitions KV and reduces the partials + (``partition_size``, ``exp_sums``/``max_logits``/``tmp_out``), so the + partition count moves with the cached length; AITER's + ``flash_attn_varlen_func`` takes a ``block_table`` but has no + ``num_splits`` knob to pin, unlike CUDA's FA4. Either way the strict + contract could not prove Split-KV disabled. + + So the pages are gathered into logical KV order and handed to the same + dense core the prefill path uses, at the same one-launch-per + ``(batch row, KV group)`` granularity. The arithmetic is then identical + to a CP=1 prefill over the same logical sequence, which is what makes + decode replay comparable against it. The cost is materializing the + cached KV; a native paged kernel would avoid that, and can replace this + once AITER can pin its split count. + """ + + self._require_rocm(q) + self._validate_paged_inputs( + q, + k_cache, + v_cache, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_k=max_seqlen_k, + ) + if out is not None: + if out.shape != q.shape: + raise ValueError("paged Attention out must have the same shape as q") + if out.dtype != q.dtype or out.device != q.device: + raise ValueError("paged Attention out must match the Q dtype and device") + if not out.is_contiguous(): + raise ValueError("paged Attention out must be contiguous") + + row_outs: list[torch.Tensor] = [] + row_lses: list[torch.Tensor] = [] + core_provenance: dict[str, Any] | None = None + launches = 0 + for row in range(q.size(0)): + cached_length = int(seqused_k[row].item()) + if cached_length <= 0 or cached_length > max_seqlen_k: + raise ValueError( + "seqused_k entries must be positive and within max_seqlen_k; " + f"row {row} requested {cached_length}" + ) + k_row, v_row = self._gather_paged_row( + k_cache, + v_cache, + page_table[row], + cached_length, + ) + # Decode attends over the whole cached prefix, so the mask is not + # causal within this launch. The logical positions are still passed + # for provenance-grade auditing of what each launch consumed. + key_positions = torch.arange( + cached_length, + dtype=torch.int64, + device=q.device, + ).unsqueeze(0) + query_positions = key_positions[:, -q.size(2) :] + row_out, row_lse, row_provenance, row_launches = self._run_core( + q[row : row + 1], + k_row, + v_row, + causal=False, + scale=scale, + query_position_ids=query_positions, + key_position_ids=key_positions, + output_dtype=q.dtype, + ) + row_outs.append(row_out) + row_lses.append(row_lse) + launches += row_launches + if core_provenance is None: + core_provenance = row_provenance + + if core_provenance is None: + raise RuntimeError("strict ROCm paged Attention executed no core launch") + + result_out = torch.cat(row_outs, dim=0) + result_lse = torch.cat(row_lses, dim=0) + if out is not None: + out.copy_(result_out) + result_out = out + self.communication_executed = False + + backend = ( + core_provenance.get("attention_backend") + or core_provenance.get("actual_backend") + or getattr(self._core, "backend_id", None) + ) + return StrictRocmAttentionResult( + out=result_out, + lse=result_lse, + provenance={ + "strict_core_id": self.core_id, + "strict_schedule": self.strict_schedule, + "actual_backend": self.backend_id, + "communication_backend": "none", + "communication_executed": False, + "native_attention_arithmetic": True, + "production_ready": True, + "fallback": False, + "fallback_reason": None, + "reference_only": False, + "split_kv": "disabled", + "query_schedule": "paged_single_query_batch", + # The dense core runs; the pages are gathered first. Recorded so + # a reader never mistakes this for a native paged kernel. + "paged_execution": "logical_kv_gather_then_dense_core", + "paged_kernel": "none", + "launch_granularity": "one_batch_row_one_kv_group", + "tp_degree_invariant": True, + "invariance_mechanism": "one_kv_group_per_launch", + "core_row_count": q.size(0) * q.size(2), + "core_launch_count": launches, + "core_batch_size": q.size(0), + "core_query_length": q.size(2), + "core_actual_backends": [] if backend is None else [str(backend)], + "core": core_provenance, + }, + ) + + @staticmethod + def _gather_paged_row( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_row: torch.Tensor, + cached_length: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Materialize one row's cached KV in logical order as ``[1, H, S, D]``. + + Physical page order never reaches the core: the pages are read through + the page table, so the launch sees the same logical sequence a prefill + over the same tokens would have seen. + """ + + page_size = k_cache.size(1) + page_count = (cached_length + page_size - 1) // page_size + if page_count > page_row.numel(): + raise ValueError("page_table row is shorter than the cached length requires") + pages = page_row[:page_count].to(dtype=torch.int64) + if int(pages.min().item()) < 0 or int(pages.max().item()) >= k_cache.size(0): + raise ValueError("page_table entries are outside the KV cache") + + def _gather(cache: torch.Tensor) -> torch.Tensor: + selected = cache.index_select(0, pages) + flat = selected.reshape(page_count * page_size, cache.size(2), cache.size(3)) + return flat[:cached_length].permute(1, 0, 2).unsqueeze(0).contiguous() + + return _gather(k_cache), _gather(v_cache) + + @staticmethod + def _validate_paged_inputs( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + *, + page_table: torch.Tensor, + seqused_k: torch.Tensor, + max_seqlen_k: int, + ) -> None: + if q.ndim != 4: + raise ValueError("paged q must use [B, H, S, D]") + if k_cache.ndim != 4 or v_cache.shape != k_cache.shape: + raise ValueError("paged k/v must use [pages, page_size, H, D]") + if q.size(1) % k_cache.size(2) != 0 or q.size(3) != k_cache.size(3): + raise ValueError("paged q/k head counts or head dimensions are incompatible") + if q.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict paged Attention supports FP16/BF16 only") + if k_cache.dtype != q.dtype or v_cache.dtype != q.dtype: + raise ValueError("paged q/k/v must share one dtype") + if not (q.device == k_cache.device == v_cache.device): + raise ValueError("paged q/k/v must be on one ROCm device") + if page_table.ndim != 2 or page_table.size(0) != q.size(0): + raise ValueError("page_table must be 2-D with one row per query") + if page_table.dtype not in (torch.int32, torch.int64): + raise ValueError("page_table must be an integer tensor") + if seqused_k.shape != (q.size(0),): + raise ValueError("seqused_k must carry one cached length per query") + if seqused_k.dtype not in (torch.int32, torch.int64): + raise ValueError("seqused_k must be an integer tensor") + if page_table.device != q.device or seqused_k.device != q.device: + raise ValueError("paged Attention metadata must be on the Q device") + if max_seqlen_k <= 0 or max_seqlen_k > page_table.size(1) * k_cache.size(1): + raise ValueError("max_seqlen_k exceeds the page table capacity") + + def _run_core( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool, + scale: float | None, + query_position_ids: torch.Tensor, + key_position_ids: torch.Tensor, + output_dtype: torch.dtype, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any], int]: + """Launch the core once per ``(batch row, KV group)`` and concatenate. + + Every launch therefore sees exactly one KV group and its Q heads, so + the result does not depend on the TP degree that produced the shard. + """ + + local_kv_heads = k.size(1) + if local_kv_heads <= 0 or q.size(1) % local_kv_heads: + raise RuntimeError( + f"local Q heads={q.size(1)} must be divisible by local KV heads={local_kv_heads}" + ) + group_size = q.size(1) // local_kv_heads + + row_outs: list[torch.Tensor] = [] + row_lses: list[torch.Tensor] = [] + core_provenance: dict[str, Any] | None = None + launches = 0 + for row in range(q.size(0)): + row_query_positions = query_position_ids[row : row + 1] + row_key_positions = key_position_ids[row : row + 1] + group_outs: list[torch.Tensor] = [] + group_lses: list[torch.Tensor] = [] + for group in range(local_kv_heads): + q_lo, q_hi = group * group_size, (group + 1) * group_size + result = self._core.forward_with_lse( + q[row : row + 1, q_lo:q_hi], + k[row : row + 1, group : group + 1], + v[row : row + 1, group : group + 1], + causal=causal, + scale=scale, + key_padding_mask=None, + query_position_ids=row_query_positions if causal else None, + key_position_ids=row_key_positions if causal else None, + output_dtype=output_dtype, + ) + group_outs.append(result.out) + group_lses.append(result.lse) + launches += 1 + if core_provenance is None: + core_provenance = dict(result.provenance) + row_outs.append(torch.cat(group_outs, dim=1)) + row_lses.append(torch.cat(group_lses, dim=1)) + + if core_provenance is None: + raise RuntimeError("strict ROCm Attention runtime executed no core launch") + return ( + torch.cat(row_outs, dim=0), + torch.cat(row_lses, dim=0), + core_provenance, + launches, + ) + + @staticmethod + def _require_rocm(tensor: torch.Tensor) -> None: + if tensor.device.type != "cuda" or torch.version.hip is None: + raise RuntimeError("strict ROCm Attention requires ROCm GPU tensors") + + @staticmethod + def _communication_plan( + contract: AttentionContract, + local_q_tokens: int, + local_kv_tokens: int, + ) -> AttentionCPCommunicationPlan: + sharding = contract.sharding + parallel = AttentionParallelSpec( + tp_world_size=sharding.tp_world_size, + tp_rank=sharding.tp_rank, + cp_world_size=sharding.cp_world_size, + cp_rank=sharding.cp_rank, + ) + query_ranges = tuple( + (rank * local_q_tokens, (rank + 1) * local_q_tokens) + for rank in range(sharding.cp_world_size) + ) + blocks = tuple( + AttentionCPBlockMetadata( + global_block_index=rank, + kv_block_start=rank * local_kv_tokens, + kv_block_end=(rank + 1) * local_kv_tokens, + owner_cp_rank=rank, + owner_tp_rank=sharding.tp_rank, + ) + for rank in range(sharding.cp_world_size) + ) + return AttentionCPCommunicationPlan( + parallel=parallel, + backend="rccl_ag_rs", + status="implemented", + expected_blocks=blocks, + expected_kv_token_range=(0, local_kv_tokens * sharding.cp_world_size), + query_token_ranges=query_ranges, + ) + + +__all__ = ["StrictRocmAttentionResult", "StrictRocmAttentionRuntime"] diff --git a/rl_engine/kernels/ops/triton/attention/__init__.py b/rl_engine/kernels/ops/triton/attention/__init__.py index 220b6c95..7df3d2ba 100644 --- a/rl_engine/kernels/ops/triton/attention/__init__.py +++ b/rl_engine/kernels/ops/triton/attention/__init__.py @@ -1,6 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors +from rl_engine.kernels.ops.triton.attention.deterministic_attn import ( + BITWISE_LIBM_PARITY, + TritonDeterministicAttentionOp, + triton_deterministic_attention, + triton_deterministic_attention_backward, + triton_deterministic_attention_forward, + triton_deterministic_attention_fp32, + triton_deterministic_attention_with_lse, +) from rl_engine.kernels.ops.triton.attention.standard_attn import ( TritonBatchInvariantAttentionOp, triton_batch_invariant_attention, @@ -8,7 +17,14 @@ ) __all__ = [ + "BITWISE_LIBM_PARITY", "TritonBatchInvariantAttentionOp", + "TritonDeterministicAttentionOp", "triton_batch_invariant_attention", "triton_batch_invariant_attention_with_lse", + "triton_deterministic_attention", + "triton_deterministic_attention_backward", + "triton_deterministic_attention_forward", + "triton_deterministic_attention_fp32", + "triton_deterministic_attention_with_lse", ] diff --git a/rl_engine/kernels/ops/triton/attention/deterministic_attn.py b/rl_engine/kernels/ops/triton/attention/deterministic_attn.py new file mode 100644 index 00000000..7914342a --- /dev/null +++ b/rl_engine/kernels/ops/triton/attention/deterministic_attn.py @@ -0,0 +1,878 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Triton port of the deterministic standard-softmax attention core (issue #147). + +This is a *bitwise* re-implementation of ``csrc/cuda/attention/deterministic_attention.cu`` +(exposed as ``_C.deterministic_attention_forward`` / ``_C.deterministic_attention_backward``). +Every reduction here reproduces the C++ kernel's floating-point order exactly, so the +Triton path and the native path return bit-identical ``out``/``lse``/``dQ``/``dK``/``dV`` +for the same inputs on the same device. + +The pipeline mirrors the native one 1:1: + + forward : QK -> masked softmax+LSE -> PV + backward: dP -> softmax backward -> dQ -> dK -> dV + +The three arithmetic contracts that have to be honoured for bitwise parity are: + +1. **Dot products are sequential FMA chains.** The C++ kernels accumulate + ``acc += (float)a[i] * (float)b[i]`` over ascending ``i`` in a single thread, which + hipcc/nvcc contract into a chain of FMAs. Every reduction below loops over the + contraction index one element at a time and uses :func:`tl.fma`, so no vector + tree reduction is ever introduced. That is why the contraction index is the *loop* + and the head dim / output tile is the *vector*: the opposite (and much faster) + arrangement would reassociate the sum. +2. **Row softmax uses the 256-lane partial + binary-tree layout.** The C++ softmax + assigns key ``k`` to thread ``k % 256``, sums each lane's keys in ascending order, + and then folds the 256 partials with ``stride = 128, 64, ... 1``. :func:`_tree_sum_256` + reproduces that fold exactly by repeatedly reshaping to ``(2, n)`` and summing axis 0. +3. **Transcendentals reproduce the vendor libm.** Every Triton exp/log intrinsic lowers + to a bare hardware ``v_exp_f32``/``v_log_f32``, which is ~1 ULP away from the + ``expf``/``logf`` the C++ kernel calls. :func:`_expf` and :func:`_logf` below + re-emit the vendor argument reduction instruction for instruction instead. + +Performance note: like the native reference core this materialises the full FP32 +``[B, Hq, Sq, Skv]`` score matrix and runs scalar-order reductions. It is a +correctness/parity core, not a FlashAttention replacement. +""" + +from __future__ import annotations + +import math +from typing import Optional + +import torch +import triton +import triton.language as tl +from torch.autograd import Function +from torch.autograd.function import once_differentiable + +_HEAD_DIM = 128 + +_IS_ROCM = torch.version.hip is not None + +# --------------------------------------------------------------------------- +# Vendor-exact expf / logf +# --------------------------------------------------------------------------- +# The softmax is the only place this core evaluates a transcendental, and it is +# also the only place where "write the obvious Triton code" is not enough: +# ``tl.exp`` / ``tl.math.exp`` / ``libdevice.exp`` all lower to ``llvm.exp.f32``, +# which the AMDGPU backend expands to a bare ``v_exp_f32``. The HIP ``expf`` the +# C++ kernel calls does a two-term argument reduction around that same hardware +# ``v_exp_f32``, so the two differ by ~1 ULP on most inputs. +# +# The sequences below reproduce, instruction for instruction, what hipcc emits +# for ``expf`` / ``logf`` on gfx9. They are verified bitwise against the vendor +# result over 4M+ random and edge-case inputs, including subnormals, +/-inf and NaN. +# +# ``_fp32_barrier`` is load-bearing: without it LLVM folds ``x * L2E_HI`` and the +# following subtract back into a single FMA, which silently changes the reduced +# argument. Inline asm is opaque to that folding. + +if _IS_ROCM: + from triton.language.extra.hip import libdevice as _ocml + + @triton.jit + def _fp32_barrier(x): + """Opaque move: stops LLVM re-associating across this point.""" + return tl.inline_asm_elementwise( + "v_mov_b32 $0, $1", "=v,v", [x], dtype=tl.float32, is_pure=True, pack=1 + ) + + @triton.jit + def _expf(x): + """Bitwise-exact HIP ``expf`` for FP32.""" + t = _fp32_barrier(x * 1.4426950216293335) + err = tl.fma(x, 1.4426950216293335, -t) + n = _ocml.rint(t) + err = tl.fma(x, 1.925962855864327e-08, err) + r = _fp32_barrier(_fp32_barrier(t - n) + err) + y = _ocml.ldexp(_ocml.exp2(r), n.to(tl.int32)) + # Written as "constant compared against x" so an unordered (NaN) compare + # falls through to the NaN result, matching v_cmp_ngt / v_cmp_nlt. + y = tl.where(-103.2789306640625 > x, 0.0, y) + return tl.where(88.72283935546875 < x, float("inf"), y) + + @triton.jit + def _logf(x): + """Bitwise-exact HIP ``logf`` for FP32.""" + small = x < 1.1754943508222875e-38 + log2_x = _ocml.log2(_ocml.ldexp(x, tl.where(small, 32, 0))) + t = _fp32_barrier(log2_x * 0.6931471228599548) + err = tl.fma(log2_x, 0.6931471228599548, -t) + err = tl.fma(log2_x, 5.769998878690785e-08, err) + r = _fp32_barrier(t + err) + r = tl.where(tl.abs(log2_x) < float("inf"), r, log2_x) + return r - tl.where(small, 22.180709838867188, 0.0) + +else: + + @triton.jit + def _expf(x): + return tl.math.exp(x) + + @triton.jit + def _logf(x): + return tl.math.log(x) + + +#: True when :func:`_expf` / :func:`_logf` reproduce this platform's libm bitwise. +#: The nvcc ``expf``/``logf`` sequences have not been ported, so on CUDA the ops +#: below refuse to run unless the caller opts out explicitly. +BITWISE_LIBM_PARITY = _IS_ROCM + +# Mirrors kSoftmaxThreads in csrc/cuda/attention/deterministic_attention.cu. The +# value is part of the arithmetic contract, not a tuning knob: changing it changes +# which keys land in which partial sum and therefore changes the result bitwise. +_SOFTMAX_LANES = 256 + +# Tile shapes. These only affect scheduling, never the reduction order, because +# every reduction is a per-output-element sequential loop. +_QK_BLOCK_Q = 16 +_QK_BLOCK_K = 64 + + +@triton.jit +def _tree_sum_256(vals): + """Fold 256 partials the way the C++ shared-memory tree reduction does. + + The C++ loop is ``for (stride = 128; stride > 0; stride >>= 1) s[i] += s[i + stride]``. + Reshaping to ``(2, n)`` and summing axis 0 is that same pairing: row-major + ``reshape(2, n)[0] == vals[:n]`` and ``[1] == vals[n:]``, so each step is + ``vals[i] + vals[i + n]``. The steps are written out because a Triton loop + cannot carry a value whose shape changes. + """ + total = tl.sum(tl.reshape(vals, (2, 128)), axis=0) + total = tl.sum(tl.reshape(total, (2, 64)), axis=0) + total = tl.sum(tl.reshape(total, (2, 32)), axis=0) + total = tl.sum(tl.reshape(total, (2, 16)), axis=0) + total = tl.sum(tl.reshape(total, (2, 8)), axis=0) + total = tl.sum(tl.reshape(total, (2, 4)), axis=0) + total = tl.sum(tl.reshape(total, (2, 2)), axis=0) + total = tl.sum(tl.reshape(total, (2, 1)), axis=0) + return tl.sum(total, axis=0) + + +# --------------------------------------------------------------------------- +# Forward +# --------------------------------------------------------------------------- + + +@triton.jit +def _qk_kernel( + q_ptr, + k_ptr, + scores_ptr, + scale, + Hq, + Hkv, + Sq, + Skv, + D: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """scores[b, hq, q, k] = scale * sum_{d ascending} Q[b,hq,q,d] * K[b,kv,k,d].""" + pid_k = tl.program_id(0) + pid_q = tl.program_id(1) + bh = tl.program_id(2) + b = bh // Hq + hq = bh % Hq + kv_head = hq // (Hq // Hkv) + + offs_q = pid_q * BLOCK_Q + tl.arange(0, BLOCK_Q) + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + q_in = offs_q < Sq + k_in = offs_k < Skv + + q_rows = q_ptr + (b * Hq + hq).to(tl.int64) * Sq * D + offs_q.to(tl.int64) * D + k_rows = k_ptr + (b * Hkv + kv_head).to(tl.int64) * Skv * D + offs_k.to(tl.int64) * D + + acc = tl.zeros((BLOCK_Q, BLOCK_K), dtype=tl.float32) + for d in range(0, D): + qv = tl.load(q_rows + d, mask=q_in, other=0.0).to(tl.float32) + kv = tl.load(k_rows + d, mask=k_in, other=0.0).to(tl.float32) + acc = tl.fma(qv[:, None], kv[None, :], acc) + + dst = ( + scores_ptr + + (b * Hq + hq).to(tl.int64) * Sq * Skv + + offs_q.to(tl.int64)[:, None] * Skv + + offs_k[None, :] + ) + tl.store(dst, scale * acc, mask=q_in[:, None] & k_in[None, :]) + + +@triton.jit +def _masked_softmax_lse_kernel( + scores_ptr, + lse_ptr, + mask_ptr, + Hq, + Sq, + Skv, + CAUSAL: tl.constexpr, + HAS_MASK: tl.constexpr, + LANES: tl.constexpr, +): + """Mask, softmax and LSE one ``(b, hq, q)`` row in place, C++ reduction order.""" + row = tl.program_id(0) + b = row // (Hq * Sq) + q = row % Sq + row_base = scores_ptr + row.to(tl.int64) * Skv + + if CAUSAL: + causal_limit = Skv - Sq + q + 1 + else: + causal_limit = Skv + + lane = tl.arange(0, LANES) + neg_inf = float("-inf") + minus_inf_vec = tl.full((LANES,), neg_inf, tl.float32) + + # Phase 1: write -inf over masked entries and take the row max. Max is + # associative and commutative in IEEE-754, so the tree shape is irrelevant here. + lane_max = minus_inf_vec + for start in range(0, Skv, LANES): + cols = start + lane + in_range = cols < Skv + valid = in_range & (cols < causal_limit) + if HAS_MASK: + keep = tl.load(mask_ptr + b.to(tl.int64) * Skv + cols, mask=in_range, other=0) + valid = valid & (keep != 0) + scores = tl.load(row_base + cols, mask=in_range, other=neg_inf) + tl.store(row_base + cols, minus_inf_vec, mask=in_range & ~valid) + lane_max = tl.maximum(lane_max, tl.where(valid, scores, neg_inf)) + row_max = tl.max(lane_max, axis=0) + + # Phase 2: exponentiate in place. Lane ``t`` sums keys t, t+LANES, ... ascending, + # exactly like thread ``t`` in the C++ kernel; masked lanes contribute +0.0, which + # is bitwise neutral for this non-negative sum. + lane_sum = tl.zeros((LANES,), dtype=tl.float32) + for start in range(0, Skv, LANES): + cols = start + lane + in_range = cols < Skv + valid = in_range & (cols < causal_limit) + if HAS_MASK: + keep = tl.load(mask_ptr + b.to(tl.int64) * Skv + cols, mask=in_range, other=0) + valid = valid & (keep != 0) + scores = tl.load(row_base + cols, mask=in_range, other=neg_inf) + probs = tl.where(valid, _expf(scores - row_max), 0.0) + tl.store(row_base + cols, probs, mask=in_range) + lane_sum += probs + row_sum = _tree_sum_256(lane_sum) + + # Phase 3: normalise and emit the LSE. A fully masked row already holds zeros + # from phase 2, so dividing it by 1.0 reproduces the C++ zero-fill branch. + is_empty = row_sum == 0.0 + denom = tl.where(is_empty, 1.0, row_sum) + for start in range(0, Skv, LANES): + cols = start + lane + in_range = cols < Skv + probs = tl.load(row_base + cols, mask=in_range, other=0.0) + tl.store(row_base + cols, probs / denom, mask=in_range) + + lse_val = tl.where(is_empty, neg_inf, row_max + _logf(row_sum)) + tl.store(lse_ptr + row, lse_val) + + +@triton.jit +def _pv_kernel( + p_ptr, + v_ptr, + out_ptr, + Hq, + Hkv, + Sq, + Skv, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """out[b, hq, q, d] = sum_{k ascending} P[b,hq,q,k] * V[b,kv,k,d].""" + row = tl.program_id(0) + bh = row // Sq + b = bh // Hq + hq = bh % Hq + kv_head = hq // (Hq // Hkv) + + offs_d = tl.arange(0, BLOCK_D) + d_in = offs_d < D + p_base = p_ptr + row.to(tl.int64) * Skv + v_base = v_ptr + (b * Hkv + kv_head).to(tl.int64) * Skv * D + + acc = tl.zeros((BLOCK_D,), dtype=tl.float32) + for col in range(0, Skv): + p = tl.load(p_base + col) + vv = tl.load(v_base + col.to(tl.int64) * D + offs_d, mask=d_in, other=0.0).to(tl.float32) + acc = tl.fma(p, vv, acc) + + tl.store(out_ptr + row.to(tl.int64) * D + offs_d, acc, mask=d_in) + + +# --------------------------------------------------------------------------- +# Backward +# --------------------------------------------------------------------------- + + +@triton.jit +def _dp_kernel( + do_ptr, + v_ptr, + dp_ptr, + Hq, + Hkv, + Sq, + Skv, + D: tl.constexpr, + BLOCK_Q: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """dP[b, hq, q, k] = sum_{d ascending} dO[b,hq,q,d] * V[b,kv,k,d].""" + pid_k = tl.program_id(0) + pid_q = tl.program_id(1) + bh = tl.program_id(2) + b = bh // Hq + hq = bh % Hq + kv_head = hq // (Hq // Hkv) + + offs_q = pid_q * BLOCK_Q + tl.arange(0, BLOCK_Q) + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + q_in = offs_q < Sq + k_in = offs_k < Skv + + do_rows = do_ptr + (b * Hq + hq).to(tl.int64) * Sq * D + offs_q.to(tl.int64) * D + v_rows = v_ptr + (b * Hkv + kv_head).to(tl.int64) * Skv * D + offs_k.to(tl.int64) * D + + acc = tl.zeros((BLOCK_Q, BLOCK_K), dtype=tl.float32) + for d in range(0, D): + dov = tl.load(do_rows + d, mask=q_in, other=0.0).to(tl.float32) + vv = tl.load(v_rows + d, mask=k_in, other=0.0).to(tl.float32) + acc = tl.fma(dov[:, None], vv[None, :], acc) + + dst = ( + dp_ptr + + (b * Hq + hq).to(tl.int64) * Sq * Skv + + offs_q.to(tl.int64)[:, None] * Skv + + offs_k[None, :] + ) + tl.store(dst, acc, mask=q_in[:, None] & k_in[None, :]) + + +@triton.jit +def _softmax_backward_kernel( + ds_ptr, + p_ptr, + Skv, + LANES: tl.constexpr, +): + """delta = sum_k dP*P (C++ tree order); then dS = P * (dP - delta) in place.""" + row = tl.program_id(0) + ds_base = ds_ptr + row.to(tl.int64) * Skv + p_base = p_ptr + row.to(tl.int64) * Skv + lane = tl.arange(0, LANES) + + lane_delta = tl.zeros((LANES,), dtype=tl.float32) + for start in range(0, Skv, LANES): + cols = start + lane + in_range = cols < Skv + dp = tl.load(ds_base + cols, mask=in_range, other=0.0) + p = tl.load(p_base + cols, mask=in_range, other=0.0) + lane_delta = tl.fma(dp, p, lane_delta) + delta = _tree_sum_256(lane_delta) + + for start in range(0, Skv, LANES): + cols = start + lane + in_range = cols < Skv + dp = tl.load(ds_base + cols, mask=in_range, other=0.0) + p = tl.load(p_base + cols, mask=in_range, other=0.0) + tl.store(ds_base + cols, p * (dp - delta), mask=in_range) + + +@triton.jit +def _dq_kernel( + ds_ptr, + k_ptr, + dq_ptr, + scale, + Hq, + Hkv, + Sq, + Skv, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """dQ[b, hq, q, d] = scale * sum_{k ascending} dS[b,hq,q,k] * K[b,kv,k,d].""" + row = tl.program_id(0) + bh = row // Sq + b = bh // Hq + hq = bh % Hq + kv_head = hq // (Hq // Hkv) + + offs_d = tl.arange(0, BLOCK_D) + d_in = offs_d < D + ds_base = ds_ptr + row.to(tl.int64) * Skv + k_base = k_ptr + (b * Hkv + kv_head).to(tl.int64) * Skv * D + + acc = tl.zeros((BLOCK_D,), dtype=tl.float32) + for col in range(0, Skv): + ds = tl.load(ds_base + col) + kv = tl.load(k_base + col.to(tl.int64) * D + offs_d, mask=d_in, other=0.0).to(tl.float32) + acc = tl.fma(ds, kv, acc) + + tl.store(dq_ptr + row.to(tl.int64) * D + offs_d, scale * acc, mask=d_in) + + +@triton.jit +def _dk_kernel( + ds_ptr, + q_ptr, + dk_ptr, + scale, + Hq, + Hkv, + Sq, + Skv, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """dK[b, hkv, k, d] = scale * sum_{group head, then q, both ascending} dS * Q.""" + k_idx = tl.program_id(0) + b_hkv = tl.program_id(1) + b = b_hkv // Hkv + hkv = b_hkv % Hkv + group = Hq // Hkv + + offs_d = tl.arange(0, BLOCK_D) + d_in = offs_d < D + + acc = tl.zeros((BLOCK_D,), dtype=tl.float32) + for local in range(0, group): + hq = hkv * group + local + ds_head = ds_ptr + (b * Hq + hq).to(tl.int64) * Sq * Skv + k_idx + q_head = q_ptr + (b * Hq + hq).to(tl.int64) * Sq * D + for qi in range(0, Sq): + ds = tl.load(ds_head + qi.to(tl.int64) * Skv) + qv = tl.load(q_head + qi.to(tl.int64) * D + offs_d, mask=d_in, other=0.0).to(tl.float32) + acc = tl.fma(ds, qv, acc) + + dst = dk_ptr + (b * Hkv + hkv).to(tl.int64) * Skv * D + k_idx.to(tl.int64) * D + offs_d + tl.store(dst, scale * acc, mask=d_in) + + +@triton.jit +def _dv_kernel( + p_ptr, + do_ptr, + dv_ptr, + Hq, + Hkv, + Sq, + Skv, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """dV[b, hkv, k, d] = sum_{group head, then q, both ascending} P * dO.""" + k_idx = tl.program_id(0) + b_hkv = tl.program_id(1) + b = b_hkv // Hkv + hkv = b_hkv % Hkv + group = Hq // Hkv + + offs_d = tl.arange(0, BLOCK_D) + d_in = offs_d < D + + acc = tl.zeros((BLOCK_D,), dtype=tl.float32) + for local in range(0, group): + hq = hkv * group + local + p_head = p_ptr + (b * Hq + hq).to(tl.int64) * Sq * Skv + k_idx + do_head = do_ptr + (b * Hq + hq).to(tl.int64) * Sq * D + for qi in range(0, Sq): + p = tl.load(p_head + qi.to(tl.int64) * Skv) + dov = tl.load(do_head + qi.to(tl.int64) * D + offs_d, mask=d_in, other=0.0).to( + tl.float32 + ) + acc = tl.fma(p, dov, acc) + + dst = dv_ptr + (b * Hkv + hkv).to(tl.int64) * Skv * D + k_idx.to(tl.int64) * D + offs_d + tl.store(dst, acc, mask=d_in) + + +# --------------------------------------------------------------------------- +# Launchers +# --------------------------------------------------------------------------- + + +def _dummy_mask(reference: torch.Tensor) -> torch.Tensor: + return reference.new_empty((1,), dtype=torch.bool) + + +def triton_deterministic_attention_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool, + scale: float, + key_padding_mask: Optional[torch.Tensor], + output_fp32: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return ``(out, lse, P)`` matching ``_C.deterministic_attention_forward`` bitwise.""" + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + mask = key_padding_mask.contiguous() if key_padding_mask is not None else None + + b, hq, sq, d = q.shape + hkv, skv = k.shape[1], k.shape[2] + + scores = torch.empty((b, hq, sq, skv), device=q.device, dtype=torch.float32) + lse = torch.empty((b, hq, sq), device=q.device, dtype=torch.float32) + out = torch.empty_like(q, dtype=torch.float32 if output_fp32 else q.dtype) + + _qk_kernel[(triton.cdiv(skv, _QK_BLOCK_K), triton.cdiv(sq, _QK_BLOCK_Q), b * hq)]( + q, + k, + scores, + float(scale), + hq, + hkv, + sq, + skv, + D=d, + BLOCK_Q=_QK_BLOCK_Q, + BLOCK_K=_QK_BLOCK_K, + num_warps=4, + ) + _masked_softmax_lse_kernel[(b * hq * sq,)]( + scores, + lse, + mask if mask is not None else _dummy_mask(q), + hq, + sq, + skv, + CAUSAL=causal, + HAS_MASK=mask is not None, + LANES=_SOFTMAX_LANES, + num_warps=4, + ) + _pv_kernel[(b * hq * sq,)]( + scores, + v, + out, + hq, + hkv, + sq, + skv, + D=d, + BLOCK_D=d, + num_warps=4, + ) + return out, lse, scores + + +def triton_deterministic_attention_backward( + grad_output: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + p: torch.Tensor, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return ``(dQ, dK, dV)`` matching ``_C.deterministic_attention_backward`` bitwise.""" + do = grad_output.contiguous() + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + p = p.contiguous() + + b, hq, sq, d = q.shape + hkv, skv = k.shape[1], k.shape[2] + + # dS reuses the dP buffer exactly like the native backward does. + ds = torch.empty((b, hq, sq, skv), device=q.device, dtype=torch.float32) + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + + _dp_kernel[(triton.cdiv(skv, _QK_BLOCK_K), triton.cdiv(sq, _QK_BLOCK_Q), b * hq)]( + do, + v, + ds, + hq, + hkv, + sq, + skv, + D=d, + BLOCK_Q=_QK_BLOCK_Q, + BLOCK_K=_QK_BLOCK_K, + num_warps=4, + ) + _softmax_backward_kernel[(b * hq * sq,)]( + ds, + p, + skv, + LANES=_SOFTMAX_LANES, + num_warps=4, + ) + _dq_kernel[(b * hq * sq,)]( + ds, + k, + dq, + float(scale), + hq, + hkv, + sq, + skv, + D=d, + BLOCK_D=d, + num_warps=4, + ) + _dk_kernel[(skv, b * hkv)]( + ds, + q, + dk, + float(scale), + hq, + hkv, + sq, + skv, + D=d, + BLOCK_D=d, + num_warps=4, + ) + _dv_kernel[(skv, b * hkv)]( + p, + do, + dv, + hq, + hkv, + sq, + skv, + D=d, + BLOCK_D=d, + num_warps=4, + ) + return dq, dk, dv + + +class _TritonDeterministicAttentionFn(Function): + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + causal: bool, + scale: float, + key_padding_mask: Optional[torch.Tensor], + output_fp32: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + q_c = q.contiguous() + k_c = k.contiguous() + v_c = v.contiguous() + mask_c = key_padding_mask.contiguous() if key_padding_mask is not None else None + + out, lse, p = triton_deterministic_attention_forward( + q_c, k_c, v_c, causal, float(scale), mask_c, output_fp32 + ) + + ctx.save_for_backward(q_c, k_c, v_c, p, mask_c) + ctx.causal = causal + ctx.scale = scale + ctx.mark_non_differentiable(lse) + + return out, lse + + @staticmethod + @once_differentiable + def backward(ctx, grad_out: torch.Tensor, grad_lse: torch.Tensor): + q_c, k_c, v_c, p, _mask_c = ctx.saved_tensors + + if grad_out.dtype != q_c.dtype: + grad_out = grad_out.to(q_c.dtype) + dq, dk, dv = triton_deterministic_attention_backward( + grad_out.contiguous(), q_c, k_c, v_c, p, float(ctx.scale) + ) + return dq, dk, dv, None, None, None, None + + +class TritonDeterministicAttentionOp: + """Triton twin of :class:`DeterministicAttentionOp`, bitwise identical to it. + + The public surface matches the native op so either can be dropped into the + strict-attention harness. Validation is duplicated rather than imported so the + Triton path stays usable when the native extension is not built. + """ + + backend_id = "rlkernel.triton.deterministic_attention" + + def __init__(self, *, require_bitwise_libm: bool = True) -> None: + """``require_bitwise_libm=False`` trades bitwise parity for portability. + + The softmax needs a bitwise-exact ``expf``/``logf``; only the HIP sequences + are ported (see :data:`BITWISE_LIBM_PARITY`). Opting out keeps the kernel + deterministic and batch-invariant but no longer bit-identical to the + native core, so it is never the default. + """ + if require_bitwise_libm and not BITWISE_LIBM_PARITY: + raise RuntimeError( + "Triton deterministic attention is bitwise-identical to " + "_C.deterministic_attention_* only on ROCm: the nvcc expf/logf " + "argument reduction has not been ported. Construct with " + "require_bitwise_libm=False to run the non-bitwise fallback." + ) + self.bitwise_libm = BITWISE_LIBM_PARITY + + 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, + ) -> torch.Tensor: + return self.forward(q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask) + + 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, + ) -> torch.Tensor: + out, _lse = self.forward_with_lse( + q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask + ) + 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, + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_inputs(q, k, v, key_padding_mask) + resolved_scale = scale if scale is not None else (1.0 / math.sqrt(q.shape[-1])) + return _TritonDeterministicAttentionFn.apply( + q, k, v, causal, resolved_scale, key_padding_mask, False + ) + + 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, + ) -> torch.Tensor: + self._validate_inputs(q, k, v, key_padding_mask) + resolved_scale = scale if scale is not None else (1.0 / math.sqrt(q.shape[-1])) + out, _lse = _TritonDeterministicAttentionFn.apply( + q, k, v, causal, resolved_scale, key_padding_mask, True + ) + return out + + @staticmethod + def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + key_padding_mask: Optional[torch.Tensor], + ) -> None: + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + raise ValueError( + f"q/k/v must be 4-D [B, H, S, D], got q={tuple(q.shape)}, " + f"k={tuple(k.shape)}, v={tuple(v.shape)}" + ) + b, hq, sq, d = q.shape + hkv, skv = k.shape[1], k.shape[2] + if k.shape[0] != b or v.shape[0] != b: + raise ValueError("batch size mismatch between q/k/v") + if v.shape[1] != hkv or v.shape[2] != skv or k.shape[3] != d or v.shape[3] != d: + raise ValueError( + f"k/v shape mismatch: k={tuple(k.shape)}, v={tuple(v.shape)}, " + f"expected k/v [B={b}, Hkv, Skv, D={d}]" + ) + if d != _HEAD_DIM: + raise ValueError(f"head dim D must be {_HEAD_DIM}, got {d}") + if hq % hkv != 0: + raise ValueError(f"Hq={hq} not divisible by Hkv={hkv} (GQA group)") + if q.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"only FP16/BF16 supported, got {q.dtype}") + if k.dtype != q.dtype or v.dtype != q.dtype: + raise ValueError("q, k, v must share the same dtype") + if not (q.is_cuda and k.is_cuda and v.is_cuda): + raise ValueError("q, k, v must be GPU tensors") + if key_padding_mask is not None: + if key_padding_mask.dtype != torch.bool: + raise ValueError("key_padding_mask must be bool") + if key_padding_mask.shape != (b, skv): + raise ValueError( + f"key_padding_mask must be [B, Skv]=[{b}, {skv}], " + f"got {tuple(key_padding_mask.shape)}" + ) + if sq < 1 or skv < 1: + raise ValueError(f"Sq and Skv must be positive, got Sq={sq}, Skv={skv}") + + +def triton_deterministic_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return TritonDeterministicAttentionOp().forward( + q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask + ) + + +def triton_deterministic_attention_with_lse( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + return TritonDeterministicAttentionOp().forward_with_lse( + q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask + ) + + +def triton_deterministic_attention_fp32( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + causal: bool = True, + scale: Optional[float] = None, + key_padding_mask: Optional[torch.Tensor] = None, +) -> torch.Tensor: + return TritonDeterministicAttentionOp().forward_fp32( + q, k, v, causal=causal, scale=scale, key_padding_mask=key_padding_mask + ) + + +__all__ = [ + "TritonDeterministicAttentionOp", + "triton_deterministic_attention", + "triton_deterministic_attention_backward", + "triton_deterministic_attention_forward", + "triton_deterministic_attention_fp32", + "triton_deterministic_attention_with_lse", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 12ea9b21..f715bad8 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -71,6 +71,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ROCM_AITER = "rl_engine.kernels.ops.rocm.aiter.AiterOp" ROCM_CK = "rl_engine.kernels.ops.rocm.composable_kernel.CKOp" ROCM_FLASH_ATTN = "rl_engine.kernels.ops.rocm.attention.flash_attn.RocmFlashAttentionOp" + # WS2 strict ROCm attention core (AITER/CK dense MHA, Split-KV disabled). + # Reachable only through ``get_attention_op``: it is not a drop-in for the + # SDPA-shaped ``attn`` wrappers and must never be a silent fallback. + ROCM_STRICT_ATTENTION = ( + "rl_engine.kernels.ops.rocm.attention.flash_attn.StrictRocmAiterCKAttentionCore" + ) # GRPO loss (group reward normalization + clipped surrogate + KL) TRITON_GRPO_LOSS = "rl_engine.kernels.ops.triton.loss.grpo_loss.TritonGRPOLossOp" @@ -346,6 +352,27 @@ def resolve_logp_op_type( return op_type +def _rocm_strict_attention_available() -> bool: + """Return whether the strict ROCm AITER/CK attention core can actually load. + + True only when this process can really execute the strict arithmetic, so an + unavailable vendor stack leaves the backend unregistered rather than + registered-but-failing at materialization time. + """ + + if torch.version.hip is None: + return False + try: + from rl_engine.kernels.ops.rocm.attention.flash_attn import _load_aiter_ck_ops + except ImportError: + return False + try: + _load_aiter_ck_ops() + except Exception: # StrictRocmAttentionUnavailable and vendor import errors + return False + return True + + class KernelRegistry: """ Central dispatcher for high-performance kernels. @@ -667,6 +694,83 @@ def __init__(self): prepend=True, ) + # Strict ROCm production core: AITER/CK dense MHA, Split-KV disabled. + # Registered only when the vendor entry points genuinely load, so an + # explicit request on a machine without AITER fails loudly in + # ``get_attention_op`` instead of resolving to different arithmetic. + if _rocm_strict_attention_available(): + self.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, + AttentionBackendCapability( + backend_id="aiter.rocm.ck_dense_mha", + roles=frozenset({AttentionRole.TRAIN, AttentionRole.INFER}), + # DECODE is deliberately absent. StrictRocmAttentionRuntime + # has a paged entry point, but no caller routes to it: the + # Vime request carries no page table and builds its contract + # with kv_cache=None. Declaring the mode before a dispatch + # path reaches it would let the binding layer pass on a + # decode path nothing executes. + modes=frozenset({AttentionMode.PREFILL, AttentionMode.CHUNKED_PREFILL}), + dtypes=frozenset({AttentionDType.BF16, AttentionDType.FP16}), + # The core itself is single-rank arithmetic. CP is supplied + # by StrictRocmAttentionRuntime, which wraps this core in + # the RCCL AG/RS transport; the sizes mirror the world + # sizes RCCLDeterministicCollective accepts. The merge + # order is that collective's fixed balanced rank tree, not + # RCCL's own reduction, so the CP merge is deterministic. + cp_world_sizes=(1, 2, 4, 8), + tp_world_sizes=None, + exports_attention_lse=True, + deterministic_cp_merge=True, + supports_packed_varlen=False, + supports_kv_cache=False, + supports_rope_metadata=True, + supports_fused_rope_attention=False, + supports_split_kv_disabled=True, + supports_split_kv_fixed=False, + supports_split_kv_auto=False, + reports_actual_split_kv_plan=True, + implementation_kind="production", + ), + platform="rocm", + prepend=True, + ) + + def register_attention_backend( + self, + backend: OpBackend, + capability: AttentionBackendCapability, + *, + platform: Optional[str] = None, + prepend: bool = False, + ) -> None: + """Register (or replace) a backend for WS2 contract-aware attention dispatch. + + The supported seam for a backend that only exists on some machines: the + static ``ws2_attention`` priority lists cannot express "present only when + the vendor stack loads", so a conditional backend registers itself here. + Re-registering the same backend replaces its capability without + duplicating the candidate entry. + """ + + if not isinstance(backend, OpBackend): + raise AttentionContractError("backend must be an OpBackend") + if not isinstance(capability, AttentionBackendCapability): + raise AttentionContractError("capability must be an AttentionBackendCapability") + resolved_platform = platform if platform is not None else self._platform() + if resolved_platform not in self._priority_map: + raise AttentionContractError( + f"unsupported platform {resolved_platform!r}; expected one of " + f"{sorted(self._priority_map)}" + ) + self._attention_capabilities[backend] = capability + candidates = self._priority_map[resolved_platform].setdefault("ws2_attention", []) + if backend not in candidates: + if prepend: + candidates.insert(0, backend) + else: + candidates.append(backend) + def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() if rocm_attn_backend in {"flash_attn", "flash-attn", "flash_attention"}: @@ -931,6 +1035,13 @@ def get_attention_op( if not isinstance(requested_backend, str) or not requested_backend.strip(): raise AttentionContractError("requested_backend must be a non-empty string") requested_backend = requested_backend.strip().lower() + if requested_backend == "auto" and contract.sharding.cp_world_size > 1: + raise AttentionContractError( + "Unsafe dispatch: requested_backend='auto' is not permitted when " + "cp_world_size > 1 without explicit cross-rank preflighting; name a " + "policy or backend id and agree on " + "AttentionContract.cross_rank_fingerprint() across ranks." + ) platform = self._platform() candidates = self._priority_map.get(platform, {}).get("ws2_attention", []) diff --git a/scripts/ws2_p2p_nccl_attention_reference_check.py b/scripts/ws2_p2p_nccl_attention_reference_check.py index 7b799d74..e04c4c06 100644 --- a/scripts/ws2_p2p_nccl_attention_reference_check.py +++ b/scripts/ws2_p2p_nccl_attention_reference_check.py @@ -15,6 +15,7 @@ import json import math import os +import subprocess import sys from pathlib import Path from types import SimpleNamespace @@ -30,6 +31,9 @@ from rl_engine.kernels.attention_contract import ( # noqa: E402 STRICT_ATTENTION_FA4_SCHEDULE_ID, STRICT_ATTENTION_PRODUCTION_CORE_ID, + STRICT_ATTENTION_RING_SCHEDULE_ID, + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, ) from rl_engine.kernels.ops.cuda.attention.cp_comm import ( # noqa: E402 AttentionCPBlockMetadata, @@ -39,6 +43,7 @@ AttentionParallelSpec, CUDAAGRSAttentionCPCommunication, P2PNCCLAttentionCPCommunication, + RCCLAGRSAttentionCPCommunication, ) from rl_engine.kernels.ops.cuda.attention.flash_attn import StrictFlashAttention4Core # noqa: E402 from rl_engine.kernels.ops.cuda.attention.flashinfer_paged_attention import ( # noqa: E402 @@ -46,7 +51,6 @@ FlashInferQwen3PagedAttentionOp, _apply_strict_rope, ) -from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op # noqa: E402 from rl_engine.kernels.ops.pytorch.attention.cp_attention import ( # noqa: E402 AttentionPartialState, DeterministicCPAttentionReferenceOp, @@ -68,9 +72,9 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument("--final-write-atol", type=float, default=2.0e-2) parser.add_argument( "--transport", - choices=("p2p_nccl_reference", "cuda_ag_rs"), + choices=("p2p_nccl_reference", "cuda_ag_rs", "rccl_ag_rs"), default="p2p_nccl_reference", - help="P2P is the correctness reference; cuda_ag_rs selects PR311/PR312", + help="P2P is the reference; cuda_ag_rs and rccl_ag_rs are self-owned transports", ) parser.add_argument( "--strict-shared-core", @@ -78,22 +82,128 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: help="run AG(Q/K/V/positions) -> shared CUDA core -> RS(Out/LSE) with backward", ) parser.add_argument("--output", type=Path) + parser.add_argument( + "--run-rocm-matrix", + action="store_true", + help="run the complete 1/2/4/8-GPU ROCm acceptance matrix", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("results/rocm-attention"), + help="directory for --run-rocm-matrix logs and JSON reports", + ) args = parser.parse_args(argv) - if args.strict_shared_core and args.transport != "cuda_ag_rs": - parser.error("--strict-shared-core requires --transport cuda_ag_rs") + if args.strict_shared_core and args.transport not in {"cuda_ag_rs", "rccl_ag_rs"}: + parser.error("--strict-shared-core requires a self-owned AG/RS transport") + if args.run_rocm_matrix and (args.strict_shared_core or args.output is not None): + parser.error("--run-rocm-matrix cannot be combined with single-run output options") return args +def _run_rocm_matrix(output_dir: Path) -> int: + if torch.version.hip is None or torch.cuda.device_count() < 8: + raise RuntimeError("the formal acceptance matrix requires 8 visible ROCm GPUs") + + repo = Path(__file__).resolve().parents[1] + output = (repo / output_dir).resolve() if not output_dir.is_absolute() else output_dir + output.mkdir(parents=True, exist_ok=True) + script = Path(__file__).resolve() + commands: list[tuple[str, list[str]]] = [ + ( + "single_gpu", + [sys.executable, "-m", "pytest", "-q", "tests/test_deterministic_attention_cuda.py"], + ), + ( + "adapter_cp_contracts", + [ + sys.executable, + "-m", + "pytest", + "-q", + "tests/test_flashinfer_pr7_attention.py", + "tests/test_cp_attention.py", + "tests/test_attention_comparison.py", + ], + ), + ] + for transport, strict in (("p2p_nccl_reference", False), ("rccl_ag_rs", True)): + for ranks in (2, 4, 8): + name = f"{transport}_{ranks}r" + command = [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + f"--nproc-per-node={ranks}", + str(script), + "--transport", + transport, + "--output", + str(output / f"{name}.json"), + ] + if strict: + command.append("--strict-shared-core") + commands.append((name, command)) + + steps: list[dict[str, object]] = [] + for name, command in commands: + completed = subprocess.run( + command, + cwd=repo, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + check=False, + ) + (output / f"{name}.log").write_text(completed.stdout, encoding="utf-8") + steps.append({"name": name, "command": command, "returncode": completed.returncode}) + + summary = { + "schema_version": "ws2_rocm_attention_acceptance/v1", + "git_commit": _current_git_commit(repo), + "platform": "rocm", + "torch": str(torch.__version__), + "hip": str(torch.version.hip), + "collective": list(torch.cuda.nccl.version()), + "device_count": torch.cuda.device_count(), + "device_name": torch.cuda.get_device_name(0), + "steps": steps, + "passed": all(step["returncode"] == 0 for step in steps), + } + (output / "summary.json").write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + print(json.dumps(summary, indent=2, sort_keys=True)) + return 0 if summary["passed"] else 1 + + +def _current_git_commit(repo: Path) -> str: + """Bind an acceptance artifact to the checkout that generated it.""" + + try: + return subprocess.check_output( + ["git", "-C", str(repo), "rev-parse", "HEAD"], + text=True, + stderr=subprocess.STDOUT, + ).strip() + except (OSError, subprocess.CalledProcessError) as exc: + raise RuntimeError("ROCm Attention acceptance requires a Git checkout") from exc + + def main(argv: Sequence[str] | None = None) -> int: args = parse_args(argv) + if args.run_rocm_matrix: + return _run_rocm_matrix(args.output_dir) if not torch.cuda.is_available() or torch.cuda.device_count() < 2: - raise RuntimeError("this check requires at least two visible CUDA devices") + raise RuntimeError("this check requires at least two visible CUDA/ROCm devices") dist.init_process_group("nccl", init_method="env://") try: world_size = dist.get_world_size() global_rank = dist.get_rank() if world_size not in {2, 4, 8}: - raise RuntimeError("this check requires 2, 4, or 8 NCCL ranks") + raise RuntimeError("this check requires 2, 4, or 8 NCCL/RCCL ranks") if torch.cuda.device_count() < world_size: raise RuntimeError("this single-node check requires one visible GPU per NCCL rank") local_rank = int(os.environ.get("LOCAL_RANK", str(global_rank))) @@ -129,12 +239,18 @@ def main(argv: Sequence[str] | None = None) -> int: dist.all_gather_object(reports, result) if global_rank == 0: report = { - "schema_version": ( - "ws2_p2p_nccl_attention_reference/v1" - if args.transport == "p2p_nccl_reference" - else "ws2_cuda_ag_rs_attention/v1" - ), + "schema_version": f"ws2_{args.transport}_attention/v2", + "git_commit": _current_git_commit(REPO_ROOT), "backend": str(dist.get_backend()), + "platform": "rocm" if torch.version.hip is not None else "cuda", + "torch_version": str(torch.__version__), + "runtime_version": ( + str(torch.version.hip) + if torch.version.hip is not None + else str(torch.version.cuda) + ), + "collective_version": list(torch.cuda.nccl.version()), + "device_name": torch.cuda.get_device_name(0), "transport": args.transport, "world_size": world_size, "tp_world_size": 1 if world_size == 2 else 2, @@ -210,11 +326,17 @@ def run_check( expected_kv_token_range=(0, args.seq_len), query_token_ranges=owner_ranges, ) - communication: P2PNCCLAttentionCPCommunication | CUDAAGRSAttentionCPCommunication + communication: ( + P2PNCCLAttentionCPCommunication + | CUDAAGRSAttentionCPCommunication + | RCCLAGRSAttentionCPCommunication + ) if args.transport == "p2p_nccl_reference": communication = P2PNCCLAttentionCPCommunication(process_group=cp_group) - else: + elif args.transport == "cuda_ag_rs": communication = CUDAAGRSAttentionCPCommunication(process_group=cp_group) + else: + communication = RCCLAGRSAttentionCPCommunication(process_group=cp_group) query_start, query_end = owner_ranges[cp_rank] q_local = q[:, :, query_start:query_end, :].contiguous() @@ -348,7 +470,11 @@ def _run_strict_shared_core_check( args: argparse.Namespace, *, plan: AttentionCPCommunicationPlan, - communication: CUDAAGRSAttentionCPCommunication | P2PNCCLAttentionCPCommunication, + communication: ( + CUDAAGRSAttentionCPCommunication + | RCCLAGRSAttentionCPCommunication + | P2PNCCLAttentionCPCommunication + ), q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -381,16 +507,14 @@ def _run_strict_shared_core_check( q_ref = q.detach().clone().requires_grad_() k_ref = k.detach().clone().requires_grad_() v_ref = v.detach().clone().requires_grad_() - rope = RoPESM90Op() + rope = _strict_rope_op() q_ready = _apply_strict_rope(rope, q_ref, positions, config.rope.rope_theta) k_ready = _apply_strict_rope(rope, k_ref, positions, config.rope.rope_theta) - reference = StrictFlashAttention4Core().forward_with_lse( + reference = _strict_attention_reference_rows( q_ready, k_ready, v_ref, - causal=True, - query_position_ids=positions, - key_position_ids=positions, + positions=positions, output_dtype=q.dtype, ) out_ref = reference.out[:, :, start:end, :] @@ -435,23 +559,18 @@ def _run_strict_shared_core_check( repeat_lse_bitwise = repeat_lse_bitwise and torch.equal(repeated.lse, distributed.lse) provenance = distributed.provenance - identity_valid = ( - provenance.get("strict_core_id") == STRICT_ATTENTION_PRODUCTION_CORE_ID - and provenance.get("strict_schedule") == STRICT_ATTENTION_FA4_SCHEDULE_ID - and provenance.get("strict_mode") is True - and provenance.get("native_attention_arithmetic") is True - and provenance.get("num_splits") == 1 - and provenance.get("deterministic_backward") is True - and provenance.get("fa_api_source") == "flash_attn.cute.interface" - and provenance.get("fallback") is False - and provenance.get("strict_split_kv") == "disabled" - and provenance.get("strict_comm_autograd") is True - and provenance.get("production_ready") is True + identity_errors = _strict_shared_core_identity_errors( + provenance, + transport=args.transport, + is_rocm=torch.version.hip is not None, ) return { "executed": True, "passed": ( - all(bitwise.values()) and repeat_out_bitwise and repeat_lse_bitwise and identity_valid + all(bitwise.values()) + and repeat_out_bitwise + and repeat_lse_bitwise + and not identity_errors ), "strict_core_id": provenance.get("strict_core_id"), "strict_schedule": provenance.get("strict_schedule"), @@ -463,6 +582,8 @@ def _run_strict_shared_core_check( "fallback": provenance.get("fallback"), "split_kv_policy": provenance.get("strict_split_kv"), "communication_autograd": provenance.get("strict_comm_autograd"), + "strict_provenance": provenance, + "identity_errors": identity_errors, "bitwise": bitwise, "max_abs": max_abs, "repeat_out_bitwise": repeat_out_bitwise, @@ -470,5 +591,112 @@ def _run_strict_shared_core_check( } +def _strict_shared_core_identity_errors( + provenance: dict[str, object], + *, + transport: str, + is_rocm: bool, +) -> list[str]: + """Return every strict-contract provenance mismatch for the rank report.""" + + expected_core = ( + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID if is_rocm else STRICT_ATTENTION_PRODUCTION_CORE_ID + ) + expected_schedule = ( + STRICT_ATTENTION_ROCM_SCHEDULE_ID if is_rocm else STRICT_ATTENTION_FA4_SCHEDULE_ID + ) + expected_backend = "aiter.rocm.ck_dense_mha" if is_rocm else "flash_attention_4.cute" + expected_rope = "rlkernel.rocm.deterministic_rope" if is_rocm else "rlkernel.cuda.rope_sm90" + expected_communication = "rccl_ag_rs" if is_rocm else "self_owned_cuda_ag_rs" + required = { + "strict_core_id": expected_core, + "strict_schedule": expected_schedule, + "attention_backend": expected_backend, + "actual_backend": expected_backend, + "rope_backend": expected_rope, + "strict_mode": True, + "native_attention_arithmetic": True, + "num_splits": 1, + "deterministic_backward": True, + "reference_only": False, + "fallback": False, + "strict_split_kv": "disabled", + "strict_comm_autograd": True, + "communication_backend": expected_communication, + "production_ready": True, + "strict_full_qkv_all_gather": True, + "strict_position_ids_all_gather": True, + "compute_communication": "decoupled", + "compute_schedule": STRICT_ATTENTION_RING_SCHEDULE_ID, + "communication_overlap": "disabled", + "ring_schedule_default": True, + "ring_partial_arithmetic": False, + "rope_fusion": False, + "q_rope_state": "post_rope", + "k_cache_rope_state": "post_rope", + } + errors = [ + f"{name}={provenance.get(name)!r}, expected {expected!r}" + for name, expected in required.items() + if provenance.get(name) != expected + ] + if is_rocm: + if provenance.get("split_kv_control") != "dense_non_split_api": + errors.append("split_kv_control must prove the AITER dense non-Split-K API") + if provenance.get("aiter_api_source") != "aiter.ops.mha": + errors.append("aiter_api_source must identify aiter.ops.mha") + if not provenance.get("aiter_source_sha256"): + errors.append("aiter_source_sha256 is missing") + else: + if provenance.get("fa_api_source") != "flash_attn.cute.interface": + errors.append("fa_api_source must identify the FA4 CuTe API") + return errors + + +def _strict_attention_core(): + if torch.version.hip is not None: + from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore + + return StrictRocmAiterCKAttentionCore() + return StrictFlashAttention4Core() + + +def _strict_attention_reference_rows( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + positions: torch.Tensor, + output_dtype: torch.dtype, +) -> SimpleNamespace: + """Run the production core with its one-logical-row execution contract.""" + + core = _strict_attention_core() + rows = [ + core.forward_with_lse( + q[index : index + 1], + k[index : index + 1], + v[index : index + 1], + causal=True, + query_position_ids=positions[index : index + 1], + key_position_ids=positions[index : index + 1], + output_dtype=output_dtype, + ) + for index in range(q.size(0)) + ] + return SimpleNamespace( + out=torch.cat([row.out for row in rows], dim=0), + lse=torch.cat([row.lse for row in rows], dim=0), + ) + + +def _strict_rope_op(): + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RocmDeterministicRoPEOp, RoPESM90Op + + if torch.version.hip is not None: + return RocmDeterministicRoPEOp() + return RoPESM90Op() + + if __name__ == "__main__": raise SystemExit(main()) diff --git a/scripts/ws2_pr7_flashinfer_attention_check.py b/scripts/ws2_pr7_flashinfer_attention_check.py index 7ab64842..7a11863c 100644 --- a/scripts/ws2_pr7_flashinfer_attention_check.py +++ b/scripts/ws2_pr7_flashinfer_attention_check.py @@ -1,12 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""PR7 FlashInfer RoPE-fused paged attention validation entry point. +"""PR7 strict paged-layout Attention validation entry point. -The default dry-run mode is CI/local friendly: it builds the FlashInfer page plan -and provenance without importing FlashInfer or requiring CUDA. On a CUDA host -with FlashInfer installed, omit ``--dry-run`` to run the opt-in PR7 candidate and -compare it with the PR6 full logical KV reference. +The default dry-run mode is CI/local friendly: it builds the paged-KV plan and +provenance without requiring a GPU vendor backend. On a GPU host, omit +``--dry-run`` to run the platform production core and its strict reference. """ from __future__ import annotations @@ -28,6 +27,8 @@ from rl_engine.kernels.attention_contract import ( # noqa: E402 STRICT_ATTENTION_FA4_SCHEDULE_ID, STRICT_ATTENTION_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, ) from rl_engine.kernels.ops.cuda.attention.cp_comm import ( # noqa: E402 AttentionCPCommunicationPlan, @@ -44,7 +45,6 @@ _materialize_strict_logical_kv, build_flashinfer_paged_kv_plan, ) -from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RoPESM90Op # noqa: E402 from rl_engine.testing.attention_comparison import ( # noqa: E402 AttentionPathResult, DecodeAttentionInputs, @@ -69,7 +69,7 @@ def main(argv: Sequence[str] | None = None) -> int: report: dict[str, Any] = { "status": "dry_run" if args.dry_run else "executed", "pr": "PR7", - "target": "Qwen3-8B TP-local FlashInfer candidate; CP transport validated separately", + "target": "Qwen3-8B TP-local paged Attention; CP transport validated separately", "mode": config.mode, "device": str(device), "shape": { @@ -98,11 +98,11 @@ def main(argv: Sequence[str] | None = None) -> int: | {"cp_comm_required": config.require_cp_comm}, "paged_kv_plan": plan.provenance(), "tests_expected": [ - "FlashInfer ROPE_LLAMA vs NativeRoPEOp + full logical KV reference", + "platform production core vs direct same-core reference", "split-K disabled/fixed policy drift", "batch composition/position invariant sweep", "attention-domain LSE export drift", - "strict shared CUDA core with separate multi-rank AG/RS forward/backward evidence", + "strict platform core with separate multi-rank AG/RS forward/backward evidence", ], "thresholds": { "out_max_abs": args.out_atol, @@ -152,7 +152,7 @@ def main(argv: Sequence[str] | None = None) -> int: ) pytorch_reference = run_decode_full_prefill_reference(reference_inputs) reference = ( - _run_strict_cuda_reference(inputs, config, plan) + _run_strict_platform_reference(inputs, config, plan) if config.strict_mode else pytorch_reference ) @@ -180,12 +180,9 @@ def main(argv: Sequence[str] | None = None) -> int: "lse": lse_stats, "dlogp": dlogp_stats, } - report["reference_backend"] = ( - "rlkernel.cuda.deterministic_attention" - if config.strict_mode - else "rlkernel.pytorch.full_logical_kv_reference" - ) + report["reference_backend"] = "rlkernel.pytorch.full_logical_kv_reference" if config.strict_mode: + report.update(_strict_execution_report_fields(candidate.provenance)) report["diagnostic_drift_vs_pytorch"] = { "out": _drift_stats(candidate.out, pytorch_reference.out), "lse": _drift_stats(candidate.lse, pytorch_reference.lse), @@ -391,15 +388,15 @@ def _make_inputs(args: argparse.Namespace, device: torch.device) -> DecodeAttent return DecodeAttentionInputs(q=q, k_cache=k_cache, v_cache=v_cache, metadata=metadata) -def _run_strict_cuda_reference( +def _run_strict_platform_reference( inputs: DecodeAttentionInputs, config: FlashInferPagedAttentionConfig, paged_plan: Any, ) -> AttentionPathResult: - """Call the production FA4 core directly on the logical KV sequence.""" + """Call the platform production core directly on the logical KV sequence.""" - core = config.deterministic_core or StrictFlashAttention4Core(split_kv=config.split_kv) - rope = config.strict_rope_op or RoPESM90Op() + core = config.deterministic_core or _strict_attention_core(config.split_kv) + rope = config.strict_rope_op or _strict_rope_op() logical_k, logical_v, key_positions = _materialize_strict_logical_kv( inputs.k_cache, inputs.v_cache, @@ -429,7 +426,7 @@ def _run_strict_cuda_reference( outputs.append(result.out) lses.append(result.lse) return AttentionPathResult( - name="direct_flash_attention4_num_splits1", + name="direct_platform_strict_attention", out=torch.cat(outputs, dim=0), lse=torch.cat(lses, dim=0), provenance={ @@ -440,6 +437,28 @@ def _run_strict_cuda_reference( ) +def _strict_execution_report_fields(provenance: dict[str, Any]) -> dict[str, Any]: + """Describe the backend that actually executed, independent of the host running tests.""" + + backend = provenance.get("actual_backend", "unknown") + return { + "target": ( + f"Qwen3-8B TP-local {backend} strict production core; " + "CP transport validated separately" + ), + "reference_backend": backend, + "rope": { + "rope_backend": provenance.get("rope_backend"), + "rope_fusion": provenance.get("rope_fusion"), + "rope_fusion_boundary": provenance.get("rope_fusion_boundary"), + "rope_theta": provenance.get("rope_theta"), + "rotary_dim": provenance.get("rotary_dim"), + "q_rope_state": provenance.get("q_rope_state"), + "k_cache_rope_state": provenance.get("k_cache_rope_state"), + }, + } + + def _run_batch_invariance_sweep( op: FlashInferQwen3PagedAttentionOp, inputs: DecodeAttentionInputs, @@ -645,33 +664,58 @@ def _acceptance_errors(report: dict[str, Any], args: argparse.Namespace) -> list if provenance.get("fallback") is not False: errors.append("FlashInfer execution used or omitted fallback provenance") if args.strict: + platform = provenance.get("platform", "cuda") + if platform not in {"cuda", "rocm"}: + errors.append("strict runtime platform provenance is invalid") + is_rocm = platform == "rocm" + expected_core = ( + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID + if is_rocm + else STRICT_ATTENTION_PRODUCTION_CORE_ID + ) + expected_schedule = ( + STRICT_ATTENTION_ROCM_SCHEDULE_ID if is_rocm else STRICT_ATTENTION_FA4_SCHEDULE_ID + ) + expected_backend = "aiter.rocm.ck_dense_mha" if is_rocm else "flash_attention_4.cute" if provenance.get("strict_mode") is not True: errors.append("strict runtime did not execute the shared Attention core") - if provenance.get("strict_core_id") != STRICT_ATTENTION_PRODUCTION_CORE_ID: - errors.append("strict runtime did not execute the FA4 production core") - if provenance.get("strict_schedule") != STRICT_ATTENTION_FA4_SCHEDULE_ID: + if provenance.get("strict_core_id") != expected_core: + errors.append("strict runtime core identity is invalid") + if provenance.get("strict_schedule") != expected_schedule: errors.append("strict runtime arithmetic schedule is invalid") - if provenance.get("actual_backend") != "flash_attention_4.cute": - errors.append("strict runtime backend is not FlashAttention-4 CuTe") + if provenance.get("actual_backend") != expected_backend: + errors.append("strict runtime backend identity is invalid") if provenance.get("native_attention_arithmetic") is not True: - errors.append("strict runtime did not execute native FA4 Attention arithmetic") + errors.append("strict runtime did not execute the native production arithmetic") if provenance.get("num_splits") != 1: - errors.append("strict runtime did not fix FA4 num_splits=1") + errors.append("strict runtime did not prove one reduction partition") if provenance.get("deterministic_backward") is not True: - errors.append("strict runtime did not request deterministic FA4 backward") - if provenance.get("fa_api_source") != "flash_attn.cute.interface": - errors.append("strict runtime did not prove the FA4 CuTe API source") + errors.append("strict runtime did not request deterministic backward") if provenance.get("reference_only") is not False: errors.append("strict runtime selected the reference core") + if is_rocm: + if provenance.get("split_kv_control") != "dense_non_split_api": + errors.append("strict ROCm runtime did not use AITER dense non-Split-K MHA") + if provenance.get("aiter_api_source") != "aiter.ops.mha": + errors.append("strict ROCm runtime did not prove the AITER API source") + if not provenance.get("aiter_source_sha256"): + errors.append("strict ROCm runtime did not fingerprint AITER MHA") + elif provenance.get("fa_api_source") != "flash_attn.cute.interface": + errors.append("strict CUDA runtime did not prove the FA4 CuTe API source") strict_plans = provenance.get("strict_core_row_plans") if not isinstance(strict_plans, list) or not strict_plans: errors.append("strict no-Split-K execution plans are missing") elif any(plan.get("actual_split_kv_policy") != "disabled" for plan in strict_plans): errors.append("strict runtime did not keep Split-KV disabled") - if provenance.get("rope_backend") not in { - "rlkernel.cuda.rope_sm90", - "rlkernel.cuda.rope_sm90_op", - }: + expected_rope_backends = ( + {"rlkernel.rocm.deterministic_rope"} + if is_rocm + else { + "rlkernel.cuda.rope_sm90", + "rlkernel.cuda.rope_sm90_op", + } + ) + if provenance.get("rope_backend") not in expected_rope_backends: errors.append("strict runtime did not use the RL-Kernel WS1 RoPE operator") elif provenance.get("pos_encoding_mode") != "ROPE_LLAMA": errors.append("FlashInfer runtime did not use ROPE_LLAMA") @@ -693,6 +737,20 @@ def _acceptance_errors(report: dict[str, Any], args: argparse.Namespace) -> list return errors +def _strict_attention_core(split_kv): + if torch.version.hip is not None: + from rl_engine.kernels.ops.rocm.attention.flash_attn import StrictRocmAiterCKAttentionCore + + return StrictRocmAiterCKAttentionCore(split_kv=split_kv) + return StrictFlashAttention4Core(split_kv=split_kv) + + +def _strict_rope_op(): + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RocmDeterministicRoPEOp, RoPESM90Op + + return RocmDeterministicRoPEOp() if torch.version.hip is not None else RoPESM90Op() + + def _select_batch_row(inputs: DecodeAttentionInputs, batch_index: int) -> DecodeAttentionInputs: metadata = inputs.metadata cp_block_owners = ( diff --git a/setup.py b/setup.py index 79f882d9..94c0062a 100644 --- a/setup.py +++ b/setup.py @@ -1,313 +1,425 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -import importlib.util -import os -import warnings -from pathlib import Path - -from setuptools import find_packages, setup - - -def _load_envs_module(): - envs_path = Path(__file__).with_name("envs.py") - spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) - if spec is None or spec.loader is None: - raise RuntimeError(f"failed to load environment helpers from {envs_path}") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -envs = _load_envs_module() - - -def _load_torch_extension_tools(): - try: - import torch - except ModuleNotFoundError as exc: - if exc.name != "torch": - raise - return None, None, None - - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - - # CUDAExtension is also the supported extension entry point for ROCm - # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when - # torch.version.hip is set. - return torch, BuildExtension, CUDAExtension - - -def _native_extension_required() -> bool: - """Whether the caller explicitly requested a native extension build.""" - return ( - envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) - or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) - or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) - or envs.env_flag("FORCE_CUDA") - ) - - -def _cuda_define_from_env(name: str, macro: str) -> list[str]: - value = os.environ.get(name) - if value is None: - return [] - parsed = int(value) - if parsed <= 0: - raise ValueError(f"{name} must be positive, got {value!r}") - return [f"-D{macro}={parsed}"] - - -_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( - "-Xfatbin", - "-compress-all", - "-gencode", - "--generate-code", - "--expt-", - "-lineinfo", - "-allow-unsupported-compiler", - "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", -) -_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { - "-Xfatbin", - "-gencode", - "--generate-code", -} - - -def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: - """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" - filtered_flags = [] - skip_next = False - for flag in flags: - if skip_next: - skip_next = False - continue - if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: - skip_next = True - continue - if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): - continue - filtered_flags.append(flag) - return filtered_flags - - -def get_extensions(): - torch, _, CUDAExtension = _load_torch_extension_tools() - if torch is None: - message = ( - "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " - "CUDA/ROCm PyTorch build first, then run " - "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." - ) - if _native_extension_required(): - raise RuntimeError(message) - warnings.warn( - f"{message} Continuing with the pure-Python fallback because no native extension " - "was explicitly requested.", - RuntimeWarning, - stacklevel=2, - ) - return [] - - extensions = [] - torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") - torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] - if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": - torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") - is_rocm = getattr(torch.version, "hip", None) is not None - - # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, - # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also - # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add - # --offload-arch. Do not require a visible GPU when a ROCm target was - # explicitly selected. - no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() - if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: - raise RuntimeError( - "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " - "Set one or more ';'-separated targets, for example " - "PYTORCH_ROCM_ARCH='gfx942;gfx950'." - ) - - if is_rocm or torch.cuda.is_available(): - cuda_sources = [ - "csrc/ops.cpp", - "csrc/fused_logp_kernel.cu", - "csrc/deterministic_logp_kernel.cu", - "csrc/cuda/gemm/det_gemm_kernel.cu", - "csrc/cuda/rmsnorm.cu", - "csrc/cuda/activation.cu", - "csrc/cuda/attention/deterministic_attention.cu", - "csrc/cuda/distributed/deterministic_collective.cu", - ] - if not is_rocm: - # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). - # The ROCm dispatcher falls back to PyTorch SDPA for this operator. - cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") - - nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): - nvcc_flags.append("--use_fast_math") - if not is_rocm: - cc_major, cc_minor = torch.cuda.get_device_capability() - enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" - if not enable_sm90: - # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. - nvcc_flags.append( - f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" - ) - nvcc_flags.append("--expt-relaxed-constexpr") - nvcc_flags.append("--expt-extended-lambda") - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - "FUSED_LOGP_TWOPASS_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", - ) - ) - nvcc_flags.extend( - _cuda_define_from_env( - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", - ) - ) - if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): - nvcc_flags.append("-lineinfo") - if ( - not is_rocm - and os.name == "nt" - and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) - ): - nvcc_flags.append("-allow-unsupported-compiler") - nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") - - cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] - extra_link_args = list(torch_rpath) - if os.name != "nt": - # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). - extra_link_args.append("-lcuda") - - if not is_rocm: - sm90_srcs = [ - "csrc/cuda/fused_logp_sm90.cu", - "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob - "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp - "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build - # Single-card batch-invariant embedding/lm-head. - "csrc/cuda/embedding_lm_head_sm90.cu", - ] - enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) - present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] - if enable_sm90 and present_sm90: - tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant - cuda_sources.extend(present_sm90) +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import importlib.util +import os +import sysconfig +import warnings +from distutils.errors import CompileError +from distutils.spawn import find_executable +from pathlib import Path + +from setuptools import Extension, find_packages, setup + + +def _load_envs_module(): + envs_path = Path(__file__).with_name("envs.py") + spec = importlib.util.spec_from_file_location("_rl_kernel_envs", envs_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"failed to load environment helpers from {envs_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +envs = _load_envs_module() + + +def _load_torch_extension_tools(): + try: + import torch + except ModuleNotFoundError as exc: + if exc.name != "torch": + raise + return None, None, None + + from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + # CUDAExtension is also the supported extension entry point for ROCm + # PyTorch builds. BuildExtension dispatches .cu/.hip sources to hipcc when + # torch.version.hip is set. + return torch, BuildExtension, CUDAExtension + + +def _native_extension_required() -> bool: + """Whether the caller explicitly requested a native extension build.""" + return ( + envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) + or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) + or envs.env_flag("FORCE_CUDA") + ) + + +def _cuda_define_from_env(name: str, macro: str) -> list[str]: + value = os.environ.get(name) + if value is None: + return [] + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive, got {value!r}") + return [f"-D{macro}={parsed}"] + + +_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES = ( + "-Xfatbin", + "-compress-all", + "-gencode", + "--generate-code", + "--expt-", + "-lineinfo", + "-allow-unsupported-compiler", + "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH", +) +_ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE = { + "-Xfatbin", + "-gencode", + "--generate-code", +} + + +def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: + """Remove CUDA-only device compiler flags before BuildExtension calls hipcc.""" + filtered_flags = [] + skip_next = False + for flag in flags: + if skip_next: + skip_next = False + continue + if flag in _ROCM_NVCC_FLAGS_WITH_SEPARATE_VALUE: + skip_next = True + continue + if flag.startswith(_ROCM_UNSUPPORTED_NVCC_FLAG_PREFIXES): + continue + filtered_flags.append(flag) + return filtered_flags + + +def get_extensions(): + torch, _, CUDAExtension = _load_torch_extension_tools() + if torch is None: + message = ( + "PyTorch is unavailable, so rl_engine._C cannot be built. Install a matching " + "CUDA/ROCm PyTorch build first, then run " + "`RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e .`." + ) + if _native_extension_required(): + raise RuntimeError(message) + warnings.warn( + f"{message} Continuing with the pure-Python fallback because no native extension " + "was explicitly requested.", + RuntimeWarning, + stacklevel=2, + ) + return [] + + extensions = [] + torch_lib_dir = os.path.join(os.path.dirname(torch.__file__), "lib") + torch_rpath = ["-Wl,-rpath,$ORIGIN/../torch/lib"] + if os.environ.get("KERNEL_ALIGN_DEV_RPATH") == "1": + torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") + is_rocm = getattr(torch.version, "hip", None) is not None + + # CUDAExtension is intentionally used for both CUDA and ROCm. On ROCm, + # PyTorch's BuildExtension hipifies CUDA sources and invokes hipcc; it also + # consumes PYTORCH_ROCM_ARCH (one or more ';'-separated gfx targets) to add + # --offload-arch. Do not require a visible GPU when a ROCm target was + # explicitly selected. + no_rocm_arch = not os.environ.get("PYTORCH_ROCM_ARCH", "").strip() + if is_rocm and no_rocm_arch and torch.cuda.device_count() == 0: + raise RuntimeError( + "ROCm builds without a visible GPU require PYTORCH_ROCM_ARCH. " + "Set one or more ';'-separated targets, for example " + "PYTORCH_ROCM_ARCH='gfx942;gfx950'." + ) + + if is_rocm or torch.cuda.is_available(): + cuda_sources = [ + "csrc/ops.cpp", + "csrc/fused_logp_kernel.cu", + "csrc/deterministic_logp_kernel.cu", + "csrc/cuda/gemm/det_gemm_kernel.cu", + "csrc/cuda/rmsnorm.cu", + "csrc/cuda/activation.cu", + "csrc/cuda/attention/deterministic_attention.cu", + ] + if not is_rocm: + # prefix_shared_attention contains NVIDIA PTX (cp.async, ldmatrix, + # and mma.sync); the ROCm dispatcher falls back to PyTorch SDPA for + # it. The CUDA collective owns CUDA IPC handles and driver API calls. + cuda_sources.extend( + [ + "csrc/cuda/attention/prefix_shared_attention.cu", + "csrc/cuda/distributed/deterministic_collective.cu", + ] + ) + else: + # RCCL stays transport-only on ROCm; this HIP kernel performs the + # fixed balanced-tree arithmetic after AllGather. + cuda_sources.append("csrc/rocm/distributed/deterministic_collective.hip") + + nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] + if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + nvcc_flags.append("--use_fast_math") + if not is_rocm: + cc_major, cc_minor = torch.cuda.get_device_capability() + enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" + if not enable_sm90: + # SM90 build emits 90a below; mixing plain compute_90 breaks TMA ptxas. + nvcc_flags.append( + f"-gencode=arch=compute_{cc_major}{cc_minor},code=sm_{cc_major}{cc_minor}" + ) + nvcc_flags.append("--expt-relaxed-constexpr") + nvcc_flags.append("--expt-extended-lambda") + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + "FUSED_LOGP_TWOPASS_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + "FUSED_LOGP_ONLINE_SPARSE_LARGE_VOCAB_BLOCK_SIZE", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + "FUSED_LOGP_ONLINE_LARGE_ROW_BYTES_THRESHOLD", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_NUMERATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + "FUSED_LOGP_ONLINE_SPARSE_DENSITY_DENOMINATOR", + ) + ) + nvcc_flags.extend( + _cuda_define_from_env( + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + "FUSED_LOGP_ONLINE_MIN_BLOCKS_PER_SM", + ) + ) + if not is_rocm and envs.env_flag(envs.KERNEL_ALIGN_NCU_LINEINFO): + nvcc_flags.append("-lineinfo") + if ( + not is_rocm + and os.name == "nt" + and envs.env_flag(envs.KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC) + ): + nvcc_flags.append("-allow-unsupported-compiler") + nvcc_flags.append("-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH") + + platform_define = "-DKERNEL_ALIGN_WITH_ROCM" if is_rocm else "-DKERNEL_ALIGN_WITH_CUDA" + cxx_flags = ["-O3", "-std=c++17", platform_define] + extra_link_args = list(torch_rpath) + if not is_rocm and os.name != "nt": + # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). + extra_link_args.append("-lcuda") + + if not is_rocm: + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + "csrc/cuda/batch_invariant_logp_kernel_sm90.cu", # TMA batch-invariant logp + "csrc/cuda/rope_sm90.cu", # RoPE rotate-half apply, gated to SM90 build + # Single-card batch-invariant embedding/lm-head. + "csrc/cuda/embedding_lm_head_sm90.cu", + ] + enable_sm90 = envs.env_flag(envs.KERNEL_ALIGN_FORCE_SM90) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + if enable_sm90 and present_sm90: + tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant + cuda_sources.extend(present_sm90) nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") - cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - - # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp - # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in - # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. - enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" - if enable_det_gemm_sm90: - tma_arch = f"{cc_major}{cc_minor}a" - arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" - if arch_flag not in nvcc_flags: - nvcc_flags.append(arch_flag) - if "-lcuda" not in extra_link_args: - extra_link_args.append("-lcuda") - nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") - cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") - - if is_rocm: - nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) - - extensions.append( - CUDAExtension( - name="rl_engine._C", - sources=cuda_sources, - include_dirs=[], - extra_compile_args={ - "cxx": cxx_flags, - "nvcc": nvcc_flags, - }, - extra_link_args=extra_link_args, - ) - ) - - if _native_extension_required() and not extensions: - raise RuntimeError( - "rl_engine._C was requested but no CUDA/ROCm build environment is available. " - "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " - "PYTORCH_ROCM_ARCH to the target architecture." - ) - - return extensions - - -def get_cmdclass(): - _, BuildExtension, _ = _load_torch_extension_tools() - if BuildExtension is None: - return {} - return {"build_ext": BuildExtension} - - -setup( - name="rl-engine", - version="0.1.0", - packages=find_packages(include=["rl_engine", "rl_engine.*"]), - install_requires=[ - "torch>=2.4.1", - "tabulate", - "numpy", - "accelerate", - "transformers==5.13.1", - ], - ext_modules=get_extensions(), - cmdclass=get_cmdclass(), - extras_require={ - "cuda": ["flashinfer"], - "rocm": ["aiter"], - "vllm": ["vllm>=0.6.0"], - "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], - }, - entry_points={ - "console_scripts": [ - "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", - ], - }, - python_requires=">=3.10", - include_package_data=True, - zip_safe=False, -) + cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + + # det_gemm SM90 (mma.sync + TMA) path: independent of the fused_logp + # SM90 sources, which currently fail ptxas on CUDA 12.4 (shared::cta in + # the shared tma_utils.cuh). det_gemm uses its own gemm/det_gemm_tma.cuh. + enable_det_gemm_sm90 = os.environ.get("KERNEL_ALIGN_DET_GEMM_SM90") == "1" + if enable_det_gemm_sm90: + tma_arch = f"{cc_major}{cc_minor}a" + arch_flag = f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}" + if arch_flag not in nvcc_flags: + nvcc_flags.append(arch_flag) + if "-lcuda" not in extra_link_args: + extra_link_args.append("-lcuda") + nvcc_flags.append("-DRL_KERNEL_ENABLE_SM90") + cxx_flags.append("-DRL_KERNEL_ENABLE_SM90") + + if is_rocm: + nvcc_flags = _filter_rocm_incompatible_nvcc_flags(nvcc_flags) + + extensions.append( + CUDAExtension( + name="rl_engine._C", + sources=cuda_sources, + include_dirs=[], + extra_compile_args={ + "cxx": cxx_flags, + "nvcc": nvcc_flags, + }, + extra_link_args=extra_link_args, + ) + ) + + extensions.extend(_ascend_extensions()) + + if _native_extension_required() and not extensions: + raise RuntimeError( + "rl_engine._C was requested but no CUDA/ROCm build environment is available. " + "Use a matching GPU-enabled PyTorch build; for a GPU-less ROCm build, set " + "PYTORCH_ROCM_ARCH to the target architecture." + ) + + return extensions + + +def _ascend_extensions(): + """Ascend C (CANN) kernels, built with bisheng. Gated on KERNEL_ALIGN_FORCE_ASCEND=1. + + Follows the official torch_npu cpp_extension_asc pattern: .asc sources + (kernel + host + pybind) are compiled by the CANN bisheng compiler into a + single rl_engine._C_npu extension module. Requires CANN toolkit (bisheng on + PATH or ASCEND_HOME_PATH set) and torch_npu. + """ + if not envs.env_flag(envs.KERNEL_ALIGN_FORCE_ASCEND): + return [] + try: + import torch # noqa: F401 + import torch_npu # noqa: F401 + except ImportError as e: + raise RuntimeError( + "KERNEL_ALIGN_FORCE_ASCEND=1 requires torch and torch_npu to be installed" + ) from e + + asc_srcs = sorted(str(p) for p in Path("csrc/ascend").glob("*.asc")) + if not asc_srcs: + raise RuntimeError("KERNEL_ALIGN_FORCE_ASCEND=1 but no .asc sources under csrc/ascend/") + return [Extension(name="rl_engine._C_npu", sources=asc_srcs, language="asc")] + + +def _bisheng_compile_cmd(ext, ext_fullpath): + """Single-command bisheng build for an Ascend C extension (see op-plugin example).""" + import torch + import torch.utils.cpp_extension as cpp_extension + import torch_npu + + if find_executable("bisheng") is None: + raise RuntimeError( + "bisheng compiler not found on PATH; source the CANN toolkit environment first" + ) + + soc = os.environ.get(envs.KERNEL_ALIGN_ASCEND_ARCH, "dav-2201") # A2/A3; A5: dav-3510 + abi_value = "1" if torch._C._GLIBCXX_USE_CXX11_ABI else "0" + module_name = ext.name.rsplit(".", 1)[-1] + + torch_npu_dir = os.path.dirname(os.path.realpath(torch_npu.__file__)) + ascend_home = os.environ.get("ASCEND_HOME_PATH", "/usr/local/Ascend/ascend-toolkit/latest") + + include_dirs = [ + *cpp_extension.include_paths(), + sysconfig.get_config_var("INCLUDEPY"), + os.path.join(torch_npu_dir, "include"), + os.path.join(torch_npu_dir, "include", "third_party", "acl", "inc"), + os.path.join(ascend_home, "include"), + ] + lib_dirs = [ + sysconfig.get_config_var("LIBDIR"), + os.path.join(os.path.dirname(torch.__file__), "lib"), + os.path.join(torch_npu_dir, "lib"), + os.path.join(ascend_home, "lib64"), + ] + + cmd = [ + "bisheng", + "-x", + "asc", + f"--npu-arch={soc}", + "-shared", + "-fPIC", + "-std=c++17", + "-O2", + f"-D_GLIBCXX_USE_CXX11_ABI={abi_value}", + f"-DTORCH_EXTENSION_NAME={module_name}", + "-lascendcl", + "-ltorch_npu", + "-ltorch", + "-ltorch_cpu", + "-ltorch_python", + "-lc10", + *ext.sources, + "-o", + ext_fullpath, + ] + cmd += [f"-I{d}" for d in include_dirs if d] + cmd += [f"-L{d}" for d in lib_dirs if d] + return cmd + + +def get_cmdclass(): + _, BuildExtension, _ = _load_torch_extension_tools() + if BuildExtension is None: + return {} + + class AscendBuildExtension(BuildExtension): + """torch BuildExtension + bisheng path for language="asc" extensions.""" + + def build_extension(self, ext): + if getattr(ext, "language", None) != "asc": + super().build_extension(ext) + return + ext_fullpath = self.get_ext_fullpath(ext.name) + os.makedirs(os.path.dirname(ext_fullpath), exist_ok=True) + try: + self.spawn(_bisheng_compile_cmd(ext, ext_fullpath)) + except Exception as e: + raise CompileError(str(e)) from e + + return {"build_ext": AscendBuildExtension} + + +setup( + name="rl-engine", + version="0.1.0", + packages=find_packages(include=["rl_engine", "rl_engine.*"]), + install_requires=[ + "torch>=2.4.1", + "tabulate", + "numpy", + "accelerate", + "transformers==5.13.1", + ], + ext_modules=get_extensions(), + cmdclass=get_cmdclass(), + extras_require={ + "cuda": ["flashinfer"], + "rocm": ["aiter"], + "vllm": ["vllm>=0.6.0"], + "drift-viewer": ["Pillow>=10", "PySide6>=6.6"], + }, + entry_points={ + "console_scripts": [ + "rlk-drift-view=rl_engine.alignment.cross_config.drift_viewer:main", + ], + }, + python_requires=">=3.10", + include_package_data=True, + zip_safe=False, +) diff --git a/tests/conftest.py b/tests/conftest.py index 55935a4d..426bad48 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,6 +5,8 @@ import pathlib import sys +import pytest + def _add_windows_dll_dirs(): if sys.platform != "win32" or not hasattr(os, "add_dll_directory"): @@ -27,3 +29,45 @@ def _add_windows_dll_dirs(): _add_windows_dll_dirs() + + +def _is_rocm() -> bool: + """Whether this interpreter is running a ROCm PyTorch build. + + ``torch.cuda`` is the device API on ROCm too, so ``cuda.is_available()`` and + ``device_count()`` cannot distinguish the platforms; ``torch.version.hip`` + can. + """ + + try: + import torch + except ImportError: + return False + return torch.version.hip is not None + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "cuda_only: test depends on CUDA-exclusive functionality and is skipped on ROCm", + ) + + +def pytest_collection_modifyitems(config, items): + """Skip CUDA-exclusive tests on ROCm instead of failing them. + + Some kernels are compiled out of ROCm builds on purpose (the CUDA-IPC + deterministic collectives, for one), so their tests cannot pass there. + Because ROCm reports GPUs through the CUDA device API, a + ``cuda.device_count()`` guard does not exclude them and they fail instead of + skipping - which makes a ROCm run look broken rather than out of scope. + """ + + if not _is_rocm(): + return + skip_rocm = pytest.mark.skip( + reason="CUDA-exclusive functionality; not built for ROCm (torch.version.hip is set)" + ) + for item in items: + if "cuda_only" in item.keywords: + item.add_marker(skip_rocm) diff --git a/tests/distributed/test_det_gemm_simulated_tp.py b/tests/distributed/test_det_gemm_simulated_tp.py index c4824e85..26d49f9b 100644 --- a/tests/distributed/test_det_gemm_simulated_tp.py +++ b/tests/distributed/test_det_gemm_simulated_tp.py @@ -46,6 +46,13 @@ def _left_fold(parts: list[torch.Tensor]) -> torch.Tensor: return acc +# NOTE: these two encode the *CUDA* det_gemm kernel's K-reduction tree. On ROCm +# det_gemm dispatches to TritonDetGemmOp, whose K-tree differs, so a K-split sum +# is not bitwise equal to the unsplit GEMM there. ``get_device_capability()`` +# returns (9, 4) on gfx942, so the SM80 guard above does not exclude ROCm. +# Whether the Triton path *should* be K-split invariant is a separate question +# for the det_gemm owners; skipping here does not settle it. +@pytest.mark.cuda_only def test_simulated_tp2_matches_full(): # Two shards: AllReduce is a+b, and BF16 add is commutative. torch.manual_seed(8) @@ -69,6 +76,7 @@ def test_simulated_tp8_left_fold_does_not_match_full(): assert n_mismatch > 0, "TP=8 left-fold unexpectedly matched TP=1" +@pytest.mark.cuda_only def test_simulated_tp2_is_batch_invariant(): torch.manual_seed(9) k, n = 256, 64 diff --git a/tests/distributed/test_deterministic_all_gather.py b/tests/distributed/test_deterministic_all_gather.py index d6058dd3..7c9fe3e3 100644 --- a/tests/distributed/test_deterministic_all_gather.py +++ b/tests/distributed/test_deterministic_all_gather.py @@ -12,13 +12,17 @@ import torch.distributed as dist import torch.multiprocessing as mp -from rl_engine.distributed import DeterministicCollective +from rl_engine.distributed import create_deterministic_collective _MAX_WORLD_SIZE = 8 _TP_SIZES = (1, 2, 4, 8) _EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) pytestmark = [ + # The deterministic collectives are CUDA-IPC kernels, compiled out of ROCm + # builds on purpose. ROCm reports GPUs through the CUDA device API, so the + # device_count guard below does not exclude it. + pytest.mark.cuda_only, pytest.mark.skipif( _EXTERNAL_WORLD_SIZE != 1, reason="this cross-TP test owns its worker processes; run pytest directly", @@ -70,7 +74,7 @@ def _worker(rank: int, port: int) -> None: groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} for tp_size, group in groups.items(): if rank < tp_size: - with DeterministicCollective( + with create_deterministic_collective( group=group, device=device, max_size_bytes=1024 * 1024, @@ -95,6 +99,11 @@ def _worker(rank: int, port: int) -> None: returned = collective.all_gather(input, out=provided) assert returned is provided assert torch.equal(provided, expected) + + empty_input = torch.empty((0, 7), dtype=dtype, device=device) + empty_output = collective.all_gather(empty_input) + assert empty_output.shape == (0, 7) + assert empty_output.numel() == 0 dist.barrier() finally: dist.destroy_process_group() diff --git a/tests/distributed/test_deterministic_all_reduce.py b/tests/distributed/test_deterministic_all_reduce.py index eb406881..0da1fa2f 100644 --- a/tests/distributed/test_deterministic_all_reduce.py +++ b/tests/distributed/test_deterministic_all_reduce.py @@ -12,13 +12,17 @@ import torch.distributed as dist import torch.multiprocessing as mp -from rl_engine.distributed import DeterministicCollective +from rl_engine.distributed import create_deterministic_collective _MAX_WORLD_SIZE = 8 _TP_SIZES = (1, 2, 4, 8) _EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) pytestmark = [ + # The deterministic collectives are CUDA-IPC kernels, compiled out of ROCm + # builds on purpose. ROCm reports GPUs through the CUDA device API, so the + # device_count guard below does not exclude it. + pytest.mark.cuda_only, pytest.mark.skipif( _EXTERNAL_WORLD_SIZE != 1, reason="this cross-TP test owns its worker processes; run pytest directly", @@ -61,7 +65,7 @@ def _worker(rank: int, port: int) -> None: groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} for tp_size, group in groups.items(): if rank < tp_size: - with DeterministicCollective( + with create_deterministic_collective( group=group, device=device, max_size_bytes=1024 * 1024, @@ -96,6 +100,25 @@ def _worker(rank: int, port: int) -> None: returned = collective.all_reduce(inplace, out=inplace) assert returned is inplace assert torch.equal(inplace, expected) + + if dtype in (torch.float16, torch.bfloat16): + packed_input = torch.cat((input, input[:1])) + packed_expected = torch.cat((expected, expected[:1])) + packed_output = collective.all_reduce(packed_input) + assert torch.equal(packed_output, packed_expected) + + output_storage = torch.empty( + packed_expected.numel() + 1, + dtype=dtype, + device=device, + ) + misaligned_output = output_storage[1:].view_as(packed_expected) + returned = collective.all_reduce( + packed_input, + out=misaligned_output, + ) + assert returned is misaligned_output + assert torch.equal(misaligned_output, packed_expected) dist.barrier() finally: dist.destroy_process_group() diff --git a/tests/distributed/test_deterministic_reduce_scatter.py b/tests/distributed/test_deterministic_reduce_scatter.py index 3125f6c3..56512bec 100644 --- a/tests/distributed/test_deterministic_reduce_scatter.py +++ b/tests/distributed/test_deterministic_reduce_scatter.py @@ -12,13 +12,17 @@ import torch.distributed as dist import torch.multiprocessing as mp -from rl_engine.distributed import DeterministicCollective +from rl_engine.distributed import create_deterministic_collective _MAX_WORLD_SIZE = 8 _TP_SIZES = (1, 2, 4, 8) _EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) pytestmark = [ + # The deterministic collectives are CUDA-IPC kernels, compiled out of ROCm + # builds on purpose. ROCm reports GPUs through the CUDA device API, so the + # device_count guard below does not exclude it. + pytest.mark.cuda_only, pytest.mark.skipif( _EXTERNAL_WORLD_SIZE != 1, reason="this cross-TP test owns its worker processes; run pytest directly", @@ -61,7 +65,7 @@ def _worker(rank: int, port: int) -> None: groups = {tp_size: dist.new_group(ranks=list(range(tp_size))) for tp_size in _TP_SIZES} for tp_size, group in groups.items(): if rank < tp_size: - with DeterministicCollective( + with create_deterministic_collective( group=group, device=device, max_size_bytes=1024 * 1024, @@ -98,6 +102,49 @@ def _worker(rank: int, port: int) -> None: returned = collective.reduce_scatter(input, out=provided) assert returned is provided assert torch.equal(provided, expected) + + if dtype in (torch.float16, torch.bfloat16): + output_storage = torch.empty( + expected.numel() + 1, + dtype=dtype, + device=device, + ) + misaligned_output = output_storage[1:].view_as(expected) + returned = collective.reduce_scatter( + input, + out=misaligned_output, + ) + assert returned is misaligned_output + assert torch.equal(misaligned_output, expected) + + other_generator = torch.Generator().manual_seed(20260817) + other_leaves_tensor = torch.randn( + _MAX_WORLD_SIZE, + _MAX_WORLD_SIZE * 17, + 19, + dtype=torch.float32, + generator=other_generator, + ).to(device=device, dtype=dtype) + other_leaves = list(other_leaves_tensor.unbind()) + other_input = _fixed_tree_reference( + other_leaves[start : start + leaves_per_rank] + ) + other_reduced = _fixed_tree_reference(other_leaves) + other_expected = other_reduced.chunk(tp_size, dim=0)[group_rank] + many_outs = (torch.empty_like(expected), torch.empty_like(other_expected)) + many_returned = collective.reduce_scatter_many( + (input, other_input), + outs=many_outs, + ) + assert many_returned[0] is many_outs[0] + assert many_returned[1] is many_outs[1] + assert torch.equal(many_returned[0], expected) + assert torch.equal(many_returned[1], other_expected) + many_baseline = tuple(value.clone() for value in many_returned) + for _ in range(3): + many_repeated = collective.reduce_scatter_many((input, other_input)) + assert torch.equal(many_repeated[0], many_baseline[0]) + assert torch.equal(many_repeated[1], many_baseline[1]) dist.barrier() finally: dist.destroy_process_group() diff --git a/tests/distributed/test_rocm_attention_transport.py b/tests/distributed/test_rocm_attention_transport.py new file mode 100644 index 00000000..d7ce9564 --- /dev/null +++ b/tests/distributed/test_rocm_attention_transport.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Any + +import pytest +import torch + +from rl_engine.distributed import collectives +from rl_engine.kernels.ops.cuda.attention.cp_comm import ( + AttentionCPBlockMetadata, + AttentionCPCommunicationPlan, + AttentionCPCommunicationUnavailable, + AttentionParallelSpec, + CUDAAGRSAttentionCPCommunication, + RCCLAGRSAttentionCPCommunication, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + + +class _FakeDist: + class group: + WORLD = "world-group" + + +def _plan(*, backend: str = "rccl_ag_rs") -> AttentionCPCommunicationPlan: + return AttentionCPCommunicationPlan( + parallel=AttentionParallelSpec(tp_world_size=2, cp_world_size=2), + backend=backend, # type: ignore[arg-type] + status="implemented", + expected_blocks=( + AttentionCPBlockMetadata(0, 0, 2, 0, 0), + AttentionCPBlockMetadata(1, 2, 4, 1, 0), + ), + expected_kv_token_range=(0, 4), + query_token_ranges=((0, 1), (1, 2)), + ) + + +class _FakeRCCLTransport: + world_size = 2 + + def __init__(self) -> None: + self.gather_calls = 0 + self.scatter_calls = 0 + + def all_gather(self, tensor: torch.Tensor) -> torch.Tensor: + self.gather_calls += 1 + return torch.cat((tensor, tensor), dim=0) + + def scatter(self, tensor: torch.Tensor) -> torch.Tensor: + self.scatter_calls += 1 + return tensor.chunk(self.world_size, dim=0)[0].contiguous() + + +def test_rccl_plan_reports_transport_only_runtime() -> None: + provenance = _plan().provenance() + + assert provenance["cp_comm_runtime"] == "rccl" + assert provenance["cp_comm_attention_numeric_reduction"] is False + + +def test_rccl_adapter_uses_root_scatter_and_reports_no_fusion( + monkeypatch: pytest.MonkeyPatch, +) -> None: + transport = _FakeRCCLTransport() + communication = RCCLAGRSAttentionCPCommunication(collective=transport) + monkeypatch.setattr(torch.version, "hip", "test", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + plan = _plan() + + local_q = torch.zeros(1, 2, 1, 4, dtype=torch.bfloat16) + assert communication.all_gather_query(local_q, plan).shape == (1, 2, 2, 4) + + full_out = torch.zeros(1, 2, 2, 4, dtype=torch.bfloat16) + full_lse = torch.zeros(1, 2, 2, dtype=torch.float32) + shard = communication.reduce_scatter_strict_result(full_out, full_lse, plan) + + assert shard.out.shape == (1, 2, 1, 4) + assert shard.lse.shape == (1, 2, 1) + assert transport.gather_calls == 1 + assert transport.scatter_calls == 2 + assert communication.transport_only is True + assert communication.supports_async_overlap is False + assert communication.supports_compute_communication_fusion is False + + +def test_rccl_adapter_fails_closed_outside_rocm(monkeypatch: pytest.MonkeyPatch) -> None: + communication = RCCLAGRSAttentionCPCommunication(collective=_FakeRCCLTransport()) + monkeypatch.setattr(torch.version, "hip", None, raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + with pytest.raises(AttentionCPCommunicationUnavailable, match="ROCm device"): + communication.all_gather_query(torch.zeros(1, 2, 1, 4), _plan()) + + +def test_rccl_adapter_rejects_cuda_plan(monkeypatch: pytest.MonkeyPatch) -> None: + communication = RCCLAGRSAttentionCPCommunication(collective=_FakeRCCLTransport()) + monkeypatch.setattr(torch.version, "hip", "test", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + with pytest.raises(AttentionCPCommunicationUnavailable, match="rccl_ag_rs plan"): + communication.all_gather_query(torch.zeros(1, 2, 1, 4), _plan(backend="cuda_ag_rs")) + + +def test_rccl_adapter_shares_the_cuda_collective_resolution() -> None: + """ROCm must not own a second transport implementation. + + CUDA and ROCm run the same balanced rank tree only because both resolve + their collective through ``collective_for_group``. A ROCm-side override + would let the two reduction orders drift apart silently, so pin that the + two adapters share one implementation. + """ + + assert ( + RCCLAGRSAttentionCPCommunication._get_collective + is CUDAAGRSAttentionCPCommunication._get_collective + ) + + +def test_rccl_adapter_resolves_the_shared_deterministic_collective( + monkeypatch: pytest.MonkeyPatch, +) -> None: + resolved = _FakeRCCLTransport() + calls: list[Any] = [] + + def _fake_collective_for_group(*, group: Any, device: Any) -> Any: + calls.append((group, device)) + return resolved + + monkeypatch.setattr(collectives, "collective_for_group", _fake_collective_for_group) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + + communication = RCCLAGRSAttentionCPCommunication(process_group="cp-group") + monkeypatch.setattr(communication, "_dist", lambda: _FakeDist()) + + assert communication._get_collective(_plan()) is resolved + assert calls == [("cp-group", torch.device("cuda", 0))] + + +def test_rccl_adapter_world_sizes_match_the_shared_collective() -> None: + capability = KernelRegistry()._attention_capabilities[OpBackend.ROCM_STRICT_ATTENTION] + + assert tuple(capability.cp_world_sizes) == collectives._SUPPORTED_WORLD_SIZES diff --git a/tests/distributed/test_rocm_strict_attention_cp.py b/tests/distributed/test_rocm_strict_attention_cp.py new file mode 100644 index 00000000..f7984ae6 --- /dev/null +++ b/tests/distributed/test_rocm_strict_attention_cp.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""End-to-end CP coverage for the strict ROCm attention provider. + +Acceptance is bitwise against a CP=1 run of the same core on the same logical +sequence. CP performs no arithmetic of its own here: the runtime all-gathers +Q/K/V, runs the core once over the full sequence, and scatters the root's +authoritative ``(out, lse)`` back to each rank's query range. Anything other +than bit equality means the CP path introduced a second merge order. + +These run through ``attention_provider`` rather than the transport directly, so +they also pin that CP is reachable from the production dispatch path. +""" + +from __future__ import annotations + +import os +import socket +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from rl_engine.kernels.registry import _rocm_strict_attention_available + +_MAX_WORLD_SIZE = 8 +_CP_SIZES = (2, 4, 8) +_EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + +# Qwen3-8B dense head layout at TP=1. +_GLOBAL_Q_HEADS = 32 +_GLOBAL_KV_HEADS = 8 +_HEAD_DIM = 128 +_GLOBAL_SEQ = 256 + +pytestmark = [ + pytest.mark.skipif( + _EXTERNAL_WORLD_SIZE != 1, + reason="this cross-CP test owns its worker processes; run pytest directly", + ), + pytest.mark.skipif( + torch.cuda.device_count() < _MAX_WORLD_SIZE, + reason="requires eight visible ROCm GPUs", + ), + pytest.mark.skipif( + not _rocm_strict_attention_available(), + reason="strict ROCm attention requires a ROCm device with aiter.ops.mha", + ), +] + + +def _global_qkv(device: torch.device) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Identical global Q/K/V on every rank, generated on CPU for bit equality.""" + + def _make(shape: tuple[int, ...], seed: int) -> torch.Tensor: + generator = torch.Generator(device="cpu").manual_seed(seed) + value = torch.randn(*shape, generator=generator, dtype=torch.float32) * 0.02 + return value.to(device=device, dtype=torch.bfloat16) + + q = _make((1, _GLOBAL_Q_HEADS, _GLOBAL_SEQ, _HEAD_DIM), 11) + k = _make((1, _GLOBAL_KV_HEADS, _GLOBAL_SEQ, _HEAD_DIM), 22) + v = _make((1, _GLOBAL_KV_HEADS, _GLOBAL_SEQ, _HEAD_DIM), 33) + return q, k, v + + +def _request( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cp_world_size: int, + cp_rank: int, + cp_layout: str, + token_start: int, + cp_group: object | None, +) -> SimpleNamespace: + kv_len = k.shape[2] + positions = torch.arange( + token_start, + token_start + kv_len, + device=q.device, + dtype=torch.int64, + ) + return SimpleNamespace( + query=q, + key=k, + value=v, + key_padding_mask=None, + # TP=1 must be stated explicitly: with torch.distributed initialized, + # a None TP group resolves to the global group, i.e. TP=8 here. + tensor_parallel_group=SimpleNamespace(rank=lambda: 0, size=lambda: 1), + context_parallel=SimpleNamespace( + world_size=cp_world_size, + rank=cp_rank, + layout=cp_layout, + ), + context_parallel_group=cp_group, + metadata={ + "global_q_heads": _GLOBAL_Q_HEADS, + "global_kv_heads": _GLOBAL_KV_HEADS, + "tp_rank": 0, + "tp_world_size": 1, + "attention_mode": "prefill", + "role": "train", + "causal": True, + "key_position_ids": positions, + }, + ) + + +def _check_one_cp_degree(group: object, cp_size: int, rank: int, device: torch.device) -> None: + from rl_engine.integrations.vime.attention import attention_provider + + q_global, k_global, v_global = _global_qkv(device) + + # CP=1 on the whole logical sequence is the acceptance reference. + reference = attention_provider( + _request( + q_global, + k_global, + v_global, + cp_world_size=1, + cp_rank=0, + cp_layout="single", + token_start=0, + cp_group=None, + ) + ) + + local_seq = _GLOBAL_SEQ // cp_size + lo, hi = rank * local_seq, (rank + 1) * local_seq + result = attention_provider( + _request( + q_global[:, :, lo:hi].contiguous(), + k_global[:, :, lo:hi].contiguous(), + v_global[:, :, lo:hi].contiguous(), + cp_world_size=cp_size, + cp_rank=rank, + cp_layout="allgather", + token_start=lo, + cp_group=group, + ) + ) + + assert result.out.shape == (1, _GLOBAL_Q_HEADS, local_seq, _HEAD_DIM) + assert result.lse.shape == (1, _GLOBAL_Q_HEADS, local_seq) + + expected_out = reference.out[:, :, lo:hi] + expected_lse = reference.lse[:, :, lo:hi] + out_mismatch = int((result.out != expected_out).sum().item()) + lse_mismatch = int((result.lse != expected_lse).sum().item()) + assert out_mismatch == 0, f"CP={cp_size} rank={rank}: {out_mismatch} out elements differ" + assert lse_mismatch == 0, f"CP={cp_size} rank={rank}: {lse_mismatch} lse elements differ" + + provenance = result.provenance + assert provenance["cp_row_ownership"]["cp_is_merge_axis"] is True + assert provenance["cp_row_ownership"]["cp_merge_owner"] == "rccl_ag_rs" + assert provenance["cp_row_ownership"]["cp_world_size"] == cp_size + # The core still runs once per (batch row, KV group), now over the gathered + # global sequence rather than this rank's shard. + assert provenance["execution"]["core_launches"] == _GLOBAL_KV_HEADS + + # Repeating the CP call must reproduce its own bits. + repeated = attention_provider( + _request( + q_global[:, :, lo:hi].contiguous(), + k_global[:, :, lo:hi].contiguous(), + v_global[:, :, lo:hi].contiguous(), + cp_world_size=cp_size, + cp_rank=rank, + cp_layout="allgather", + token_start=lo, + cp_group=group, + ) + ) + assert torch.equal(repeated.out, result.out) + assert torch.equal(repeated.lse, result.lse) + + +def _worker(rank: int, port: int) -> None: + torch.cuda.set_device(rank) + device = torch.device("cuda", rank) + dist.init_process_group( + backend="nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=_MAX_WORLD_SIZE, + timeout=timedelta(minutes=10), + ) + try: + for cp_size in _CP_SIZES: + group = dist.new_group(ranks=list(range(cp_size))) + if rank < cp_size: + _check_one_cp_degree(group, cp_size, rank, device) + dist.barrier() + finally: + dist.destroy_process_group() + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def test_strict_rocm_attention_cp_is_bitwise_against_cp1() -> None: + mp.spawn( + _worker, + args=(_find_free_port(),), + nprocs=_MAX_WORLD_SIZE, + join=True, + ) diff --git a/tests/distributed/test_transport_deterministic_collective.py b/tests/distributed/test_transport_deterministic_collective.py new file mode 100644 index 00000000..5c9ffb84 --- /dev/null +++ b/tests/distributed/test_transport_deterministic_collective.py @@ -0,0 +1,500 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Any + +import pytest +import torch + +import rl_engine.distributed as distributed +import rl_engine.distributed.collectives as collectives +from rl_engine.distributed import ( + RCCLDeterministicCollective, + TorchDistributedDeterministicCollective, +) + + +class _FakeDistributed: + """Single-process model of rank-ordered AllGather transport.""" + + def __init__( + self, + peer_inputs: list[torch.Tensor], + *, + rank: int = 0, + backend: str = "gloo", + peer_signatures: list[tuple[Any, ...]] | None = None, + peer_capacities: list[int] | None = None, + ) -> None: + self.peer_inputs = peer_inputs + self.rank = rank + self.backend = backend + self.peer_signatures = peer_signatures + self.peer_capacities = peer_capacities + self.into_tensor_calls = 0 + self.list_transport_calls = 0 + self.object_gather_calls = 0 + self.last_transport_output: torch.Tensor | None = None + + @property + def tensor_transport_calls(self) -> int: + return self.into_tensor_calls + self.list_transport_calls + + @staticmethod + def is_available() -> bool: + return True + + @staticmethod + def is_initialized() -> bool: + return True + + def get_rank(self, *, group: Any) -> int: + return self.rank + + def get_world_size(self, *, group: Any) -> int: + return len(self.peer_inputs) + + def get_backend(self, group: Any) -> str: + return self.backend + + def all_gather_object(self, output: list[Any], value: Any, *, group: Any) -> None: + self.object_gather_calls += 1 + if isinstance(value, int): + values = self.peer_capacities or [value] * len(self.peer_inputs) + else: + values = self.peer_signatures or [value] * len(self.peer_inputs) + output[:] = values + + def all_gather_into_tensor( + self, + output: torch.Tensor, + input: torch.Tensor, + *, + group: Any, + ) -> None: + self.into_tensor_calls += 1 + self.last_transport_output = output + gathered = torch.cat([peer.reshape(-1) for peer in self.peer_inputs]) + output.copy_(gathered) + + def all_gather( + self, + output: list[torch.Tensor], + input: torch.Tensor, + *, + group: Any, + ) -> None: + self.list_transport_calls += 1 + self.last_transport_output = output[0]._base + for destination, peer in zip(output, self.peer_inputs, strict=True): + destination.copy_(peer.reshape(-1)) + + +def test_public_exports_use_canonical_collectives_module() -> None: + assert distributed.DeterministicCollective is collectives.DeterministicCollective + assert distributed.RCCLDeterministicCollective is collectives.RCCLDeterministicCollective + assert ( + distributed.TorchDistributedDeterministicCollective + is collectives.TorchDistributedDeterministicCollective + ) + assert ( + distributed.create_deterministic_collective is collectives.create_deterministic_collective + ) + + +def _make_collective( + monkeypatch: pytest.MonkeyPatch, + peer_inputs: list[torch.Tensor], + *, + rank: int = 0, + backend: str = "gloo", + peer_signatures: list[tuple[Any, ...]] | None = None, + peer_capacities: list[int] | None = None, + max_size_bytes: int = 1024, +) -> tuple[TorchDistributedDeterministicCollective, _FakeDistributed]: + fake_dist = _FakeDistributed( + peer_inputs, + rank=rank, + backend=backend, + peer_signatures=peer_signatures, + peer_capacities=peer_capacities, + ) + monkeypatch.setattr(collectives, "dist", fake_dist) + collective = TorchDistributedDeterministicCollective( + group=object(), + device="cpu", + max_size_bytes=max_size_bytes, + ) + return collective, fake_dist + + +@pytest.mark.parametrize("world_size", [1, 2, 4, 8]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_all_reduce_supports_fixed_world_sizes_and_dtypes( + monkeypatch: pytest.MonkeyPatch, + world_size: int, + dtype: torch.dtype, +) -> None: + peers = [torch.full((2, 3), rank + 1, dtype=dtype) for rank in range(world_size)] + collective, fake_dist = _make_collective(monkeypatch, peers) + + provided = torch.empty_like(peers[0]) + returned = collective.all_reduce(peers[0], out=provided) + + assert returned is provided + assert torch.equal(provided, torch.full_like(provided, world_size * (world_size + 1) // 2)) + assert fake_dist.tensor_transport_calls == (0 if world_size == 1 else 1) + + +def test_all_reduce_uses_balanced_not_rank_ordered_left_fold( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Balanced: (1e20 + 1) + (-1e20 + 1) == 0 in FP32. + # Left fold: ((1e20 + 1) + -1e20) + 1 == 1 in FP32. + peers = [torch.tensor([value], dtype=torch.float32) for value in (1.0e20, 1.0, -1.0e20, 1.0)] + collective, _ = _make_collective(monkeypatch, peers) + + output = collective.all_reduce(peers[0]) + + assert torch.equal(output, torch.zeros_like(output)) + + +def test_all_reduce_stages_before_writing_in_place(monkeypatch: pytest.MonkeyPatch) -> None: + peers = [torch.full((2, 3), rank + 1, dtype=torch.float32) for rank in range(4)] + collective, _ = _make_collective(monkeypatch, peers) + local = peers[0].clone() + + returned = collective.all_reduce(local, out=local) + + assert returned is local + assert torch.equal(local, torch.full_like(local, 10)) + + +def test_all_gather_is_rank_ordered_and_transport_only( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.arange(rank * 6, (rank + 1) * 6).reshape(2, 3) for rank in range(4)] + collective, fake_dist = _make_collective(monkeypatch, peers, rank=2) + + output = collective.all_gather(peers[2]) + + assert torch.equal(output, torch.cat(peers, dim=0)) + assert fake_dist.tensor_transport_calls == 1 + + +def test_nccl_backend_uses_all_gather_into_tensor_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((2, 3), rank) for rank in range(2)] + collective, fake_dist = _make_collective(monkeypatch, peers, backend="nccl") + + output = collective.all_gather(peers[0]) + + assert torch.equal(output, torch.cat(peers, dim=0)) + assert fake_dist.into_tensor_calls == 1 + assert fake_dist.list_transport_calls == 0 + + +def test_all_gather_writes_directly_to_provided_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((2, 3), rank) for rank in range(2)] + collective, fake_dist = _make_collective(monkeypatch, peers, backend="nccl") + provided = torch.empty(4, 3, dtype=peers[0].dtype) + + returned = collective.all_gather(peers[0], out=provided) + + assert returned is provided + assert fake_dist.last_transport_output is not None + assert fake_dist.last_transport_output.data_ptr() == provided.data_ptr() + assert collective._workspace is None + + +def test_reduction_workspace_grows_once_and_is_reused( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((2, 3), rank + 1, dtype=torch.float32) for rank in range(2)] + collective, _ = _make_collective(monkeypatch, peers, backend="nccl") + + collective.all_reduce(peers[0]) + assert collective._workspace is not None + assert collective.workspace_size_bytes == sum(peer.numel() for peer in peers) * 4 + first_pointer = collective._workspace.data_ptr() + collective.all_reduce(peers[0]) + + assert collective._workspace.data_ptr() == first_pointer + collective.close() + assert collective._workspace is None + assert collective.workspace_size_bytes == 0 + + +def test_matching_signature_is_validated_once_per_hot_path( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((2, 3), rank + 1, dtype=torch.float32) for rank in range(2)] + collective, fake_dist = _make_collective(monkeypatch, peers, backend="nccl") + constructor_object_gathers = fake_dist.object_gather_calls + + collective.all_reduce(peers[0]) + collective.all_reduce(peers[0]) + + assert fake_dist.object_gather_calls == constructor_object_gathers + 1 + + +def test_latest_collective_api_can_skip_signature_handshakes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((4, 2), rank + 1, dtype=torch.float32) for rank in range(2)] + collective, fake_dist = _make_collective(monkeypatch, peers, backend="nccl") + constructor_object_gathers = fake_dist.object_gather_calls + + collective.all_reduce(peers[0], validate_signature=False) + collective.all_gather(peers[0], validate_signature=False) + collective.reduce_scatter(peers[0], validate_signature=False) + gathered = collective.all_gather_many( + (peers[0], peers[0]), + validate_signature=False, + ) + fake_dist.peer_inputs = [torch.cat((peer, peer), dim=-1) for peer in peers] + scattered = collective.reduce_scatter_many( + (peers[0], peers[0]), + validate_signature=False, + ) + + assert len(gathered) == 2 + assert len(scattered) == 2 + assert fake_dist.object_gather_calls == constructor_object_gathers + + +def test_reduce_scatter_reduces_then_selects_local_leading_shard( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.full((8, 3), rank + 1, dtype=torch.bfloat16) for rank in range(4)] + collective, _ = _make_collective(monkeypatch, peers, rank=2) + + output = collective.reduce_scatter(peers[2]) + + expected_full = torch.full_like(peers[0], 10) + assert torch.equal(output, expected_full.chunk(4, dim=0)[2]) + + +def test_reduce_scatter_many_packs_lanes_and_transports_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Each lane must retain its own balanced tree. A left fold would produce + # one for lane 0 and two for lane 1, while the fixed tree produces zero for + # both lanes in FP32. + lane_values = [ + ( + 1.0e20, + 1.0, + -1.0e20, + 1.0, + ), + ( + 1.0e20, + 2.0, + -1.0e20, + 2.0, + ), + ] + # Give each rank distinct values while preserving the cancellation pattern + # in every row. The fake transport returns these packed rank inputs. + lane_peers = [ + ( + torch.full((8, 1), lane_values[0][rank], dtype=torch.float32), + torch.full((8, 1), lane_values[1][rank], dtype=torch.float32), + ) + for rank in range(4) + ] + packed_peers = [torch.cat(lanes, dim=-1) for lanes in lane_peers] + collective, fake_dist = _make_collective(monkeypatch, packed_peers, rank=2) + + local_lanes = lane_peers[2] + outputs = (torch.empty(2, 1), torch.empty(2, 1)) + returned = collective.reduce_scatter_many(local_lanes, outs=outputs) + + assert returned[0] is outputs[0] + assert returned[1] is outputs[1] + assert torch.equal(outputs[0], torch.zeros_like(outputs[0])) + assert torch.equal(outputs[1], torch.zeros_like(outputs[1])) + assert fake_dist.tensor_transport_calls == 1 + + +def test_reduce_scatter_many_rejects_oversized_packed_input( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.ones(4, 2, dtype=torch.float32) for _ in range(2)] + collective, fake_dist = _make_collective( + monkeypatch, + peers, + max_size_bytes=32, + ) + monkeypatch.setattr(collectives, "_PACKED_REDUCE_SCATTER_MAX_BYTES", 1024) + + with pytest.raises(ValueError, match="packed input requires"): + collective.reduce_scatter_many((peers[0], peers[0])) + assert fake_dist.tensor_transport_calls == 0 + + +def test_reduce_scatter_many_uses_separate_calls_for_large_payloads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.ones(4, 2, dtype=torch.float32) for _ in range(2)] + collective, fake_dist = _make_collective(monkeypatch, peers) + monkeypatch.setattr(collectives, "_PACKED_REDUCE_SCATTER_MAX_BYTES", 1) + + outputs = collective.reduce_scatter_many((peers[0], peers[0])) + + assert len(outputs) == 2 + assert all(torch.equal(output, torch.full_like(output, 2)) for output in outputs) + assert fake_dist.tensor_transport_calls == 2 + + +def test_matching_signature_is_checked_before_tensor_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.ones(2, 3), torch.ones(2, 3)] + signatures = [ + ("all_reduce", (2, 3), "torch.float32", 6), + ("reduce_scatter", (2, 3), "torch.float32", 6), + ] + collective, fake_dist = _make_collective( + monkeypatch, + peers, + peer_signatures=signatures, + ) + + with pytest.raises(ValueError, match="matching shapes and dtypes"): + collective.all_reduce(peers[0]) + assert fake_dist.tensor_transport_calls == 0 + + +def test_capacity_must_match_on_every_rank(monkeypatch: pytest.MonkeyPatch) -> None: + peers = [torch.ones(2, 3), torch.ones(2, 3)] + + with pytest.raises(ValueError, match="same max_size_bytes"): + _make_collective( + monkeypatch, + peers, + max_size_bytes=1024, + peer_capacities=[1024, 2048], + ) + + +def test_validation_fails_closed_before_transport(monkeypatch: pytest.MonkeyPatch) -> None: + peers = [torch.ones(4, 2), torch.ones(4, 2)] + collective, fake_dist = _make_collective(monkeypatch, peers, max_size_bytes=31) + + with pytest.raises(TypeError, match="float32, float16, and bfloat16"): + collective.all_reduce(torch.ones(1, dtype=torch.int32)) + with pytest.raises(ValueError, match="max_size_bytes"): + collective.all_reduce(peers[0]) + with pytest.raises(ValueError, match=r"input.size\(0\).+divisible"): + collective.reduce_scatter(torch.ones(3, 2)) + with pytest.raises(ValueError, match="at least one dimension"): + collective.all_gather(torch.tensor(1.0)) + assert fake_dist.tensor_transport_calls == 0 + + +def test_lifecycle_is_idempotent_and_context_manager_closes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + peers = [torch.ones(2, 3)] + collective, _ = _make_collective(monkeypatch, peers) + + with collective as entered: + assert entered is collective + assert not collective.closed + assert torch.equal(collective.all_reduce(peers[0]), peers[0]) + + assert collective.closed + collective.close() + with pytest.raises(RuntimeError, match="closed"): + collective.all_reduce(peers[0]) + + +@pytest.mark.parametrize("world_size", [3, 16]) +def test_unsupported_world_size_is_rejected( + monkeypatch: pytest.MonkeyPatch, + world_size: int, +) -> None: + fake_dist = _FakeDistributed([torch.ones(1)] * world_size) + monkeypatch.setattr(collectives, "dist", fake_dist) + + with pytest.raises(ValueError, match="world_size in"): + TorchDistributedDeterministicCollective(group=object(), device="cpu") + + +def test_rccl_class_requires_rocm_build(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(collectives.torch.version, "hip", None, raising=False) + + with pytest.raises(RuntimeError, match="ROCm PyTorch build"): + RCCLDeterministicCollective(group=object(), device="cuda:0") + + +def test_rccl_class_requires_nccl_process_group(monkeypatch: pytest.MonkeyPatch) -> None: + fake_dist = _FakeDistributed([torch.ones(1)], backend="gloo") + monkeypatch.setattr(collectives, "dist", fake_dist) + monkeypatch.setattr(collectives.torch.version, "hip", "6.3", raising=False) + monkeypatch.setattr(collectives.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(collectives.torch.cuda, "current_device", lambda: 0) + + with pytest.raises(RuntimeError, match="NCCL process-group API"): + RCCLDeterministicCollective(group=object(), device="cuda:0") + + +def test_rccl_class_rejects_cpu_before_process_group_exchange( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(collectives.torch.version, "hip", "6.3", raising=False) + monkeypatch.setattr(collectives.torch.cuda, "is_available", lambda: True) + + with pytest.raises(ValueError, match="ROCm device"): + RCCLDeterministicCollective(group=object(), device="cpu") + + +def test_factory_dispatches_rocm_to_rccl(monkeypatch: pytest.MonkeyPatch) -> None: + sentinel = object() + calls: list[dict[str, Any]] = [] + + def fake_rccl(**kwargs: Any) -> object: + calls.append(kwargs) + return sentinel + + monkeypatch.setattr(collectives.torch.version, "hip", "6.3", raising=False) + monkeypatch.setattr(collectives, "RCCLDeterministicCollective", fake_rccl) + + group = object() + result = collectives.create_deterministic_collective( + group=group, + device="cuda:3", + max_size_bytes=1234, + ) + + assert result is sentinel + assert calls == [{"group": group, "device": "cuda:3", "max_size_bytes": 1234}] + + +def test_factory_preserves_existing_cuda_collective(monkeypatch: pytest.MonkeyPatch) -> None: + sentinel = object() + calls: list[dict[str, Any]] = [] + + def fake_cuda(**kwargs: Any) -> object: + calls.append(kwargs) + return sentinel + + monkeypatch.setattr(collectives.torch.version, "hip", None, raising=False) + monkeypatch.setattr(collectives, "DeterministicCollective", fake_cuda) + + group = object() + result = collectives.create_deterministic_collective( + group=group, + device="cuda:1", + max_size_bytes=4321, + ) + + assert result is sentinel + assert calls == [{"group": group, "device": "cuda:1", "max_size_bytes": 4321}] diff --git a/tests/test_attention_correctness.py b/tests/test_attention_correctness.py index d68ad3cb..a77e195c 100644 --- a/tests/test_attention_correctness.py +++ b/tests/test_attention_correctness.py @@ -8,6 +8,15 @@ import torch import torch.nn.functional as F +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, +) +from rl_engine.kernels.ops.rocm.attention.flash_attn import ( + StrictRocmAiterCKAttentionCore, + StrictRocmAttentionUnavailable, +) + try: from torch.nn.attention import SDPBackend, sdpa_kernel except ImportError: @@ -440,3 +449,143 @@ def test_native_attention_rejects_invalid_gqa_head_ratio(): with pytest.raises(ValueError, match="q heads must be divisible"): NativeAttentionOp()(q, k, v) + + +def test_strict_rocm_aiter_ck_core_fixes_forward_and_backward_contract(monkeypatch): + calls = [] + + def fake_fwd( + q, + k, + v, + dropout_p, + softmax_scale, + causal, + window_left, + window_right, + sink_size, + return_lse, + return_dropout_mask, + ): + calls.append( + ( + "forward", + dropout_p, + softmax_scale, + causal, + window_left, + window_right, + sink_size, + return_lse, + return_dropout_mask, + ) + ) + return ( + q.clone(), + torch.zeros(q.size(0), q.size(2), q.size(1), dtype=torch.float32), + torch.empty(0), + torch.zeros(2, dtype=torch.int64), + ) + + def fake_bwd( + dout, + q, + k, + v, + out, + lse, + dropout_p, + softmax_scale, + causal, + window_left, + window_right, + deterministic, + **kwargs, + ): + calls.append( + ( + "backward", + dropout_p, + softmax_scale, + causal, + window_left, + window_right, + deterministic, + kwargs["rng_state"].shape, + ) + ) + return torch.ones_like(q), torch.ones_like(k), torch.ones_like(v), torch.empty(0) + + core = StrictRocmAiterCKAttentionCore( + _mha_fwd=fake_fwd, + _mha_bwd=fake_bwd, + _source_sha256="a" * 64, + ) + monkeypatch.setattr(core, "_validate_inputs", lambda *_args: None) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: type("Props", (), {"name": "test-gpu", "gcnArchName": "gfx-test"})(), + ) + q = torch.randn(1, 4, 2, 8, dtype=torch.bfloat16, requires_grad=True) + k = torch.randn(1, 2, 3, 8, dtype=torch.bfloat16, requires_grad=True) + v = torch.randn(1, 2, 3, 8, dtype=torch.bfloat16, requires_grad=True) + result = core.forward_with_lse( + q, + k, + v, + causal=True, + scale=0.125, + query_position_ids=torch.tensor([[1, 2]]), + key_position_ids=torch.tensor([[0, 1, 2]]), + ) + result.out.float().sum().backward() + + assert calls == [ + ("forward", 0.0, 0.125, True, -1, -1, 0, True, False), + ("backward", 0.0, 0.125, True, -1, -1, True, torch.Size([2])), + ] + assert result.out.shape == q.shape + assert result.lse.shape == q.shape[:3] + assert result.lse.dtype is torch.float32 + assert result.provenance["strict_core_id"] == STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID + assert result.provenance["strict_schedule"] == STRICT_ATTENTION_ROCM_SCHEDULE_ID + assert result.provenance["attention_backend"] == "aiter.rocm.ck_dense_mha" + assert result.provenance["split_kv_control"] == "dense_non_split_api" + assert result.provenance["num_splits"] == 1 + assert result.provenance["deterministic_backward"] is True + assert result.provenance["aiter_source_sha256"] == "a" * 64 + assert q.grad is not None and k.grad is not None and v.grad is not None + + +def test_strict_rocm_aiter_ck_core_rejects_non_fp32_lse(monkeypatch): + def fake_fwd(q, k, v, *_args): + return ( + q, + torch.zeros(q.size(0), q.size(2), q.size(1), dtype=q.dtype), + torch.empty(0), + torch.empty(2), + ) + + core = StrictRocmAiterCKAttentionCore( + _mha_fwd=fake_fwd, + _mha_bwd=lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr(core, "_validate_inputs", lambda *_args: None) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: type("Props", (), {"name": "test-gpu", "gcnArchName": "gfx-test"})(), + ) + q = torch.randn(1, 4, 1, 8, dtype=torch.bfloat16) + k = torch.randn(1, 2, 1, 8, dtype=torch.bfloat16) + + with pytest.raises(StrictRocmAttentionUnavailable, match="FP32 LSE"): + core.forward_with_lse( + q, + k, + k, + causal=True, + query_position_ids=torch.tensor([[0]]), + key_position_ids=torch.tensor([[0]]), + ) diff --git a/tests/test_attention_dispatch.py b/tests/test_attention_dispatch.py new file mode 100644 index 00000000..5c3d296f --- /dev/null +++ b/tests/test_attention_dispatch.py @@ -0,0 +1,499 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Contract-aware attention dispatch: registration, policy, and fail-closed. + +These cases use a fresh ``KernelRegistry`` so they never mutate the process +singleton, and they assert the property that motivates a separate dispatch +entry point: a WS2 attention caller must never be served by a backend that +declares different reduction, Split-KV, or LSE-export semantics. +""" + +from __future__ import annotations + +import pytest + +from rl_engine.kernels.attention_contract import ( + AttentionBackendCapability, + AttentionContract, + AttentionContractError, + AttentionDType, + AttentionMode, + AttentionRole, + ReductionSpec, + ShardingSpec, + SplitKVSpec, + validate_cross_config_alignment, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend, _rocm_strict_attention_available + + +def _contract(*, cp_world_size: int = 1, seq_len: int = 128) -> AttentionContract: + sharding = ShardingSpec( + tp_rank=0, + tp_world_size=1, + cp_rank=0, + cp_world_size=cp_world_size, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=0, + local_q_heads=32, + local_kv_head_start=0, + local_kv_heads=8, + global_sequence_length=seq_len, + local_sequence_length=seq_len, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, seq_len), + ) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=seq_len, + head_dim=128, + causal=True, + causal_offsets=(0,), + sharding=sharding, + reduction=ReductionSpec(), + split_kv=SplitKVSpec.disabled(), + export_lse=True, + ) + + +def _capability(**overrides) -> AttentionBackendCapability: + fields = { + "backend_id": "test.strict.core", + "roles": frozenset({AttentionRole.TRAIN, AttentionRole.INFER}), + "modes": frozenset({AttentionMode.PREFILL}), + "dtypes": frozenset({AttentionDType.BF16}), + "cp_world_sizes": (1,), + "exports_attention_lse": True, + "reports_actual_split_kv_plan": True, + "implementation_kind": "production", + } + fields.update(overrides) + return AttentionBackendCapability(**fields) + + +class _FakeCore: + pass + + +@pytest.fixture() +def registry(monkeypatch): + fresh = KernelRegistry() + # Serve a stand-in instance so dispatch never imports a vendor stack. + monkeypatch.setattr(fresh, "_get_or_create_backend", lambda backend: _FakeCore()) + return fresh + + +def _platform(registry) -> str: + return registry._platform() + + +def test_strict_rocm_core_is_registered_only_when_the_vendor_stack_loads(): + """The strict core is conditional; every other candidate is static. + + ``ws2_attention`` is a static priority list, so a backend that exists only + on some machines cannot be declared there. It registers itself at runtime, + and only when ``aiter.ops.mha`` really loaded - otherwise dispatch would + offer a backend that fails at materialization. + """ + + fresh = KernelRegistry() + expected = _rocm_strict_attention_available() + + rocm = fresh._priority_map["rocm"].get("ws2_attention", []) + assert (OpBackend.ROCM_STRICT_ATTENTION in rocm) is expected + if expected: + # It must lead: a strict caller should not land on a reference first. + assert rocm[0] is OpBackend.ROCM_STRICT_ATTENTION + assert OpBackend.ROCM_STRICT_ATTENTION in fresh._attention_capabilities + + # It is a ROCm backend and must never appear on another platform. + for platform in ("cuda", "cpu"): + assert OpBackend.ROCM_STRICT_ATTENTION not in fresh._priority_map[platform].get( + "ws2_attention", [] + ) + + +def test_rocm_strict_cp_capability_matches_the_transport(): + """The declared CP degrees must be the ones the RCCL transport accepts. + + CP is supplied by StrictRocmAttentionRuntime wrapping this core in the RCCL + AG/RS transport. Declaring a degree the transport rejects would make + dispatch hand back a backend that fails at materialization; declaring fewer + would hide working CP behind an "unsupported" rejection. + """ + + if not _rocm_strict_attention_available(): + pytest.skip("strict ROCm attention requires a ROCm device with aiter.ops.mha") + + capability = KernelRegistry()._attention_capabilities[OpBackend.ROCM_STRICT_ATTENTION] + + assert capability.cp_world_sizes == (1, 2, 4, 8) + # The merge order is the transport's fixed balanced rank tree, so the CP + # combine is deterministic even though the core is single-rank arithmetic. + assert capability.deterministic_cp_merge is True + assert capability.exports_attention_lse is True + + +def test_registered_backend_resolves_with_provenance(registry): + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, _capability(), platform=_platform(registry) + ) + + result = registry.get_attention_op(_contract(), requested_backend="test.strict.core") + + assert result.capability.backend_id == "test.strict.core" + assert result.provenance["actual_backend"] == "test.strict.core" + assert result.provenance["fallback"] is False + assert result.provenance["requested_backend"] == "test.strict.core" + assert result.provenance["contract"]["lse_domain"] == "attention" + + +def test_unregistered_contract_fails_loudly_instead_of_falling_back(registry): + # With no attention candidate at all, dispatch must raise rather than reach + # into the legacy priority lists for something with other semantics. + for ops in registry._priority_map.values(): + ops["ws2_attention"] = [] + + with pytest.raises(RuntimeError, match="No attention backend supports"): + registry.get_attention_op(_contract(), requested_backend="auto") + + +def test_explicit_backend_id_never_resolves_to_a_different_backend(registry): + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, _capability(), platform=_platform(registry) + ) + + with pytest.raises(RuntimeError, match="does not match requested_backend"): + registry.get_attention_op(_contract(), requested_backend="some.other.backend") + + +def test_capability_mismatch_is_rejected_rather_than_approximated(registry): + # A backend that cannot export attention-domain LSE must not serve a + # contract that requires it. + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, + _capability(exports_attention_lse=False), + platform=_platform(registry), + ) + + with pytest.raises(RuntimeError, match="LSE export is unsupported"): + registry.get_attention_op(_contract(), requested_backend="test.strict.core") + + +def test_cp_contract_requires_deterministic_merge_support(registry): + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, + _capability(cp_world_sizes=(1, 2)), + platform=_platform(registry), + ) + + with pytest.raises(RuntimeError, match="deterministic CP"): + registry.get_attention_op(_contract(cp_world_size=2), requested_backend="test.strict.core") + + +def test_auto_is_rejected_under_context_parallelism(registry): + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, _capability(), platform=_platform(registry) + ) + + with pytest.raises(AttentionContractError, match="Unsafe dispatch"): + registry.get_attention_op(_contract(cp_world_size=2), requested_backend="auto") + + +def test_deterministic_is_a_valid_attention_policy(registry): + """``deterministic`` is a real ``implementation_kind`` for attention. + + It is not a policy for logprob dispatch, but ``AttentionBackendCapability`` + admits it, and it is the default here, so requesting it must select a + deterministic backend rather than being rejected. + """ + + platform = _platform(registry) + registry._priority_map[platform]["ws2_attention"] = [OpBackend.ROCM_STRICT_ATTENTION] + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, + _capability(implementation_kind="deterministic"), + platform=platform, + ) + + result = registry.get_attention_op(_contract(), requested_backend="deterministic") + assert result.capability.implementation_kind == "deterministic" + + with pytest.raises(RuntimeError, match="does not satisfy requested_backend=production"): + registry.get_attention_op(_contract(), requested_backend="production") + + +def test_implementation_kind_policy_filters_without_marking_fallback(registry): + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, _capability(), platform=_platform(registry) + ) + + result = registry.get_attention_op(_contract(), requested_backend="production") + assert result.provenance["fallback"] is False + + with pytest.raises(RuntimeError, match="does not satisfy requested_backend=reference"): + registry.get_attention_op(_contract(), requested_backend="reference") + + +def test_reregistration_replaces_capability_without_duplicating(registry): + platform = _platform(registry) + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, _capability(), platform=platform + ) + first = list(registry._priority_map[platform]["ws2_attention"]) + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, + _capability(implementation_kind="reference"), + platform=platform, + ) + second = registry._priority_map[platform]["ws2_attention"] + + assert first == second + assert second.count(OpBackend.ROCM_STRICT_ATTENTION) == 1 + assert ( + registry._attention_capabilities[OpBackend.ROCM_STRICT_ATTENTION].implementation_kind + == "reference" + ) + + +def test_register_rejects_wrong_types_and_unknown_platforms(registry): + with pytest.raises(AttentionContractError, match="must be an OpBackend"): + registry.register_attention_backend("not-a-backend", _capability()) + with pytest.raises(AttentionContractError, match="must be an AttentionBackendCapability"): + registry.register_attention_backend(OpBackend.ROCM_STRICT_ATTENTION, object()) + with pytest.raises(AttentionContractError, match="unsupported platform"): + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, _capability(), platform="quantum" + ) + + +def test_registration_touches_only_the_ws2_attention_list(registry): + """Registering must not perturb any legacy dispatch key. + + ``ws2_attention`` lives inside the priority map, so registration does write + there - but the SDPA-shaped ``attn`` / ``attention`` keys that legacy + ``get_op`` callers resolve through must be left exactly as they were. + """ + + platform = _platform(registry) + before = {op: list(v) for op, v in registry._priority_map[platform].items()} + + registry.register_attention_backend( + OpBackend.ROCM_STRICT_ATTENTION, _capability(), platform=platform + ) + + after = {op: list(v) for op, v in registry._priority_map[platform].items()} + changed = {op for op in after if before.get(op) != after[op]} + assert changed <= {"ws2_attention"} + for legacy in ("attn", "attention", "cp_attention", "kv_cache_attention"): + if legacy in before: + assert before[legacy] == after[legacy] + assert OpBackend.ROCM_STRICT_ATTENTION not in after[legacy] + + +def test_contract_fingerprint_is_rank_independent(): + left = _contract() + right_sharding = ShardingSpec( + tp_rank=1, + tp_world_size=2, + cp_rank=0, + cp_world_size=1, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=16, + local_q_heads=16, + local_kv_head_start=4, + local_kv_heads=4, + global_sequence_length=128, + local_sequence_length=128, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 128), + ) + left_tp2 = AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=128, + head_dim=128, + causal=True, + causal_offsets=(0,), + sharding=ShardingSpec( + tp_rank=0, + tp_world_size=2, + cp_rank=0, + cp_world_size=1, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=0, + local_q_heads=16, + local_kv_head_start=0, + local_kv_heads=4, + global_sequence_length=128, + local_sequence_length=128, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 128), + ), + reduction=ReductionSpec(), + split_kv=SplitKVSpec.disabled(), + export_lse=True, + ) + right_tp2 = AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=128, + head_dim=128, + causal=True, + causal_offsets=(0,), + sharding=right_sharding, + reduction=ReductionSpec(), + split_kv=SplitKVSpec.disabled(), + export_lse=True, + ) + + # Both TP ranks of one logical invocation agree; a different TP degree does not. + assert left_tp2.cross_rank_fingerprint() == right_tp2.cross_rank_fingerprint() + assert left.cross_rank_fingerprint() != left_tp2.cross_rank_fingerprint() + + +# --------------------------------------------------------------------------- +# Cross-config binding: train and rollout must use the same parallel degrees +# --------------------------------------------------------------------------- + + +def _tp_contract(tp_world_size: int, tp_rank: int = 0, seq_len: int = 128) -> AttentionContract: + """A contract for one TP rank of a 32Q/8KV layout.""" + + local_q = 32 // tp_world_size + local_kv = 8 // tp_world_size + sharding = ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + cp_rank=0, + cp_world_size=1, + global_q_heads=32, + global_kv_heads=8, + local_q_head_start=tp_rank * local_q, + local_q_heads=local_q, + local_kv_head_start=tp_rank * local_kv, + local_kv_heads=local_kv, + global_sequence_length=seq_len, + local_sequence_length=seq_len, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, seq_len), + ) + return AttentionContract( + role=AttentionRole.TRAIN, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=seq_len, + head_dim=128, + causal=True, + causal_offsets=(0,), + sharding=sharding, + reduction=ReductionSpec(), + split_kv=SplitKVSpec.disabled(), + export_lse=True, + ) + + +def test_matching_contracts_pass_cross_config_alignment(): + validate_cross_config_alignment(_tp_contract(4, tp_rank=0), _tp_contract(4, tp_rank=3)) + + +def test_differing_tp_degrees_are_comparable(): + """A TP-degree difference must NOT be rejected. + + The provider pins every launch to one batch row and one KV group, which + makes a head shard's result independent of the TP degree that produced it + (verified bitwise at TP=1/2/4/8 on MI300X). Rejecting the comparison would + refuse results that are in fact identical. + """ + + validate_cross_config_alignment(_tp_contract(4), _tp_contract(8)) + validate_cross_config_alignment(_tp_contract(1), _tp_contract(8)) + + +def test_head_layout_mismatch_fails_closed(): + train = _tp_contract(2) + rollout_sharding = ShardingSpec( + tp_rank=0, + tp_world_size=2, + cp_rank=0, + cp_world_size=1, + global_q_heads=16, + global_kv_heads=8, + local_q_head_start=0, + local_q_heads=8, + local_kv_head_start=0, + local_kv_heads=4, + global_sequence_length=128, + local_sequence_length=128, + global_block_indices=(0,), + global_block_token_starts=(0,), + local_block_offsets=(0, 128), + ) + rollout = AttentionContract( + role=AttentionRole.INFER, + mode=AttentionMode.PREFILL, + dtype=AttentionDType.BF16, + batch_size=1, + query_sequence_length=128, + head_dim=128, + causal=True, + causal_offsets=(0,), + sharding=rollout_sharding, + reduction=ReductionSpec(), + split_kv=SplitKVSpec.disabled(), + export_lse=True, + ) + with pytest.raises(AttentionContractError, match="global head layouts"): + validate_cross_config_alignment(train, rollout) + + +def test_contract_fingerprint_is_a_per_invocation_rank_preflight(): + """The fingerprint agrees across ranks of one invocation, and only there. + + It still separates TP degrees, which is correct for its purpose: every rank + of a single logical invocation must agree on one topology. It is not a + train-vs-rollout equality token -- those may legitimately run different TP + degrees and still be bitwise equal. + """ + + assert ( + _tp_contract(4, tp_rank=0).cross_rank_fingerprint() + == _tp_contract(4, tp_rank=3).cross_rank_fingerprint() + ) + assert _tp_contract(4).cross_rank_fingerprint() != _tp_contract(8).cross_rank_fingerprint() + + +def test_dtype_and_split_kv_mismatches_are_explained(): + train = _tp_contract(2) + rollout = AttentionContract( + role=train.role, + mode=train.mode, + dtype=AttentionDType.FP16, + batch_size=train.batch_size, + query_sequence_length=train.query_sequence_length, + head_dim=train.head_dim, + causal=train.causal, + causal_offsets=train.causal_offsets, + sharding=train.sharding, + reduction=train.reduction, + split_kv=train.split_kv, + export_lse=True, + ) + with pytest.raises(AttentionContractError, match="dtype"): + validate_cross_config_alignment(train, rollout) diff --git a/tests/test_build_platform_collectives.py b/tests/test_build_platform_collectives.py new file mode 100644 index 00000000..19d3a890 --- /dev/null +++ b/tests/test_build_platform_collectives.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import runpy +from typing import Any + +import setuptools +import torch +from torch.utils import cpp_extension + + +def _load_extension_config(monkeypatch, *, hip: str | None) -> dict[str, Any]: + captured: dict[str, Any] = {} + + def fake_setup(**kwargs: Any) -> None: + captured.update(kwargs) + + def fake_extension(**kwargs: Any) -> dict[str, Any]: + return kwargs + + monkeypatch.setattr(setuptools, "setup", fake_setup) + monkeypatch.setattr(cpp_extension, "CUDAExtension", fake_extension) + monkeypatch.setattr(torch.version, "hip", hip, raising=False) + monkeypatch.delenv("KERNEL_ALIGN_FORCE_SM90", raising=False) + monkeypatch.delenv("KERNEL_ALIGN_DET_GEMM_SM90", raising=False) + if hip is None: + monkeypatch.delenv("PYTORCH_ROCM_ARCH", raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (8, 0)) + else: + monkeypatch.setenv("PYTORCH_ROCM_ARCH", "gfx942") + + runpy.run_path("setup.py", run_name=f"rl_kernel_setup_probe_{hip or 'cuda'}") + return captured["ext_modules"][0] + + +def test_rocm_build_excludes_cuda_ipc_collective_and_driver(monkeypatch) -> None: + extension = _load_extension_config(monkeypatch, hip="test") + + assert "csrc/cuda/distributed/deterministic_collective.cu" not in extension["sources"] + assert "csrc/rocm/distributed/deterministic_collective.hip" in extension["sources"] + assert "-DKERNEL_ALIGN_WITH_ROCM" in extension["extra_compile_args"]["cxx"] + assert "-DKERNEL_ALIGN_WITH_CUDA" not in extension["extra_compile_args"]["cxx"] + assert "-lcuda" not in extension["extra_link_args"] + + +def test_cuda_build_keeps_existing_ipc_collective(monkeypatch) -> None: + extension = _load_extension_config(monkeypatch, hip=None) + + assert "csrc/cuda/distributed/deterministic_collective.cu" in extension["sources"] + assert "-DKERNEL_ALIGN_WITH_CUDA" in extension["extra_compile_args"]["cxx"] + assert "-DKERNEL_ALIGN_WITH_ROCM" not in extension["extra_compile_args"]["cxx"] + assert "-lcuda" in extension["extra_link_args"] diff --git a/tests/test_deterministic_attention_cuda.py b/tests/test_deterministic_attention_cuda.py index f08315a8..f981e840 100644 --- a/tests/test_deterministic_attention_cuda.py +++ b/tests/test_deterministic_attention_cuda.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Deterministic standard-softmax attention CUDA tests (issue #147). +"""Deterministic standard-softmax Attention tests for CUDA and ROCm. Covers (per §7 and §8 of the implementation plan): - Forward correctness via #108 harness (run_operator_suite) @@ -23,28 +23,43 @@ import pytest import torch -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +from rl_engine.platforms.device import device_ctx + +IS_ROCM = device_ctx.is_rocm +IS_CUDA = device_ctx.device_type == "cuda" +IS_GPU = IS_CUDA or IS_ROCM +BACKEND = "rocm" if IS_ROCM else "cuda" + +pytestmark = pytest.mark.skipif(not IS_GPU, reason="CUDA/ROCm GPU not available") +ROCM_ONLY = pytest.mark.skipif(not IS_ROCM, reason="ROCm-only acceptance check") try: + from rl_engine.kernels.attention_contract import SplitKVSpec from rl_engine.kernels.gtest.op_checks import CandidateSpec, OperatorCase, run_operator_suite - from rl_engine.kernels.ops.cuda.attention.deterministic_attn import DeterministicAttentionOp + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionOp, + RLKernelDeterministicAttentionCore, + ) from rl_engine.kernels.ops.pytorch.attention.standard_attn import NativeAttentionOp _OP_AVAILABLE = True except (ImportError, RuntimeError): _OP_AVAILABLE = False +if IS_ROCM: + from rl_engine.kernels.ops.cuda.rotary_embedding.rope import RocmDeterministicRoPEOp + pytestmark = [ pytestmark, - pytest.mark.skipif(not _OP_AVAILABLE, reason="CUDA attention op not built"), + pytest.mark.skipif(not _OP_AVAILABLE, reason=f"{BACKEND} attention op not built"), ] -DEVICE = "cuda" +DEVICE = device_ctx.device D = 128 @pytest.fixture -def cuda_op(): +def attention_op(): return DeterministicAttentionOp() @@ -109,8 +124,8 @@ def _build_harness_cases(): def test_harness_forward(): """§8.4: run_operator_suite forward — candidate vs gold (accuracy tolerance).""" - cuda = DeterministicAttentionOp() - candidate = CandidateSpec(name="cuda-attention", fn=cuda, backend="cuda") + attention = DeterministicAttentionOp() + candidate = CandidateSpec(name=f"{BACKEND}-attention", fn=attention, backend=BACKEND) report = run_operator_suite("attention", candidates=[candidate], cases=_build_harness_cases()) for cr in report.candidates: for case in cr.cases: @@ -126,8 +141,8 @@ def test_harness_forward(): def test_harness_backward(): """§8.4: run_operator_suite backward — grad comparison vs gold fp32 autograd.""" - cuda = DeterministicAttentionOp() - candidate = CandidateSpec(name="cuda-attention", fn=cuda, backend="cuda") + attention = DeterministicAttentionOp() + candidate = CandidateSpec(name=f"{BACKEND}-attention", fn=attention, backend=BACKEND) report = run_operator_suite( "attention", candidates=[candidate], @@ -163,18 +178,18 @@ def test_harness_backward(): @pytest.mark.parametrize("dtype,hq,hkv,sq,skv,causal", SWEEP_CONFIGS) -def test_forward_sweep(cuda_op, gold_op, dtype, hq, hkv, sq, skv, causal): +def test_forward_sweep(attention_op, gold_op, dtype, hq, hkv, sq, skv, causal): B = 2 torch.manual_seed(42) q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=dtype) k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=dtype) v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=dtype) - out_cuda = cuda_op.forward(q, k, v, causal=causal) + out_actual = attention_op.forward(q, k, v, causal=causal) out_gold = gold_op.forward_fp32(q, k, v, causal=causal) atol, rtol = _tol(dtype) - torch.testing.assert_close(out_cuda.float(), out_gold.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(out_actual.float(), out_gold.float(), atol=atol, rtol=rtol) # ============================================================================= @@ -183,18 +198,18 @@ def test_forward_sweep(cuda_op, gold_op, dtype, hq, hkv, sq, skv, causal): @pytest.mark.parametrize("scale", [None, 0.0, 0.05]) -def test_scale(cuda_op, gold_op, scale): +def test_scale(attention_op, gold_op, scale): B, hq, hkv, sq, skv = 2, 4, 1, 8, 16 torch.manual_seed(42) q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) - out_cuda = cuda_op.forward(q, k, v, causal=True, scale=scale) + out_actual = attention_op.forward(q, k, v, causal=True, scale=scale) out_gold = gold_op.forward_fp32(q, k, v, causal=True, scale=scale) atol, rtol = _tol(torch.bfloat16) - torch.testing.assert_close(out_cuda.float(), out_gold.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(out_actual.float(), out_gold.float(), atol=atol, rtol=rtol) # ============================================================================= @@ -203,7 +218,7 @@ def test_scale(cuda_op, gold_op, scale): @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) -def test_forward_with_padding(cuda_op, gold_op, dtype): +def test_forward_with_padding(attention_op, gold_op, dtype): B, hq, hkv, sq, skv = 2, 4, 1, 8, 16 torch.manual_seed(7) q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=dtype) @@ -213,21 +228,21 @@ def test_forward_with_padding(cuda_op, gold_op, dtype): mask[0, 10:] = False mask[1, 12:] = False - out_cuda = cuda_op.forward(q, k, v, causal=True, key_padding_mask=mask) + out_actual = attention_op.forward(q, k, v, causal=True, key_padding_mask=mask) out_gold = gold_op.forward_fp32(q, k, v, causal=True, key_padding_mask=mask) atol, rtol = _tol(dtype) - torch.testing.assert_close(out_cuda.float(), out_gold.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(out_actual.float(), out_gold.float(), atol=atol, rtol=rtol) -def test_fully_masked_row(cuda_op): +def test_fully_masked_row(attention_op): B, hq, hkv, sq, skv = 1, 1, 1, 2, 4 q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) mask = torch.zeros(B, skv, device=DEVICE, dtype=torch.bool) - out, lse = cuda_op.forward_with_lse(q, k, v, causal=False, key_padding_mask=mask) + out, lse = attention_op.forward_with_lse(q, k, v, causal=False, key_padding_mask=mask) assert (out == 0).all() assert (lse == float("-inf")).all() @@ -237,14 +252,14 @@ def test_fully_masked_row(cuda_op): # ============================================================================= -def test_lse_correctness(cuda_op): +def test_lse_correctness(attention_op): B, hq, hkv, sq, skv = 2, 4, 1, 8, 16 torch.manual_seed(99) q = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.bfloat16) k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) - _, lse_cuda = cuda_op.forward_with_lse(q, k, v, causal=True) + _, lse_actual = attention_op.forward_with_lse(q, k, v, causal=True) scale = 1.0 / math.sqrt(D) g = hq // hkv @@ -256,7 +271,7 @@ def test_lse_correctness(cuda_op): scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf")) lse_gold = torch.logsumexp(scores, dim=-1) - torch.testing.assert_close(lse_cuda, lse_gold, atol=5e-2, rtol=2e-2) + torch.testing.assert_close(lse_actual, lse_gold, atol=5e-2, rtol=2e-2) # ============================================================================= @@ -264,7 +279,7 @@ def test_lse_correctness(cuda_op): # ============================================================================= -def test_valid_only_vs_padded_accuracy(cuda_op): +def test_valid_only_vs_padded_accuracy(attention_op): """Padding changes reduction width; result is near-equal, not bitwise. We use causal=False here so that all valid keys are equally visible @@ -295,8 +310,8 @@ def test_valid_only_vs_padded_accuracy(cuda_op): mask = torch.ones(B, skv_padded, device=DEVICE, dtype=torch.bool) mask[:, skv_valid:] = False - out_valid = cuda_op.forward(q, k_valid, v_valid, causal=False) - out_padded = cuda_op.forward(q, k_padded, v_padded, causal=False, key_padding_mask=mask) + out_valid = attention_op.forward(q, k_valid, v_valid, causal=False) + out_padded = attention_op.forward(q, k_padded, v_padded, causal=False, key_padding_mask=mask) atol, rtol = _tol(torch.bfloat16) torch.testing.assert_close(out_valid.float(), out_padded.float(), atol=atol, rtol=rtol) @@ -307,7 +322,7 @@ def test_valid_only_vs_padded_accuracy(cuda_op): # ============================================================================= -def test_batch_invariance_single(cuda_op): +def test_batch_invariance_single(attention_op): """Same sample in full batch vs extracted single — bitwise.""" B, hq, hkv, sq, skv = 4, 4, 1, 8, 16 torch.manual_seed(11) @@ -315,17 +330,17 @@ def test_batch_invariance_single(cuda_op): k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) - out_full, lse_full = cuda_op.forward_with_lse(q, k, v, causal=True) + out_full, lse_full = attention_op.forward_with_lse(q, k, v, causal=True) for i in range(B): - out_single, lse_single = cuda_op.forward_with_lse( + out_single, lse_single = attention_op.forward_with_lse( q[i : i + 1], k[i : i + 1], v[i : i + 1], causal=True ) assert torch.equal(out_full[i : i + 1], out_single), f"Output batch invariance failed i={i}" assert torch.equal(lse_full[i : i + 1], lse_single), f"LSE batch invariance failed i={i}" -def test_batch_invariance_position_permutation(cuda_op): +def test_batch_invariance_position_permutation(attention_op): """Same sample at different batch positions — bitwise.""" B, hq, hkv, sq, skv = 4, 32, 8, 8, 16 torch.manual_seed(12) @@ -333,13 +348,13 @@ def test_batch_invariance_position_permutation(cuda_op): k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) - out_full = cuda_op.forward(q, k, v, causal=True) + out_full = attention_op.forward(q, k, v, causal=True) perm = [2, 0, 3, 1] q_perm = q[perm] k_perm = k[perm] v_perm = v[perm] - out_perm = cuda_op.forward(q_perm, k_perm, v_perm, causal=True) + out_perm = attention_op.forward(q_perm, k_perm, v_perm, causal=True) for new_pos, orig_pos in enumerate(perm): assert torch.equal( @@ -347,7 +362,7 @@ def test_batch_invariance_position_permutation(cuda_op): ), f"Position permutation invariance failed: orig={orig_pos} new={new_pos}" -def test_batch_invariance_chunk(cuda_op): +def test_batch_invariance_chunk(attention_op): """Batch-dim chunking — bitwise.""" B, hq, hkv, sq, skv = 4, 4, 1, 8, 16 torch.manual_seed(13) @@ -355,11 +370,11 @@ def test_batch_invariance_chunk(cuda_op): k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) - out_full = cuda_op.forward(q, k, v, causal=True) + out_full = attention_op.forward(q, k, v, causal=True) out_chunk = torch.cat( [ - cuda_op.forward(q[:2], k[:2], v[:2], causal=True), - cuda_op.forward(q[2:], k[2:], v[2:], causal=True), + attention_op.forward(q[:2], k[:2], v[:2], causal=True), + attention_op.forward(q[2:], k[2:], v[2:], causal=True), ], dim=0, ) @@ -381,20 +396,22 @@ def test_batch_invariance_chunk(cuda_op): (3, 32, 8), ], ) -def test_chunked_prefill(cuda_op, chunk_size, hq, hkv): +def test_chunked_prefill(attention_op, chunk_size, hq, hkv): B, T = 1, 16 torch.manual_seed(22) q = torch.randn(B, hq, T, D, device=DEVICE, dtype=torch.bfloat16) k = torch.randn(B, hkv, T, D, device=DEVICE, dtype=torch.bfloat16) v = torch.randn(B, hkv, T, D, device=DEVICE, dtype=torch.bfloat16) - out_full = cuda_op.forward(q, k, v, causal=True) + out_full = attention_op.forward(q, k, v, causal=True) outs = [] for t in range(0, T, chunk_size): c = min(chunk_size, T - t) outs.append( - cuda_op.forward(q[:, :, t : t + c], k[:, :, : t + c], v[:, :, : t + c], causal=True) + attention_op.forward( + q[:, :, t : t + c], k[:, :, : t + c], v[:, :, : t + c], causal=True + ) ) out_chunked = torch.cat(outs, dim=2) assert torch.equal( @@ -403,7 +420,7 @@ def test_chunked_prefill(cuda_op, chunk_size, hq, hkv): @pytest.mark.parametrize("chunk_size", [1, 3, 8]) -def test_chunked_prefill_with_padding(cuda_op, chunk_size): +def test_chunked_prefill_with_padding(attention_op, chunk_size): """§7.4: chunked-prefill with key_padding_mask (mask sliced with Skv).""" B, hq, hkv, T = 1, 4, 1, 16 torch.manual_seed(23) @@ -413,13 +430,13 @@ def test_chunked_prefill_with_padding(cuda_op, chunk_size): mask = torch.ones(B, T, device=DEVICE, dtype=torch.bool) mask[0, 12:] = False - out_full = cuda_op.forward(q, k, v, causal=True, key_padding_mask=mask) + out_full = attention_op.forward(q, k, v, causal=True, key_padding_mask=mask) outs = [] for t in range(0, T, chunk_size): c = min(chunk_size, T - t) outs.append( - cuda_op.forward( + attention_op.forward( q[:, :, t : t + c], k[:, :, : t + c], v[:, :, : t + c], @@ -434,7 +451,7 @@ def test_chunked_prefill_with_padding(cuda_op, chunk_size): @pytest.mark.parametrize("chunk_size", [1, 3]) -def test_chunked_prefill_lse(cuda_op, chunk_size): +def test_chunked_prefill_lse(attention_op, chunk_size): """§7.4: LSE chunked-prefill invariance.""" B, hq, hkv, T = 1, 4, 1, 12 torch.manual_seed(24) @@ -442,12 +459,12 @@ def test_chunked_prefill_lse(cuda_op, chunk_size): k = torch.randn(B, hkv, T, D, device=DEVICE, dtype=torch.bfloat16) v = torch.randn(B, hkv, T, D, device=DEVICE, dtype=torch.bfloat16) - _, lse_full = cuda_op.forward_with_lse(q, k, v, causal=True) + _, lse_full = attention_op.forward_with_lse(q, k, v, causal=True) lses = [] for t in range(0, T, chunk_size): c = min(chunk_size, T - t) - _, lse_chunk = cuda_op.forward_with_lse( + _, lse_chunk = attention_op.forward_with_lse( q[:, :, t : t + c], k[:, :, : t + c], v[:, :, : t + c], causal=True ) lses.append(lse_chunk) @@ -460,7 +477,7 @@ def test_chunked_prefill_lse(cuda_op, chunk_size): # ============================================================================= -def test_prefill_decode_slice(cuda_op): +def test_prefill_decode_slice(attention_op): """§7.5.1: prefill[:, :, -1:] == decode(q[-1:], k_full, v_full).""" B, hq, hkv, sq, skv = 1, 4, 1, 8, 16 torch.manual_seed(33) @@ -468,8 +485,8 @@ def test_prefill_decode_slice(cuda_op): k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16) - prefill = cuda_op.forward(q, k, v, causal=True) - decode = cuda_op.forward(q[:, :, -1:], k, v, causal=True) + prefill = attention_op.forward(q, k, v, causal=True) + decode = attention_op.forward(q[:, :, -1:], k, v, causal=True) assert torch.equal(prefill[:, :, -1:], decode) @@ -482,7 +499,7 @@ def test_prefill_decode_slice(cuda_op): (3, 32, 8), ], ) -def test_kv_cache_handoff(cuda_op, S_new, hq, hkv): +def test_kv_cache_handoff(attention_op, S_new, hq, hkv): """§7.5.2: cat(k_cache, k_new) handoff == prefill tail.""" B, S_past = 1, 12 torch.manual_seed(44) @@ -491,12 +508,12 @@ def test_kv_cache_handoff(cuda_op, S_new, hq, hkv): v_full = torch.randn(B, hkv, S_past + S_new, D, device=DEVICE, dtype=torch.bfloat16) q_new = q_full[:, :, -S_new:] - prefill_tail = cuda_op.forward(q_full, k_full, v_full, causal=True)[:, :, -S_new:] - decode_path = cuda_op.forward(q_new, k_full, v_full, causal=True) + prefill_tail = attention_op.forward(q_full, k_full, v_full, causal=True)[:, :, -S_new:] + decode_path = attention_op.forward(q_new, k_full, v_full, causal=True) assert torch.equal(decode_path, prefill_tail) -def test_kv_cache_handoff_with_padding(cuda_op): +def test_kv_cache_handoff_with_padding(attention_op): """§7.5.2: cat handoff with padding mask.""" B, hq, hkv, S_past, S_new = 1, 4, 1, 12, 1 Skv = S_past + S_new @@ -508,10 +525,10 @@ def test_kv_cache_handoff_with_padding(cuda_op): mask[0, 8:10] = False q_new = q_full[:, :, -S_new:] - prefill_tail = cuda_op.forward(q_full, k_full, v_full, causal=True, key_padding_mask=mask)[ + prefill_tail = attention_op.forward(q_full, k_full, v_full, causal=True, key_padding_mask=mask)[ :, :, -S_new: ] - decode_path = cuda_op.forward(q_new, k_full, v_full, causal=True, key_padding_mask=mask) + decode_path = attention_op.forward(q_new, k_full, v_full, causal=True, key_padding_mask=mask) assert torch.equal(decode_path, prefill_tail) @@ -520,7 +537,7 @@ def test_kv_cache_handoff_with_padding(cuda_op): # ============================================================================= -def test_backward_smoke(cuda_op): +def test_backward_smoke(attention_op): """Backward runs and produces gradients with correct shapes.""" B, hq, hkv, sq, skv = 1, 4, 1, 4, 8 torch.manual_seed(55) @@ -528,7 +545,7 @@ def test_backward_smoke(cuda_op): k = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16, requires_grad=True) v = torch.randn(B, hkv, skv, D, device=DEVICE, dtype=torch.bfloat16, requires_grad=True) - out = cuda_op.forward(q, k, v, causal=True) + out = attention_op.forward(q, k, v, causal=True) loss = out.sum() loss.backward() @@ -537,7 +554,7 @@ def test_backward_smoke(cuda_op): assert v.grad is not None and v.grad.shape == v.shape -def test_backward_fp64_reference(cuda_op): +def test_backward_fp64_reference(attention_op): """§7.6: FP64 high-precision gradient comparison.""" B, hq, hkv, sq, skv = 1, 4, 1, 4, 8 torch.manual_seed(56) @@ -547,11 +564,11 @@ def test_backward_fp64_reference(cuda_op): grad_out = torch.randn(B, hq, sq, D, device=DEVICE, dtype=torch.float32) - out_cuda = cuda_op.forward(q_bf16, k_bf16, v_bf16, causal=True) - out_cuda.backward(grad_out.to(out_cuda.dtype)) - dq_cuda = q_bf16.grad.float() - dk_cuda = k_bf16.grad.float() - dv_cuda = v_bf16.grad.float() + out_actual = attention_op.forward(q_bf16, k_bf16, v_bf16, causal=True) + out_actual.backward(grad_out.to(out_actual.dtype)) + dq_actual = q_bf16.grad.float() + dk_actual = k_bf16.grad.float() + dv_actual = v_bf16.grad.float() scale = 1.0 / math.sqrt(D) g = hq // hkv @@ -572,12 +589,12 @@ def test_backward_fp64_reference(cuda_op): dk_gold = k64.grad.float() dv_gold = v64.grad.float() - torch.testing.assert_close(dq_cuda, dq_gold, atol=5e-2, rtol=2e-2) - torch.testing.assert_close(dk_cuda, dk_gold, atol=5e-2, rtol=2e-2) - torch.testing.assert_close(dv_cuda, dv_gold, atol=5e-2, rtol=2e-2) + torch.testing.assert_close(dq_actual, dq_gold, atol=5e-2, rtol=2e-2) + torch.testing.assert_close(dk_actual, dk_gold, atol=5e-2, rtol=2e-2) + torch.testing.assert_close(dv_actual, dv_gold, atol=5e-2, rtol=2e-2) -def test_gradient_batch_invariance(cuda_op): +def test_gradient_batch_invariance(attention_op): """§7.6: dQ/dK/dV bitwise identical for same sample at different batch positions.""" B, hq, hkv, sq, skv = 3, 4, 1, 4, 8 torch.manual_seed(57) @@ -589,21 +606,21 @@ def test_gradient_batch_invariance(cuda_op): q_full = q.clone().requires_grad_(True) k_full = k.clone().requires_grad_(True) v_full = v.clone().requires_grad_(True) - out_full = cuda_op.forward(q_full, k_full, v_full, causal=True) + out_full = attention_op.forward(q_full, k_full, v_full, causal=True) out_full.backward(grad_out) for i in range(B): qi = q[i : i + 1].clone().requires_grad_(True) ki = k[i : i + 1].clone().requires_grad_(True) vi = v[i : i + 1].clone().requires_grad_(True) - out_i = cuda_op.forward(qi, ki, vi, causal=True) + out_i = attention_op.forward(qi, ki, vi, causal=True) out_i.backward(grad_out[i : i + 1]) assert torch.equal(q_full.grad[i : i + 1], qi.grad), f"dQ batch invariance failed i={i}" assert torch.equal(k_full.grad[i : i + 1], ki.grad), f"dK batch invariance failed i={i}" assert torch.equal(v_full.grad[i : i + 1], vi.grad), f"dV batch invariance failed i={i}" -def test_gqa_dk_dv_order(cuda_op): +def test_gqa_dk_dv_order(attention_op): """§7.6/§4.1: GQA dK/dV must follow fixed (hq_local, query_index) order. Two checks: @@ -623,13 +640,13 @@ def test_gqa_dk_dv_order(cuda_op): q1 = q_data.clone().requires_grad_(True) k1 = k_data.clone().requires_grad_(True) v1 = v_data.clone().requires_grad_(True) - cuda_op.forward(q1, k1, v1, causal=True).backward(grad_data) + attention_op.forward(q1, k1, v1, causal=True).backward(grad_data) q_batch = q_data.expand(4, -1, -1, -1).contiguous().clone().requires_grad_(True) k_batch = k_data.expand(4, -1, -1, -1).contiguous().clone().requires_grad_(True) v_batch = v_data.expand(4, -1, -1, -1).contiguous().clone().requires_grad_(True) grad_batch = grad_data.expand(4, -1, -1, -1).contiguous() - cuda_op.forward(q_batch, k_batch, v_batch, causal=True).backward(grad_batch) + attention_op.forward(q_batch, k_batch, v_batch, causal=True).backward(grad_batch) assert torch.equal(k1.grad, k_batch.grad[0:1]), "dK GQA order depends on batch size" assert torch.equal(v1.grad, v_batch.grad[0:1]), "dV GQA order depends on batch size" @@ -655,3 +672,119 @@ def test_gqa_dk_dv_order(cuda_op): torch.testing.assert_close(k1.grad.float(), dk_gold, atol=5e-2, rtol=2e-2) torch.testing.assert_close(v1.grad.float(), dv_gold, atol=5e-2, rtol=2e-2) + + +# ============================================================================= +# Strict core acceptance and ROCm HIP RoPE acceptance +# ============================================================================= + + +def _strict_qkv(*, batch: int = 2, sequence: int = 7): + generator = torch.Generator(device="cpu").manual_seed(942) + q = torch.randn(batch, 4, sequence, D, dtype=torch.bfloat16, generator=generator).to(DEVICE) + k = torch.randn(batch, 1, sequence, D, dtype=torch.bfloat16, generator=generator).to(DEVICE) + v = torch.randn(batch, 1, sequence, D, dtype=torch.bfloat16, generator=generator).to(DEVICE) + return q, k, v + + +def test_strict_attention_core_is_repeat_bitwise_and_no_fallback(): + q, k, v = _strict_qkv() + positions = torch.arange(q.size(2), device=q.device).expand(q.size(0), -1) + core = RLKernelDeterministicAttentionCore() + first = core.forward_with_lse( + q, + k, + v, + query_position_ids=positions, + key_position_ids=positions, + ) + second = core.forward_with_lse( + q, + k, + v, + query_position_ids=positions, + key_position_ids=positions, + ) + assert torch.equal(first.out, second.out) + assert torch.equal(first.lse, second.lse) + assert first.provenance["attention_backend"] == f"rlkernel.{BACKEND}.deterministic_attention" + assert first.provenance["fallback"] is False + assert first.provenance["split_kv"]["actual_split_kv_policy"] == "disabled" + + +def test_strict_attention_forward_backward_train_rollout_bitwise(): + q, k, v = (tensor.requires_grad_() for tensor in _strict_qkv(batch=1, sequence=5)) + positions = torch.arange(q.size(2), device=q.device).expand(q.size(0), -1) + core = RLKernelDeterministicAttentionCore() + train = core.forward_with_lse( + q, + k, + v, + query_position_ids=positions, + key_position_ids=positions, + ) + grad = torch.randn(train.out.shape, dtype=train.out.dtype, device="cpu").to(DEVICE) + (train.out.float() * grad.float()).sum().backward() + train_grads = tuple(tensor.grad.detach().clone() for tensor in (q, k, v)) + + q2, k2, v2 = (tensor.detach().clone().requires_grad_() for tensor in (q, k, v)) + rollout = core.forward_with_lse( + q2, + k2, + v2, + query_position_ids=positions, + key_position_ids=positions, + ) + (rollout.out.float() * grad.float()).sum().backward() + assert torch.equal(train.out, rollout.out) + assert torch.equal(train.lse, rollout.lse) + assert all( + torch.equal(expected, actual.grad) + for expected, actual in zip(train_grads, (q2, k2, v2), strict=True) + ) + + +@ROCM_ONLY +def test_rocm_rope_is_batch_invariant_and_backward_repeat_bitwise(): + generator = torch.Generator(device="cpu").manual_seed(714) + x = torch.randn(3, 4, 6, D, dtype=torch.bfloat16, generator=generator).to(DEVICE) + positions = torch.arange(6, device=x.device).expand(3, -1) + rope = RocmDeterministicRoPEOp() + together = rope(x, positions) + separate = torch.cat([rope(x[index : index + 1], positions[index]) for index in range(3)]) + assert torch.equal(together, separate) + assert torch.equal(together, rope(x, positions)) + + grad = torch.randn(together.shape, dtype=together.dtype, device="cpu").to(DEVICE) + x1 = x.detach().clone().requires_grad_() + x2 = x.detach().clone().requires_grad_() + (rope(x1, positions).float() * grad.float()).sum().backward() + (rope(x2, positions).float() * grad.float()).sum().backward() + assert torch.equal(x1.grad, x2.grad) + + +@ROCM_ONLY +def test_rocm_rope_matches_fp32_rotate_half_reference(): + x = torch.randn(1, 2, 4, D, dtype=torch.bfloat16, device=DEVICE) + positions = torch.arange(4, device=x.device).expand(1, -1) + actual = RocmDeterministicRoPEOp()(x, positions) + half = x.size(-1) // 2 + inv_freq = 1.0 / ( + 1_000_000.0 ** (torch.arange(half, dtype=torch.float32, device=x.device) / half) + ) + frequency = positions.float().unsqueeze(-1) * inv_freq + cos = frequency.cos().unsqueeze(1) + sin = frequency.sin().unsqueeze(1) + reference = torch.cat( + ( + x[..., :half].float() * cos - x[..., half:].float() * sin, + x[..., half:].float() * cos + x[..., :half].float() * sin, + ), + dim=-1, + ).to(x.dtype) + torch.testing.assert_close(actual, reference, atol=0, rtol=0) + + +def test_strict_core_rejects_split_k(): + with pytest.raises(ValueError, match="Split-KV"): + RLKernelDeterministicAttentionCore(split_kv=SplitKVSpec.fixed(32)) diff --git a/tests/test_flashinfer_pr7_attention.py b/tests/test_flashinfer_pr7_attention.py index fa237196..03c15455 100644 --- a/tests/test_flashinfer_pr7_attention.py +++ b/tests/test_flashinfer_pr7_attention.py @@ -3,9 +3,12 @@ from __future__ import annotations +import importlib.util import json +import sys import types from dataclasses import replace +from pathlib import Path import pytest import torch @@ -14,6 +17,9 @@ STRICT_ATTENTION_CORE_ID, STRICT_ATTENTION_FA4_SCHEDULE_ID, STRICT_ATTENTION_PRODUCTION_CORE_ID, + STRICT_ATTENTION_RING_SCHEDULE_ID, + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, STRICT_ATTENTION_SCHEDULE_ID, SplitKVSpec, ) @@ -48,8 +54,21 @@ materialize_flashinfer_paged_kv_cache, ) from rl_engine.testing.attention_comparison import DecodeKVCacheMetadata -from scripts import ws2_p2p_nccl_attention_reference_check as p2p_check_script -from scripts import ws2_pr7_flashinfer_attention_check as check_script + + +def _load_repo_script(name: str): + path = Path(__file__).resolve().parents[1] / "scripts" / f"{name}.py" + spec = importlib.util.spec_from_file_location(f"rl_kernel_{name}", path) + if spec is None or spec.loader is None: + raise ImportError(f"cannot load repository script {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +p2p_check_script = _load_repo_script("ws2_p2p_nccl_attention_reference_check") +check_script = _load_repo_script("ws2_pr7_flashinfer_attention_check") class _FakeFlashInferWrapper: @@ -1717,7 +1736,7 @@ def test_pr7_check_acceptance_errors_require_all_drift_and_invariance_fields(): assert check_script._acceptance_errors(report, args) == ["batch_invariant_sweep failed"] -def test_pr7_check_accepts_strict_fa4_production_core(): +def test_pr7_check_accepts_strict_cuda_production_core(): args = check_script._parse_args(["--strict", "--device", "cuda"]) report = { "device": "cuda:0", @@ -1752,6 +1771,44 @@ def test_pr7_check_accepts_strict_fa4_production_core(): assert check_script._acceptance_errors(report, args) == [] +def test_pr7_check_accepts_strict_rocm_production_core(): + args = check_script._parse_args(["--strict", "--device", "cuda"]) + report = { + "device": "cuda:0", + "shape": {"q_heads": 16, "kv_heads": 4, "head_dim": 128}, + "candidate_provenance": { + "attention_mode": "decode", + "fallback": False, + "strict_mode": True, + "strict_core_id": STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + "strict_schedule": STRICT_ATTENTION_ROCM_SCHEDULE_ID, + "actual_backend": "aiter.rocm.ck_dense_mha", + "platform": "rocm", + "native_attention_arithmetic": True, + "num_splits": 1, + "deterministic_backward": True, + "reference_only": False, + "split_kv_control": "dense_non_split_api", + "aiter_api_source": "aiter.ops.mha", + "aiter_source_sha256": "a" * 64, + "strict_core_row_plans": [{"actual_split_kv_policy": "disabled"}], + "rope_backend": "rlkernel.rocm.deterministic_rope", + "rope_theta": 1_000_000.0, + "rotary_dim": 128, + "arithmetic_semantics_verified": True, + }, + "drift": { + "out": {"max_abs": 0.0}, + "lse": {"max_abs": 0.0}, + "dlogp": {"max_abs": 0.0}, + }, + "batch_invariant_sweep": {"passed": True}, + "page_layout_invariant_sweep": {"passed": True}, + } + + assert check_script._acceptance_errors(report, args) == [] + + def test_pr7_check_rejects_reference_core_as_production(): args = check_script._parse_args(["--strict", "--device", "cuda"]) report = { @@ -1782,10 +1839,44 @@ def test_pr7_check_rejects_reference_core_as_production(): } errors = check_script._acceptance_errors(report, args) - assert "strict runtime did not execute the FA4 production core" in errors + assert "strict runtime did not execute the native production arithmetic" in errors assert "strict runtime selected the reference core" in errors +@pytest.mark.parametrize( + ("backend", "rope_backend"), + [ + ("flash_attention_4.cute", "rlkernel.cuda.rope_sm90"), + ("aiter.rocm.ck_dense_mha", "rlkernel.rocm.deterministic_rope"), + ], +) +def test_strict_report_uses_executed_platform_provenance(backend, rope_backend): + fields = check_script._strict_execution_report_fields( + { + "actual_backend": backend, + "rope_backend": rope_backend, + "rope_fusion": False, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1_000_000.0, + "rotary_dim": 128, + "q_rope_state": "post_rope", + "k_cache_rope_state": "post_rope", + } + ) + + assert fields["reference_backend"] == backend + assert backend in fields["target"] + assert fields["rope"] == { + "rope_backend": rope_backend, + "rope_fusion": False, + "rope_fusion_boundary": "rlkernel_rope_then_attention", + "rope_theta": 1_000_000.0, + "rotary_dim": 128, + "q_rope_state": "post_rope", + "k_cache_rope_state": "post_rope", + } + + def test_pr7_check_rejects_nonfinite_drift_and_wrong_tp_local_shape(): args = check_script._parse_args([]) report = { @@ -1845,6 +1936,116 @@ def test_strict_shared_core_entrypoint_requires_self_owned_ag_rs(): assert args.strict_shared_core is True +def test_strict_shared_core_reference_executes_one_logical_row(monkeypatch): + calls = [] + + class _RowCore: + def forward_with_lse(self, q, k, v, **kwargs): + calls.append((q.shape[0], kwargs["query_position_ids"].clone())) + return types.SimpleNamespace( + out=q + k + v, + lse=(q + k).sum(dim=-1), + ) + + monkeypatch.setattr(p2p_check_script, "_strict_attention_core", _RowCore) + q = torch.randn(2, 1, 3, 4, requires_grad=True) + k = torch.randn(2, 1, 3, 4, requires_grad=True) + v = torch.randn(2, 1, 3, 4, requires_grad=True) + positions = torch.arange(3).expand(2, -1) + + result = p2p_check_script._strict_attention_reference_rows( + q, + k, + v, + positions=positions, + output_dtype=q.dtype, + ) + + assert [batch for batch, _positions in calls] == [1, 1] + assert [item.tolist() for _batch, item in calls] == [[[0, 1, 2]], [[0, 1, 2]]] + assert torch.equal(result.out, q + k + v) + assert torch.equal(result.lse, (q + k).sum(dim=-1)) + (result.out.sum() + result.lse.sum()).backward() + assert q.grad is not None + assert k.grad is not None + assert v.grad is not None + + +def _strict_acceptance_provenance(**overrides): + provenance = { + "strict_core_id": STRICT_ATTENTION_PRODUCTION_CORE_ID, + "strict_schedule": STRICT_ATTENTION_FA4_SCHEDULE_ID, + "attention_backend": "flash_attention_4.cute", + "actual_backend": "flash_attention_4.cute", + "rope_backend": "rlkernel.cuda.rope_sm90", + "strict_mode": True, + "native_attention_arithmetic": True, + "num_splits": 1, + "deterministic_backward": True, + "reference_only": False, + "fa_api_source": "flash_attn.cute.interface", + "fallback": False, + "strict_split_kv": "disabled", + "strict_comm_autograd": True, + "communication_backend": "self_owned_cuda_ag_rs", + "production_ready": True, + "strict_full_qkv_all_gather": True, + "strict_position_ids_all_gather": True, + "compute_communication": "decoupled", + "compute_schedule": STRICT_ATTENTION_RING_SCHEDULE_ID, + "communication_overlap": "disabled", + "ring_schedule_default": True, + "ring_partial_arithmetic": False, + "rope_fusion": False, + "q_rope_state": "post_rope", + "k_cache_rope_state": "post_rope", + } + provenance.update(overrides) + return provenance + + +@pytest.mark.parametrize( + ("field", "invalid"), + [ + ("compute_schedule", "dynamic_ring"), + ("communication_overlap", "enabled"), + ("ring_schedule_default", False), + ("ring_partial_arithmetic", True), + ("actual_backend", "flashinfer"), + ("rope_backend", "native_rope"), + ("strict_comm_autograd", False), + ], +) +def test_strict_shared_core_acceptance_rejects_provenance_drift(field, invalid): + errors = p2p_check_script._strict_shared_core_identity_errors( + _strict_acceptance_provenance(**{field: invalid}), + transport="cuda_ag_rs", + is_rocm=False, + ) + + assert any(error.startswith(f"{field}=") for error in errors) + + +def test_strict_shared_core_acceptance_requires_rocm_backend_and_rope(): + provenance = _strict_acceptance_provenance( + strict_core_id=STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + strict_schedule=STRICT_ATTENTION_ROCM_SCHEDULE_ID, + attention_backend="aiter.rocm.ck_dense_mha", + actual_backend="aiter.rocm.ck_dense_mha", + rope_backend="rlkernel.rocm.deterministic_rope", + communication_backend="rccl_ag_rs", + split_kv_control="dense_non_split_api", + aiter_api_source="aiter.ops.mha", + aiter_source_sha256="a" * 64, + ) + + assert not p2p_check_script._strict_shared_core_identity_errors( + provenance, + transport="rccl_ag_rs", + is_rocm=True, + ) + + @pytest.mark.parametrize( ("argv", "message"), [ @@ -2007,6 +2208,10 @@ def fake_fa4( assert result.lse.dtype is torch.float32 +# FA4 is the CUDA production core. On ROCm the strict default correctly resolves +# to the AITER CK core instead, so the StrictFlashAttention4Core monkeypatch below +# is never consulted and the assertions cannot hold there. +@pytest.mark.cuda_only def test_strict_paged_default_selects_fa4_production_core(monkeypatch): core = _RecordingStrictCore() core.core_id = STRICT_ATTENTION_PRODUCTION_CORE_ID diff --git a/tests/test_qwen_ffn.py b/tests/test_qwen_ffn.py index 63b994fe..e461c97f 100644 --- a/tests/test_qwen_ffn.py +++ b/tests/test_qwen_ffn.py @@ -137,6 +137,48 @@ def _close_ffn_collectives() -> None: ffn_module._COLLECTIVES.clear() +def test_ffn_collective_creation_uses_platform_factory(monkeypatch): + import rl_engine.distributed.collectives as collectives + + sentinel = object() + calls = [] + + def fake_factory(**kwargs): + calls.append(kwargs) + return sentinel + + monkeypatch.setattr(collectives, "create_deterministic_collective", fake_factory) + monkeypatch.setattr(collectives.dist, "get_rank", lambda group: 0) + monkeypatch.setattr(collectives.dist, "get_world_size", lambda group: 2) + monkeypatch.setattr(collectives.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(ffn_module, "_COLLECTIVE_MIN_CAPACITY_BYTES", 2048) + group = object() + + ffn_module._COLLECTIVES.clear() + result = ffn_module._collective_for_group(group, min_size_bytes=1234) + + assert result is sentinel + assert calls == [{"group": group, "device": 0, "max_size_bytes": 2048}] + ffn_module._COLLECTIVES.clear() + + +def test_packed_tp_inference_fails_closed_on_rocm(monkeypatch): + class _FakeDist: + @staticmethod + def get_world_size(*, group): + return 2 + + monkeypatch.setattr(ffn_module, "_require_parallel_group", lambda group, name: _FakeDist()) + monkeypatch.setattr(torch.version, "hip", "6.3", raising=False) + + with pytest.raises(RuntimeError, match="not available with the ROCm/RCCL transport"): + ffn_module.Qwen3FFNOp().prepare_packed_inference( + torch.empty(2, 2), + torch.empty(2, 2), + tp_group=object(), + ) + + def _shard_ranges( rank: int, *, diff --git a/tests/test_rocm_aiter_api_contract.py b/tests/test_rocm_aiter_api_contract.py new file mode 100644 index 00000000..f7ef6dce --- /dev/null +++ b/tests/test_rocm_aiter_api_contract.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Schema validation for the AITER entry points the strict ROCm core calls. + +The CUDA core validates the FA4 CuTe API by parameter name before it runs. +AITER hides its signature behind a JIT wrapper (``inspect.signature`` reports +``(*args, **kwargs)``), so the equivalent check reads the registered Torch +schema. The strict calls are positional, which makes argument *order* part of +the contract too: an upstream insertion would silently reinterpret everything +after it while every call still type-checks. +""" + +from __future__ import annotations + +import pytest + +from rl_engine.kernels.ops.rocm.attention.flash_attn import ( + _AITER_BWD_POSITIONAL_CONTRACT, + _AITER_BWD_REQUIRED_KEYWORDS, + _AITER_FWD_POSITIONAL_CONTRACT, + StrictRocmAttentionUnavailable, + _validate_aiter_schema, +) + + +def _aiter_available() -> bool: + try: + import aiter.ops.mha # noqa: F401 + except Exception: + return False + return True + + +requires_aiter = pytest.mark.skipif(not _aiter_available(), reason="AITER is not installed") + + +def test_positional_contract_matches_the_strict_call_sites() -> None: + """The tuples must stay in step with what the autograd Function passes. + + ``_AiterCKAttentionFn`` calls both ops positionally. If someone edits a + call site without editing the contract, the schema check would still pass + while the call means something else. + """ + + assert _AITER_FWD_POSITIONAL_CONTRACT[:6] == ( + "q", + "k", + "v", + "dropout_p", + "softmax_scale", + "is_causal", + ) + # The forward passes (-1, -1, 0, True, False) after is_causal. + assert _AITER_FWD_POSITIONAL_CONTRACT[6:] == ( + "window_size_left", + "window_size_right", + "sink_size", + "return_softmax_lse", + "return_dropout_randval", + ) + # The backward pins determinism positionally, so its slot must not move. + assert _AITER_BWD_POSITIONAL_CONTRACT[-1] == "deterministic" + assert _AITER_BWD_POSITIONAL_CONTRACT.index("softmax_lse") == 5 + assert "rng_state" in _AITER_BWD_REQUIRED_KEYWORDS + + +@requires_aiter +def test_installed_aiter_satisfies_the_strict_contract() -> None: + _validate_aiter_schema("mha_fwd", _AITER_FWD_POSITIONAL_CONTRACT) + _validate_aiter_schema( + "mha_bwd", + _AITER_BWD_POSITIONAL_CONTRACT, + required_keywords=_AITER_BWD_REQUIRED_KEYWORDS, + ) + + +@requires_aiter +def test_reordered_positional_contract_fails_closed() -> None: + """A swap the installed schema does not have must be rejected.""" + + swapped = ("q", "k", "v", "softmax_scale", "dropout_p") + with pytest.raises(StrictRocmAttentionUnavailable, match="positional contract changed"): + _validate_aiter_schema("mha_fwd", swapped) + + +@requires_aiter +def test_missing_keyword_control_fails_closed() -> None: + with pytest.raises(StrictRocmAttentionUnavailable, match="missing strict controls"): + _validate_aiter_schema( + "mha_fwd", + _AITER_FWD_POSITIONAL_CONTRACT, + required_keywords=frozenset({"num_splits"}), + ) + + +def test_unregistered_operator_fails_closed() -> None: + with pytest.raises(StrictRocmAttentionUnavailable, match="cannot read the Torch schema"): + _validate_aiter_schema("mha_fwd_that_does_not_exist", ("q",)) diff --git a/tests/test_rocm_collective_benchmark.py b/tests/test_rocm_collective_benchmark.py new file mode 100644 index 00000000..e4eb8643 --- /dev/null +++ b/tests/test_rocm_collective_benchmark.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest +import torch + +_BENCHMARK_PATH = Path(__file__).parents[1] / "benchmarks" / "benchmark_rocm_collectives.py" +_SPEC = importlib.util.spec_from_file_location( + "rlkernel_rocm_collective_benchmark", _BENCHMARK_PATH +) +assert _SPEC is not None and _SPEC.loader is not None +benchmark = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(benchmark) + + +def test_rocm_collective_benchmark_parser() -> None: + args = benchmark.parse_args( + [ + "--size-bytes", + "1024", + "4096", + "--dtype", + "fp32", + "--operations", + "all_reduce", + "reduce_scatter", + "--warmup", + "2", + "--iterations", + "3", + "--samples", + "4", + ] + ) + + benchmark._validate_args(args) + assert args.size_bytes == [1024, 4096] + assert args.dtype == "fp32" + assert args.operations == ["all_reduce", "reduce_scatter"] + assert (args.warmup, args.iterations, args.samples) == (2, 3, 4) + + +@pytest.mark.parametrize( + "argv", + ( + ["--size-bytes", "0"], + ["--warmup", "-1"], + ["--iterations", "0"], + ["--samples", "0"], + ), +) +def test_rocm_collective_benchmark_rejects_invalid_counts(argv: list[str]) -> None: + with pytest.raises(ValueError): + benchmark._validate_args(benchmark.parse_args(argv)) + + +def test_rocm_collective_benchmark_aligns_reduce_scatter_input() -> None: + tensor, actual_bytes = benchmark._make_inputs( + size_bytes=35, + dtype=torch.float32, + world_size=4, + rank=2, + device=torch.device("cpu"), + ) + + assert tensor.is_contiguous() + assert tensor.numel() % 4 == 0 + assert actual_bytes == tensor.numel() * tensor.element_size() diff --git a/tests/test_rocm_strict_paged_attention.py b/tests/test_rocm_strict_paged_attention.py new file mode 100644 index 00000000..53ef1463 --- /dev/null +++ b/tests/test_rocm_strict_paged_attention.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Decode-stage paged Attention on the strict ROCm runtime. + +The core is injected, so these run without ROCm. What they pin is the part +that is ours: the page table decides logical KV order, the cached rows reach +the core exactly as a dense prefill over the same tokens would, and the +provenance never claims a native paged kernel. +""" + +from __future__ import annotations + +from typing import Any + +import pytest +import torch + +from rl_engine.kernels.attention_contract import ( + STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID, + STRICT_ATTENTION_ROCM_SCHEDULE_ID, +) +from rl_engine.kernels.ops.rocm.attention.strict_runtime import StrictRocmAttentionRuntime + +_HEAD_DIM = 8 +_PAGE_SIZE = 4 + + +class _RecordingCore: + """Dense core stand-in that records exactly what each launch consumed.""" + + core_id = STRICT_ATTENTION_ROCM_PRODUCTION_CORE_ID + strict_schedule = STRICT_ATTENTION_ROCM_SCHEDULE_ID + backend_id = "aiter.rocm.ck_dense_mha" + + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def forward_with_lse(self, q, k, v, **kwargs) -> Any: + self.calls.append( + { + "q": q, + "k": k.clone(), + "v": v.clone(), + "causal": kwargs.get("causal"), + "query_position_ids": kwargs.get("query_position_ids"), + "key_position_ids": kwargs.get("key_position_ids"), + } + ) + + class _Result: + out = torch.zeros(q.size(0), q.size(1), q.size(2), _HEAD_DIM, dtype=q.dtype) + lse = torch.zeros(q.size(0), q.size(1), q.size(2), dtype=torch.float32) + provenance = {"attention_backend": "aiter.rocm.ck_dense_mha"} + + return _Result() + + +def _runtime() -> StrictRocmAttentionRuntime: + return StrictRocmAttentionRuntime(core=_RecordingCore()) + + +def _cache(pages: int, kv_heads: int = 1) -> torch.Tensor: + total = pages * _PAGE_SIZE * kv_heads * _HEAD_DIM + return ( + torch.arange(total, dtype=torch.float32) + .reshape(pages, _PAGE_SIZE, kv_heads, _HEAD_DIM) + .to(torch.bfloat16) + ) + + +def _paged_call(runtime, *, page_table, seqused_k, q_heads=1, kv_heads=1, pages=4): + k_cache = _cache(pages, kv_heads) + v_cache = _cache(pages, kv_heads) + 1 + q = torch.zeros(page_table.size(0), q_heads, 1, _HEAD_DIM, dtype=torch.bfloat16) + return ( + runtime.forward_paged_with_lse( + q, + k_cache, + v_cache, + page_table=page_table, + seqused_k=seqused_k, + max_seqlen_k=page_table.size(1) * _PAGE_SIZE, + scale=None, + ), + k_cache, + v_cache, + ) + + +@pytest.fixture(autouse=True) +def _pretend_rocm(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(StrictRocmAttentionRuntime, "_require_rocm", staticmethod(lambda t: None)) + + +def test_paged_decode_gathers_kv_in_logical_not_physical_order() -> None: + """A shuffled page table must still produce logical KV order. + + This is the property that makes decode replay comparable with prefill: if + physical page order leaked through, the same logical sequence would produce + different arithmetic depending on how the allocator handed out pages. + """ + + core = _RecordingCore() + runtime = StrictRocmAttentionRuntime(core=core) + # Logical tokens 0..7 live on physical pages 3 then 1. + page_table = torch.tensor([[3, 1]], dtype=torch.int32) + _result, k_cache, _v_cache = _paged_call( + runtime, + page_table=page_table, + seqused_k=torch.tensor([8], dtype=torch.int32), + ) + + assert len(core.calls) == 1 + gathered_k = core.calls[0]["k"] + assert gathered_k.shape == (1, 1, 8, _HEAD_DIM) + + expected = torch.cat((k_cache[3], k_cache[1]), dim=0) # [8, H, D] logical order + assert torch.equal(gathered_k[0].permute(1, 0, 2), expected) + + +def test_paged_decode_truncates_to_the_cached_length() -> None: + core = _RecordingCore() + runtime = StrictRocmAttentionRuntime(core=core) + page_table = torch.tensor([[0, 1]], dtype=torch.int32) + + _paged_call( + runtime, + page_table=page_table, + seqused_k=torch.tensor([5], dtype=torch.int32), + ) + + # Five cached tokens span two pages but must not expose the page tail. + assert core.calls[0]["k"].shape == (1, 1, 5, _HEAD_DIM) + assert core.calls[0]["v"].shape == (1, 1, 5, _HEAD_DIM) + # The launch is non-causal, so the core is handed no position ids; the + # truncation above is what bounds the launch to the cached prefix. + assert core.calls[0]["key_position_ids"] is None + + +def test_paged_decode_is_not_causal_within_a_launch() -> None: + """Decode attends over the whole cached prefix, so the launch is not causal.""" + + core = _RecordingCore() + runtime = StrictRocmAttentionRuntime(core=core) + + _paged_call( + runtime, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([4], dtype=torch.int32), + ) + + assert core.calls[0]["causal"] is False + + +def test_paged_decode_keeps_one_kv_group_per_launch() -> None: + """The TP-degree invariance mechanism must survive into the paged path.""" + + core = _RecordingCore() + runtime = StrictRocmAttentionRuntime(core=core) + + result, _k, _v = _paged_call( + runtime, + page_table=torch.tensor([[0], [1]], dtype=torch.int32), + seqused_k=torch.tensor([4, 4], dtype=torch.int32), + q_heads=4, + kv_heads=2, + ) + + # Two rows x two KV groups. + assert len(core.calls) == 4 + assert result.provenance["core_launch_count"] == 4 + for call in core.calls: + assert call["k"].size(1) == 1 # exactly one KV group per launch + assert call["q"].size(1) == 2 # its two Q heads + assert result.provenance["launch_granularity"] == "one_batch_row_one_kv_group" + assert result.provenance["tp_degree_invariant"] is True + + +def test_paged_decode_provenance_does_not_claim_a_paged_kernel() -> None: + """The gather is the implementation; the provenance must say so.""" + + runtime = _runtime() + result, _k, _v = _paged_call( + runtime, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([4], dtype=torch.int32), + ) + + assert result.provenance["paged_kernel"] == "none" + assert result.provenance["paged_execution"] == "logical_kv_gather_then_dense_core" + assert result.provenance["split_kv"] == "disabled" + assert result.provenance["strict_schedule"] == STRICT_ATTENTION_ROCM_SCHEDULE_ID + assert result.provenance["communication_executed"] is False + assert result.provenance["query_schedule"] == "paged_single_query_batch" + + +@pytest.mark.parametrize( + ("page_table", "seqused_k", "match"), + [ + (torch.tensor([[9]], dtype=torch.int32), torch.tensor([4], dtype=torch.int32), "outside"), + (torch.tensor([[0]], dtype=torch.int32), torch.tensor([0], dtype=torch.int32), "positive"), + ( + torch.tensor([[0]], dtype=torch.int32), + torch.tensor([9], dtype=torch.int32), + "within max_seqlen_k", + ), + ], +) +def test_paged_decode_fails_closed_on_bad_metadata(page_table, seqused_k, match) -> None: + runtime = _runtime() + with pytest.raises(ValueError, match=match): + _paged_call(runtime, page_table=page_table, seqused_k=seqused_k) + + +def test_paged_decode_rejects_a_mismatched_out_buffer() -> None: + runtime = _runtime() + k_cache = _cache(2) + q = torch.zeros(1, 1, 1, _HEAD_DIM, dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="same shape as q"): + runtime.forward_paged_with_lse( + q, + k_cache, + k_cache + 1, + page_table=torch.tensor([[0]], dtype=torch.int32), + seqused_k=torch.tensor([4], dtype=torch.int32), + max_seqlen_k=_PAGE_SIZE, + scale=None, + out=torch.zeros(2, 1, 1, _HEAD_DIM, dtype=torch.bfloat16), + ) + + +def test_rocm_registry_does_not_claim_decode_before_a_caller_routes_to_it() -> None: + """The paged entry point exists, but nothing dispatches to it yet. + + The Vime provider always calls ``forward_with_lse`` and builds its contract + with ``kv_cache=None``, so no decode request can reach the paged path. + Claiming the mode here would let the cross-config binding accept a decode + path that never runs. Flip this together with the dispatch wiring. + """ + + from rl_engine.kernels.attention_contract import AttentionMode + from rl_engine.kernels.registry import KernelRegistry, OpBackend + + capabilities = KernelRegistry()._attention_capabilities + capability = capabilities.get(OpBackend.ROCM_STRICT_ATTENTION) + if capability is None: + pytest.skip("AITER is unavailable, so the strict ROCm backend is not registered") + + assert AttentionMode.DECODE not in capability.modes + assert capability.supports_kv_cache is False + # Whatever the modes, the gather must keep the Split-KV claims intact. + assert capability.supports_split_kv_disabled is True + assert capability.supports_split_kv_fixed is False + assert capability.supports_split_kv_auto is False diff --git a/tests/test_triton_deterministic_attention.py b/tests/test_triton_deterministic_attention.py new file mode 100644 index 00000000..8e87ecf9 --- /dev/null +++ b/tests/test_triton_deterministic_attention.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Bitwise parity between the Triton and native deterministic Attention cores. + +The Triton core in ``rl_engine.kernels.ops.triton.attention.deterministic_attn`` is a +port of ``csrc/cuda/attention/deterministic_attention.cu``. Its contract is stronger +than "numerically close": every tensor it returns must be bit-identical to the native +kernel's. These tests pin that contract, plus the two properties the port depends on: +the vendor-exact ``expf``/``logf`` helpers, and batch invariance. +""" + +import math + +import pytest +import torch + +from rl_engine.platforms.device import device_ctx + +IS_ROCM = device_ctx.is_rocm +IS_GPU = device_ctx.device_type == "cuda" or IS_ROCM + +pytestmark = pytest.mark.skipif(not IS_GPU, reason="CUDA/ROCm GPU not available") + +try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + from rl_engine.kernels.ops.cuda.attention.deterministic_attn import ( + DeterministicAttentionOp, + ) + from rl_engine.kernels.ops.triton.attention.deterministic_attn import ( + BITWISE_LIBM_PARITY, + TritonDeterministicAttentionOp, + triton_deterministic_attention_backward, + triton_deterministic_attention_forward, + ) + + _IMPORTED = True +except (ImportError, RuntimeError): # pragma: no cover - import guard + _IMPORTED = False + +_NATIVE = _IMPORTED and _EXT_AVAILABLE and hasattr(_C, "deterministic_attention_forward") + +needs_native = pytest.mark.skipif( + not _NATIVE, reason="native deterministic attention kernel not built" +) +needs_bitwise_libm = pytest.mark.skipif( + not (_IMPORTED and BITWISE_LIBM_PARITY), + reason="bitwise expf/logf sequence is only ported for ROCm", +) + +DEVICE = device_ctx.device +D = 128 + +# (B, Hq, Hkv, Sq, Skv, causal, mask_kind, scale) +_CASES = [ + (1, 1, 1, 1, 1, True, None, None), + (1, 2, 2, 1, 64, True, None, None), # decode step + (1, 8, 2, 64, 64, True, None, None), # GQA, group 4 + (2, 4, 1, 128, 128, True, None, None), # MQA + (1, 2, 2, 256, 256, True, None, None), # exactly one softmax lane chunk + (1, 2, 2, 257, 257, True, None, None), # one past the chunk boundary + (1, 2, 2, 100, 512, False, None, None), # two full lane chunks + (2, 4, 2, 64, 700, True, None, None), # ragged multi-chunk + (1, 2, 2, 32, 32, True, "right", None), + (1, 2, 2, 32, 32, False, "left", None), + (2, 2, 2, 16, 300, True, "right", None), + (2, 2, 2, 8, 8, True, "allfalse", None), # fully masked row -> lse == -inf + (1, 2, 2, 16, 16, True, None, 0.0), + (1, 2, 2, 16, 16, True, None, 3.7), + (1, 2, 2, 16, 16, False, None, -1.25), +] + + +def _make_inputs(case, dtype, seed): + b, hq, hkv, sq, skv, causal, mask_kind, scale = case + gen = torch.Generator(device=DEVICE).manual_seed(seed) + q = torch.randn(b, hq, sq, D, device=DEVICE, dtype=dtype, generator=gen) + k = torch.randn(b, hkv, skv, D, device=DEVICE, dtype=dtype, generator=gen) + v = torch.randn(b, hkv, skv, D, device=DEVICE, dtype=dtype, generator=gen) + + mask = None + if mask_kind is not None: + mask = torch.ones(b, skv, device=DEVICE, dtype=torch.bool) + if mask_kind == "right": + mask[:, skv // 2 :] = False + elif mask_kind == "left": + mask[:, : skv // 3] = False + elif mask_kind == "allfalse": + mask[0, :] = False + else: # pragma: no cover - guards the parametrisation itself + raise AssertionError(f"unknown mask kind {mask_kind}") + + resolved_scale = 1.0 / math.sqrt(D) if scale is None else scale + return q, k, v, causal, resolved_scale, mask + + +def _case_id(case): + b, hq, hkv, sq, skv, causal, mask_kind, scale = case + return f"b{b}_hq{hq}_hkv{hkv}_sq{sq}_skv{skv}_causal{int(causal)}_{mask_kind}_{scale}" + + +@needs_native +@needs_bitwise_libm +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +@pytest.mark.parametrize("case", _CASES, ids=[_case_id(c) for c in _CASES]) +def test_triton_forward_is_bitwise_identical_to_native(case, dtype): + q, k, v, causal, scale, mask = _make_inputs(case, dtype, seed=11) + + ref_out, ref_lse, ref_p = _C.deterministic_attention_forward(q, k, v, causal, scale, mask) + out, lse, p = triton_deterministic_attention_forward(q, k, v, causal, scale, mask, False) + + assert torch.equal(out, ref_out) + assert torch.equal(lse, ref_lse) + assert torch.equal(p, ref_p) + + +@needs_native +@needs_bitwise_libm +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +@pytest.mark.parametrize("case", _CASES, ids=[_case_id(c) for c in _CASES]) +def test_triton_backward_is_bitwise_identical_to_native(case, dtype): + q, k, v, causal, scale, mask = _make_inputs(case, dtype, seed=12) + + _ref_out, _ref_lse, ref_p = _C.deterministic_attention_forward(q, k, v, causal, scale, mask) + _out, _lse, p = triton_deterministic_attention_forward(q, k, v, causal, scale, mask, False) + grad_out = torch.randn_like(q) + + ref_dq, ref_dk, ref_dv = _C.deterministic_attention_backward( + grad_out, q, k, v, ref_p, causal, scale, mask + ) + dq, dk, dv = triton_deterministic_attention_backward(grad_out, q, k, v, p, scale) + + assert torch.equal(dq, ref_dq) + assert torch.equal(dk, ref_dk) + assert torch.equal(dv, ref_dv) + + +@needs_native +@needs_bitwise_libm +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +def test_triton_autograd_matches_native_autograd_bitwise(dtype): + b, hq, hkv, sq, skv = 2, 8, 2, 96, 300 + gen = torch.Generator(device=DEVICE).manual_seed(13) + base = ( + torch.randn(b, hq, sq, D, device=DEVICE, dtype=dtype, generator=gen), + torch.randn(b, hkv, skv, D, device=DEVICE, dtype=dtype, generator=gen), + torch.randn(b, hkv, skv, D, device=DEVICE, dtype=dtype, generator=gen), + ) + grad_out = torch.randn(b, hq, sq, D, device=DEVICE, dtype=dtype, generator=gen) + + results = [] + for op in (DeterministicAttentionOp(), TritonDeterministicAttentionOp()): + tensors = [t.clone().requires_grad_(True) for t in base] + out, lse = op.forward_with_lse(*tensors, causal=True) + out.backward(grad_out) + results.append((out.detach(), lse, [t.grad for t in tensors])) + + native, triton_result = results + assert torch.equal(triton_result[0], native[0]) + assert torch.equal(triton_result[1], native[1]) + for triton_grad, native_grad in zip(triton_result[2], native[2]): + assert torch.equal(triton_grad, native_grad) + + +@needs_bitwise_libm +def test_triton_fp32_output_downcasts_to_the_native_dtype_result(): + """``output_fp32=True`` must expose the same accumulator the bf16 path rounds.""" + q, k, v, causal, scale, mask = _make_inputs(_CASES[6], torch.bfloat16, seed=14) + + out, _lse, _p = triton_deterministic_attention_forward(q, k, v, causal, scale, mask, False) + out_fp32, _lse32, _p32 = triton_deterministic_attention_forward( + q, k, v, causal, scale, mask, True + ) + + assert out_fp32.dtype is torch.float32 + assert torch.equal(out_fp32.to(torch.bfloat16), out) + + +@needs_bitwise_libm +def test_triton_fully_masked_row_is_zero_with_neg_inf_lse(): + q, k, v, causal, scale, mask = _make_inputs(_CASES[11], torch.bfloat16, seed=15) + + out, lse, p = triton_deterministic_attention_forward(q, k, v, causal, scale, mask, False) + + assert torch.equal(out[0], torch.zeros_like(out[0])) + assert torch.equal(p[0], torch.zeros_like(p[0])) + assert torch.isneginf(lse[0]).all() + assert torch.isfinite(lse[1]).all() + + +@needs_bitwise_libm +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +def test_triton_batch_slice_is_bitwise_invariant(dtype): + """A row's result must not depend on what else was in the batch.""" + b, hq, hkv, sq, skv = 4, 4, 2, 48, 192 + gen = torch.Generator(device=DEVICE).manual_seed(16) + q = torch.randn(b, hq, sq, D, device=DEVICE, dtype=dtype, generator=gen) + k = torch.randn(b, hkv, skv, D, device=DEVICE, dtype=dtype, generator=gen) + v = torch.randn(b, hkv, skv, D, device=DEVICE, dtype=dtype, generator=gen) + scale = 1.0 / math.sqrt(D) + + batched, batched_lse, _ = triton_deterministic_attention_forward( + q, k, v, True, scale, None, False + ) + single, single_lse, _ = triton_deterministic_attention_forward( + q[2:3], k[2:3], v[2:3], True, scale, None, False + ) + + assert torch.equal(single[0], batched[2]) + assert torch.equal(single_lse[0], batched_lse[2]) + + +@needs_bitwise_libm +def test_expf_and_logf_match_the_vendor_libm_bitwise(): + """The softmax parity rests on these two helpers; pin them independently.""" + import triton + + import triton.language as tl # isort: skip + from rl_engine.kernels.ops.triton.attention import deterministic_attn as mod + + @triton.jit + def _exp_probe(x_ptr, out_ptr, n_elem, EXPF: tl.constexpr): + offs = tl.program_id(0) * 256 + tl.arange(0, 256) + keep = offs < n_elem + value = tl.load(x_ptr + offs, mask=keep, other=0.0) + tl.store(out_ptr + offs, EXPF(value), mask=keep) + + gen = torch.Generator(device=DEVICE).manual_seed(17) + edge = torch.tensor( + [0.0, -0.0, 1.0, -1.0, 88.72283935546875, 88.73, -103.2789306640625, -104.0], + device=DEVICE, + dtype=torch.float32, + ) + xs = torch.cat( + [torch.rand(1 << 18, device=DEVICE, generator=gen) * 240.0 - 130.0, edge] + ).contiguous() + got = torch.empty_like(xs) + _exp_probe[(triton.cdiv(xs.numel(), 256),)](xs, got, xs.numel(), EXPF=mod._expf) + assert torch.equal(got, torch.exp(xs)) + + positives = torch.cat( + [ + torch.rand(1 << 18, device=DEVICE, generator=gen) * 1e3, + torch.rand(1 << 12, device=DEVICE, generator=gen) * 1e-38, # subnormal inputs + torch.tensor([1.0, 1e-45, 3.4e38], device=DEVICE), + ] + ).contiguous() + got = torch.empty_like(positives) + _exp_probe[(triton.cdiv(positives.numel(), 256),)]( + positives, got, positives.numel(), EXPF=mod._logf + ) + assert torch.equal(got, torch.log(positives)) + + +def test_op_refuses_to_run_without_a_bitwise_libm(): + """On a platform with no ported expf/logf the op must fail loudly, not silently.""" + if BITWISE_LIBM_PARITY: + op = TritonDeterministicAttentionOp() + assert op.bitwise_libm is True + else: # pragma: no cover - exercised on CUDA + with pytest.raises(RuntimeError, match="bitwise-identical"): + TritonDeterministicAttentionOp() + assert TritonDeterministicAttentionOp(require_bitwise_libm=False).bitwise_libm is False + + +@needs_bitwise_libm +@pytest.mark.parametrize( + "kwargs, message", + [ + ({"head_dim": 64}, "head dim D must be 128"), + ({"dtype": torch.float32}, "only FP16/BF16 supported"), + ({"gqa_mismatch": True}, "not divisible"), + ], +) +def test_triton_op_validation_mirrors_the_native_op(kwargs, message): + head_dim = kwargs.get("head_dim", D) + dtype = kwargs.get("dtype", torch.bfloat16) + hkv = 3 if kwargs.get("gqa_mismatch") else 2 + q = torch.randn(1, 2, 8, head_dim, device=DEVICE, dtype=dtype) + k = torch.randn(1, hkv, 8, head_dim, device=DEVICE, dtype=dtype) + v = torch.randn(1, hkv, 8, head_dim, device=DEVICE, dtype=dtype) + + with pytest.raises(ValueError, match=message): + TritonDeterministicAttentionOp().forward(q, k, v, causal=True) diff --git a/tests/test_vime_attention_provider.py b/tests/test_vime_attention_provider.py new file mode 100644 index 00000000..ec94bff7 --- /dev/null +++ b/tests/test_vime_attention_provider.py @@ -0,0 +1,368 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Coverage for the optional Vime WS2 strict attention adapter. + +The validation-shaped cases need a real ROCm device with AITER present and are +skipped elsewhere. The contract/fail-closed cases are pure metadata checks and +run anywhere, so a CUDA or CPU CI job still catches a provider that silently +widens what it accepts. +""" + +from __future__ import annotations + +import math +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.integrations.vime.attention import AttentionProviderUnavailable, attention_provider +from rl_engine.kernels.registry import _rocm_strict_attention_available + +STRICT_ROCM = _rocm_strict_attention_available() +requires_strict_rocm = pytest.mark.skipif( + not STRICT_ROCM, + reason="strict ROCm attention requires a ROCm device with aiter.ops.mha", +) + +# Qwen3-8B dense head layout, TP=1 local view. +GLOBAL_Q_HEADS = 32 +GLOBAL_KV_HEADS = 8 +HEAD_DIM = 128 + + +def _metadata(**overrides): + metadata = { + "global_q_heads": GLOBAL_Q_HEADS, + "global_kv_heads": GLOBAL_KV_HEADS, + "tp_rank": 0, + "tp_world_size": 1, + "attention_mode": "prefill", + "role": "train", + "causal": True, + } + metadata.update(overrides) + return metadata + + +def _request( + *, + batch_size: int = 1, + seq_len: int = 128, + query_len: int | None = None, + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", + cp_world_size: int = 1, + cp_rank: int = 0, + cp_layout: str = "single", + seed: int = 0, + **metadata_overrides, +): + generator = torch.Generator(device=device).manual_seed(seed) + shape_q = (batch_size, GLOBAL_Q_HEADS, query_len or seq_len, HEAD_DIM) + shape_kv = (batch_size, GLOBAL_KV_HEADS, seq_len, HEAD_DIM) + return SimpleNamespace( + query=torch.randn(*shape_q, generator=generator, device=device, dtype=dtype), + key=torch.randn(*shape_kv, generator=generator, device=device, dtype=dtype), + value=torch.randn(*shape_kv, generator=generator, device=device, dtype=dtype), + key_padding_mask=None, + tensor_parallel_group=None, + context_parallel=SimpleNamespace(world_size=cp_world_size, rank=cp_rank, layout=cp_layout), + metadata=_metadata(**metadata_overrides), + ) + + +def _cpu_request(**kwargs): + """A request whose rejection happens before any device work.""" + + kwargs.setdefault("device", "cpu") + kwargs.setdefault("seq_len", 8) + return _request(**kwargs) + + +# --------------------------------------------------------------------------- +# Contract and fail-closed behavior (device independent) +# --------------------------------------------------------------------------- + + +def test_cp_zigzag_layout_fails_closed(): + """Only the contiguous-per-rank layout matches the strict CP block plan.""" + + request = _cpu_request(cp_world_size=2, cp_rank=1, cp_layout="zigzag") + + with pytest.raises(AttentionProviderUnavailable, match="requires the 'allgather' layout"): + attention_provider(request) + + +def test_cp_without_a_process_group_fails_closed(): + """CP>1 needs a group to build the RCCL transport; it must not fall back.""" + + request = _cpu_request(cp_world_size=2, cp_rank=1, cp_layout="allgather") + + with pytest.raises(AttentionProviderUnavailable, match="context_parallel_group"): + attention_provider(request) + + +def test_cp_contract_describes_one_contiguous_block_per_rank(): + """The CP sharding must place this rank's block at its global offset.""" + + from rl_engine.integrations.vime.attention import _contract_for_request + + seq_len = 8 + request = _cpu_request( + seq_len=seq_len, + cp_world_size=4, + cp_rank=2, + cp_layout="allgather", + ) + contract, *_ = _contract_for_request(request) + sharding = contract.sharding + + assert sharding.cp_world_size == 4 + assert sharding.global_sequence_length == seq_len * 4 + assert sharding.local_sequence_length == seq_len + assert sharding.global_block_indices == (2,) + assert sharding.global_block_token_starts == (2 * seq_len,) + assert sharding.local_block_offsets == (0, seq_len) + + +def test_decode_without_kv_cache_identity_fails_closed(): + request = _cpu_request(attention_mode="decode") + + with pytest.raises(AttentionProviderUnavailable, match="KV-cache identity"): + attention_provider(request) + + +@pytest.mark.parametrize( + "overrides", + [ + {"dropout_p": 0.1}, + {"sliding_window": 128}, + {"logit_soft_cap": 30.0}, + {"alibi_slopes": [0.1]}, + {"window_size": (256, 0)}, + ], +) +def test_distribution_changing_knobs_fail_closed(overrides): + request = _cpu_request(**overrides) + + with pytest.raises(AttentionProviderUnavailable): + attention_provider(request) + + +def test_key_padding_mask_is_refused(): + request = _cpu_request() + request.key_padding_mask = torch.ones(1, 8, dtype=torch.bool) + + with pytest.raises(AttentionProviderUnavailable, match="unpadded logical row"): + attention_provider(request) + + +def test_fp32_is_refused(): + request = _cpu_request(dtype=torch.float32) + + with pytest.raises(AttentionProviderUnavailable, match="BF16/FP16"): + attention_provider(request) + + +def test_head_counts_must_cover_the_tp_group_exactly(): + request = _cpu_request(global_q_heads=GLOBAL_Q_HEADS * 2) + + with pytest.raises(AttentionProviderUnavailable, match="do not cover global_q_heads"): + attention_provider(request) + + +def test_declared_tp_rank_must_agree_with_the_group(): + request = _cpu_request(tp_rank=3) + + with pytest.raises(AttentionProviderUnavailable, match="disagrees with TP group rank"): + attention_provider(request) + + +def test_cp_layout_must_describe_local_ownership(): + request = _cpu_request(cp_layout="unknown") + + with pytest.raises(AttentionProviderUnavailable, match="local CP token ownership"): + attention_provider(request) + + +def test_non_contiguous_key_positions_are_refused(): + request = _cpu_request(key_position_ids=[0, 1, 2, 3, 9, 10, 11, 12]) + + with pytest.raises(AttentionProviderUnavailable, match="contiguous increasing"): + attention_provider(request) + + +# --------------------------------------------------------------------------- +# Strict arithmetic (requires a ROCm device with AITER) +# --------------------------------------------------------------------------- + + +@requires_strict_rocm +def test_provider_exports_attention_lse_and_strict_provenance(): + result = attention_provider(_request(seq_len=256)) + + assert result.backend_id == "aiter.rocm.ck_dense_mha" + assert result.out.shape == (1, GLOBAL_Q_HEADS, 256, HEAD_DIM) + assert result.lse.shape == (1, GLOBAL_Q_HEADS, 256) + assert result.lse.dtype == torch.float32 + assert result.provenance["fallback"] is False + assert result.provenance["actual_backend"] == "aiter.rocm.ck_dense_mha" + assert result.provenance["lse_domain"] == "attention" + + core = result.provenance["core"] + assert core["native_attention_arithmetic"] is True + assert core["deterministic_backward"] is True + assert core["num_splits"] == 1 + assert core["fallback"] is False + assert core["merge_order"] == "global_block_index" + assert core["accum_dtype"] == "fp32" + assert core["downcast_at"] == "final_write" + + +@requires_strict_rocm +def test_training_and_rollout_roles_are_bitwise_identical(): + train = attention_provider(_request(seq_len=256, role="train", seed=3)) + rollout = attention_provider(_request(seq_len=256, role="infer", seed=3)) + + assert torch.equal(train.out, rollout.out) + assert torch.equal(train.lse, rollout.lse) + + +@requires_strict_rocm +@pytest.mark.parametrize("batch_size", [2, 4]) +@pytest.mark.parametrize("seq_len", [256, 512, 2048]) +def test_batch_composition_is_bitwise_invariant(batch_size, seq_len): + """A batch must equal the same rows submitted one at a time. + + Raw AITER does not provide this for every shape: measured on MI300X it is + batch-composition sensitive in BF16 at ``S=256`` (B=4) and ``S=512`` (B=2 + and B=4), while holding at 128/1024/2048/4096. Shape-dependent breakage is + exactly what a per-row rule has to defend against, because the shapes that + hold would otherwise make the bug look absent. ``S=512`` is kept in this + parametrization deliberately. + """ + + batched_request = _request(batch_size=batch_size, seq_len=seq_len, seed=11) + batched = attention_provider(batched_request) + + for row in range(batch_size): + single = SimpleNamespace( + query=batched_request.query[row : row + 1], + key=batched_request.key[row : row + 1], + value=batched_request.value[row : row + 1], + key_padding_mask=None, + tensor_parallel_group=None, + context_parallel=batched_request.context_parallel, + metadata=batched_request.metadata, + ) + row_result = attention_provider(single) + assert torch.equal(batched.out[row : row + 1], row_result.out) + assert torch.equal(batched.lse[row : row + 1], row_result.lse) + + +@requires_strict_rocm +def test_repeated_invocations_are_bitwise_identical(): + first = attention_provider(_request(seq_len=512, seed=5)) + second = attention_provider(_request(seq_len=512, seed=5)) + + assert torch.equal(first.out, second.out) + assert torch.equal(first.lse, second.lse) + + +@requires_strict_rocm +def test_backward_gradients_are_deterministic(): + def run(): + request = _request(seq_len=256, seed=17) + request.query.requires_grad_(True) + request.key.requires_grad_(True) + request.value.requires_grad_(True) + result = attention_provider(request) + result.out.backward(torch.ones_like(result.out)) + return request.query.grad, request.key.grad, request.value.grad + + first = run() + second = run() + for lhs, rhs in zip(first, second, strict=True): + assert torch.equal(lhs, rhs) + + +@requires_strict_rocm +def test_contract_fingerprint_is_rank_independent(): + left = attention_provider(_request(seq_len=128, seed=2)) + right = attention_provider(_request(seq_len=128, seed=2)) + + assert left.contract_id == right.contract_id + assert len(left.contract_id) == 64 + + +@requires_strict_rocm +def test_explicit_scale_is_honored(): + scale = 1.0 / math.sqrt(HEAD_DIM) + default = attention_provider(_request(seq_len=128, seed=8)) + explicit = attention_provider(_request(seq_len=128, seed=8, softmax_scale=scale)) + + assert torch.equal(default.out, explicit.out) + + +@requires_strict_rocm +def test_provenance_records_the_launch_pinning(): + """Launch granularity is the mechanism behind the bitwise claim. + + A reader of a strict report has to be able to tell that the result came + from pinned one-row/one-KV-group launches rather than a batched call that + happened to agree. + """ + + result = attention_provider(_request(seq_len=256)) + execution = result.provenance["execution"] + binding = result.provenance["cross_config_binding"] + + assert execution["launch_granularity"] == "one_batch_row_one_kv_group" + assert execution["kv_groups_materialized_independently"] is True + assert execution["batch_rows_materialized_independently"] is True + # B=1 request over a 32Q/8KV layout -> one launch per KV group. + assert execution["core_launches"] == GLOBAL_KV_HEADS + assert binding["tp_degree_invariant"] is True + assert binding["invariance_mechanism"] == "one_kv_group_per_launch" + + +@requires_strict_rocm +@pytest.mark.parametrize("tp", [2, 4, 8]) +@pytest.mark.parametrize("seq_len", [512, 2048]) +def test_tp_degree_is_bitwise_invariant(tp, seq_len): + """A TP head shard must equal the same slice of an unsharded run. + + TP performs no cross-rank reduction in attention, so this has to hold for + train and rollout to compare across TP degrees. Raw AITER does not provide + it (up to 7.8125e-03 drift); the per-KV-group launch rule is what does. + """ + + class _Group: + def __init__(self, rank, size): + self._rank, self._size = rank, size + + def rank(self): + return self._rank + + def size(self): + return self._size + + base = _request(seq_len=seq_len, seed=3) + full = attention_provider(base) + + local_q, local_kv = GLOBAL_Q_HEADS // tp, GLOBAL_KV_HEADS // tp + for rank in range(tp): + shard = SimpleNamespace( + query=base.query[:, rank * local_q : (rank + 1) * local_q], + key=base.key[:, rank * local_kv : (rank + 1) * local_kv], + value=base.value[:, rank * local_kv : (rank + 1) * local_kv], + key_padding_mask=None, + tensor_parallel_group=_Group(rank, tp), + context_parallel=base.context_parallel, + metadata=_metadata(tp_rank=rank, tp_world_size=tp), + ) + result = attention_provider(shard) + assert torch.equal(result.out, full.out[:, rank * local_q : (rank + 1) * local_q]) + assert torch.equal(result.lse, full.lse[:, rank * local_q : (rank + 1) * local_q])