diff --git a/.github/workflows/qwenimage-rope.yml b/.github/workflows/qwenimage-rope.yml new file mode 100644 index 000000000..7819acb9d --- /dev/null +++ b/.github/workflows/qwenimage-rope.yml @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +name: Qwen-Image RoPE + +on: + pull_request: + branches: [test-qwenimage] + types: [opened, synchronize, reopened, ready_for_review] + workflow_dispatch: + +permissions: + contents: read + +jobs: + cpu-contracts: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ['3.10', '3.12'] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - name: Install CPU reference dependencies + run: | + python -m pip install pip==25.3 setuptools==80.9.0 wheel==0.45.1 + python -m pip install torch==2.9.1 --index-url https://download.pytorch.org/whl/cpu + python -m pip install --no-build-isolation -e . pytest + - name: Run CPU contracts, input registration and build-profile checks + env: + OMP_NUM_THREADS: '1' + MKL_NUM_THREADS: '1' + run: | + python -m pytest -q -rs \ + tests/ops/rope/test_multi_axis_rope.py \ + tests/validation/operators/test_operator_inputs.py \ + tests/build/test_multi_axis_rope_build.py \ + --junitxml=artifacts/rope-cpu.xml + - uses: actions/upload-artifact@v4 + if: always() + with: + name: rope-cpu-${{ matrix.python-version }} + path: artifacts/rope-cpu.xml + + static-and-docs: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + cache: pip + - run: python -m pip install pre-commit==4.3.0 -r requirements/docs.txt + - name: Check changed files + run: | + python - <<'PY' + import json + import os + import re + import subprocess + import sys + + with open(os.environ['GITHUB_EVENT_PATH'], encoding='utf-8') as stream: + event = json.load(stream) + base = event.get('pull_request', {}).get('base', {}).get('sha') + if base: + assert re.fullmatch(r'[0-9a-f]{40}', base) + else: + base = 'HEAD^' + changed = subprocess.check_output([ + 'git', 'diff', '--name-only', '--diff-filter=ACMR', '-z', base, 'HEAD' + ]).decode('utf-8').split('\0') + files = [('./' + path if path.startswith('-') else path) for path in changed if path] + if files: + subprocess.check_call([sys.executable, '-m', 'pre_commit', 'run', '--files', *files]) + PY + - name: Build documentation + run: mkdocs build --strict -f mkdocs.yaml diff --git a/MANIFEST.in b/MANIFEST.in index 5ce1ee23f..67165e195 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -2,3 +2,4 @@ include pyproject.toml setup.py README.md LICENSE recursive-include build_tools *.py recursive-include csrc *.cpp *.cu *.cuh *.h *.hip *.asc recursive-include rl_engine *.json *.jsonl *.sh *.patch *.pyi +recursive-include rl_engine FLASH_ATTENTION_LICENSE diff --git a/benchmarks/operators/rope/benchmark_multi_axis_rope.py b/benchmarks/operators/rope/benchmark_multi_axis_rope.py new file mode 100644 index 000000000..316a75ba7 --- /dev/null +++ b/benchmarks/operators/rope/benchmark_multi_axis_rope.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Standalone RoPE: contiguous/packed QKV, eager/graph, optional pinned FlashAttention.""" + +from __future__ import annotations + +import argparse +import ast +import gc +import hashlib +import importlib.util +import json +import platform +import statistics +import subprocess +import sys +import time +from pathlib import Path +from typing import Optional, Union + +import torch + +ROOT = Path(__file__).resolve().parents[3] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from rl_engine.ops.rope.multi_axis_rope import ( # noqa: E402 + build_multi_axis_cos_sin, + get_multi_axis_rope_op, + qwen_image_positions, +) + +FLASH_REVISION = "98c1cee0ed59307833841cdbf567fe336ceec13a" +FLASH_SHA256 = "d2037e623854c702fb075784fe4292d4a89d6b9ac1ae33cab3d6788f8f94f240" +FLASH_LAYER_SHA256 = "a79cb42dfc33e126227bdba27ae7a024d7d5cc5c8fd7f7a5f26c68d501024e0a" + + +def upstream_complex(x, cis): + """Diffusers use_real=False arithmetic; table assembly is outside timing.""" + z = torch.view_as_complex(x.float().reshape(*x.shape[:-1], 64, 2)) + return torch.view_as_real(z * cis.unsqueeze(-2)).flatten(-2).to(x.dtype) + + +def load_flash_baselines(source_path, layer_path): + """Load checksum-pinned upstream kernel and autograd class without installing flash-attn.""" + source = source_path.read_text(encoding="utf-8") + layer = layer_path.read_text(encoding="utf-8") + for text, expected in ((source, FLASH_SHA256), (layer, FLASH_LAYER_SHA256)): + if hashlib.sha256(text.encode()).hexdigest() != expected: + raise ValueError("FlashAttention source checksum mismatch (normalized UTF-8 text)") + node = next(n for n in ast.parse(layer).body if getattr(n, "name", None) == "ApplyRotaryEmb") + directory = ROOT / "artifacts/flash_rope_baseline" + directory.mkdir(parents=True, exist_ok=True) + result = {} + assert source.count("BLOCK_H=2,") == 1 + for variant in ("default", "no_fma"): + code = source + if variant == "no_fma": + code = code.replace("BLOCK_H=2,", "BLOCK_H=2,\n enable_fp_fusion=False,") + path = directory / f"flash_rotary_{variant}.py" + path.write_text(code, encoding="utf-8") + name = f"rlk_flash_baseline_{variant}" + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + namespace = { + "__name__": name + "_autograd", + "torch": torch, + "Tensor": torch.Tensor, + "Optional": Optional, + "Union": Union, + "apply_rotary": module.apply_rotary, + } + exec(compile(ast.Module(body=[node], type_ignores=[]), str(layer_path), "exec"), namespace) + cls = namespace["ApplyRotaryEmb"] + result[f"flash_{variant}"] = lambda x, c, s, cls=cls: cls.apply(x, c, s, True) + return result + + +def measure(fn, args): + for _ in range(args.warmup): + result = fn() + torch.cuda.synchronize() + del result + wall = [] + for _ in range(args.trials): + start = time.perf_counter() + for _ in range(args.iterations): + result = fn() + torch.cuda.synchronize() + wall.append((time.perf_counter() - start) * 1000 / args.iterations) + del result + + # Collect graph owners before capture, so GC cannot destroy another graph + # during stream capture. Preserve the caller's GC state. + gc.collect() + enabled = gc.isenabled() + gc.disable() + graph = torch.cuda.CUDAGraph() + try: + with torch.cuda.graph(graph, stream=torch.cuda.Stream()): + for _ in range(args.graph_calls): + result = fn() + for _ in range(args.warmup): + graph.replay() + events = [ + (torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)) + for _ in range(args.iterations) + ] + for start, end in events: + start.record() + graph.replay() + end.record() + torch.cuda.synchronize() + times = [start.elapsed_time(end) / args.graph_calls for start, end in events] + del graph, result + finally: + if enabled: + gc.enable() + + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + result = fn() + torch.cuda.synchronize() + extra = torch.cuda.max_memory_allocated() - before + del result + return { + "graph_median_ms": statistics.median(times), + "graph_min_ms": min(times), + "graph_max_ms": max(times), + "eager_wall_median_ms": statistics.median(wall), + "eager_wall_trials_ms": wall, + "extra_peak_bytes": extra, + } + + +def byte_differences(a, b): + return int((a.contiguous().view(torch.uint8) != b.contiguous().view(torch.uint8)).sum()) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--iterations", type=int, default=30) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--trials", type=int, default=3) + parser.add_argument("--graph-calls", type=int, default=20) + parser.add_argument("--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16") + parser.add_argument( + "--layouts", + nargs="+", + choices=["contiguous", "packed_qkv"], + default=["contiguous", "packed_qkv"], + ) + parser.add_argument("--flash-source", type=Path) + parser.add_argument("--flash-layer-source", type=Path) + args = parser.parse_args() + if min(args.iterations, args.warmup, args.trials, args.graph_calls) < 1: + parser.error("timing counts must be positive") + if bool(args.flash_source) != bool(args.flash_layer_source): + parser.error("both FlashAttention source files are required") + dtype = getattr(torch, args.dtype) + torch.manual_seed(386) + ours = {name: get_multi_axis_rope_op(name, device="cuda") for name in ("cuda", "triton")} + flash = ( + load_flash_baselines(args.flash_source, args.flash_layer_source) + if args.flash_source + else {} + ) + reference = get_multi_axis_rope_op("reference") + rows = [] + correctness = [] + for height, width in ((2, 3), (64, 64), (83, 83), (104, 58)): + positions, _ = qwen_image_positions(height, width, 0) + cos, sin = build_multi_axis_cos_sin(positions, device="cuda") + cis = torch.complex(cos, sin) + + def upstream_op(x, _cos, _sin, cis=cis): + return upstream_complex(x, cis) + + ops = {"upstream_complex": upstream_op, **ours, **flash} + for layout in args.layouts: + if layout == "packed_qkv": + packed = torch.randn(1, height * width, 3, 24, 128, device="cuda", dtype=dtype) + x = packed[:, :, 0].detach().requires_grad_() + else: + x = torch.randn( + 1, height * width, 24, 128, device="cuda", dtype=dtype, requires_grad=True + ) + grad = torch.randn_like(x) + gold = reference(x, cos, sin) + gold_dx = torch.autograd.grad(gold, x, grad)[0] + for name, op in ops.items(): + result = op(x, cos, sin) + dx = torch.autograd.grad(result, x, grad)[0] + diff = [byte_differences(result, gold), byte_differences(dx, gold_dx)] + if name in ("cuda", "triton", "flash_no_fma"): + assert diff == [0, 0], (name, layout, x.shape, diff) + else: + tol = {torch.float32: 2e-6, torch.float16: 0.004, torch.bfloat16: 0.04}[dtype] + torch.testing.assert_close(result, gold, atol=tol, rtol=tol) + torch.testing.assert_close(dx, gold_dx, atol=tol, rtol=tol) + correctness.append( + { + "backend": name, + "layout": layout, + "shape": list(x.shape), + "different_bytes_y_dx": diff, + } + ) + del result, dx, gold, gold_dx + for name, op in ops.items(): + for phase in ("forward", "forward_backward"): + + def call(op=op, phase=phase, x=x, cos=cos, sin=sin, grad=grad): + if phase == "forward": + with torch.no_grad(): + return op(x, cos, sin) + return torch.autograd.grad(op(x, cos, sin), x, grad)[0] + + stats = measure(call, args) + rows.append( + { + "backend": name, + "phase": phase, + "layout": layout, + "shape": list(x.shape), + "input_stride": list(x.stride()), + **stats, + } + ) + print( + f"{height}x{width} {layout:11s} {name:16s} {phase:16s} " + f"graph={stats['graph_median_ms']:.5f}ms " + f"eager={stats['eager_wall_median_ms']:.5f}ms", + flush=True, + ) + sources = [ + "csrc/cuda/rope/multi_axis_rope.cu", + "rl_engine/backends/shared/triton/rope/_flash_attn_rotary.py", + "rl_engine/backends/shared/triton/rope/multi_axis_rope.py", + "rl_engine/backends/cuda/rope/multi_axis_rope.py", + "rl_engine/ops/rope/_registration.py", + "rl_engine/ops/rope/multi_axis_rope.py", + "rl_engine/reference/rope/multi_axis_rope.py", + "benchmarks/operators/rope/benchmark_multi_axis_rope.py", + ] + import triton + + evidence = { + "gpu": torch.cuda.get_device_name(), + "capability": torch.cuda.get_device_capability(), + "os": platform.platform(), + "python": sys.version, + "torch": torch.__version__, + "cuda_runtime": torch.version.cuda, + "triton": triton.__version__, + "dtype": args.dtype, + "base_commit": subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=ROOT, text=True + ).strip(), + "source_sha256": {p: hashlib.sha256((ROOT / p).read_bytes()).hexdigest() for p in sources}, + "timing": "CUDA graph ms/call and synchronized eager wall ms/call; warm allocation; " + "excludes tables, compilation, and initial allocation of inputs", + "memory": "extra peak allocated bytes excluding inputs/tables/grad; includes returned " + "output or gradient; packed QKV backing storage allocated before timing", + "iterations": args.iterations, + "warmup": args.warmup, + "trials": args.trials, + "graph_calls": args.graph_calls, + "flash_baseline": { + "included": bool(flash), + "revision": FLASH_REVISION, + "kernel_sha256_normalized": FLASH_SHA256, + "layer_sha256_normalized": FLASH_LAYER_SHA256, + "mode": "upstream rotary_kernel, apply_rotary and ApplyRotaryEmb, interleaved=True", + "no_fma_change": "only add enable_fp_fusion=False at the original kernel launch", + }, + "trace": {name: op.trace() for name, op in ours.items()}, + "qualification": "Exact bytes vs unfused FP32 gold for CUDA, Triton and flash_no_fma; " + "declared tolerance for upstream complex and default FlashAttention", + "correctness": correctness, + "results": rows, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(evidence, indent=2) + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/build_tools/extensions.py b/build_tools/extensions.py index db079e828..2f04f5cdd 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -105,6 +105,7 @@ def _native_extension_required() -> bool: """Whether the caller explicitly requested a native extension build.""" return ( envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or envs.env_flag("RL_KERNEL_ROPE_ONLY") or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) or envs.env_flag("FORCE_CUDA") @@ -182,6 +183,33 @@ def get_extensions(): torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") is_rocm = getattr(torch.version, "hip", None) is not None + # A normal wheel build for the independently usable RoPE library. This + # explicit profile avoids compiling unrelated Hopper/IPC operators and is + # also useful on Windows. Default builds still include RoPE in the full _C. + if envs.env_flag("RL_KERNEL_ROPE_ONLY"): + if is_rocm or getattr(torch.version, "cuda", None) is None: + raise RuntimeError("RL_KERNEL_ROPE_ONLY requires CUDA-enabled PyTorch") + from torch.utils.cpp_extension import CUDA_HOME + + library_dirs = [] + if os.name == "nt" and CUDA_HOME and (Path(CUDA_HOME) / "lib/cudart.lib").is_file(): + library_dirs.append(str(Path(CUDA_HOME) / "lib")) + return [ + GpuExtension( + name="rl_engine._rope_cuda", + sources=[ + "csrc/bindings/rope_module.cpp", + "csrc/cuda/rope/multi_axis_rope.cu", + ], + library_dirs=library_dirs, + extra_compile_args={ + "cxx": ["/O2", "/std:c++17"] if os.name == "nt" else ["-O3", "-std=c++17"], + "nvcc": ["-O3", "-std=c++17", "--fmad=false"], + }, + extra_link_args=[] if os.name == "nt" else list(torch_rpath), + ) + ] + if _musa_build_available(torch): extensions.append( GpuExtension( @@ -233,6 +261,7 @@ def get_extensions(): ] ) else: + cuda_sources.append("csrc/cuda/rope/multi_axis_rope.cu") # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. cuda_sources.append("csrc/cuda/collectives/deterministic_collective.cu") diff --git a/ci/scripts/run_qwenimage_rope_ws1.sh b/ci/scripts/run_qwenimage_rope_ws1.sh new file mode 100644 index 000000000..447e93f14 --- /dev/null +++ b/ci/scripts/run_qwenimage_rope_ws1.sh @@ -0,0 +1,89 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# Operator-only gate: requires an ordinary, complete rl_engine._C installation. +set -euo pipefail + +ROPE_ROOT="$(cd "$(dirname "$0")/../.." && pwd)" +cd "$ROPE_ROOT" +ROPE_PYTHON="${ROPE_PYTHON:-python3}" +ROPE_WS1_OUTPUT="${ROPE_WS1_OUTPUT:-${TMPDIR:-/tmp}/rlk-qwenimage-rope}" +mkdir -p "$ROPE_WS1_OUTPUT" +export RL_KERNEL_REQUIRE_EXT=1 + +"$ROPE_PYTHON" - "$ROPE_WS1_OUTPUT/environment.json" <<'PY' +import hashlib +import json +import platform +import sys +from pathlib import Path + +import torch +import triton + +from build_tools.envs import env_flag +from rl_engine import _C + +if env_flag("RL_KERNEL_ROPE_ONLY"): + raise SystemExit("The complete _C build is required; unset RL_KERNEL_ROPE_ONLY") +if env_flag("KERNEL_ALIGN_USE_FAST_MATH"): + raise SystemExit("This numerical profile does not qualify fast-math builds") +if torch.version.hip is not None or not torch.cuda.is_available(): + raise SystemExit("A real NVIDIA CUDA device is required") +if not torch.cuda.is_bf16_supported(): + raise SystemExit("The WS1 BF16 profile requires BF16-capable hardware") +if not hasattr(torch.ops.rl_kernel_native, "multi_axis_rope"): + raise SystemExit("The imported _C library does not register multi_axis_rope") +library = Path(_C.__file__).resolve() +report = { + "scope": "Qwen-Image multi_axis_rope operator; ordinary full _C build", + "python": platform.python_version(), + "platform": platform.platform(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "triton": triton.__version__, + "gpu": torch.cuda.get_device_name(), + "capability": list(torch.cuda.get_device_capability()), + "native_library": str(library), + "native_library_sha256": hashlib.sha256(library.read_bytes()).hexdigest(), +} +Path(sys.argv[1]).write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") +print(json.dumps(report)) +PY + +"$ROPE_PYTHON" -m pytest tests/ops/rope/test_multi_axis_rope.py -q -rs \ + --junitxml="$ROPE_WS1_OUTPUT/operator-tests.xml" + +# GPU/backward/compile cases must execute. Only the two-device guard may skip +# on a single-GPU runner; missing CUDA/Triton support is not acceptance evidence. +"$ROPE_PYTHON" - "$ROPE_WS1_OUTPUT/operator-tests.xml" <<'PY' +import sys +import xml.etree.ElementTree as ET + +root = ET.parse(sys.argv[1]).getroot() +unexpected = [case.get("name") for case in root.iter("testcase") + if case.find("skipped") is not None and case.get("name") != "test_device_guard"] +if unexpected: + raise SystemExit(f"Required RoPE tests skipped: {unexpected}") +PY + +for backend in cuda triton; do + "$ROPE_PYTHON" tools/validation/operators/check_operator.py \ + --op multi_axis_rope --candidate "$backend" --device cuda --dtype bf16 \ + --batch 2 --seq 7 --n-heads 3 --check-grad --json \ + > "$ROPE_WS1_OUTPUT/operator-$backend.json" + "$ROPE_PYTHON" tools/validation/operators/validate_multi_axis_rope.py \ + --ws1 "$backend" --output "$ROPE_WS1_OUTPUT/ws1-$backend.json" +done + +"$ROPE_PYTHON" - "$ROPE_WS1_OUTPUT" <<'PY' +import json +import sys +from pathlib import Path + +root = Path(sys.argv[1]) +for backend in ("cuda", "triton"): + report = json.loads((root / f"ws1-{backend}.json").read_text(encoding="utf-8")) + if not report["passed"] or report["trace"]["actual_backend"] != f"rlkernel.{backend}.qwenimage_multi_axis_rope.v1": + raise SystemExit(f"WS1 failed or selected an unexpected backend: {backend}") +print("Qwen-Image RoPE full-extension operator gate passed") +PY diff --git a/csrc/bindings/rope_module.cpp b/csrc/bindings/rope_module.cpp new file mode 100644 index 000000000..c12e5b90d --- /dev/null +++ b/csrc/bindings/rope_module.cpp @@ -0,0 +1,8 @@ +// SPDX-License-Identifier: Apache-2.0 +// Python module entry point for the optional RoPE-only wheel. Operators are +// registered by multi_axis_rope.cu through the PyTorch dispatcher. +#include + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.doc() = "Qwen-Image RoPE dispatcher library"; +} diff --git a/csrc/cuda/rope/multi_axis_rope.cu b/csrc/cuda/rope/multi_axis_rope.cu new file mode 100644 index 000000000..27459c5fb --- /dev/null +++ b/csrc/cuda/rope/multi_axis_rope.cu @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// Qwen-Image adjacent-pair rotation. Reads strided BSHD inputs/tables directly. +#include +#include +#include +#include +#include +#include +#include + +namespace { +constexpr int64_t kHeadDim = 128; +constexpr int64_t kPairsPerHead = kHeadDim / 2; + +struct Strides { + int64_t xb, xs, xh, xd; + int64_t cb, cs, cd, sb, ss, sd; +}; + +template +__global__ void rotate_pairs(const scalar_t* x, const float* cos, const float* sin, + scalar_t* y, int64_t n, int64_t seq, int64_t heads, + Strides strides, bool transpose) { + const int64_t step = static_cast(blockDim.x) * gridDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + i < n; i += step) { + const int64_t pair = i % kPairsPerHead; + const int64_t row = i / kPairsPerHead; + const int64_t head = row % heads; + const int64_t token = row / heads; + const int64_t batch = token / seq; + const int64_t position = token % seq; + const int64_t xi = batch * strides.xb + position * strides.xs + + head * strides.xh + 2 * pair * strides.xd; + const float c = cos[batch * strides.cb + position * strides.cs + pair * strides.cd]; + const float sn = sin[batch * strides.sb + position * strides.ss + pair * strides.sd]; + const float s = transpose ? -sn : sn; + const float a = static_cast(x[xi]); + const float b = static_cast(x[xi + strides.xd]); + y[2 * i] = static_cast(__fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s))); + y[2 * i + 1] = static_cast(__fadd_rn(__fmul_rn(b, c), __fmul_rn(a, s))); + } +} + +at::Tensor multi_axis_rope_apply(const at::Tensor& x, const at::Tensor& cos, + const at::Tensor& sin, bool transpose) { + TORCH_CHECK(x.is_cuda() && cos.is_cuda() && sin.is_cuda(), "multi_axis_rope: CUDA tensors required"); + TORCH_CHECK(x.device() == cos.device() && x.device() == sin.device(), "multi_axis_rope: device mismatch"); + TORCH_CHECK(x.dim() == 4 && x.size(3) == kHeadDim && x.size(2) > 0, + "multi_axis_rope: x must be [B, S, H, 128], H > 0"); + TORCH_CHECK(x.scalar_type() == at::kFloat || x.scalar_type() == at::kHalf || + x.scalar_type() == at::kBFloat16, "multi_axis_rope: unsupported input dtype"); + TORCH_CHECK(cos.scalar_type() == at::kFloat && sin.scalar_type() == at::kFloat, + "multi_axis_rope: tables must be fp32"); + TORCH_CHECK(cos.sizes() == sin.sizes(), "multi_axis_rope: table shape mismatch"); + TORCH_CHECK((cos.dim() == 2 && cos.size(0) == x.size(1) && cos.size(1) == kPairsPerHead) || + (cos.dim() == 3 && cos.size(0) == x.size(0) && cos.size(1) == x.size(1) && cos.size(2) == kPairsPerHead), + "multi_axis_rope: tables must be [S, 64] or [B, S, 64]"); + TORCH_CHECK(!cos.requires_grad() && !sin.requires_grad(), "multi_axis_rope: table gradients unsupported"); + const c10::cuda::CUDAGuard guard(x.device()); + auto y = at::empty(x.sizes(), x.options()); + const int64_t n = x.numel() / 2; + if (n == 0) return y; + const Strides strides{x.stride(0), x.stride(1), x.stride(2), x.stride(3), + cos.dim() == 3 ? cos.stride(0) : 0, cos.stride(-2), cos.stride(-1), + sin.dim() == 3 ? sin.stride(0) : 0, sin.stride(-2), sin.stride(-1)}; + const int blocks = static_cast(std::min((n + 255) / 256, 65535)); + auto stream = c10::cuda::getCurrentCUDAStream(x.get_device()); + AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, + x.scalar_type(), "multi_axis_rope", [&] { + rotate_pairs<<>>(x.data_ptr(), cos.data_ptr(), + sin.data_ptr(), y.data_ptr(), n, x.size(1), x.size(2), strides, transpose); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return y; +} +} // namespace + +// The focused validator and normal package build load this same registration. +TORCH_LIBRARY_FRAGMENT(rl_kernel_native, m) { + m.def("multi_axis_rope(Tensor x, Tensor cos, Tensor sin, bool transpose) -> Tensor"); +} +TORCH_LIBRARY_IMPL(rl_kernel_native, CUDA, m) { + m.impl("multi_axis_rope", &multi_axis_rope_apply); +} diff --git a/docs/.nav.yml b/docs/.nav.yml index 29b5ead37..d03ade331 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -29,6 +29,7 @@ nav: - operators/sampling.md - operators/det-gemm.md - operators/embedding.md + - operators/multi-axis-rope.md - Developer Guide: - contributing/README.md - Contributor Guide: contributing/contributor-guide.md diff --git a/docs/operators/README.md b/docs/operators/README.md index 00f4cbb45..2d3b27d5d 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -26,6 +26,7 @@ Every operator page should include: - [Fused Linear LogP TP Test Runbook](linear-logp-tp-test.md) - [GRPO Loss](grpo-loss.md) - [RoPE](rope.md) +- [Qwen-Image Multi-Axis RoPE](multi-axis-rope.md) - [LM Head](lm_head.md) - [Policy Ratio + KL Penalty](ratio-kl.md) - [Pack and Pad](pack-and-pad.md) diff --git a/docs/operators/multi-axis-rope.md b/docs/operators/multi-axis-rope.md new file mode 100644 index 000000000..a7c8f6495 --- /dev/null +++ b/docs/operators/multi-axis-rope.md @@ -0,0 +1,233 @@ +# Qwen-Image Multi-Axis RoPE + +This operator implements the `multi_axis_rope` CUDA row of +[RFC #386](https://github.com/RL-Align/RL-Kernel/issues/386). It rotates **adjacent +channel pairs** with three independently indexed axes of widths `[16, 56, 56]`. +The existing GPT-NeoX rotate-half operator has different semantics. + +## Reference and scope + +Semantics are pinned to Diffusers +[`40f63690da440e8c382ab53e67508835deb1ed2f`](https://github.com/huggingface/diffusers/blob/40f63690da440e8c382ab53e67508835deb1ed2f/src/diffusers/models/transformers/transformer_qwenimage.py): +`QwenEmbedRope(scale_rope=True)` and `apply_rotary_emb_qwen(use_real=False)`. +The position helper covers single-frame text-to-image at packed latent resolution. +For a 1024×1024 image the packed grid is 64×64. Multi-image editing, video, +QwenEmbedLayer3DRope, distributed execution and full-model integration are outside +this operator's qualification. + +Image coordinates are `(0, h, w)` in row-major order. The centered coordinate +range is `[-ceil(size/2), floor(size/2))`, so an odd grid has one extra negative +coordinate. Text starts at `max(height//2, width//2)` on all three axes. The helper +returns image and text positions separately: the caller determines stream order. + +## API + +```python +import torch +from rl_engine.ops.rope.multi_axis_rope import ( + qwen_image_positions, build_multi_axis_cos_sin, get_multi_axis_rope_op, +) + +image_pos, text_pos = qwen_image_positions(64, 64, 128) +cos, sin = build_multi_axis_cos_sin(image_pos, device="cuda") +x = torch.randn(1, 4096, 24, 128, device="cuda", dtype=torch.bfloat16, + requires_grad=True) +op = get_multi_axis_rope_op("cuda", device=x.device) +y = op(x, cos, sin) +y.float().square().mean().backward() +print(op.trace()) +``` + +| Argument | Contract | +| --- | --- | +| `x` | `[B, S, H, 128]`, H > 0; FP16, BF16 or FP32 | +| `cos`, `sin` | Matching FP32 `[S, 64]` shared tables or `[B, S, 64]` per-sample tables | +| Devices | Input and tables on the same device | +| Layout | BSHD, adjacent pairs; input and independent table strides read directly | +| Output | Input shape/dtype/device; contiguous | +| Empty input | B=0 and/or S=0 supported, including backward | +| Gradients | Input gradients and higher-order input VJPs; table gradients rejected | +| Position builder | int32/int64 `[S, 3]` or `[B, S, 3]`, within `[-4096, 4095]` | + +Table construction follows the upstream CPU FP32 `pow` and `polar` operations, +then copies the results to the target device. Build and retain tables before CUDA +Graph capture. Reusing these tables also prevents per-batch transcendental drift. +The CPU factories select their device explicitly, including inside a CUDA or +meta default-device context. Packed QKV views require no materializing copy. + +## Numerical contract + +Contract ID: `qwenimage-interleaved-rope-fp32-v1`. + +For each adjacent pair `(a, b)`, compute `(a*c - b*s, b*c + a*s)` in FP32 with +separate rounded multiplications and additions, followed by one output cast. +There is no cross-element reduction, split-K, atomic accumulation or TF32. +CUDA uses round-to-nearest arithmetic intrinsics; Triton disables FP contraction. +Backward uses the transpose rotation (negate sine), including recursive VJPs. + +CUDA and Triton must agree in raw bytes within a pinned environment. Compare +batch placement, token chunking, padding and graph replay the same way. +Independent FP64 complex arithmetic provides an accuracy oracle. Upstream's GPU +complex multiplication may contract arithmetic, so its comparison uses declared +tolerances, not an unverified bitwise claim. Representative test tolerances are +`8e-7` for FP32, `0.002` for FP16 and `0.02` for BF16 (absolute and relative). +Non-finite propagation follows the explicit arithmetic; NaN payload equality is +not a cross-backend guarantee. + +## Backends and dispatch + +- `cuda`: production function in `csrc/cuda/rope/multi_axis_rope.cu`, registered + in `rl_engine._C` by ordinary CUDA builds. Does not require Hopper-specific PTX. +- `triton`: adapts FlashAttention's existing interleaved rotary kernel, with + explicit shared/per-sample table strides and `enable_fp_fusion=False`. + The upstream tile geometry is retained; this profile is limited to NVIDIA CUDA. +- `reference`: PyTorch FP32 arithmetic, including CPU. +- `auto`: CUDA → Triton → reference on NVIDIA, reference on other platforms. + +Explicit `cuda`/`triton` requests fail when unavailable; they never silently switch +implementations. Strict callers select an explicit backend and retain `op.trace()`. +`auto` is for diagnostic use and is not a strict-route certification. The legacy +registry also exposes `kernel_registry.get_op("multi_axis_rope", device=...)`. +This PR does not add an Ascend backend or certify ROCm. Existing Ascend work in +PR #411 can share semantics once its reference contract is aligned. + +Both GPU paths expose a PyTorch custom operator with fake-tensor and autograd +registration. `torch.compile(op, fullgraph=True)` supports forward and backward. +The native implementation is registered through `TORCH_LIBRARY`; Python does +not attach symbols to an unrelated extension at validation time. + +## Related work and reuse + +The Triton kernel comes from FlashAttention +[`98c1cee`](https://github.com/Dao-AILab/flash-attention/blob/98c1cee0ed59307833841cdbf567fe336ceec13a/flash_attn/ops/triton/rotary.py), +with its BSD-3-Clause copyright and full license retained beside the adapted +source and included in wheels/source distributions. Its default fused arithmetic +can differ in low bits; disabling contraction gives the declared profile. +The benchmark loads the original kernel and autograd wrapper from checksum-pinned +source files to compare the default and contraction-disabled variants. + +[Ascend PR #411](https://github.com/RL-Align/RL-Kernel/pull/411) covers the same +roadmap row on another platform. Its current half-split reference and uncentered +positions differ from the pinned Diffusers contract here. Combining the two +references requires semantic alignment. Existing Qwen-Image implementations in +[SGLang](https://github.com/sgl-project/sglang/tree/a1966e1b5d8b750b494d7e53cca1d2464f960038/python/sglang/multimodal_gen) +and [vLLM-Omni](https://github.com/vllm-project/vllm-omni/tree/c9c288e3f5828f2a2a9b07acb16b3b709906059a) +also motivate stride support and future QK-norm/RoPE fusion. This contribution +qualifies the unfused operator; it makes no new-algorithm or fused-model claim. + +## Validation and benchmarks + +```bash +python -m pytest tests/ops/rope/test_multi_axis_rope.py -q -rs +python tools/validation/operators/check_operator.py --op multi_axis_rope \ + --candidate cuda --device cuda --dtype bf16 --batch 2 --seq 7 --n-heads 3 --check-grad +python -m benchmarks.operators.rope.benchmark_multi_axis_rope --output rope.json +python tools/validation/operators/validate_multi_axis_rope.py --ws1 cuda --output ws1-cuda.json +python tools/validation/operators/validate_multi_axis_rope.py --ws1 triton --output ws1-triton.json +mkdocs build --strict -f mkdocs.yaml +``` + +For operator-only native compilation, without building every unrelated extension: + +```bash +python tools/validation/operators/validate_multi_axis_rope.py -q -rs +python tools/validation/operators/validate_multi_axis_rope.py --benchmark --output rope.json +``` + +To audit against the actual pinned upstream symbols without installing Diffusers, +download the source file linked above and run: + +```bash +python tools/validation/operators/check_qwenimage_rope_upstream.py \ + --source transformer_qwenimage.py --output upstream.json +``` + +The audit checks the source checksum before loading the two relevant definitions, +then compares image/text tables and forward/backward across three dtypes. + +The focused runner loads the **same production CUDA source and dispatcher +registration**, without a generated binding or extension monkeypatch. It reuses +an installed native library when available. Set a matching `CUDA_HOME`, CUDA +architecture and host compiler before invoking it. + +Ordinary CUDA builds include the source in `rl_engine._C`. An explicit operator +wheel profile builds `rl_engine._rope_cuda` with the same dispatcher registration: + +```bash +RL_KERNEL_ROPE_ONLY=1 python -m pip wheel --no-build-isolation --no-deps . -w dist +python -m pip install --force-reinstall --no-deps dist/rl_kernel-*.whl +# Change to a directory outside the source checkout before this command: +python /path/to/RL-Kernel/tools/validation/operators/check_multi_axis_rope_install.py \ + --package-root /path/to/site-packages --output installed-wheel.json +``` + +On Windows set `$env:RL_KERNEL_ROPE_ONLY='1'` in PowerShell and enter the matching +MSVC build environment. Select a CUDA toolkit matching `torch.version.cuda`. +The operator profile does not certify the complete collection of unrelated +native operators. CPU-only/ROCm builds reject this explicit CUDA wheel profile. + +`multi_axis_rope` has an `OP_SPECS` entry, standard input builder and gradient +adapter. `--ws1` runs the existing forward and gradient accuracy/invariance APIs +with C1's elementwise tolerance profile and C2's logical configuration transforms. +This gate qualifies the Qwen-Image operator and does not certify the Dense Qwen3 +chain, distributed execution, or model assembly. + +The standard CLI and WS1 forward/gradient adapters share a centered image-grid +prefix from `qwen_image_positions`. The square grid contains the full logical +sample; a nonsquare sequence length uses its row-major prefix. WS1 selects rows +by `(sample_id, token_position)` from that sample's grid, so physical chunking, +reordering and padding cannot change active-token coordinates. Padding uses +zero coordinates. These are operator fixtures, not a model's image/text assembly. + +Backward provenance is recorded inside the opaque custom-op implementation, +outside fake/AOT tracing. Its Python `execution_count` is not a CUDA Graph replay +counter: replay executes the captured GPU kernels without re-entering Python. + +For a complete native package installation on Linux, run the dedicated gate: + +```bash +export CUDA_HOME=/usr/local/cuda-12.8 +export PATH="$CUDA_HOME/bin:$PATH" +RL_KERNEL_REQUIRE_EXT=1 python -m pip install --no-build-isolation -e . +ROPE_PYTHON=python bash ci/scripts/run_qwenimage_rope_ws1.sh +``` + +The gate requires a real NVIDIA GPU, Triton and an ordinary `rl_engine._C` +containing the native RoPE registration. It rejects the operator-only profile +and fast-math mode, checks required tests actually execute, and permits only +the two-device guard to skip on a single-GPU host. Results include the loaded +binary fingerprint, ordinary operator CLI checks and both WS1 backend reports. +Set `ROPE_WS1_OUTPUT` to select the artifact directory. + +On WSL2, keep the checkout and Python environment in the Linux filesystem. +GNU linking of the package's CUDA driver calls may need +`LIBRARY_PATH=/usr/lib/wsl/lib`. Record the WSL kernel and physical GPU in the +evidence; a WSL2 profile does not qualify a bare-metal Linux server or another +GPU architecture. The `Qwen-Image RoPE` workflow runs CPU contracts and changed-file +checks for PRs targeting `test-qwenimage`; GPU qualification uses the gate above. + +Tests cover independent forward/VJP accuracy, adjacent-pair layout, shared and +batched positions, odd and rectangular grids, batch/padding/chunk invariance, +non-contiguous tensors, packed QKV views without copies, default-device contexts, +empty inputs, boundary rejection, higher derivatives, CUDA/Triton byte equality, +graph replay, the current CUDA stream, `torch.library.opcheck` and full-graph +compilation. A two-device guard test skips when only one GPU is present. + +The benchmark covers a short synthetic case and 64×64, 83×83 and 104×58 grids, +with 24 heads and both contiguous/packed-QKV layouts. It checks forward and input +gradients before timing and records CUDA Graph and synchronized eager wall time, +extra peak allocation, source hashes and actual-backend traces. Table generation +and compilation are excluded. These measurements establish no model speedup. + +To include the mature FlashAttention baselines, download `flash_attn/ops/triton/rotary.py` +and `flash_attn/layers/rotary.py` at the linked revision, then pass both paths: + +```bash +python tools/validation/operators/validate_multi_axis_rope.py --benchmark \ + --flash-source /path/to/ops/rotary.py --flash-layer-source /path/to/layers/rotary.py \ + --output rope-with-flash.json +``` + +The normalized source checksums are enforced before loading the audited symbols. +The default variant retains upstream arithmetic; `flash_no_fma` changes only the +kernel launch contraction setting and must match the FP32 gold's output bytes. diff --git a/pyproject.toml b/pyproject.toml index 207606282..f4f8b8ae9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,7 @@ where = ["."] include = ["rl_engine*"] [tool.setuptools.package-data] -rl_engine = ["**/*.json", "**/*.jsonl", "**/*.sh", "**/*.patch", "*.pyi"] +rl_engine = ["**/*.json", "**/*.jsonl", "**/*.sh", "**/*.patch", "*.pyi", "**/FLASH_ATTENTION_LICENSE"] [tool.ruff] line-length = 100 diff --git a/reports/experiments/qwenimage_multi_axis_rope_sm120.md b/reports/experiments/qwenimage_multi_axis_rope_sm120.md new file mode 100644 index 000000000..a827bb42e --- /dev/null +++ b/reports/experiments/qwenimage_multi_axis_rope_sm120.md @@ -0,0 +1,153 @@ +# Qwen-Image multi-axis RoPE: SM120 validation + +Recorded 2026-10-28 for [PR #623](https://github.com/RL-Align/RL-Kernel/pull/623). +The rotation arithmetic is based on `ac17ab87e8d93dabe93dfdc134e9853d47de7af1`. +The [operator documentation](../../docs/operators/multi-axis-rope.md) defines the +contract, upstream reuse, API and detailed reproduction commands. + +## Environments + +Both profiles use one RTX 5080 (SM120, 16 GB), Windows driver 617.42, +PyTorch 2.9.1+cu128 and CUDA Toolkit 12.8. They are separate recorded runs, +not a controlled Windows-versus-WSL performance comparison. + +| Item | WSL2 Linux | Windows | +| --- | --- | --- | +| OS | Ubuntu 24.04.5, glibc 2.39; WSL 3.0.3 on Windows 11 | Windows | +| Kernel | 6.18.54.1-microsoft-standard-WSL2 | Native Windows | +| Python | 3.12.3 | 3.13.9 | +| Compiler | CUDA 12.8.93, GCC 13.3 | CUDA 12.8, MSVC 2022 x64 | +| Triton | 3.5.1 | triton-windows 3.5.1.post24 | +| Native build | Ordinary complete `rl_engine._C`, all 10 configured translation units | Explicit operator-only `rl_engine._rope_cuda` | + +The WSL checkout, environments and caches are in the Linux filesystem. +The normal full wheel was built with `RL_KERNEL_ROPE_ONLY`, forced SM90 and +fast-math flags unset. Linking existing CUDA driver calls used +`LIBRARY_PATH=/usr/lib/wsl/lib`; the Windows driver supplies CUDA driver access. + +## Initial qualification + +| Check | WSL2 Linux | Windows | +| --- | --- | --- | +| Public operator suite | 81 passed, 1 two-GPU guard skip | 81 passed, 1 two-GPU guard skip | +| Related regressions | 195 passed, 7 Ascend skips | 195 passed, 7 Ascend skips | +| CUDA/Triton WS1 forward/gradient accuracy and invariance | Both passed with backend provenance | Both passed with backend provenance | +| Standard operator CLI forward/input gradient | Both backends passed | Both backends passed | +| Pinned actual Diffusers symbols | 30 cases passed; image/text tables exact | 30 cases passed; image/text tables exact | +| Compute Sanitizer | 34 selected cases passed, zero errors | 34 selected cases passed, zero errors | +| Fresh installed wheel outside checkout | Full wheel: both backends' eager/compiled forward and input-gradient bytes exact | Operator wheel: same checks passed | +| CPU-only PyTorch 2.9.1+cpu, Python 3.10.22 and 3.12.3 | Each: 57 passed, 52 expected GPU skips; installed reference forward/backward passed | Not recorded | +| Repository layout | 37 passed, zero skips/failures | 31 passed, 5 failures, 1 skip | +| Changed-file static checks and strict docs | Passed locally | Passed locally | + +The operator suite covers FP32/FP16/BF16, independent accuracy, byte invariance, +packed-QKV/no-copy and table strides, empty inputs, higher input VJPs, stream +selection, graph replay, opcheck and fullgraph forward/backward. The installed +wheels contain the FlashAttention BSD license; the CPU packages contain no CUDA +binary. Local CPU checks are separate from hosted CI and GPU acceptance. + +The five historical Windows layout failures also reproduce on its clean base +(27 passed, 5 failures, 1 skip): driver linking and Windows symlink privileges or +placeholders. The later Linux run preserved all 27 tracked symlinks and passed. +The earlier Windows logs remain unchanged in the archive. + +## Automated-review follow-up + +The standard CLI and WS1 adapters now share a centered image-grid prefix, +selected by full logical sample length and token identity. Eleven CPU regression +cases cover empty/square/nonsquare lengths and packed/permuted/chunked/padded +layouts. Before the change, ten of those cases failed; all eleven now pass. +Additional tests pin non-CUDA reference routing, library-load error propagation, +shared-table broadcasting, and head-dimension stride 2 or outermost storage. +Native shape checks and indexing share named dimension constants; the rotation +arithmetic is unchanged. Backward telemetry documents its graph-replay limits. + +The following are new local WSL2 runs of the updated implementation and fixtures: + +| Check | Result | +| --- | --- | +| Rebuilt ordinary complete CUDA wheel; fresh installed-package eager/fullgraph forward/input gradient | Passed for both GPU backends | +| Public operator suite | 105 passed, 1 two-GPU guard skip | +| CUDA/Triton WS1 accuracy/invariance and ordinary CLI checks | Both passed with explicit backend provenance | +| Related regressions, including the eleven new input cases | 206 passed, 7 Ascend skips | +| Repository layout with native extension enforcement | 37 passed | +| CPU-only PyTorch 2.9.1+cpu, Python 3.10.22 and 3.12.3 | Each: 80 passed, 64 expected GPU skips; installed reference forward/backward passed | +| Expanded Compute Sanitizer subset, including the new stride cases | 52 passed, zero errors | +| Changed-file checks and strict docs | Passed locally | + +The full wheel contains the shared validation helper and the FlashAttention BSD +license. Wheel SHA256: +`9feb8200b948021c268c933f93efcd424b6b68d6c55bc24c3004d6336deb3a13`. +Non-CUDA routing tests exercise candidate selection with CPU tensors; they do +not certify NPU/ROCm/MUSA hardware. The benchmark below retains its initial-run +measurements; no new timing comparison was performed for this follow-up. + +## Standalone benchmark + +WSL2 BF16, batch 1, 24 heads, D=128. CUDA Graph medians in ms/call; +each cell is **forward / forward+backward**. + +| Grid / layout | Native CUDA | Adapted Triton | FlashAttention, no FMA | Diffusers complex | +| --- | --- | --- | --- | --- | +| 64x64 contiguous | 0.05780 / 0.12967 | 0.04929 / 0.12414 | 0.04898 / 0.12407 | 0.24446 / 0.51281 | +| 64x64 packed QKV | 0.05868 / 0.13079 | 0.04731 / 0.11841 | 0.04722 / 0.11847 | 0.24754 / 0.51583 | +| 83x83 contiguous | 0.11198 / 0.22392 | 0.10720 / 0.22538 | 0.10739 / 0.21420 | 0.52579 / 1.06721 | +| 83x83 packed QKV | 0.11233 / 0.22486 | 0.10745 / 0.21478 | 0.10744 / 0.21469 | 0.52820 / 1.06947 | +| 104x58 contiguous | 0.09797 / 0.19708 | 0.09366 / 0.18497 | 0.09362 / 0.18458 | 0.44992 / 0.90110 | +| 104x58 packed QKV | 0.09610 / 0.19612 | 0.09310 / 0.18828 | 0.09323 / 0.18817 | 0.44979 / 0.90276 | + +All 40 pre-timing comparisons passed; CUDA, adapted Triton and the no-FMA +baseline had zero differing output/input-gradient bytes against the FP32 gold. +The archive includes 80 measurements, the short case, default FlashAttention, +eager wall timings and peak allocation. At 4096 tokens, either GPU path's forward +extra peak is 24 MiB, the returned output size, for both layouts. + +Each graph contains 20 calls, with five warm calls/replays and 30 measured +replays. Eager wall timings use three synchronized trials of 30 calls. +Tables, compilation and input/backing allocation are excluded; extra peak +includes the returned output or gradient. The upstream source/autograd wrapper +is pinned to `98c1cee0ed59307833841cdbf567fe336ceec13a`; the no-FMA variant changes +only contraction. These single-process runs use a shared GPU without clock +pinning and establish no advantage over FlashAttention or model speedup. + +## Reproduction + +Use a matching CUDA 12.8/PyTorch environment in a Linux checkout: + +```bash +export CUDA_HOME=/usr/local/cuda-12.8 +export PATH="$CUDA_HOME/bin:$PATH" +export LIBRARY_PATH="/usr/lib/wsl/lib${LIBRARY_PATH:+:$LIBRARY_PATH}" +export MAX_JOBS=2 TORCH_CUDA_ARCH_LIST=12.0 RL_KERNEL_REQUIRE_EXT=1 +unset RL_KERNEL_ROPE_ONLY KERNEL_ALIGN_FORCE_SM90 KERNEL_ALIGN_USE_FAST_MATH +python -m pip wheel --no-build-isolation --no-deps . -w /tmp/rope-full-wheel +python -m pip install --force-reinstall --no-deps /tmp/rope-full-wheel/rl_kernel-*.whl +python setup.py build_ext --inplace +ROPE_PYTHON=python ROPE_WS1_OUTPUT=/tmp/rope-ws1 bash ci/scripts/run_qwenimage_rope_ws1.sh +python ci/run.py layout -q -rs +compute-sanitizer --tool memcheck --error-exitcode 1 python -m pytest \ + tests/ops/rope/test_multi_axis_rope.py -q \ + -k "forward_backward_accuracy or noncontiguous_input_tables_and_gradient or empty_forward_backward or packed_qkv_and_independent_table_strides or native_entrypoint_checks" +``` + +The [GPU gate](../../ci/scripts/run_qwenimage_rope_ws1.sh) requires the full +native extension and Triton, records loaded binary/backend identity, and permits +only the two-GPU guard to skip. The [CPU workflow](../../.github/workflows/qwenimage-rope.yml) +contains the Python 3.10/3.12 contract commands. The operator documentation covers +installed-wheel probes, pinned upstream audits and benchmark reproduction. +Hosted CI status is tracked on the PR. + +## Archived evidence and qualification limits + +Detailed per-case JSON, JUnit, build/test logs, dependency versions, and +source/binary SHA256 manifests remain available at immutable commit +`84751473006401c19c1c8c17844e8ef8b1c44375`: + +- [WSL2 complete-extension evidence](https://github.com/dandan20031031/RL-Kernel/tree/84751473006401c19c1c8c17844e8ef8b1c44375/reports/experiments/qwenimage_multi_axis_rope_wsl2_sm120) +- [Historical Windows operator-profile evidence](https://github.com/dandan20031031/RL-Kernel/tree/84751473006401c19c1c8c17844e8ef8b1c44375/reports/experiments/qwenimage_multi_axis_rope_v2_sm120) + +Those archives identify the initial runs; the follow-up table records subsequent +reruns. Bare-metal Linux, A100/H100, two-GPU guards, +QK-norm fusion, model/block assembly and trajectories remain unqualified. +Building the full library does not certify every native operator or Dense Qwen3. +No distributed or cross-version/device byte-equality claim is made. diff --git a/rl_engine/backends/cuda/rope/multi_axis_rope.py b/rl_engine/backends/cuda/rope/multi_axis_rope.py new file mode 100644 index 000000000..83ef741ca --- /dev/null +++ b/rl_engine/backends/cuda/rope/multi_axis_rope.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Architecture-independent CUDA implementation; no Hopper-only instructions.""" + +import importlib.util + +import torch + +from rl_engine.backends import extension # noqa: F401 - loads the normal native library +from rl_engine.ops.rope._registration import register_rotation +from rl_engine.ops.rope.multi_axis_rope import MultiAxisRoPEMetadata, validate_inputs + + +def _apply(x, cos, sin, transpose): + return torch.ops.rl_kernel_native.multi_axis_rope(x, cos, sin, transpose) + + +def native_available(): + if torch.version.hip is not None: + return False + if not hasattr(torch.ops.rl_kernel_native, "multi_axis_rope"): + library = importlib.util.find_spec("rl_engine._rope_cuda") + if library is not None: + torch.ops.load_library(library.origin) + return hasattr(torch.ops.rl_kernel_native, "multi_axis_rope") + + +_rotation = register_rotation( + "qwen_image_rope_cuda", + _apply, + kernel_id="torch.ops.rl_kernel_native.multi_axis_rope", + family="cuda", +) + + +class MultiAxisRoPECudaOp(MultiAxisRoPEMetadata): + backend_id = "rlkernel.cuda.qwenimage_multi_axis_rope.v1" + kernel_id = "torch.ops.rl_kernel_native.multi_axis_rope" + + def __init__(self): + if not native_available(): + raise RuntimeError( + "CUDA multi_axis_rope extension is unavailable; rebuild rl_engine._C" + ) + + def __call__(self, x, cos, sin): + validate_inputs(x, cos, sin) + if x.device.type != "cuda": + raise ValueError("CUDA multi_axis_rope requires CUDA tensors") + return _rotation(x, cos, sin, False) diff --git a/rl_engine/backends/shared/triton/rope/FLASH_ATTENTION_LICENSE b/rl_engine/backends/shared/triton/rope/FLASH_ATTENTION_LICENSE new file mode 100644 index 000000000..5860e4b33 --- /dev/null +++ b/rl_engine/backends/shared/triton/rope/FLASH_ATTENTION_LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/rl_engine/backends/shared/triton/rope/_flash_attn_rotary.py b/rl_engine/backends/shared/triton/rope/_flash_attn_rotary.py new file mode 100644 index 000000000..bd617c8c1 --- /dev/null +++ b/rl_engine/backends/shared/triton/rope/_flash_attn_rotary.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: BSD-3-Clause +# Copyright (c) 2025, Tri Dao. +# Adapted from Dao-AILab/flash-attention, revision +# 98c1cee0ed59307833841cdbf567fe336ceec13a, flash_attn/ops/triton/rotary.py. +# Changes: explicit shared/per-sample table strides. Launch policy is in the +# sibling wrapper. See FLASH_ATTENTION_LICENSE for the retained BSD-3 terms. + +import triton +import triton.language as tl + + +@triton.jit +def rotary_kernel( + OUT, # Pointers to matrices + X, + COS, + SIN, + CU_SEQLENS, + SEQLEN_OFFSETS, # this could be int or a pointer + # Matrix dimensions + seqlen, + nheads, + seqlen_ro, + # strides + stride_out_batch, + stride_out_seqlen, + stride_out_nheads, + stride_out_headdim, + stride_x_batch, + stride_x_seqlen, + stride_x_nheads, + stride_x_headdim, + stride_cos_batch, + stride_cos_seqlen, + stride_cos_dim, + stride_sin_batch, + stride_sin_seqlen, + stride_sin_dim, + # Meta-parameters + # We want ROTARY_DIM to be constexpr, otherwise the triton compiler doesn't know that + # the mask is constant every 8 elements, and it will generate LDG.16 instead of LDG.128 + ROTARY_DIM: tl.constexpr, + IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr, + IS_VARLEN: tl.constexpr, + INTERLEAVED: tl.constexpr, + CONJUGATE: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_M: tl.constexpr, +): + BLOCK_K: tl.constexpr = triton.next_power_of_2(ROTARY_DIM) + ROTARY_DIM_HALF = ROTARY_DIM // 2 + pid_head = tl.program_id(axis=0) + pid_m = tl.program_id(axis=1) + pid_batch = tl.program_id(axis=2) + + if not IS_VARLEN: + X = X + pid_batch * stride_x_batch + OUT = OUT + pid_batch * stride_out_batch + else: + start_idx = tl.load(CU_SEQLENS + pid_batch) + seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx + X = X + start_idx * stride_x_seqlen + OUT = OUT + start_idx * stride_out_seqlen + + if pid_m * BLOCK_M >= seqlen: + return + + rh = pid_head * BLOCK_H + tl.arange(0, BLOCK_H) + rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + if not IS_SEQLEN_OFFSETS_TENSOR: + rm_cs = rm + SEQLEN_OFFSETS + else: + rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch) + + rk_half = tl.arange(0, BLOCK_K // 2) + COS = ( + COS + + pid_batch * stride_cos_batch + + (rm_cs[:, None] * stride_cos_seqlen + rk_half[None, :] * stride_cos_dim) + ) + SIN = ( + SIN + + pid_batch * stride_sin_batch + + (rm_cs[:, None] * stride_sin_seqlen + rk_half[None, :] * stride_sin_dim) + ) + mask_cs = (rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < ROTARY_DIM_HALF) + cos = tl.load(COS, mask=mask_cs, other=1.0).to(tl.float32) + sin = tl.load(SIN, mask=mask_cs, other=0.0).to(tl.float32) + if CONJUGATE: + sin = -sin + + if not INTERLEAVED: + # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT + X = X + ( + rh[:, None, None] * stride_x_nheads + + rm[None, :, None] * stride_x_seqlen + + rk_half[None, None, :] * stride_x_headdim + ) + OUT = OUT + ( + rh[:, None, None] * stride_out_nheads + + rm[None, :, None] * stride_out_seqlen + + rk_half[None, None, :] * stride_out_headdim + ) + mask = ( + (rh[:, None, None] < nheads) + & (rm[None, :, None] < seqlen) + & (rk_half[None, None, :] < ROTARY_DIM_HALF) + ) + x0 = tl.load(X, mask=mask, other=0.0).to(tl.float32) + x1 = tl.load( + X + ROTARY_DIM_HALF * stride_x_headdim, + mask=mask, + other=0.0, + ).to(tl.float32) + o0 = x0 * cos - x1 * sin + o1 = x0 * sin + x1 * cos + tl.store(OUT, o0, mask=mask) + tl.store(OUT + ROTARY_DIM_HALF * stride_out_headdim, o1, mask=mask) + else: + rk = tl.arange(0, BLOCK_K) + X = X + ( + rh[:, None, None] * stride_x_nheads + + rm[None, :, None] * stride_x_seqlen + + rk[None, None, :] * stride_x_headdim + ) + OUT = OUT + ( + rh[:, None, None] * stride_out_nheads + + rm[None, :, None] * stride_out_seqlen + + rk[None, None, :] * stride_out_headdim + ) + mask = ( + (rh[:, None, None] < nheads) + & (rm[None, :, None] < seqlen) + & (rk[None, None, :] < ROTARY_DIM) + ) + x = tl.load(X, mask=mask, other=0.0).to(tl.float32) + x0, x1 = tl.split(tl.reshape(x, [BLOCK_H, BLOCK_M, BLOCK_K // 2, 2])) + o0 = x0 * cos - x1 * sin + o1 = x0 * sin + x1 * cos + o = tl.reshape(tl.join(o0, o1), [BLOCK_H, BLOCK_M, BLOCK_K]) + tl.store(OUT, o, mask=mask) diff --git a/rl_engine/backends/shared/triton/rope/multi_axis_rope.py b/rl_engine/backends/shared/triton/rope/multi_axis_rope.py new file mode 100644 index 000000000..1a37836fe --- /dev/null +++ b/rl_engine/backends/shared/triton/rope/multi_axis_rope.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FlashAttention-derived, stride-aware rotation under a frozen FP32 profile.""" + +import torch +import triton + +from rl_engine.ops.rope._registration import register_rotation +from rl_engine.ops.rope.multi_axis_rope import MultiAxisRoPEMetadata, validate_inputs + +from ._flash_attn_rotary import rotary_kernel + + +def _apply(x, cos, sin, transpose): + output = torch.empty(x.shape, device=x.device, dtype=x.dtype) + batch, seq, heads, _ = x.shape + if batch and seq: + with torch.cuda.device(x.device): + rotary_kernel[(triton.cdiv(heads, 2), triton.cdiv(seq, 8), batch)]( + output, + x, + cos, + sin, + None, + 0, + seq, + heads, + seq, + *output.stride(), + *x.stride(), + cos.stride(0) if cos.ndim == 3 else 0, + cos.stride(-2), + cos.stride(-1), + sin.stride(0) if sin.ndim == 3 else 0, + sin.stride(-2), + sin.stride(-1), + 128, + False, + False, + True, + transpose, + BLOCK_H=2, + BLOCK_M=8, + enable_fp_fusion=False, + ) + return output + + +_rotation = register_rotation( + "qwen_image_rope_triton", + _apply, + kernel_id="rl_engine.backends.shared.triton.rope._flash_attn_rotary.rotary_kernel", + family="triton", +) + + +class MultiAxisRoPETritonOp(MultiAxisRoPEMetadata): + backend_id = "rlkernel.triton.qwenimage_multi_axis_rope.v1" + kernel_id = "rl_engine.backends.shared.triton.rope._flash_attn_rotary.rotary_kernel" + + def __init__(self): + if torch.version.hip is not None or not torch.cuda.is_available(): + raise RuntimeError("this multi_axis_rope Triton profile requires NVIDIA CUDA") + + def __call__(self, x, cos, sin): + validate_inputs(x, cos, sin) + if x.device.type != "cuda": + raise ValueError("Triton multi_axis_rope requires CUDA tensors") + return _rotation(x, cos, sin, False) diff --git a/rl_engine/ops/rope/_registration.py b/rl_engine/ops/rope/_registration.py new file mode 100644 index 000000000..33798a03e --- /dev/null +++ b/rl_engine/ops/rope/_registration.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Dispatcher, fake and recursive input-VJP registration for the RoPE profile.""" + +import torch + +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.ops.rope.multi_axis_rope import validate_inputs + + +def register_rotation(name, implementation, *, kernel_id, family): + @torch.library.custom_op(f"rl_kernel::{name}", mutates_args=(), device_types="cuda") + def rotation( + x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, transpose: bool + ) -> torch.Tensor: + validate_inputs(x, cos, sin) + result = implementation(x, cos, sin, transpose) + if transpose: + # Keep telemetry inside the opaque custom op; fake/AOT tracing + # must not execute this Python lock. CUDA Graph replay runs only + # captured kernels, so execution_count is not a replay counter. + record_backward( + "multi_axis_rope", + kernel_id=kernel_id, + impl="interleaved_fp32_transpose", + family=family, + ) + return result + + @rotation.register_fake + def fake(x, cos, sin, transpose): + validate_inputs(x, cos, sin) + return torch.empty(x.shape, dtype=x.dtype, device=x.device) + + def setup_context(ctx, inputs, output): + _, cos, sin, transpose = inputs + ctx.save_for_backward(cos, sin) + ctx.transpose = transpose + + def backward(ctx, grad): + cos, sin = ctx.saved_tensors + return rotation(grad, cos, sin, not ctx.transpose), None, None, None + + rotation.register_autograd(backward, setup_context=setup_context) + return rotation diff --git a/rl_engine/ops/rope/multi_axis_rope.py b/rl_engine/ops/rope/multi_axis_rope.py new file mode 100644 index 000000000..ff95adff2 --- /dev/null +++ b/rl_engine/ops/rope/multi_axis_rope.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen-Image interleaved three-axis RoPE (RFC #386). + +Tables are immutable inputs, built before capture/replay. The rotation accepts +BSHD tensors, unlike the legacy BHSD, rotate-half RoPE operator. +""" + +from __future__ import annotations + +import torch + +CONTRACT = "qwenimage-interleaved-rope-fp32-v1" +AXES_DIMS = (16, 56, 56) +UPSTREAM_REVISION = "40f63690da440e8c382ab53e67508835deb1ed2f" + + +def qwen_image_positions( + height: int, width: int, text_length: int +) -> tuple[torch.Tensor, torch.Tensor]: + """CPU int64 positions for one image's *packed latent* H/W, then text. + + Matches QwenEmbedRope(scale_rope=True), single-frame text-to-image. Odd + grids have the extra negative coordinate: H=3 gives [-2, -1, 0]. + Return separate tables so callers explicitly choose their stream order. + """ + for name, value in ( + ("height", height), + ("width", width), + ("text_length", text_length), + ): + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + if not 1 <= height <= 8192 or not 1 <= width <= 8192 or text_length < 0: + raise ValueError("height/width must be in [1, 8192], text_length must be nonnegative") + start = max(height // 2, width // 2) + if start + text_length > 4096: + raise ValueError("text positions exceed the pinned upstream 4096-entry table") + h = torch.arange(-(height - height // 2), height // 2, device="cpu") + w = torch.arange(-(width - width // 2), width // 2, device="cpu") + hh, ww = torch.meshgrid(h, w, indexing="ij") + image = torch.stack((torch.zeros_like(hh), hh, ww), dim=-1).reshape(-1, 3) + text = ( + torch.arange(start, start + text_length, device="cpu")[:, None].expand(-1, 3).contiguous() + ) + return image, text + + +def build_multi_axis_cos_sin( + positions: torch.Tensor, *, device=None +) -> tuple[torch.Tensor, torch.Tensor]: + """Build canonical FP32 tables on CPU, then copy to the requested device. + + positions: int32/int64 [S, 3] or [B, S, 3], range [-4096, 4095]. + The CPU torch.polar construction follows the pinned Diffusers reference; + no GPU transcendental or per-batch table recomputation enters the hot path. + """ + if positions.dtype not in (torch.int32, torch.int64): + raise TypeError("positions must be int32 or int64") + if positions.ndim not in (2, 3) or positions.shape[-1] != 3: + raise ValueError("positions must be [S, 3] or [B, S, 3]") + cpu = positions.detach().to(device="cpu", dtype=torch.float32) + if torch.any((cpu < -4096) | (cpu >= 4096)): + raise ValueError("positions exceed the pinned upstream [-4096, 4095] range") + phases = [] + for axis, dim in enumerate(AXES_DIMS): + inv_freq = 1.0 / torch.pow(10000, torch.arange(0, dim, 2, device="cpu").float().div(dim)) + phases.append(cpu[..., axis, None] * inv_freq) + phase = torch.cat(phases, dim=-1) + cis = torch.polar(torch.ones_like(phase), phase) + target = positions.device if device is None else device + return cis.real.contiguous().to(target), cis.imag.contiguous().to(target) + + +def validate_inputs(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> None: + if x.ndim != 4 or x.shape[-1] != 128 or x.shape[2] <= 0: + raise ValueError("x must be [B, S, H, 128] with H > 0") + if x.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError("x must be float16, bfloat16 or float32") + shared = (x.shape[1], 64) + batched = (x.shape[0], x.shape[1], 64) + if cos.shape not in (shared, batched) or sin.shape != cos.shape: + raise ValueError("cos/sin must have matching [S, 64] or [B, S, 64] shapes") + if cos.dtype != torch.float32 or sin.dtype != torch.float32: + raise TypeError("cos/sin must be float32") + if x.device != cos.device or x.device != sin.device: + raise ValueError("x, cos and sin must be on the same device") + if cos.requires_grad or sin.requires_grad: + raise ValueError("cos/sin are fixed tables; table gradients are unsupported") + + +def get_multi_axis_rope_op(backend: str = "auto", *, device="cpu"): + """Resolve an op exposing trace(); explicit backends never silently fall back. + + Strict callers must request cuda or triton and retain the returned trace. + auto is a diagnostic/convenience mode with a PyTorch reference fallback. + """ + if backend == "auto": + from rl_engine.runtime.registry import kernel_registry + + return kernel_registry.get_op("multi_axis_rope", device=device) + if backend == "cuda": + from rl_engine.backends.cuda.rope.multi_axis_rope import MultiAxisRoPECudaOp + + return MultiAxisRoPECudaOp() + if backend == "triton": + from rl_engine.backends.shared.triton.rope.multi_axis_rope import MultiAxisRoPETritonOp + + return MultiAxisRoPETritonOp() + if backend == "reference": + from rl_engine.reference.rope.multi_axis_rope import NativeMultiAxisRoPEOp + + return NativeMultiAxisRoPEOp() + raise ValueError(f"unknown multi_axis_rope backend: {backend!r}") + + +def multi_axis_rope(x, cos, sin, *, backend="auto"): + return get_multi_axis_rope_op(backend, device=x.device)(x, cos, sin) + + +class MultiAxisRoPEMetadata: + backend_id: str + kernel_id: str + + def trace(self) -> dict: + return { + "contract": CONTRACT, + "actual_backend": self.backend_id, + "kernel_id": self.kernel_id, + "layout": "BSHD, adjacent channel pairs", + "arithmetic": "fp32 separate multiply/add, one output cast, no reduction", + "table_revision": UPSTREAM_REVISION, + } diff --git a/rl_engine/reference/rope/multi_axis_rope.py b/rl_engine/reference/rope/multi_axis_rope.py new file mode 100644 index 000000000..353eb6a00 --- /dev/null +++ b/rl_engine/reference/rope/multi_axis_rope.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Explicit FP32 reference for the Qwen-Image interleaved rotation.""" + +import torch + +from rl_engine.ops.rope.multi_axis_rope import MultiAxisRoPEMetadata, validate_inputs + + +class NativeMultiAxisRoPEOp(MultiAxisRoPEMetadata): + op_class = "elementwise" + backend_id = "rlkernel.reference.qwenimage_multi_axis_rope.v1" + kernel_id = "rl_engine.reference.rope.multi_axis_rope.NativeMultiAxisRoPEOp" + + def __call__(self, x, cos, sin): + return self.forward(x, cos, sin) + + def forward(self, x, cos, sin): + return self.forward_fp32(x, cos, sin).to(x.dtype) + + def forward_fp32(self, x, cos, sin): + validate_inputs(x, cos, sin) + c, s = cos.unsqueeze(-2), sin.unsqueeze(-2) + even, odd = x[..., 0::2].float(), x[..., 1::2].float() + out = torch.stack((even * c - odd * s, odd * c + even * s), dim=-1) + return out.flatten(-2) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 7751d5901..326081f44 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -167,6 +167,11 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): PYTORCH_NATIVE_ROPE = "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp" TRITON_ROPE = "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp" CUDA_ROPE_SM90 = "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op" + CUDA_MULTI_AXIS_ROPE = "rl_engine.backends.cuda.rope.multi_axis_rope.MultiAxisRoPECudaOp" + TRITON_MULTI_AXIS_ROPE = ( + "rl_engine.backends.shared.triton.rope.multi_axis_rope.MultiAxisRoPETritonOp" + ) + PYTORCH_MULTI_AXIS_ROPE = "rl_engine.reference.rope.multi_axis_rope.NativeMultiAxisRoPEOp" ASCEND_ROPE = "rl_engine.kernels.ops.ascend.rotary_embedding.rope.RoPEAscendOp" PYTORCH_NATIVE_SILU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp" PYTORCH_NATIVE_SWIGLU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp" @@ -737,6 +742,15 @@ def __init__(self): ], }, } + # CUDA qualification for the new interleaved Qwen-Image contract. + # Other platforms retain an explicit diagnostic reference, not logp. + for candidates in self._priority_map.values(): + candidates["multi_axis_rope"] = [OpBackend.PYTORCH_MULTI_AXIS_ROPE] + self._priority_map["cuda"]["multi_axis_rope"] = [ + OpBackend.CUDA_MULTI_AXIS_ROPE, + OpBackend.TRITON_MULTI_AXIS_ROPE, + OpBackend.PYTORCH_MULTI_AXIS_ROPE, + ] # Preserve the former CPU fallback behavior for every operator on NPU, # then override only the operators with an Ascend-specific backend. self._priority_map["npu"] = { diff --git a/rl_engine/validation/operators/gradient_adapters.py b/rl_engine/validation/operators/gradient_adapters.py index f22d80285..48da77045 100644 --- a/rl_engine/validation/operators/gradient_adapters.py +++ b/rl_engine/validation/operators/gradient_adapters.py @@ -32,6 +32,7 @@ MissingBackwardError, ) from rl_engine.validation.operators.operator_specs import OP_SPECS, _load_object +from rl_engine.validation.operators.rope_inputs import qwen_image_grid_prefix REPO_ROOT = Path(__file__).resolve().parents[3] @@ -220,6 +221,21 @@ def to_dict(self) -> dict[str, Any]: "csrc/cuda/logprob/fused_linear_logp_sm90.cu", ), ), + # Optional Qwen-Image extension: do not certify it as a Dense Qwen3 chain node. + "multi_axis_rope": GradientAdapterSpec( + op_name="multi_axis_rope", + chain_node="multi_axis_rope", + op_class="elementwise", + spec_name="multi_axis_rope", + tensors=(_DX,), + requirement="optional_fused", + source_files=( + "rl_engine/backends/cuda/rope/multi_axis_rope.py", + "rl_engine/backends/shared/triton/rope/multi_axis_rope.py", + "rl_engine/backends/shared/triton/rope/_flash_attn_rotary.py", + "csrc/cuda/rope/multi_axis_rope.cu", + ), + ), "rope": GradientAdapterSpec( op_name="rope", chain_node="rope", @@ -712,6 +728,27 @@ def _row_inputs( "weight": params["weight"], "bias": None, } + if op_name == "multi_axis_rope": + from rl_engine.ops.rope.multi_axis_rope import build_multi_axis_cos_sin + + if head_dim != 128: + raise ValueError("multi_axis_rope gradient adapter requires head_dim=128") + rows = _stack_rows(keys, leading, (n_heads, head_dim), device=device, dtype=dtype) + # Use the complete logical sample, not this physical call's rows, to + # keep grid centers stable under chunks, permutations and padding. + lengths: dict[str, int] = {} + for sample_id, position in tokens: + lengths[sample_id] = max(lengths.get(sample_id, 0), position + 1) + grids: dict[str, torch.Tensor] = {} + coords = torch.zeros((len(keys), 3), device="cpu", dtype=torch.int64) + for row, key in enumerate(keys): + if key is not None: + sample_id, position = key + if sample_id not in grids: + grids[sample_id] = qwen_image_grid_prefix(lengths[sample_id]) + coords[row] = grids[sample_id][position] + cos, sin = build_multi_axis_cos_sin(coords, device=device) + return {"x": rows.unsqueeze(0), "cos": cos, "sin": sin} if op_name == "rope": rows = _stack_rows(keys, leading, (n_heads, head_dim), device=device, dtype=dtype) return { @@ -724,6 +761,10 @@ def _row_inputs( def _to_rows(op_name: str, value: torch.Tensor, n_rows: int) -> torch.Tensor: """Normalize an operator output / input-grad back to (n_rows, *tail).""" + if op_name == "multi_axis_rope": + if value.shape[0] != 1 or value.shape[1] != n_rows: + raise ValueError("multi_axis_rope row adapter expects [1, tokens, heads, 128]") + return value.squeeze(0) if op_name == "rope": # RoPE runs as (1, heads, tokens, head_dim); tokens is the row axis. permuted = value.permute(0, 2, 1, 3) diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index ca3b7c120..0306a683d 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -8,6 +8,8 @@ import torch +from rl_engine.validation.operators.rope_inputs import qwen_image_grid_prefix + DEFAULT_HIDDEN = 4096 DEFAULT_N_HEADS = 32 DEFAULT_N_KV_HEADS = 8 @@ -37,6 +39,7 @@ def make_operator_inputs( "linear_logp": _make_linear_logp_inputs, "batch_invariant_logp": _make_batch_invariant_logp_inputs, "rope": _make_rope_inputs, + "multi_axis_rope": _make_multi_axis_rope_inputs, "silu": _make_silu_inputs, "swiglu": _make_swiglu_inputs, "embedding": _make_embedding_inputs, @@ -67,6 +70,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "batch_invariant_logp": f"{batch}x{seq}x{vocab}", "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", + "multi_axis_rope": f"{batch}x{seq}x{_arg_int(args, 'n_heads', 24)}x128", "silu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", "swiglu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", @@ -250,6 +254,22 @@ def _make_batch_invariant_logp_inputs( } +def _make_multi_axis_rope_inputs(args, dtype, device): + from rl_engine.ops.rope.multi_axis_rope import build_multi_axis_cos_sin + + batch, seq = _batch_seq(args) + if _arg_int(args, "head_dim", 128) != 128: + raise ValueError("multi_axis_rope requires head_dim=128") + heads = _arg_int(args, "n_heads", 24) + positions = qwen_image_grid_prefix(seq) + cos, sin = build_multi_axis_cos_sin(positions, device=device) + return { + "x": _floating_tensor((batch, seq, heads, 128), args, dtype, device, offset=0), + "cos": cos, + "sin": sin, + } + + def _make_rope_inputs( args: argparse.Namespace, dtype: torch.dtype, device: torch.device ) -> dict[str, Any]: diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 2323f4787..bc23ab28a 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -197,6 +197,20 @@ def _load_object(path: str) -> Any: }, grad_input_names=("a", "b"), ), + "multi_axis_rope": OperatorSpec( + name="multi_axis_rope", + op_class="elementwise", + gold_path="rl_engine.reference.rope.multi_axis_rope.NativeMultiAxisRoPEOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.reference.rope.multi_axis_rope.NativeMultiAxisRoPEOp", + "cuda": "rl_engine.backends.cuda.rope.multi_axis_rope.MultiAxisRoPECudaOp", + "triton": ( + "rl_engine.backends.shared.triton.rope.multi_axis_rope.MultiAxisRoPETritonOp" + ), + }, + grad_input_names=("x",), + ), "rope": OperatorSpec( name="rope", op_class="elementwise", diff --git a/rl_engine/validation/operators/qwen_image_rope.py b/rl_engine/validation/operators/qwen_image_rope.py new file mode 100644 index 000000000..bc605bd32 --- /dev/null +++ b/rl_engine/validation/operators/qwen_image_rope.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen-Image operator-only four-judgment gate using the standard C3/C4 APIs. + +Reuses C2's logical sample/config transforms and C1's elementwise tolerances. +It does not certify the Dense Qwen3 chain, distributed execution or Qwen-Image +model assembly. The operator has its own OP_SPECS entry and gradient adapter. +""" + +import torch + +from rl_engine.contracts.numerical import BackendProvenance +from rl_engine.ops.rope.multi_axis_rope import get_multi_axis_rope_op +from rl_engine.reference.rope.multi_axis_rope import NativeMultiAxisRoPEOp +from rl_engine.validation.operators.forward_invariance import assert_forward_batch_invariant +from rl_engine.validation.operators.gradient_adapters import ( + get_adapter, + make_forward_runner, + make_gradient_runner, +) +from rl_engine.validation.operators.gradient_invariance import assert_gradient_batch_invariant + + +def run_multi_axis_rope_ws1(backend): + if backend not in ("cuda", "triton"): + raise ValueError("WS1 GPU evidence requires an explicit cuda or triton backend") + device = torch.device("cuda") + profile = "cuda_bf16" if backend == "cuda" else "triton_cuda_bf16" + op = get_multi_axis_rope_op(backend, device=device) + gold = NativeMultiAxisRoPEOp().forward_fp32 + trace = op.trace() + dtype = torch.bfloat16 + provenance = BackendProvenance( + backend_profile=profile, + requested_backend=backend, + actual_backend=backend, + execution_dtype="bfloat16", + accumulation_dtype="float32", + output_dtype="bfloat16", + reference_dtype="float32", + candidate_tf32_enabled=False, + reference_tf32_enabled=False, + ) + metadata = { + "backend_profile": profile, + "provenance": provenance, + "op_class": "elementwise", + "dtype": dtype, + "op_name": "multi_axis_rope", + "candidate_id": trace["actual_backend"], + "device": str(device), + "compute_capability": "sm%d%d" % torch.cuda.get_device_capability(), + "observed_actual_backend": backend, + "observed_kernel_id": trace["kernel_id"], + "observed_output_dtype": "bfloat16", + } + runner_options = {"device": device, "dtype": dtype, "n_heads": 3, "head_dim": 128} + candidate_options = {"backend_family": backend, "kernel_id": trace["kernel_id"]} + forward = assert_forward_batch_invariant( + make_forward_runner( + "multi_axis_rope", op, reference=False, **runner_options, **candidate_options + ), + gold_fn=make_forward_runner("multi_axis_rope", gold, reference=True, **runner_options), + include_logprob_smoke=False, + **metadata, + ) + gradient = assert_gradient_batch_invariant( + make_gradient_runner( + "multi_axis_rope", op, reference=False, **runner_options, **candidate_options + ), + gold_fn=make_gradient_runner("multi_axis_rope", gold, reference=True, **runner_options), + grad_tensors=get_adapter("multi_axis_rope").tensors, + **metadata, + ) + return { + "scope": "Qwen-Image multi_axis_rope operator; C1 elementwise and C2 config transforms", + "trace": trace, + "passed": forward.passed and gradient.passed, + "forward": forward.to_dict(), + "gradient": gradient.to_dict(), + } diff --git a/rl_engine/validation/operators/rope_inputs.py b/rl_engine/validation/operators/rope_inputs.py new file mode 100644 index 000000000..5e22f977e --- /dev/null +++ b/rl_engine/validation/operators/rope_inputs.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared Qwen-Image coordinates for operator and logical-layout validation.""" + +import math + +import torch + +from rl_engine.ops.rope.multi_axis_rope import qwen_image_positions + + +def qwen_image_grid_prefix(sequence_length: int) -> torch.Tensor: + """Use a centered square-grid prefix for a complete logical sequence. + + CLI sequence lengths need not be square. Grid size depends on the full + sample length, never the physical batch, padding or current chunk length. + """ + if sequence_length < 0: + raise ValueError("sequence_length must be nonnegative") + side = math.isqrt(max(0, sequence_length - 1)) + 1 + image, _ = qwen_image_positions(side, side, 0) + return image[:sequence_length] diff --git a/tests/build/test_multi_axis_rope_build.py b/tests/build/test_multi_axis_rope_build.py new file mode 100644 index 000000000..c9d06791e --- /dev/null +++ b/tests/build/test_multi_axis_rope_build.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The explicit operator wheel profile must build native code or fail closed.""" + +from types import SimpleNamespace + +import pytest +import torch + +from build_tools import extensions + + +def test_rope_only_wheel_selects_production_sources(monkeypatch): + monkeypatch.setenv("RL_KERNEL_ROPE_ONLY", "1") + fake_torch = SimpleNamespace( + __file__=torch.__file__, version=SimpleNamespace(cuda="12.8", hip=None) + ) + monkeypatch.setattr( + extensions, + "_load_torch_extension_tools", + lambda: (fake_torch, None, lambda **kwargs: kwargs), + ) + (extension,) = extensions.get_extensions() + assert extension["name"] == "rl_engine._rope_cuda" + assert set(extension["sources"]) == { + "csrc/bindings/rope_module.cpp", + "csrc/cuda/rope/multi_axis_rope.cu", + } + assert "--fmad=false" in extension["extra_compile_args"]["nvcc"] + + +@pytest.mark.parametrize("hip", [None, "6.2"]) +def test_rope_only_wheel_rejects_non_cuda_builds(monkeypatch, hip): + monkeypatch.setenv("RL_KERNEL_ROPE_ONLY", "1") + fake_torch = SimpleNamespace( + __file__=torch.__file__, version=SimpleNamespace(cuda=None, hip=hip) + ) + monkeypatch.setattr(extensions, "_load_torch_extension_tools", lambda: (fake_torch, None, None)) + with pytest.raises(RuntimeError, match="requires CUDA-enabled PyTorch"): + extensions.get_extensions() + + +def test_rope_only_wheel_rejects_missing_torch(monkeypatch): + monkeypatch.setenv("RL_KERNEL_ROPE_ONLY", "1") + monkeypatch.setattr(extensions, "_load_torch_extension_tools", lambda: (None, None, None)) + with pytest.raises(RuntimeError, match="PyTorch is unavailable"): + extensions.get_extensions() diff --git a/tests/ops/rope/test_multi_axis_rope.py b/tests/ops/rope/test_multi_axis_rope.py new file mode 100644 index 000000000..9f30c35b0 --- /dev/null +++ b/tests/ops/rope/test_multi_axis_rope.py @@ -0,0 +1,517 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Accuracy is independent FP64 complex algebra; invariance compares raw bytes.""" + +import gc +import math +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.ops.rope.multi_axis_rope import ( + build_multi_axis_cos_sin, + get_multi_axis_rope_op, + multi_axis_rope, + qwen_image_positions, +) + + +def assert_bits(a, b): + assert a.dtype == b.dtype and a.shape == b.shape + assert torch.equal(a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8)) + + +def complex_reference(x, cos, sin, *, transpose=False): + pairs = x.double().reshape(*x.shape[:-1], 64, 2) + z = torch.view_as_complex(pairs.contiguous()) + phase = torch.complex(cos.double(), sin.double()).unsqueeze(-2) + if transpose: + phase = phase.conj() + return torch.view_as_real(z * phase).flatten(-2) + + +@pytest.fixture(params=["reference", "cuda", "triton"]) +def implementation(request): + backend = request.param + device = "cpu" if backend == "reference" else "cuda" + if device == "cuda" and (not torch.cuda.is_available() or torch.version.hip is not None): + pytest.skip("NVIDIA CUDA required") + if backend == "cuda": + from rl_engine.backends.cuda.rope.multi_axis_rope import native_available + + if not native_available(): + pytest.skip("build native multi_axis_rope or use the focused CUDA validation runner") + if backend == "triton": + pytest.importorskip("triton") + return get_multi_axis_rope_op(backend, device=device), device + + +def inputs(device, dtype=torch.float32, *, batch=2, seq=7, heads=3, batched=False): + generator = torch.Generator().manual_seed(386) + x = torch.randn(batch, seq, heads, 128, generator=generator).to(device, dtype) + shape = (batch, seq, 3) if batched else (seq, 3) + pos = torch.randint(-42, 53, shape, generator=generator) + cos, sin = build_multi_axis_cos_sin(pos, device=device) + return x, cos, sin + + +@pytest.mark.parametrize("height,width", [(2, 4), (3, 5), (64, 64), (83, 83), (104, 58)]) +def test_positions_and_axis_frequencies(height, width): + image, text = qwen_image_positions(height, width, 3) + # Explicit flattened endpoints and a row boundary catch swapped H/W and odd-grid shifts. + assert image[0].tolist() == [0, -((height + 1) // 2), -((width + 1) // 2)] + assert image[-1].tolist() == [0, height // 2 - 1, width // 2 - 1] + assert image[width].tolist() == [0, -((height + 1) // 2) + 1, -((width + 1) // 2)] + assert text[:, 0].tolist() == list( + range(max(height // 2, width // 2), max(height // 2, width // 2) + 3) + ) + cos, sin = build_multi_axis_cos_sin(image[[0, -1]]) + # Independent scalar FP64 trig oracle, per-axis exponents reset at 8 and 36 pairs. + for row, position in enumerate(image[[0, -1]].tolist()): + offset = 0 + for axis, width_axis in enumerate((16, 56, 56)): + for pair in range(width_axis // 2): + angle = position[axis] / (10000.0 ** (2 * pair / width_axis)) + assert abs(cos[row, offset + pair].item() - math.cos(angle)) < 4e-6 + assert abs(sin[row, offset + pair].item() - math.sin(angle)) < 4e-6 + offset += width_axis // 2 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("batched", [False, True]) +def test_forward_backward_accuracy(implementation, dtype, batched): + op, device = implementation + x, cos, sin = inputs(device, dtype, batched=batched) + x.requires_grad_() + grad = torch.randn_like(x) + result = op(x, cos, sin) + dx = torch.autograd.grad(result, x, grad)[0] + tolerance = {torch.float32: 8e-7, torch.float16: 0.002, torch.bfloat16: 0.02}[dtype] + # Make the oracle's batch axis explicit when the candidate reads shared tables. + c = cos if batched else cos.unsqueeze(0).expand(x.shape[0], -1, -1) + s = sin if batched else sin.unsqueeze(0).expand(x.shape[0], -1, -1) + torch.testing.assert_close( + result.double(), complex_reference(x, c, s), atol=tolerance, rtol=tolerance + ) + torch.testing.assert_close( + dx.double(), + complex_reference(grad, c, s, transpose=True), + atol=tolerance, + rtol=tolerance, + ) + + +def test_adjacent_pair_layout_not_rotate_half(implementation): + op, device = implementation + x = torch.arange(128, dtype=torch.float32, device=device).reshape(1, 1, 1, 128) + cos = torch.zeros(1, 64, device=device) + sin = torch.ones_like(cos) + y = op(x, cos, sin) + assert_bits(y[..., 0::2], -x[..., 1::2]) + assert_bits(y[..., 1::2], x[..., 0::2]) + + +def test_batch_position_padding_and_chunk_invariance(implementation): + op, device = implementation + x, cos, sin = inputs(device, batched=True, batch=3, seq=19) + x.requires_grad_() + grad = torch.randn_like(x) + y = op(x, cos, sin) + dx = torch.autograd.grad(y, x, grad)[0] + # Reordering samples with different positions must not reuse sample zero's table. + permutation = torch.tensor([2, 0, 1], device=device) + assert_bits(op(x[permutation], cos[permutation], sin[permutation]), y[permutation]) + for b in range(3): + for lo, hi in [(0, 1), (1, 8), (8, 19)]: + xx = x[b : b + 1, lo:hi].detach().requires_grad_() + yy = op(xx, cos[b, lo:hi], sin[b, lo:hi]) + dd = torch.autograd.grad(yy, xx, grad[b : b + 1, lo:hi])[0] + assert_bits(yy, y[b : b + 1, lo:hi]) + assert_bits(dd, dx[b : b + 1, lo:hi]) + padded = torch.cat((x.detach(), torch.randn(3, 5, 3, 128, device=device)), dim=1) + pc = torch.cat((cos, torch.ones(3, 5, 64, device=device)), dim=1) + ps = torch.cat((sin, torch.zeros(3, 5, 64, device=device)), dim=1) + assert_bits(op(padded, pc, ps)[:, :19], y) + + +def test_noncontiguous_input_tables_and_gradient(implementation): + op, device = implementation + base, cos, sin = inputs(device, batched=True, seq=14, heads=4) + x = base[:, ::2, ::2].detach().requires_grad_() + c, s = cos[:, ::2], sin[:, ::2] + grad = torch.randn_like(base)[:, ::2, ::2] + result = op(x, c, s) + dx = torch.autograd.grad(result, x, grad)[0] + xx = x.detach().contiguous().requires_grad_() + yy = op(xx, c.contiguous(), s.contiguous()) + dd = torch.autograd.grad(yy, xx, grad.contiguous())[0] + assert_bits(result, yy) + assert_bits(dx, dd) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("layout", ["head_dim_stride_2", "head_dim_outermost"]) +def test_arbitrary_head_dimension_stride_and_gradient(implementation, dtype, layout): + op, device = implementation + _, cos, sin = inputs(device, dtype) + if layout == "head_dim_stride_2": + x = torch.randn(2, 7, 3, 256, device=device, dtype=dtype)[..., ::2] + assert x.stride(-1) == 2 + else: + x = torch.randn(128, 2, 7, 3, device=device, dtype=dtype).permute(1, 2, 3, 0) + assert x.stride(-1) > max(x.stride()[:-1]) + x.requires_grad_() + grad = torch.randn(2, 7, 3, 256, device=device, dtype=dtype)[..., ::2] + result = op(x, cos, sin) + dx = torch.autograd.grad(result, x, grad)[0] + tolerance = {torch.float32: 8e-7, torch.float16: 0.002, torch.bfloat16: 0.02}[dtype] + torch.testing.assert_close( + result.double(), complex_reference(x, cos, sin), atol=tolerance, rtol=tolerance + ) + torch.testing.assert_close( + dx.double(), + complex_reference(grad, cos, sin, transpose=True), + atol=tolerance, + rtol=tolerance, + ) + + +@pytest.mark.parametrize("batch,seq", [(0, 7), (2, 0), (0, 0)]) +def test_empty_forward_backward(implementation, batch, seq): + op, device = implementation + x, cos, sin = inputs(device, batch=batch, seq=seq) + x.requires_grad_() + result = op(x, cos, sin) + result.sum().backward() + assert result.shape == x.shape and x.grad.shape == x.shape + + +def test_second_derivative(implementation): + op, device = implementation + x, cos, sin = inputs(device, batch=1, seq=2) + x.requires_grad_() + y = op(x, cos, sin) + first = torch.autograd.grad(y.square().sum(), x, create_graph=True)[0] + second = torch.autograd.grad(first.sum(), x)[0] + # R^T R = (cos^2 + sin^2) I, including the small table rounding error. + expected = ( + (2 * (cos.double().square() + sin.double().square())) + .repeat_interleave(2, -1)[None, :, None, :] + .expand_as(x) + ) + torch.testing.assert_close(second.double(), expected, atol=8e-7, rtol=8e-7) + + +@pytest.mark.parametrize( + "problem", ["shape", "dtype", "table_shape", "table_dtype", "table_grad", "device"] +) +def test_reject_invalid_inputs(implementation, problem): + op, device = implementation + x, cos, sin = inputs(device) + if problem == "shape": + x = x[..., :64] + elif problem == "dtype": + x = x.double() + elif problem == "table_shape": + sin = sin[:-1] + elif problem == "table_dtype": + cos = cos.half() + elif problem == "table_grad": + cos.requires_grad_() + elif problem == "device": + sin = torch.empty_like(sin, device="meta") + with pytest.raises((ValueError, TypeError)): + op(x, cos, sin) + + +def test_explicit_backend_and_reference_dispatch(monkeypatch): + from rl_engine.runtime.registry import KernelRegistry + + op = KernelRegistry().get_op("multi_axis_rope", device="cpu") + assert op.trace()["actual_backend"].startswith("rlkernel.reference.") + from rl_engine.backends.cuda.rope import multi_axis_rope as cuda_module + + monkeypatch.setattr(cuda_module, "native_available", lambda: False) + with pytest.raises(RuntimeError, match="unavailable"): + get_multi_axis_rope_op("cuda") + with pytest.raises(ValueError, match="unknown"): + get_multi_axis_rope_op("typo") + x, cos, sin = inputs("cpu") + assert_bits(multi_axis_rope(x, cos, sin), op(x, cos, sin)) + + +@pytest.mark.parametrize("platform", ["cpu", "npu", "rocm", "musa"]) +def test_non_cuda_registry_keeps_multi_axis_reference(platform, monkeypatch): + """Exercise candidate selection on CPU; no accelerator execution is certified.""" + from rl_engine.runtime.registry import KernelRegistry, OpBackend + + registry = KernelRegistry() + assert registry._priority_map[platform]["multi_axis_rope"] == [ + OpBackend.PYTORCH_MULTI_AXIS_ROPE + ] + monkeypatch.setattr(registry, "_platform", lambda: platform) + op = registry.get_op("multi_axis_rope") + assert op.trace()["actual_backend"] == "rlkernel.reference.qwenimage_multi_axis_rope.v1" + x, cos, sin = inputs("cpu") + assert_bits(op(x, cos, sin), get_multi_axis_rope_op("reference")(x, cos, sin)) + + +@pytest.mark.parametrize("exception_type", [OSError, RuntimeError]) +def test_operator_wheel_load_error_propagates(exception_type, monkeypatch): + from rl_engine.backends.cuda.rope import multi_axis_rope as cuda_module + + error = exception_type("probe library load failure") + + def fail_load(path): + raise error + + monkeypatch.setattr( + cuda_module, + "torch", + SimpleNamespace( + version=SimpleNamespace(hip=None), + ops=SimpleNamespace(rl_kernel_native=SimpleNamespace(), load_library=fail_load), + ), + ) + monkeypatch.setattr( + cuda_module.importlib.util, + "find_spec", + lambda name: SimpleNamespace(origin="probe.so"), + ) + with pytest.raises(exception_type, match="probe library load failure") as caught: + cuda_module.native_available() + assert caught.value is error + + +def test_cuda_registry_uses_native_when_present(): + from rl_engine.backends.cuda.rope.multi_axis_rope import native_available + from rl_engine.runtime.registry import KernelRegistry + + if not torch.cuda.is_available() or not native_available(): + pytest.skip("native CUDA multi_axis_rope extension required") + op = KernelRegistry().get_op("multi_axis_rope", device="cuda") + assert op.trace()["actual_backend"] == "rlkernel.cuda.qwenimage_multi_axis_rope.v1" + x, cos, sin = inputs("cuda") + torch.testing.assert_close( + op(x, cos, sin).double(), complex_reference(x, cos, sin), atol=8e-7, rtol=8e-7 + ) + + +def test_cuda_registry_falls_back_to_triton(monkeypatch): + from rl_engine.runtime.registry import KernelRegistry + + if not torch.cuda.is_available() or torch.version.hip is not None: + pytest.skip("NVIDIA CUDA required") + pytest.importorskip("triton") + from rl_engine.backends.cuda.rope import multi_axis_rope as cuda_module + + monkeypatch.setattr(cuda_module, "native_available", lambda: False) + op = KernelRegistry().get_op("multi_axis_rope", device="cuda") + assert op.trace()["actual_backend"] == "rlkernel.triton.qwenimage_multi_axis_rope.v1" + x, cos, sin = inputs("cuda") + torch.testing.assert_close( + op(x, cos, sin).double(), complex_reference(x, cos, sin), atol=8e-7, rtol=8e-7 + ) + + +@pytest.mark.parametrize( + "bad", + [torch.zeros(2, 3), torch.zeros(2, 4, dtype=torch.int64), torch.full((2, 3), 4096)], +) +def test_reject_invalid_positions(bad): + with pytest.raises((TypeError, ValueError)): + build_multi_axis_cos_sin(bad) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_cuda_triton_byte_parity_and_graph(dtype): + if not torch.cuda.is_available() or torch.version.hip is not None: + pytest.skip("NVIDIA CUDA required") + from rl_engine.backends.cuda.rope.multi_axis_rope import native_available + + if not native_available(): + pytest.skip("native multi_axis_rope extension required") + pytest.importorskip("triton") + cuda, triton = get_multi_axis_rope_op("cuda"), get_multi_axis_rope_op("triton") + x, cos, sin = inputs("cuda", dtype, batched=True, seq=17, heads=24) + x.requires_grad_() + grad = torch.randn_like(x) + yc, yt = cuda(x, cos, sin), triton(x, cos, sin) + dc = torch.autograd.grad(yc, x, grad)[0] + dt = torch.autograd.grad(yt, x, grad)[0] + assert_bits(yc, yt) + assert_bits(dc, dt) + reference = get_multi_axis_rope_op("reference") + yr = reference(x, cos, sin) + dr = torch.autograd.grad(yr, x, grad)[0] + assert_bits(yc, yr) + assert_bits(dc, dr) + # Capture forward + autograd after warmup on a side stream. + for op in (cuda, triton): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + static_x = x.detach().clone().requires_grad_() + static_grad = grad.clone() + for _ in range(3): + torch.autograd.grad(op(static_x, cos, sin), static_x, static_grad) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + gc.collect() + was_enabled = gc.isenabled() + gc.disable() + try: + with torch.cuda.graph(graph, stream=stream): + result = op(static_x, cos, sin) + derivative = torch.autograd.grad(result, static_x, static_grad)[0] + finally: + if was_enabled: + gc.enable() + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + assert_bits(result, yc) + assert_bits(derivative, dc) + # Replay must consume updated static buffers, not just retain captured outputs. + with torch.no_grad(): + static_x.copy_(x * 2) + static_grad.copy_(grad * 0.5) + graph.replay() + torch.cuda.synchronize() + expected = op(static_x, cos, sin) + expected_dx = torch.autograd.grad(expected, static_x, static_grad)[0] + assert_bits(result, expected) + assert_bits(derivative, expected_dx) + + +def test_native_entrypoint_checks_and_nondefault_stream(): + if not torch.cuda.is_available() or torch.version.hip is not None: + pytest.skip("NVIDIA CUDA required") + from rl_engine.backends.cuda.rope.multi_axis_rope import native_available + + if not native_available(): + pytest.skip("native multi_axis_rope extension required") + x, cos, sin = inputs("cuda") + native = torch.ops.rl_kernel_native.multi_axis_rope + for c, s in [(cos.half(), sin), (cos, sin[:-1]), (cos.cpu(), sin)]: + with pytest.raises(RuntimeError): + native(x, c, s, False) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + xx = x * 3 + y = native(xx, cos, sin, False) + torch.cuda.current_stream().wait_stream(stream) + torch.testing.assert_close(y.double(), complex_reference(xx, cos, sin), atol=2e-6, rtol=2e-6) + + +def test_device_guard(): + if torch.cuda.device_count() < 2 or torch.version.hip is not None: + pytest.skip("requires two NVIDIA GPUs") + from rl_engine.backends.cuda.rope.multi_axis_rope import native_available + + if not native_available(): + pytest.skip("native multi_axis_rope extension required") + with torch.cuda.device(0): + x, cos, sin = inputs("cuda:1") + for backend in ("cuda", "triton"): + y = get_multi_axis_rope_op(backend)(x, cos, sin) + assert y.device == x.device + assert torch.cuda.current_device() == 0 + torch.testing.assert_close( + y.double(), complex_reference(x, cos, sin), atol=8e-7, rtol=8e-7 + ) + + +@pytest.mark.parametrize("default_device", ["meta", "cuda"]) +def test_canonical_tables_ignore_default_device(default_device): + if default_device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA context requires CUDA") + positions, text = qwen_image_positions(3, 5, 7) + expected = build_multi_axis_cos_sin(positions) + with torch.device(default_device): + image_again, text_again = qwen_image_positions(3, 5, 7) + actual = build_multi_axis_cos_sin(positions, device="cpu") + assert image_again.device.type == "cpu" and text_again.device.type == "cpu" + assert_bits(image_again, positions) + assert_bits(text_again, text) + for a, b in zip(actual, expected, strict=True): + assert_bits(a, b) + + +def test_packed_qkv_and_independent_table_strides(implementation): + op, device = implementation + packed = torch.randn(2, 17, 3, 3, 256, device=device) + x = packed[:, :, 0, :, ::2].detach().requires_grad_() + cos_base = torch.randn(2, 17, 128, device=device) + cos = cos_base[..., ::2] + sin = torch.randn(17, 2, 64, device=device).transpose(0, 1) + grad = torch.randn_like(x) + result = op(x, cos, sin) + dx = torch.autograd.grad(result, x, grad)[0] + assert result.is_contiguous() + torch.testing.assert_close( + result.double(), complex_reference(x, cos, sin), atol=2e-6, rtol=2e-6 + ) + torch.testing.assert_close( + dx.double(), complex_reference(grad, cos, sin, transpose=True), atol=2e-6, rtol=2e-6 + ) + if device == "cuda": + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU]) as profile: + op(x, cos, sin) + copies = {event.key for event in profile.key_averages()} & { + "aten::contiguous", + "aten::clone", + "aten::copy_", + } + assert not copies, f"RoPE materialized its strided inputs: {copies}" + + +@pytest.mark.parametrize("backend", ["cuda", "triton"]) +def test_dispatcher_fake_autograd_and_fullgraph_compile(backend): + if not torch.cuda.is_available() or torch.version.hip is not None: + pytest.skip("NVIDIA CUDA required") + if backend == "cuda": + from rl_engine.backends.cuda.rope.multi_axis_rope import _rotation, native_available + + if not native_available(): + pytest.skip("build native RoPE library") + else: + pytest.importorskip("triton") + from rl_engine.backends.shared.triton.rope.multi_axis_rope import _rotation + op = get_multi_axis_rope_op(backend) + x, cos, sin = inputs("cuda", batched=True, seq=5, heads=2) + x.requires_grad_() + checks = torch.library.opcheck(_rotation, (x, cos, sin, False)) + assert all(value == "SUCCESS" for value in checks.values()) + compiled = torch.compile(op, fullgraph=True) + grad = torch.randn_like(x) + expected = op(x, cos, sin) + expected_dx = torch.autograd.grad(expected, x, grad)[0] + result = compiled(x, cos, sin) + dx = torch.autograd.grad(result, x, grad)[0] + assert_bits(result, expected) + assert_bits(dx, expected_dx) + + +@pytest.mark.parametrize("backend", ["cuda", "triton"]) +@pytest.mark.ws1 +def test_standard_ws1_four_judgments(backend): + if not torch.cuda.is_available() or torch.version.hip is not None: + pytest.skip("NVIDIA CUDA required") + if backend == "cuda": + from rl_engine.backends.cuda.rope.multi_axis_rope import native_available + + if not native_available(): + pytest.skip("build native RoPE library") + else: + pytest.importorskip("triton") + from rl_engine.validation.operators.qwen_image_rope import run_multi_axis_rope_ws1 + + result = run_multi_axis_rope_ws1(backend) + assert result["passed"] + assert result["trace"]["actual_backend"].startswith(f"rlkernel.{backend}.") + for phase in ("forward", "gradient"): + assert result[phase]["accuracy_reports"] + assert result[phase]["invariance_reports"] + assert all(report["passed"] for report in result[phase]["invariance_reports"]) diff --git a/tests/validation/operators/test_operator_inputs.py b/tests/validation/operators/test_operator_inputs.py index 9e5015bf4..0416d1540 100644 --- a/tests/validation/operators/test_operator_inputs.py +++ b/tests/validation/operators/test_operator_inputs.py @@ -8,6 +8,9 @@ import pytest import torch +from rl_engine.config.workload import LogicalSample +from rl_engine.ops.rope.multi_axis_rope import build_multi_axis_cos_sin, qwen_image_positions +from rl_engine.validation.operators.gradient_adapters import _row_inputs from rl_engine.validation.operators.operator_inputs import make_operator_inputs, operator_shape_name from rl_engine.validation.operators.operator_specs import ( make_candidate, @@ -103,6 +106,86 @@ def test_cp_attention_operator_spec_registers_backward_grad_inputs(): assert candidate.name == "pytorch-cp_attention" +def test_multi_axis_rope_operator_spec_and_seeded_inputs(): + args = _args(op="multi_axis_rope", batch=2, seq=7, n_heads=3, input_mode="random") + case = make_operator_case(args, torch.float32, torch.device("cpu")) + inputs = make_operator_inputs("multi_axis_rope", args, torch.float32, torch.device("cpu")) + repeated = make_operator_inputs("multi_axis_rope", args, torch.float32, torch.device("cpu")) + assert case.op_class == "elementwise" + assert case.grad_input_names == ("x",) + assert inputs["x"].shape == (2, 7, 3, 128) + assert inputs["cos"].shape == inputs["sin"].shape == (7, 64) + for name in inputs: + assert torch.equal(inputs[name], repeated[name]) + assert operator_shape_name("multi_axis_rope", args) == "2x7x3x128" + with pytest.raises(ValueError, match="128"): + make_operator_inputs( + "multi_axis_rope", _args(head_dim=64), torch.float32, torch.device("cpu") + ) + + +@pytest.mark.parametrize( + "seq,side", [(0, 1), (1, 1), (7, 3), (15, 4), (16, 4), (4096, 64), (6889, 83)] +) +def test_multi_axis_rope_validation_uses_centered_image_grid_prefix(seq, side): + inputs = make_operator_inputs( + "multi_axis_rope", _args(seq=seq, n_heads=1), torch.float32, torch.device("cpu") + ) + image, _ = qwen_image_positions(side, side, 0) + cos, sin = build_multi_axis_cos_sin(image[:seq]) + assert torch.equal(inputs["cos"], cos) + assert torch.equal(inputs["sin"], sin) + + +@pytest.mark.parametrize("layout", ["packed", "permuted", "chunked", "padded"]) +def test_multi_axis_rope_validation_preserves_logical_tables_across_layouts(layout): + samples = ( + LogicalSample("short", tuple(range(7)), prompt_len=1, seq_len=7), + LogicalSample("long", tuple(range(15)), prompt_len=2, seq_len=15), + ) + tokens = { + (token.sample_id, token.token_position): token + for sample in samples + for token in sample.tokens() + } + expected = { + sample.sample_id: make_operator_inputs( + "multi_axis_rope", + _args(seq=sample.seq_len, n_heads=1), + torch.float32, + torch.device("cpu"), + ) + for sample in samples + } + keys = list(tokens) + if layout == "permuted": + keys.reverse() + elif layout == "chunked": + keys = [("short", 5), ("short", 6), ("long", 12)] + elif layout == "padded": + keys = [None, ("long", 4), None, ("short", 6), ("short", 0), None] + inputs = _row_inputs( + "multi_axis_rope", + keys, + tokens, + {}, + device=torch.device("cpu"), + dtype=torch.float32, + hidden=1, + vocab_size=17, + n_heads=1, + head_dim=128, + ) + for row, key in enumerate(keys): + if key is None: + assert torch.equal(inputs["cos"][row], torch.ones(64)) + assert torch.equal(inputs["sin"][row], torch.zeros(64)) + else: + sample_id, position = key + for name in ("cos", "sin"): + assert torch.equal(inputs[name][row], expected[sample_id][name][position]) + + def test_constant_linear_logp_inputs_match_operator_contract(): args = _args(input_mode="constant", constant_value=0.5, token_value=3) inputs = make_operator_inputs("linear_logp", args, torch.float32, torch.device("cpu")) diff --git a/tools/validation/operators/check_multi_axis_rope_install.py b/tools/validation/operators/check_multi_axis_rope_install.py new file mode 100644 index 000000000..4fa42d15f --- /dev/null +++ b/tools/validation/operators/check_multi_axis_rope_install.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Run from outside the checkout after installing a RoPE-enabled wheel. + +Does not change sys.path or compile/load a JIT extension. --package-root makes +the expected installation location an explicit assertion in package validation. +""" + +import argparse +import importlib.util +import json +from pathlib import Path + +import torch + +import rl_engine +from rl_engine.backends import extension +from rl_engine.ops.rope.multi_axis_rope import ( + build_multi_axis_cos_sin, + get_multi_axis_rope_op, + qwen_image_positions, +) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--package-root", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + package = Path(rl_engine.__file__).resolve().parent + assert package.is_relative_to(args.package_root.resolve()), package + license_path = package / "backends/shared/triton/rope/FLASH_ATTENTION_LICENSE" + assert "Redistribution and use" in license_path.read_text(encoding="utf-8") + positions, _ = qwen_image_positions(3, 5, 0) + cos, sin = build_multi_axis_cos_sin(positions, device="cuda") + packed = torch.randn(2, 15, 3, 3, 128, device="cuda", dtype=torch.bfloat16) + x = packed[:, :, 0].detach().requires_grad_() + grad = torch.randn_like(x) + gold = get_multi_axis_rope_op("reference")(x, cos, sin) + gold_dx = torch.autograd.grad(gold, x, grad)[0] + outcomes = {} + for backend in ("cuda", "triton"): + op = get_multi_axis_rope_op(backend, device=x.device) + for mode, fn in (("eager", op), ("compiled", torch.compile(op, fullgraph=True))): + result = fn(x, cos, sin) + dx = torch.autograd.grad(result, x, grad)[0] + assert torch.equal(result.view(torch.uint8), gold.contiguous().view(torch.uint8)) + assert torch.equal(dx.view(torch.uint8), gold_dx.contiguous().view(torch.uint8)) + outcomes[f"{backend}_{mode}"] = "passed" + outcomes[f"{backend}_trace"] = op.trace() + native = importlib.util.find_spec("rl_engine._rope_cuda") + report = { + "package": str(package), + "companion_library": native.origin if native else None, + "full_extension_library": getattr(extension._C, "__file__", None), + "native_dispatcher_registered": hasattr(torch.ops.rl_kernel_native, "multi_axis_rope"), + "torch": torch.__version__, + "gpu": torch.cuda.get_device_name(), + "license_present": True, + "packed_qkv_stride": list(x.stride()), + "checks": outcomes, + "passed": True, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/operators/check_qwenimage_rope_upstream.py b/tools/validation/operators/check_qwenimage_rope_upstream.py new file mode 100644 index 000000000..153000904 --- /dev/null +++ b/tools/validation/operators/check_qwenimage_rope_upstream.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Audit table and rotation semantics against a checksum-pinned upstream file. + +Supply transformer_qwenimage.py from Diffusers revision +40f63690da440e8c382ab53e67508835deb1ed2f. Only the two audited symbol definitions +are executed, after verifying the source hash. No Diffusers installation or model +checkpoint is required. This checks CPU semantics, not full-model integration. +""" + +from __future__ import annotations + +import argparse +import ast +import functools +import hashlib +import json +import logging +import sys +from pathlib import Path + +import torch + +ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(ROOT)) +SOURCE_SHA256 = "0bad11901deae69f05aec01fd464e51e38582d455068f2851ee665036179321f" + + +def main(): + from rl_engine.ops.rope.multi_axis_rope import ( + build_multi_axis_cos_sin, + get_multi_axis_rope_op, + qwen_image_positions, + ) + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + source = args.source.read_text(encoding="utf-8") + digest = hashlib.sha256(source.encode()).hexdigest() + if digest != SOURCE_SHA256: + raise ValueError("upstream source does not match the pinned revision checksum") + symbols = {"QwenEmbedRope", "apply_rotary_emb_qwen"} + nodes = [n for n in ast.parse(source).body if getattr(n, "name", None) in symbols] + if len(nodes) != 2: + raise ValueError("missing audited upstream definitions") + namespace = { + "torch": torch, + "nn": torch.nn, + "lru_cache_unless_export": functools.lru_cache, + "logger": logging.getLogger("upstream_audit"), + } + exec( + compile(ast.Module(body=nodes, type_ignores=[]), str(args.source), "exec"), + namespace, + ) + embed = namespace["QwenEmbedRope"](theta=10000, axes_dim=[16, 56, 56], scale_rope=True) + apply_upstream = namespace["apply_rotary_emb_qwen"] + op = get_multi_axis_rope_op("reference") + torch.manual_seed(386) + results = [] + for height, width in ((2, 3), (3, 5), (64, 64), (83, 83), (104, 58)): + positions = qwen_image_positions(height, width, 7) + upstream_tables = embed((1, height, width), device=torch.device("cpu"), max_txt_seq_len=7) + for stream, pos, upstream in zip( + ("image", "text"), positions, upstream_tables, strict=True + ): + cos, sin = build_multi_axis_cos_sin(pos) + assert torch.equal(cos, upstream.real) and torch.equal(sin, upstream.imag) + for dtype in (torch.float32, torch.float16, torch.bfloat16): + x = torch.randn(1, len(pos), 2, 128, dtype=dtype, requires_grad=True) + grad = torch.randn_like(x) + expected = apply_upstream(x, upstream, use_real=False) + actual = op(x, cos, sin) + eg = torch.autograd.grad(expected, x, grad)[0] + ag = torch.autograd.grad(actual, x, grad)[0] + tol = {torch.float32: 8e-7, torch.float16: 0.002, torch.bfloat16: 0.02}[dtype] + torch.testing.assert_close(actual, expected, atol=tol, rtol=tol) + torch.testing.assert_close(ag, eg, atol=tol, rtol=tol) + results.append( + { + "grid": [height, width], + "stream": stream, + "dtype": str(dtype), + "table_bitwise": True, + "forward_max_abs": (actual.float() - expected.float()).abs().max().item(), + "dx_max_abs": (ag.float() - eg.float()).abs().max().item(), + } + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps({"source_sha256": digest, "cases": results}, indent=2) + "\n") + print( + f"{len(results)} upstream forward/backward cases passed; " + "all image/text tables exactly equal" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/operators/validate_multi_axis_rope.py b/tools/validation/operators/validate_multi_axis_rope.py new file mode 100644 index 000000000..d18ef7cbb --- /dev/null +++ b/tools/validation/operators/validate_multi_axis_rope.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Compile only the production RoPE CUDA source for focused GPU validation. + +Loads the production TORCH_LIBRARY registration directly. Normal package builds +include the same source in rl_engine._C. No generated binding or monkeypatch is +used. This validates the operator library, not every unrelated native operator. +""" + +from __future__ import annotations + +import argparse +import json +import os +import runpy +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(ROOT)) + + +def load_focused_extension(): + import torch + from torch.utils.cpp_extension import CUDA_HOME, load + + from rl_engine.backends.cuda.rope.multi_axis_rope import native_available + + if native_available(): + print("Using the installed production RoPE dispatcher library") + return + + build = ROOT / "artifacts" / "multi_axis_rope_native" + build.mkdir(parents=True, exist_ok=True) + # Conda CUDA toolkits use lib/, whereas the Windows SDK convention is lib/x64/. + link_flags = [] + if os.name == "nt" and CUDA_HOME and (Path(CUDA_HOME) / "lib/cudart.lib").is_file(): + link_flags.append(f'/LIBPATH:"{Path(CUDA_HOME) / "lib"}"') + os.environ["PATH"] = str(Path(sys.executable).parent) + os.pathsep + os.environ["PATH"] + load( + name="rlk_multi_axis_rope_native", + sources=[str(ROOT / "csrc/cuda/rope/multi_axis_rope.cu")], + build_directory=str(build), + extra_cflags=(["/O2", "/std:c++17"] if os.name == "nt" else ["-O3", "-std=c++17"]), + extra_cuda_cflags=["-O3", "-std=c++17", "--fmad=false"], + extra_ldflags=link_flags, + is_python_module=False, + verbose=True, + ) + assert hasattr(torch.ops.rl_kernel_native, "multi_axis_rope") + print( + f"Focused CUDA source build; torch={torch.__version__}, GPU={torch.cuda.get_device_name()}" + ) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("--ws1", choices=("cuda", "triton")) + parser.add_argument("--output", type=Path) + args, remaining = parser.parse_known_args() + load_focused_extension() + if args.ws1: + from rl_engine.validation.operators.qwen_image_rope import run_multi_axis_rope_ws1 + + result = run_multi_axis_rope_ws1(args.ws1) + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") + print(f"WS1 four judgments backend={args.ws1}: passed={result['passed']}") + raise SystemExit(0 if result["passed"] else 1) + elif args.benchmark: + sys.argv = ["benchmark_multi_axis_rope", *remaining] + if args.output: + sys.argv.extend(["--output", str(args.output)]) + runpy.run_module("benchmarks.operators.rope.benchmark_multi_axis_rope", run_name="__main__") + else: + import pytest + + raise SystemExit( + pytest.main([str(ROOT / "tests/ops/rope/test_multi_axis_rope.py"), *remaining]) + ) + + +if __name__ == "__main__": + main()