From 5f6a9bdae28e1481a3f6a75c6eb005cc115b10aa Mon Sep 17 00:00:00 2001 From: Schatten Date: Thu, 13 Aug 2026 13:32:15 +0800 Subject: [PATCH] feat(ws2): add Qwen3 FFN orchestration and validation Signed-off-by: Schatten --- benchmarks/benchmark_qwen3_ffn.py | 292 +++++++++++++++++++ rl_engine/kernels/__init__.py | 24 ++ rl_engine/kernels/ffn.py | 382 +++++++++++++++++++++++++ tests/test_qwen3_ffn.py | 459 ++++++++++++++++++++++++++++++ 4 files changed, 1157 insertions(+) create mode 100644 benchmarks/benchmark_qwen3_ffn.py create mode 100644 rl_engine/kernels/ffn.py create mode 100644 tests/test_qwen3_ffn.py diff --git a/benchmarks/benchmark_qwen3_ffn.py b/benchmarks/benchmark_qwen3_ffn.py new file mode 100644 index 00000000..b4bff5d8 --- /dev/null +++ b/benchmarks/benchmark_qwen3_ffn.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Compare Qwen3 FFN consistent and fast paths on SM90. + +The default shape is the TP=2 rank-local Qwen3-8B FFN from the #239 +architecture. ``--intermediate-size 12288`` benchmarks the unsharded FFN; +communication is intentionally outside this PR2 benchmark. + +Examples: + python benchmarks/benchmark_qwen3_ffn.py + python benchmarks/benchmark_qwen3_ffn.py --backend triton --profile-stages + python benchmarks/benchmark_qwen3_ffn.py --intermediate-size 12288 + python benchmarks/benchmark_qwen3_ffn.py --tokens 512 +""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Callable + +import torch +from tabulate import tabulate + +from rl_engine.kernels.ffn import ( + QWEN3_8B_HIDDEN_SIZE, + QWEN3_8B_INTERMEDIATE_SIZE, + QWEN3_8B_TP2_INTERMEDIATE_SIZE, + build_qwen3_ffn, + qwen3_ffn_fp32_reference, +) +from rl_engine.testing.reference_ops import summarize_kernel_drift + + +@dataclass(frozen=True) +class FFNBenchmarkResult: + path: str + gemm_backend: str + activation_backend: str + forward_ms: float + forward_backward_ms: float + max_abs_error: float + mean_abs_error: float + peak_memory_mb: float + stage_ms: dict[str, float] + + +def _time_ms(fn: Callable[[], object], *, warmup: int, iters: int) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return float(start.elapsed_time(end)) / iters + + +def _make_case(args, path: str): + device, dtype = torch.device("cuda"), torch.bfloat16 + h, i = args.hidden_size, args.intermediate_size + generator = torch.Generator(device="cpu").manual_seed(args.seed) + + def make(shape, scale): + value = torch.randn(shape, generator=generator, dtype=torch.float32) + value.mul_(scale) + return value.to(device=device, dtype=dtype) + + gate_weight = make((h, i), args.weight_std) + up_weight = make((h, i), args.weight_std) + down_weight = make((i, h), args.weight_std) + x = make((args.tokens, h), args.input_std) + dy = make((args.tokens, h), 1.0) + backend = args.backend if path == "consistent" else "pytorch" + module = build_qwen3_ffn( + gate_weight, + up_weight, + down_weight, + path=path, + backend=backend, + ) + return module, x, dy + + +def _stage_times(module, x, *, warmup: int, iters: int) -> dict[str, float]: + with torch.no_grad(): + gate = module.gemm_op(x, module.gate_weight) + up = module.gemm_op(x, module.up_weight) + hidden = module.swiglu_op(gate, up) + closures = { + "gate_gemm": lambda: module.gemm_op(x, module.gate_weight), + "up_gemm": lambda: module.gemm_op(x, module.up_weight), + "swiglu": lambda: module.swiglu_op(gate, up), + "down_gemm": lambda: module.gemm_op(hidden, module.down_weight), + } + return { + name: _time_ms(closure, warmup=warmup, iters=iters) + for name, closure in closures.items() + } + + +def _run_path(args, path: str) -> FFNBenchmarkResult: + module, x, dy = _make_case(args, path) + + with torch.no_grad(): + candidate = module(x) + reference = qwen3_ffn_fp32_reference( + x, module.gate_weight, module.up_weight, module.down_weight + ).output + drift = summarize_kernel_drift(candidate, reference) + del candidate, reference + + def forward(): + with torch.no_grad(): + return module(x) + + def forward_backward(): + module.zero_grad(set_to_none=True) + x.grad = None + module(x).backward(dy) + return x.grad + + x.requires_grad_(True) + forward_ms = _time_ms(forward, warmup=args.warmup, iters=args.iters) + forward_backward_ms = _time_ms(forward_backward, warmup=args.warmup, iters=args.iters) + + torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + forward_backward() + torch.cuda.synchronize() + peak_memory_mb = torch.cuda.max_memory_allocated() / (1024**2) + + stage_ms = ( + _stage_times(module, x.detach(), warmup=args.warmup, iters=args.iters) + if args.profile_stages + else {} + ) + provenance = module.provenance + return FFNBenchmarkResult( + path=provenance.path, + gemm_backend=provenance.gemm_backend, + activation_backend=provenance.activation_backend, + forward_ms=forward_ms, + forward_backward_ms=forward_backward_ms, + max_abs_error=float(drift["max_abs_error"]), + mean_abs_error=float(drift["mean_abs_error"]), + peak_memory_mb=float(peak_memory_mb), + stage_ms=stage_ms, + ) + + +def _metadata(args) -> dict[str, object]: + capability = torch.cuda.get_device_capability() + return { + "model": "Qwen3-8B dense FFN", + "device": torch.cuda.get_device_name(), + "compute_capability": f"SM{capability[0]}{capability[1]}", + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "dtype": "bfloat16", + "accumulation_dtype": "float32", + "tokens": args.tokens, + "hidden_size": args.hidden_size, + "intermediate_size": args.intermediate_size, + "intermediate_scope": ( + "tp2_local" + if args.intermediate_size == QWEN3_8B_TP2_INTERMEDIATE_SIZE + else "custom_or_unsharded" + ), + "seed": args.seed, + "data_generator": "CPU FP32 MT19937, then quantized to CUDA BF16", + "input_std": args.input_std, + "weight_std": args.weight_std, + "warmup": args.warmup, + "iters": args.iters, + "scope": "rank-local FFN arithmetic; no collective communication", + } + + +def render_results(results: list[FFNBenchmarkResult]) -> str: + fast = next(result for result in results if result.path == "fast") + fast_backward_ms = fast.forward_backward_ms + rows = [] + for result in results: + backward_overhead = result.forward_backward_ms / fast_backward_ms + rows.append( + [ + result.path, + result.gemm_backend, + result.activation_backend, + f"{result.forward_ms:.3f}", + f"{result.forward_backward_ms:.3f}", + f"{result.forward_ms / fast.forward_ms:.2f}x", + f"{backward_overhead:.2f}x", + f"{result.max_abs_error:.6g}", + f"{result.peak_memory_mb:.0f}", + ] + ) + return tabulate( + rows, + headers=( + "path", + "GEMM", + "activation", + "forward ms", + "fwd+bwd ms", + "forward / fast", + "fwd+bwd / fast", + "max abs vs FP32", + "peak MB", + ), + tablefmt="github", + ) + + +def run(args) -> dict[str, object]: + if not torch.cuda.is_available(): + raise RuntimeError("Qwen3 FFN benchmark requires a CUDA SM90 GPU") + capability = torch.cuda.get_device_capability() + if capability[0] != 9: + actual_sm = f"SM{capability[0]}{capability[1]}" + raise RuntimeError(f"FFN benchmark targets SM90, got {actual_sm}") + dimensions = (args.tokens, args.hidden_size, args.intermediate_size) + if any(value <= 0 for value in dimensions): + raise ValueError("tokens and FFN dimensions must be positive") + if args.warmup < 0 or args.iters <= 0: + raise ValueError("warmup must be non-negative; iters must be positive") + + torch.backends.cuda.matmul.allow_tf32 = False + results = [_run_path(args, "consistent"), _run_path(args, "fast")] + payload = { + "metadata": _metadata(args), + "results": [asdict(result) for result in results], + } + print(json.dumps(payload["metadata"], indent=2)) + print() + print(render_results(results)) + if args.profile_stages: + print("\nStage breakdown (ms)") + stage_rows = [ + [ + result.path, + *[f"{result.stage_ms[name]:.3f}" for name in result.stage_ms], + ] + for result in results + ] + stage_names = list(results[0].stage_ms) + headers = ("path", *stage_names) + print(tabulate(stage_rows, headers=headers, tablefmt="github")) + if args.json_out is not None: + args.json_out.parent.mkdir(parents=True, exist_ok=True) + content = json.dumps(payload, indent=2) + "\n" + args.json_out.write_text(content, encoding="utf-8") + return payload + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + backends = ("cuda", "triton") + parser.add_argument("--backend", choices=backends, default="cuda") + parser.add_argument("--tokens", type=int, default=128) + hidden_default = QWEN3_8B_HIDDEN_SIZE + parser.add_argument("--hidden-size", type=int, default=hidden_default) + parser.add_argument( + "--intermediate-size", + type=int, + default=QWEN3_8B_TP2_INTERMEDIATE_SIZE, + help=( + "rank-local intermediate width; defaults to Qwen3-8B TP=2 " + f"({QWEN3_8B_TP2_INTERMEDIATE_SIZE}); use " + f"{QWEN3_8B_INTERMEDIATE_SIZE} for the unsharded FFN" + ), + ) + parser.add_argument("--seed", type=int, default=239) + parser.add_argument("--input-std", type=float, default=1.0) + parser.add_argument("--weight-std", type=float, default=0.02) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--profile-stages", action="store_true") + parser.add_argument("--json-out", type=Path) + return parser.parse_args(argv) + + +if __name__ == "__main__": + run(parse_args()) diff --git a/rl_engine/kernels/__init__.py b/rl_engine/kernels/__init__.py index e69de29b..400e64ad 100644 --- a/rl_engine/kernels/__init__.py +++ b/rl_engine/kernels/__init__.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from rl_engine.kernels.ffn import ( + QWEN3_8B_HIDDEN_SIZE, + QWEN3_8B_INTERMEDIATE_SIZE, + QWEN3_8B_TP2_INTERMEDIATE_SIZE, + Qwen3FFN, + Qwen3FFNProvenance, + Qwen3FFNStages, + build_qwen3_ffn, + qwen3_ffn_fp32_reference, +) + +__all__ = [ + "QWEN3_8B_HIDDEN_SIZE", + "QWEN3_8B_INTERMEDIATE_SIZE", + "QWEN3_8B_TP2_INTERMEDIATE_SIZE", + "Qwen3FFN", + "Qwen3FFNProvenance", + "Qwen3FFNStages", + "build_qwen3_ffn", + "qwen3_ffn_fp32_reference", +] diff --git a/rl_engine/kernels/ffn.py b/rl_engine/kernels/ffn.py new file mode 100644 index 00000000..f3311c5e --- /dev/null +++ b/rl_engine/kernels/ffn.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Qwen3 FFN orchestration for the WS2 consistency and fast paths. + +The module owns no process groups and launches no collectives. It consumes the +rank-local token tensor supplied by the distributed wrapper (after any SP +ownership conversion) and returns the local Down projection partial before the +wrapper's TP AllReduce or sequence-parallel ReduceScatter:: + + gate = GEMM(x, W_gate) + up = GEMM(x, W_up) + hidden = SwiGLU(gate, up) + output = GEMM(hidden, W_down) + +For an unsharded Qwen3-8B FFN, ``intermediate_size`` is 12288. Under TP=2 +it is 6144, so the same orchestration object is reusable by the distributed +wrappers without hiding communication inside the FFN. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +GemmCallable = Callable[[Tensor, Tensor], Tensor] +SwiGLUCallable = Callable[[Tensor, Tensor], Tensor] + +QWEN3_8B_HIDDEN_SIZE = 4096 +QWEN3_8B_INTERMEDIATE_SIZE = 12288 +QWEN3_8B_TP2_INTERMEDIATE_SIZE = 6144 + + +@dataclass(frozen=True) +class Qwen3FFNProvenance: + """The selected execution path, recorded by tests and benchmarks.""" + + path: str + gemm_backend: str + activation_backend: str + + +@dataclass(frozen=True) +class Qwen3FFNStages: + """Observable FFN boundaries used by tolerance and backward validation.""" + + gate: Tensor + up: Tensor + hidden: Tensor + output: Tensor + + +def _validate_weights(gate_weight: Tensor, up_weight: Tensor, down_weight: Tensor) -> None: + if gate_weight.ndim != 2 or up_weight.ndim != 2 or down_weight.ndim != 2: + raise ValueError("Qwen3 FFN weights must all be rank-2 GEMM matrices") + if gate_weight.shape != up_weight.shape: + raise ValueError( + "gate and up weights must share [hidden, intermediate] shape, got " + f"{tuple(gate_weight.shape)} and {tuple(up_weight.shape)}" + ) + hidden_size, intermediate_size = gate_weight.shape + if down_weight.shape != (intermediate_size, hidden_size): + expected = (intermediate_size, hidden_size) + actual = tuple(down_weight.shape) + raise ValueError( + "down weight must have [intermediate, hidden] shape " f"{expected}, got {actual}" + ) + if not gate_weight.is_floating_point(): + dtype = gate_weight.dtype + raise TypeError(f"FFN weights must be floating point, got {dtype}") + if not (gate_weight.dtype == up_weight.dtype == down_weight.dtype): + raise TypeError( + "Qwen3 FFN weights must share dtype, got " + f"{gate_weight.dtype}, {up_weight.dtype}, and {down_weight.dtype}" + ) + if not (gate_weight.device == up_weight.device == down_weight.device): + devices = (gate_weight.device, up_weight.device, down_weight.device) + raise RuntimeError(f"FFN weights must share device, got {devices}") + + +def _make_parameter(value: Tensor, *, trainable: bool) -> nn.Parameter: + if isinstance(value, nn.Parameter) and value.requires_grad == trainable: + return value + return nn.Parameter(value.detach(), requires_grad=trainable) + + +class Qwen3FFN(nn.Module): + """Backend-agnostic Qwen3 FFN over explicit GEMM-layout weights. + + Weight layouts follow ``A[M,K] @ B[K,N]`` instead of ``nn.Linear``'s + transposed storage: + + - ``gate_weight`` and ``up_weight``: ``[hidden, intermediate_local]`` + - ``down_weight``: ``[intermediate_local, hidden]`` + + The leading input dimensions are flattened into the GEMM row dimension and + restored on output. No bias is present in the Qwen3 MLP contract. + """ + + def __init__( + self, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + gemm_op: GemmCallable, + swiglu_op: SwiGLUCallable, + provenance: Qwen3FFNProvenance, + trainable: bool = True, + ) -> None: + super().__init__() + _validate_weights(gate_weight, up_weight, down_weight) + self.gate_weight = _make_parameter(gate_weight, trainable=trainable) + self.up_weight = _make_parameter(up_weight, trainable=trainable) + self.down_weight = _make_parameter(down_weight, trainable=trainable) + self.gemm_op = gemm_op + self.swiglu_op = swiglu_op + self.provenance = provenance + + @property + def hidden_size(self) -> int: + return int(self.gate_weight.shape[0]) + + @property + def intermediate_size(self) -> int: + """Return local intermediate width: 12288 full or 6144 at TP=2.""" + + return int(self.gate_weight.shape[1]) + + def _validate_input(self, x: Tensor) -> None: + if x.ndim < 2: + shape = tuple(x.shape) + raise ValueError(f"FFN input must be rank 2+, got {shape}") + if x.shape[-1] != self.hidden_size: + raise ValueError( + f"Qwen3FFN input last dimension must be {self.hidden_size}, " + f"got shape {tuple(x.shape)}" + ) + if x.numel() == 0: + raise ValueError("Qwen3FFN does not support an empty token axis") + if x.dtype != self.gate_weight.dtype: + raise TypeError( + f"Qwen3FFN input and weights must share dtype, got {x.dtype} " + f"and {self.gate_weight.dtype}" + ) + if x.device != self.gate_weight.device: + weight_device = self.gate_weight.device + raise RuntimeError( + f"FFN input/weights must share device, got " f"{x.device} and {weight_device}" + ) + + def forward_with_stages(self, x: Tensor) -> Qwen3FFNStages: + """Run the FFN and expose every arithmetic boundary for validation.""" + + self._validate_input(x) + leading_shape = x.shape[:-1] + x_2d = x.reshape(-1, self.hidden_size).contiguous() + intermediate_shape = (*leading_shape, self.intermediate_size) + output_shape = (*leading_shape, self.hidden_size) + + gate = self.gemm_op(x_2d, self.gate_weight).reshape(intermediate_shape) + up = self.gemm_op(x_2d, self.up_weight).reshape(intermediate_shape) + hidden = self.swiglu_op(gate, up) + hidden_2d = hidden.reshape(-1, self.intermediate_size).contiguous() + output_2d = self.gemm_op(hidden_2d, self.down_weight) + + return Qwen3FFNStages( + gate=gate, + up=up, + hidden=hidden, + output=output_2d.reshape(output_shape), + ) + + def forward(self, x: Tensor) -> Tensor: + return self.forward_with_stages(x).output + + +def _fast_swiglu(gate: Tensor, up: Tensor) -> Tensor: + """Framework-native fast path without an invariance claim.""" + + return F.silu(gate) * up + + +def _resolve_ops( + path: str, backend: str +) -> tuple[GemmCallable, SwiGLUCallable, Qwen3FFNProvenance]: + normalized_path = path.strip().lower().replace("-", "_") + normalized_backend = backend.strip().lower().replace("-", "_") + + if normalized_path == "fast": + if normalized_backend not in {"pytorch", "torch"}: + raise ValueError( + "the fast FFN path currently requires backend='pytorch'; " f"got {backend!r}" + ) + return ( + torch.matmul, + _fast_swiglu, + Qwen3FFNProvenance( + path="fast", + gemm_backend="pytorch.matmul", + activation_backend="torch.nn.functional.silu", + ), + ) + + if normalized_path != "consistent": + expected = "'consistent' or 'fast'" + raise ValueError(f"unsupported FFN path {path!r}; expected {expected}") + + if normalized_backend == "cuda": + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + from rl_engine.kernels.ops.cuda.activation.swiglu import SwiGLUCudaOp + from rl_engine.kernels.ops.cuda.matmul.det_gemm import DetGemmOp + + required_symbols = ( + "det_gemm_fwd", + "det_gemm_da", + "det_gemm_db", + "swiglu_forward", + "swiglu_backward", + ) + missing_symbols = ( + required_symbols + if not _EXT_AVAILABLE or _C is None + else tuple(symbol for symbol in required_symbols if not hasattr(_C, symbol)) + ) + if missing_symbols: + missing = ", ".join(missing_symbols) + raise RuntimeError( + "the CUDA consistent FFN backend requires the compiled " + f"forward/backward symbols; missing: {missing}" + ) + + gemm = DetGemmOp() + activation = SwiGLUCudaOp() + return ( + gemm, + activation, + Qwen3FFNProvenance( + path="consistent", + gemm_backend="cuda.det_gemm", + activation_backend="cuda.swiglu", + ), + ) + + if normalized_backend == "triton": + try: + from rl_engine.kernels.ops.triton.activation.swiglu import TritonSwiGLUOp + from rl_engine.kernels.ops.triton.matmul.det_gemm import TritonDetGemmOp + except ImportError as error: + raise RuntimeError("the Triton consistent FFN backend is unavailable") from error + + gemm = TritonDetGemmOp() + activation = TritonSwiGLUOp() + return ( + gemm, + activation, + Qwen3FFNProvenance( + path="consistent", + gemm_backend="triton.det_gemm", + activation_backend="triton.swiglu", + ), + ) + + raise ValueError( + "the consistent FFN path requires backend='cuda' or backend='triton'; " f"got {backend!r}" + ) + + +def build_qwen3_ffn( + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, + *, + path: str, + backend: str, + trainable: bool = True, +) -> Qwen3FFN: + """Construct an explicit consistent or fast Qwen3 FFN implementation. + + Backend resolution is fail-closed. A requested deterministic CUDA/Triton + implementation is never replaced by ``torch.matmul`` implicitly. + """ + + _validate_weights(gate_weight, up_weight, down_weight) + normalized_path = path.strip().lower().replace("-", "_") + normalized_backend = backend.strip().lower().replace("-", "_") + valid_pair = (normalized_path, normalized_backend) in { + ("consistent", "cuda"), + ("consistent", "triton"), + ("fast", "pytorch"), + ("fast", "torch"), + } + if not valid_pair: + # _resolve_ops owns the stable, user-facing error messages. + _resolve_ops(path, backend) + + if normalized_path == "consistent": + if gate_weight.dtype != torch.bfloat16: + raise TypeError( + "the consistent Qwen3 FFN path requires BF16 weights; " f"got {gate_weight.dtype}" + ) + if gate_weight.device.type != "cuda": + raise RuntimeError( + "the consistent Qwen3 FFN path requires CUDA SM90 weights; " + f"got {gate_weight.device}" + ) + capability = torch.cuda.get_device_capability(gate_weight.device) + if capability[0] != 9: + raise RuntimeError( + "the consistent Qwen3 FFN path targets SM90; " + f"got SM{capability[0]}{capability[1]}" + ) + + gemm_op, swiglu_op, provenance = _resolve_ops(path, backend) + return Qwen3FFN( + gate_weight, + up_weight, + down_weight, + gemm_op=gemm_op, + swiglu_op=swiglu_op, + provenance=provenance, + trainable=trainable, + ) + + +def qwen3_ffn_fp32_reference( + x: Tensor, + gate_weight: Tensor, + up_weight: Tensor, + down_weight: Tensor, +) -> Qwen3FFNStages: + """Uninterrupted FP32 reference over the exact quantized input values.""" + + _validate_weights(gate_weight, up_weight, down_weight) + if x.ndim < 2 or x.shape[-1] != gate_weight.shape[0]: + raise ValueError( + f"reference input must end in hidden size {gate_weight.shape[0]}, " + f"got {tuple(x.shape)}" + ) + if x.device != gate_weight.device: + raise RuntimeError("reference input and weights must share device") + + old_tf32: Any = None + if x.device.type == "cuda": + old_tf32 = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + try: + leading_shape = x.shape[:-1] + hidden_size, intermediate_size = gate_weight.shape + x_fp32 = x.float().reshape(-1, hidden_size) + intermediate_shape = (*leading_shape, intermediate_size) + gate = (x_fp32 @ gate_weight.float()).reshape(intermediate_shape) + up = (x_fp32 @ up_weight.float()).reshape(intermediate_shape) + hidden = gate * torch.sigmoid(gate) * up + output_2d = hidden.reshape(-1, intermediate_size) @ down_weight.float() + finally: + if old_tf32 is not None: + torch.backends.cuda.matmul.allow_tf32 = old_tf32 + + output_shape = (*leading_shape, hidden_size) + return Qwen3FFNStages( + gate=gate, + up=up, + hidden=hidden, + output=output_2d.reshape(output_shape), + ) + + +__all__ = [ + "QWEN3_8B_HIDDEN_SIZE", + "QWEN3_8B_INTERMEDIATE_SIZE", + "QWEN3_8B_TP2_INTERMEDIATE_SIZE", + "Qwen3FFN", + "Qwen3FFNProvenance", + "Qwen3FFNStages", + "build_qwen3_ffn", + "qwen3_ffn_fp32_reference", +] diff --git a/tests/test_qwen3_ffn.py b/tests/test_qwen3_ffn.py new file mode 100644 index 00000000..4785e14d --- /dev/null +++ b/tests/test_qwen3_ffn.py @@ -0,0 +1,459 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 #239 PR2: complete Qwen3 FFN tolerance and batch-invariance checks.""" + +from __future__ import annotations + +import pytest +import torch + +from benchmarks.benchmark_qwen3_ffn import FFNBenchmarkResult, parse_args, render_results +from rl_engine.kernels.ffn import ( + QWEN3_8B_HIDDEN_SIZE, + QWEN3_8B_INTERMEDIATE_SIZE, + QWEN3_8B_TP2_INTERMEDIATE_SIZE, + Qwen3FFN, + Qwen3FFNProvenance, + build_qwen3_ffn, + qwen3_ffn_fp32_reference, +) +from rl_engine.kernels.gtest.tolerance import load_contract +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + +_HAS_CUDA_CONSISTENT = bool( + torch.cuda.is_available() + and _EXT_AVAILABLE + and _C is not None + and all( + hasattr(_C, symbol) + for symbol in ( + "det_gemm_fwd", + "det_gemm_da", + "det_gemm_db", + "swiglu_forward", + "swiglu_backward", + ) + ) +) + +try: + import triton # noqa: F401 + + _HAS_TRITON = True +except ImportError: + _HAS_TRITON = False + +_IS_SM90 = bool(torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 9) + + +def _consistent_backends(): + return [ + pytest.param( + "cuda", + marks=pytest.mark.skipif( + not (_IS_SM90 and _HAS_CUDA_CONSISTENT), + reason=("CUDA consistent FFN tests require SM90 and compiled " "GEMM/SwiGLU ops"), + ), + ), + pytest.param( + "triton", + marks=pytest.mark.skipif( + not (_IS_SM90 and _HAS_TRITON), + reason="Triton consistent FFN tests require SM90 and Triton", + ), + ), + ] + + +def _randn(shape, *, seed: int, dtype=torch.float32, device="cpu", scale=1.0): + generator = torch.Generator(device="cpu").manual_seed(seed) + value = torch.randn(shape, generator=generator) * scale + return value.to(device=device, dtype=dtype) + + +def _weights( + hidden: int, + intermediate: int, + *, + dtype=torch.float32, + device="cpu", +): + return ( + _randn( + (hidden, intermediate), + seed=11, + dtype=dtype, + device=device, + scale=0.02, + ), + _randn( + (hidden, intermediate), + seed=12, + dtype=dtype, + device=device, + scale=0.02, + ), + _randn( + (intermediate, hidden), + seed=13, + dtype=dtype, + device=device, + scale=0.02, + ), + ) + + +def _reduction_tolerance(dtype: torch.dtype) -> tuple[float, float]: + contract = load_contract() + key = { + torch.float32: "float32", + torch.bfloat16: "bfloat16", + torch.float16: "float16", + }[dtype] + values = contract["accuracy"]["default"]["reduction"][key] + return float(values["atol"]), float(values["rtol"]) + + +def _assert_close(actual: torch.Tensor, expected: torch.Tensor, dtype: torch.dtype) -> None: + atol, rtol = _reduction_tolerance(dtype) + actual_fp32, expected_fp32 = actual.float(), expected.float() + assert_close = torch.testing.assert_close + assert_close(actual_fp32, expected_fp32, atol=atol, rtol=rtol) + + +def test_fast_ffn_constructs_complete_unsharded_qwen3_flow(): + assert QWEN3_8B_HIDDEN_SIZE == 4096 + assert QWEN3_8B_INTERMEDIATE_SIZE == 12288 + assert QWEN3_8B_TP2_INTERMEDIATE_SIZE == 6144 + + hidden, intermediate = 16, 48 + module = build_qwen3_ffn(*_weights(hidden, intermediate), path="fast", backend="pytorch") + x = _randn((2, 3, hidden), seed=20) + stages = module.forward_with_stages(x) + + assert stages.gate.shape == (2, 3, intermediate) + assert stages.up.shape == (2, 3, intermediate) + assert stages.hidden.shape == (2, 3, intermediate) + assert stages.output.shape == x.shape + assert module.hidden_size == hidden + assert module.intermediate_size == intermediate + assert module.provenance.path == "fast" + + +def test_ffn_benchmark_defaults_to_qwen3_8b_tp2_and_renders_comparison(): + args = parse_args([]) + assert args.hidden_size == QWEN3_8B_HIDDEN_SIZE + assert args.intermediate_size == QWEN3_8B_TP2_INTERMEDIATE_SIZE + assert args.seed == 239 + + fast = FFNBenchmarkResult( + path="fast", + gemm_backend="pytorch.matmul", + activation_backend="torch.nn.functional.silu", + forward_ms=1.0, + forward_backward_ms=3.0, + max_abs_error=0.01, + mean_abs_error=0.001, + peak_memory_mb=100.0, + stage_ms={}, + ) + consistent = FFNBenchmarkResult( + path="consistent", + gemm_backend="cuda.det_gemm", + activation_backend="cuda.swiglu", + forward_ms=2.0, + forward_backward_ms=6.0, + max_abs_error=0.0, + mean_abs_error=0.0, + peak_memory_mb=120.0, + stage_ms={}, + ) + table = render_results([consistent, fast]) + assert "consistent" in table and "fast" in table + assert "2.00x" in table + + +def test_same_ffn_contract_accepts_tp_local_intermediate_shard(): + # Qwen3-8B uses I=12288 globally and I_local=6144 under TP=2. Small values + # exercise the same ownership contract without model-scale allocations. + hidden, intermediate_local = 8, 6 + module = build_qwen3_ffn(*_weights(hidden, intermediate_local), path="fast", backend="pytorch") + stages = module.forward_with_stages(_randn((4, hidden), seed=21)) + assert stages.hidden.shape == (4, intermediate_local) + # This is a Down partial; the outer distributed wrapper reduces it. + assert stages.output.shape == (4, hidden) + + +def test_fast_ffn_forward_and_all_boundaries_match_fp32_reference(): + hidden, intermediate = 12, 28 + weights = _weights(hidden, intermediate) + x = _randn((2, 5, hidden), seed=22) + module = build_qwen3_ffn(*weights, path="fast", backend="pytorch") + + actual = module.forward_with_stages(x) + expected = qwen3_ffn_fp32_reference(x, *weights) + for actual_tensor, expected_tensor in zip( + (actual.gate, actual.up, actual.hidden, actual.output), + (expected.gate, expected.up, expected.hidden, expected.output), + strict=True, + ): + _assert_close(actual_tensor, expected_tensor, torch.float32) + + +def test_fast_ffn_complete_backward_matches_fp32_reference(): + hidden, intermediate = 10, 24 + weights = _weights(hidden, intermediate) + x = _randn((2, 4, hidden), seed=23).requires_grad_(True) + dy = _randn((2, 4, hidden), seed=24) + module = build_qwen3_ffn(*weights, path="fast", backend="pytorch") + actual = module.forward_with_stages(x) + for tensor in (actual.gate, actual.up, actual.hidden): + tensor.retain_grad() + actual.output.backward(dy) + + x_ref = x.detach().clone().requires_grad_(True) + weight_refs = tuple(weight.detach().clone().requires_grad_(True) for weight in weights) + expected = qwen3_ffn_fp32_reference(x_ref, *weight_refs) + for tensor in (expected.gate, expected.up, expected.hidden): + tensor.retain_grad() + expected.output.backward(dy) + + for actual_tensor, expected_tensor in zip( + ( + actual.output, + actual.gate.grad, + actual.up.grad, + actual.hidden.grad, + x.grad, + module.gate_weight.grad, + module.up_weight.grad, + module.down_weight.grad, + ), + ( + expected.output, + expected.gate.grad, + expected.up.grad, + expected.hidden.grad, + x_ref.grad, + weight_refs[0].grad, + weight_refs[1].grad, + weight_refs[2].grad, + ), + strict=True, + ): + assert actual_tensor is not None and expected_tensor is not None + _assert_close(actual_tensor, expected_tensor, torch.float32) + + +def test_ffn_rejects_invalid_shapes_dtype_device_and_backend(): + hidden, intermediate = 8, 12 + gate, up, down = _weights(hidden, intermediate) + provenance = Qwen3FFNProvenance("test", "test", "test") + + with pytest.raises(ValueError, match="gate and up weights"): + Qwen3FFN( + gate, + up[:, :-1], + down, + gemm_op=torch.matmul, + swiglu_op=lambda g, u: g * u, + provenance=provenance, + ) + with pytest.raises(ValueError, match="down weight"): + Qwen3FFN( + gate, + up, + down[:-1], + gemm_op=torch.matmul, + swiglu_op=lambda g, u: g * u, + provenance=provenance, + ) + backend_error = "requires backend='cuda' or backend='triton'" + with pytest.raises(ValueError, match=backend_error): + build_qwen3_ffn(gate, up, down, path="consistent", backend="pytorch") + with pytest.raises(ValueError, match="requires backend='pytorch'"): + build_qwen3_ffn(gate, up, down, path="fast", backend="cuda") + with pytest.raises(TypeError, match="requires BF16 weights"): + build_qwen3_ffn(gate, up, down, path="consistent", backend="cuda") + bf16_weights = tuple(weight.bfloat16() for weight in (gate, up, down)) + with pytest.raises(RuntimeError, match="requires CUDA SM90 weights"): + build_qwen3_ffn(*bf16_weights, path="consistent", backend="cuda") + + module = build_qwen3_ffn(gate, up, down, path="fast", backend="pytorch") + with pytest.raises(ValueError, match="last dimension"): + module(torch.ones(2, hidden + 1)) + with pytest.raises(TypeError, match="share dtype"): + module(torch.ones(2, hidden, dtype=torch.float64)) + with pytest.raises(ValueError, match="empty token"): + module(torch.empty(0, hidden)) + + +@pytest.mark.parametrize("backend", _consistent_backends()) +def test_consistent_ffn_qwen3_tp2_local_shape_forward_backward(backend): + hidden = QWEN3_8B_HIDDEN_SIZE + intermediate_local = QWEN3_8B_TP2_INTERMEDIATE_SIZE + dtype, device = torch.bfloat16, "cuda" + module = build_qwen3_ffn( + *_weights(hidden, intermediate_local, dtype=dtype, device=device), + path="consistent", + backend=backend, + ) + # M=32 keeps dW on the aligned SM90 path while limiting test memory/time. + x = _randn((32, hidden), seed=28, dtype=dtype, device=device) + x.requires_grad_(True) + dy = _randn((32, hidden), seed=29, dtype=dtype, device=device) + + stages = module.forward_with_stages(x) + assert stages.gate.shape == (32, intermediate_local) + assert stages.up.shape == (32, intermediate_local) + assert stages.hidden.shape == (32, intermediate_local) + assert stages.output.shape == (32, hidden) + + stages.output.backward(dy) + expected_grads = ( + (x.grad, x.shape), + (module.gate_weight.grad, (hidden, intermediate_local)), + (module.up_weight.grad, (hidden, intermediate_local)), + (module.down_weight.grad, (intermediate_local, hidden)), + ) + for gradient, expected_shape in expected_grads: + assert gradient is not None + assert gradient.shape == expected_shape + assert torch.isfinite(gradient).all() + + +@pytest.mark.parametrize("backend", _consistent_backends()) +def test_consistent_ffn_forward_backward_tolerance_against_fp32(backend): + hidden, intermediate = 128, 256 + dtype, device = torch.bfloat16, "cuda" + weights = _weights(hidden, intermediate, dtype=dtype, device=device) + module = build_qwen3_ffn(*weights, path="consistent", backend=backend) + # M=64 exercises the aligned det_gemm_db reduction path as well as dA. + x = _randn((2, 32, hidden), seed=30, dtype=dtype, device=device) + x.requires_grad_(True) + dy = _randn((2, 32, hidden), seed=31, dtype=dtype, device=device) + + actual = module.forward_with_stages(x) + for tensor in (actual.gate, actual.up, actual.hidden): + tensor.retain_grad() + actual.output.backward(dy) + + x_ref = x.detach().float().requires_grad_(True) + weight_refs = tuple(weight.detach().float().requires_grad_(True) for weight in weights) + expected = qwen3_ffn_fp32_reference(x_ref, *weight_refs) + for tensor in (expected.gate, expected.up, expected.hidden): + tensor.retain_grad() + expected.output.backward(dy.float()) + + for actual_tensor, expected_tensor in zip( + ( + actual.gate, + actual.up, + actual.hidden, + actual.output, + actual.gate.grad, + actual.up.grad, + actual.hidden.grad, + x.grad, + module.gate_weight.grad, + module.up_weight.grad, + module.down_weight.grad, + ), + ( + expected.gate, + expected.up, + expected.hidden, + expected.output, + expected.gate.grad, + expected.up.grad, + expected.hidden.grad, + x_ref.grad, + weight_refs[0].grad, + weight_refs[1].grad, + weight_refs[2].grad, + ), + strict=True, + ): + assert actual_tensor is not None and expected_tensor is not None + _assert_close(actual_tensor, expected_tensor, dtype) + + +def _run_activation_backward(module, x, dy): + x_run = x.detach().clone().requires_grad_(True) + stages = module.forward_with_stages(x_run) + for tensor in (stages.gate, stages.up, stages.hidden): + tensor.retain_grad() + stages.output.backward(dy) + return ( + stages.output.detach(), + stages.gate.grad.detach(), + stages.up.grad.detach(), + stages.hidden.grad.detach(), + x_run.grad.detach(), + ) + + +@pytest.mark.parametrize("backend", _consistent_backends()) +def test_consistent_ffn_batch_chunk_padding_invariance(backend): + hidden, intermediate = 128, 256 + dtype, device = torch.bfloat16, "cuda" + module = build_qwen3_ffn( + *_weights(hidden, intermediate, dtype=dtype, device=device), + path="consistent", + backend=backend, + trainable=False, + ) + x = _randn((8, hidden), seed=40, dtype=dtype, device=device) + dy = _randn((8, hidden), seed=41, dtype=dtype, device=device) + + full = _run_activation_backward(module, x, dy) + singleton = _run_activation_backward(module, x[:1], dy[:1]) + chunked_parts = [ + _run_activation_backward(module, x[:3], dy[:3]), + _run_activation_backward(module, x[3:6], dy[3:6]), + _run_activation_backward(module, x[6:], dy[6:]), + ] + chunked = tuple( + torch.cat([part[index] for part in chunked_parts], dim=0) for index in range(len(full)) + ) + + padding_x = _randn((3, hidden), seed=42, dtype=dtype, device=device) + padding_dy = _randn((3, hidden), seed=43, dtype=dtype, device=device) + padded = _run_activation_backward( + module, torch.cat((x, padding_x)), torch.cat((dy, padding_dy)) + ) + + for full_tensor, singleton_tensor, chunked_tensor, padded_tensor in zip( + full, singleton, chunked, padded, strict=True + ): + assert torch.equal(full_tensor[:1], singleton_tensor) + assert torch.equal(full_tensor, chunked_tensor) + assert torch.equal(full_tensor, padded_tensor[: x.shape[0]]) + + +@pytest.mark.parametrize("backend", _consistent_backends()) +def test_consistent_ffn_weight_gradients_repeat_bitwise(backend): + # dW reduces over tokens. PR2 claims repeat determinism and FP32 tolerance, + # not bitwise equality after changing token partition/reduction boundaries. + hidden, intermediate = 128, 256 + dtype, device = torch.bfloat16, "cuda" + module = build_qwen3_ffn( + *_weights(hidden, intermediate, dtype=dtype, device=device), + path="consistent", + backend=backend, + ) + x = _randn((64, hidden), seed=50, dtype=dtype, device=device) + dy = _randn((64, hidden), seed=51, dtype=dtype, device=device) + + def run(): + module.zero_grad(set_to_none=True) + module(x).backward(dy) + grads = (parameter.grad for parameter in module.parameters()) + return tuple(grad.detach().clone() for grad in grads) + + expected = run() + for _ in range(3): + actual = run() + pairs = zip(expected, actual, strict=True) + assert all(torch.equal(expected_grad, actual_grad) for expected_grad, actual_grad in pairs)