diff --git a/benchmarks/benchmark_rocm_det_gemm_leaf.py b/benchmarks/benchmark_rocm_det_gemm_leaf.py new file mode 100644 index 00000000..eb00a05b --- /dev/null +++ b/benchmarks/benchmark_rocm_det_gemm_leaf.py @@ -0,0 +1,593 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Offline tile/occupancy sweep for the strict ROCm deterministic GEMM leaf. + +The benchmark launches the production Triton kernels directly with preallocated +buffers. Every candidate is checked against the pinned 64x64/4-warp baseline at +both the complete leaf workspace and final tree root before timings are reported. + +Example: + + python benchmarks/benchmark_rocm_det_gemm_leaf.py \ + --device 4 \ + --cases fwd_gate_m32,fwd_down_m32 \ + --configs 16x128x4,32x64x2xn,32x128x4xn,64x64x4 \ + --output /tmp/rlk_leaf_sweep.json +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import platform +import statistics +import subprocess +import sys +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Callable + +import torch +import triton + +from rl_engine.kernels.ops.triton.matmul.det_gemm import ( + _copy_tree_root_kernel, + _copy_tree_root_transposed_kernel, + _det_gemm_tree_leaf_kernel, + _det_gemm_tree_reduce_kernel, + _device_tree_plan, +) + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_KERNEL_SOURCE = _REPO_ROOT / "rl_engine/kernels/ops/triton/matmul/det_gemm.py" + + +@dataclass(frozen=True) +class LeafConfig: + block_m: int + block_n: int + num_warps: int + order: str = "leaf" + waves_per_eu: int = 0 + + @property + def slug(self) -> str: + return ( + f"{self.block_m}x{self.block_n}x{self.num_warps}x" + f"{self.order}xw{self.waves_per_eu}" + ) + + +@dataclass(frozen=True) +class LeafCase: + name: str + m_size: int + k_size: int + n_size: int + transposed_a: bool = False + transpose_output: bool = False + + +def _leaf_cases() -> dict[str, LeafCase]: + cases: list[LeafCase] = [] + for token_count in (1, 8, 16, 32): + cases.extend( + ( + LeafCase(f"fwd_gate_m{token_count}", token_count, 4096, 12288), + LeafCase(f"fwd_down_m{token_count}", token_count, 12288, 4096), + LeafCase( + f"wgrad_gate_m{token_count}", + 4096, + token_count, + 12288, + transposed_a=True, + transpose_output=True, + ), + LeafCase( + f"wgrad_down_m{token_count}", + 12288, + token_count, + 4096, + transposed_a=True, + transpose_output=True, + ), + ) + ) + for tp_size in (2, 4, 8): + local_intermediate = 12288 // tp_size + suffix = f"tp{tp_size}_m{token_count}" + cases.extend( + ( + LeafCase( + f"fwd_gate_{suffix}", + token_count, + 4096, + local_intermediate, + ), + LeafCase( + f"fwd_down_{suffix}", + token_count, + local_intermediate, + 4096, + ), + LeafCase( + f"wgrad_gate_{suffix}", + 4096, + token_count, + local_intermediate, + transposed_a=True, + transpose_output=True, + ), + LeafCase( + f"wgrad_down_{suffix}", + local_intermediate, + token_count, + 4096, + transposed_a=True, + transpose_output=True, + ), + ) + ) + return {case.name: case for case in cases} + + +_CASES = _leaf_cases() + +_BASELINE = LeafConfig(64, 64, 4) + + +def _parse_csv(value: str) -> list[str]: + return [item.strip() for item in value.split(",") if item.strip()] + + +def _parse_config(value: str) -> LeafConfig: + try: + parts = value.split("x") + if len(parts) not in (3, 4, 5): + raise ValueError + block_m, block_n, num_warps = (int(part) for part in parts[:3]) + order = parts[3] if len(parts) >= 4 else "leaf" + waves_per_eu = int(parts[4]) if len(parts) == 5 else 0 + except (TypeError, ValueError) as error: + raise argparse.ArgumentTypeError( + "config must be BLOCK_MxBLOCK_NxNUM_WARPS[xORDER[xWAVES_PER_EU]], " + f"got {value!r}" + ) from error + if ( + block_m <= 0 + or block_n <= 0 + or num_warps not in (1, 2, 4, 8) + or order not in ("leaf", "n") + or waves_per_eu not in (0, 1, 2, 4) + ): + raise argparse.ArgumentTypeError(f"invalid leaf config {value!r}") + return LeafConfig(block_m, block_n, num_warps, order, waves_per_eu) + + +def _git_output(*arguments: str) -> str: + try: + result = subprocess.run( + ["git", *arguments], + cwd=_REPO_ROOT, + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return "" + return result.stdout.strip() + + +def _command_output(*arguments: str) -> str: + try: + result = subprocess.run( + list(arguments), + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return "" + return result.stdout.strip() + + +def _sha256(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + raw = tensor.detach().contiguous().view(torch.uint8).cpu().numpy() + return hashlib.sha256(memoryview(raw)).hexdigest() + + +def _percentile(values: list[float], percentile: float) -> float: + ordered = sorted(values) + position = (len(ordered) - 1) * percentile + lower = int(position) + upper = min(lower + 1, len(ordered) - 1) + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def _summary(samples_ms: list[float]) -> dict[str, float | list[float]]: + return { + "samples_ms": samples_ms, + "median_ms": statistics.median(samples_ms), + "p95_ms": _percentile(samples_ms, 0.95), + "min_ms": min(samples_ms), + "max_ms": max(samples_ms), + } + + +def _measure( + launch: Callable[[], None], + *, + warmup: int, + samples: int, +) -> dict[str, float | list[float]]: + for _ in range(warmup): + launch() + torch.cuda.synchronize() + events = [ + (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + for _ in range(samples) + ] + for start, end in events: + start.record() + launch() + end.record() + torch.cuda.synchronize() + return _summary([float(start.elapsed_time(end)) for start, end in events]) + + +def _inputs(case: LeafCase, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + generator = torch.Generator(device=device).manual_seed(4100 + case.k_size + case.m_size) + if case.transposed_a: + source = torch.randn( + (case.k_size, case.m_size), + generator=generator, + device=device, + dtype=torch.bfloat16, + ) + a = source.t() + if any(stride <= 0 for stride in a.stride()): + raise RuntimeError("wgrad benchmark requires a positive-stride transpose view") + if case.k_size > 1 and a.is_contiguous(): + raise RuntimeError("non-degenerate wgrad benchmark requires a transpose view") + else: + a = torch.randn( + (case.m_size, case.k_size), + generator=generator, + device=device, + dtype=torch.bfloat16, + ) + b = torch.randn( + (case.k_size, case.n_size), + generator=generator, + device=device, + dtype=torch.bfloat16, + ) + return a, b + + +def _launch_leaf( + case: LeafCase, + config: LeafConfig, + a: torch.Tensor, + b: torch.Tensor, + workspace: torch.Tensor, + plan, +) -> None: + tiles_m = triton.cdiv(case.m_size, config.block_m) + tiles_n = triton.cdiv(case.n_size, config.block_n) + grid = ( + (tiles_n, tiles_m, len(plan.host.leaf_nodes)) + if config.order == "n" + else (len(plan.host.leaf_nodes), tiles_m, tiles_n) + ) + launch_options = {"num_warps": config.num_warps} + if config.waves_per_eu: + launch_options["waves_per_eu"] = config.waves_per_eu + _det_gemm_tree_leaf_kernel[grid]( + a, + b, + workspace, + plan.leaf_starts, + plan.leaf_lengths, + plan.leaf_nodes, + M=case.m_size, + N=case.n_size, + K=case.k_size, + stride_am=a.stride(0), + stride_ak=a.stride(1), + stride_bk=b.stride(0), + stride_bn=b.stride(1), + BLOCK_M=config.block_m, + BLOCK_N=config.block_n, + BLOCK_K=32, + N_FASTEST=config.order == "n", + **launch_options, + ) + + +def _launch_tree( + case: LeafCase, + config: LeafConfig, + a: torch.Tensor, + b: torch.Tensor, + workspace: torch.Tensor, + output: torch.Tensor, + plan, +) -> None: + _launch_leaf(case, config, a, b, workspace, plan) + reduction_block = 256 + for operations, (lower, upper, result) in zip( + plan.host.reduction_levels, + plan.reduction_levels, + strict=True, + ): + grid = ( + len(operations), + triton.cdiv(case.m_size * case.n_size, reduction_block), + ) + _det_gemm_tree_reduce_kernel[grid]( + workspace, + lower, + upper, + result, + M=case.m_size, + N=case.n_size, + BLOCK=reduction_block, + ) + + if case.transpose_output: + block = 32 + grid = ( + triton.cdiv(case.m_size, block), + triton.cdiv(case.n_size, block), + ) + _copy_tree_root_transposed_kernel[grid]( + workspace, + output, + plan.host.root, + M=case.m_size, + N=case.n_size, + BLOCK_M=block, + BLOCK_N=block, + ) + else: + block = 256 + _copy_tree_root_kernel[(triton.cdiv(output.numel(), block),)]( + workspace, + output, + plan.host.root, + output.numel(), + BLOCK=block, + ) + + +def _same_raw_bytes(actual: torch.Tensor, expected: torch.Tensor) -> bool: + if actual.shape != expected.shape or actual.dtype != expected.dtype: + return False + return bool(torch.equal(actual.contiguous().view(torch.uint8), expected.view(torch.uint8))) + + +def _run_case( + case: LeafCase, + configs: list[LeafConfig], + *, + device: torch.device, + warmup: int, + samples: int, +) -> dict[str, object]: + print( + f"{case.name}: A=({case.m_size}, {case.k_size}), " + f"B=({case.k_size}, {case.n_size})", + flush=True, + ) + a, b = _inputs(case, device) + plan = _device_tree_plan(case.k_size, device) + workspace_shape = (plan.host.node_count, case.m_size, case.n_size) + output_shape = ( + (case.n_size, case.m_size) + if case.transpose_output + else (case.m_size, case.n_size) + ) + reference_workspace = torch.empty(workspace_shape, dtype=torch.bfloat16, device=device) + reference_output = torch.empty(output_shape, dtype=torch.bfloat16, device=device) + _launch_tree( + case, + _BASELINE, + a, + b, + reference_workspace, + reference_output, + plan, + ) + torch.cuda.synchronize() + leaf_indices = plan.leaf_nodes.to(torch.int64) + reference_leaves = reference_workspace.index_select(0, leaf_indices) + reference_fingerprints = { + "leaf_workspace_sha256_raw_bytes": _tensor_sha256(reference_leaves), + "root_sha256_raw_bytes": _tensor_sha256(reference_output), + "leaf_workspace_nbytes": reference_leaves.numel() + * reference_leaves.element_size(), + "root_nbytes": reference_output.numel() * reference_output.element_size(), + } + + results: list[dict[str, object]] = [] + for config in configs: + workspace = torch.empty(workspace_shape, dtype=torch.bfloat16, device=device) + output = torch.empty(output_shape, dtype=torch.bfloat16, device=device) + _launch_tree(case, config, a, b, workspace, output, plan) + _launch_leaf(case, config, a, b, workspace, plan) + torch.cuda.synchronize() + candidate_leaves = workspace.index_select(0, leaf_indices) + leaf_raw_bytes_equal = _same_raw_bytes(candidate_leaves, reference_leaves) + + _launch_tree(case, config, a, b, workspace, output, plan) + torch.cuda.synchronize() + root_raw_bytes_equal = _same_raw_bytes(output, reference_output) + if not leaf_raw_bytes_equal or not root_raw_bytes_equal: + raise RuntimeError( + f"{case.name}/{config.slug} changed strict GEMM raw bytes: " + f"leaf={leaf_raw_bytes_equal}, root={root_raw_bytes_equal}" + ) + + leaf_timing = _measure( + lambda: _launch_leaf(case, config, a, b, workspace, plan), + warmup=warmup, + samples=samples, + ) + tree_timing = _measure( + lambda: _launch_tree(case, config, a, b, workspace, output, plan), + warmup=warmup, + samples=samples, + ) + result = { + "config": asdict(config), + "slug": config.slug, + "leaf_raw_bytes_equal": leaf_raw_bytes_equal, + "root_raw_bytes_equal": root_raw_bytes_equal, + "leaf_timing": leaf_timing, + "tree_timing": tree_timing, + } + results.append(result) + print( + f" {config.slug}: leaf={leaf_timing['median_ms']:.4f} ms, " + f"tree={tree_timing['median_ms']:.4f} ms", + flush=True, + ) + del candidate_leaves, workspace, output + + baseline_result = next(result for result in results if result["slug"] == _BASELINE.slug) + baseline_leaf = float(baseline_result["leaf_timing"]["median_ms"]) + baseline_tree = float(baseline_result["tree_timing"]["median_ms"]) + for result in results: + leaf_median = float(result["leaf_timing"]["median_ms"]) + tree_median = float(result["tree_timing"]["median_ms"]) + result["leaf_speedup_vs_baseline"] = baseline_leaf / leaf_median + result["tree_speedup_vs_baseline"] = baseline_tree / tree_median + results.sort(key=lambda result: float(result["leaf_timing"]["median_ms"])) + return { + "case": asdict(case), + "tree": { + "leaf_count": len(plan.host.leaf_nodes), + "node_count": plan.host.node_count, + "reduction_levels": len(plan.host.reduction_levels), + }, + "baseline_fingerprints": reference_fingerprints, + "results": results, + } + + +def _validate_args(args: argparse.Namespace) -> None: + if getattr(torch.version, "hip", None) is None: + raise RuntimeError("this benchmark requires a ROCm PyTorch build") + if not torch.cuda.is_available(): + raise RuntimeError("no ROCm GPU is available") + if args.device < 0 or args.device >= torch.cuda.device_count(): + raise ValueError(f"--device must be in [0, {torch.cuda.device_count() - 1}]") + if args.warmup < 0 or args.samples <= 0: + raise ValueError("--warmup must be non-negative and --samples must be positive") + unknown = sorted(set(args.cases) - _CASES.keys()) + if unknown: + raise ValueError(f"unknown cases: {', '.join(unknown)}") + if _BASELINE not in args.configs: + args.configs.append(_BASELINE) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", type=int, default=0) + parser.add_argument( + "--cases", + type=_parse_csv, + default=["fwd_gate_m32", "fwd_down_m32"], + help=f"comma-separated case names; choices: {','.join(_CASES)}", + ) + parser.add_argument( + "--configs", + type=lambda value: [_parse_config(item) for item in _parse_csv(value)], + default=[ + LeafConfig(16, 64, 2), + LeafConfig(16, 128, 4), + LeafConfig(32, 64, 2, "n"), + LeafConfig(32, 128, 4, "n"), + _BASELINE, + ], + ) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--samples", type=int, default=30) + parser.add_argument("--output", type=Path) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + _validate_args(args) + torch.cuda.set_device(args.device) + torch.backends.cuda.matmul.allow_tf32 = False + device = torch.device("cuda", args.device) + properties = torch.cuda.get_device_properties(args.device) + tracked_diff = _git_output("diff", "--binary") + benchmark_source = Path(__file__).read_bytes() + payload = { + "environment": { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "command": [sys.executable, *sys.argv], + "hostname": platform.node(), + "python": platform.python_version(), + "torch": torch.__version__, + "hip": torch.version.hip, + "triton": triton.__version__, + "gpu_index": args.device, + "gpu": properties.name, + "architecture": getattr(properties, "gcnArchName", ""), + "git_commit": _git_output("rev-parse", "HEAD"), + "git_status": _git_output("status", "--short").splitlines(), + "tracked_diff_sha256": _sha256(tracked_diff.encode()), + "kernel_source_sha256": _sha256(_KERNEL_SOURCE.read_bytes()), + "benchmark_source_sha256": _sha256(benchmark_source), + "rocm_smi_snapshot": _command_output( + "rocm-smi", + "--showuse", + "--showmemuse", + "--showtemp", + "--showclocks", + ), + }, + "methodology": { + "timing": "GPU events around preallocated direct kernel launches", + "correctness": "raw bytes of every leaf node and final BF16 root", + "baseline": asdict(_BASELINE), + "warmup": args.warmup, + "samples": args.samples, + }, + "cases": [], + } + for case_name in args.cases: + payload["cases"].append( + _run_case( + _CASES[case_name], + args.configs, + device=device, + warmup=args.warmup, + samples=args.samples, + ) + ) + torch.cuda.empty_cache() + + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + print(f"results: {args.output.resolve()}") + else: + print(json.dumps(payload, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_rocm_ffn.py b/benchmarks/benchmark_rocm_ffn.py new file mode 100644 index 00000000..8b0b342b --- /dev/null +++ b/benchmarks/benchmark_rocm_ffn.py @@ -0,0 +1,1925 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""ROCm benchmark for the deterministic distributed Triton Qwen3 FFN. + +Distributed performance compares four paths at the same TP/CP/SP topology: +H100 official/deterministic and MI300X official/deterministic. Determinism +compares every Triton TP/CP/SP layout bitwise with Triton TP=1. A separate +single-GPU section retains the official Hugging Face Qwen3MLP TP=1 context. No +model checkpoint or serving engine is used. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import queue +import statistics +import tempfile +import time +import traceback +from datetime import timedelta +from pathlib import Path +from typing import Any, Callable + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F +from transformers import __version__ as transformers_version +from transformers.models.qwen3.configuration_qwen3 import Qwen3Config +from transformers.models.qwen3.modeling_qwen3 import Qwen3MLP + +import rl_engine.kernels.ops.triton.ffn.ffn as ffn_module +from rl_engine.kernels.ops.triton.ffn import ( + pack_qwen3_ffn_forward_weights, + qwen3_ffn, +) + +_DISTRIBUTED_CONFIGS: dict[int, tuple[tuple[str, int, int, bool], ...]] = { + 2: (("tp2", 2, 1, False), ("tp2_sp", 2, 1, True)), + 4: ( + ("tp4", 4, 1, False), + ("tp2_cp2", 2, 2, False), + ("tp2_cp2_sp", 2, 2, True), + ), + 8: ( + ("tp8", 8, 1, False), + ("tp4_cp2", 4, 2, False), + ("tp4_cp2_sp", 4, 2, True), + ), +} +_TP1_FORWARD_CACHE_BYTES = 3 * 4096 * 12288 * torch.bfloat16.itemsize +_COMMUNICATION_CONTRACT = { + "forward": {"all_gather": 1, "reduce_scatter": 1, "total": 2}, + "train_fwd_bwd": { + "all_gather": 7, + "reduce_scatter_logical_lanes": 3, + "logical_total": 10, + "collective_invocations": 9, + }, +} +_PREVIOUS_DISTRIBUTED_TRITON_MS = { + ("tp2", "forward"): 0.9039933793246746, + ("tp2", "train_fwd_bwd"): 2.6996671222150326, + ("tp2_sp", "forward"): 1.0561398230493069, + ("tp2_sp", "train_fwd_bwd"): 2.9646214097738266, + ("tp4", "forward"): 0.8358820341527462, + ("tp4", "train_fwd_bwd"): 2.6998785324394703, + ("tp2_cp2", "forward"): 0.9015901014208794, + ("tp2_cp2", "train_fwd_bwd"): 3.5605919547379017, + ("tp2_cp2_sp", "forward"): 1.1296039447188377, + ("tp2_cp2_sp", "train_fwd_bwd"): 3.992859274148941, + ("tp8", "forward"): 1.0859542526304722, + ("tp8", "train_fwd_bwd"): 2.5620022788643837, + ("tp4_cp2", "forward"): 1.0536308400332928, + ("tp4_cp2", "train_fwd_bwd"): 3.259910736232996, + ("tp4_cp2_sp", "forward"): 1.1927778832614422, + ("tp4_cp2_sp", "train_fwd_bwd"): 3.7497887387871742, +} + + +def _percentile(values: list[float], percentile: float) -> float: + ordered = sorted(values) + 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().float() + expected_float = expected.detach().float() + denominator = torch.linalg.vector_norm(expected_float) + numerator = torch.linalg.vector_norm(actual_float - expected_float) + if denominator.item() == 0.0: + return float(numerator.item()) + return float((numerator / denominator).item()) + + +def _accuracy(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: + difference = actual.detach().float() - expected.detach().float() + return { + "max_abs": float(difference.abs().max().item()), + "mean_abs": float(difference.abs().mean().item()), + "relative_l2": _relative_l2(actual, expected), + "exact_fraction": float( + (actual.detach() == expected.detach()).float().mean().item() + ), + } + + +def _mismatches(left: torch.Tensor, right: torch.Tensor) -> int: + return int((left.detach() != right.detach()).sum().item()) + + +def _randn( + shape: tuple[int, ...], + *, + seed: int, + device: torch.device, + scale: float = 0.02, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + generator = torch.Generator(device="cpu").manual_seed(seed) + value = torch.randn(shape, generator=generator, dtype=torch.float32) * scale + return value.to(device=device, dtype=dtype) + + +def _gpu_event_samples( + function: Callable[[], Any], + *, + warmup: int, + samples: int, +) -> list[float]: + for _ in range(warmup): + function() + torch.cuda.synchronize() + events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] + for _ in range(samples): + 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 _official_qwen3_mlp( + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, +) -> Qwen3MLP: + """Build the upstream Transformers Qwen3 FFN with the benchmark weights.""" + config = Qwen3Config( + hidden_size=gate_weight.size(1), + intermediate_size=gate_weight.size(0), + hidden_act="silu", + ) + module = Qwen3MLP(config).to( + device=gate_weight.device, + dtype=gate_weight.dtype, + ) + with torch.no_grad(): + module.gate_proj.weight.copy_(gate_weight) + module.up_proj.weight.copy_(up_weight) + module.down_proj.weight.copy_(down_weight) + return module + + +def _official_inference(module: Qwen3MLP, hidden: torch.Tensor) -> torch.Tensor: + with torch.no_grad(): + return module(hidden) + + +def _training_step( + function: Callable[..., torch.Tensor], + inputs: list[torch.Tensor], + grad_output: torch.Tensor, +) -> torch.Tensor: + for value in inputs: + value.grad = None + output = function(*inputs) + output.backward(grad_output) + return output + + +class _NativeAllReduce(torch.autograd.Function): + """Autograd-aware native ProcessGroup all-reduce used by the baseline.""" + + @staticmethod + def forward(ctx: Any, input: torch.Tensor, group: Any) -> torch.Tensor: + ctx.group = group + output = input.contiguous().clone() + dist.all_reduce(output, group=group) + return output + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]: + return grad_output, None + + +class _NativeCopyToTensorParallel(torch.autograd.Function): + """Identity in forward and native all-reduce in backward.""" + + @staticmethod + def forward(ctx: Any, input: torch.Tensor, group: Any) -> torch.Tensor: + ctx.group = group + return input + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]: + grad_input = grad_output.contiguous().clone() + dist.all_reduce(grad_input, group=ctx.group) + return grad_input, None + + +class _NativeAllGather(torch.autograd.Function): + """Autograd-aware native first-dimension all-gather baseline.""" + + @staticmethod + def forward(ctx: Any, input: torch.Tensor, group: Any) -> torch.Tensor: + world_size = dist.get_world_size(group=group) + ctx.group = group + ctx.input_shape = tuple(input.shape) + input_flat = input.contiguous().view(-1) + output_flat = torch.empty( + world_size * input_flat.numel(), + dtype=input.dtype, + device=input.device, + ) + dist.all_gather_into_tensor(output_flat, input_flat, group=group) + return output_flat.view(world_size * input.size(0), *input.shape[1:]) + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]: + grad_output_flat = grad_output.contiguous().view(-1) + grad_input_flat = torch.empty( + math.prod(ctx.input_shape), + dtype=grad_output.dtype, + device=grad_output.device, + ) + dist.reduce_scatter_tensor( + grad_input_flat, + grad_output_flat, + group=ctx.group, + ) + return grad_input_flat.view(ctx.input_shape), None + + +class _NativeReduceScatter(torch.autograd.Function): + """Autograd-aware native first-dimension reduce-scatter baseline.""" + + @staticmethod + def forward(ctx: Any, input: torch.Tensor, group: Any) -> torch.Tensor: + world_size = dist.get_world_size(group=group) + if input.size(0) % world_size != 0: + raise ValueError("native reduce-scatter requires divisible dimension 0") + ctx.group = group + ctx.input_shape = tuple(input.shape) + output_shape = (input.size(0) // world_size, *input.shape[1:]) + output_flat = torch.empty( + math.prod(output_shape), + dtype=input.dtype, + device=input.device, + ) + dist.reduce_scatter_tensor( + output_flat, + input.contiguous().view(-1), + group=group, + ) + return output_flat.view(output_shape) + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]: + grad_output_flat = grad_output.contiguous().view(-1) + grad_input_flat = torch.empty( + math.prod(ctx.input_shape), + dtype=grad_output.dtype, + device=grad_output.device, + ) + dist.all_gather_into_tensor( + grad_input_flat, + grad_output_flat, + group=ctx.group, + ) + return grad_input_flat.view(ctx.input_shape), None + + +def _official_distributed_ffn( + hidden: torch.Tensor, + gate_weight: torch.Tensor, + up_weight: torch.Tensor, + down_weight: torch.Tensor, + *, + tp_group: Any, + sequence_parallel: bool, +) -> torch.Tensor: + """Upstream Qwen3 FFN math with native distributed collectives.""" + full_hidden = ( + _NativeAllGather.apply(hidden, tp_group) + if sequence_parallel + else _NativeCopyToTensorParallel.apply(hidden, tp_group) + ) + activated = F.silu(F.linear(full_hidden, gate_weight)) * F.linear( + full_hidden, up_weight + ) + partial_output = F.linear(activated, down_weight) + if sequence_parallel: + return _NativeReduceScatter.apply(partial_output, tp_group) + return _NativeAllReduce.apply(partial_output, tp_group) + + +def _official_distributed_training_step( + inputs: list[torch.Tensor], + grad_output: torch.Tensor, + *, + tp_group: Any, + cp_group: Any, + sequence_parallel: bool, +) -> torch.Tensor: + for value in inputs: + value.grad = None + output = _official_distributed_ffn( + *inputs, + tp_group=tp_group, + sequence_parallel=sequence_parallel, + ) + output.backward(grad_output) + if cp_group is not None: + for weight in inputs[1:]: + dist.all_reduce(weight.grad, group=cp_group) + return output + + +def _single_gpu_benchmarks( + *, + warmup: int, + samples: int, + training_samples: int, +) -> dict[str, list[dict[str, Any]]]: + device = torch.device("cuda", 0) + torch.cuda.set_device(device) + torch.backends.cuda.matmul.allow_tf32 = False + results: dict[str, list[dict[str, Any]]] = {"speed": [], "dtype_accuracy": []} + gate_weight = _randn((12288, 4096), seed=3000, device=device) + up_weight = _randn((12288, 4096), seed=3001, device=device) + down_weight = _randn((4096, 12288), seed=3002, device=device) + forward_weights = pack_qwen3_ffn_forward_weights( + gate_weight, + up_weight, + down_weight, + ) + official = _official_qwen3_mlp(gate_weight, up_weight, down_weight) + for index, tokens in enumerate((1, 8, 32)): + hidden = _randn((tokens, 4096), seed=3010 + index * 2, device=device) + grad_output = _randn( + (tokens, 4096), seed=3011 + index * 2, device=device + ) + weights = (gate_weight, up_weight, down_weight) + official_timing = _summary_ms( + _gpu_event_samples( + lambda: _official_inference(official, hidden), + warmup=warmup, + samples=samples, + ) + ) + triton_timing = _summary_ms( + _gpu_event_samples( + lambda: qwen3_ffn( + hidden, + *weights, + forward_weights=forward_weights, + ), + warmup=warmup, + samples=samples, + ) + ) + results["speed"].append( + { + "name": f"(M,H,I)=({tokens},4096,12288), forward", + "direction": "forward", + "tokens": tokens, + "hidden": 4096, + "intermediate": 12288, + "dtype": "bfloat16", + "weight_layout": "packed_forward_cache", + "official_tp1": official_timing, + "triton": triton_timing, + "latency_ratio_vs_official_tp1": ( + triton_timing["median_ms"] / official_timing["median_ms"] + ), + } + ) + + official_hidden = hidden.detach().clone().requires_grad_(True) + triton_inputs = [ + value.detach().clone().requires_grad_(True) + for value in (hidden, *weights) + ] + triton_forward_weights = pack_qwen3_ffn_forward_weights( + *triton_inputs[1:] + ) + + def triton_training_step() -> torch.Tensor: + return _training_step( + lambda *values: qwen3_ffn( + *values, + forward_weights=triton_forward_weights, + ), + triton_inputs, + grad_output, + ) + + def official_training_step() -> torch.Tensor: + official.zero_grad(set_to_none=True) + official_hidden.grad = None + output = official(official_hidden) + output.backward(grad_output) + return output + + official_train_timing = _summary_ms( + _gpu_event_samples( + official_training_step, + warmup=max(1, warmup // 2), + samples=training_samples, + ) + ) + triton_train_timing = _summary_ms( + _gpu_event_samples( + triton_training_step, + warmup=max(1, warmup // 2), + samples=training_samples, + ) + ) + results["speed"].append( + { + "name": f"(M,H,I)=({tokens},4096,12288), forward+backward", + "direction": "train_fwd_bwd", + "tokens": tokens, + "hidden": 4096, + "intermediate": 12288, + "dtype": "bfloat16", + "weight_layout": "packed_forward_cache", + "official_tp1": official_train_timing, + "triton": triton_train_timing, + "latency_ratio_vs_official_tp1": ( + triton_train_timing["median_ms"] + / official_train_timing["median_ms"] + ), + } + ) + del ( + hidden, + grad_output, + official_hidden, + triton_inputs, + triton_forward_weights, + ) + torch.cuda.empty_cache() + + # This is intentionally separate from the determinism and speed results. + del official, forward_weights, gate_weight, up_weight, down_weight + torch.cuda.empty_cache() + tokens = 8 + fp32_hidden = _randn( + (tokens, 4096), seed=3100, device=device, dtype=torch.float32 + ) + fp32_gate = _randn( + (12288, 4096), seed=3101, device=device, dtype=torch.float32 + ) + fp32_up = _randn( + (12288, 4096), seed=3102, device=device, dtype=torch.float32 + ) + fp32_down = _randn( + (4096, 12288), seed=3103, device=device, dtype=torch.float32 + ) + official_fp32 = _official_qwen3_mlp(fp32_gate, fp32_up, fp32_down) + with torch.no_grad(): + fp32_output = official_fp32(fp32_hidden) + del official_fp32 + torch.cuda.empty_cache() + official_fp16 = _official_qwen3_mlp( + fp32_gate.half(), fp32_up.half(), fp32_down.half() + ) + with torch.no_grad(): + fp16_output = official_fp16(fp32_hidden.half()) + results["dtype_accuracy"].append( + { + "name": "Official Qwen3MLP TP=1 FP16 vs FP32", + "tokens": tokens, + "hidden": 4096, + "intermediate": 12288, + "candidate_dtype": "float16", + "reference_dtype": "float32", + **_accuracy(fp16_output, fp32_output), + } + ) + return results + + + + +def _mesh_groups( + world_size: int, + tp_size: int, + cp_size: int, +) -> tuple[list[Any], list[Any]]: + if tp_size * cp_size != world_size: + raise ValueError("TP size times CP size must equal world size") + if tp_size == world_size and cp_size == 1: + return [dist.group.WORLD], [] + tp_groups = [] + if tp_size > 1: + for cp_rank in range(cp_size): + ranks = list(range(cp_rank * tp_size, (cp_rank + 1) * tp_size)) + tp_groups.append(dist.new_group(ranks=ranks)) + cp_groups = [] + if cp_size > 1: + for tp_rank in range(tp_size): + ranks = [cp_rank * tp_size + tp_rank for cp_rank in range(cp_size)] + cp_groups.append(dist.new_group(ranks=ranks)) + return tp_groups, cp_groups + + +def _shard_ranges( + rank: int, + *, + tp_size: int, + cp_size: int, + sequence_parallel: bool, + token_count: int, + intermediate_size: int, +) -> tuple[int, int, int, int]: + tp_rank = rank % tp_size + cp_rank = rank // tp_size + cp_tokens = token_count // cp_size + local_tokens = cp_tokens // tp_size if sequence_parallel else cp_tokens + token_start = cp_rank * cp_tokens + if sequence_parallel: + token_start += tp_rank * local_tokens + token_end = token_start + local_tokens + local_intermediate = intermediate_size // tp_size + feature_start = tp_rank * local_intermediate + feature_end = feature_start + local_intermediate + return token_start, token_end, feature_start, feature_end + + +def _distributed_wall_samples( + function: Callable[[], Any], + *, + group: Any, + warmup: int, + samples: int, +) -> list[float]: + for _ in range(warmup): + function() + torch.cuda.synchronize() + dist.barrier(group=group) + timings = [] + for _ in range(samples): + torch.cuda.synchronize() + start = time.perf_counter() + function() + torch.cuda.synchronize() + timings.append((time.perf_counter() - start) * 1000.0) + dist.barrier(group=group) + return timings + + +def _slowest_rank_summary( + local_timings: list[float], group: Any +) -> dict[str, float]: + world_size = dist.get_world_size(group=group) + gathered: list[list[float] | None] = [None] * world_size + dist.all_gather_object(gathered, local_timings, group=group) + slowest = [ + max(float(rank_values[index]) for rank_values in gathered if rank_values) + for index in range(len(local_timings)) + ] + return _summary_ms(slowest) + + + + +def _distributed_ffn_benchmark( + rank: int, + world_size: int, + configs: tuple[tuple[str, int, int, bool], ...], + *, + warmup: int, + samples: int, + training_samples: int, +) -> list[dict[str, Any]]: + device = torch.device("cuda", rank) + token_count = 32 + hidden_size = 4096 + intermediate_size = 12288 + hidden_full = _randn((token_count, hidden_size), seed=5000, device=device) + gate_full = _randn( + (intermediate_size, hidden_size), seed=5001, device=device + ) + up_full = _randn((intermediate_size, hidden_size), seed=5002, device=device) + down_full = _randn( + (hidden_size, intermediate_size), seed=5003, device=device + ) + grad_output_full = _randn( + (token_count, hidden_size), seed=5004, device=device + ) + + # Exactness reference: the same deterministic Triton implementation at TP=1. + full_values = (hidden_full, gate_full, up_full, down_full) + tp1_forward_weights = pack_qwen3_ffn_forward_weights(*full_values[1:]) + with torch.no_grad(): + tp1_forward = qwen3_ffn( + *full_values, + forward_weights=tp1_forward_weights, + ).detach().clone() + tp1_inputs = [ + value.detach().clone().requires_grad_(True) for value in full_values + ] + tp1_training_weights = pack_qwen3_ffn_forward_weights(*tp1_inputs[1:]) + tp1_train = _training_step( + lambda *values: qwen3_ffn( + *values, + forward_weights=tp1_training_weights, + ), + tp1_inputs, + grad_output_full, + ).detach().clone() + tp1_grads = [value.grad.detach().clone() for value in tp1_inputs] + del tp1_inputs, tp1_forward_weights, tp1_training_weights + torch.cuda.empty_cache() + + meshes: dict[tuple[int, int], tuple[list[Any], list[Any]]] = {} + results = [] + + for name, tp_size, cp_size, sequence_parallel in configs: + mesh_key = (tp_size, cp_size) + if mesh_key not in meshes: + meshes[mesh_key] = _mesh_groups(world_size, tp_size, cp_size) + tp_groups, cp_groups = meshes[mesh_key] + tp_rank = rank % tp_size + cp_rank = rank // tp_size + tp_group = tp_groups[cp_rank] if tp_size > 1 else None + cp_group = cp_groups[tp_rank] if cp_size > 1 else None + token_start, token_end, feature_start, feature_end = _shard_ranges( + rank, + tp_size=tp_size, + cp_size=cp_size, + sequence_parallel=sequence_parallel, + token_count=token_count, + intermediate_size=intermediate_size, + ) + shard = ( + hidden_full[token_start:token_end].contiguous(), + gate_full[feature_start:feature_end].contiguous(), + up_full[feature_start:feature_end].contiguous(), + down_full[:, feature_start:feature_end].contiguous(), + ) + shard_forward_weights = pack_qwen3_ffn_forward_weights(*shard[1:]) + local_grad_output = grad_output_full[token_start:token_end].contiguous() + + def official_distributed_forward(): + with torch.no_grad(): + return _official_distributed_ffn( + *shard, + tp_group=tp_group, + sequence_parallel=sequence_parallel, + ) + + official_forward_summary = _slowest_rank_summary( + _distributed_wall_samples( + official_distributed_forward, + group=dist.group.WORLD, + warmup=warmup, + samples=samples, + ), + dist.group.WORLD, + ) + official_inputs = [ + value.detach().clone().requires_grad_(True) for value in shard + ] + official_train_summary = _slowest_rank_summary( + _distributed_wall_samples( + lambda: _official_distributed_training_step( + official_inputs, + local_grad_output, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ), + group=dist.group.WORLD, + warmup=max(1, warmup // 2), + samples=training_samples, + ), + dist.group.WORLD, + ) + + def triton_forward(): + return qwen3_ffn( + *shard, + forward_weights=shard_forward_weights, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + + with torch.no_grad(): + triton_output = triton_forward() + triton_repeat = triton_forward() + triton_forward_summary = _slowest_rank_summary( + _distributed_wall_samples( + triton_forward, + group=dist.group.WORLD, + warmup=warmup, + samples=samples, + ), + dist.group.WORLD, + ) + + triton_inputs = [ + value.detach().clone().requires_grad_(True) for value in shard + ] + repeat_inputs = [ + value.detach().clone().requires_grad_(True) for value in shard + ] + triton_forward_weights = pack_qwen3_ffn_forward_weights(*triton_inputs[1:]) + repeat_forward_weights = pack_qwen3_ffn_forward_weights(*repeat_inputs[1:]) + + def triton_training_step(inputs, packed_weights): + for value in inputs: + value.grad = None + output = qwen3_ffn( + *inputs, + forward_weights=packed_weights, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + output.backward(local_grad_output) + return output + + triton_train = triton_training_step( + triton_inputs, + triton_forward_weights, + ).detach().clone() + triton_grads = [value.grad.detach().clone() for value in triton_inputs] + repeat_train = triton_training_step( + repeat_inputs, + repeat_forward_weights, + ).detach().clone() + repeat_grads = [value.grad.detach().clone() for value in repeat_inputs] + triton_train_summary = _slowest_rank_summary( + _distributed_wall_samples( + lambda: triton_training_step( + triton_inputs, + triton_forward_weights, + ), + group=dist.group.WORLD, + warmup=max(1, warmup // 2), + samples=training_samples, + ), + dist.group.WORLD, + ) + + expected_forward = tp1_forward[token_start:token_end] + expected_train = tp1_train[token_start:token_end] + expected_grads = ( + tp1_grads[0][token_start:token_end], + tp1_grads[1][feature_start:feature_end], + tp1_grads[2][feature_start:feature_end], + tp1_grads[3][:, feature_start:feature_end], + ) + local_exactness = { + "tp1_forward_output": _mismatches(triton_output, expected_forward), + "tp1_training_output": _mismatches(triton_train, expected_train), + "tp1_hidden_gradient": _mismatches( + triton_grads[0], expected_grads[0] + ), + "tp1_weight_gradient": sum( + _mismatches(actual, expected) + for actual, expected in zip( + triton_grads[1:], expected_grads[1:], strict=True + ) + ), + "repeat_forward": _mismatches(triton_output, triton_repeat), + "train_infer_mismatch_count": _mismatches(triton_output, triton_train), + "repeat_training": _mismatches(triton_train, repeat_train) + + sum( + _mismatches(actual, repeat) + for actual, repeat in zip(triton_grads, repeat_grads, strict=True) + ), + } + gathered: list[dict[str, Any] | None] = [None] * world_size + dist.all_gather_object(gathered, local_exactness) + if rank == 0: + valid = [value for value in gathered if value is not None] + common = { + "name": name, + "world_size": world_size, + "tp_size": tp_size, + "cp_size": cp_size, + "sequence_parallel": sequence_parallel, + "tokens": token_count, + "hidden": hidden_size, + "intermediate": intermediate_size, + "weight_layout": "packed_forward_cache", + } + results.extend( + ( + { + **common, + "direction": "forward", + "official_distributed": official_forward_summary, + "triton": triton_forward_summary, + "latency_ratio_triton_vs_official_distributed": ( + triton_forward_summary["median_ms"] + / official_forward_summary["median_ms"] + ), + "tp1_mismatch": { + "forward_output": sum( + value["tp1_forward_output"] for value in valid + ), + }, + "repeat_mismatch_count": sum( + value["repeat_forward"] for value in valid + ), + "train_infer_mismatch_count": sum( + value["train_infer_mismatch_count"] for value in valid + ), + }, + { + **common, + "direction": "train_fwd_bwd", + "official_distributed": official_train_summary, + "triton": triton_train_summary, + "latency_ratio_triton_vs_official_distributed": ( + triton_train_summary["median_ms"] + / official_train_summary["median_ms"] + ), + "tp1_mismatch": { + "training_output": sum( + value["tp1_training_output"] for value in valid + ), + "hidden_gradient": sum( + value["tp1_hidden_gradient"] for value in valid + ), + "weight_gradient": sum( + value["tp1_weight_gradient"] for value in valid + ), + }, + "repeat_mismatch_count": sum( + value["repeat_training"] for value in valid + ), + "train_infer_mismatch_count": sum( + value["train_infer_mismatch_count"] for value in valid + ), + }, + ) + ) + del ( + shard, + shard_forward_weights, + official_inputs, + triton_inputs, + triton_forward_weights, + repeat_inputs, + repeat_forward_weights, + ) + torch.cuda.empty_cache() + dist.barrier() + + for collective in list(ffn_module._COLLECTIVES.values()): + collective.close() + ffn_module._COLLECTIVES.clear() + return results + + +def _distributed_worker( + rank: int, + world_size: int, + init_method: str, + result_queue: Any, + warmup: int, + samples: int, + training_samples: int, +) -> None: + try: + try: + available_cpus = sorted(os.sched_getaffinity(0)) + numa_span = max(1, len(available_cpus) // 2) + numa_index = 0 if rank < 4 else 1 + local_rank = rank % 4 + cpu_index = min( + numa_index * numa_span + local_rank, + len(available_cpus) - 1, + ) + os.sched_setaffinity(0, {available_cpus[cpu_index]}) + except (AttributeError, OSError): + pass + torch.set_num_threads(1) + 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=15), + ) + distributed_ffn = _distributed_ffn_benchmark( + rank, + world_size, + _DISTRIBUTED_CONFIGS[world_size], + warmup=warmup, + samples=samples, + training_samples=training_samples, + ) + if rank == 0: + result_queue.put( + { + "ok": True, + "world_size": world_size, + "distributed_ffn": distributed_ffn, + } + ) + except Exception: + result_queue.put( + { + "ok": False, + "rank": rank, + "world_size": world_size, + "traceback": traceback.format_exc(), + } + ) + raise + finally: + for collective in list(ffn_module._COLLECTIVES.values()): + collective.close() + ffn_module._COLLECTIVES.clear() + if dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +def _run_distributed_world( + world_size: int, + *, + warmup: int, + samples: int, + training_samples: int, +) -> dict[str, Any]: + 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_worker, + args=( + rank, + world_size, + init_method, + result_queue, + warmup, + samples, + training_samples, + ), + ) + for rank in range(world_size) + ] + for process in processes: + process.start() + result = None + try: + result = result_queue.get(timeout=1800) + if not result["ok"]: + for process in processes: + if process.is_alive(): + process.terminate() + except queue.Empty as exc: + for process in processes: + if process.is_alive(): + process.terminate() + raise RuntimeError( + f"timed out waiting for world_size={world_size} benchmark" + ) from exc + finally: + for process in processes: + process.join(timeout=60) + if process.is_alive(): + process.terminate() + process.join(timeout=30) + result_queue.close() + result_queue.join_thread() + if result is None: + raise RuntimeError(f"world_size={world_size} returned no result") + if not result["ok"]: + raise RuntimeError(result.get("traceback", str(result))) + for process in processes: + if process.exitcode != 0: + raise RuntimeError( + f"world_size={world_size} worker exited with {process.exitcode}" + ) + return result + + + + +def _topology_exactness_rows( + distributed_rows: list[dict[str, Any]], +) -> list[dict[str, Any]]: + merged: dict[str, dict[str, Any]] = {} + for row in distributed_rows: + entry = merged.setdefault( + row["name"], + { + "name": row["name"], + "world_size": row["world_size"], + "tp_size": row["tp_size"], + "cp_size": row["cp_size"], + "sequence_parallel": row["sequence_parallel"], + "forward_output": 0, + "training_output": 0, + "hidden_gradient": 0, + "weight_gradient": 0, + "repeat": 0, + "train_infer": 0, + }, + ) + for key, value in row["tp1_mismatch"].items(): + entry[key] += value + entry["repeat"] += row["repeat_mismatch_count"] + entry["train_infer"] += row["train_infer_mismatch_count"] + return list(merged.values()) + + +def _load_cuda_cpu_comparison( + output_directory: Path, +) -> dict[str, Any] | None: + comparison_path = output_directory / "cuda_cpu_comparison.json" + if not comparison_path.exists(): + return None + return json.loads(comparison_path.read_text(encoding="utf-8")) + + +def _distributed_platform_comparison_rows( + current_rows: list[dict[str, Any]], + comparison_payload: dict[str, Any] | None, +) -> list[dict[str, Any]]: + """Join H100 and MI300X distributed timings by topology and direction.""" + if comparison_payload is None: + return [] + h100_lookup = { + (row["name"], row["direction"]): row + for row in comparison_payload["distributed"] + } + rows: list[dict[str, Any]] = [] + for current in current_rows: + key = (current["name"], current["direction"]) + h100 = h100_lookup.get(key) + if h100 is None: + continue + h100_official_ms = float(h100["official_h100_ms"]) + h100_deterministic_ms = float(h100["cuda_h100_ms"]) + mi300x_official_ms = float( + current["official_distributed"]["median_ms"] + ) + mi300x_deterministic_ms = float(current["triton"]["median_ms"]) + rows.append( + { + "name": current["name"], + "direction": current["direction"], + "tp_size": current["tp_size"], + "cp_size": current["cp_size"], + "sequence_parallel": current["sequence_parallel"], + "h100_official_distributed_ms": h100_official_ms, + "h100_deterministic_cuda_ms": h100_deterministic_ms, + "mi300x_official_distributed_ms": mi300x_official_ms, + "mi300x_deterministic_triton_ms": mi300x_deterministic_ms, + "h100_deterministic_over_official_ratio": ( + h100_deterministic_ms / h100_official_ms + ), + "mi300x_deterministic_over_official_ratio": ( + mi300x_deterministic_ms / mi300x_official_ms + ), + "deterministic_mi300x_over_h100_ratio": ( + mi300x_deterministic_ms / h100_deterministic_ms + ), + } + ) + return rows + + +def _previous_deterministic_comparison_rows( + current_rows: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Compare the current MI300X run with the previous checked report data.""" + rows: list[dict[str, Any]] = [] + for current in current_rows: + key = (current["name"], current["direction"]) + previous_ms = _PREVIOUS_DISTRIBUTED_TRITON_MS.get(key) + if previous_ms is None: + continue + current_ms = float(current["triton"]["median_ms"]) + rows.append( + { + "name": current["name"], + "direction": current["direction"], + "previous_ms": previous_ms, + "current_ms": current_ms, + "latency_reduction_ratio": 1.0 - current_ms / previous_ms, + } + ) + return rows + + +def _write_report( + payload: dict[str, Any], + output_directory: Path, + comparison_payload: dict[str, Any] | None = None, +) -> None: + environment = payload["environment"] + methodology = payload["methodology"] + single_speed = payload["single_gpu"]["speed"] + dtype_rows = payload["single_gpu"]["dtype_accuracy"] + distributed_speed = payload["distributed_ffn"] + platform_comparison = _distributed_platform_comparison_rows( + distributed_speed, comparison_payload + ) + previous_comparison = _previous_deterministic_comparison_rows( + distributed_speed + ) + previous_reductions = [ + row["latency_reduction_ratio"] for row in previous_comparison + ] + exactness_rows = _topology_exactness_rows(distributed_speed) + single_ratios = [ + row["latency_ratio_vs_official_tp1"] for row in single_speed + ] + platform_ratios = [ + row["deterministic_mi300x_over_h100_ratio"] + for row in platform_comparison + ] + total_tp1_mismatch = sum( + row[key] + for row in exactness_rows + for key in ( + "forward_output", + "training_output", + "hidden_gradient", + "weight_gradient", + ) + ) + total_repeat_mismatch = sum(row["repeat"] for row in exactness_rows) + total_train_infer_mismatch = sum(row["train_infer"] for row in exactness_rows) + dtype_row = dtype_rows[0] + + lines = [ + "# PR #325 ROCm deterministic Triton FFN report", + "", + "This is an operator-only MI300X report. It does not load or benchmark a " + "model checkpoint.", + "", + "## Comparison contract", + "", + "1. **Determinism:** every Triton TP/CP/SP result is compared bitwise with " + "the same deterministic Triton FFN at **TP=1**. The reported metric is " + "element mismatch count; acceptance requires 0.", + "2. **FP16/FP32:** one separate, simple output comparison runs official " + "Hugging Face `Qwen3MLP` at TP=1 in FP16 and FP32. FP32 is the reference.", + "3. **Speed:** single-GPU speed retains the official Qwen3MLP TP=1 " + "context. Distributed speed compares four same-topology paths: H100 " + "official/deterministic and MI300X official/deterministic.", + "", + "## Environment", + "", + "| Field | Value |", + "|---|---|", + ] + for key, value in environment.items(): + lines.append(f"| {key} | {value} |") + lines.extend( + ( + "", + "## Methodology", + "", + "- Operator shape: H=4096, I=12288; BF16 is used for all speed and " + "determinism measurements.", + "- Single-GPU shapes use M=1/8/32. Distributed cases use the same full " + "logical M=32 input for TP2/4/8, TP+CP, and sequence parallelism.", + "- The distributed comparison joins rows by topology and direction: " + "H100 and MI300X use the same logical M=32 workload and TP/CP/SP layout. " + "Official TP=1 latency is neither collected nor used in that distributed " + "ratio.", + "- MI300X official distributed uses upstream Qwen3 FFN math with native " + "PyTorch BF16 GEMMs and native RCCL collectives over the same shards; " + "the deterministic path uses the current Triton FFN and fixed-order " + "transport.", + "- Deterministic Triton timings use the explicit prepacked forward-weight " + "cache. Packing happens once outside the timed region; canonical source " + "weights remain the autograd and optimizer source of truth.", + f"- The TP=1 cache adds {_TP1_FORWARD_CACHE_BYTES / 2**20:.0f} MiB; " + "each TP rank holds that amount divided by TP size. Refresh cost is " + "excluded because the benchmark measures the steady-state FFN call.", + "- The distributed exactness baseline is the PR's deterministic Triton " + "FFN at TP=1. Local outputs, dHidden, and sharded dWeights are compared " + "against their exact TP=1 slices.", + "- Communication contract for TP+CP+SP: forward uses 1 TP AllGather " + "plus 1 TP ReduceScatter (2 calls). Forward+backward retains 7 " + "AllGathers plus 3 logical ReduceScatter lanes; PR #357 merges the " + "two independent backward gate/up lanes into one " + "`reduce_scatter_many` call, for 9 collective invocations.", + "- Implementation note: the current ROCm deterministic communication " + "operator is adopted from PR #357. This changes the implementation " + "under test, not the benchmark comparison contract.", + f"- Single-GPU timing: {methodology['single_gpu_timing']}; distributed " + f"timing: {methodology['distributed_timing']}.", + f"- Distributed workers: {methodology['distributed_worker_cpu_affinity']} " + "to reduce host-scheduler noise in synchronized wall-clock samples.", + f"- {methodology['warmup']} warmups, {methodology['samples']} measured " + f"forward samples, and {methodology['training_samples']} measured " + "forward+backward samples.", + "- `NCCL_IB_DISABLE=1` keeps the distributed run on intra-node XGMI. " + "Median, p95, min, and max values are available in `results.json`.", + "", + "Reproduce from the repository root:", + "", + "```bash", + "python benchmarks/benchmark_rocm_ffn.py \\", + f" --warmup {methodology['warmup']} \\", + f" --samples {methodology['samples']} \\", + f" --training-samples {methodology['training_samples']} \\", + " --output-dir benchmarks/results/pr325_rocm_mi300x", + "```", + "", + "## Results summary", + "", + f"- TP=1 exactness baseline: **{total_tp1_mismatch} mismatched " + "elements** across topology forward outputs, training outputs, " + "dHidden, and dWeights.", + f"- Repeat mismatch: **{total_repeat_mismatch}**; training/inference " + f"forward mismatch: **{total_train_infer_mismatch}**.", + f"- Single-GPU deterministic Triton packed-cache latency is " + f"**{min(single_ratios):.2f}-{max(single_ratios):.2f}x** the official " + "Qwen3MLP TP=1 latency across M=1/8/32 and forward/training.", + ( + f"- MI300X deterministic Triton latency is **{min(platform_ratios):.2f}-" + f"{max(platform_ratios):.2f}x** the H100 deterministic CUDA latency " + "for the same distributed layouts. Both official distributed " + "paths are reported alongside them." + if platform_ratios + else "- No matching H100 distributed rows were supplied." + ), + ( + f"- Versus the previous deterministic MI300X benchmark, PR #357 " + f"improves **{sum(value > 0 for value in previous_reductions)}/" + f"{len(previous_reductions)}** rows, with a mean latency reduction " + f"of **{statistics.mean(previous_reductions) * 100:.1f}%**." + if previous_reductions + else "- No previous deterministic MI300X rows were available." + ), + f"- The separate official-Qwen3MLP FP16 versus FP32 observation has " + f"relative-L2 error **{dtype_row['relative_l2']:.3e}** for " + "(M,H,I)=(8,4096,12288).", + "", + "## Single-GPU FFN speed", + "", + "Performance only; no official-versus-Triton accuracy metric is " + "reported here.", + "", + "| Shape / direction | Official Qwen3MLP TP=1 (ms) | Deterministic " + "Triton, packed (ms) | Triton / official TP=1 |", + "|---|---:|---:|---:|", + ) + ) + for row in single_speed: + lines.append( + f"| {row['name']} | {row['official_tp1']['median_ms']:.4f} | " + f"{row['triton']['median_ms']:.4f} | " + f"{row['latency_ratio_vs_official_tp1']:.2f}x |" + ) + + lines.extend(("", "## Distributed FFN speed", "")) + if platform_comparison: + lines.extend( + ( + "Every row compares the same distributed topology and direction. " + "No TP=1 latency is used in this table.", + "", + "| Parallel layout | Direction | H100 official distributed (ms) | " + "H100 deterministic CUDA (ms) | MI300X official distributed (ms) | " + "MI300X deterministic Triton (ms) | H100 det / official | " + "MI300X det / official |", + "|---|---|---:|---:|---:|---:|---:|---:|", + ) + ) + for row in platform_comparison: + lines.append( + f"| {row['name']} | {row['direction']} | " + f"{row['h100_official_distributed_ms']:.4f} | " + f"{row['h100_deterministic_cuda_ms']:.4f} | " + f"{row['mi300x_official_distributed_ms']:.4f} | " + f"{row['mi300x_deterministic_triton_ms']:.4f} | " + f"{row['h100_deterministic_over_official_ratio']:.2f}x | " + f"{row['mi300x_deterministic_over_official_ratio']:.2f}x |" + ) + else: + lines.extend( + ( + "No matching H100 distributed data was supplied; current MI300X " + "latency is reported without a TP=1 ratio.", + "", + "| Parallel layout | Direction | MI300X official distributed (ms) | " + "MI300X deterministic Triton (ms) | Deterministic / official |", + "|---|---|---:|---:|---:|", + ) + ) + for row in distributed_speed: + lines.append( + f"| {row['name']} | {row['direction']} | " + f"{row['official_distributed']['median_ms']:.4f} | " + f"{row['triton']['median_ms']:.4f} | " + f"{row['latency_ratio_triton_vs_official_distributed']:.2f}x |" + ) + + if previous_comparison: + lines.extend( + ( + "", + "### PR #357 latency change versus the previous benchmark", + "", + "This comparison changes only the deterministic ROCm communication " + "implementation. It is not included as another series in the main " + "four-path figure.", + "", + "| Parallel layout | Direction | Previous (ms) | Current (ms) | " + "Latency reduction |", + "|---|---|---:|---:|---:|", + ) + ) + for row in previous_comparison: + lines.append( + f"| {row['name']} | {row['direction']} | " + f"{row['previous_ms']:.4f} | {row['current_ms']:.4f} | " + f"{row['latency_reduction_ratio'] * 100:.1f}% |" + ) + + if comparison_payload is not None: + comparison_environment = comparison_payload["environment"] + comparison_source = comparison_payload["source"] + local_single = { + (row["tokens"], row["direction"]): row for row in single_speed + } + lines.extend( + ( + "", + "## CUDA GPU and CPU performance context", + "", + f"The additional measurements come from [{comparison_source['report']}]" + f"({comparison_source['pull_request']}) at CUDA commit " + f"`{comparison_source['cuda_commit']}`. H100 Triton replays use " + "this PR's code at " + f"`{comparison_source['h100_triton_replay_commit']}`.", + "", + "The same-H100 CUDA/Triton ratio is the hardware-matched comparison. " + "CPU and MI300X columns provide absolute-latency context only; they " + "are not hardware-normalized speed claims.", + "", + "| Comparison environment | Value |", + "|---|---|", + f"| CUDA GPU | {comparison_environment['gpu']} " + f"({comparison_environment['architecture']}) |", + f"| CUDA / PyTorch | {comparison_environment['cuda']} / " + f"{comparison_environment['torch']} |", + f"| CPU | {comparison_environment['cpu']}, " + f"{comparison_environment['cpu_threads']} intra-op threads |", + f"| Transformers | {comparison_environment['transformers']} |", + "", + "### Single-GPU and CPU absolute latency", + "", + "| Shape / direction | CPU official (ms) | H100 official TP=1 " + "(ms) | H100 Triton replay (ms) | H100 CUDA (ms) | CUDA / " + "Triton H100 | MI300X official TP=1 (ms) | MI300X Triton (ms) |", + "|---|---:|---:|---:|---:|---:|---:|---:|", + ) + ) + for row in comparison_payload["single_gpu"]: + local = local_single[(row["tokens"], row["direction"])] + direction = ( + "forward" + if row["direction"] == "forward" + else "forward+backward" + ) + lines.append( + f"| M={row['tokens']}, {direction} | " + f"{row['official_cpu_ms']:.4f} | " + f"{row['official_h100_ms']:.4f} | " + f"{row['triton_h100_ms']:.4f} | " + f"{row['cuda_h100_ms']:.4f} | " + f"{row['cuda_h100_ms'] / row['triton_h100_ms']:.2f}x | " + f"{local['official_tp1']['median_ms']:.4f} | " + f"{local['triton']['median_ms']:.4f} |" + ) + lines.extend( + ( + "", + "Both H100 columns used in the main distributed table are the " + "user-supplied distributed timings. They are joined directly with " + "MI300X rows of the same topology and direction; TP=1 values are " + "excluded from all four columns.", + ) + ) + + lines.extend( + ( + "", + "## Topology exactness versus Triton TP=1", + "", + "All columns are element mismatch counts. This table does not compare " + "against the official FFN.", + "", + "| Parallel layout | Forward output | Training output | dHidden | " + "dWeights | Repeat | Train/infer |", + "|---|---:|---:|---:|---:|---:|---:|", + ) + ) + for row in exactness_rows: + lines.append( + f"| {row['name']} | {row['forward_output']} | " + f"{row['training_output']} | {row['hidden_gradient']} | " + f"{row['weight_gradient']} | {row['repeat']} | " + f"{row['train_infer']} |" + ) + + lines.extend( + ( + "", + "## Simple FP16 versus FP32 observation", + "", + "This is an official `Qwen3MLP` TP=1 output comparison only; it is not " + "used to judge deterministic Triton and is not included in speed " + "ratios.", + "", + "| Shape | Candidate | Reference | Max abs | Mean abs | Relative L2 |", + "|---|---|---|---:|---:|---:|", + f"| (M,H,I)=({dtype_row['tokens']},{dtype_row['hidden']}," + f"{dtype_row['intermediate']}) | FP16 | FP32 | " + f"{dtype_row['max_abs']:.3e} | {dtype_row['mean_abs']:.3e} | " + f"{dtype_row['relative_l2']:.3e} |", + "", + "## Deterministic communication overlap", + "", + "The current timing includes the fixed-order communication schedule " + "and makes no overlap claim. Forward SP all-gather must finish before " + "gate/up projection, and TP reduction consumes the down-projection " + "output, so those edges are hard dependencies.", + "", + "In backward, the gate and up contributions to dHidden are independent " + "until their final ordered addition. A future implementation can place " + "the fixed-rank reduction of one contribution on a second stream while " + "computing the other, but it must preserve rank order, reduction tree, " + "wait points, and gate-then-up addition order. Any optimization is " + "accepted only if every TP=1 mismatch column remains zero.", + "", + "## Figures", + "", + "![Single-GPU CUDA, packed Triton, and CPU latency]" + "(single_gpu_overhead.png)", + "", + "![Topology mismatch versus Triton TP=1](collective_overhead.png)", + "", + "![Distributed H100 CUDA and MI300X packed Triton latency]" + "(distributed_ffn_overhead.png)", + "", + ) + ) + (output_directory / "report.md").write_text( + "\n".join(lines), encoding="utf-8" + ) + + +def _write_figures( + payload: dict[str, Any], + output_directory: Path, + comparison_payload: dict[str, Any] | None = None, +) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + import numpy as np + + plt.style.use("seaborn-v0_8-whitegrid") + plt.rcParams.update( + { + "font.size": 18, + "axes.titlesize": 25, + "axes.labelsize": 21, + "xtick.labelsize": 15, + "ytick.labelsize": 16, + "legend.fontsize": 17, + } + ) + + single_rows = payload["single_gpu"]["speed"] + if comparison_payload is None: + single_labels = [ + f"M={row['tokens']}\n" + f"{'FWD' if row['direction'] == 'forward' else 'FWD+BWD'}" + for row in single_rows + ] + positions = np.arange(len(single_rows)) + width = 0.37 + figure, axis = plt.subplots(figsize=(17, 10)) + official_values = [ + row["official_tp1"]["median_ms"] for row in single_rows + ] + triton_values = [row["triton"]["median_ms"] for row in single_rows] + official_bars = axis.bar( + positions - width / 2, + official_values, + width, + label="Official Qwen3MLP TP=1", + color="#2563eb", + ) + triton_bars = axis.bar( + positions + width / 2, + triton_values, + width, + label="Deterministic Triton FFN (packed)", + color="#7c3aed", + ) + for bars, values in ( + (official_bars, official_values), + (triton_bars, triton_values), + ): + axis.bar_label( + bars, + labels=[f"{value:.2f}" for value in values], + padding=4, + fontsize=13, + rotation=90, + ) + axis.set_yscale("log") + axis.set_xlabel( + "Token count M and measured direction\nH=4096, I=12288, BF16" + ) + axis.set_ylabel("Median latency (ms, log scale)") + axis.set_title("MI300X single-GPU FFN speed: official TP=1 vs Triton") + axis.set_xticks(positions, single_labels) + axis.legend(loc="upper left") + figure.tight_layout() + else: + local_lookup = { + (row["tokens"], row["direction"]): row for row in single_rows + } + comparison_lookup = { + (row["tokens"], row["direction"]): row + for row in comparison_payload["single_gpu"] + } + figure, axes = plt.subplots(1, 2, figsize=(24, 10), sharey=True) + series = ( + ("CPU official BF16", "official_cpu_ms", "#6b7280"), + ("H100 official TP=1", "official_h100_ms", "#60a5fa"), + ("H100 Triton replay", "triton_h100_ms", "#06b6d4"), + ("H100 deterministic CUDA", "cuda_h100_ms", "#dc2626"), + ("MI300X official TP=1", "official_mi300x_ms", "#86efac"), + ( + "MI300X deterministic Triton (packed)", + "triton_mi300x_ms", + "#7c3aed", + ), + ) + width = 0.13 + for axis, direction, direction_label in ( + (axes[0], "forward", "Forward"), + (axes[1], "train_fwd_bwd", "Forward + backward"), + ): + tokens = (1, 8, 32) + positions = np.arange(len(tokens)) + combined_rows = [] + for token_count in tokens: + external = comparison_lookup[(token_count, direction)] + local = local_lookup[(token_count, direction)] + combined_rows.append( + { + **external, + "official_mi300x_ms": local["official_tp1"]["median_ms"], + "triton_mi300x_ms": local["triton"]["median_ms"], + } + ) + for series_index, (label, key, color) in enumerate(series): + values = [row[key] for row in combined_rows] + offset = (series_index - (len(series) - 1) / 2) * width + bars = axis.bar( + positions + offset, + values, + width, + label=label, + color=color, + ) + axis.bar_label( + bars, + labels=[f"{value:.2f}" for value in values], + padding=3, + fontsize=9, + rotation=90, + ) + axis.set_yscale("log") + axis.set_xlabel("Token count M\nH=4096, I=12288, BF16") + axis.set_title(direction_label) + axis.set_xticks(positions, [f"M={value}" for value in tokens]) + axes[0].set_ylabel("Median latency (ms, log scale)") + axes[0].set_ylim(0.07, 140.0) + handles, labels = axes[0].get_legend_handles_labels() + figure.legend( + handles, + labels, + loc="upper center", + bbox_to_anchor=(0.5, 0.91), + ncol=3, + ) + figure.suptitle( + "Single-GPU / CPU FFN absolute latency across platforms", + y=0.99, + ) + figure.text( + 0.5, + 0.01, + "Cross-hardware values are context only; H100 CUDA vs H100 Triton " + "is the hardware-matched comparison.", + ha="center", + fontsize=14, + ) + figure.tight_layout(rect=(0.0, 0.06, 1.0, 0.82)) + figure.savefig(output_directory / "single_gpu_overhead.png", dpi=180) + plt.close(figure) + + exactness_rows = _topology_exactness_rows(payload["distributed_ffn"]) + mismatch_keys = ( + "forward_output", + "training_output", + "hidden_gradient", + "weight_gradient", + ) + mismatch_labels = ( + "Forward\noutput", + "Training\noutput", + "dHidden", + "dWeights", + ) + matrix = np.asarray( + [[row[key] for key in mismatch_keys] for row in exactness_rows], + dtype=float, + ) + figure, axis = plt.subplots(figsize=(12, 10)) + image = axis.imshow( + matrix, + aspect="auto", + cmap="RdYlGn_r", + vmin=0, + vmax=max(1.0, matrix.max()), + ) + for y_position in range(matrix.shape[0]): + for x_position in range(matrix.shape[1]): + axis.text( + x_position, + y_position, + str(int(matrix[y_position, x_position])), + ha="center", + va="center", + fontsize=18, + fontweight="bold", + ) + axis.set_xticks(range(len(mismatch_labels)), mismatch_labels) + axis.set_yticks( + range(len(exactness_rows)), + [ + f"TP={row['tp_size']}, CP={row['cp_size']}, " + f"SP={'on' if row['sequence_parallel'] else 'off'}" + for row in exactness_rows + ], + ) + axis.set_xlabel("Compared tensor category") + axis.set_ylabel("Distributed parallel configuration") + axis.set_title("Topology mismatch vs Triton TP=1") + colorbar = figure.colorbar(image, ax=axis, fraction=0.03, pad=0.03) + colorbar.set_label("Mismatched elements") + figure.tight_layout() + figure.savefig(output_directory / "collective_overhead.png", dpi=180) + plt.close(figure) + + rows = payload["distributed_ffn"] + platform_comparison = _distributed_platform_comparison_rows( + rows, comparison_payload + ) + figure, axes = plt.subplots(1, 2, figsize=(26, 10), sharey=True) + if platform_comparison: + comparison_lookup = { + (row["name"], row["direction"]): row + for row in platform_comparison + } + distributed_series = ( + ("H100 official distributed", "h100_official_distributed_ms", "#60a5fa"), + ("H100 deterministic CUDA", "h100_deterministic_cuda_ms", "#dc2626"), + ("MI300X official distributed", "mi300x_official_distributed_ms", "#86efac"), + ( + "MI300X deterministic Triton", + "mi300x_deterministic_triton_ms", + "#7c3aed", + ), + ) + width = 0.19 + else: + distributed_series = ( + ("MI300X official distributed", "mi300x_official_distributed_ms", "#86efac"), + ( + "MI300X deterministic Triton", + "mi300x_deterministic_triton_ms", + "#7c3aed", + ), + ) + width = 0.34 + for axis, direction, direction_label in ( + (axes[0], "forward", "Forward"), + (axes[1], "train_fwd_bwd", "Forward + backward"), + ): + direction_rows = [row for row in rows if row["direction"] == direction] + positions = np.arange(len(direction_rows)) + labels = [row["name"].upper().replace("_", "\n") for row in direction_rows] + combined_rows = [] + for row in direction_rows: + combined = { + "mi300x_official_distributed_ms": row["official_distributed"]["median_ms"], + "mi300x_deterministic_triton_ms": row["triton"]["median_ms"], + } + if platform_comparison: + combined.update(comparison_lookup[(row["name"], direction)]) + combined_rows.append(combined) + for series_index, (label, key, color) in enumerate(distributed_series): + values = [row[key] for row in combined_rows] + offset = ( + series_index - (len(distributed_series) - 1) / 2 + ) * width + bars = axis.bar( + positions + offset, + values, + width, + label=label, + color=color, + ) + axis.bar_label( + bars, + labels=[f"{value:.2f}" for value in values], + padding=3, + fontsize=10, + rotation=90, + ) + axis.set_yscale("log") + axis.set_xlabel("Parallel configuration\nM=32, H=4096, I=12288, BF16") + axis.set_title(direction_label) + axis.set_xticks(positions, labels) + axes[0].set_ylabel("Median latency (ms, log scale)") + handles, labels = axes[0].get_legend_handles_labels() + if not platform_comparison: + axes[0].legend(loc="upper left") + figure.suptitle( + "MI300X distributed FFN latency", + y=0.99, + ) + figure.tight_layout(rect=(0.0, 0.0, 1.0, 0.95)) + else: + axes[0].set_ylim(0.1, 30.0) + figure.legend( + handles, + labels, + loc="upper center", + bbox_to_anchor=(0.5, 0.91), + ncol=4, + ) + figure.suptitle( + "Distributed FFN latency: H100 CUDA vs MI300X Triton", + y=0.99, + ) + figure.text( + 0.5, + 0.01, + "All four paths use the same M=32 workload, direction, and TP/CP/SP " + "topology; no TP=1 latency is included.", + ha="center", + fontsize=14, + ) + figure.tight_layout(rect=(0.0, 0.06, 1.0, 0.82)) + figure.savefig(output_directory / "distributed_ffn_overhead.png", dpi=180) + plt.close(figure) + + +def _environment() -> dict[str, Any]: + properties = torch.cuda.get_device_properties(0) + return { + "gpu": torch.cuda.get_device_name(0), + "gpu_count": torch.cuda.device_count(), + "architecture": properties.gcnArchName, + "torch": torch.__version__, + "transformers": transformers_version, + "hip": torch.version.hip, + "python": os.sys.version.split()[0], + "git_commit": os.popen("git rev-parse HEAD").read().strip(), + "single_gpu_speed_context": "Hugging Face Transformers Qwen3MLP, TP=1", + "distributed_speed_comparison": "four same-topology H100/MI300X paths", + "deterministic_compute": "ROCm-native Triton", + "deterministic_transport": "fixed-tree HIP IPC with RCCL fallback on ROCm", + "NCCL_IB_DISABLE": os.environ.get("NCCL_IB_DISABLE", ""), + } + + +def _validate_environment() -> None: + if getattr(torch.version, "hip", None) is None: + raise RuntimeError("this benchmark requires a ROCm PyTorch build") + if not torch.cuda.is_available() or torch.cuda.device_count() < 8: + raise RuntimeError("this benchmark requires eight visible ROCm GPUs") + if not dist.is_available() or not dist.is_nccl_available(): + raise RuntimeError("PyTorch RCCL/ProcessGroupNCCL support is unavailable") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("benchmarks/results/pr325_rocm_mi300x"), + ) + parser.add_argument( + "--world-sizes", + type=int, + nargs="+", + choices=(2, 4, 8), + default=(2, 4, 8), + help="distributed world sizes to benchmark (default: 2 4 8)", + ) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--samples", type=int, default=10) + parser.add_argument("--training-samples", type=int, default=5) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + _validate_environment() + os.environ.setdefault("NCCL_IB_DISABLE", "1") + args.output_dir.mkdir(parents=True, exist_ok=True) + payload: dict[str, Any] = { + "environment": _environment(), + "methodology": { + "single_gpu_timing": "GPU events, median and p95", + "distributed_timing": "synchronized wall clock, slowest rank/sample", + "distributed_worker_cpu_affinity": "one NUMA-local CPU per GPU rank", + "warmup": args.warmup, + "samples": args.samples, + "training_samples": args.training_samples, + "operator_only": True, + "triton_weight_layout": "packed_forward_cache_outside_timed_region", + "tp1_forward_cache_bytes": _TP1_FORWARD_CACHE_BYTES, + }, + "communication_contract": _COMMUNICATION_CONTRACT, + "single_gpu": _single_gpu_benchmarks( + warmup=args.warmup, + samples=args.samples, + training_samples=args.training_samples, + ), + "distributed_ffn": [], + } + torch.cuda.empty_cache() + for world_size in args.world_sizes: + result = _run_distributed_world( + world_size, + warmup=args.warmup, + samples=args.samples, + training_samples=args.training_samples, + ) + payload["distributed_ffn"].extend(result["distributed_ffn"]) + comparison_payload = _load_cuda_cpu_comparison(args.output_dir) + platform_comparison = _distributed_platform_comparison_rows( + payload["distributed_ffn"], comparison_payload + ) + if platform_comparison: + payload["distributed_platform_comparison"] = { + "source": "cuda_cpu_comparison.json", + "contract": "same M=32 workload, direction, and TP/CP/SP topology", + "rows": platform_comparison, + } + previous_comparison = _previous_deterministic_comparison_rows( + payload["distributed_ffn"] + ) + if previous_comparison: + payload["previous_deterministic_comparison"] = { + "source": "previous checked MI300X benchmark before PR #357", + "rows": previous_comparison, + } + (args.output_dir / "results.json").write_text( + json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8" + ) + _write_report(payload, args.output_dir, comparison_payload) + _write_figures(payload, args.output_dir, comparison_payload) + print(json.dumps({"output_dir": str(args.output_dir), "status": "ok"})) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/profile_rocm_ffn.py b/benchmarks/profile_rocm_ffn.py new file mode 100644 index 00000000..099ffe0f --- /dev/null +++ b/benchmarks/profile_rocm_ffn.py @@ -0,0 +1,784 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Collect torch.profiler data for the strict ROCm Triton Qwen3 FFN. + +This script profiles the public ``qwen3_ffn`` API without changing or wrapping +its internal arithmetic. Profiler timings are intended for attribution only; +uninstrumented GPU-event samples are saved separately for latency comparisons. + +Example: + + python benchmarks/profile_rocm_ffn.py \ + --direction both \ + --tokens 32 \ + --warmup 3 \ + --active-steps 5 \ + --latency-samples 20 \ + --output-dir /tmp/rlk_torch_profiler_m32 +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import os +import platform +import statistics +import subprocess +import sys +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import torch +from torch.profiler import ProfilerActivity, profile, record_function + +from rl_engine.kernels.ops.triton.ffn import ( + Qwen3FFNForwardWeights, + pack_qwen3_ffn_forward_weights, + qwen3_ffn, +) + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_INPUT_SEEDS = {1: 3010, 8: 3012, 32: 3014} +_KEY_AVERAGE_FIELDS = ( + "name", + "count", + "device_type", + "self_cpu_time_us", + "cpu_time_us", + "self_device_time_us", + "device_time_us", + "self_cpu_memory_bytes", + "cpu_memory_bytes", + "self_device_memory_bytes", + "device_memory_bytes", + "input_shapes", + "stack", +) + + +@dataclass +class FFNCase: + direction: str + hidden: torch.Tensor + gate_weight: torch.Tensor + up_weight: torch.Tensor + down_weight: torch.Tensor + grad_output: torch.Tensor + forward_weights: Qwen3FFNForwardWeights | None + + @property + def slug(self) -> str: + return self.direction.replace("-", "_") + + @property + def training(self) -> bool: + return self.direction == "forward-backward" + + def clear_gradients(self) -> None: + for tensor in ( + self.hidden, + self.gate_weight, + self.up_weight, + self.down_weight, + ): + tensor.grad = None + + def run(self, *, use_forward_weights: bool = True) -> torch.Tensor: + forward_weights = self.forward_weights if use_forward_weights else None + if not self.training: + with torch.no_grad(): + return qwen3_ffn( + self.hidden, + self.gate_weight, + self.up_weight, + self.down_weight, + forward_weights=forward_weights, + ) + + self.clear_gradients() + output = qwen3_ffn( + self.hidden, + self.gate_weight, + self.up_weight, + self.down_weight, + forward_weights=forward_weights, + ) + output.backward(self.grad_output) + return output + + def result_tensors(self, output: torch.Tensor) -> dict[str, torch.Tensor]: + result = {"output": output} + if not self.training: + return result + + for name, tensor in ( + ("dhidden", self.hidden), + ("dgate_weight", self.gate_weight), + ("dup_weight", self.up_weight), + ("ddown_weight", self.down_weight), + ): + if tensor.grad is None: + raise RuntimeError(f"{name} was not produced by backward") + result[name] = tensor.grad + return result + + +def _randn( + shape: tuple[int, ...], + *, + seed: int, + device: torch.device, + requires_grad: bool, +) -> 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).requires_grad_(requires_grad) + + +def _build_case(args: argparse.Namespace, direction: str) -> FFNCase: + device = torch.device("cuda", args.device) + training = direction == "forward-backward" + input_seed = ( + args.input_seed + if args.input_seed is not None + else _INPUT_SEEDS.get(args.tokens, 3010) + ) + hidden = _randn( + (args.tokens, args.hidden_size), + seed=input_seed, + device=device, + requires_grad=training, + ) + gate_weight = _randn( + (args.intermediate_size, args.hidden_size), + seed=args.weight_seed, + device=device, + requires_grad=training, + ) + up_weight = _randn( + (args.intermediate_size, args.hidden_size), + seed=args.weight_seed + 1, + device=device, + requires_grad=training, + ) + down_weight = _randn( + (args.hidden_size, args.intermediate_size), + seed=args.weight_seed + 2, + device=device, + requires_grad=training, + ) + forward_weights = ( + pack_qwen3_ffn_forward_weights(gate_weight, up_weight, down_weight) + if args.weight_layout == "packed" + else None + ) + return FFNCase( + direction=direction, + hidden=hidden, + gate_weight=gate_weight, + up_weight=up_weight, + down_weight=down_weight, + grad_output=_randn( + (args.tokens, args.hidden_size), + seed=input_seed + 1, + device=device, + requires_grad=False, + ), + forward_weights=forward_weights, + ) + + +def _weight_layout_metadata(case: FFNCase) -> dict[str, Any]: + metadata: dict[str, Any] = {"additional_forward_weight_bytes": 0} + for name, weight in ( + ("gate_weight", case.gate_weight), + ("up_weight", case.up_weight), + ("down_weight", case.down_weight), + ): + metadata[name] = { + "shape": list(weight.shape), + "stride": list(weight.stride()), + "is_contiguous": weight.is_contiguous(), + } + if case.forward_weights is not None: + packed_weights = ( + ("gate_weight_t", case.forward_weights.gate_weight_t), + ("up_weight_t", case.forward_weights.up_weight_t), + ("down_weight_t", case.forward_weights.down_weight_t), + ) + packed_tensors = { + name: { + "shape": list(weight.shape), + "stride": list(weight.stride()), + "is_contiguous": weight.is_contiguous(), + "requires_grad": weight.requires_grad, + "nbytes": weight.numel() * weight.element_size(), + } + for name, weight in packed_weights + } + additional_bytes = sum( + weight.numel() * weight.element_size() for _, weight in packed_weights + ) + metadata["additional_forward_weight_bytes"] = additional_bytes + metadata["packed_forward_weights"] = { + "additional_bytes": additional_bytes, + "tensors": packed_tensors, + } + return metadata + + +def _git_output(*arguments: str) -> str: + try: + completed = subprocess.run( + ["git", *arguments], + cwd=_REPO_ROOT, + check=True, + capture_output=True, + text=True, + ) + except (OSError, subprocess.CalledProcessError): + return "" + return completed.stdout.strip() + + +def _write_json(path: Path, payload: Any) -> None: + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + + +def _percentile(values: list[float], percentile: float) -> float: + ordered = sorted(values) + position = (len(ordered) - 1) * percentile + lower = int(position) + upper = min(lower + 1, len(ordered) - 1) + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def _latency_summary(samples_ms: list[float]) -> dict[str, Any]: + if not samples_ms: + return {"samples_ms": []} + return { + "samples_ms": samples_ms, + "median_ms": statistics.median(samples_ms), + "p95_ms": _percentile(samples_ms, 0.95), + "min_ms": min(samples_ms), + "max_ms": max(samples_ms), + } + + +def _gpu_event_samples(case: FFNCase, samples: int) -> list[float]: + if samples == 0: + return [] + events = [ + ( + torch.cuda.Event(enable_timing=True), + torch.cuda.Event(enable_timing=True), + ) + for _ in range(samples) + ] + for start, end in events: + start.record() + output = case.run() + end.record() + del output + torch.cuda.synchronize() + case.clear_gradients() + return [float(start.elapsed_time(end)) for start, end in events] + + +def _tensor_fingerprint(tensor: torch.Tensor) -> dict[str, Any]: + detached = tensor.detach().contiguous() + raw = detached.view(torch.uint8).cpu().numpy() + return { + "shape": list(detached.shape), + "dtype": str(detached.dtype), + "nbytes": detached.numel() * detached.element_size(), + "sha256_raw_bytes": hashlib.sha256(memoryview(raw)).hexdigest(), + } + + +def _run_and_fingerprint( + case: FFNCase, + *, + use_forward_weights: bool = True, +) -> dict[str, Any]: + output = case.run(use_forward_weights=use_forward_weights) + torch.cuda.synchronize() + fingerprints = { + name: _tensor_fingerprint(tensor) + for name, tensor in case.result_tensors(output).items() + } + del output + case.clear_gradients() + return fingerprints + + +def _json_safe(value: Any) -> Any: + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, (list, tuple)): + return [_json_safe(item) for item in value] + if isinstance(value, dict): + return {str(key): _json_safe(item) for key, item in value.items()} + return str(value) + + +def _key_average_rows(profiler: profile, *, group_by_input_shape: bool) -> list[dict[str, Any]]: + rows: list[dict[str, Any]] = [] + for event in profiler.key_averages(group_by_input_shape=group_by_input_shape): + rows.append( + { + "name": event.key, + "count": event.count, + "device_type": str(event.device_type), + "self_cpu_time_us": event.self_cpu_time_total, + "cpu_time_us": event.cpu_time_total, + "self_device_time_us": event.self_device_time_total, + "device_time_us": event.device_time_total, + "self_cpu_memory_bytes": event.self_cpu_memory_usage, + "cpu_memory_bytes": event.cpu_memory_usage, + "self_device_memory_bytes": event.self_device_memory_usage, + "device_memory_bytes": event.device_memory_usage, + "input_shapes": _json_safe(event.input_shapes), + "stack": _json_safe(event.stack), + } + ) + return sorted( + rows, + key=lambda row: ( + float(row["self_device_time_us"]), + float(row["self_cpu_time_us"]), + ), + reverse=True, + ) + + +def _write_key_averages(output_dir: Path, slug: str, rows: list[dict[str, Any]]) -> None: + _write_json(output_dir / f"{slug}.key_averages.json", rows) + with (output_dir / f"{slug}.key_averages.csv").open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=_KEY_AVERAGE_FIELDS) + writer.writeheader() + for row in rows: + csv_row = dict(row) + csv_row["input_shapes"] = json.dumps(row["input_shapes"]) + csv_row["stack"] = json.dumps(row["stack"]) + writer.writerow(csv_row) + + +def _kernel_breakdown(rows: list[dict[str, Any]]) -> dict[str, Any]: + device_rows = [ + row + for row in rows + if str(row["device_type"]).endswith(("CUDA", "HIP")) + # Kineto also emits a synthetic device-side aggregate for a user + # annotation. It overlaps all child kernels and must not be summed. + and not str(row["name"]).startswith("rl_kernel::") + ] + categories: dict[str, dict[str, float | int]] = { + "layout_copy": {"count": 0, "self_device_time_us": 0.0}, + "gemm_leaf": {"count": 0, "self_device_time_us": 0.0}, + "gemm_reduce": {"count": 0, "self_device_time_us": 0.0}, + "gemm_root_copy": {"count": 0, "self_device_time_us": 0.0}, + "swiglu": {"count": 0, "self_device_time_us": 0.0}, + "elementwise_add": {"count": 0, "self_device_time_us": 0.0}, + "other_device_kernels": {"count": 0, "self_device_time_us": 0.0}, + } + for row in device_rows: + name = str(row["name"]).lower() + if "direct_copy_kernel" in name: + category = "layout_copy" + elif "det_gemm_tree_leaf" in name: + category = "gemm_leaf" + elif "det_gemm_tree_reduce" in name: + category = "gemm_reduce" + elif "copy_tree_root" in name: + category = "gemm_root_copy" + elif "swiglu" in name: + category = "swiglu" + elif "functor_add" in name: + category = "elementwise_add" + else: + category = "other_device_kernels" + categories[category]["count"] += int(row["count"]) + categories[category]["self_device_time_us"] += float( + row["self_device_time_us"] + ) + + total_device_us = sum( + float(category["self_device_time_us"]) for category in categories.values() + ) + for category in categories.values(): + category["percent_of_device_kernel_time"] = ( + 100.0 * float(category["self_device_time_us"]) / total_device_us + if total_device_us + else 0.0 + ) + return { + "accounting": ( + "sum of device kernel events; synthetic rl_kernel:: annotation " + "device aggregates are excluded to avoid double counting" + ), + "total_device_kernel_time_us": total_device_us, + "categories": categories, + "top_device_kernels": device_rows[:25], + } + + +def _environment(args: argparse.Namespace) -> dict[str, Any]: + properties = torch.cuda.get_device_properties(args.device) + try: + import triton + + triton_version = triton.__version__ + except (ImportError, AttributeError): + triton_version = "" + status = _git_output("status", "--short") + return { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "command": [sys.executable, *sys.argv], + "hostname": platform.node(), + "python": platform.python_version(), + "torch": torch.__version__, + "hip": torch.version.hip, + "triton": triton_version, + "gpu_index": args.device, + "gpu": torch.cuda.get_device_name(args.device), + "architecture": getattr(properties, "gcnArchName", ""), + "total_memory_bytes": properties.total_memory, + "gpu_count": torch.cuda.device_count(), + "git_commit": _git_output("rev-parse", "HEAD"), + "git_branch": _git_output("branch", "--show-current"), + "git_dirty": bool(status), + "git_status": status.splitlines(), + "environment": { + name: os.environ.get(name, "") + for name in ( + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "CUDA_VISIBLE_DEVICES", + "NCCL_IB_DISABLE", + ) + }, + } + + +def _profile_case( + case: FFNCase, + args: argparse.Namespace, + output_dir: Path, +) -> dict[str, Any]: + for _ in range(args.warmup): + output = case.run() + del output + torch.cuda.synchronize() + case.clear_gradients() + + latency = _latency_summary(_gpu_event_samples(case, args.latency_samples)) + _write_json(output_dir / f"{case.slug}.latency.json", latency) + + # The reference deliberately reuses the exact canonical tensor objects but + # passes no forward-weight cache. In packed mode this exercises the original + # per-call transpose path independently of the profiled candidate. + standard_reference_fingerprints = ( + _run_and_fingerprint(case, use_forward_weights=False) + if not args.skip_bitwise_hash + else {} + ) + before_profile_fingerprints = ( + _run_and_fingerprint(case) if not args.skip_bitwise_hash else {} + ) + torch.cuda.synchronize() + + device = torch.device("cuda", args.device) + torch.cuda.reset_peak_memory_stats(device) + baseline_allocated = torch.cuda.memory_allocated(device) + baseline_reserved = torch.cuda.memory_reserved(device) + + effective_record_shapes = args.record_shapes or args.profile_memory + effective_with_stack = args.with_stack or args.profile_memory + with profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=effective_record_shapes, + profile_memory=args.profile_memory, + with_stack=effective_with_stack, + acc_events=True, + ) as profiler: + for _ in range(args.active_steps): + with record_function(f"rl_kernel::{case.direction}"): + output = case.run() + del output + profiler.step() + torch.cuda.synchronize() + + peak_allocated = torch.cuda.max_memory_allocated(device) + peak_reserved = torch.cuda.max_memory_reserved(device) + memory = { + "baseline_allocated_bytes": baseline_allocated, + "baseline_reserved_bytes": baseline_reserved, + "peak_allocated_bytes": peak_allocated, + "peak_reserved_bytes": peak_reserved, + "peak_allocated_delta_bytes": max(0, peak_allocated - baseline_allocated), + "peak_reserved_delta_bytes": max(0, peak_reserved - baseline_reserved), + } + + trace_path = output_dir / f"{case.slug}.trace.json.gz" + profiler.export_chrome_trace(str(trace_path)) + rows = _key_average_rows( + profiler, + group_by_input_shape=effective_record_shapes, + ) + _write_key_averages(output_dir, case.slug, rows) + kernel_breakdown = _kernel_breakdown(rows) + _write_json(output_dir / f"{case.slug}.kernel_breakdown.json", kernel_breakdown) + + key_averages = profiler.key_averages( + group_by_input_shape=effective_record_shapes + ) + summary = "\n\n".join( + ( + "Sorted by self device time\n" + + key_averages.table( + sort_by="self_device_time_total", + row_limit=args.row_limit, + ), + "Sorted by self CPU time\n" + + key_averages.table( + sort_by="self_cpu_time_total", + row_limit=args.row_limit, + ), + ) + ) + (output_dir / f"{case.slug}.summary.txt").write_text(summary) + + memory_timeline_error = "" + if args.profile_memory: + try: + profiler.export_memory_timeline( + str(output_dir / f"{case.slug}.memory.raw.json.gz"), + device=str(device), + ) + except (AssertionError, RuntimeError, ValueError) as error: + memory_timeline_error = f"{type(error).__name__}: {error}" + + after_profile_fingerprints = ( + _run_and_fingerprint(case) if not args.skip_bitwise_hash else {} + ) + before_matches_reference = { + name: before_profile_fingerprints.get(name) == fingerprint + for name, fingerprint in standard_reference_fingerprints.items() + } + after_matches_reference = { + name: after_profile_fingerprints.get(name) == fingerprint + for name, fingerprint in standard_reference_fingerprints.items() + } + after_matches_before = { + name: after_profile_fingerprints.get(name) == fingerprint + for name, fingerprint in before_profile_fingerprints.items() + } + all_match = ( + None + if args.skip_bitwise_hash + else ( + all(before_matches_reference.values()) + and all(after_matches_reference.values()) + and all(after_matches_before.values()) + ) + ) + correctness = { + "comparison": "SHA256 over the contiguous raw tensor bytes", + "reference_mode": ( + "standard qwen3_ffn call with forward_weights=None using the same " + "canonical input and weight tensors" + ), + "candidate_mode": args.weight_layout, + "skipped": args.skip_bitwise_hash, + "all_match": all_match, + "matches_standard_reference_before_profile": before_matches_reference, + "matches_standard_reference_after_profile": after_matches_reference, + "matches_before_profile_after_profile": after_matches_before, + "standard_reference": standard_reference_fingerprints, + "before_profile": before_profile_fingerprints, + "after_profile": after_profile_fingerprints, + } + _write_json(output_dir / f"{case.slug}.correctness.json", correctness) + case.clear_gradients() + + return { + "direction": case.direction, + "weight_layout": args.weight_layout, + "weight_layout_metadata": _weight_layout_metadata(case), + "files": { + "trace": trace_path.name, + "key_averages_json": f"{case.slug}.key_averages.json", + "key_averages_csv": f"{case.slug}.key_averages.csv", + "summary": f"{case.slug}.summary.txt", + "kernel_breakdown": f"{case.slug}.kernel_breakdown.json", + "latency": f"{case.slug}.latency.json", + "correctness": f"{case.slug}.correctness.json", + "memory_timeline": ( + f"{case.slug}.memory.raw.json.gz" if args.profile_memory else "" + ), + }, + "profiler": { + "active_steps": args.active_steps, + "record_shapes": effective_record_shapes, + "profile_memory": args.profile_memory, + "with_stack": effective_with_stack, + "trace_size_bytes": trace_path.stat().st_size, + "memory_timeline_error": memory_timeline_error, + }, + "latency": latency, + "memory": memory, + "correctness_all_match": correctness["all_match"], + "kernel_breakdown": kernel_breakdown, + } + + +def _validate_args(args: argparse.Namespace) -> None: + if getattr(torch.version, "hip", None) is None: + raise RuntimeError("this profiler requires a ROCm PyTorch build") + if not torch.cuda.is_available(): + raise RuntimeError("no ROCm GPU is available") + if args.device < 0 or args.device >= torch.cuda.device_count(): + raise ValueError( + f"--device must be in [0, {torch.cuda.device_count() - 1}], got {args.device}" + ) + for name in ( + "tokens", + "hidden_size", + "intermediate_size", + "active_steps", + ): + if getattr(args, name) <= 0: + raise ValueError(f"--{name.replace('_', '-')} must be positive") + if args.warmup < 0 or args.latency_samples < 0: + raise ValueError("--warmup and --latency-samples must be non-negative") + if args.profile_memory and args.active_steps > 1: + print( + "warning: memory/shape/stack profiling adds overhead; " + "prefer --active-steps 1 for a memory run", + file=sys.stderr, + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("/tmp/rl_kernel_torch_profiler"), + ) + parser.add_argument( + "--direction", + choices=("forward", "forward-backward", "both"), + default="both", + ) + parser.add_argument("--tokens", type=int, default=32) + parser.add_argument("--hidden-size", type=int, default=4096) + parser.add_argument("--intermediate-size", type=int, default=12288) + parser.add_argument("--device", type=int, default=0) + parser.add_argument( + "--weight-layout", + choices=("standard", "packed"), + default="standard", + help=( + "standard materializes forward transposes in every FFN call; packed " + "materializes detached forward-only weights once before measurement" + ), + ) + parser.add_argument("--weight-seed", type=int, default=3000) + parser.add_argument("--input-seed", type=int) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--active-steps", type=int, default=5) + parser.add_argument("--latency-samples", type=int, default=20) + parser.add_argument("--row-limit", type=int, default=100) + parser.add_argument("--record-shapes", action="store_true") + parser.add_argument("--profile-memory", action="store_true") + parser.add_argument("--with-stack", action="store_true") + parser.add_argument("--skip-bitwise-hash", action="store_true") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + _validate_args(args) + torch.cuda.set_device(args.device) + torch.backends.cuda.matmul.allow_tf32 = False + args.output_dir.mkdir(parents=True, exist_ok=True) + + directions = ( + ("forward", "forward-backward") + if args.direction == "both" + else (args.direction,) + ) + manifest: dict[str, Any] = { + "environment": _environment(args), + "workload": { + "operator": "rl_engine.kernels.ops.triton.ffn.qwen3_ffn", + "tokens": args.tokens, + "hidden_size": args.hidden_size, + "intermediate_size": args.intermediate_size, + "dtype": "torch.bfloat16", + "weight_layout": args.weight_layout, + "weight_seeds": [args.weight_seed + offset for offset in range(3)], + "input_seed": ( + args.input_seed + if args.input_seed is not None + else _INPUT_SEEDS.get(args.tokens, 3010) + ), + "grad_output_seed": ( + args.input_seed + 1 + if args.input_seed is not None + else _INPUT_SEEDS.get(args.tokens, 3010) + 1 + ), + "warmup": args.warmup, + "latency_samples": args.latency_samples, + "profiler_active_steps": args.active_steps, + }, + "methodology": { + "profiler_use": "attribution and launch analysis only", + "latency_use": "uninstrumented GPU events", + "jit_and_tree_plan": "warmed before profiler starts", + "weight_packing": ( + "performed once during case construction outside all timing" + ), + "bitwise_fingerprint": "SHA256 over raw BF16 bytes", + "packed_bitwise_reference": ( + "uncached standard qwen3_ffn path using the same canonical tensors" + ), + }, + "profiles": [], + } + _write_json(args.output_dir / "manifest.json", manifest) + + for direction in directions: + print(f"profiling {direction} on cuda:{args.device} ...", flush=True) + case = _build_case(args, direction) + result = _profile_case(case, args, args.output_dir) + manifest["profiles"].append(result) + _write_json(args.output_dir / "manifest.json", manifest) + latency = result["latency"] + if latency.get("samples_ms"): + print( + f" uninstrumented median={latency['median_ms']:.4f} ms, " + f"p95={latency['p95_ms']:.4f} ms", + flush=True, + ) + print( + f" raw-bit fingerprints match: {result['correctness_all_match']}", + flush=True, + ) + del case + torch.cuda.empty_cache() + + print(f"profiler artifacts: {args.output_dir.resolve()}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/results/pr325_rocm_mi300x/collective_overhead.png b/benchmarks/results/pr325_rocm_mi300x/collective_overhead.png new file mode 100644 index 00000000..93f440bc Binary files /dev/null and b/benchmarks/results/pr325_rocm_mi300x/collective_overhead.png differ diff --git a/benchmarks/results/pr325_rocm_mi300x/cuda_cpu_comparison.json b/benchmarks/results/pr325_rocm_mi300x/cuda_cpu_comparison.json new file mode 100644 index 00000000..44baf8b3 --- /dev/null +++ b/benchmarks/results/pr325_rocm_mi300x/cuda_cpu_comparison.json @@ -0,0 +1,79 @@ +{ + "source": { + "report": "PR #321 deterministic CUDA FFN performance report", + "pull_request": "https://github.com/RL-Align/RL-Kernel/pull/321", + "cuda_commit": "8576fa4bf449734ae99e9b50be8756bb282a8916", + "h100_triton_replay_commit": "e64abab904880b877d26d04c0cfad020b992aa51", + "note": "User-supplied measurements; cross-hardware values are context only and are not hardware-normalized." + }, + "environment": { + "gpu": "NVIDIA H100 80GB HBM3", + "gpu_count": 8, + "architecture": "sm_90", + "cuda": "13.0", + "torch": "2.13.0+cu130", + "transformers": "5.13.1", + "python": "3.11.15", + "cpu": "Intel(R) Xeon(R) Platinum 8468", + "cpu_threads": 96, + "deterministic_compute": "native CUDA kernels", + "deterministic_transport": "fixed-order CUDA IPC" + }, + "methodology": { + "tokens": [1, 8, 32], + "hidden": 4096, + "intermediate": 12288, + "dtype": "bfloat16", + "warmup": 3, + "samples": 10, + "training_samples": 5, + "cpu_timing": "wall clock with 96 PyTorch intra-op threads", + "gpu_timing": "GPU events for single GPU; synchronized slowest-rank wall clock for distributed" + }, + "single_gpu": [ + {"tokens": 1, "direction": "forward", "official_h100_ms": 0.1193, "cuda_h100_ms": 3.9988, "triton_h100_ms": 1.7381, "official_cpu_ms": 12.9558}, + {"tokens": 1, "direction": "train_fwd_bwd", "official_h100_ms": 0.5184, "cuda_h100_ms": 9.0704, "triton_h100_ms": 4.2997, "official_cpu_ms": 58.5765}, + {"tokens": 8, "direction": "forward", "official_h100_ms": 0.1239, "cuda_h100_ms": 3.9923, "triton_h100_ms": 1.8293, "official_cpu_ms": 12.7676}, + {"tokens": 8, "direction": "train_fwd_bwd", "official_h100_ms": 0.5414, "cuda_h100_ms": 9.4500, "triton_h100_ms": 4.5028, "official_cpu_ms": 82.3669}, + {"tokens": 32, "direction": "forward", "official_h100_ms": 0.1311, "cuda_h100_ms": 4.0277, "triton_h100_ms": 2.2529, "official_cpu_ms": 8.5392}, + {"tokens": 32, "direction": "train_fwd_bwd", "official_h100_ms": 0.7842, "cuda_h100_ms": 9.1635, "triton_h100_ms": 5.4261, "official_cpu_ms": 65.9689} + ], + "distributed": [ + {"name": "tp2", "direction": "forward", "official_h100_ms": 0.1552, "cuda_h100_ms": 2.7896}, + {"name": "tp2", "direction": "train_fwd_bwd", "official_h100_ms": 0.6301, "cuda_h100_ms": 7.8125}, + {"name": "tp2_sp", "direction": "forward", "official_h100_ms": 0.1552, "cuda_h100_ms": 3.3310}, + {"name": "tp2_sp", "direction": "train_fwd_bwd", "official_h100_ms": 0.6301, "cuda_h100_ms": 11.7695}, + {"name": "tp4", "direction": "forward", "official_h100_ms": 0.1530, "cuda_h100_ms": 2.1687}, + {"name": "tp4", "direction": "train_fwd_bwd", "official_h100_ms": 0.5714, "cuda_h100_ms": 9.3954}, + {"name": "tp2_cp2", "direction": "forward", "official_h100_ms": 0.1530, "cuda_h100_ms": 2.8261}, + {"name": "tp2_cp2", "direction": "train_fwd_bwd", "official_h100_ms": 0.5714, "cuda_h100_ms": 16.1722}, + {"name": "tp2_cp2_sp", "direction": "forward", "official_h100_ms": 0.1530, "cuda_h100_ms": 3.5814}, + {"name": "tp2_cp2_sp", "direction": "train_fwd_bwd", "official_h100_ms": 0.5714, "cuda_h100_ms": 16.7894}, + {"name": "tp8", "direction": "forward", "official_h100_ms": 0.1537, "cuda_h100_ms": 2.3206}, + {"name": "tp8", "direction": "train_fwd_bwd", "official_h100_ms": 0.5006, "cuda_h100_ms": 10.1375}, + {"name": "tp4_cp2", "direction": "forward", "official_h100_ms": 0.1537, "cuda_h100_ms": 2.2141}, + {"name": "tp4_cp2", "direction": "train_fwd_bwd", "official_h100_ms": 0.5006, "cuda_h100_ms": 16.5617}, + {"name": "tp4_cp2_sp", "direction": "forward", "official_h100_ms": 0.1537, "cuda_h100_ms": 3.2972}, + {"name": "tp4_cp2_sp", "direction": "train_fwd_bwd", "official_h100_ms": 0.5006, "cuda_h100_ms": 17.1457} + ], + "dtype_accuracy": { + "tokens": 8, + "hidden": 4096, + "intermediate": 12288, + "candidate_dtype": "float16", + "reference_dtype": "float32", + "max_abs": 2.077e-6, + "mean_abs": 3.736e-7, + "relative_l2": 6.540e-4 + }, + "exactness": { + "reference": "deterministic CUDA TP=1", + "topologies": ["tp2", "tp2_sp", "tp4", "tp2_cp2", "tp2_cp2_sp", "tp8", "tp4_cp2", "tp4_cp2_sp"], + "forward_output_mismatch": 0, + "training_output_mismatch": 0, + "hidden_gradient_mismatch": 0, + "weight_gradient_mismatch": 0, + "repeat_mismatch": 0, + "train_infer_mismatch": 0 + } +} diff --git a/benchmarks/results/pr325_rocm_mi300x/distributed_ffn_overhead.png b/benchmarks/results/pr325_rocm_mi300x/distributed_ffn_overhead.png new file mode 100644 index 00000000..28584265 Binary files /dev/null and b/benchmarks/results/pr325_rocm_mi300x/distributed_ffn_overhead.png differ diff --git a/benchmarks/results/pr325_rocm_mi300x/report.md b/benchmarks/results/pr325_rocm_mi300x/report.md new file mode 100644 index 00000000..2c371bb4 --- /dev/null +++ b/benchmarks/results/pr325_rocm_mi300x/report.md @@ -0,0 +1,184 @@ +# PR #325 ROCm deterministic Triton FFN report + +This is an operator-only MI300X report. It does not load or benchmark a model checkpoint. + +## Comparison contract + +1. **Determinism:** every Triton TP/CP/SP result is compared bitwise with the same deterministic Triton FFN at **TP=1**. The reported metric is element mismatch count; acceptance requires 0. +2. **FP16/FP32:** one separate, simple output comparison runs official Hugging Face `Qwen3MLP` at TP=1 in FP16 and FP32. FP32 is the reference. +3. **Speed:** single-GPU speed retains the official Qwen3MLP TP=1 context. Distributed speed compares four same-topology paths: H100 official/deterministic and MI300X official/deterministic. + +## Environment + +| Field | Value | +|---|---| +| NCCL_IB_DISABLE | 1 | +| architecture | gfx942:sramecc+:xnack- | +| deterministic_compute | ROCm-native Triton | +| deterministic_transport | fixed-tree HIP IPC with RCCL fallback on ROCm | +| distributed_speed_comparison | four same-topology H100/MI300X paths | +| git_commit | caef501101a3906c733076f31f3b5a9870169d16 | +| gpu | AMD Instinct MI300X | +| gpu_count | 8 | +| hip | 7.14.60850 | +| python | 3.12.3 | +| single_gpu_speed_context | Hugging Face Transformers Qwen3MLP, TP=1 | +| torch | 2.12.0+rocm7.14.0a20260608 | +| transformers | 5.10.4 | + +## Methodology + +- Operator shape: H=4096, I=12288; BF16 is used for all speed and determinism measurements. +- Single-GPU shapes use M=1/8/32. Distributed cases use the same full logical M=32 input for TP2/4/8, TP+CP, and sequence parallelism. +- The distributed comparison joins rows by topology and direction: H100 and MI300X use the same logical M=32 workload and TP/CP/SP layout. Official TP=1 latency is neither collected nor used in that distributed ratio. +- MI300X official distributed uses upstream Qwen3 FFN math with native PyTorch BF16 GEMMs and native RCCL collectives over the same shards; the deterministic path uses the current Triton FFN and fixed-order transport. +- Deterministic Triton timings use the explicit prepacked forward-weight cache. Packing happens once outside the timed region; canonical source weights remain the autograd and optimizer source of truth. +- The TP=1 cache adds 288 MiB; each TP rank holds that amount divided by TP size. Refresh cost is excluded because the benchmark measures the steady-state FFN call. +- The distributed exactness baseline is the PR's deterministic Triton FFN at TP=1. Local outputs, dHidden, and sharded dWeights are compared against their exact TP=1 slices. +- Communication contract for TP+CP+SP: forward uses 1 TP AllGather plus 1 TP ReduceScatter (2 calls). Forward+backward retains 7 AllGathers plus 3 logical ReduceScatter lanes; PR #357 merges the two independent backward gate/up lanes into one `reduce_scatter_many` call, for 9 collective invocations. +- Implementation note: the current ROCm deterministic communication operator is adopted from PR #357. This changes the implementation under test, not the benchmark comparison contract. +- Single-GPU timing: GPU events, median and p95; distributed timing: synchronized wall clock, slowest rank/sample. +- Distributed workers: one NUMA-local CPU per GPU rank to reduce host-scheduler noise in synchronized wall-clock samples. +- 10 warmups, 50 measured forward samples, and 20 measured forward+backward samples. +- `NCCL_IB_DISABLE=1` keeps the distributed run on intra-node XGMI. Median, p95, min, and max values are available in `results.json`. + +Reproduce from the repository root: + +```bash +python benchmarks/benchmark_rocm_ffn.py \ + --warmup 10 \ + --samples 50 \ + --training-samples 20 \ + --output-dir benchmarks/results/pr325_rocm_mi300x +``` + +## Results summary + +- TP=1 exactness baseline: **0 mismatched elements** across topology forward outputs, training outputs, dHidden, and dWeights. +- Repeat mismatch: **0**; training/inference forward mismatch: **0**. +- Single-GPU deterministic Triton packed-cache latency is **3.92-7.64x** the official Qwen3MLP TP=1 latency across M=1/8/32 and forward/training. +- MI300X deterministic Triton latency is **0.14-0.38x** the H100 deterministic CUDA latency for the same distributed layouts. Both official distributed paths are reported alongside them. +- Versus the previous deterministic MI300X benchmark, PR #357 improves **16/16** rows, with a mean latency reduction of **22.1%**. +- The separate official-Qwen3MLP FP16 versus FP32 observation has relative-L2 error **6.544e-04** for (M,H,I)=(8,4096,12288). + +## Single-GPU FFN speed + +Performance only; no official-versus-Triton accuracy metric is reported here. + +| Shape / direction | Official Qwen3MLP TP=1 (ms) | Deterministic Triton, packed (ms) | Triton / official TP=1 | +|---|---:|---:|---:| +| (M,H,I)=(1,4096,12288), forward | 0.1004 | 0.6053 | 6.03x | +| (M,H,I)=(1,4096,12288), forward+backward | 0.4371 | 1.7138 | 3.92x | +| (M,H,I)=(8,4096,12288), forward | 0.1077 | 0.6054 | 5.62x | +| (M,H,I)=(8,4096,12288), forward+backward | 0.6476 | 2.8767 | 4.44x | +| (M,H,I)=(32,4096,12288), forward | 0.1146 | 0.8761 | 7.64x | +| (M,H,I)=(32,4096,12288), forward+backward | 0.4182 | 2.3188 | 5.54x | + +## Distributed FFN speed + +Every row compares the same distributed topology and direction. No TP=1 latency is used in this table. + +| Parallel layout | Direction | H100 official distributed (ms) | H100 deterministic CUDA (ms) | MI300X official distributed (ms) | MI300X deterministic Triton (ms) | H100 det / official | MI300X det / official | +|---|---|---:|---:|---:|---:|---:|---:| +| tp2 | forward | 0.1552 | 2.7896 | 0.2323 | 0.8759 | 17.97x | 3.77x | +| tp2 | train_fwd_bwd | 0.6301 | 7.8125 | 0.6984 | 2.0802 | 12.40x | 2.98x | +| tp2_sp | forward | 0.1552 | 3.3310 | 0.3066 | 0.8885 | 21.46x | 2.90x | +| tp2_sp | train_fwd_bwd | 0.6301 | 11.7695 | 1.4217 | 2.5187 | 18.68x | 1.77x | +| tp4 | forward | 0.1530 | 2.1687 | 0.2380 | 0.8324 | 14.17x | 3.50x | +| tp4 | train_fwd_bwd | 0.5714 | 9.3954 | 0.8280 | 2.2969 | 16.44x | 2.77x | +| tp2_cp2 | forward | 0.1530 | 2.8261 | 0.2317 | 0.7562 | 18.47x | 3.26x | +| tp2_cp2 | train_fwd_bwd | 0.5714 | 16.1722 | 4.4830 | 2.5281 | 28.30x | 0.56x | +| tp2_cp2_sp | forward | 0.1530 | 3.5814 | 0.3262 | 0.8067 | 23.41x | 2.47x | +| tp2_cp2_sp | train_fwd_bwd | 0.5714 | 16.7894 | 4.5870 | 2.5462 | 29.38x | 0.56x | +| tp8 | forward | 0.1537 | 2.3206 | 0.2075 | 0.7065 | 15.10x | 3.41x | +| tp8 | train_fwd_bwd | 0.5006 | 10.1375 | 0.8685 | 2.0507 | 20.25x | 2.36x | +| tp4_cp2 | forward | 0.1537 | 2.2141 | 0.2424 | 0.8306 | 14.41x | 3.43x | +| tp4_cp2 | train_fwd_bwd | 0.5006 | 16.5617 | 3.0071 | 2.4620 | 33.08x | 0.82x | +| tp4_cp2_sp | forward | 0.1537 | 3.2972 | 0.3153 | 0.7618 | 21.45x | 2.42x | +| tp4_cp2_sp | train_fwd_bwd | 0.5006 | 17.1457 | 2.9346 | 2.4232 | 34.25x | 0.83x | + +### PR #357 latency change versus the previous benchmark + +This comparison changes only the deterministic ROCm communication implementation. It is not included as another series in the main four-path figure. + +| Parallel layout | Direction | Previous (ms) | Current (ms) | Latency reduction | +|---|---|---:|---:|---:| +| tp2 | forward | 0.9040 | 0.8759 | 3.1% | +| tp2 | train_fwd_bwd | 2.6997 | 2.0802 | 22.9% | +| tp2_sp | forward | 1.0561 | 0.8885 | 15.9% | +| tp2_sp | train_fwd_bwd | 2.9646 | 2.5187 | 15.0% | +| tp4 | forward | 0.8359 | 0.8324 | 0.4% | +| tp4 | train_fwd_bwd | 2.6999 | 2.2969 | 14.9% | +| tp2_cp2 | forward | 0.9016 | 0.7562 | 16.1% | +| tp2_cp2 | train_fwd_bwd | 3.5606 | 2.5281 | 29.0% | +| tp2_cp2_sp | forward | 1.1296 | 0.8067 | 28.6% | +| tp2_cp2_sp | train_fwd_bwd | 3.9929 | 2.5462 | 36.2% | +| tp8 | forward | 1.0860 | 0.7065 | 34.9% | +| tp8 | train_fwd_bwd | 2.5620 | 2.0507 | 20.0% | +| tp4_cp2 | forward | 1.0536 | 0.8306 | 21.2% | +| tp4_cp2 | train_fwd_bwd | 3.2599 | 2.4620 | 24.5% | +| tp4_cp2_sp | forward | 1.1928 | 0.7618 | 36.1% | +| tp4_cp2_sp | train_fwd_bwd | 3.7498 | 2.4232 | 35.4% | + +## CUDA GPU and CPU performance context + +The additional measurements come from [PR #321 deterministic CUDA FFN performance report](https://github.com/RL-Align/RL-Kernel/pull/321) at CUDA commit `8576fa4bf449734ae99e9b50be8756bb282a8916`. H100 Triton replays use this PR's code at `e64abab904880b877d26d04c0cfad020b992aa51`. + +The same-H100 CUDA/Triton ratio is the hardware-matched comparison. CPU and MI300X columns provide absolute-latency context only; they are not hardware-normalized speed claims. + +| Comparison environment | Value | +|---|---| +| CUDA GPU | NVIDIA H100 80GB HBM3 (sm_90) | +| CUDA / PyTorch | 13.0 / 2.13.0+cu130 | +| CPU | Intel(R) Xeon(R) Platinum 8468, 96 intra-op threads | +| Transformers | 5.13.1 | + +### Single-GPU and CPU absolute latency + +| Shape / direction | CPU official (ms) | H100 official TP=1 (ms) | H100 Triton replay (ms) | H100 CUDA (ms) | CUDA / Triton H100 | MI300X official TP=1 (ms) | MI300X Triton (ms) | +|---|---:|---:|---:|---:|---:|---:|---:| +| M=1, forward | 12.9558 | 0.1193 | 1.7381 | 3.9988 | 2.30x | 0.1004 | 0.6053 | +| M=1, forward+backward | 58.5765 | 0.5184 | 4.2997 | 9.0704 | 2.11x | 0.4371 | 1.7138 | +| M=8, forward | 12.7676 | 0.1239 | 1.8293 | 3.9923 | 2.18x | 0.1077 | 0.6054 | +| M=8, forward+backward | 82.3669 | 0.5414 | 4.5028 | 9.4500 | 2.10x | 0.6476 | 2.8767 | +| M=32, forward | 8.5392 | 0.1311 | 2.2529 | 4.0277 | 1.79x | 0.1146 | 0.8761 | +| M=32, forward+backward | 65.9689 | 0.7842 | 5.4261 | 9.1635 | 1.69x | 0.4182 | 2.3188 | + +Both H100 columns used in the main distributed table are the user-supplied distributed timings. They are joined directly with MI300X rows of the same topology and direction; TP=1 values are excluded from all four columns. + +## Topology exactness versus Triton TP=1 + +All columns are element mismatch counts. This table does not compare against the official FFN. + +| Parallel layout | Forward output | Training output | dHidden | dWeights | Repeat | Train/infer | +|---|---:|---:|---:|---:|---:|---:| +| tp2 | 0 | 0 | 0 | 0 | 0 | 0 | +| tp2_sp | 0 | 0 | 0 | 0 | 0 | 0 | +| tp4 | 0 | 0 | 0 | 0 | 0 | 0 | +| tp2_cp2 | 0 | 0 | 0 | 0 | 0 | 0 | +| tp2_cp2_sp | 0 | 0 | 0 | 0 | 0 | 0 | +| tp8 | 0 | 0 | 0 | 0 | 0 | 0 | +| tp4_cp2 | 0 | 0 | 0 | 0 | 0 | 0 | +| tp4_cp2_sp | 0 | 0 | 0 | 0 | 0 | 0 | + +## Simple FP16 versus FP32 observation + +This is an official `Qwen3MLP` TP=1 output comparison only; it is not used to judge deterministic Triton and is not included in speed ratios. + +| Shape | Candidate | Reference | Max abs | Mean abs | Relative L2 | +|---|---|---|---:|---:|---:| +| (M,H,I)=(8,4096,12288) | FP16 | FP32 | 2.046e-06 | 3.742e-07 | 6.544e-04 | + +## Deterministic communication overlap + +The current timing includes the fixed-order communication schedule and makes no overlap claim. Forward SP all-gather must finish before gate/up projection, and TP reduction consumes the down-projection output, so those edges are hard dependencies. + +In backward, the gate and up contributions to dHidden are independent until their final ordered addition. A future implementation can place the fixed-rank reduction of one contribution on a second stream while computing the other, but it must preserve rank order, reduction tree, wait points, and gate-then-up addition order. Any optimization is accepted only if every TP=1 mismatch column remains zero. + +## Figures + +![Single-GPU CUDA, packed Triton, and CPU latency](single_gpu_overhead.png) + +![Topology mismatch versus Triton TP=1](collective_overhead.png) + +![Distributed H100 CUDA and MI300X packed Triton latency](distributed_ffn_overhead.png) diff --git a/benchmarks/results/pr325_rocm_mi300x/results.json b/benchmarks/results/pr325_rocm_mi300x/results.json new file mode 100644 index 00000000..575e685c --- /dev/null +++ b/benchmarks/results/pr325_rocm_mi300x/results.json @@ -0,0 +1,1036 @@ +{ + "communication_contract": { + "forward": { + "all_gather": 1, + "reduce_scatter": 1, + "total": 2 + }, + "train_fwd_bwd": { + "all_gather": 7, + "collective_invocations": 9, + "logical_total": 10, + "reduce_scatter_logical_lanes": 3 + } + }, + "distributed_ffn": [ + { + "cp_size": 1, + "direction": "forward", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 3.770714554916725, + "name": "tp2", + "official_distributed": { + "max_ms": 1.3578161597251892, + "median_ms": 0.23228488862514496, + "min_ms": 0.2129673957824707, + "p95_ms": 0.8294882718473673 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": false, + "tokens": 32, + "tp1_mismatch": { + "forward_output": 0 + }, + "tp_size": 2, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 1.9319094717502594, + "median_ms": 0.8758800104260445, + "min_ms": 0.8196169510483742, + "p95_ms": 1.4833177905529737 + }, + "weight_layout": "packed_forward_cache", + "world_size": 2 + }, + { + "cp_size": 1, + "direction": "train_fwd_bwd", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 2.978428538282444, + "name": "tp2", + "official_distributed": { + "max_ms": 1.2473920360207558, + "median_ms": 0.6984230130910873, + "min_ms": 0.6383182480931282, + "p95_ms": 1.2350523378700018 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": false, + "tokens": 32, + "tp1_mismatch": { + "hidden_gradient": 0, + "training_output": 0, + "weight_gradient": 0 + }, + "tp_size": 2, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 3.5044336691498756, + "median_ms": 2.0802030339837074, + "min_ms": 1.9138418138027191, + "p95_ms": 3.1109470874071126 + }, + "weight_layout": "packed_forward_cache", + "world_size": 2 + }, + { + "cp_size": 1, + "direction": "forward", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 2.8983204548458, + "name": "tp2_sp", + "official_distributed": { + "max_ms": 2.685968764126301, + "median_ms": 0.3065601922571659, + "min_ms": 0.28007570654153824, + "p95_ms": 0.8993070106953382 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": true, + "tokens": 32, + "tp1_mismatch": { + "forward_output": 0 + }, + "tp_size": 2, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 2.250421792268753, + "median_ms": 0.888509675860405, + "min_ms": 0.8365819230675697, + "p95_ms": 2.004028530791401 + }, + "weight_layout": "packed_forward_cache", + "world_size": 2 + }, + { + "cp_size": 1, + "direction": "train_fwd_bwd", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 1.771568250466244, + "name": "tp2_sp", + "official_distributed": { + "max_ms": 2.6939501985907555, + "median_ms": 1.4217207208275795, + "min_ms": 0.8287103846669197, + "p95_ms": 2.582016121596098 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": true, + "tokens": 32, + "tp1_mismatch": { + "hidden_gradient": 0, + "training_output": 0, + "weight_gradient": 0 + }, + "tp_size": 2, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 5.430372431874275, + "median_ms": 2.5186752900481224, + "min_ms": 1.9955337047576904, + "p95_ms": 3.6994490772485746 + }, + "weight_layout": "packed_forward_cache", + "world_size": 2 + }, + { + "cp_size": 1, + "direction": "forward", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 3.4979071801625334, + "name": "tp4", + "official_distributed": { + "max_ms": 0.8178036659955978, + "median_ms": 0.23796828463673592, + "min_ms": 0.21947640925645828, + "p95_ms": 0.7959454320371151 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": false, + "tokens": 32, + "tp1_mismatch": { + "forward_output": 0 + }, + "tp_size": 4, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 2.6544714346528053, + "median_ms": 0.8323909714818001, + "min_ms": 0.6374167278409004, + "p95_ms": 1.891128625720739 + }, + "weight_layout": "packed_forward_cache", + "world_size": 4 + }, + { + "cp_size": 1, + "direction": "train_fwd_bwd", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 2.7738258924780523, + "name": "tp4", + "official_distributed": { + "max_ms": 2.372283488512039, + "median_ms": 0.82804961130023, + "min_ms": 0.6606811657547951, + "p95_ms": 1.4657761901617057 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": false, + "tokens": 32, + "tp1_mismatch": { + "hidden_gradient": 0, + "training_output": 0, + "weight_gradient": 0 + }, + "tp_size": 4, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 3.301742486655712, + "median_ms": 2.296865452080965, + "min_ms": 1.6987714916467667, + "p95_ms": 2.88137081079185 + }, + "weight_layout": "packed_forward_cache", + "world_size": 4 + }, + { + "cp_size": 2, + "direction": "forward", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 3.264261952381431, + "name": "tp2_cp2", + "official_distributed": { + "max_ms": 0.8327467367053032, + "median_ms": 0.23166416212916374, + "min_ms": 0.21467916667461395, + "p95_ms": 0.652501266449689 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": false, + "tokens": 32, + "tp1_mismatch": { + "forward_output": 0 + }, + "tp_size": 2, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 2.562844194471836, + "median_ms": 0.7562125101685524, + "min_ms": 0.6980272009968758, + "p95_ms": 2.0815798547118893 + }, + "weight_layout": "packed_forward_cache", + "world_size": 4 + }, + { + "cp_size": 2, + "direction": "train_fwd_bwd", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 0.5639282606270454, + "name": "tp2_cp2", + "official_distributed": { + "max_ms": 5.045952275395393, + "median_ms": 4.483005963265896, + "min_ms": 4.005086608231068, + "p95_ms": 5.044202227145433 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": false, + "tokens": 32, + "tp1_mismatch": { + "hidden_gradient": 0, + "training_output": 0, + "weight_gradient": 0 + }, + "tp_size": 2, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 5.543430335819721, + "median_ms": 2.5280937552452087, + "min_ms": 1.921333372592926, + "p95_ms": 4.515600809827447 + }, + "weight_layout": "packed_forward_cache", + "world_size": 4 + }, + { + "cp_size": 2, + "direction": "forward", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 2.473328555193229, + "name": "tp2_cp2_sp", + "official_distributed": { + "max_ms": 0.9094811975955963, + "median_ms": 0.32615475356578827, + "min_ms": 0.2819793298840523, + "p95_ms": 0.8830895647406578 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": true, + "tokens": 32, + "tp1_mismatch": { + "forward_output": 0 + }, + "tp_size": 2, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 2.9444824904203415, + "median_ms": 0.8066878654062748, + "min_ms": 0.7346402853727341, + "p95_ms": 1.7819570843130346 + }, + "weight_layout": "packed_forward_cache", + "world_size": 4 + }, + { + "cp_size": 2, + "direction": "train_fwd_bwd", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 0.5550966504336481, + "name": "tp2_cp2_sp", + "official_distributed": { + "max_ms": 6.967155262827873, + "median_ms": 4.586996044963598, + "min_ms": 4.187238402664661, + "p95_ms": 6.129474937915803 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": true, + "tokens": 32, + "tp1_mismatch": { + "hidden_gradient": 0, + "training_output": 0, + "weight_gradient": 0 + }, + "tp_size": 2, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 4.86080814152956, + "median_ms": 2.546226140111685, + "min_ms": 1.9829655066132545, + "p95_ms": 3.2532437238842262 + }, + "weight_layout": "packed_forward_cache", + "world_size": 4 + }, + { + "cp_size": 1, + "direction": "forward", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 3.4050468145461728, + "name": "tp8", + "official_distributed": { + "max_ms": 1.0286271572113037, + "median_ms": 0.20749308168888092, + "min_ms": 0.1928461715579033, + "p95_ms": 0.5358625203371044 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": false, + "tokens": 32, + "tp1_mismatch": { + "forward_output": 0 + }, + "tp_size": 8, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 1.8369676545262337, + "median_ms": 0.7065236568450928, + "min_ms": 0.6407918408513069, + "p95_ms": 1.6900700516998768 + }, + "weight_layout": "packed_forward_cache", + "world_size": 8 + }, + { + "cp_size": 1, + "direction": "train_fwd_bwd", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 2.3611604677472107, + "name": "tp8", + "official_distributed": { + "max_ms": 2.177082933485508, + "median_ms": 0.8685095235705376, + "min_ms": 0.6101867184042931, + "p95_ms": 1.8530789297074082 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": false, + "tokens": 32, + "tp1_mismatch": { + "hidden_gradient": 0, + "training_output": 0, + "weight_gradient": 0 + }, + "tp_size": 8, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 5.000975914299488, + "median_ms": 2.0506903529167175, + "min_ms": 1.6332557424902916, + "p95_ms": 4.608689062297344 + }, + "weight_layout": "packed_forward_cache", + "world_size": 8 + }, + { + "cp_size": 2, + "direction": "forward", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 3.4271073503642904, + "name": "tp4_cp2", + "official_distributed": { + "max_ms": 1.757740043103695, + "median_ms": 0.2423599362373352, + "min_ms": 0.2230815589427948, + "p95_ms": 0.8416332770138979 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": false, + "tokens": 32, + "tp1_mismatch": { + "forward_output": 0 + }, + "tp_size": 4, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 2.3825177922844887, + "median_ms": 0.8305935189127922, + "min_ms": 0.6722584366798401, + "p95_ms": 1.6696299426257601 + }, + "weight_layout": "packed_forward_cache", + "world_size": 8 + }, + { + "cp_size": 2, + "direction": "train_fwd_bwd", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 0.8187486604980407, + "name": "tp4_cp2", + "official_distributed": { + "max_ms": 5.41381910443306, + "median_ms": 3.0070655047893524, + "min_ms": 2.4463413283228874, + "p95_ms": 4.813990509137511 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": false, + "tokens": 32, + "tp1_mismatch": { + "hidden_gradient": 0, + "training_output": 0, + "weight_gradient": 0 + }, + "tp_size": 4, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 4.88690659403801, + "median_ms": 2.462030854076147, + "min_ms": 1.8442394211888313, + "p95_ms": 4.497065208852291 + }, + "weight_layout": "packed_forward_cache", + "world_size": 8 + }, + { + "cp_size": 2, + "direction": "forward", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 2.416452985140764, + "name": "tp4_cp2_sp", + "official_distributed": { + "max_ms": 0.9128358215093613, + "median_ms": 0.31526200473308563, + "min_ms": 0.28138794004917145, + "p95_ms": 0.8974010124802589 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": true, + "tokens": 32, + "tp1_mismatch": { + "forward_output": 0 + }, + "tp_size": 4, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 2.155771479010582, + "median_ms": 0.7618158124387264, + "min_ms": 0.6957026198506355, + "p95_ms": 1.7150717787444583 + }, + "weight_layout": "packed_forward_cache", + "world_size": 8 + }, + { + "cp_size": 2, + "direction": "train_fwd_bwd", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_triton_vs_official_distributed": 0.8257421849400354, + "name": "tp4_cp2_sp", + "official_distributed": { + "max_ms": 4.409045912325382, + "median_ms": 2.9345820657908916, + "min_ms": 2.5761546567082405, + "p95_ms": 4.245032416656614 + }, + "repeat_mismatch_count": 0, + "sequence_parallel": true, + "tokens": 32, + "tp1_mismatch": { + "hidden_gradient": 0, + "training_output": 0, + "weight_gradient": 0 + }, + "tp_size": 4, + "train_infer_mismatch_count": 0, + "triton": { + "max_ms": 4.953155294060707, + "median_ms": 2.4232082068920135, + "min_ms": 1.9790586084127426, + "p95_ms": 3.922184882685543 + }, + "weight_layout": "packed_forward_cache", + "world_size": 8 + } + ], + "distributed_platform_comparison": { + "contract": "same M=32 workload, direction, and TP/CP/SP topology", + "rows": [ + { + "cp_size": 1, + "deterministic_mi300x_over_h100_ratio": 0.3139805027337412, + "direction": "forward", + "h100_deterministic_cuda_ms": 2.7896, + "h100_deterministic_over_official_ratio": 17.97422680412371, + "h100_official_distributed_ms": 0.1552, + "mi300x_deterministic_over_official_ratio": 3.770714554916725, + "mi300x_deterministic_triton_ms": 0.8758800104260445, + "mi300x_official_distributed_ms": 0.23228488862514496, + "name": "tp2", + "sequence_parallel": false, + "tp_size": 2 + }, + { + "cp_size": 1, + "deterministic_mi300x_over_h100_ratio": 0.26626598834991455, + "direction": "train_fwd_bwd", + "h100_deterministic_cuda_ms": 7.8125, + "h100_deterministic_over_official_ratio": 12.398825583240756, + "h100_official_distributed_ms": 0.6301, + "mi300x_deterministic_over_official_ratio": 2.978428538282444, + "mi300x_deterministic_triton_ms": 2.0802030339837074, + "mi300x_official_distributed_ms": 0.6984230130910873, + "name": "tp2", + "sequence_parallel": false, + "tp_size": 2 + }, + { + "cp_size": 1, + "deterministic_mi300x_over_h100_ratio": 0.26673962049246625, + "direction": "forward", + "h100_deterministic_cuda_ms": 3.331, + "h100_deterministic_over_official_ratio": 21.46262886597938, + "h100_official_distributed_ms": 0.1552, + "mi300x_deterministic_over_official_ratio": 2.8983204548458, + "mi300x_deterministic_triton_ms": 0.888509675860405, + "mi300x_official_distributed_ms": 0.3065601922571659, + "name": "tp2_sp", + "sequence_parallel": true, + "tp_size": 2 + }, + { + "cp_size": 1, + "deterministic_mi300x_over_h100_ratio": 0.21400019457480116, + "direction": "train_fwd_bwd", + "h100_deterministic_cuda_ms": 11.7695, + "h100_deterministic_over_official_ratio": 18.678781145849868, + "h100_official_distributed_ms": 0.6301, + "mi300x_deterministic_over_official_ratio": 1.771568250466244, + "mi300x_deterministic_triton_ms": 2.5186752900481224, + "mi300x_official_distributed_ms": 1.4217207208275795, + "name": "tp2_sp", + "sequence_parallel": true, + "tp_size": 2 + }, + { + "cp_size": 1, + "deterministic_mi300x_over_h100_ratio": 0.38382024783593865, + "direction": "forward", + "h100_deterministic_cuda_ms": 2.1687, + "h100_deterministic_over_official_ratio": 14.174509803921568, + "h100_official_distributed_ms": 0.153, + "mi300x_deterministic_over_official_ratio": 3.4979071801625334, + "mi300x_deterministic_triton_ms": 0.8323909714818001, + "mi300x_official_distributed_ms": 0.23796828463673592, + "name": "tp4", + "sequence_parallel": false, + "tp_size": 4 + }, + { + "cp_size": 1, + "deterministic_mi300x_over_h100_ratio": 0.2444670213169173, + "direction": "train_fwd_bwd", + "h100_deterministic_cuda_ms": 9.3954, + "h100_deterministic_over_official_ratio": 16.44277213860693, + "h100_official_distributed_ms": 0.5714, + "mi300x_deterministic_over_official_ratio": 2.7738258924780523, + "mi300x_deterministic_triton_ms": 2.296865452080965, + "mi300x_official_distributed_ms": 0.82804961130023, + "name": "tp4", + "sequence_parallel": false, + "tp_size": 4 + }, + { + "cp_size": 2, + "deterministic_mi300x_over_h100_ratio": 0.26758165322124217, + "direction": "forward", + "h100_deterministic_cuda_ms": 2.8261, + "h100_deterministic_over_official_ratio": 18.47124183006536, + "h100_official_distributed_ms": 0.153, + "mi300x_deterministic_over_official_ratio": 3.264261952381431, + "mi300x_deterministic_triton_ms": 0.7562125101685524, + "mi300x_official_distributed_ms": 0.23166416212916374, + "name": "tp2_cp2", + "sequence_parallel": false, + "tp_size": 2 + }, + { + "cp_size": 2, + "deterministic_mi300x_over_h100_ratio": 0.15632342880036165, + "direction": "train_fwd_bwd", + "h100_deterministic_cuda_ms": 16.1722, + "h100_deterministic_over_official_ratio": 28.302765138256913, + "h100_official_distributed_ms": 0.5714, + "mi300x_deterministic_over_official_ratio": 0.5639282606270454, + "mi300x_deterministic_triton_ms": 2.5280937552452087, + "mi300x_official_distributed_ms": 4.483005963265896, + "name": "tp2_cp2", + "sequence_parallel": false, + "tp_size": 2 + }, + { + "cp_size": 2, + "deterministic_mi300x_over_h100_ratio": 0.22524372184237304, + "direction": "forward", + "h100_deterministic_cuda_ms": 3.5814, + "h100_deterministic_over_official_ratio": 23.4078431372549, + "h100_official_distributed_ms": 0.153, + "mi300x_deterministic_over_official_ratio": 2.473328555193229, + "mi300x_deterministic_triton_ms": 0.8066878654062748, + "mi300x_official_distributed_ms": 0.32615475356578827, + "name": "tp2_cp2_sp", + "sequence_parallel": true, + "tp_size": 2 + }, + { + "cp_size": 2, + "deterministic_mi300x_over_h100_ratio": 0.15165676796738922, + "direction": "train_fwd_bwd", + "h100_deterministic_cuda_ms": 16.7894, + "h100_deterministic_over_official_ratio": 29.3829191459573, + "h100_official_distributed_ms": 0.5714, + "mi300x_deterministic_over_official_ratio": 0.5550966504336481, + "mi300x_deterministic_triton_ms": 2.546226140111685, + "mi300x_official_distributed_ms": 4.586996044963598, + "name": "tp2_cp2_sp", + "sequence_parallel": true, + "tp_size": 2 + }, + { + "cp_size": 1, + "deterministic_mi300x_over_h100_ratio": 0.3044573200228789, + "direction": "forward", + "h100_deterministic_cuda_ms": 2.3206, + "h100_deterministic_over_official_ratio": 15.098243331164607, + "h100_official_distributed_ms": 0.1537, + "mi300x_deterministic_over_official_ratio": 3.4050468145461728, + "mi300x_deterministic_triton_ms": 0.7065236568450928, + "mi300x_official_distributed_ms": 0.20749308168888092, + "name": "tp8", + "sequence_parallel": false, + "tp_size": 8 + }, + { + "cp_size": 1, + "deterministic_mi300x_over_h100_ratio": 0.20228758105220396, + "direction": "train_fwd_bwd", + "h100_deterministic_cuda_ms": 10.1375, + "h100_deterministic_over_official_ratio": 20.25069916100679, + "h100_official_distributed_ms": 0.5006, + "mi300x_deterministic_over_official_ratio": 2.3611604677472107, + "mi300x_deterministic_triton_ms": 2.0506903529167175, + "mi300x_official_distributed_ms": 0.8685095235705376, + "name": "tp8", + "sequence_parallel": false, + "tp_size": 8 + }, + { + "cp_size": 2, + "deterministic_mi300x_over_h100_ratio": 0.3751382136817633, + "direction": "forward", + "h100_deterministic_cuda_ms": 2.2141, + "h100_deterministic_over_official_ratio": 14.4053350683149, + "h100_official_distributed_ms": 0.1537, + "mi300x_deterministic_over_official_ratio": 3.4271073503642904, + "mi300x_deterministic_triton_ms": 0.8305935189127922, + "mi300x_official_distributed_ms": 0.2423599362373352, + "name": "tp4_cp2", + "sequence_parallel": false, + "tp_size": 4 + }, + { + "cp_size": 2, + "deterministic_mi300x_over_h100_ratio": 0.14865809995810497, + "direction": "train_fwd_bwd", + "h100_deterministic_cuda_ms": 16.5617, + "h100_deterministic_over_official_ratio": 33.083699560527364, + "h100_official_distributed_ms": 0.5006, + "mi300x_deterministic_over_official_ratio": 0.8187486604980407, + "mi300x_deterministic_triton_ms": 2.462030854076147, + "mi300x_official_distributed_ms": 3.0070655047893524, + "name": "tp4_cp2", + "sequence_parallel": false, + "tp_size": 4 + }, + { + "cp_size": 2, + "deterministic_mi300x_over_h100_ratio": 0.23104931834245007, + "direction": "forward", + "h100_deterministic_cuda_ms": 3.2972, + "h100_deterministic_over_official_ratio": 21.45217957059206, + "h100_official_distributed_ms": 0.1537, + "mi300x_deterministic_over_official_ratio": 2.416452985140764, + "mi300x_deterministic_triton_ms": 0.7618158124387264, + "mi300x_official_distributed_ms": 0.31526200473308563, + "name": "tp4_cp2_sp", + "sequence_parallel": true, + "tp_size": 4 + }, + { + "cp_size": 2, + "deterministic_mi300x_over_h100_ratio": 0.1413303747815495, + "direction": "train_fwd_bwd", + "h100_deterministic_cuda_ms": 17.1457, + "h100_deterministic_over_official_ratio": 34.250299640431486, + "h100_official_distributed_ms": 0.5006, + "mi300x_deterministic_over_official_ratio": 0.8257421849400354, + "mi300x_deterministic_triton_ms": 2.4232082068920135, + "mi300x_official_distributed_ms": 2.9345820657908916, + "name": "tp4_cp2_sp", + "sequence_parallel": true, + "tp_size": 4 + } + ], + "source": "cuda_cpu_comparison.json" + }, + "environment": { + "NCCL_IB_DISABLE": "1", + "architecture": "gfx942:sramecc+:xnack-", + "deterministic_compute": "ROCm-native Triton", + "deterministic_transport": "fixed-tree HIP IPC with RCCL fallback on ROCm", + "distributed_speed_comparison": "four same-topology H100/MI300X paths", + "git_commit": "caef501101a3906c733076f31f3b5a9870169d16", + "gpu": "AMD Instinct MI300X", + "gpu_count": 8, + "hip": "7.14.60850", + "python": "3.12.3", + "single_gpu_speed_context": "Hugging Face Transformers Qwen3MLP, TP=1", + "torch": "2.12.0+rocm7.14.0a20260608", + "transformers": "5.10.4" + }, + "methodology": { + "distributed_timing": "synchronized wall clock, slowest rank/sample", + "distributed_worker_cpu_affinity": "one NUMA-local CPU per GPU rank", + "operator_only": true, + "samples": 50, + "single_gpu_timing": "GPU events, median and p95", + "tp1_forward_cache_bytes": 301989888, + "training_samples": 20, + "triton_weight_layout": "packed_forward_cache_outside_timed_region", + "warmup": 10 + }, + "previous_deterministic_comparison": { + "rows": [ + { + "current_ms": 0.8758800104260445, + "direction": "forward", + "latency_reduction_ratio": 0.031099087163262373, + "name": "tp2", + "previous_ms": 0.9039933793246746 + }, + { + "current_ms": 2.0802030339837074, + "direction": "train_fwd_bwd", + "latency_reduction_ratio": 0.22945943339972408, + "name": "tp2", + "previous_ms": 2.6996671222150326 + }, + { + "current_ms": 0.888509675860405, + "direction": "forward", + "latency_reduction_ratio": 0.158719653904269, + "name": "tp2_sp", + "previous_ms": 1.0561398230493069 + }, + { + "current_ms": 2.5186752900481224, + "direction": "train_fwd_bwd", + "latency_reduction_ratio": 0.1504226199863158, + "name": "tp2_sp", + "previous_ms": 2.9646214097738266 + }, + { + "current_ms": 0.8323909714818001, + "direction": "forward", + "latency_reduction_ratio": 0.00417650162140959, + "name": "tp4", + "previous_ms": 0.8358820341527462 + }, + { + "current_ms": 2.296865452080965, + "direction": "train_fwd_bwd", + "latency_reduction_ratio": 0.1492708192299168, + "name": "tp4", + "previous_ms": 2.6998785324394703 + }, + { + "current_ms": 0.7562125101685524, + "direction": "forward", + "latency_reduction_ratio": 0.16124577124706252, + "name": "tp2_cp2", + "previous_ms": 0.9015901014208794 + }, + { + "current_ms": 2.5280937552452087, + "direction": "train_fwd_bwd", + "latency_reduction_ratio": 0.28997936652606293, + "name": "tp2_cp2", + "previous_ms": 3.5605919547379017 + }, + { + "current_ms": 0.8066878654062748, + "direction": "forward", + "latency_reduction_ratio": 0.28586663566666093, + "name": "tp2_cp2_sp", + "previous_ms": 1.1296039447188377 + }, + { + "current_ms": 2.546226140111685, + "direction": "train_fwd_bwd", + "latency_reduction_ratio": 0.36230506379306326, + "name": "tp2_cp2_sp", + "previous_ms": 3.992859274148941 + }, + { + "current_ms": 0.7065236568450928, + "direction": "forward", + "latency_reduction_ratio": 0.34939832397754955, + "name": "tp8", + "previous_ms": 1.0859542526304722 + }, + { + "current_ms": 2.0506903529167175, + "direction": "train_fwd_bwd", + "latency_reduction_ratio": 0.19957512534856403, + "name": "tp8", + "previous_ms": 2.5620022788643837 + }, + { + "current_ms": 0.8305935189127922, + "direction": "forward", + "latency_reduction_ratio": 0.21168450338208877, + "name": "tp4_cp2", + "previous_ms": 1.0536308400332928 + }, + { + "current_ms": 2.462030854076147, + "direction": "train_fwd_bwd", + "latency_reduction_ratio": 0.24475513187788767, + "name": "tp4_cp2", + "previous_ms": 3.259910736232996 + }, + { + "current_ms": 0.7618158124387264, + "direction": "forward", + "latency_reduction_ratio": 0.3613095756305654, + "name": "tp4_cp2_sp", + "previous_ms": 1.1927778832614422 + }, + { + "current_ms": 2.4232082068920135, + "direction": "train_fwd_bwd", + "latency_reduction_ratio": 0.3537747388734832, + "name": "tp4_cp2_sp", + "previous_ms": 3.7497887387871742 + } + ], + "source": "previous checked MI300X benchmark before PR #357" + }, + "single_gpu": { + "dtype_accuracy": [ + { + "candidate_dtype": "float16", + "exact_fraction": 3.0517578125e-05, + "hidden": 4096, + "intermediate": 12288, + "max_abs": 2.046101144514978e-06, + "mean_abs": 3.7421963838824013e-07, + "name": "Official Qwen3MLP TP=1 FP16 vs FP32", + "reference_dtype": "float32", + "relative_l2": 0.0006544404313899577, + "tokens": 8 + } + ], + "speed": [ + { + "direction": "forward", + "dtype": "bfloat16", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_vs_official_tp1": 6.029729218534371, + "name": "(M,H,I)=(1,4096,12288), forward", + "official_tp1": { + "max_ms": 0.24884900450706482, + "median_ms": 0.10038900002837181, + "min_ms": 0.09506099671125412, + "p95_ms": 0.15199404805898656 + }, + "tokens": 1, + "triton": { + "max_ms": 0.7080109715461731, + "median_ms": 0.6053184866905212, + "min_ms": 0.5777779817581177, + "p95_ms": 0.6747872471809386 + }, + "weight_layout": "packed_forward_cache" + }, + { + "direction": "train_fwd_bwd", + "dtype": "bfloat16", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_vs_official_tp1": 3.9212252358156867, + "name": "(M,H,I)=(1,4096,12288), forward+backward", + "official_tp1": { + "max_ms": 0.5410429835319519, + "median_ms": 0.43706849217414856, + "min_ms": 0.42226698994636536, + "p95_ms": 0.5354484349489211 + }, + "tokens": 1, + "triton": { + "max_ms": 1.910356044769287, + "median_ms": 1.7138440012931824, + "min_ms": 1.6589829921722412, + "p95_ms": 1.8108766913414003 + }, + "weight_layout": "packed_forward_cache" + }, + { + "direction": "forward", + "dtype": "bfloat16", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_vs_official_tp1": 5.61974558812295, + "name": "(M,H,I)=(8,4096,12288), forward", + "official_tp1": { + "max_ms": 0.14878100156784058, + "median_ms": 0.10771999880671501, + "min_ms": 0.10279300063848495, + "p95_ms": 0.11455849818885325 + }, + "tokens": 8, + "triton": { + "max_ms": 0.6983559727668762, + "median_ms": 0.6053589880466461, + "min_ms": 0.5815439820289612, + "p95_ms": 0.6279941111803056 + }, + "weight_layout": "packed_forward_cache" + }, + { + "direction": "train_fwd_bwd", + "dtype": "bfloat16", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_vs_official_tp1": 4.442236747459528, + "name": "(M,H,I)=(8,4096,12288), forward+backward", + "official_tp1": { + "max_ms": 0.7275599837303162, + "median_ms": 0.6475815176963806, + "min_ms": 0.5028650164604187, + "p95_ms": 0.672568279504776 + }, + "tokens": 8, + "triton": { + "max_ms": 3.3488519191741943, + "median_ms": 2.8767104148864746, + "min_ms": 2.6616709232330322, + "p95_ms": 3.018102836608887 + }, + "weight_layout": "packed_forward_cache" + }, + { + "direction": "forward", + "dtype": "bfloat16", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_vs_official_tp1": 7.6426809923032115, + "name": "(M,H,I)=(32,4096,12288), forward", + "official_tp1": { + "max_ms": 0.16067799925804138, + "median_ms": 0.11462999880313873, + "min_ms": 0.11064399778842926, + "p95_ms": 0.11765709854662418 + }, + "tokens": 32, + "triton": { + "max_ms": 0.9773309826850891, + "median_ms": 0.8760805130004883, + "min_ms": 0.8697310090065002, + "p95_ms": 0.8829010009765625 + }, + "weight_layout": "packed_forward_cache" + }, + { + "direction": "train_fwd_bwd", + "dtype": "bfloat16", + "hidden": 4096, + "intermediate": 12288, + "latency_ratio_vs_official_tp1": 5.544395107362785, + "name": "(M,H,I)=(32,4096,12288), forward+backward", + "official_tp1": { + "max_ms": 0.48680299520492554, + "median_ms": 0.41822099685668945, + "min_ms": 0.409527987241745, + "p95_ms": 0.4841772079467773 + }, + "tokens": 32, + "triton": { + "max_ms": 2.744915008544922, + "median_ms": 2.3187824487686157, + "min_ms": 2.2827088832855225, + "p95_ms": 2.4811455845832824 + }, + "weight_layout": "packed_forward_cache" + } + ] + } +} diff --git a/benchmarks/results/pr325_rocm_mi300x/single_gpu_overhead.png b/benchmarks/results/pr325_rocm_mi300x/single_gpu_overhead.png new file mode 100644 index 00000000..c0bb0d4e Binary files /dev/null and b/benchmarks/results/pr325_rocm_mi300x/single_gpu_overhead.png differ diff --git a/benchmarks/results/ws2_cpu/report.md b/benchmarks/results/ws2_cpu/report.md index e5c5b510..3f74f014 100644 --- a/benchmarks/results/ws2_cpu/report.md +++ b/benchmarks/results/ws2_cpu/report.md @@ -147,4 +147,3 @@ A head shard computed under TP=N versus the same slice of an unsharded run. TP p ![Bitwise exactness matrix](exactness_matrix.png) ![TP-degree invariance](tp_degree_invariance.png) - diff --git a/benchmarks/results/ws2_rocm_mi300x/report.md b/benchmarks/results/ws2_rocm_mi300x/report.md index 36ce7166..1b0256df 100644 --- a/benchmarks/results/ws2_rocm_mi300x/report.md +++ b/benchmarks/results/ws2_rocm_mi300x/report.md @@ -280,4 +280,3 @@ Schedule: all-gather Q/K/V and the position ids over the CP group, run the stric ![TP-degree invariance](tp_degree_invariance.png) ![Distributed CP latency](distributed_cp_latency.png) - diff --git a/csrc/cuda/distributed/deterministic_collective.cu b/csrc/cuda/distributed/deterministic_collective.cu index 72d3f874..e7d836a3 100644 --- a/csrc/cuda/distributed/deterministic_collective.cu +++ b/csrc/cuda/distributed/deterministic_collective.cu @@ -949,7 +949,10 @@ class DeterministicCollectiveState { has_staged_input_ = true; } - void all_reduce(torch::Tensor& output, cudaStream_t stream) { + void all_reduce( + torch::Tensor& output, + cudaStream_t stream, + bool allow_owner_path = true) { check_tensor(output, "output"); TORCH_CHECK(has_staged_input_, "stage() must be called before all_reduce()"); TORCH_CHECK( @@ -960,7 +963,10 @@ class DeterministicCollectiveState { "all-reduce output size must match the staged input size"); const int64_t element_count = output.numel(); - if (staged_owner_path_) { + // Graph-replayed calls must avoid the owner-push branch because it writes + // to remote IPC frames. Callers that use the fused ABI pass false here; + // direct staged collectives retain the eager owner optimization. + if (allow_owner_path && staged_owner_path_) { if (rank_ == 0) { // Logical rank 0 is the topology-favorable reader in both traced TP // engines. It evaluates the original fixed tree exactly once, then @@ -1022,6 +1028,58 @@ class DeterministicCollectiveState { return; } + // Small collectives are launch-bound. The fast kernel folds the peer + // stage wait, fixed-tree reduction, and completion publication into one + // launch while preserving the exact reduction order. + if (staged_fast_path_ && staged_bytes_ <= kSingleBlockFastPathMaxBytes) { + if (element_count > 0) { + switch (output.scalar_type()) { + case at::ScalarType::Float: + launch_all_reduce_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + element_count, + world_size_, + stream); + break; + case at::ScalarType::Half: + launch_all_reduce_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + element_count, + world_size_, + stream); + break; +#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) + case at::ScalarType::BFloat16: + launch_all_reduce_fast( + peers_, + local_stage_sequence_, + local_done_sequence_, + static_cast(output.data_ptr()), + element_count, + world_size_, + stream); + break; +#endif + default: + TORCH_CHECK( + false, + "deterministic all-reduce supports float32, float16, and bfloat16; got ", + output.scalar_type()); + } + AT_CUDA_CHECK(cudaGetLastError()); + } else { + publish_done(stream); + } + has_staged_input_ = false; + return; + } + wait_for_staged_peers(stream); if (element_count == 0) { publish_done(stream); @@ -1090,10 +1148,11 @@ class DeterministicCollectiveState { output.numel() == input.numel(), "all-reduce output size must match the input size"); - // Route both ABI entry points through the same staged protocol so the - // topology-aware owner reduction is always applied. + // Use the graph-safe staged protocol for every message size. The fused + // two-slot protocol is intentionally kept available in the extension for + // experiments, but is not safe to replay across vLLM's many graph shapes. stage(input, stream); - all_reduce(output, stream); + all_reduce(output, stream, /*allow_owner_path=*/false); } void all_gather_fused( diff --git a/docs/benchmarking/README.md b/docs/benchmarking/README.md index 99da15c0..849a82d0 100644 --- a/docs/benchmarking/README.md +++ b/docs/benchmarking/README.md @@ -14,9 +14,30 @@ python benchmarks/profiler.py --format json --output reports/profile.json python benchmarks/benchmark_sampling.py python benchmarks/benchmark_grpo_op.py python benchmarks/benchmark_pack.py --smoke +python benchmarks/benchmark_rocm_ffn.py --help python scripts/run_perf.py ``` +## Unified ROCm deterministic FFN benchmark + +The RCCL/FFN benchmark has one entry point. It measures the deterministic +Triton FFN across TP, CP, and sequence-parallel layouts, checks bitwise +exactness against deterministic TP=1, and compares four distributed paths at +the same topology: H100 official/deterministic and MI300X +official/deterministic. TP=1 latency is not part of the distributed figure. + +```bash +HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 NCCL_IB_DISABLE=1 PYTHONPATH=. \ + python benchmarks/benchmark_rocm_ffn.py \ + --warmup 5 --samples 20 --training-samples 10 \ + --output-dir benchmarks/results/pr325_rocm_mi300x +``` + +The checked-in report is +`benchmarks/results/pr325_rocm_mi300x/report.md`, with the machine-readable +data in `results.json` and the updated topology figure in +`distributed_ffn_overhead.png`. + The automated profiler records one row per workload shape with: - `tokens_per_sec`: active tokens divided by median latency. diff --git a/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh b/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh index 4b3a601c..eb1cf61b 100755 --- a/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh +++ b/examples/vime_qwen3_8b_tp2_cp2/aligned_python_entrypoint.sh @@ -12,14 +12,37 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then # vLLM hooks; otherwise it silently falls back to native attention/FFN and # loses the R/R performance path. Preserve explicit ablation selections. strict_linear_logp=0 + rollout_batch_size="" + n_samples_per_prompt="1" + explicit_vllm_execution_config=0 previous_arg="" for current_arg in "$@"; do if [[ "${previous_arg}" == "--linear-logp-provider-mode" && "${current_arg}" == "strict" ]]; then strict_linear_logp=1 - break + elif [[ "${previous_arg}" == "--rollout-batch-size" ]]; then + rollout_batch_size="${current_arg}" + elif [[ "${previous_arg}" == "--n-samples-per-prompt" ]]; then + n_samples_per_prompt="${current_arg}" fi + + case "${current_arg}" in + --linear-logp-provider-mode=strict) + strict_linear_logp=1 + ;; + --rollout-batch-size=*) + rollout_batch_size="${current_arg#*=}" + ;; + --n-samples-per-prompt=*) + n_samples_per_prompt="${current_arg#*=}" + ;; + --vllm-enforce-eager|--vllm-optimization-level|--vllm-optimization-level=*|--vllm-compilation-config|--vllm-compilation-config=*) + explicit_vllm_execution_config=1 + ;; + esac previous_arg="${current_arg}" done + + strict_cudagraph_args=() if [[ "${strict_linear_logp}" == "1" ]]; then export RL_KERNEL_VLLM_INTEGRATION="${RL_KERNEL_VLLM_INTEGRATION:-1}" export RL_KERNEL_CUDA_ONLY="${RL_KERNEL_CUDA_ONLY:-1}" @@ -27,6 +50,42 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then export RL_KERNEL_ATTENTION_CASE="${RL_KERNEL_ATTENTION_CASE:-R/R}" export RL_KERNEL_FFN_CASE="${RL_KERNEL_FFN_CASE:-R/R}" export RL_KERNEL_LOGP_CASE="${RL_KERNEL_LOGP_CASE:-R/R}" + + # Strict rollout kernels preserve their arithmetic order under CUDA Graph. + # Capturing the complete decode graph removes the per-layer host-launch + # gaps that otherwise dominate small decode batches. Capture every exact + # batch size: padding a strict custom kernel to a larger sparse graph can + # access invalid slots and, more importantly, changes the tested contract. + # Explicit vLLM execution flags always win so callers can opt out. + if [[ "${explicit_vllm_execution_config}" == "0" ]]; then + if [[ "${rollout_batch_size}" =~ ^[1-9][0-9]*$ && "${n_samples_per_prompt}" =~ ^[1-9][0-9]*$ ]]; then + max_capture_size=$((rollout_batch_size * n_samples_per_prompt)) + if [[ -n "${RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE:-}" ]]; then + max_capture_size="${RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE}" + fi + if ! [[ "${max_capture_size}" =~ ^[1-9][0-9]*$ ]]; then + echo "RL_KERNEL_VLLM_CUDAGRAPH_MAX_CAPTURE_SIZE must be a positive integer" >&2 + exit 2 + fi + + capture_sizes="[" + for ((batch_size = 1; batch_size <= max_capture_size; batch_size++)); do + if ((batch_size > 1)); then + capture_sizes+="," + fi + capture_sizes+="${batch_size}" + done + capture_sizes+="]" + compilation_config="{\"cudagraph_mode\":\"FULL_DECODE_ONLY\",\"cudagraph_capture_sizes\":${capture_sizes},\"max_cudagraph_capture_size\":${max_capture_size}}" + strict_cudagraph_args=( + --vllm-optimization-level 0 + --vllm-compilation-config "${compilation_config}" + ) + echo "[RL-Kernel] strict vLLM full-decode CUDA Graph capture sizes: ${capture_sizes}" >&2 + else + echo "[RL-Kernel] strict CUDA Graph disabled: rollout batch size is unavailable" >&2 + fi + fi fi exec "${REAL_PYTHON}" "$@" \ --seed 1234 \ @@ -35,7 +94,8 @@ if [[ "${1:-}" == "train.py" || "${1:-}" == */train.py ]]; then --vllm-attention-backend flash_attn \ --vllm-disable-custom-all-reduce \ --deterministic-mode \ - --accumulate-allreduce-grads-in-fp32 + --accumulate-allreduce-grads-in-fp32 \ + "${strict_cudagraph_args[@]}" fi exec "${REAL_PYTHON}" "$@" diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index b599ad76..673828e4 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -615,16 +615,18 @@ def all_gather( def all_gather_many( self, - inputs: tuple[torch.Tensor, ...], + inputs: tuple[torch.Tensor, ...] | list[torch.Tensor], *, validate_signature: bool = True, ) -> tuple[torch.Tensor, ...]: - """Gather several tensors through the single-tensor transport ABI.""" + """Gather several tensors through the platform transport.""" - if not inputs: + values = tuple(inputs) + if not values: raise ValueError("all_gather_many requires at least one input") return tuple( - self.all_gather(input, validate_signature=validate_signature) for input in inputs + self.all_gather(value, validate_signature=validate_signature) + for value in values ) def reduce_scatter( @@ -1262,7 +1264,7 @@ def collective_for_group( minimum_capacity_bytes: int = _DEFAULT_MAX_SIZE_BYTES, device: torch.device | str | int | None = None, ) -> Any | None: - """Return the process-local RL-Kernel collective shared by hot-path ops.""" + """Return the process-local platform collective shared by hot-path ops.""" if group is None: return None @@ -1271,8 +1273,8 @@ def collective_for_group( if minimum_capacity_bytes <= 0: raise ValueError("minimum_capacity_bytes must be positive") - rank = dist.get_rank(group=group) - world_size = dist.get_world_size(group=group) + rank = int(dist.get_rank(group=group)) + world_size = int(dist.get_world_size(group=group)) if device is None: device_index = torch.cuda.current_device() else: diff --git a/rl_engine/distributed/transport_collectives.py b/rl_engine/distributed/transport_collectives.py new file mode 100644 index 00000000..b9db44ad --- /dev/null +++ b/rl_engine/distributed/transport_collectives.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Compatibility imports for the unified deterministic collectives. + +The implementation now lives in :mod:`rl_engine.distributed.collectives`. +This module remains as a stable import path for older benchmark and integration +callers; it contains no separate collective implementation. +""" + +from rl_engine.distributed.collectives import ( + RCCLDeterministicCollective, + TorchDistributedDeterministicCollective, + create_deterministic_collective, +) + +__all__ = [ + "RCCLDeterministicCollective", + "TorchDistributedDeterministicCollective", + "create_deterministic_collective", +] diff --git a/rl_engine/integrations/megatron_runtime.py b/rl_engine/integrations/megatron_runtime.py index e296e824..72ac8919 100644 --- a/rl_engine/integrations/megatron_runtime.py +++ b/rl_engine/integrations/megatron_runtime.py @@ -6,6 +6,7 @@ from __future__ import annotations import importlib +import os from collections.abc import Callable, Iterable from types import MethodType from typing import Any @@ -290,7 +291,48 @@ def initialize_from_environment(_args: Any = None) -> MegatronIntegration: from rl_engine.integrations.ablation import integration_plan_from_environment - return install_megatron_integration(integration_plan_from_environment()) + plan = integration_plan_from_environment() + integration = install_megatron_integration(plan) + if plan.implementation_for("attention", "training") is Implementation.RL_KERNEL: + _precompile_strict_attention_training(_args) + return integration + + +def _precompile_strict_attention_training(args: Any) -> None: + """Warm FA4 CuTe fwd/bwd JIT outside Vime's actor_train timer.""" + + if args is None or torch.version.hip is not None or not torch.cuda.is_available(): + return + if os.getenv("RL_KERNEL_PRECOMPILE_FA4", "1") == "0": + return + + from rl_engine.kernels.ops.cuda.attention.flash_attn import StrictFlashAttention4Core + + attention_heads = int(getattr(args, "num_attention_heads", 0) or 0) + query_groups = int(getattr(args, "num_query_groups", 0) or attention_heads) + tp_size = int(getattr(args, "tensor_model_parallel_size", 1) or 1) + head_dim = int( + getattr(args, "kv_channels", 0) + or (int(getattr(args, "hidden_size", 0) or 0) // attention_heads) + ) + if ( + attention_heads <= 0 + or query_groups <= 0 + or tp_size <= 0 + or attention_heads % tp_size + or query_groups % tp_size + or head_dim <= 0 + ): + return + + params_dtype = getattr(args, "params_dtype", None) + dtype = params_dtype if params_dtype in (torch.float16, torch.bfloat16) else torch.bfloat16 + StrictFlashAttention4Core.precompile_training( + q_heads=attention_heads // tp_size, + kv_heads=query_groups // tp_size, + head_dim=head_dim, + dtype=dtype, + ) __all__ = ["initialize_from_environment", "install_megatron_integration"] diff --git a/rl_engine/integrations/vime/attention.py b/rl_engine/integrations/vime/attention.py index d9c006bf..4eb8648f 100644 --- a/rl_engine/integrations/vime/attention.py +++ b/rl_engine/integrations/vime/attention.py @@ -365,12 +365,6 @@ def attention_provider(request: Any) -> AttentionProviderResult: """ 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) @@ -382,6 +376,12 @@ def attention_provider(request: Any) -> AttentionProviderResult: else: cp_group = None + 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] + # 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. diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index c4da28be..ade13cab 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -12,7 +12,11 @@ import torch -from rl_engine.distributed.collectives import DETERMINISTIC_ALL_REDUCE_OP +from rl_engine.distributed.collectives import ( + DETERMINISTIC_ALL_REDUCE_OP, + collective_for_group, + deterministic_all_reduce_inplace, +) from rl_engine.integrations.ablation import ( Implementation, IntegrationPlan, @@ -40,6 +44,8 @@ _STRICT_RMS_NORM_INIT_MARKER = "__rl_kernel_original_strict_rms_norm_init__" _STRICT_ROTARY_INIT_MARKER = "__rl_kernel_original_strict_rotary_init__" _STRICT_LM_HEAD_LINEAR_PATCH_MARKER = "__rl_kernel_original_lm_head_linear_apply__" +_STRICT_O_PROJ_COLLECTIVE_MARKER = "__rl_kernel_o_proj_collective__" +_STRICT_ROW_PARALLEL_PATCH_MARKER = "__rl_kernel_original_row_parallel_forward__" _RLK_ATTENTION_BACKEND: type[Any] | None = None _RLK_ATTENTION_IMPL: type[Any] | None = None _RLK_ATTENTION_BUILDER: type[Any] | None = None @@ -367,6 +373,7 @@ def _patch_qwen3_strict_model( rotary_cls: type[Any] | None = None, linear_method_cls: type[Any] | None = None, attention_cls: type[Any] | None = None, + row_parallel_cls: type[Any] | None = None, det_gemm: Any | None = None, ) -> None: """Align vLLM's RMSNorm and Attention projections with Megatron.""" @@ -374,7 +381,7 @@ def _patch_qwen3_strict_model( production_classes = rms_norm_cls is None or linear_method_cls is None or attention_cls is None if production_classes: from vllm.model_executor.layers.layernorm import RMSNorm - from vllm.model_executor.layers.linear import UnquantizedLinearMethod + from vllm.model_executor.layers.linear import RowParallelLinear, UnquantizedLinearMethod from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding from vllm.model_executor.models.qwen3 import Qwen3Attention @@ -382,6 +389,7 @@ def _patch_qwen3_strict_model( rotary_cls = RotaryEmbedding linear_method_cls = UnquantizedLinearMethod attention_cls = Qwen3Attention + row_parallel_cls = RowParallelLinear assert rms_norm_cls is not None assert linear_method_cls is not None assert attention_cls is not None @@ -435,6 +443,62 @@ def strict_rms_norm_forward_cuda( eps=instance.variance_epsilon, ) + def bind_o_proj_collective(module: Any) -> None: + if int(getattr(module, "tp_size", 1)) <= 1: + return + from vllm.distributed.parallel_state import get_tp_group + + coordinator = get_tp_group() + group = getattr(coordinator, "device_group", coordinator) + collective = collective_for_group(group) + if collective is None: + raise RuntimeError("strict rollout o_proj requires an initialized TP process group") + setattr(module, _STRICT_O_PROJ_COLLECTIVE_MARKER, collective) + + if row_parallel_cls is not None and not hasattr( + row_parallel_cls, _STRICT_ROW_PARALLEL_PATCH_MARKER + ): + row_parallel_forward = row_parallel_cls.forward + + def strict_row_parallel_forward(instance: Any, input_: torch.Tensor) -> Any: + collective = getattr(instance, _STRICT_O_PROJ_COLLECTIVE_MARKER, None) + if collective is None: + return row_parallel_forward(instance, input_) + + if instance.input_is_parallel: + input_parallel = input_ + else: + from vllm.distributed import split_tensor_along_last_dim + + input_parallel = split_tensor_along_last_dim( + input_, num_partitions=instance.tp_size + )[instance.tp_rank].contiguous() + + assert instance.quant_method is not None + bias_ = ( + None + if (instance.tp_rank > 0 or instance.skip_bias_add) + else instance.bias + ) + output_parallel = instance.quant_method.apply(instance, input_parallel, bias_) + + if instance.reduce_results and instance.tp_size > 1: + deterministic_all_reduce_inplace( + output_parallel, + collective_handle=int(collective._handle), + ) + output = output_parallel + else: + output = output_parallel + + if not instance.return_bias: + return output + output_bias = instance.bias if instance.skip_bias_add else None + return output, output_bias + + setattr(row_parallel_cls, _STRICT_ROW_PARALLEL_PATCH_MARKER, row_parallel_forward) + row_parallel_cls.forward = strict_row_parallel_forward + if not hasattr(rms_norm_cls, _STRICT_RMS_NORM_INIT_MARKER): rms_norm_init = rms_norm_cls.__init__ @@ -463,6 +527,7 @@ def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: attention_init(instance, *args, **kwargs) setattr(instance.qkv_proj, _STRICT_PROJECTION_MARKER, "qkv") setattr(instance.o_proj, _STRICT_PROJECTION_MARKER, "o_proj") + bind_o_proj_collective(instance.o_proj) setattr(attention_cls, _STRICT_MODEL_PATCH_MARKER, attention_init) linear_method_cls.apply = deterministic_linear_apply diff --git a/rl_engine/kernels/ops/cuda/attention/cp_comm.py b/rl_engine/kernels/ops/cuda/attention/cp_comm.py index da871488..6766bb67 100644 --- a/rl_engine/kernels/ops/cuda/attention/cp_comm.py +++ b/rl_engine/kernels/ops/cuda/attention/cp_comm.py @@ -589,15 +589,105 @@ 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. +class _RCCLRankOrderedTransport: + """RCCL transport with rank-ordered all-gather and root-owned scatter. - 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. + The scatter half deliberately performs no floating-point reduction. The + strict Attention arithmetic remains entirely in the deterministic core. """ + def __init__(self, *, process_group: Any, root: int) -> None: + 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" + ) + backend = str(dist.get_backend(process_group)).lower() + if "nccl" not in backend or torch.version.hip is None: + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL AG/RS requires the PyTorch NCCL API on ROCm" + ) + self.group = process_group + self.rank = int(dist.get_rank(process_group)) + self.world_size = int(dist.get_world_size(process_group)) + self.root = int(root) + if self.root < 0 or self.root >= self.world_size: + raise AttentionCPCommunicationUnavailable("RCCL scatter root is outside the group") + + def all_gather(self, local: torch.Tensor) -> torch.Tensor: + import torch.distributed as dist + + if not local.is_cuda or not local.is_contiguous(): + raise AttentionCPCommunicationUnavailable( + "RCCL AllGather requires a contiguous ROCm tensor" + ) + shape = (self.world_size * local.size(0), *local.shape[1:]) + gathered = torch.empty(shape, dtype=local.dtype, device=local.device) + dist.all_gather_into_tensor(gathered, local, group=self.group) + return gathered + + def scatter(self, full: torch.Tensor) -> torch.Tensor: + import torch.distributed as dist + + if not full.is_cuda or not full.is_contiguous(): + raise AttentionCPCommunicationUnavailable( + "RCCL ReduceScatter transport requires a contiguous ROCm tensor" + ) + if full.size(0) % self.world_size: + raise AttentionCPCommunicationUnavailable( + "RCCL scatter leading dimension must divide the CP world size" + ) + # Leading-dimension chunks are already contiguous. Avoid materializing + # copies for every rank; non-root ranks do not need a scatter list. + local_shape = (full.size(0) // self.world_size, *full.shape[1:]) + local = torch.empty(local_shape, dtype=full.dtype, device=full.device) + if self.world_size == 1: + local.copy_(full) + return local + global_root = self.root + if self.group is not None: + get_global_rank = getattr(dist, "get_global_rank", None) + if callable(get_global_rank): + global_root = int(get_global_rank(self.group, self.root)) + else: + global_root = int(dist.get_process_group_ranks(self.group)[self.root]) + scatter_list = list(full.chunk(self.world_size, dim=0)) if self.rank == self.root else None + dist.scatter(local, scatter_list=scatter_list, src=global_root, group=self.group) + return local + + def reduce_scatter(self, full: torch.Tensor) -> torch.Tensor: + """Deterministic rank-order sum followed by local scatter. + + RCCL is used for point-to-point transport. The floating-point sum is + performed locally in source-rank order, so collective reduction order + is not delegated to RCCL. + """ + if not full.is_cuda or not full.is_contiguous(): + raise AttentionCPCommunicationUnavailable( + "RCCL ReduceScatter requires a contiguous ROCm tensor" + ) + if full.size(0) % self.world_size: + raise AttentionCPCommunicationUnavailable( + "RCCL ReduceScatter leading dimension must divide the CP world size" + ) + chunks = tuple(chunk.contiguous() for chunk in full.chunk(self.world_size, dim=0)) + gathered = self.all_gather(full) + local = gathered[ + self.rank * chunks[self.rank].size(0) : (self.rank + 1) * chunks[self.rank].size(0) + ].clone() + chunk_rows = chunks[self.rank].size(0) + # Start with source rank 0, then add the remaining source ranks in + # ascending order. This avoids counting source 0 twice. + for source in range(1, self.world_size): + source_full = gathered[source * full.size(0) : (source + 1) * full.size(0)] + local.add_(source_full[self.rank * chunk_rows : (self.rank + 1) * chunk_rows]) + return local + + +class RCCLAGRSAttentionCPCommunication(CUDAAGRSAttentionCPCommunication): + """ROCm AG/RS adapter using RCCL only as rank-ordered tensor transport.""" + backend_id = "rccl_ag_rs" collective_label = "self-owned RCCL AG/RS" supports_autograd = True @@ -605,6 +695,18 @@ class RCCLAGRSAttentionCPCommunication(CUDAAGRSAttentionCPCommunication): supports_async_overlap = False supports_compute_communication_fusion = False + def _get_collective(self, plan: AttentionCPCommunicationPlan): + if self._collective is None: + self._collective = _RCCLRankOrderedTransport( + process_group=self._process_group, + root=plan.merge_root_cp_rank, + ) + if self._collective.world_size != plan.parallel.cp_world_size: + raise AttentionCPCommunicationUnavailable( + "self-owned RCCL world size does not match the CP plan" + ) + return self._collective + def _dist(self): import torch.distributed as dist diff --git a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py index 1ed6d47c..3a9084b1 100644 --- a/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/deterministic_attn.py @@ -226,9 +226,8 @@ def _validate_inputs( class RLKernelDeterministicAttentionCore: """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`` on - CUDA, and AITER CK dense MHA on ROCm. + Production uses FA4 CuTe on CUDA and AITER CK dense MHA on ROCm. This core + remains useful for correctness and capability-gap diagnosis. """ core_id = STRICT_ATTENTION_CORE_ID diff --git a/rl_engine/kernels/ops/cuda/attention/flash_attn.py b/rl_engine/kernels/ops/cuda/attention/flash_attn.py index cec8d2fe..9ad510b3 100644 --- a/rl_engine/kernels/ops/cuda/attention/flash_attn.py +++ b/rl_engine/kernels/ops/cuda/attention/flash_attn.py @@ -109,6 +109,82 @@ def __init__( self._op = op self._paged_op = paged_op + @classmethod + def precompile_training( + cls, + *, + q_heads: int, + kv_heads: int, + head_dim: int, + device: torch.device | None = None, + dtype: torch.dtype = torch.bfloat16, + sequence_length: int = 512, + ) -> None: + """Compile strict training FA4 forward/backward before timing. + + FA4 CuTe compiles its forward kernel on the first invocation and its + deterministic backward kernel on the first autograd backward. A + one-step RL workload would otherwise charge both compilations to + Vime's ``actor_train`` timer. These isolated tensors exercise the + same Qwen-style GQA and multi-block shape class without touching model + tensors, RNG state, or distributed collectives. + """ + if torch.version.hip is not None: + raise StrictFlashAttentionUnavailable( + "FA4 CUDA precompile is unavailable on ROCm" + ) + if not torch.cuda.is_available(): + raise StrictFlashAttentionUnavailable( + "FA4 CUDA precompile requires an available CUDA device" + ) + if dtype not in (torch.float16, torch.bfloat16): + raise ValueError("strict FA4 training precompile requires FP16 or BF16") + if q_heads <= 0 or kv_heads <= 0 or q_heads % kv_heads != 0: + raise ValueError("Q/KV head counts must be positive and GQA-compatible") + if head_dim <= 0 or sequence_length <= 0: + raise ValueError("head_dim and sequence_length must be positive") + + target = ( + torch.device("cuda", torch.cuda.current_device()) + if device is None + else device + ) + if target.type != "cuda": + raise ValueError("strict FA4 training precompile requires a CUDA device") + + core = cls() + # Zeros deliberately avoid consuming RNG state. The tensors are + # independent from the model and only populate FA4's process-local + # JIT caches. + q = torch.zeros( + (1, sequence_length, q_heads, head_dim), + dtype=dtype, + device=target, + requires_grad=True, + ) + k = torch.zeros( + (1, sequence_length, kv_heads, head_dim), + dtype=dtype, + device=target, + requires_grad=True, + ) + v = torch.zeros_like(k, requires_grad=True) + positions = torch.arange(sequence_length, dtype=torch.int64, device=target).expand(1, -1) + with torch.enable_grad(): + result = core.forward_bshd_with_lse( + q, + k, + v, + causal=True, + scale=head_dim**-0.5, + query_position_ids=positions, + key_position_ids=positions, + output_dtype=dtype, + ) + result.out.sum().backward() + torch.cuda.synchronize(target) + del result, q, k, v, positions, core + @staticmethod def _validate_api(op: Callable[..., Any]) -> None: try: diff --git a/rl_engine/kernels/ops/triton/activation/swiglu.py b/rl_engine/kernels/ops/triton/activation/swiglu.py index 6fb66313..ad11ecc8 100644 --- a/rl_engine/kernels/ops/triton/activation/swiglu.py +++ b/rl_engine/kernels/ops/triton/activation/swiglu.py @@ -32,7 +32,7 @@ def _silu_fwd_kernel(x_ptr, y_ptr, n_elements, BLOCK: tl.constexpr): offs = pid * BLOCK + tl.arange(0, BLOCK) mask = offs < n_elements x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) - s = 1.0 / (1.0 + tl.exp(-x)) + s = tl.div_rn(1.0, 1.0 + tl.exp(-x)) y = x * s tl.store(y_ptr + offs, y.to(y_ptr.dtype.element_ty), mask=mask) @@ -44,7 +44,7 @@ def _silu_bwd_kernel(dy_ptr, x_ptr, dx_ptr, n_elements, BLOCK: tl.constexpr): mask = offs < n_elements dy = tl.load(dy_ptr + offs, mask=mask, other=0.0).to(tl.float32) x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32) - s = 1.0 / (1.0 + tl.exp(-x)) + s = tl.div_rn(1.0, 1.0 + tl.exp(-x)) # silu'(x) = s * (1 + x * (1 - s)) dx = dy * s * (1.0 + x * (1.0 - s)) tl.store(dx_ptr + offs, dx.to(dx_ptr.dtype.element_ty), mask=mask) @@ -57,26 +57,42 @@ def _swiglu_fwd_kernel(gate_ptr, up_ptr, y_ptr, n_elements, BLOCK: tl.constexpr) mask = offs < n_elements g = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32) u = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32) - s = 1.0 / (1.0 + tl.exp(-g)) + s = tl.div_rn(1.0, 1.0 + tl.exp(-g)) y = (g * s) * u tl.store(y_ptr + offs, y.to(y_ptr.dtype.element_ty), mask=mask) @triton.jit def _swiglu_bwd_kernel( - dy_ptr, gate_ptr, up_ptr, d_gate_ptr, d_up_ptr, n_elements, BLOCK: tl.constexpr + dy_ptr, gate_ptr, silu_grad_ptr, d_up_ptr, n_elements, BLOCK: tl.constexpr ): pid = tl.program_id(0) offs = pid * BLOCK + tl.arange(0, BLOCK) mask = offs < n_elements dy = tl.load(dy_ptr + offs, mask=mask, other=0.0).to(tl.float32) g = tl.load(gate_ptr + offs, mask=mask, other=0.0).to(tl.float32) - u = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32) - s = 1.0 / (1.0 + tl.exp(-g)) + s = tl.div_rn(1.0, 1.0 + tl.exp(-g)) silu_g = g * s d_up = dy * silu_g - d_gate = dy * u * s * (1.0 + g * (1.0 - s)) + silu_grad = s * (1.0 + g * (1.0 - s)) tl.store(d_up_ptr + offs, d_up.to(d_up_ptr.dtype.element_ty), mask=mask) + tl.store(silu_grad_ptr + offs, silu_grad, mask=mask) + + +@triton.jit +def _swiglu_dgate_kernel( + dy_ptr, up_ptr, silu_grad_ptr, d_gate_ptr, n_elements, BLOCK: tl.constexpr +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + dy = tl.load(dy_ptr + offs, mask=mask, other=0.0).to(tl.float32) + u = tl.load(up_ptr + offs, mask=mask, other=0.0).to(tl.float32) + silu_grad = tl.load(silu_grad_ptr + offs, mask=mask, other=0.0) + # The native HIP source evaluates (dy * up) * silu_grad. Persisting the + # latter as FP32 prevents LLVM from reassociating the expression across + # the function boundary and changing a BF16 tie at rare elements. + d_gate = (dy * u) * silu_grad tl.store(d_gate_ptr + offs, d_gate.to(d_gate_ptr.dtype.element_ty), mask=mask) @@ -124,11 +140,27 @@ def _launch_swiglu_bwd(dy: Tensor, gate: Tensor, up: Tensor) -> tuple[Tensor, Te up_c = up.contiguous() d_gate = torch.empty_like(gate_c) d_up = torch.empty_like(up_c) + silu_grad = torch.empty_like(gate_c, dtype=torch.float32) n = gate_c.numel() if n == 0: return d_gate, d_up grid = (triton.cdiv(n, _BLOCK),) - _swiglu_bwd_kernel[grid](dy_c, gate_c, up_c, d_gate, d_up, n, BLOCK=_BLOCK) + _swiglu_bwd_kernel[grid]( + dy_c, + gate_c, + silu_grad, + d_up, + n, + BLOCK=_BLOCK, + ) + _swiglu_dgate_kernel[grid]( + dy_c, + up_c, + silu_grad, + d_gate, + n, + BLOCK=_BLOCK, + ) return d_gate, d_up diff --git a/rl_engine/kernels/ops/triton/ffn/__init__.py b/rl_engine/kernels/ops/triton/ffn/__init__.py new file mode 100644 index 00000000..16a8565c --- /dev/null +++ b/rl_engine/kernels/ops/triton/ffn/__init__.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from .ffn import ( + QWEN3_8B_HIDDEN_SIZE, + QWEN3_8B_INTERMEDIATE_SIZE, + Qwen3FFNForwardWeights, + pack_qwen3_ffn_forward_weights, + qwen3_ffn, + qwen3_ffn_triton, + refresh_qwen3_ffn_forward_weights, +) + +__all__ = [ + "QWEN3_8B_HIDDEN_SIZE", + "QWEN3_8B_INTERMEDIATE_SIZE", + "Qwen3FFNForwardWeights", + "pack_qwen3_ffn_forward_weights", + "qwen3_ffn", + "qwen3_ffn_triton", + "refresh_qwen3_ffn_forward_weights", +] diff --git a/rl_engine/kernels/ops/triton/ffn/ffn.py b/rl_engine/kernels/ops/triton/ffn/ffn.py new file mode 100644 index 00000000..f618031a --- /dev/null +++ b/rl_engine/kernels/ops/triton/ffn/ffn.py @@ -0,0 +1,636 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""ROCm-native distributed deterministic Qwen3 FFN built with Triton. + +BF16 is preserved at every GEMM/SwiGLU boundary, GEMMs use a canonical +FP32-leaf/BF16-node K tree, and TP reductions use the rank-ordered balanced +RCCL transport collective. CP gathers full token sequences before +weight-gradient GEMMs so their K tree is identical to CP=1. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch +from torch import Tensor + +from rl_engine.kernels.ops.triton.activation.swiglu import ( + _launch_swiglu_bwd, + _launch_swiglu_fwd, +) +from rl_engine.kernels.ops.triton.matmul.det_gemm import _triton_tree_gemm + +QWEN3_8B_HIDDEN_SIZE = 4096 +QWEN3_8B_INTERMEDIATE_SIZE = 12288 + +_COLLECTIVE_MIN_CAPACITY_BYTES = 64 * 1024 * 1024 +_COLLECTIVES: dict[tuple[int, int, int, int], Any] = {} + + +@dataclass(eq=False) +class Qwen3FFNForwardWeights: + """Stable-storage GEMM-ready copies produced after loading or synchronization.""" + + gate_weight_t: Tensor + up_weight_t: Tensor + down_weight_t: Tensor + _sources: tuple[Tensor, Tensor, Tensor] = field(repr=False) + _source_data_ptrs: tuple[int, int, int] = field(repr=False) + _source_versions: tuple[int | None, int | None, int | None] = field(repr=False) + _packed_data_ptrs: tuple[int, int, int] = field(repr=False) + _packed_versions: tuple[int, int, int] = field(repr=False) + _source_shapes: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]] = field( + repr=False + ) + _source_strides: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]] = field( + repr=False + ) + + def refresh_( + self, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + ) -> Qwen3FFNForwardWeights: + """Refresh values in-place while preserving CUDA Graph-visible addresses.""" + + return refresh_qwen3_ffn_forward_weights( + self, + gate_weight, + up_weight, + down_weight, + ) + + +def _require_parallel_group(group: Any, name: str): + if group is None: + return None + + import torch.distributed as dist + + if not dist.is_available(): + raise RuntimeError(f"{name}-parallel FFN requires torch.distributed.") + if not dist.is_initialized(): + raise RuntimeError(f"{name}-parallel FFN requires an initialized process group.") + if dist.get_world_size(group=group) <= 1: + raise ValueError(f"{name}_group must contain at least two ranks.") + return dist + + +def _validate_ffn_inputs( + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, +) -> None: + tensors = { + "rmsnorm_output": rmsnorm_output, + "gate_weight": gate_weight, + "up_weight": up_weight, + "down_weight": down_weight, + } + for name, tensor in tensors.items(): + if not isinstance(tensor, Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(tensor)!r}.") + + if rmsnorm_output.dim() < 1: + raise ValueError("rmsnorm_output must have at least one dimension.") + if rmsnorm_output.numel() == 0: + raise ValueError("rmsnorm_output must contain at least one token.") + for name, weight in ( + ("gate_weight", gate_weight), + ("up_weight", up_weight), + ("down_weight", down_weight), + ): + if weight.dim() != 2: + raise ValueError(f"{name} must be 2-D, got shape {tuple(weight.shape)}.") + + hidden_size = rmsnorm_output.size(-1) + intermediate_size = gate_weight.size(0) + if intermediate_size == 0: + raise ValueError("FFN intermediate size must be positive.") + expected_shapes = { + "gate_weight": (intermediate_size, hidden_size), + "up_weight": (intermediate_size, hidden_size), + "down_weight": (hidden_size, intermediate_size), + } + for name, expected in expected_shapes.items(): + actual = tuple(tensors[name].shape) + if actual != expected: + raise ValueError(f"{name} must have shape {expected}, got {actual}.") + + for name, tensor in tensors.items(): + 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/ROCm GPU device, got '{tensor.device}'." + ) + if tensor.device != rmsnorm_output.device: + raise RuntimeError( + f"all FFN inputs must be on {rmsnorm_output.device}, " + f"got {name} on {tensor.device}." + ) + + +def _tracked_tensor_version(tensor: Tensor) -> int | None: + # Inference tensors deliberately have no version counter. Packed buffers are + # allocated outside inference mode below, but a loader-owned source may be an + # inference tensor and therefore relies on the explicit refresh lifecycle. + return None if torch.is_inference(tensor) else int(tensor._version) + + +def _validate_forward_weight_sources( + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, +) -> tuple[Tensor, Tensor, Tensor]: + weights = { + "gate_weight": gate_weight, + "up_weight": up_weight, + "down_weight": down_weight, + } + for name, weight in weights.items(): + if not isinstance(weight, Tensor): + raise TypeError(f"{name} must be a torch.Tensor, got {type(weight)!r}.") + if weight.dim() != 2: + raise ValueError(f"{name} must be 2-D, got shape {tuple(weight.shape)}.") + if weight.dtype != torch.bfloat16: + raise TypeError(f"{name} must have dtype bfloat16, got {weight.dtype}.") + if not weight.is_cuda: + raise RuntimeError( + f"{name} must be on a CUDA/ROCm GPU device, got '{weight.device}'." + ) + + if tuple(up_weight.shape) != tuple(gate_weight.shape): + raise ValueError( + "up_weight must have the same shape as gate_weight, got " + f"{tuple(up_weight.shape)} and {tuple(gate_weight.shape)}." + ) + expected_down_shape = (gate_weight.size(1), gate_weight.size(0)) + if tuple(down_weight.shape) != expected_down_shape: + raise ValueError( + f"down_weight must have shape {expected_down_shape}, " + f"got {tuple(down_weight.shape)}." + ) + if any(weight.device != gate_weight.device for weight in weights.values()): + raise RuntimeError("all FFN weights must be on the same device before packing.") + return gate_weight, up_weight, down_weight + + +def pack_qwen3_ffn_forward_weights( + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, +) -> Qwen3FFNForwardWeights: + """Prepare detached forward-only transposes outside the FFN hot path. + + This is analogous to vLLM's backend-specific + ``process_weights_after_loading``/kernel-format lifecycle: prepare the + kernel-facing layout once per weight load/reload and reuse it at runtime. + The canonical weights remain the source of truth for backward and + optimization. Call ``refresh_`` after every optimizer update or external + synchronization. Freshness checks are best-effort for inference tensors and + external writers, so the loader/optimizer must refresh even when it mutates + through ``.data``, DLPack, or a custom kernel. + """ + + sources = _validate_forward_weight_sources(gate_weight, up_weight, down_weight) + # Ensure the cached storage has an ordinary version counter even when the + # model loader calls us from torch.inference_mode(). A fresh explicit copy + # also prevents degenerate transposes from aliasing source storage. + with torch.inference_mode(False), torch.no_grad(): + packed = tuple( + torch.empty( + (weight.size(1), weight.size(0)), + dtype=weight.dtype, + device=weight.device, + ) + for weight in sources + ) + for packed_weight, source in zip(packed, sources, strict=True): + packed_weight.copy_(source.t()) + return Qwen3FFNForwardWeights( + gate_weight_t=packed[0], + up_weight_t=packed[1], + down_weight_t=packed[2], + _sources=sources, + _source_data_ptrs=tuple(weight.data_ptr() for weight in sources), + _source_versions=tuple(_tracked_tensor_version(weight) for weight in sources), + _packed_data_ptrs=tuple(weight.data_ptr() for weight in packed), + _packed_versions=tuple(int(weight._version) for weight in packed), + _source_shapes=tuple(tuple(weight.shape) for weight in sources), + _source_strides=tuple(tuple(weight.stride()) for weight in sources), + ) + + +def refresh_qwen3_ffn_forward_weights( + forward_weights: Qwen3FFNForwardWeights, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, +) -> Qwen3FFNForwardWeights: + """Refresh a forward cache in-place without changing any packed data pointer. + + Stable storage is required by already-captured CUDA Graphs: replacing either + the canonical tensors or this bundle would leave graph nodes pointing at old + values. Copy reloaded values into the original canonical tensors first, then + call this function. Order both copies before graph replay, normally on the + same stream or with an explicit stream dependency. + """ + + if not isinstance(forward_weights, Qwen3FFNForwardWeights): + raise TypeError( + "forward_weights must be created by " + "pack_qwen3_ffn_forward_weights." + ) + sources = _validate_forward_weight_sources(gate_weight, up_weight, down_weight) + if any( + source is not original + for source, original in zip(sources, forward_weights._sources, strict=True) + ): + raise ValueError( + "stable refresh requires the original canonical weight tensors; " + "copy new values into those tensors, or repack and recapture CUDA Graphs." + ) + packed = ( + forward_weights.gate_weight_t, + forward_weights.up_weight_t, + forward_weights.down_weight_t, + ) + for name, target, source, expected_ptr in zip( + ("gate_weight", "up_weight", "down_weight"), + packed, + sources, + forward_weights._packed_data_ptrs, + strict=True, + ): + expected_shape = (source.size(1), source.size(0)) + if tuple(target.shape) != expected_shape: + raise ValueError( + f"packed {name} has shape {tuple(target.shape)}, but refresh " + f"requires {expected_shape}; repack and recapture CUDA Graphs." + ) + if target.dtype != source.dtype or target.device != source.device: + raise RuntimeError( + f"packed {name} dtype/device cannot change during stable refresh; " + "repack and recapture CUDA Graphs." + ) + if not target.is_contiguous() or target.data_ptr() != expected_ptr: + raise RuntimeError( + f"packed {name} storage changed; repack and recapture CUDA Graphs." + ) + + with torch.inference_mode(False), torch.no_grad(): + for target, source in zip(packed, sources, strict=True): + target.copy_(source.t()) + + forward_weights._source_versions = tuple( + _tracked_tensor_version(weight) for weight in sources + ) + forward_weights._packed_versions = tuple(int(weight._version) for weight in packed) + return forward_weights + + +def _validate_forward_weights( + forward_weights: Qwen3FFNForwardWeights, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, +) -> None: + if not isinstance(forward_weights, Qwen3FFNForwardWeights): + raise TypeError( + "forward_weights must be created by " + "pack_qwen3_ffn_forward_weights." + ) + + sources = (gate_weight, up_weight, down_weight) + names = ("gate_weight", "up_weight", "down_weight") + for index, (name, source) in enumerate(zip(names, sources, strict=True)): + if forward_weights._sources[index] is not source: + raise ValueError(f"forward_weights was not packed from this {name} tensor.") + if forward_weights._source_data_ptrs[index] != source.data_ptr(): + raise RuntimeError(f"{name} storage changed after forward weights were packed.") + if forward_weights._source_shapes[index] != tuple(source.shape): + raise RuntimeError(f"{name} shape changed after forward weights were packed.") + if forward_weights._source_strides[index] != tuple(source.stride()): + raise RuntimeError(f"{name} strides changed after forward weights were packed.") + if forward_weights._source_versions[index] != _tracked_tensor_version(source): + raise RuntimeError( + f"{name} changed after forward weights were packed; refresh before FFN." + ) + + expected = ( + (gate_weight.size(1), gate_weight.size(0)), + (up_weight.size(1), up_weight.size(0)), + (down_weight.size(1), down_weight.size(0)), + ) + packed = ( + forward_weights.gate_weight_t, + forward_weights.up_weight_t, + forward_weights.down_weight_t, + ) + for name, weight, shape in zip(names, packed, expected, strict=True): + if not isinstance(weight, Tensor): + raise TypeError(f"packed {name} must be a torch.Tensor.") + if tuple(weight.shape) != shape: + raise ValueError( + f"packed {name} must have shape {shape}, got {tuple(weight.shape)}." + ) + if weight.dtype != torch.bfloat16: + raise TypeError(f"packed {name} must have dtype bfloat16, got {weight.dtype}.") + if weight.device != gate_weight.device: + raise RuntimeError( + f"packed {name} must be on {gate_weight.device}, got {weight.device}." + ) + if not weight.is_contiguous(): + raise ValueError(f"packed {name} must be contiguous.") + if weight.requires_grad: + raise ValueError(f"packed {name} must be detached from autograd.") + for index, (name, weight) in enumerate(zip(names, packed, strict=True)): + if forward_weights._packed_data_ptrs[index] != weight.data_ptr(): + raise RuntimeError(f"packed {name} storage changed; repack before FFN.") + if forward_weights._packed_versions[index] != int(weight._version): + raise RuntimeError(f"packed {name} changed; refresh before FFN.") + + +def _create_collective(*, group: Any, max_size_bytes: int): + try: + from rl_engine.distributed import create_deterministic_collective + except ImportError as exc: + raise RuntimeError( + "parallel Triton FFN requires the deterministic collective factory" + ) from exc + return create_deterministic_collective( + group=group, + max_size_bytes=max_size_bytes, + ) + + +def _collective_for_group(group: Any, *, min_size_bytes: int): + if group is None: + return None + + import torch.distributed as dist + + rank = dist.get_rank(group=group) + world_size = dist.get_world_size(group=group) + device_index = torch.cuda.current_device() + key = (id(group), rank, world_size, device_index) + cached = _COLLECTIVES.get(key) + if cached is not None and cached.max_size_bytes >= min_size_bytes: + return cached + if cached is not None: + cached.close() + + collective = _create_collective( + group=group, + max_size_bytes=max(_COLLECTIVE_MIN_CAPACITY_BYTES, min_size_bytes), + ) + _COLLECTIVES[key] = collective + return collective + + +def _all_gather_tokens(tensor: Tensor, collective: Any) -> Tensor: + return collective.all_gather(tensor.contiguous()) + + +def _reduce_scatter_tokens(tensor: Tensor, collective: Any) -> Tensor: + world_size = collective.world_size + if tensor.size(0) % world_size != 0: + raise ValueError( + "the gathered token count must be divisible by the tensor-parallel " + f"world size, got {tensor.size(0)} and {world_size}." + ) + return collective.reduce_scatter(tensor.contiguous()) + + +def _all_reduce_inplace(tensor: Tensor, collective: Any) -> Tensor: + return collective.all_reduce(tensor, out=tensor) + + +def _gemm(a: Tensor, b: Tensor) -> Tensor: + return _triton_tree_gemm(a, b) + + +def _gemm_db(a: Tensor, grad_output: Tensor) -> Tensor: + grad_weight = torch.empty( + (grad_output.size(1), a.size(1)), + dtype=torch.bfloat16, + device=a.device, + ) + # Mirror only the buffer placement used by Transformer Engine's + # fuse_wgrad_accumulation and Megatron's gradient_accumulation_fusion: write + # the canonical GEMM root into the final [out, in] gradient buffer. The + # strict operand order, K tree, BF16 nodes, and rounding remain unchanged. + # Unlike their fused main_grad paths, this returns one fresh autograd dW and + # does not fuse or reorder microbatch accumulation. + return _triton_tree_gemm( + a.t(), + grad_output, + transpose_output=True, + out=grad_weight, + preserve_a_strides=True, + ) + + +class _TritonDeterministicFFNFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + gate_weight_t: Tensor | None, + up_weight_t: Tensor | None, + down_weight_t: Tensor | None, + tp_group: Any, + cp_group: Any, + sequence_parallel: bool, + ) -> Tensor: + tp_dist = _require_parallel_group(tp_group, "tensor") + _require_parallel_group(cp_group, "context") + if sequence_parallel and tp_dist is None: + raise ValueError("sequence_parallel requires a tensor-parallel group.") + + input_shape = rmsnorm_output.shape + rmsnorm_output_2d = rmsnorm_output.reshape(-1, input_shape[-1]).contiguous() + 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 independent gate/up input + # gradient lanes in one fixed-tree transport call. + reduction_bytes = token_hidden_bytes * (2 if sequence_parallel else 1) + min_size_bytes = max( + reduction_bytes, + gemm_tokens * gate_weight.size(0) * element_size, + gate_weight.numel() * element_size, + up_weight.numel() * element_size, + down_weight.numel() * element_size, + ) + tp_collective = _collective_for_group(tp_group, min_size_bytes=min_size_bytes) + cp_collective = _collective_for_group(cp_group, min_size_bytes=min_size_bytes) + + if sequence_parallel: + rmsnorm_output_2d = _all_gather_tokens(rmsnorm_output_2d, tp_collective) + + # A loader integration can prepare/cache this backend-specific layout + # once per load/reload, analogous to vLLM's + # process_weights_after_loading lifecycle. Canonical weights stay + # untouched and are saved below for backward. + gate = _gemm( + rmsnorm_output_2d, + gate_weight.t().contiguous() if gate_weight_t is None else gate_weight_t, + ) + up = _gemm( + rmsnorm_output_2d, + up_weight.t().contiguous() if up_weight_t is None else up_weight_t, + ) + activated = _launch_swiglu_fwd(gate, up) + output = _gemm( + activated, + down_weight.t().contiguous() if down_weight_t is None else down_weight_t, + ) + + if sequence_parallel: + output = _reduce_scatter_tokens(output, tp_collective) + elif tp_collective is not None: + output = _all_reduce_inplace(output, tp_collective) + + ctx.save_for_backward( + rmsnorm_output_2d, + gate, + up, + activated, + gate_weight, + up_weight, + down_weight, + ) + ctx.input_shape = input_shape + ctx.tp_collective = tp_collective + ctx.cp_collective = cp_collective + ctx.sequence_parallel = sequence_parallel + return output.reshape(*input_shape[:-1], output.size(-1)) + + @staticmethod + def backward(ctx: Any, grad_output: Tensor) -> tuple[Any, ...]: + ( + rmsnorm_output, + gate, + up, + activated, + gate_weight, + up_weight, + down_weight, + ) = ctx.saved_tensors + tp_collective = ctx.tp_collective + cp_collective = ctx.cp_collective + grad_output = grad_output.reshape(-1, grad_output.size(-1)).contiguous() + if ctx.sequence_parallel: + grad_output = _all_gather_tokens(grad_output, tp_collective) + + if cp_collective is not None: + activated_full = _all_gather_tokens(activated, cp_collective) + grad_output_full = _all_gather_tokens(grad_output, cp_collective) + grad_down_weight = _gemm_db(activated_full, grad_output_full) + else: + grad_down_weight = _gemm_db(activated, grad_output) + + grad_activated = _gemm(grad_output, down_weight) + grad_gate, grad_up = _launch_swiglu_bwd(grad_activated, gate, up) + + if cp_collective is not None: + rmsnorm_full = _all_gather_tokens(rmsnorm_output, cp_collective) + grad_gate_full = _all_gather_tokens(grad_gate, cp_collective) + grad_up_full = _all_gather_tokens(grad_up, cp_collective) + grad_gate_weight = _gemm_db(rmsnorm_full, grad_gate_full) + grad_up_weight = _gemm_db(rmsnorm_full, grad_up_full) + else: + grad_gate_weight = _gemm_db(rmsnorm_output, grad_gate) + grad_up_weight = _gemm_db(rmsnorm_output, grad_up) + + grad_rmsnorm_from_gate = _gemm(grad_gate, gate_weight) + grad_rmsnorm_from_up = _gemm(grad_up, up_weight) + if ctx.sequence_parallel: + 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, + ) + + grad_rmsnorm_output = grad_rmsnorm_from_gate.add_(grad_rmsnorm_from_up) + return ( + grad_rmsnorm_output.reshape(ctx.input_shape), + grad_gate_weight, + grad_up_weight, + grad_down_weight, + None, + None, + None, + None, + None, + None, + ) + + +def qwen3_ffn( + rmsnorm_output: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + forward_weights: Qwen3FFNForwardWeights | None = None, + tp_group: Any = None, + cp_group: Any = None, + sequence_parallel: bool = False, +) -> Tensor: + """Apply the distributed deterministic Qwen3 FFN with Triton kernels. + + ``forward_weights`` is an optional, forward-only cache. Refresh it in-place + after every optimizer update or external weight synchronization. The + canonical weight arguments always remain the autograd/optimizer source of + truth. Refresh is mandatory for inference tensors and external writers, + whose mutations cannot always be discovered from a PyTorch version counter. + """ + + _validate_ffn_inputs(rmsnorm_output, gate_weight, up_weight, down_weight) + if forward_weights is not None: + _validate_forward_weights( + forward_weights, + gate_weight, + up_weight, + down_weight, + ) + if not isinstance(sequence_parallel, bool): + raise TypeError( + f"sequence_parallel must be a bool, got {type(sequence_parallel)!r}." + ) + return _TritonDeterministicFFNFunction.apply( + rmsnorm_output, + gate_weight, + up_weight, + down_weight, + None if forward_weights is None else forward_weights.gate_weight_t, + None if forward_weights is None else forward_weights.up_weight_t, + None if forward_weights is None else forward_weights.down_weight_t, + tp_group, + cp_group, + sequence_parallel, + ) + + +# Keep the explicit suffix for callers that select an implementation by name. +qwen3_ffn_triton = qwen3_ffn diff --git a/rl_engine/kernels/ops/triton/matmul/det_gemm.py b/rl_engine/kernels/ops/triton/matmul/det_gemm.py index 50025db2..078167b0 100644 --- a/rl_engine/kernels/ops/triton/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/triton/matmul/det_gemm.py @@ -1,12 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Batch-invariant deterministic GEMM, Triton path (WS1). +"""Batch- and TP-invariant deterministic GEMM, Triton path. -Portable implementation with the SAME invariance guarantees as the CUDA path: -autotune disabled, BLOCK sizes pinned, no split-K, fixed K-loop order, FP32 -accumulation, no TF32. Used as the cross-backend reference and the ROCm/portable -fallback. Slower than a tuned GEMM by design. +BF16 outputs use the same arithmetic graph as the native deterministic kernel: +each canonical K-tree leaf contains at most 32 values, accumulates in FP32, and +rounds to BF16. Separate Triton kernels then evaluate the canonical midpoint +tree with BF16 nodes. The strict ROCm leaf keeps the native gfx942 scalar FMA +order so both implementations have zero mismatch; a future MFMA leaf requires +a separately versioned arithmetic contract because its internal accumulation +does not match the scalar reference at every BF16 rounding boundary. + +Autotuning and split-K are intentionally disabled. FP32-output calls preserve +the earlier fixed-order, no-split-K contract and do not use BF16 tree nodes. """ + +from __future__ import annotations + +import functools +import threading +from dataclasses import dataclass + import torch try: @@ -22,12 +35,357 @@ # Pinned. NOT autotuned (autotune picks per-shape configs -> breaks invariance). _BLOCK_M, _BLOCK_N, _BLOCK_K = 64, 64, 32 +_TREE_PLANS: dict[tuple[int, int], "_DeviceTreePlan"] = {} +_TREE_PLAN_LOCK = threading.Lock() + + +@dataclass(frozen=True) +class _TreeLeafConfig: + block_m: int + block_n: int + num_warps: int + n_fastest: bool + + +_DEFAULT_TREE_LEAF_CONFIG = _TreeLeafConfig(_BLOCK_M, _BLOCK_N, 4, False) + +# Offline-swept on ROCm gfx942 with Triton 3.7. These entries deliberately +# cover only the Qwen3-8B TP1 logical shapes and stride modes used by the FFN. +# Other architectures and shapes retain the established 64x64 specialization. +_GFX942_QWEN_FORWARD_LEAF_CONFIGS = { + 1: _TreeLeafConfig(1, 128, 1, True), + 8: _TreeLeafConfig(8, 128, 2, True), + 32: _TreeLeafConfig(32, 128, 4, True), +} +_GFX942_QWEN_WGRAD_LEAF_CONFIGS = { + 1: _TreeLeafConfig(128, 64, 2, True), + 8: _TreeLeafConfig(128, 64, 2, True), + 32: _TreeLeafConfig(64, 64, 2, True), +} +_QWEN_FORWARD_GEMM_SHAPES = {(4096, 12288), (12288, 4096)} +_QWEN_WGRAD_OUTPUT_SHAPES = {(4096, 12288), (12288, 4096)} + +# Qwen3-8B TP2 local shards. TP4/8 candidates improved isolated leaves but did +# not clear the distributed end-to-end promotion threshold, so they deliberately +# retain the fallback. These entries were swept independently from the TP1 table +# because smaller local N/K changes both grid occupancy and the +# point where tree-reduction launch cost dominates leaf time. Keep the key +# exact: nearby shapes and non-target token counts retain the established +# fallback instead of inheriting a configuration from a different TP graph. +_GFX942_QWEN_TP_SHARD_FORWARD_LEAF_CONFIGS = { + (16, 4096, 6144): _TreeLeafConfig(8, 64, 1, True), + (16, 6144, 4096): _TreeLeafConfig(8, 64, 1, True), + (32, 4096, 6144): _TreeLeafConfig(32, 128, 4, True), + (32, 6144, 4096): _TreeLeafConfig(32, 128, 4, True), +} +_GFX942_QWEN_TP_SHARD_WGRAD_LEAF_CONFIGS = { + (4096, 32, 6144): _TreeLeafConfig(128, 64, 2, True), + (6144, 32, 4096): _TreeLeafConfig(128, 64, 2, True), +} + + +def _gfx942_qwen_tree_leaf_config( + m_size: int, + k_size: int, + n_size: int, + *, + transpose_output: bool, + preserve_a_strides: bool, +) -> _TreeLeafConfig: + logical_shape = (m_size, k_size, n_size) + if ( + not transpose_output + and not preserve_a_strides + and (k_size, n_size) in _QWEN_FORWARD_GEMM_SHAPES + ): + return _GFX942_QWEN_FORWARD_LEAF_CONFIGS.get( + m_size, + _DEFAULT_TREE_LEAF_CONFIG, + ) + if not transpose_output and not preserve_a_strides: + return _GFX942_QWEN_TP_SHARD_FORWARD_LEAF_CONFIGS.get( + logical_shape, + _DEFAULT_TREE_LEAF_CONFIG, + ) + if ( + transpose_output + and preserve_a_strides + and (m_size, n_size) in _QWEN_WGRAD_OUTPUT_SHAPES + ): + return _GFX942_QWEN_WGRAD_LEAF_CONFIGS.get( + k_size, + _DEFAULT_TREE_LEAF_CONFIG, + ) + if transpose_output and preserve_a_strides: + return _GFX942_QWEN_TP_SHARD_WGRAD_LEAF_CONFIGS.get( + logical_shape, + _DEFAULT_TREE_LEAF_CONFIG, + ) + return _DEFAULT_TREE_LEAF_CONFIG + + +@functools.lru_cache(maxsize=None) +def _device_arch(device_index: int) -> str: + if getattr(torch.version, "hip", None) is None: + return "" + properties = torch.cuda.get_device_properties(device_index) + return str(getattr(properties, "gcnArchName", "")).partition(":")[0] + + +@functools.lru_cache(maxsize=None) +def _tree_leaf_config( + device: torch.device, + m_size: int, + k_size: int, + n_size: int, + *, + transpose_output: bool, + preserve_a_strides: bool, +) -> _TreeLeafConfig: + device_index = device.index if device.index is not None else torch.cuda.current_device() + if _device_arch(device_index) != "gfx942": + return _DEFAULT_TREE_LEAF_CONFIG + return _gfx942_qwen_tree_leaf_config( + m_size, + k_size, + n_size, + transpose_output=transpose_output, + preserve_a_strides=preserve_a_strides, + ) + + +@dataclass(frozen=True) +class _TreePlan: + leaf_starts: tuple[int, ...] + leaf_lengths: tuple[int, ...] + leaf_nodes: tuple[int, ...] + reduction_levels: tuple[tuple[tuple[int, int, int], ...], ...] + root: int + node_count: int + + +@dataclass(frozen=True) +class _DeviceTreePlan: + host: _TreePlan + leaf_starts: torch.Tensor + leaf_lengths: torch.Tensor + leaf_nodes: torch.Tensor + reduction_levels: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...] + + +def _build_tree_plan(k_size: int) -> _TreePlan: + if k_size <= 0: + raise ValueError(f"deterministic GEMM K must be positive, got {k_size}") + + leaf_starts: list[int] = [] + leaf_lengths: list[int] = [] + leaf_nodes: list[int] = [] + reductions_by_height: dict[int, list[tuple[int, int, int]]] = {} + next_node = 0 + + def visit(begin: int, end: int) -> tuple[int, int]: + nonlocal next_node + if end - begin <= _BLOCK_K: + node = next_node + next_node += 1 + leaf_starts.append(begin) + leaf_lengths.append(end - begin) + leaf_nodes.append(node) + return node, 0 + + midpoint = begin + (end - begin) // 2 + lower, lower_height = visit(begin, midpoint) + upper, upper_height = visit(midpoint, end) + node = next_node + next_node += 1 + height = max(lower_height, upper_height) + 1 + reductions_by_height.setdefault(height, []).append((lower, upper, node)) + return node, height + + root, max_height = visit(0, k_size) + reduction_levels = tuple( + tuple(reductions_by_height.get(height, ())) + for height in range(1, max_height + 1) + ) + return _TreePlan( + leaf_starts=tuple(leaf_starts), + leaf_lengths=tuple(leaf_lengths), + leaf_nodes=tuple(leaf_nodes), + reduction_levels=reduction_levels, + root=root, + node_count=next_node, + ) + + +def _device_tree_plan(k_size: int, device: torch.device) -> _DeviceTreePlan: + device_index = device.index if device.index is not None else torch.cuda.current_device() + key = (device_index, k_size) + with _TREE_PLAN_LOCK: + cached = _TREE_PLANS.get(key) + if cached is not None: + return cached + + host = _build_tree_plan(k_size) + + def indices(values: tuple[int, ...]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int32, device=device) + + levels = [] + for operations in host.reduction_levels: + lower, upper, output = zip(*operations, strict=True) + levels.append((indices(lower), indices(upper), indices(output))) + result = _DeviceTreePlan( + host=host, + leaf_starts=indices(host.leaf_starts), + leaf_lengths=indices(host.leaf_lengths), + leaf_nodes=indices(host.leaf_nodes), + reduction_levels=tuple(levels), + ) + _TREE_PLANS[key] = result + return result if _TRITON_AVAILABLE: @triton.jit - def _det_gemm_kernel( + def _det_gemm_tree_leaf_kernel( + a_ptr, + b_ptr, + workspace_ptr, + leaf_starts_ptr, + leaf_lengths_ptr, + leaf_nodes_ptr, + M: tl.constexpr, + N: tl.constexpr, + K: tl.constexpr, + stride_am: tl.constexpr, + stride_ak: tl.constexpr, + stride_bk: tl.constexpr, + stride_bn: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + N_FASTEST: tl.constexpr, + ): + if N_FASTEST: + pid_n = tl.program_id(0) + pid_m = tl.program_id(1) + leaf = tl.program_id(2) + else: + leaf = tl.program_id(0) + pid_m = tl.program_id(1) + pid_n = tl.program_id(2) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + leaf_start = tl.load(leaf_starts_ptr + leaf) + leaf_length = tl.load(leaf_lengths_ptr + leaf) + leaf_node = tl.load(leaf_nodes_ptr + leaf) + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + # Keep the leaf's ascending scalar FMA order identical to the native + # gfx942 correctness kernel. A tl.dot/MFMA leaf is topology-stable but + # differs from the scalar reference at rare BF16 rounding boundaries. + for offset in tl.static_range(0, BLOCK_K): + k_offset = leaf_start + offset + active = offset < leaf_length + a = tl.load( + a_ptr + offs_m * stride_am + k_offset * stride_ak, + mask=(offs_m < M) & active, + other=0.0, + ).to(tl.float32) + b = tl.load( + b_ptr + k_offset * stride_bk + offs_n * stride_bn, + mask=(offs_n < N) & active, + other=0.0, + ).to(tl.float32) + acc += a[:, None] * b[None, :] + output_offsets = leaf_node * M * N + offs_m[:, None] * N + offs_n[None, :] + output_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + tl.store( + workspace_ptr + output_offsets, + acc.to(workspace_ptr.dtype.element_ty), + mask=output_mask, + ) + + @triton.jit + def _det_gemm_tree_reduce_kernel( + workspace_ptr, + lower_nodes_ptr, + upper_nodes_ptr, + output_nodes_ptr, + M: tl.constexpr, + N: tl.constexpr, + BLOCK: tl.constexpr, + ): + operation = tl.program_id(0) + block = tl.program_id(1) + offsets = block * BLOCK + tl.arange(0, BLOCK) + elements = M * N + mask = offsets < elements + lower_node = tl.load(lower_nodes_ptr + operation) + upper_node = tl.load(upper_nodes_ptr + operation) + output_node = tl.load(output_nodes_ptr + operation) + lower = tl.load( + workspace_ptr + lower_node * elements + offsets, + mask=mask, + other=0.0, + ).to(tl.float32) + upper = tl.load( + workspace_ptr + upper_node * elements + offsets, + mask=mask, + other=0.0, + ).to(tl.float32) + result = lower + upper + tl.store( + workspace_ptr + output_node * elements + offsets, + result.to(workspace_ptr.dtype.element_ty), + mask=mask, + ) + + @triton.jit + def _copy_tree_root_kernel( + workspace_ptr, + output_ptr, + root, + elements, + BLOCK: tl.constexpr, + ): + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < elements + values = tl.load(workspace_ptr + root * elements + offsets, mask=mask) + tl.store(output_ptr + offsets, values, mask=mask) + + @triton.jit + def _copy_tree_root_transposed_kernel( + workspace_ptr, + output_ptr, + root, + M: tl.constexpr, + N: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + ): + offsets_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) + offsets_n = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) + elements = M * N + mask = (offsets_m[:, None] < M) & (offsets_n[None, :] < N) + values = tl.load( + workspace_ptr + + root * elements + + offsets_m[:, None] * N + + offsets_n[None, :], + mask=mask, + ) + # Store the already-rounded BF16 root directly in [N, M] layout. + # This is a pure address permutation: the canonical GEMM leaves, tree, + # operand order, and every rounding boundary remain unchanged. + tl.store( + output_ptr + offsets_n[:, None] * M + offsets_m[None, :], + tl.trans(values), + mask=tl.trans(mask), + ) + + @triton.jit + def _det_gemm_fp32_kernel( a_ptr, b_ptr, c_ptr, @@ -72,18 +430,13 @@ def _det_gemm_kernel( tl.store(c_ptrs, c, mask=mask) -def _triton_gemm(a, b, *, output_dtype=None): +def _triton_gemm_fp32(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: a, b = a.contiguous(), b.contiguous() M, K = a.shape _, N = b.shape - c = torch.empty( - (M, N), device=a.device, dtype=a.dtype if output_dtype is None else output_dtype - ) + c = torch.empty((M, N), device=a.device, dtype=torch.float32) grid = (triton.cdiv(M, _BLOCK_M), triton.cdiv(N, _BLOCK_N)) - promote_inputs = ( - output_dtype == torch.float32 or a.dtype == torch.float32 or b.dtype == torch.float32 - ) - _det_gemm_kernel[grid]( + _det_gemm_fp32_kernel[grid]( a, b, c, @@ -99,31 +452,165 @@ def _triton_gemm(a, b, *, output_dtype=None): BLOCK_M=_BLOCK_M, BLOCK_N=_BLOCK_N, BLOCK_K=_BLOCK_K, - PROMOTE_INPUTS=promote_inputs, + PROMOTE_INPUTS=True, ) return c +def _triton_tree_gemm( + a: torch.Tensor, + b: torch.Tensor, + *, + transpose_output: bool = False, + out: torch.Tensor | None = None, + preserve_a_strides: bool = False, +) -> torch.Tensor: + if not _TRITON_AVAILABLE: + raise RuntimeError("Triton is unavailable") + if a.dim() != 2 or b.dim() != 2: + raise ValueError("Triton deterministic GEMM expects two 2-D tensors") + if a.size(1) != b.size(0): + raise ValueError(f"Triton deterministic GEMM K mismatch: {a.size(1)} and {b.size(0)}") + if a.dtype != torch.bfloat16 or b.dtype != torch.bfloat16: + raise TypeError("Triton tree GEMM requires BF16 inputs") + if not a.is_cuda or not b.is_cuda or a.device != b.device: + raise RuntimeError("Triton tree GEMM inputs must share one CUDA/ROCm device") + + # The leaf kernel accepts positive arbitrary strides. Wgrad uses this to + # consume an activation transpose view without materializing another copy. + # Other callers retain the established contiguous-input behavior. + if not preserve_a_strides: + a = a.contiguous() + b = b.contiguous() + m_size, k_size = a.shape + n_size = b.size(1) + result_shape = (n_size, m_size) if transpose_output else (m_size, n_size) + if out is None: + result = torch.empty(result_shape, dtype=torch.bfloat16, device=a.device) + else: + if tuple(out.shape) != result_shape: + raise ValueError( + f"Triton tree GEMM output must have shape {result_shape}, " + f"got {tuple(out.shape)}" + ) + if out.dtype != torch.bfloat16: + raise TypeError(f"Triton tree GEMM output must be BF16, got {out.dtype}") + if out.device != a.device: + raise RuntimeError( + f"Triton tree GEMM output must be on {a.device}, got {out.device}" + ) + if not out.is_contiguous(): + raise ValueError("Triton tree GEMM output buffer must be contiguous") + if out.requires_grad: + raise ValueError("Triton tree GEMM output buffer must not require gradients") + result = out + plan = _device_tree_plan(k_size, a.device) + workspace = torch.empty( + (plan.host.node_count, m_size, n_size), + dtype=torch.bfloat16, + device=a.device, + ) + leaf_config = _tree_leaf_config( + a.device, + m_size, + k_size, + n_size, + transpose_output=transpose_output, + preserve_a_strides=preserve_a_strides, + ) + tiles_m = triton.cdiv(m_size, leaf_config.block_m) + tiles_n = triton.cdiv(n_size, leaf_config.block_n) + leaf_grid = ( + (tiles_n, tiles_m, len(plan.host.leaf_nodes)) + if leaf_config.n_fastest + else (len(plan.host.leaf_nodes), tiles_m, tiles_n) + ) + _det_gemm_tree_leaf_kernel[leaf_grid]( + a, + b, + workspace, + plan.leaf_starts, + plan.leaf_lengths, + plan.leaf_nodes, + M=m_size, + N=n_size, + K=k_size, + stride_am=a.stride(0), + stride_ak=a.stride(1), + stride_bk=b.stride(0), + stride_bn=b.stride(1), + BLOCK_M=leaf_config.block_m, + BLOCK_N=leaf_config.block_n, + BLOCK_K=_BLOCK_K, + N_FASTEST=leaf_config.n_fastest, + num_warps=leaf_config.num_warps, + ) + reduction_block = 256 + for operations, (lower, upper, output) in zip( + plan.host.reduction_levels, + plan.reduction_levels, + strict=True, + ): + grid = (len(operations), triton.cdiv(m_size * n_size, reduction_block)) + _det_gemm_tree_reduce_kernel[grid]( + workspace, + lower, + upper, + output, + M=m_size, + N=n_size, + BLOCK=reduction_block, + ) + + copy_block = 256 + if transpose_output: + transpose_block = 32 + transpose_grid = ( + triton.cdiv(m_size, transpose_block), + triton.cdiv(n_size, transpose_block), + ) + _copy_tree_root_transposed_kernel[transpose_grid]( + workspace, + result, + plan.host.root, + M=m_size, + N=n_size, + BLOCK_M=transpose_block, + BLOCK_N=transpose_block, + ) + else: + _copy_tree_root_kernel[(triton.cdiv(result.numel(), copy_block),)]( + workspace, + result, + plan.host.root, + result.numel(), + BLOCK=copy_block, + ) + if out is not None: + # Triton mutates caller-owned storage outside PyTorch's dispatcher. + # Keep saved-tensor and cache version checks semantically correct. + torch.autograd.graph.increment_version(result) + return result + + class _TritonDetGemmFn(torch.autograd.Function): @staticmethod def forward(ctx, a, b, output_fp32=False): ctx.save_for_backward(a, b) ctx.output_fp32 = bool(output_fp32) - return _triton_gemm(a, b, output_dtype=torch.float32 if output_fp32 else None) + return _triton_gemm_fp32(a, b) if output_fp32 else _triton_tree_gemm(a, b) @staticmethod def backward(ctx, grad_out): a, b = ctx.saved_tensors grad_out = grad_out.contiguous() da = ( - _triton_gemm(grad_out, b.t().contiguous(), output_dtype=torch.float32) - .reshape_as(a) - .to(a.dtype) + _triton_tree_gemm(grad_out.to(torch.bfloat16), b.t().contiguous()).reshape_as(a) if ctx.needs_input_grad[0] else None ) db = ( - _triton_gemm(a.t().contiguous(), grad_out, output_dtype=torch.float32).to(b.dtype) + _triton_tree_gemm(a.t().contiguous(), grad_out.to(torch.bfloat16)) if ctx.needs_input_grad[1] else None ) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index f715bad8..03bb9ca9 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -1046,23 +1046,35 @@ def get_attention_op( platform = self._platform() candidates = self._priority_map.get(platform, {}).get("ws2_attention", []) rejected: list[str] = [] + # A candidate skipped only because it does not satisfy the caller's + # explicit backend/policy request is not a fallback. Count only an + # otherwise eligible candidate that failed capability or loading. + capability_rejections = 0 for backend in candidates: capability = self._attention_capabilities.get(backend) if capability is None: rejected.append(f"{backend.name}: no AttentionBackendCapability declared") + capability_rejections += 1 continue - incompatibilities = list(capability.incompatibilities(contract)) policy_mismatch = self._attention_policy_mismatch(requested_backend, capability) + incompatibilities = list(capability.incompatibilities(contract)) if policy_mismatch is not None: + # Still report capability details for diagnostics, but this + # candidate was excluded by the caller's policy, so those + # details must not turn the selected backend into a fallback. incompatibilities.append(policy_mismatch) + rejected.append(f"{backend.name}: " + "; ".join(incompatibilities)) + continue if incompatibilities: rejected.append(f"{backend.name}: " + "; ".join(incompatibilities)) + capability_rejections += 1 continue op = self._get_or_create_backend(backend) if op is None: rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + capability_rejections += 1 continue return AttentionDispatchResult( @@ -1073,7 +1085,7 @@ def get_attention_op( "actual_backend": capability.backend_id, "backend_enum": backend.name, "platform": platform, - "fallback": bool(rejected), + "fallback": capability_rejections > 0, "prior_rejections": list(rejected), "contract": contract.to_dict(), "capability": capability.to_dict(), diff --git a/setup.py b/setup.py index 94c0062a..ef1ab10d 100644 --- a/setup.py +++ b/setup.py @@ -226,7 +226,7 @@ def get_extensions(): 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": + if os.name != "nt" and not is_rocm: # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). extra_link_args.append("-lcuda") diff --git a/tests/distributed/test_qwen_ffn_topology.py b/tests/distributed/test_qwen_ffn_topology.py new file mode 100644 index 00000000..5cb72a09 --- /dev/null +++ b/tests/distributed/test_qwen_ffn_topology.py @@ -0,0 +1,497 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Real multi-GPU topology checks for the ROCm-native Triton Qwen3 FFN. + +PyTorch exposes RCCL through the ``nccl`` process-group backend. The FFN uses +ROCm-native Triton compute and a rank-ordered RCCL transport collective. +""" + +from __future__ import annotations + +import os +import queue +import tempfile +import traceback +from datetime import timedelta +from pathlib import Path +from typing import Any + +import pytest +import torch +import torch.multiprocessing as mp + +import rl_engine.kernels.ops.triton.ffn.ffn as ffn_module +from rl_engine.kernels.ops.triton.ffn import ( + QWEN3_8B_HIDDEN_SIZE, + QWEN3_8B_INTERMEDIATE_SIZE, + pack_qwen3_ffn_forward_weights, + qwen3_ffn, +) +_IS_ROCM = getattr(torch.version, "hip", None) is not None +_EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + +# I=512 keeps every TP=1/2/4/8 shard aligned to the 32-wide GEMM K-tree and +# keeps the SM90 N dimension tile-aligned even at TP=8. +_TOKENS = 32 +_HIDDEN = 64 +_INTERMEDIATE = 512 + +_WORLD2_CONFIGS = ( + ("tp2", 2, 1, False), + ("tp2_sp", 2, 1, True), +) +_WORLD4_CONFIGS = ( + ("tp4", 4, 1, False), + ("tp2_cp2", 2, 2, False), + ("tp2_cp2_sp", 2, 2, True), +) +_WORLD8_CONFIGS = (("tp8", 8, 1, False),) + +pytestmark = pytest.mark.skipif( + _EXTERNAL_WORLD_SIZE != 1, + reason="topology tests own their local worker processes; run pytest directly", +) + + +def _has_topology_devices(count: int) -> bool: + return bool( + _IS_ROCM + and torch.cuda.is_available() + and torch.distributed.is_available() + and torch.distributed.is_nccl_available() + and torch.cuda.device_count() >= count + ) + + +def _has_qwen3_8b_capacity(count: int) -> bool: + if not _has_topology_devices(count): + return False + minimum_bytes = 8 * 1024**3 + try: + return all( + torch.cuda.get_device_properties(index).total_memory >= minimum_bytes + for index in range(count) + ) + except RuntimeError: + return False + + +def _randn( + shape: tuple[int, ...], + *, + seed: int, + device: torch.device, +) -> 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) + + +def _make_inputs( + token_count: int, + hidden_size: int, + intermediate_size: int, + device: torch.device, + *, + seed: int, +) -> tuple[torch.Tensor, ...]: + return ( + _randn((token_count, hidden_size), seed=seed, device=device), + _randn((intermediate_size, hidden_size), seed=seed + 1, device=device), + _randn((intermediate_size, hidden_size), seed=seed + 2, device=device), + _randn((hidden_size, intermediate_size), seed=seed + 3, device=device), + _randn((token_count, hidden_size), seed=seed + 4, device=device), + ) + + +def _close_ffn_collectives() -> None: + for collective in list(ffn_module._COLLECTIVES.values()): + collective.close() + ffn_module._COLLECTIVES.clear() + + +def _shard_ranges( + rank: int, + *, + tp_size: int, + cp_size: int, + sequence_parallel: bool, + token_count: int, + intermediate_size: int, +) -> tuple[int, int, int, int]: + if tp_size * cp_size <= rank: + raise ValueError("rank lies outside the requested TP/CP mesh") + if token_count % cp_size: + raise ValueError("token count must be divisible by CP size") + cp_tokens = token_count // cp_size + if sequence_parallel and cp_tokens % tp_size: + raise ValueError("each CP token shard must be divisible by TP size for SP") + if intermediate_size % tp_size: + raise ValueError("intermediate size must be divisible by TP size") + + tp_rank = rank % tp_size + cp_rank = rank // tp_size + local_tokens = cp_tokens // tp_size if sequence_parallel else cp_tokens + token_start = cp_rank * cp_tokens + if sequence_parallel: + token_start += tp_rank * local_tokens + token_end = token_start + local_tokens + + local_intermediate = intermediate_size // tp_size + feature_start = tp_rank * local_intermediate + feature_end = feature_start + local_intermediate + return token_start, token_end, feature_start, feature_end + + +def _canonical( + hidden: torch.Tensor, + gate: torch.Tensor, + up: torch.Tensor, + down: torch.Tensor, + grad_output: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]: + with torch.no_grad(): + inference = qwen3_ffn(hidden, gate, up, down) + inputs = [value.detach().clone().requires_grad_(True) for value in (hidden, gate, up, down)] + training = qwen3_ffn(*inputs) + training.backward(grad_output) + assert torch.equal(inference, training.detach()), "TP=1 train/infer forward mismatch" + return inference, training, inputs + + +def _mesh_groups( + dist: Any, + tp_size: int, + cp_size: int, +) -> tuple[list[Any], list[Any]]: + world_size = dist.get_world_size() + if tp_size * cp_size != world_size: + raise ValueError("TP size times CP size must equal the process-group world size") + if tp_size == world_size and cp_size == 1: + return [dist.group.WORLD], [] + if cp_size == world_size and tp_size == 1: + return [], [dist.group.WORLD] + + tp_groups = [] + if tp_size > 1: + for cp_rank in range(cp_size): + ranks = list(range(cp_rank * tp_size, (cp_rank + 1) * tp_size)) + tp_groups.append(dist.new_group(ranks=ranks)) + + cp_groups = [] + if cp_size > 1: + for tp_rank in range(tp_size): + ranks = [cp_rank * tp_size + tp_rank for cp_rank in range(cp_size)] + cp_groups.append(dist.new_group(ranks=ranks)) + return tp_groups, cp_groups + + +def _run_topology( + rank: int, + dist: Any, + meshes: dict[tuple[int, int], tuple[list[Any], list[Any]]], + *, + name: str, + tp_size: int, + cp_size: int, + sequence_parallel: bool, + hidden: torch.Tensor, + gate: torch.Tensor, + up: torch.Tensor, + down: torch.Tensor, + grad_output: torch.Tensor, + inference_reference: torch.Tensor, + training_reference: torch.Tensor, + reference_inputs: list[torch.Tensor], +) -> None: + mesh_key = (tp_size, cp_size) + if mesh_key not in meshes: + meshes[mesh_key] = _mesh_groups(dist, tp_size, cp_size) + tp_groups, cp_groups = meshes[mesh_key] + + tp_rank = rank % tp_size + cp_rank = rank // tp_size + tp_group = tp_groups[cp_rank] if tp_size > 1 else None + cp_group = cp_groups[tp_rank] if cp_size > 1 else None + token_start, token_end, feature_start, feature_end = _shard_ranges( + rank, + tp_size=tp_size, + cp_size=cp_size, + sequence_parallel=sequence_parallel, + token_count=hidden.size(0), + intermediate_size=gate.size(0), + ) + shard = ( + hidden[token_start:token_end].contiguous(), + gate[feature_start:feature_end].contiguous(), + up[feature_start:feature_end].contiguous(), + down[:, feature_start:feature_end].contiguous(), + ) + inference_forward_weights = pack_qwen3_ffn_forward_weights(*shard[1:]) + + with torch.no_grad(): + inference = qwen3_ffn( + *shard, + forward_weights=inference_forward_weights, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + inputs = [value.detach().clone().requires_grad_(True) for value in shard] + training_forward_weights = pack_qwen3_ffn_forward_weights(*inputs[1:]) + training = qwen3_ffn( + *inputs, + forward_weights=training_forward_weights, + tp_group=tp_group, + cp_group=cp_group, + sequence_parallel=sequence_parallel, + ) + training.backward(grad_output[token_start:token_end].contiguous()) + + expected_output = inference_reference[token_start:token_end] + assert torch.equal(inference, training.detach()), f"{name}: train/infer mismatch" + assert torch.equal(inference, expected_output), f"{name}: inference mismatch vs TP=1" + assert torch.equal(training.detach(), training_reference.detach()[token_start:token_end]), ( + f"{name}: training forward mismatch vs TP=1" + ) + assert torch.equal(inputs[0].grad, reference_inputs[0].grad[token_start:token_end]), ( + f"{name}: hidden grad mismatch vs TP=1" + ) + + expected_weight_grads = ( + (inputs[1].grad, reference_inputs[1].grad[feature_start:feature_end], "gate"), + (inputs[2].grad, reference_inputs[2].grad[feature_start:feature_end], "up"), + ( + inputs[3].grad, + reference_inputs[3].grad[:, feature_start:feature_end], + "down", + ), + ) + for actual, expected, label in expected_weight_grads: + assert torch.equal(actual, expected), f"{name}: {label} weight grad mismatch vs TP=1" + + +def _topology_worker( + rank: int, + world_size: int, + init_method: str, + result_queue: Any, + configs: tuple[tuple[str, int, int, bool], ...], +) -> None: + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(minutes=5), + ) + device = torch.device("cuda", rank) + hidden, gate, up, down, grad_output = _make_inputs( + _TOKENS, + _HIDDEN, + _INTERMEDIATE, + device, + seed=600, + ) + inference_reference, training_reference, reference_inputs = _canonical( + hidden, + gate, + up, + down, + grad_output, + ) + meshes: dict[tuple[int, int], tuple[list[Any], list[Any]]] = {} + for name, tp_size, cp_size, sequence_parallel in configs: + _run_topology( + rank, + dist, + meshes, + name=name, + tp_size=tp_size, + cp_size=cp_size, + sequence_parallel=sequence_parallel, + hidden=hidden, + gate=gate, + up=up, + down=down, + grad_output=grad_output, + inference_reference=inference_reference, + training_reference=training_reference, + reference_inputs=reference_inputs, + ) + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _qwen3_8b_tp2_worker( + rank: int, + world_size: int, + init_method: str, + result_queue: Any, +) -> None: + try: + import torch.distributed as dist + + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=init_method, + rank=rank, + world_size=world_size, + timeout=timedelta(minutes=10), + ) + device = torch.device("cuda", rank) + hidden, gate, up, down, grad_output = _make_inputs( + 2, + QWEN3_8B_HIDDEN_SIZE, + QWEN3_8B_INTERMEDIATE_SIZE, + device, + seed=800, + ) + inference_reference, training_reference, reference_inputs = _canonical( + hidden, + gate, + up, + down, + grad_output, + ) + _run_topology( + rank, + dist, + {}, + name="qwen3_8b_tp2", + tp_size=2, + cp_size=1, + sequence_parallel=False, + hidden=hidden, + gate=gate, + up=up, + down=down, + grad_output=grad_output, + inference_reference=inference_reference, + training_reference=training_reference, + reference_inputs=reference_inputs, + ) + result_queue.put({"ok": True, "rank": rank}) + except Exception: # pragma: no cover - forwarded to the parent process. + result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) + raise + finally: + _close_ffn_collectives() + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +def _spawn_workers( + worker: Any, + world_size: int, + worker_args: tuple[Any, ...] = (), + *, + timeout_seconds: int, +) -> None: + if not _has_topology_devices(world_size): + platform = "ROCm/RCCL" if _IS_ROCM else "CUDA/NCCL" + pytest.skip( + f"requires {world_size} {platform} GPUs plus FFN and deterministic collective support" + ) + + # Single-GPU tests may have populated the parent process's caching + # allocator on device 0. Release unused blocks before spawned rank 0 owns + # that device, which matters for the Qwen3-8B smoke case. + torch.cuda.empty_cache() + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as temporary_directory: + init_method = (Path(temporary_directory) / "nccl_init").as_uri() + result_queue = ctx.Queue() + processes = [ + ctx.Process( + target=worker, + args=(rank, world_size, init_method, result_queue, *worker_args), + ) + for rank in range(world_size) + ] + for process in processes: + process.start() + + results = [] + try: + for _ in processes: + result = result_queue.get(timeout=timeout_seconds) + results.append(result) + if not result["ok"]: + for process in processes: + if process.is_alive(): + process.terminate() + break + except queue.Empty: + for process in processes: + if process.is_alive(): + process.terminate() + pytest.fail(f"timed out waiting for {world_size} FFN topology workers") + finally: + for process in processes: + # RCCL teardown can take longer than ten seconds after every + # worker has already reported a successful result, especially + # for the eight-rank topology on ROCm hosts with unusable RDMA + # interfaces. Allow cleanup to finish instead of turning a + # successful numerical check into a SIGTERM false failure. + process.join(timeout=30) + if process.is_alive(): + process.terminate() + process.join(timeout=30) + result_queue.close() + result_queue.join_thread() + + for result in sorted(results, key=lambda item: item["rank"]): + assert result["ok"], result.get("traceback") + for process in processes: + assert process.exitcode == 0 + + +def test_triton_qwen3_ffn_tp2_and_tp_sp_match_tp1_bitwise() -> None: + _spawn_workers( + _topology_worker, + 2, + (_WORLD2_CONFIGS,), + timeout_seconds=180, + ) + + +def test_triton_qwen3_ffn_tp4_tp_cp_and_tp_cp_sp_match_tp1_bitwise() -> None: + _spawn_workers( + _topology_worker, + 4, + (_WORLD4_CONFIGS,), + timeout_seconds=240, + ) + + +def test_triton_qwen3_ffn_tp8_matches_tp1_bitwise() -> None: + _spawn_workers( + _topology_worker, + 8, + (_WORLD8_CONFIGS,), + timeout_seconds=360, + ) + + +def test_triton_qwen3_8b_ffn_tp2_smoke_matches_tp1_bitwise() -> None: + if os.environ.get("RL_KERNEL_SKIP_QWEN3_8B_TOPOLOGY") == "1": + pytest.skip("Qwen3-8B topology smoke disabled by environment") + if not _has_qwen3_8b_capacity(2): + pytest.skip("Qwen3-8B TP=2 smoke requires two GPUs with at least 8 GiB each") + _spawn_workers( + _qwen3_8b_tp2_worker, + 2, + timeout_seconds=600, + ) diff --git a/tests/test_det_gemm.py b/tests/test_det_gemm.py index 633009ef..c1d0992f 100644 --- a/tests/test_det_gemm.py +++ b/tests/test_det_gemm.py @@ -2,11 +2,12 @@ # Copyright (c) 2026 RL-Kernel Contributors """Invariance + correctness tests for det_gemm (WS1). -Runs against both deterministic backends — the hand-written CUDA kernel and the -Triton path — each of which must independently satisfy the invariance contract. +Runs the ROCm-native Triton path directly. CUDA continues to exercise both its +existing native kernel and Triton, each independently satisfying the contract. The PyTorch path (torch.matmul) is intentionally NOT tested here: it is the non-deterministic reference baseline and would fail batch-invariance by design. """ + import pytest import torch @@ -15,7 +16,20 @@ from rl_engine.kernels.ops.cuda.matmul import deterministic_gemm try: + import triton + from rl_engine.kernels.ops.triton.matmul import deterministic_gemm_triton + from rl_engine.kernels.ops.triton.matmul.det_gemm import ( + _DEFAULT_TREE_LEAF_CONFIG, + _GFX942_QWEN_FORWARD_LEAF_CONFIGS, + _GFX942_QWEN_TP_SHARD_FORWARD_LEAF_CONFIGS, + _GFX942_QWEN_TP_SHARD_WGRAD_LEAF_CONFIGS, + _GFX942_QWEN_WGRAD_LEAF_CONFIGS, + _det_gemm_tree_leaf_kernel, + _device_tree_plan, + _gfx942_qwen_tree_leaf_config, + _triton_tree_gemm, + ) _HAS_TRITON = True except ImportError: @@ -23,14 +37,21 @@ torch.backends.cuda.matmul.allow_tf32 = False DEV = "cuda" +IS_ROCM = getattr(torch.version, "hip", None) is not None +HAS_SUPPORTED_GPU = torch.cuda.is_available() and ( + IS_ROCM or torch.cuda.get_device_capability()[0] >= 8 +) +IS_GFX942 = IS_ROCM and torch.cuda.is_available() and str( + getattr(torch.cuda.get_device_properties(0), "gcnArchName", "") +).startswith("gfx942") pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, - reason="det_gemm requires CUDA SM80+", + not HAS_SUPPORTED_GPU, + reason="det_gemm requires a ROCm GPU or CUDA SM80+", ) -# Each deterministic backend is validated independently. -_BACKENDS = [("cuda", deterministic_gemm)] +# ROCm acceptance intentionally depends only on the Triton implementation. +_BACKENDS = [] if IS_ROCM else [("cuda", deterministic_gemm)] if _HAS_TRITON: _BACKENDS.append(("triton", deterministic_gemm_triton)) @@ -39,22 +60,239 @@ def _rand(*shape): return torch.randn(*shape, device=DEV, dtype=torch.bfloat16) -# Matches det_gemm_kernel.cu: mid-split K-tree, FP32 leaf width 32, BF16 internal adds. +def _assert_same_raw_bytes(actual: torch.Tensor, expected: torch.Tensor) -> None: + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype == torch.bfloat16 + assert actual.is_contiguous() + assert expected.is_contiguous() + assert torch.equal( + actual.reshape(-1).view(torch.uint8), + expected.reshape(-1).view(torch.uint8), + ) + + +def _leaf_workspace( + a: torch.Tensor, + b: torch.Tensor, + config, +) -> torch.Tensor: + m_size, k_size = a.shape + n_size = b.size(1) + plan = _device_tree_plan(k_size, a.device) + workspace = torch.empty( + (plan.host.node_count, m_size, n_size), + dtype=torch.bfloat16, + device=a.device, + ) + tiles_m = triton.cdiv(m_size, config.block_m) + tiles_n = triton.cdiv(n_size, config.block_n) + grid = ( + (tiles_n, tiles_m, len(plan.host.leaf_nodes)) + if config.n_fastest + else (len(plan.host.leaf_nodes), tiles_m, tiles_n) + ) + _det_gemm_tree_leaf_kernel[grid]( + a, + b, + workspace, + plan.leaf_starts, + plan.leaf_lengths, + plan.leaf_nodes, + M=m_size, + N=n_size, + K=k_size, + stride_am=a.stride(0), + stride_ak=a.stride(1), + stride_bk=b.stride(0), + stride_bn=b.stride(1), + BLOCK_M=config.block_m, + BLOCK_N=config.block_n, + BLOCK_K=_K_TREE_LEAF, + N_FASTEST=config.n_fastest, + num_warps=config.num_warps, + ) + return workspace.index_select(0, plan.leaf_nodes.to(torch.int64)) + + +def _special_bf16(shape: tuple[int, ...], *, offset: int = 0) -> torch.Tensor: + bits = torch.tensor( + ( + 0x0000, + 0x8000, + 0x0001, + 0x8001, + 0x007F, + 0x807F, + 0x0080, + 0x8080, + 0x3F80, + 0xBF80, + 0x7F7F, + 0xFF7F, + 0x7F80, + 0xFF80, + 0x7F81, + 0x7FC1, + 0xFF81, + 0xFFFF, + ), + dtype=torch.uint16, + ) + elements = 1 + for size in shape: + elements *= size + indices = (torch.arange(elements, dtype=torch.int64) + offset) % bits.numel() + return bits[indices].view(torch.bfloat16).reshape(shape).to(DEV) + + _K_TREE_LEAF = 32 def _k_tree_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Canonical FP32-leaf/BF16-node midpoint tree used by Triton.""" + a = a.detach().contiguous() b = b.detach().contiguous() - k = a.shape[1] - def rec(lo: int, hi: int) -> torch.Tensor: + def reduce_range(lo: int, hi: int) -> torch.Tensor: if hi - lo <= _K_TREE_LEAF: - return (a[:, lo:hi].float() @ b[lo:hi, :].float()).to(dtype=torch.bfloat16) - mid = lo + (hi - lo) // 2 - return rec(lo, mid) + rec(mid, hi) + return (a[:, lo:hi].float() @ b[lo:hi, :].float()).to(torch.bfloat16) + midpoint = lo + (hi - lo) // 2 + return reduce_range(lo, midpoint) + reduce_range(midpoint, hi) + + return reduce_range(0, a.size(1)) - return rec(0, k) + +def _balanced_tree_sum(parts: list[torch.Tensor]) -> torch.Tensor: + level = parts + while len(level) > 1: + level = [level[index] + level[index + 1] for index in range(0, len(level), 2)] + return level[0] + + +@pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") +@pytest.mark.parametrize( + ("shape", "transpose_output", "preserve_a_strides", "expected"), + ( + ((1, 4096, 12288), False, False, (1, 128, 1, True)), + ((8, 4096, 12288), False, False, (8, 128, 2, True)), + ((32, 12288, 4096), False, False, (32, 128, 4, True)), + ((4096, 1, 12288), True, True, (128, 64, 2, True)), + ((12288, 8, 4096), True, True, (128, 64, 2, True)), + ((4096, 32, 12288), True, True, (64, 64, 2, True)), + ((16, 4096, 12288), False, False, (64, 64, 4, False)), + ((32, 4096, 4096), False, False, (64, 64, 4, False)), + ((4096, 16, 12288), True, True, (64, 64, 4, False)), + ((4096, 8, 12288), True, False, (64, 64, 4, False)), + ((32, 4096, 6144), False, False, (32, 128, 4, True)), + ((32, 6144, 4096), False, False, (32, 128, 4, True)), + ((16, 4096, 6144), False, False, (8, 64, 1, True)), + ((16, 6144, 4096), False, False, (8, 64, 1, True)), + ((32, 4096, 3072), False, False, (64, 64, 4, False)), + ((32, 3072, 4096), False, False, (64, 64, 4, False)), + ((32, 4096, 1536), False, False, (64, 64, 4, False)), + ((32, 1536, 4096), False, False, (64, 64, 4, False)), + ((4096, 32, 6144), True, True, (128, 64, 2, True)), + ((6144, 32, 4096), True, True, (128, 64, 2, True)), + ((4096, 32, 3072), True, True, (64, 64, 4, False)), + ((3072, 32, 4096), True, True, (64, 64, 4, False)), + ((4096, 16, 6144), True, True, (64, 64, 4, False)), + ), +) +def test_gfx942_qwen_leaf_config_table_is_exact( + shape, + transpose_output, + preserve_a_strides, + expected, +): + config = _gfx942_qwen_tree_leaf_config( + *shape, + transpose_output=transpose_output, + preserve_a_strides=preserve_a_strides, + ) + assert ( + config.block_m, + config.block_n, + config.num_warps, + config.n_fastest, + ) == expected + + +@pytest.mark.skipif( + not (_HAS_TRITON and IS_GFX942), + reason="leaf specialization raw-byte tests require ROCm gfx942", +) +@pytest.mark.parametrize("k_size", (1, 8, 31, 32, 33, 65, 4096, 12288)) +def test_gfx942_leaf_configs_preserve_special_value_workspace_raw_bytes(k_size): + a = _special_bf16((3, k_size)) + b = _special_bf16((k_size, 17), offset=7) + expected = _leaf_workspace(a, b, _DEFAULT_TREE_LEAF_CONFIG) + configs = tuple( + dict.fromkeys( + ( + *_GFX942_QWEN_FORWARD_LEAF_CONFIGS.values(), + *_GFX942_QWEN_TP_SHARD_FORWARD_LEAF_CONFIGS.values(), + *_GFX942_QWEN_TP_SHARD_WGRAD_LEAF_CONFIGS.values(), + *_GFX942_QWEN_WGRAD_LEAF_CONFIGS.values(), + ) + ) + ) + + for config in configs: + _assert_same_raw_bytes(_leaf_workspace(a, b, config), expected) + + +@pytest.mark.skipif( + not (_HAS_TRITON and IS_GFX942), + reason="leaf specialization raw-byte tests require ROCm gfx942", +) +def test_gfx942_leaf_configs_preserve_transposed_a_workspace_raw_bytes(): + source = _special_bf16((65, 3)) + a = source.t() + b = _special_bf16((65, 17), offset=11) + assert not a.is_contiguous() + assert all(stride > 0 for stride in a.stride()) + expected = _leaf_workspace(a, b, _DEFAULT_TREE_LEAF_CONFIG) + configs = tuple( + dict.fromkeys( + ( + *_GFX942_QWEN_FORWARD_LEAF_CONFIGS.values(), + *_GFX942_QWEN_TP_SHARD_FORWARD_LEAF_CONFIGS.values(), + *_GFX942_QWEN_TP_SHARD_WGRAD_LEAF_CONFIGS.values(), + *_GFX942_QWEN_WGRAD_LEAF_CONFIGS.values(), + ) + ) + ) + + for config in configs: + _assert_same_raw_bytes(_leaf_workspace(a, b, config), expected) + + +@pytest.mark.parametrize("tp_size", (2, 4, 8)) +@pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") +def test_forward_matches_balanced_contiguous_k_shards_bitwise(tp_size): + """The GEMM K-tree and the TP collective rank tree must be the same graph.""" + + torch.manual_seed(8) + # Qwen3-8B's down projection uses K=12288. Its 512 24-wide leaves also + # exercise the non-power-of-two midpoint tree used by the SM90 kernel. + m, k, n = 4, 12288, 64 + a, b = _rand(m, k), _rand(k, n) + width = k // tp_size + parts = [ + deterministic_gemm_triton( + a[:, rank * width : (rank + 1) * width].contiguous(), + b[rank * width : (rank + 1) * width].contiguous(), + ) + for rank in range(tp_size) + ] + + full = deterministic_gemm_triton(a, b) + sharded = _balanced_tree_sum(parts) + assert torch.equal(full, sharded), ( + f"full GEMM differed from balanced TP={tp_size} shards at " + f"{int((full != sharded).sum().item())} elements" + ) @pytest.mark.parametrize( @@ -175,10 +413,7 @@ def test_forward_correctness(name, gemm): M, K, N = 128, 2048, 2048 a, b = _rand(M, K), _rand(K, N) out = gemm(a, b).float() - if name == "cuda": - ref = _k_tree_gemm(a, b).float() - else: - ref = a.float() @ b.float() + ref = _k_tree_gemm(a, b).float() contract = load_contract() thresholds = contract["accuracy"]["default"]["reduction"]["bfloat16"] torch.testing.assert_close(out, ref, atol=thresholds["atol"], rtol=thresholds["rtol"]) @@ -208,28 +443,21 @@ def test_backward_correctness(name, gemm): b = _rand(K, N).requires_grad_(True) g = _rand(M, N) gemm(a, b).backward(g) - if name == "cuda": - da = _k_tree_gemm(g, b.t().contiguous()) - db = _k_tree_gemm(a.detach().t().contiguous(), g) - contract = load_contract() - thresholds = contract["accuracy"]["default"]["reduction"]["bfloat16"] - torch.testing.assert_close( - a.grad.float(), da.float(), atol=thresholds["atol"], rtol=thresholds["rtol"] - ) - torch.testing.assert_close( - b.grad.float(), db.float(), atol=thresholds["atol"], rtol=thresholds["rtol"] - ) - return - af = a.detach().float().requires_grad_(True) - bf = b.detach().float().requires_grad_(True) - (af @ bf).backward(g.float()) + expected_da = _k_tree_gemm(g, b.detach().t().contiguous()) + expected_db = _k_tree_gemm(a.detach().t().contiguous(), g) contract = load_contract() thresholds = contract["accuracy"]["default"]["reduction"]["bfloat16"] torch.testing.assert_close( - a.grad.float(), af.grad, atol=thresholds["atol"], rtol=thresholds["rtol"] + a.grad.float(), + expected_da.float(), + atol=thresholds["atol"], + rtol=thresholds["rtol"], ) torch.testing.assert_close( - b.grad.float(), bf.grad, atol=thresholds["atol"], rtol=thresholds["rtol"] + b.grad.float(), + expected_db.float(), + atol=thresholds["atol"], + rtol=thresholds["rtol"], ) @@ -252,9 +480,9 @@ def test_target_shapes_invariance(name, gemm, shape): row = _rand(1, K) big = _rand(64, K) big[0] = row[0] - assert torch.equal( - gemm(row, b)[0], gemm(big, b)[0] - ), f"{name}: batch-invariance broken at shape {shape}" + assert torch.equal(gemm(row, b)[0], gemm(big, b)[0]), ( + f"{name}: batch-invariance broken at shape {shape}" + ) @pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") @@ -266,10 +494,117 @@ def test_triton_ragged_tiles_mask_all_axes(): out = deterministic_gemm_triton(a, b) torch.cuda.synchronize() assert tuple(out.shape) == (80, 129) - torch.testing.assert_close( - out.float(), a.detach().float() @ b.detach().float(), atol=5e-2, rtol=2e-2 - ) + assert torch.equal(out, _k_tree_gemm(a, b)) out.backward(_rand(80, 129)) torch.cuda.synchronize() assert torch.isfinite(a.grad).all() assert torch.isfinite(b.grad).all() + + +@pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") +@pytest.mark.parametrize("shape", ((4, 65, 129), (8, 4096, 128), (4, 12288, 64))) +def test_triton_tree_matches_python_reference_forward_bitwise(shape): + """Triton evaluates the canonical FP32-leaf/BF16-node tree exactly.""" + + torch.manual_seed(47) + m_size, k_size, n_size = shape + a = _rand(m_size, k_size) + b = _rand(k_size, n_size) + + expected = _k_tree_gemm(a, b) + triton_output = deterministic_gemm_triton(a, b) + + assert torch.equal(expected, triton_output), ( + f"Triton/reference mismatch at {shape}: " + f"{int((expected != triton_output).sum().item())} elements" + ) + + +@pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") +def test_triton_tree_matches_python_reference_backward_bitwise(): + torch.manual_seed(48) + a = _rand(32, 512) + b = _rand(512, 128) + grad_output = _rand(32, 128) + triton_inputs = [value.detach().clone().requires_grad_(True) for value in (a, b)] + + deterministic_gemm_triton(*triton_inputs).backward(grad_output) + + expected = ( + _k_tree_gemm(grad_output, b.t().contiguous()), + _k_tree_gemm(a.t().contiguous(), grad_output), + ) + for expected_grad, triton_input in zip(expected, triton_inputs, strict=True): + assert torch.equal(expected_grad, triton_input.grad) + + +@pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") +@pytest.mark.parametrize( + "shape", + ( + (1, 1, 3), + (3, 31, 5), + (17, 32, 33), + (33, 33, 17), + (65, 65, 31), + ), +) +def test_triton_tree_transposed_out_matches_legacy_root_copy_raw_bytes(shape): + """Root placement may transpose addresses, but never the BF16 values.""" + + torch.manual_seed(49) + m_size, k_size, n_size = shape + a = _rand(m_size, k_size) + b = _rand(k_size, n_size) + + legacy = _triton_tree_gemm(a, b).t().contiguous() + output_buffer = torch.empty( + (n_size, m_size), + dtype=torch.bfloat16, + device=a.device, + ) + version_before = output_buffer._version + transposed = _triton_tree_gemm( + a, + b, + transpose_output=True, + out=output_buffer, + ) + + assert transposed is output_buffer + assert output_buffer._version == version_before + 1 + assert transposed.stride() == (m_size, 1) + _assert_same_raw_bytes(transposed, legacy) + + +@pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") +def test_triton_wgrad_reads_positive_stride_transpose_view_raw_bytes(): + """The copy-free a.T wgrad path must preserve the legacy GEMM bit graph.""" + + torch.manual_seed(50) + token_count, input_size, output_size = 33, 65, 17 + activations = _rand(token_count, input_size) + grad_output = _rand(token_count, output_size) + activation_t = activations.t() + assert not activation_t.is_contiguous() + assert all(stride > 0 for stride in activation_t.stride()) + + legacy = _triton_tree_gemm( + activation_t.contiguous(), + grad_output, + ).t().contiguous() + output_buffer = torch.empty( + (output_size, input_size), + dtype=torch.bfloat16, + device=activations.device, + ) + direct = _triton_tree_gemm( + activation_t, + grad_output, + transpose_output=True, + out=output_buffer, + preserve_a_strides=True, + ) + + assert direct is output_buffer + _assert_same_raw_bytes(direct, legacy) diff --git a/tests/test_flashinfer_pr7_attention.py b/tests/test_flashinfer_pr7_attention.py index 03c15455..f7d9286b 100644 --- a/tests/test_flashinfer_pr7_attention.py +++ b/tests/test_flashinfer_pr7_attention.py @@ -2243,8 +2243,8 @@ def production_forward(*args, **kwargs): core.forward_with_lse = production_forward monkeypatch.setattr( paged_attention_module, - "StrictFlashAttention4Core", - lambda *, split_kv: core, + "_resolve_strict_core", + lambda _config: core, ) q, k, v = (tensor.to(torch.bfloat16) for tensor in _qkv(query_len=1)) result = FlashInferQwen3PagedAttentionOp(flashinfer_module=_fake_flashinfer())( diff --git a/tests/test_qwen_ffn.py b/tests/test_qwen_ffn.py index e461c97f..f0dade06 100644 --- a/tests/test_qwen_ffn.py +++ b/tests/test_qwen_ffn.py @@ -1,1041 +1,237 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Tests for the deterministic Qwen3 dense FFN. - -Covers single-GPU correctness, token boundaries, Qwen3-8B shapes, TP/CP/SP -bitwise alignment, and DeterministicCollective cache lifetime. -""" +"""Single-GPU checks for the ROCm-native deterministic Triton Qwen3 FFN.""" from __future__ import annotations -import queue -import tempfile -import traceback -from datetime import timedelta -from pathlib import Path +import inspect import pytest import torch -import torch.multiprocessing as mp import torch.nn.functional as F -import rl_engine.kernels.ops.pytorch.ffn.ffn as ffn_module -from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE -from rl_engine.kernels.ops.pytorch.ffn.ffn import ( +import rl_engine.kernels.ops.triton.ffn.ffn as ffn_module +from rl_engine.kernels.ops.triton.ffn import ( QWEN3_8B_HIDDEN_SIZE, QWEN3_8B_INTERMEDIATE_SIZE, + Qwen3FFNForwardWeights, + pack_qwen3_ffn_forward_weights, qwen3_ffn, + qwen3_ffn_triton, + refresh_qwen3_ffn_forward_weights, ) +from rl_engine.platforms.device import device_ctx -_REQUIRED_SYMBOLS = ( - "det_gemm_fwd", - "det_gemm_fwd_rhs_transposed", - "det_gemm_db_transposed", - "swiglu_forward", - "swiglu_backward", -) -_HIDDEN = 64 -_INTERMEDIATE = 512 -_TOPOLOGY_TOKENS = 256 -_CP_TOKEN_COUNTS = (8, 32, 64, 96, 128, 256) -_TOKEN_BOUNDARY_COUNTS = (8, 31, 32, 33, 64, 96, 128) -_WORLD2_CONFIGS = ( - ("tp2_sp", 2, 1, True, _TOPOLOGY_TOKENS), - ("cp2", 1, 2, False, _TOPOLOGY_TOKENS), - *((f"cp2_T{token_count}", 1, 2, False, token_count) for token_count in _CP_TOKEN_COUNTS), -) -_WORLD4_CONFIGS = ( - ("tp4", 4, 1, False, _TOPOLOGY_TOKENS), - ("cp4", 1, 4, False, _TOPOLOGY_TOKENS), - ("tp2_cp2", 2, 2, False, _TOPOLOGY_TOKENS), - ("tp2_cp2_sp", 2, 2, True, _TOPOLOGY_TOKENS), - *((f"cp4_T{token_count}", 1, 4, False, token_count) for token_count in _CP_TOKEN_COUNTS), -) -_WORLD8_WORLD_GROUP_CONFIGS = ( - ("tp8", 8, 1, False, _TOPOLOGY_TOKENS), - ("tp8_sp", 8, 1, True, _TOPOLOGY_TOKENS), - ("cp8", 1, 8, False, _TOPOLOGY_TOKENS), - *((f"cp8_T{token_count}", 1, 8, False, token_count) for token_count in _CP_TOKEN_COUNTS), -) -_WORLD8_TP2_CP4_CONFIGS = (("tp2_cp4", 2, 4, False, _TOPOLOGY_TOKENS),) -_WORLD8_TP4_CP2_CONFIGS = ( - ("tp4_cp2", 4, 2, False, _TOPOLOGY_TOKENS), - ("tp4_cp2_sp", 4, 2, True, _TOPOLOGY_TOKENS), -) - - -def _has_sm90_ffn_devices(count: int) -> bool: - return ( - _EXT_AVAILABLE - and torch.distributed.is_available() - and torch.distributed.is_nccl_available() - and torch.cuda.device_count() >= count - and all(torch.cuda.get_device_capability(index)[0] == 9 for index in range(count)) - and all(hasattr(_C, name) for name in _REQUIRED_SYMBOLS) - ) - - -def _has_sm90_ffn() -> bool: - return ( - torch.cuda.is_available() - and torch.cuda.get_device_capability()[0] == 9 - and _EXT_AVAILABLE - and all(hasattr(_C, name) for name in _REQUIRED_SYMBOLS) - ) - +_IS_ROCM = getattr(torch.version, "hip", None) is not None +_HAS_GPU = torch.cuda.is_available() and device_ctx.device.type == "cuda" -requires_cuda_ffn = pytest.mark.skipif( - not _has_sm90_ffn(), - reason="FFN optimized-path validation requires SM90 and the GEMM/SwiGLU extension", +requires_rocm = pytest.mark.skipif( + not (_IS_ROCM and _HAS_GPU), + reason="the Triton FFN acceptance tests require a ROCm GPU", ) -class _TorchKernelStub: - def __init__(self) -> None: - self.calls: list[str] = [] - - def det_gemm_fwd(self, a, b): - self.calls.append("det_gemm_fwd") - return a @ b - - def det_gemm_fwd_rhs_transposed(self, a, bt): - self.calls.append("det_gemm_fwd_rhs_transposed") - return a @ bt.t() - - def det_gemm_db_transposed(self, a, grad_output): - self.calls.append("det_gemm_db_transposed") - return grad_output.t() @ a +def _randn(shape, *, seed: int, dtype=torch.float32, device="cpu"): + generator = torch.Generator(device="cpu").manual_seed(seed) + value = torch.randn(*shape, generator=generator, dtype=torch.float32) * 0.1 + return value.to(device=device, dtype=dtype) - def swiglu_forward(self, gate, up): - self.calls.append("swiglu_forward") - return gate * torch.sigmoid(gate) * up - def swiglu_backward(self, grad_output, gate, up): - self.calls.append("swiglu_backward") - sigmoid = torch.sigmoid(gate) - grad_gate = grad_output * up * sigmoid * (1.0 + gate * (1.0 - sigmoid)) - grad_up = grad_output * gate * sigmoid - return grad_gate, grad_up +def _assert_same_raw_bytes(actual: torch.Tensor, expected: torch.Tensor) -> None: + actual = actual.detach() + expected = expected.detach() + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype == torch.bfloat16 + assert actual.is_contiguous() + assert expected.is_contiguous() + assert torch.equal( + actual.reshape(-1).view(torch.uint8), + expected.reshape(-1).view(torch.uint8), + ) def _reference(hidden_states, gate_weight, up_weight, down_weight): gate = hidden_states @ gate_weight.t() up = hidden_states @ up_weight.t() - activated = F.silu(gate) * up - return (activated @ down_weight.t()), gate, up, activated - - -def _randn(shape, *, seed, device="cpu", dtype=torch.float32): - 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=dtype) - - -def _close_ffn_collectives() -> None: - for collective in list(ffn_module._COLLECTIVES.values()): - collective.close() - 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, - *, - tp_size: int, - cp_size: int, - sequence_parallel: bool, - token_count: int, - intermediate_size: int, -) -> tuple[int, int, int, int]: - tp_rank = rank % tp_size - cp_rank = rank // tp_size - cp_tokens = token_count // cp_size - local_tokens = cp_tokens // tp_size if sequence_parallel else cp_tokens - token_start = cp_rank * cp_tokens - if sequence_parallel: - token_start += tp_rank * local_tokens - token_end = token_start + local_tokens - local_i = intermediate_size // tp_size - feat_start = tp_rank * local_i - feat_end = feat_start + local_i - return token_start, token_end, feat_start, feat_end - - -def _spawn_nccl_workers(worker, world_size: int, worker_args=(), *, timeout: int = 180) -> None: - if not _has_sm90_ffn_devices(world_size): - pytest.skip(f"requires {world_size} SM90 GPUs, NCCL, and the GEMM/SwiGLU extension") - - ctx = mp.get_context("spawn") - with tempfile.TemporaryDirectory() as tmpdir: - init_method = (Path(tmpdir) / "nccl_init").as_uri() - result_queue = ctx.Queue() - processes = [ - ctx.Process( - target=worker, - args=(rank, world_size, init_method, result_queue, *worker_args), - ) - for rank in range(world_size) - ] - for process in processes: - process.start() - results = [] - try: - for _ in processes: - results.append(result_queue.get(timeout=timeout)) - except queue.Empty: - for process in processes: - if process.is_alive(): - process.terminate() - pytest.fail(f"timed out waiting for {world_size} FFN workers") - finally: - for process in processes: - process.join(timeout=10) - if process.is_alive(): - process.terminate() - - for result in sorted(results, key=lambda item: item["rank"]): - assert result["ok"], result.get("traceback") or result.get("failures") - for process in processes: - assert process.exitcode == 0 - - -def _distributed_ffn_backward_nccl_worker( - rank, - world_size, - init_method, - result_queue, - cp_size, - sequence_parallel, -): - try: - import torch.distributed as dist - - torch.cuda.set_device(rank) - dist.init_process_group( - backend="nccl", - init_method=init_method, - rank=rank, - world_size=world_size, - ) - - tp_size = world_size // cp_size - tp_groups = [ - dist.new_group(list(range(cp_rank * tp_size, (cp_rank + 1) * tp_size))) - for cp_rank in range(cp_size) - ] - cp_groups = [ - dist.new_group([cp_rank * tp_size + tp_rank for cp_rank in range(cp_size)]) - for tp_rank in range(tp_size) - ] - tp_rank = rank % tp_size - cp_rank = rank // tp_size - tp_group = tp_groups[cp_rank] - cp_group = cp_groups[tp_rank] if cp_size > 1 else None - - token_count, hidden_size, intermediate_size = 8, 64, 128 - token_start, token_end, feature_start, feature_end = _shard_ranges( - rank, - tp_size=tp_size, - cp_size=cp_size, - sequence_parallel=sequence_parallel, - token_count=token_count, - intermediate_size=intermediate_size, - ) - local_tokens = token_end - token_start + return (F.silu(gate) * up) @ down_weight.t() - device = torch.device("cuda", rank) - rmsnorm_output = _randn( - (token_count, hidden_size), seed=40, device=device, dtype=torch.bfloat16 - ) - gate_weight = _randn( - (intermediate_size, hidden_size), - seed=41, - device=device, - dtype=torch.bfloat16, - ) - up_weight = _randn( - (intermediate_size, hidden_size), - seed=42, - device=device, - dtype=torch.bfloat16, - ) - down_weight = _randn( - (hidden_size, intermediate_size), - seed=43, - device=device, - dtype=torch.bfloat16, - ) - grad_output = _randn( - (token_count, hidden_size), seed=44, device=device, dtype=torch.bfloat16 - ) - reference_inputs = [ - value.detach().float().requires_grad_(True) - for value in (rmsnorm_output, gate_weight, up_weight, down_weight) - ] - reference_output, _, _, _ = _reference(*reference_inputs) - reference_output.backward(grad_output.float()) - - local_grad_output = grad_output[token_start:token_end].contiguous() - actual_inputs = [ - value.detach().clone().requires_grad_(True) - for value in ( - rmsnorm_output[token_start:token_end].contiguous(), - gate_weight[feature_start:feature_end].contiguous(), - up_weight[feature_start:feature_end].contiguous(), - down_weight[:, feature_start:feature_end].contiguous(), - ) - ] - actual_output = qwen3_ffn( - *actual_inputs, - tp_group=tp_group, - cp_group=cp_group, - sequence_parallel=sequence_parallel, - ) - actual_output.backward(local_grad_output) - - expected_grads = ( - reference_inputs[0].grad[token_start:token_end], - reference_inputs[1].grad[feature_start:feature_end], - reference_inputs[2].grad[feature_start:feature_end], - reference_inputs[3].grad[:, feature_start:feature_end], - ) - torch.testing.assert_close( - actual_output.float(), - reference_output[token_start:token_end].detach(), - atol=5e-2, - rtol=2e-2, - ) - for actual, expected in zip(actual_inputs, expected_grads, strict=True): - torch.testing.assert_close( - actual.grad.float(), - expected, - atol=5e-2, - rtol=2e-2, - ) - - slice_size = max(1, local_tokens // 2) - slice_start = (local_tokens - slice_size) // 2 - slice_end = slice_start + slice_size - slice_inputs = [ - value.detach().clone().requires_grad_(True) - for value in ( - actual_inputs[0][slice_start:slice_end].contiguous(), - actual_inputs[1], - actual_inputs[2], - actual_inputs[3], - ) - ] - slice_output = qwen3_ffn( - *slice_inputs, - tp_group=tp_group, - cp_group=cp_group, - sequence_parallel=sequence_parallel, - ) - slice_output.backward(local_grad_output[slice_start:slice_end]) - - assert torch.equal( - slice_output, - actual_output[slice_start:slice_end], - ), "FFN output changed with the local token batch size" - assert torch.equal( - slice_inputs[0].grad, - actual_inputs[0].grad[slice_start:slice_end], - ), "FFN input gradient changed with the local token batch size" - result_queue.put({"ok": True, "rank": rank}) - except Exception: # pragma: no cover - forwarded to the parent process. - result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) - raise - finally: - _close_ffn_collectives() - if torch.distributed.is_available() and torch.distributed.is_initialized(): - torch.distributed.destroy_process_group() - - -def _tp1_vs_tpn_train_infer_worker(rank, world_size, init_method, result_queue, expect_match): - try: - import torch.distributed as dist - - torch.cuda.set_device(rank) - dist.init_process_group( - backend="nccl", - init_method=init_method, - rank=rank, - world_size=world_size, - ) - device = torch.device("cuda", rank) - token_count, hidden_size, intermediate_size = 16, 64, 256 - hidden = _randn((token_count, hidden_size), seed=50, device=device, dtype=torch.bfloat16) - gate_weight = _randn( - (intermediate_size, hidden_size), - seed=51, - device=device, - dtype=torch.bfloat16, - ) - up_weight = _randn( - (intermediate_size, hidden_size), - seed=52, - device=device, - dtype=torch.bfloat16, - ) - down_weight = _randn( - (hidden_size, intermediate_size), - seed=53, - device=device, - dtype=torch.bfloat16, - ) - grad_output = _randn( - (token_count, hidden_size), seed=54, device=device, dtype=torch.bfloat16 - ) +class _ValidationTensor: + """Tensor-shaped object for validation-order tests without a GPU.""" - with torch.no_grad(): - infer_tp1 = qwen3_ffn(hidden, gate_weight, up_weight, down_weight) - - tp1_inputs = [ - value.detach().clone().requires_grad_(True) - for value in (hidden, gate_weight, up_weight, down_weight) - ] - train_tp1 = qwen3_ffn(*tp1_inputs) - train_tp1.backward(grad_output) - assert torch.equal(infer_tp1, train_tp1.detach()), "TP=1 train/infer forward mismatch" - - local_i = intermediate_size // world_size - feat_start = rank * local_i - feat_end = feat_start + local_i - shard = ( - hidden, - gate_weight[feat_start:feat_end].contiguous(), - up_weight[feat_start:feat_end].contiguous(), - down_weight[:, feat_start:feat_end].contiguous(), - ) - with torch.no_grad(): - infer_tpn = qwen3_ffn(*shard, tp_group=dist.group.WORLD) - - tpn_inputs = [value.detach().clone().requires_grad_(True) for value in shard] - train_tpn = qwen3_ffn(*tpn_inputs, tp_group=dist.group.WORLD) - train_tpn.backward(grad_output) - assert torch.equal( - infer_tpn, train_tpn.detach() - ), f"TP={world_size} train/infer forward mismatch" - - infer_match = torch.equal(infer_tp1, infer_tpn) - train_match = torch.equal(train_tp1.detach(), train_tpn.detach()) - hidden_match = torch.equal(tp1_inputs[0].grad, tpn_inputs[0].grad) - if expect_match: - assert infer_match, f"TP=1 vs TP={world_size} infer forward mismatch" - assert train_match, f"TP=1 vs TP={world_size} train forward mismatch" - assert hidden_match, f"TP=1 vs TP={world_size} hidden grad mismatch" - else: - assert not infer_match, f"TP=1 vs TP={world_size} infer forward unexpectedly matched" - assert not train_match, f"TP=1 vs TP={world_size} train forward unexpectedly matched" - assert not hidden_match, f"TP=1 vs TP={world_size} hidden grad unexpectedly matched" - - assert torch.equal( - tp1_inputs[1].grad[feat_start:feat_end], tpn_inputs[1].grad - ), f"TP=1 vs TP={world_size} gate weight grad mismatch" - assert torch.equal( - tp1_inputs[2].grad[feat_start:feat_end], tpn_inputs[2].grad - ), f"TP=1 vs TP={world_size} up weight grad mismatch" - assert torch.equal( - tp1_inputs[3].grad[:, feat_start:feat_end], tpn_inputs[3].grad - ), f"TP=1 vs TP={world_size} down weight grad mismatch" - - result_queue.put({"ok": True, "rank": rank}) - except Exception: # pragma: no cover - forwarded to the parent process. - result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) - raise - finally: - _close_ffn_collectives() - if torch.distributed.is_available() and torch.distributed.is_initialized(): - torch.distributed.destroy_process_group() - - -def _make_topology_inputs(token_count, device): - hidden = _randn((token_count, _HIDDEN), seed=60, device=device, dtype=torch.bfloat16) - gate = _randn((_INTERMEDIATE, _HIDDEN), seed=61, device=device, dtype=torch.bfloat16) - up = _randn((_INTERMEDIATE, _HIDDEN), seed=62, device=device, dtype=torch.bfloat16) - down = _randn((_HIDDEN, _INTERMEDIATE), seed=63, device=device, dtype=torch.bfloat16) - grad = _randn((token_count, _HIDDEN), seed=64, device=device, dtype=torch.bfloat16) - return hidden, gate, up, down, grad - - -def _canonical(hidden, gate, up, down, grad): - with torch.no_grad(): - infer = qwen3_ffn(hidden, gate, up, down) - inputs = [value.detach().clone().requires_grad_(True) for value in (hidden, gate, up, down)] - train = qwen3_ffn(*inputs) - train.backward(grad) - return infer, train, inputs - - -def _mesh_groups(dist, tp_size, cp_size): - world_size = dist.get_world_size() - if tp_size == world_size and cp_size == 1: - return [dist.group.WORLD], [] - if cp_size == world_size and tp_size == 1: - return [], [dist.group.WORLD] - tp_groups = [] - if tp_size > 1: - for cp_rank in range(cp_size): - ranks = list(range(cp_rank * tp_size, (cp_rank + 1) * tp_size)) - tp_groups.append(dist.new_group(ranks)) - cp_groups = [] - if cp_size > 1: - for tp_rank in range(tp_size): - ranks = [cp_rank * tp_size + tp_rank for cp_rank in range(cp_size)] - cp_groups.append(dist.new_group(ranks)) - return tp_groups, cp_groups - - -def _run_topology_config( - rank, - dist, - meshes, - *, - name, - tp_size, - cp_size, - sequence_parallel, - hidden, - gate, - up, - down, - grad, - infer_ref, - train_ref, - ref_inputs, -): - key = (tp_size, cp_size) - if key not in meshes: - meshes[key] = _mesh_groups(dist, tp_size, cp_size) - tp_groups, cp_groups = meshes[key] - tp_rank = rank % tp_size - cp_rank = rank // tp_size - tp_group = tp_groups[cp_rank] if tp_size > 1 else None - cp_group = cp_groups[tp_rank] if cp_size > 1 else None - token_start, token_end, feat_start, feat_end = _shard_ranges( - rank, - tp_size=tp_size, - cp_size=cp_size, - sequence_parallel=sequence_parallel, - token_count=hidden.size(0), - intermediate_size=_INTERMEDIATE, - ) - shard = ( - hidden[token_start:token_end].contiguous(), - gate[feat_start:feat_end].contiguous(), - up[feat_start:feat_end].contiguous(), - down[:, feat_start:feat_end].contiguous(), - ) - with torch.no_grad(): - infer = qwen3_ffn( - *shard, - tp_group=tp_group, - cp_group=cp_group, - sequence_parallel=sequence_parallel, - ) - inputs = [value.detach().clone().requires_grad_(True) for value in shard] - train = qwen3_ffn( - *inputs, - tp_group=tp_group, - cp_group=cp_group, - sequence_parallel=sequence_parallel, - ) - train.backward(grad[token_start:token_end].contiguous()) + def __init__(self, shape, dtype=torch.bfloat16) -> None: + self.shape = torch.Size(shape) + self.dtype = dtype + self.device = torch.device("cuda") + self.is_cuda = True - assert torch.equal(infer, train.detach()), f"{name}: train/infer forward mismatch" - assert torch.equal( - infer, infer_ref[token_start:token_end] - ), f"{name}: infer forward mismatch vs TP=1/CP=1" - assert torch.equal( - train.detach(), train_ref.detach()[token_start:token_end] - ), f"{name}: train forward mismatch vs TP=1/CP=1" - assert torch.equal( - inputs[0].grad, ref_inputs[0].grad[token_start:token_end] - ), f"{name}: hidden grad mismatch vs TP=1/CP=1" + def dim(self): + return len(self.shape) - weight_checks = ( - (1, ref_inputs[1].grad[feat_start:feat_end], "gate"), - (2, ref_inputs[2].grad[feat_start:feat_end], "up"), - (3, ref_inputs[3].grad[:, feat_start:feat_end], "down"), - ) - for index, expected, label in weight_checks: - assert torch.equal( - inputs[index].grad, expected - ), f"{name}: {label} weight grad mismatch vs TP=1/CP=1" - - -def _topology_worker(rank, world_size, init_method, result_queue, configs): - try: - import torch.distributed as dist - - torch.cuda.set_device(rank) - dist.init_process_group( - backend="nccl", - init_method=init_method, - rank=rank, - world_size=world_size, - timeout=timedelta(minutes=5), - device_id=torch.device("cuda", rank), - ) - device = torch.device("cuda", rank) - meshes = {} - canonical = {} - for name, tp_size, cp_size, sequence_parallel, token_count in configs: - if token_count not in canonical: - tensors = _make_topology_inputs(token_count, device) - canonical[token_count] = (*tensors, *_canonical(*tensors)) - hidden, gate, up, down, grad, infer_ref, train_ref, ref_inputs = canonical[token_count] - _run_topology_config( - rank, - dist, - meshes, - name=name, - tp_size=tp_size, - cp_size=cp_size, - sequence_parallel=sequence_parallel, - hidden=hidden, - gate=gate, - up=up, - down=down, - grad=grad, - infer_ref=infer_ref, - train_ref=train_ref, - ref_inputs=ref_inputs, - ) - result_queue.put({"ok": True, "rank": rank}) - except Exception: # pragma: no cover - forwarded to the parent process. - result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) - raise - finally: - _close_ffn_collectives() - if torch.distributed.is_available() and torch.distributed.is_initialized(): - torch.distributed.destroy_process_group() - - -def _ffn_tensors(token_count, device, seed, *, hidden=_HIDDEN, intermediate=_INTERMEDIATE): - rmsnorm = _randn((token_count, hidden), seed=seed, device=device, dtype=torch.bfloat16) - gate = _randn((intermediate, hidden), seed=seed + 1, device=device, dtype=torch.bfloat16) - up = _randn((intermediate, hidden), seed=seed + 2, device=device, dtype=torch.bfloat16) - down = _randn((hidden, intermediate), seed=seed + 3, device=device, dtype=torch.bfloat16) - return rmsnorm, gate, up, down - - -def _cache_worker(rank, world_size, init_method, result_queue): - try: - import torch.distributed as dist - - torch.cuda.set_device(rank) - dist.init_process_group( - backend="nccl", - init_method=init_method, - rank=rank, - world_size=world_size, - ) - device = torch.device("cuda", rank) - ffn_module._COLLECTIVE_MIN_CAPACITY_BYTES = 64 - _close_ffn_collectives() - - small = _ffn_tensors(8, device, seed=100, intermediate=128) - first = qwen3_ffn(*small, tp_group=dist.group.WORLD) - assert len(ffn_module._COLLECTIVES) == 1 - ((cache_key, first_collective),) = ffn_module._COLLECTIVES.items() - first_handle = first_collective._handle - first_capacity = first_collective.max_size_bytes - assert first_handle != 0 - - repeated = qwen3_ffn(*small, tp_group=dist.group.WORLD) - assert torch.equal(first, repeated) - assert len(ffn_module._COLLECTIVES) == 1 - assert ffn_module._COLLECTIVES[cache_key] is first_collective - assert first_collective._handle == first_handle - - large = _ffn_tensors(256, device, seed=110, intermediate=128) - grown = qwen3_ffn(*large, tp_group=dist.group.WORLD) - assert grown.shape[0] == 256 - assert len(ffn_module._COLLECTIVES) == 1 - grown_collective = next(iter(ffn_module._COLLECTIVES.values())) - assert grown_collective is not first_collective - assert first_collective._handle == first_handle - assert grown_collective.max_size_bytes > first_capacity - assert grown_collective._handle != 0 - - _close_ffn_collectives() - assert ffn_module._COLLECTIVES == {} - recreated = qwen3_ffn(*small, tp_group=dist.group.WORLD) - assert torch.equal(recreated, first) - assert len(ffn_module._COLLECTIVES) == 1 - recreated_collective = next(iter(ffn_module._COLLECTIVES.values())) - assert recreated_collective is not grown_collective - assert recreated_collective._handle != 0 - - rebuilt_group = dist.new_group(ranks=[0, 1]) - rebuilt = qwen3_ffn(*small, tp_group=rebuilt_group) - assert torch.equal(rebuilt, first) - assert len(ffn_module._COLLECTIVES) == 2 - - _close_ffn_collectives() - first_collective.close() - result_queue.put({"ok": True, "rank": rank}) - except Exception: # pragma: no cover - forwarded to the parent process. - result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) - raise - finally: - _close_ffn_collectives() - if torch.distributed.is_available() and torch.distributed.is_initialized(): - torch.distributed.destroy_process_group() - - -def _uneven_sp_worker(rank, world_size, init_method, result_queue): - try: - import torch.distributed as dist - - torch.cuda.set_device(rank) - dist.init_process_group( - backend="nccl", - init_method=init_method, - rank=rank, - world_size=world_size, - ) - device = torch.device("cuda", rank) - hidden, gate, up, down = _ffn_tensors(2, device, seed=130, intermediate=128) - local_i = 128 // world_size - feat_start = rank * local_i - feat_end = feat_start + local_i - local_hidden = hidden[:2] if rank == 0 else hidden[:1] - try: - qwen3_ffn( - local_hidden, - gate[feat_start:feat_end].contiguous(), - up[feat_start:feat_end].contiguous(), - down[:, feat_start:feat_end].contiguous(), - tp_group=dist.group.WORLD, - sequence_parallel=True, - ) - result_queue.put( - { - "ok": False, - "rank": rank, - "failures": "uneven SP tokens should have raised", - } - ) - except ValueError as exc: - message = str(exc) - if "matching shapes" not in message and "world_size" not in message: - raise - result_queue.put({"ok": True, "rank": rank}) - except Exception: # pragma: no cover - forwarded to the parent process. - result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) - raise - finally: - _close_ffn_collectives() - if torch.distributed.is_available() and torch.distributed.is_initialized(): - torch.distributed.destroy_process_group() - - -def _qwen3_8b_weights(device): - hidden = _randn((8, QWEN3_8B_HIDDEN_SIZE), seed=90, device=device, dtype=torch.bfloat16) - gate = _randn( - (QWEN3_8B_INTERMEDIATE_SIZE, QWEN3_8B_HIDDEN_SIZE), - seed=91, - device=device, - dtype=torch.bfloat16, - ) - up = _randn( - (QWEN3_8B_INTERMEDIATE_SIZE, QWEN3_8B_HIDDEN_SIZE), - seed=92, - device=device, - dtype=torch.bfloat16, - ) - down = _randn( - (QWEN3_8B_HIDDEN_SIZE, QWEN3_8B_INTERMEDIATE_SIZE), - seed=93, - device=device, - dtype=torch.bfloat16, - ) - grad = _randn((8, QWEN3_8B_HIDDEN_SIZE), seed=94, device=device, dtype=torch.bfloat16) - return hidden, gate, up, down, grad + def numel(self): + result = 1 + for size in self.shape: + result *= size + return result + def size(self, dim): + return self.shape[dim] -def _qwen3_8b_tp2_worker(rank, world_size, init_method, result_queue): - try: - import torch.distributed as dist - torch.cuda.set_device(rank) - dist.init_process_group( - backend="nccl", - init_method=init_method, - rank=rank, - world_size=world_size, - ) - device = torch.device("cuda", rank) - hidden, gate, up, down, grad = _qwen3_8b_weights(device) - with torch.no_grad(): - infer_tp1 = qwen3_ffn(hidden, gate, up, down) - tp1_inputs = [ - value.detach().clone().requires_grad_(True) for value in (hidden, gate, up, down) - ] - train_tp1 = qwen3_ffn(*tp1_inputs) - train_tp1.backward(grad) - - local_i = QWEN3_8B_INTERMEDIATE_SIZE // world_size - feat_start = rank * local_i - feat_end = feat_start + local_i - shard = ( - hidden, - gate[feat_start:feat_end].contiguous(), - up[feat_start:feat_end].contiguous(), - down[:, feat_start:feat_end].contiguous(), - ) - with torch.no_grad(): - infer_tp2 = qwen3_ffn(*shard, tp_group=dist.group.WORLD) - tp2_inputs = [value.detach().clone().requires_grad_(True) for value in shard] - train_tp2 = qwen3_ffn(*tp2_inputs, tp_group=dist.group.WORLD) - train_tp2.backward(grad) - - assert torch.equal(infer_tp1, infer_tp2) - assert torch.equal(train_tp1.detach(), train_tp2.detach()) - assert torch.equal(tp1_inputs[0].grad, tp2_inputs[0].grad) - assert torch.equal(tp1_inputs[1].grad[feat_start:feat_end], tp2_inputs[1].grad) - assert torch.equal(tp1_inputs[2].grad[feat_start:feat_end], tp2_inputs[2].grad) - assert torch.equal(tp1_inputs[3].grad[:, feat_start:feat_end], tp2_inputs[3].grad) - result_queue.put({"ok": True, "rank": rank}) - except Exception: # pragma: no cover - forwarded to the parent process. - result_queue.put({"ok": False, "rank": rank, "traceback": traceback.format_exc()}) - raise - finally: - _close_ffn_collectives() - if torch.distributed.is_available() and torch.distributed.is_initialized(): - torch.distributed.destroy_process_group() - - -def test_qwen_ffn_qwen3_8b_dimensions_are_pinned(): +def test_qwen3_ffn_public_api_and_model_dimensions_are_pinned(): + assert qwen3_ffn_triton is qwen3_ffn assert QWEN3_8B_HIDDEN_SIZE == 4096 assert QWEN3_8B_INTERMEDIATE_SIZE == 12288 - -def test_qwen_ffn_backward_matches_autograd_reference(monkeypatch): - stub = _TorchKernelStub() - monkeypatch.setattr(ffn_module, "_C", stub) - monkeypatch.setattr(ffn_module, "_EXT_AVAILABLE", True) - monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) - monkeypatch.setattr( - "rl_engine.kernels.ops.cuda.matmul.det_gemm._require_sm90_backend", - lambda: None, + signature = inspect.signature(qwen3_ffn) + assert tuple(signature.parameters) == ( + "rmsnorm_output", + "gate_weight", + "up_weight", + "down_weight", + "forward_weights", + "tp_group", + "cp_group", + "sequence_parallel", ) + assert signature.parameters["forward_weights"].kind is inspect.Parameter.KEYWORD_ONLY + assert signature.parameters["forward_weights"].default is None + assert signature.parameters["tp_group"].kind is inspect.Parameter.KEYWORD_ONLY + assert signature.parameters["cp_group"].kind is inspect.Parameter.KEYWORD_ONLY + assert signature.parameters["sequence_parallel"].default is False + + +@pytest.mark.parametrize("weight_name", ["gate_weight", "up_weight", "down_weight"]) +def test_qwen3_ffn_rejects_non_huggingface_weight_layout(weight_name): + tensors = { + "rmsnorm_output": torch.empty((2, 8), dtype=torch.bfloat16), + "gate_weight": torch.empty((12, 8), dtype=torch.bfloat16), + "up_weight": torch.empty((12, 8), dtype=torch.bfloat16), + "down_weight": torch.empty((8, 12), dtype=torch.bfloat16), + } + tensors[weight_name] = tensors[weight_name].t().contiguous() + + with pytest.raises(ValueError, match=rf"{weight_name} must have shape"): + qwen3_ffn(**tensors) + + +def test_qwen3_ffn_rejects_zero_intermediate_size_before_device_dispatch(): + tensors = { + "rmsnorm_output": torch.empty((2, 8), dtype=torch.bfloat16), + "gate_weight": torch.empty((0, 8), dtype=torch.bfloat16), + "up_weight": torch.empty((0, 8), dtype=torch.bfloat16), + "down_weight": torch.empty((8, 0), dtype=torch.bfloat16), + } + + with pytest.raises(ValueError, match="intermediate size must be positive"): + qwen3_ffn(**tensors) + + +@pytest.mark.parametrize( + "input_name", + ("rmsnorm_output", "gate_weight", "up_weight", "down_weight"), +) +def test_qwen3_ffn_rejects_non_bf16_inputs(input_name, monkeypatch): + monkeypatch.setattr(ffn_module, "Tensor", _ValidationTensor) + tensors = { + "rmsnorm_output": _ValidationTensor((2, 8)), + "gate_weight": _ValidationTensor((12, 8)), + "up_weight": _ValidationTensor((12, 8)), + "down_weight": _ValidationTensor((8, 12)), + } + tensors[input_name].dtype = torch.float32 - hidden = _randn((2, 3, 8), seed=0) - gate_weight = _randn((12, 8), seed=1) - up_weight = _randn((12, 8), seed=2) - down_weight = _randn((8, 12), seed=3) - grad_output = _randn(hidden.shape, seed=4) - - ref_inputs = [ - value.detach().clone().requires_grad_(True) - for value in (hidden, gate_weight, up_weight, down_weight) - ] - expected, _, _, _ = _reference(*ref_inputs) - expected.backward(grad_output) - - actual_inputs = [ - value.detach().clone().requires_grad_(True) - for value in (hidden, gate_weight, up_weight, down_weight) - ] - actual = qwen3_ffn(*actual_inputs) - actual.backward(grad_output) - - torch.testing.assert_close(actual, expected.detach()) - for actual_input, reference in zip(actual_inputs, ref_inputs, strict=True): - torch.testing.assert_close(actual_input.grad, reference.grad) - for weight in actual_inputs[1:]: - assert weight.grad is not None - assert weight.grad.is_contiguous() - assert weight.grad.stride() == weight.stride() - - assert stub.calls.count("det_gemm_fwd") == 3 - assert stub.calls.count("det_gemm_fwd_rhs_transposed") == 3 - assert stub.calls.count("det_gemm_db_transposed") == 3 - assert stub.calls.count("swiglu_forward") == 1 - assert stub.calls.count("swiglu_backward") == 1 - - -def test_qwen_ffn_disable_split_k_false_uses_torch_matmul(monkeypatch): - stub = _TorchKernelStub() - monkeypatch.setattr(ffn_module, "_C", stub) - monkeypatch.setattr(ffn_module, "_EXT_AVAILABLE", True) - monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) - - hidden = _randn((2, 3, 8), seed=0) - gate_weight = _randn((12, 8), seed=1) - up_weight = _randn((12, 8), seed=2) - down_weight = _randn((8, 12), seed=3) - grad_output = _randn(hidden.shape, seed=4) - - ref_inputs = [ - value.detach().clone().requires_grad_(True) - for value in (hidden, gate_weight, up_weight, down_weight) - ] - expected, _, _, _ = _reference(*ref_inputs) - expected.backward(grad_output) - - actual_inputs = [ - value.detach().clone().requires_grad_(True) - for value in (hidden, gate_weight, up_weight, down_weight) - ] - actual = qwen3_ffn(*actual_inputs, disable_split_k=False) - actual.backward(grad_output) - - torch.testing.assert_close(actual, expected.detach()) - for actual_input, reference in zip(actual_inputs, ref_inputs, strict=True): - torch.testing.assert_close(actual_input.grad, reference.grad) - - assert stub.calls.count("det_gemm_fwd") == 0 - assert stub.calls.count("det_gemm_fwd_rhs_transposed") == 0 - assert stub.calls.count("det_gemm_db_transposed") == 0 - assert stub.calls.count("swiglu_forward") == 1 - assert stub.calls.count("swiglu_backward") == 1 + with pytest.raises(TypeError, match=rf"{input_name} must have dtype bfloat16"): + ffn_module._validate_ffn_inputs(**tensors) -def test_qwen_ffn_deterministic_false_uses_production_gemm(monkeypatch): - modes = [] +def test_qwen3_ffn_sequence_parallel_flag_must_be_bool(monkeypatch): monkeypatch.setattr(ffn_module, "_validate_ffn_inputs", lambda *args: None) - monkeypatch.setattr(ffn_module, "_require_ffn_kernels", lambda **kwargs: modes.append(kwargs)) - monkeypatch.setattr( - ffn_module._DeterministicFFNFunction, - "apply", - lambda *args: args[-1], - ) - - tensors = [torch.empty(1)] * 4 - assert qwen3_ffn(*tensors, deterministic=False) is False - assert modes == [{"disable_split_k": False, "packed_gate_up": False}] - - -def test_qwen_ffn_rejects_conflicting_backend_switches(): - tensors = [torch.empty(1)] * 4 - - with pytest.raises(ValueError, match="conflicting FFN backends"): - qwen3_ffn(*tensors, deterministic=True, disable_split_k=False) - - -def test_qwen_ffn_rejects_non_bool_deterministic(): - tensors = [torch.empty(1)] * 4 - - with pytest.raises(TypeError, match="deterministic must be a bool or None"): - qwen3_ffn(*tensors, deterministic=1) # type: ignore[arg-type] - - -def test_qwen_ffn_rejects_non_bool_disable_split_k(): - hidden = torch.empty((2, 8), dtype=torch.bfloat16) - gate_weight = torch.empty((12, 8), dtype=torch.bfloat16) - up_weight = torch.empty((12, 8), dtype=torch.bfloat16) - down_weight = torch.empty((8, 12), dtype=torch.bfloat16) - with pytest.raises(TypeError, match="disable_split_k must be a bool"): - qwen3_ffn( - hidden, - gate_weight, - up_weight, - down_weight, - disable_split_k=1, # type: ignore[arg-type] - ) - - -def test_qwen_ffn_rejects_non_huggingface_weight_layout(): - hidden = torch.empty((2, 8), dtype=torch.bfloat16) - gate_weight = torch.empty((8, 12), dtype=torch.bfloat16) - up_weight = torch.empty((12, 8), dtype=torch.bfloat16) - down_weight = torch.empty((8, 12), dtype=torch.bfloat16) + tensors = (object(), object(), object(), object()) - with pytest.raises(ValueError, match="gate_weight must have shape"): - qwen3_ffn(hidden, gate_weight, up_weight, down_weight) + with pytest.raises(TypeError, match="sequence_parallel must be a bool"): + qwen3_ffn(*tensors, sequence_parallel=1) -@requires_cuda_ffn -@pytest.mark.parametrize("disable_split_k", [True, False]) -def test_qwen_ffn_cuda_forward_backward_matches_fp32_reference(disable_split_k): - hidden = _randn((2, 3, 64), seed=10, device="cuda", dtype=torch.bfloat16) - gate_weight = _randn((128, 64), seed=11, device="cuda", dtype=torch.bfloat16) - up_weight = _randn((128, 64), seed=12, device="cuda", dtype=torch.bfloat16) - down_weight = _randn((64, 128), seed=13, device="cuda", dtype=torch.bfloat16) - grad_output = _randn(hidden.shape, seed=14, device="cuda", dtype=torch.bfloat16) +@requires_rocm +def test_qwen3_ffn_forward_backward_matches_fp32_reference(): + device = device_ctx.device + hidden = _randn((2, 3, 32), seed=20, dtype=torch.bfloat16, device=device) + gate_weight = _randn((64, 32), seed=21, dtype=torch.bfloat16, device=device) + up_weight = _randn((64, 32), seed=22, dtype=torch.bfloat16, device=device) + down_weight = _randn((32, 64), seed=23, dtype=torch.bfloat16, device=device) + grad_output = _randn(hidden.shape, seed=24, dtype=torch.bfloat16, device=device) - ref_inputs = [ + reference_inputs = [ value.detach().cpu().float().requires_grad_(True) for value in (hidden, gate_weight, up_weight, down_weight) ] - expected, _, _, _ = _reference(*ref_inputs) + expected = _reference(*reference_inputs) expected.backward(grad_output.cpu().float()) actual_inputs = [ value.detach().clone().requires_grad_(True) for value in (hidden, gate_weight, up_weight, down_weight) ] - actual = qwen3_ffn(*actual_inputs, disable_split_k=disable_split_k) + actual = qwen3_ffn(*actual_inputs) actual.backward(grad_output) + assert actual.shape == hidden.shape + assert actual.dtype is torch.bfloat16 torch.testing.assert_close( - actual.cpu().float(), - expected.detach(), - atol=5e-2, - rtol=2e-2, + actual.cpu().float(), expected.detach(), atol=5e-2, rtol=2e-2 ) - for actual_input, reference in zip(actual_inputs, ref_inputs, strict=True): + for actual_input, reference_input in zip(actual_inputs, reference_inputs, strict=True): + assert actual_input.grad.dtype is torch.bfloat16 torch.testing.assert_close( actual_input.grad.cpu().float(), - reference.grad, + reference_input.grad, atol=5e-2, rtol=2e-2, ) -@requires_cuda_ffn -def test_qwen_ffn_cuda_forward_and_input_gradient_are_batch_invariant(): - gate_weight = _randn((128, 64), seed=20, device="cuda", dtype=torch.bfloat16) - up_weight = _randn((128, 64), seed=21, device="cuda", dtype=torch.bfloat16) - down_weight = _randn((64, 128), seed=22, device="cuda", dtype=torch.bfloat16) - hidden = _randn((6, 64), seed=23, device="cuda", dtype=torch.bfloat16) - grad_output = _randn(hidden.shape, seed=24, device="cuda", dtype=torch.bfloat16) +def _training_step(values, grad_output): + inputs = [value.detach().clone().requires_grad_(True) for value in values] + output = qwen3_ffn(*inputs) + output.backward(grad_output) + return output.detach(), [value.grad.detach() for value in inputs] + + +@requires_rocm +def test_qwen3_ffn_repeat_and_train_infer_have_zero_mismatch(): + device = device_ctx.device + values = ( + _randn((8, 64), seed=30, dtype=torch.bfloat16, device=device), + _randn((128, 64), seed=31, dtype=torch.bfloat16, device=device), + _randn((128, 64), seed=32, dtype=torch.bfloat16, device=device), + _randn((64, 128), seed=33, dtype=torch.bfloat16, device=device), + ) + grad_output = _randn((8, 64), seed=34, dtype=torch.bfloat16, device=device) + + with torch.no_grad(): + inference_first = qwen3_ffn(*values) + inference_second = qwen3_ffn(*values) + training_first, gradients_first = _training_step(values, grad_output) + training_second, gradients_second = _training_step(values, grad_output) + + assert torch.equal(inference_first, inference_second) + assert torch.equal(inference_first, training_first) + assert torch.equal(training_first, training_second) + assert all( + torch.equal(first, second) + for first, second in zip(gradients_first, gradients_second, strict=True) + ) + + +@requires_rocm +def test_qwen3_ffn_output_and_hidden_gradient_are_token_slice_invariant(): + device = device_ctx.device + hidden = _randn((6, 32), seed=40, dtype=torch.bfloat16, device=device) + gate_weight = _randn((64, 32), seed=41, dtype=torch.bfloat16, device=device) + up_weight = _randn((64, 32), seed=42, dtype=torch.bfloat16, device=device) + down_weight = _randn((32, 64), seed=43, dtype=torch.bfloat16, device=device) + grad_output = _randn(hidden.shape, seed=44, dtype=torch.bfloat16, device=device) full_hidden = hidden.detach().clone().requires_grad_(True) full_output = qwen3_ffn(full_hidden, gate_weight, up_weight, down_weight) @@ -1049,122 +245,198 @@ def test_qwen_ffn_cuda_forward_and_input_gradient_are_batch_invariant(): assert torch.equal(slice_hidden.grad, full_hidden.grad[2:4]) -@requires_cuda_ffn -def test_qwen_ffn_cuda_train_and_infer_forward_are_bitwise_identical(): - hidden = _randn((16, 64), seed=30, device="cuda", dtype=torch.bfloat16) - gate_weight = _randn((128, 64), seed=31, device="cuda", dtype=torch.bfloat16) - up_weight = _randn((128, 64), seed=32, device="cuda", dtype=torch.bfloat16) - down_weight = _randn((64, 128), seed=33, device="cuda", dtype=torch.bfloat16) - with torch.no_grad(): - infer = qwen3_ffn(hidden, gate_weight, up_weight, down_weight) - train_hidden = hidden.detach().clone().requires_grad_(True) - train = qwen3_ffn(train_hidden, gate_weight, up_weight, down_weight) - assert torch.equal(infer, train.detach()) - - -@requires_cuda_ffn -@pytest.mark.parametrize("token_count", _TOKEN_BOUNDARY_COUNTS) -def test_qwen_ffn_output_and_hidden_grad_are_batch_invariant(token_count): - device = torch.device("cuda", 0) - gate = _randn((256, _HIDDEN), seed=70, device=device, dtype=torch.bfloat16) - up = _randn((256, _HIDDEN), seed=71, device=device, dtype=torch.bfloat16) - down = _randn((_HIDDEN, 256), seed=72, device=device, dtype=torch.bfloat16) - hidden = _randn((token_count, _HIDDEN), seed=73, device=device, dtype=torch.bfloat16) - grad = _randn((token_count, _HIDDEN), seed=74, device=device, dtype=torch.bfloat16) - - full_hidden = hidden.detach().clone().requires_grad_(True) - full_output = qwen3_ffn(full_hidden, gate, up, down) - full_output.backward(grad) - slice_end = min(8, token_count) - slice_hidden = hidden[:slice_end].detach().clone().requires_grad_(True) - slice_output = qwen3_ffn(slice_hidden, gate, up, down) - slice_output.backward(grad[:slice_end]) - - assert torch.equal(slice_output, full_output[:slice_end]) - assert torch.equal(slice_hidden.grad, full_hidden.grad[:slice_end]) +@requires_rocm +def test_qwen3_ffn_forward_weight_pack_is_detached_contiguous_and_transposed(): + device = device_ctx.device + gate_weight = _randn((67, 35), seed=50, dtype=torch.bfloat16, device=device) + up_weight = _randn((67, 35), seed=51, dtype=torch.bfloat16, device=device) + down_weight = _randn((35, 67), seed=52, dtype=torch.bfloat16, device=device) + packed = pack_qwen3_ffn_forward_weights(gate_weight, up_weight, down_weight) -@requires_cuda_ffn -def test_qwen_ffn_qwen3_8b_shapes_run_and_are_batch_invariant(): - device = torch.device("cuda", 0) - hidden, gate, up, down, grad = _qwen3_8b_weights(device) + assert isinstance(packed, Qwen3FFNForwardWeights) + expected = ( + gate_weight.t().contiguous(), + up_weight.t().contiguous(), + down_weight.t().contiguous(), + ) + actual = (packed.gate_weight_t, packed.up_weight_t, packed.down_weight_t) + assert tuple(value.shape for value in actual) == ((35, 67), (35, 67), (67, 35)) + for packed_weight, expected_weight in zip(actual, expected, strict=True): + assert packed_weight.is_contiguous() + assert not packed_weight.requires_grad + assert packed_weight.grad_fn is None + assert packed_weight.grad is None + _assert_same_raw_bytes(packed_weight, expected_weight) + + +@requires_rocm +def test_qwen3_ffn_standard_and_packed_paths_match_all_raw_bytes(): + device = device_ctx.device + values = ( + _randn((7, 35), seed=60, dtype=torch.bfloat16, device=device), + _randn((67, 35), seed=61, dtype=torch.bfloat16, device=device), + _randn((67, 35), seed=62, dtype=torch.bfloat16, device=device), + _randn((35, 67), seed=63, dtype=torch.bfloat16, device=device), + ) + grad_output = _randn((7, 35), seed=64, dtype=torch.bfloat16, device=device) + inference_pack = pack_qwen3_ffn_forward_weights(*values[1:]) with torch.no_grad(): - infer = qwen3_ffn(hidden, gate, up, down) - full_hidden = hidden.detach().clone().requires_grad_(True) - full_gate = gate.detach().clone().requires_grad_(True) - full_up = up.detach().clone().requires_grad_(True) - full_down = down.detach().clone().requires_grad_(True) - train = qwen3_ffn(full_hidden, full_gate, full_up, full_down) - train.backward(grad) - assert torch.equal(infer, train.detach()) - - slice_hidden = hidden[:4].detach().clone().requires_grad_(True) - slice_out = qwen3_ffn(slice_hidden, full_gate, full_up, full_down) - slice_out.backward(grad[:4]) - assert torch.equal(slice_out, train.detach()[:4]) - assert torch.equal(slice_hidden.grad, full_hidden.grad[:4]) - - -@requires_cuda_ffn -def test_qwen_ffn_sequence_parallel_requires_tensor_parallel_group(): - device = torch.device("cuda", 0) - hidden, gate, up, down = _ffn_tensors(8, device, seed=120, intermediate=128) - with pytest.raises(ValueError, match="sequence_parallel requires a tensor-parallel group"): - qwen3_ffn(hidden, gate, up, down, sequence_parallel=True) - - -def test_qwen_ffn_tp_correctness_and_batch_invariance(): - _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 2, (1, False), timeout=90) - - -def test_qwen_ffn_tp_sp_correctness_and_batch_invariance(): - _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 2, (1, True), timeout=90) - - -def test_qwen_ffn_tp_cp_correctness_and_batch_invariance(): - _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 4, (2, False), timeout=90) - - -def test_qwen_ffn_tp_cp_sp_correctness_and_batch_invariance(): - _spawn_nccl_workers(_distributed_ffn_backward_nccl_worker, 4, (2, True), timeout=90) - - -def test_qwen_ffn_tp1_vs_tp2_train_infer_bitwise_identical(): - _spawn_nccl_workers(_tp1_vs_tpn_train_infer_worker, 2, (True,), timeout=120) - - -def test_qwen_ffn_tp1_vs_tp8_train_infer_bitwise_identical(): - _spawn_nccl_workers(_tp1_vs_tpn_train_infer_worker, 8, (True,), timeout=120) - - -def test_qwen_ffn_world2_tp_sp_and_cp_match_tp1_cp1_bitwise(): - _spawn_nccl_workers(_topology_worker, 2, (_WORLD2_CONFIGS,), timeout=120) - - -def test_qwen_ffn_world4_tp_cp_sp_match_tp1_cp1_bitwise(): - _spawn_nccl_workers(_topology_worker, 4, (_WORLD4_CONFIGS,), timeout=120) - - -def test_qwen_ffn_world8_tp8_and_cp8_match_tp1_cp1_bitwise(): - _spawn_nccl_workers(_topology_worker, 8, (_WORLD8_WORLD_GROUP_CONFIGS,), timeout=120) + standard_inference = qwen3_ffn(*values) + packed_inference = qwen3_ffn(*values, forward_weights=inference_pack) + + standard_inputs = [value.detach().clone().requires_grad_(True) for value in values] + packed_inputs = [value.detach().clone().requires_grad_(True) for value in values] + training_pack = pack_qwen3_ffn_forward_weights(*packed_inputs[1:]) + standard_training = qwen3_ffn(*standard_inputs) + packed_training = qwen3_ffn(*packed_inputs, forward_weights=training_pack) + standard_training.backward(grad_output) + packed_training.backward(grad_output) + + _assert_same_raw_bytes(packed_inference, standard_inference) + _assert_same_raw_bytes(standard_training, standard_inference) + _assert_same_raw_bytes(packed_training, standard_training) + for packed_input, standard_input in zip(packed_inputs, standard_inputs, strict=True): + assert packed_input.grad is not None + assert standard_input.grad is not None + _assert_same_raw_bytes(packed_input.grad, standard_input.grad) + for packed_weight in ( + training_pack.gate_weight_t, + training_pack.up_weight_t, + training_pack.down_weight_t, + ): + assert not packed_weight.requires_grad + assert packed_weight.grad is None + + +@pytest.mark.parametrize( + ("weight_index", "weight_name"), + ((0, "gate_weight"), (1, "up_weight"), (2, "down_weight")), +) +@requires_rocm +def test_qwen3_ffn_rejects_stale_forward_weights_after_inplace_source_update( + weight_index, + weight_name, +): + device = device_ctx.device + hidden = _randn((2, 8), seed=70, dtype=torch.bfloat16, device=device) + weights = [ + _randn((12, 8), seed=71, dtype=torch.bfloat16, device=device), + _randn((12, 8), seed=72, dtype=torch.bfloat16, device=device), + _randn((8, 12), seed=73, dtype=torch.bfloat16, device=device), + ] + packed = pack_qwen3_ffn_forward_weights(*weights) + with torch.no_grad(): + weights[weight_index].add_(1) + with pytest.raises(RuntimeError, match=rf"{weight_name} changed after .*refresh"): + qwen3_ffn(hidden, *weights, forward_weights=packed) -def test_qwen_ffn_world8_tp2_cp4_match_tp1_cp1_bitwise(): - _spawn_nccl_workers(_topology_worker, 8, (_WORLD8_TP2_CP4_CONFIGS,), timeout=120) +@requires_rocm +def test_qwen3_ffn_forward_weight_refresh_preserves_storage_and_bits(): + device = device_ctx.device + hidden = _randn((7, 35), seed=80, dtype=torch.bfloat16, device=device) + weights = [ + _randn((67, 35), seed=81, dtype=torch.bfloat16, device=device), + _randn((67, 35), seed=82, dtype=torch.bfloat16, device=device), + _randn((35, 67), seed=83, dtype=torch.bfloat16, device=device), + ] + packed = pack_qwen3_ffn_forward_weights(*weights) + packed_tensors = ( + packed.gate_weight_t, + packed.up_weight_t, + packed.down_weight_t, + ) + data_ptrs = tuple(weight.data_ptr() for weight in packed_tensors) -def test_qwen_ffn_world8_tp4_cp2_match_tp1_cp1_bitwise(): - _spawn_nccl_workers(_topology_worker, 8, (_WORLD8_TP4_CP2_CONFIGS,), timeout=120) + with torch.no_grad(): + for weight in weights: + weight.add_(torch.ones_like(weight)) + refreshed = refresh_qwen3_ffn_forward_weights(packed, *weights) + assert refreshed is packed + assert tuple(weight.data_ptr() for weight in packed_tensors) == data_ptrs + with torch.no_grad(): + standard = qwen3_ffn(hidden, *weights) + cached = qwen3_ffn(hidden, *weights, forward_weights=packed) + _assert_same_raw_bytes(cached, standard) + + +@requires_rocm +def test_qwen3_ffn_forward_weight_pack_works_inside_inference_mode(): + device = device_ctx.device + with torch.inference_mode(): + hidden = _randn((7, 35), seed=90, dtype=torch.bfloat16, device=device) + weights = ( + _randn((67, 35), seed=91, dtype=torch.bfloat16, device=device), + _randn((67, 35), seed=92, dtype=torch.bfloat16, device=device), + _randn((35, 67), seed=93, dtype=torch.bfloat16, device=device), + ) + assert all(torch.is_inference(weight) for weight in weights) + packed = pack_qwen3_ffn_forward_weights(*weights) + cached = qwen3_ffn(hidden, *weights, forward_weights=packed) + standard = qwen3_ffn(hidden, *weights) + + assert not torch.is_inference(packed.gate_weight_t) + assert not torch.is_inference(packed.up_weight_t) + assert not torch.is_inference(packed.down_weight_t) + _assert_same_raw_bytes(cached, standard) + + +@requires_rocm +def test_qwen3_ffn_rejects_mutated_packed_forward_weight(): + device = device_ctx.device + hidden = _randn((2, 8), seed=100, dtype=torch.bfloat16, device=device) + weights = ( + _randn((12, 8), seed=101, dtype=torch.bfloat16, device=device), + _randn((12, 8), seed=102, dtype=torch.bfloat16, device=device), + _randn((8, 12), seed=103, dtype=torch.bfloat16, device=device), + ) + packed = pack_qwen3_ffn_forward_weights(*weights) + with torch.no_grad(): + packed.gate_weight_t.add_(1) -def test_qwen_ffn_collective_cache_growth_preserves_borrowers_and_rebuilds_group(): - _spawn_nccl_workers(_cache_worker, 2, timeout=120) + with pytest.raises(RuntimeError, match="packed gate_weight changed; refresh"): + qwen3_ffn(hidden, *weights, forward_weights=packed) -def test_qwen_ffn_sequence_parallel_rejects_uneven_tokens(): - _spawn_nccl_workers(_uneven_sp_worker, 2, timeout=90) +@requires_rocm +def test_qwen3_ffn_refresh_updates_already_captured_cuda_graph(): + device = device_ctx.device + hidden = _randn((7, 35), seed=110, dtype=torch.bfloat16, device=device) + weights = [ + _randn((67, 35), seed=111, dtype=torch.bfloat16, device=device), + _randn((67, 35), seed=112, dtype=torch.bfloat16, device=device), + _randn((35, 67), seed=113, dtype=torch.bfloat16, device=device), + ] + packed = pack_qwen3_ffn_forward_weights(*weights) + capture_stream = torch.cuda.Stream(device=device) + capture_stream.wait_stream(torch.cuda.current_stream(device)) + with torch.cuda.stream(capture_stream), torch.no_grad(): + for _ in range(3): + qwen3_ffn(hidden, *weights, forward_weights=packed) + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream), torch.no_grad(): + static_output = qwen3_ffn(hidden, *weights, forward_weights=packed) + graph.replay() + torch.cuda.synchronize() + before_refresh = static_output.clone() + with torch.no_grad(): + weights[0].add_(1) + packed.refresh_(*weights) + graph.replay() + torch.cuda.synchronize() + after_refresh = static_output.clone() + with torch.no_grad(): + expected = qwen3_ffn(hidden, *weights) -def test_qwen_ffn_qwen3_8b_shapes_tp2_matches_tp1_bitwise(): - _spawn_nccl_workers(_qwen3_8b_tp2_worker, 2, timeout=180) + assert not torch.equal( + before_refresh.view(torch.uint8), + after_refresh.view(torch.uint8), + ) + _assert_same_raw_bytes(after_refresh, expected)