diff --git a/benchmarks/operators/norm/benchmark_adaln_modulation.py b/benchmarks/operators/norm/benchmark_adaln_modulation.py new file mode 100644 index 000000000..fc7d7dc2e --- /dev/null +++ b/benchmarks/operators/norm/benchmark_adaln_modulation.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Compare AdaLN backends on Qwen-Image image-token lengths. + +Run with --real for S=4096/6889/6032, H=3072. These S values are derived +from 8x VAE downsampling and 2x2 latent packing, not specified by #386. +""" + +import argparse +import contextlib +import hashlib +import json +import platform +import statistics +import subprocess +import sys +from pathlib import Path + +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + + +def layer_norm_baseline(x, modulation, *, eps=1e-6): + """Ordinary PyTorch math for accuracy/performance, without a replay claim.""" + shift, scale, _ = modulation.float().chunk(3, dim=-1) + norm = torch.nn.functional.layer_norm(x.float(), (x.shape[-1],), eps=eps) + return ( + (norm * (1.0 + scale[:, None]) + shift[:, None]).to(x.dtype), + modulation[:, None, 2 * x.shape[-1] :], + ) + + +def _bench(op, x, modulation, *, backward, warmup, repeat): + def run(): + x.grad = modulation.grad = None + if backward: + y, gate = op(x, modulation) + torch.autograd.backward((y, gate), (torch.ones_like(y), torch.ones_like(gate))) + else: + with torch.no_grad(): + op(x, modulation) + + for _ in range(warmup): + run() + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + allocated_before = torch.cuda.memory_allocated() + samples = [] + for _ in range(repeat): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + return { + "median_ms": statistics.median(samples), + "samples_ms": samples, + "peak_allocated_bytes": torch.cuda.max_memory_allocated(), + "allocated_before_samples_bytes": allocated_before, + } + + +def parse_args(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--real", action="store_true", help="run all three Qwen-Image resolutions") + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--repeat", type=int, default=10) + parser.add_argument("--dtype", choices=("bf16", "fp32"), default="bf16") + parser.add_argument("--seed", type=int, default=386) + parser.add_argument("--json", action="store_true", help="emit one JSON report on stdout") + args = parser.parse_args(argv) + if args.warmup < 0 or args.repeat < 1: + parser.error("warmup must be nonnegative and repeat must be positive") + return args + + +def _command_output(command): + try: + return subprocess.check_output( + command, cwd=REPO_ROOT, text=True, stderr=subprocess.DEVNULL, timeout=5 + ).strip() + except (OSError, subprocess.SubprocessError): + return None + + +def _sha256(path): + if not path or not Path(path).is_file(): + return None + with open(path, "rb") as source: + digest = hashlib.sha256() + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def main(argv=None): + args = parse_args(argv) + if not torch.cuda.is_available(): + raise SystemExit("CUDA device required") + # Imports and extension diagnostics must not corrupt the JSON stdout stream. + with contextlib.redirect_stdout(sys.stderr): + from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp + from rl_engine.runtime.adaln import SHARED_PROFILE, AdaLNDispatch + + backends = [("pytorch", NativeAdaLNModulationOp())] + skipped = [] + try: + from rl_engine.backends.shared.triton.norm.adaln_modulation import ( + TritonAdaLNModulationOp, + ) + + backends.append(("triton", TritonAdaLNModulationOp())) + except ImportError as exc: + skipped.append({"backend": "triton", "reason": str(exc)}) + try: + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + backends.append(("cuda", CudaAdaLNModulationOp())) + except RuntimeError as exc: + skipped.append({"backend": "cuda", "reason": str(exc)}) + try: + import triton + + triton_version = triton.__version__ + except ImportError: + triton_version = None + environment = { + "gpu": torch.cuda.get_device_name(), + "device": str(torch.cuda.current_device()), + "capability": list(torch.cuda.get_device_capability()), + "cuda_runtime": torch.version.cuda, + "torch": torch.__version__, + "python": platform.python_version(), + "triton": triton_version, + "driver_inventory": _command_output( + ["nvidia-smi", "--query-gpu=uuid,name,driver_version", "--format=csv,noheader"] + ), + "source_revision": _command_output(["git", "rev-parse", "HEAD"]), + "source_status": _command_output(["git", "status", "--porcelain=v1"]), + } + source_paths = [ + "benchmarks/operators/norm/benchmark_adaln_modulation.py", + "rl_engine/runtime/adaln.py", + "rl_engine/reference/norm/adaln_modulation.py", + "rl_engine/backends/cuda/norm/adaln_modulation.py", + "rl_engine/backends/shared/triton/norm/adaln_modulation.py", + "rl_engine/ops/autograd/vjp_fp32.py", + "csrc/cuda/norm/adaln_modulation.cu", + "csrc/bindings/ops.cpp", + "setup.py", + "build_tools/extensions.py", + ] + environment["source_sha256"] = {path: _sha256(REPO_ROOT / path) for path in source_paths} + backends = [ + ( + name, + AdaLNDispatch( + op, + { + "requested_backend": name, + "fallback": False, + "fallback_reason": None, + "arithmetic_profile": SHARED_PROFILE, + }, + hidden=3072, + ), + ) + for name, op in backends + ] + backends.append(("layer_norm_baseline", layer_norm_baseline)) + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 + generator = torch.Generator(device="cuda").manual_seed(args.seed) + results = [] + sizes = (4096, 6889, 6032) if args.real else (64,) + for seq in sizes: + x = torch.randn( + 1, seq, 3072, device="cuda", dtype=dtype, generator=generator, requires_grad=True + ) + modulation = torch.randn( + 1, 9216, device="cuda", dtype=dtype, generator=generator, requires_grad=True + ) + for name, op in backends: + # Observe outside timing; retain the original bare-operator event interval. + with torch.no_grad(): + op(x, modulation) + trace = ( + dict(op.trace) + if isinstance(op, AdaLNDispatch) + else { + "actual_backend": "pytorch", + "arithmetic_profile": "ordinary-layer-norm", + "observation": "successful_forward", + "strict": False, + "profiler_observed": False, + "fallback": False, + } + ) + trace["binary_sha256"] = _sha256(trace.get("binary_path")) + timed_op = op.implementation if isinstance(op, AdaLNDispatch) else op + for backward in (False, True): + measured = _bench( + timed_op, + x, + modulation, + backward=backward, + warmup=args.warmup, + repeat=args.repeat, + ) + results.append( + { + "backend": name, + "shape": list(x.shape), + "backward": backward, + "trace": trace, + **measured, + } + ) + report = { + "schema_version": "adaln-benchmark-v1", + "environment": environment, + "dtype": args.dtype, + "seed": args.seed, + "warmup": args.warmup, + "repeat": args.repeat, + "timing": "CUDA events; forward or autograd forward+backward including ones_like upstreams", + "skipped_backends": skipped, + "results": results, + } + if args.json: + print(json.dumps(report, indent=2)) + else: + print(json.dumps({key: value for key, value in report.items() if key != "results"})) + for result in results: + print( + f"S={result['shape'][1]} backend={result['backend']} " + f"backward={result['backward']} median_ms={result['median_ms']:.3f} " + f"peak_MB={result['peak_allocated_bytes'] / 2**20:.1f}" + ) + + +if __name__ == "__main__": + main() diff --git a/build_tools/extensions.py b/build_tools/extensions.py index db079e828..8de45a676 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -223,6 +223,7 @@ def get_extensions(): "csrc/cuda/activation/activation.cu", "csrc/cuda/attention/deterministic_attention.cu", ] + adaln_cuda_enabled = not is_rocm and not envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH) if is_rocm: # ROCm-tuned WS2 vocab-parallel logprob kernels; the shared # deterministic_logp_kernel.cu keeps the SM90-tuned CUDA path. @@ -235,6 +236,8 @@ def get_extensions(): else: # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. + if adaln_cuda_enabled: + cuda_sources.append("csrc/cuda/norm/adaln_modulation.cu") cuda_sources.append("csrc/cuda/collectives/deterministic_collective.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. @@ -324,6 +327,8 @@ def get_extensions(): platform_define = "-DKERNEL_ALIGN_WITH_ROCM" if is_rocm else "-DKERNEL_ALIGN_WITH_CUDA" cxx_flags = ["-O3", "-std=c++17", platform_define] + if adaln_cuda_enabled: + cxx_flags.append("-DRLK_ADALN_CUDA_ENABLED") extra_link_args = list(torch_rpath) if os.name != "nt" and not is_rocm: # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). diff --git a/ci/scripts/run_ws1_gtest.sh b/ci/scripts/run_ws1_gtest.sh index cc47dd678..c0d27c277 100755 --- a/ci/scripts/run_ws1_gtest.sh +++ b/ci/scripts/run_ws1_gtest.sh @@ -23,8 +23,20 @@ echo "[ws1-gtest] interpreter=$PY out=$OUT" tests/validation/operators/test_four_judgment_matrix.py \ tests/validation/operators/test_ws1_candidate_evidence.py \ tests/validation/operators/test_op_checks.py \ + tests/validation/operators/test_adaln_gtest.py \ + tests/backends/test_adaln_shared_bounds.py \ + tests/runtime/test_adaln_dispatch.py \ + tests/ops/norm/test_adaln_modulation.py \ tests/validation/operators/test_elementwise_inventory.py +echo "[ws1-gtest] AdaLN shared outputs and gradients against independent CPU gold" +for backend in cuda triton; do + for dtype in fp32 bf16; do + "$PY" tools/validation/operators/check_operator.py --op adaln_modulation --candidate "$backend" \ + --device cuda --dtype "$dtype" --batch 2 --seq 5 --normalized-dim 3072 --check-grad + done +done + echo "[ws1-gtest] C3/C4 CUDA + Triton smoke (silu)" "$PY" tools/validation/operators/check_forward_invariance.py \ --op silu --candidate cuda --backend-profile cuda_bf16 diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index 0c7c3e7f4..63ed671d8 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -256,6 +256,19 @@ std::vector swiglu_packed_backward_cuda( torch::Tensor dy, torch::Tensor gate_up); +#if defined(RLK_ADALN_CUDA_ENABLED) +std::vector adaln_modulation_forward_cuda( + torch::Tensor x, torch::Tensor modulation, double eps, int threads); +std::vector adaln_modulation_backward_cuda( + torch::Tensor dy, torch::Tensor dg, torch::Tensor x, + torch::Tensor modulation, double eps, int threads); +std::vector adaln_indexed_forward_cuda( + torch::Tensor x, torch::Tensor modulation, torch::Tensor index, double eps, int threads); +std::vector adaln_indexed_backward_cuda( + c10::optional dy, c10::optional dg, + torch::Tensor x, torch::Tensor modulation, torch::Tensor index, double eps, int threads); +#endif + // RMSNorm Declarations & Wrappers void rmsnorm_forward_cuda( @@ -737,6 +750,23 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Batch-invariant SwiGLU forward for [rows, 2 * intermediate]"); m.def("swiglu_packed_backward", &swiglu_packed_backward, "Batch-invariant SwiGLU backward for [rows, 2 * intermediate]"); +#if defined(RLK_ADALN_CUDA_ENABLED) + m.def("adaln_modulation_forward", &adaln_modulation_forward_cuda, + "AdaLN modulation forward CUDA", + py::arg("x"), py::arg("modulation"), py::arg("eps"), py::arg("threads") = 256); + m.def("adaln_modulation_backward", &adaln_modulation_backward_cuda, + "AdaLN modulation backward CUDA", + py::arg("dy"), py::arg("dg"), py::arg("x"), py::arg("modulation"), + py::arg("eps"), py::arg("threads") = 256); + m.def("adaln_indexed_forward", &adaln_indexed_forward_cuda, + "Indexed AdaLN small FP32 CUDA slice", + py::arg("x"), py::arg("modulation"), py::arg("index"), py::arg("eps"), + py::arg("threads") = 256); + m.def("adaln_indexed_backward", &adaln_indexed_backward_cuda, + "Indexed AdaLN FP32 ascending-token VJP", + py::arg("dy"), py::arg("dg"), py::arg("x"), py::arg("modulation"), + py::arg("index"), py::arg("eps"), py::arg("threads") = 256); +#endif // Deterministic standard-softmax attention (issue #147) m.def( diff --git a/csrc/cuda/norm/adaln_modulation.cu b/csrc/cuda/norm/adaln_modulation.cu new file mode 100644 index 000000000..a53cc0715 --- /dev/null +++ b/csrc/cuda/norm/adaln_modulation.cu @@ -0,0 +1,347 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// One CTA per logical row; H uses a fixed lower-plus-upper pairwise tree. +// Shared modulation gradients fold token 0..S-1 in FP32, without atomics. + +#include +#include +#include +#include +#include +#include + +#if defined(__CUDA_FAST_MATH__) +#error "adaln_modulation requires CUDA compilation without --use_fast_math" +#endif + +namespace { + +__device__ float fixed_sum(float* shared, int width) { + __syncthreads(); + for (int stride = width / 2; stride > 0; stride /= 2) { + for (int i = threadIdx.x; i < stride; i += blockDim.x) { + shared[i] = __fadd_rn(shared[i], shared[i + stride]); + } + __syncthreads(); + } + const float result = shared[0]; + // Every warp must read the result before callers reuse the shared buffer. + __syncthreads(); + return result; +} + +struct RowStats { + float mean; + float rstd; +}; + +template +__device__ __forceinline__ RowStats row_stats( + const scalar_t* x, float* shared, int row, int hidden, int width, float eps) { + for (int col = threadIdx.x; col < width; col += blockDim.x) { + shared[col] = col < hidden ? static_cast(x[row * hidden + col]) : 0.0f; + } + const float mean = __fdiv_rn(fixed_sum(shared, width), hidden); + for (int col = threadIdx.x; col < width; col += blockDim.x) { + const float centered = col < hidden + ? __fsub_rn(static_cast(x[row * hidden + col]), mean) : 0.0f; + shared[col] = __fmul_rn(centered, centered); + } + const float variance = __fdiv_rn(fixed_sum(shared, width), hidden); + const float rstd = __fdiv_rn(1.0f, __fsqrt_rn(__fadd_rn(variance, eps))); + return {mean, rstd}; +} + +template +__global__ void forward_kernel( + const scalar_t* x, const scalar_t* modulation, scalar_t* y, scalar_t* gate, + int seq, int hidden, int width, float eps) { + extern __shared__ float shared[]; + const int row = blockIdx.x; + const int batch = row / seq; + const auto stats = row_stats(x, shared, row, hidden, width, eps); + const float mean = stats.mean; + const float rstd = stats.rstd; + for (int col = threadIdx.x; col < hidden; col += blockDim.x) { + const float centered = __fsub_rn(static_cast(x[row * hidden + col]), mean); + const float norm = __fmul_rn(centered, rstd); + const float scale = static_cast(modulation[batch * 3 * hidden + hidden + col]); + const float shift = static_cast(modulation[batch * 3 * hidden + col]); + const float value = __fadd_rn(__fmul_rn(norm, __fadd_rn(1.0f, scale)), shift); + y[row * hidden + col] = static_cast(value); + if (row % seq == 0) gate[batch * hidden + col] = modulation[batch * 3 * hidden + 2 * hidden + col]; + } +} + +template +__global__ void backward_rows_kernel( + const scalar_t* x, const scalar_t* modulation, const scalar_t* dy, + scalar_t* dx, float* partial_shift, float* partial_scale, + int seq, int hidden, int width, float eps, + const int* index = nullptr, int batch_size = 0, bool broadcast = false) { + extern __shared__ float shared[]; + const int row = blockIdx.x; + const int batch = row / seq + (index == nullptr ? 0 : + index[broadcast ? row % seq : row] * batch_size); + const auto stats = row_stats(x, shared, row, hidden, width, eps); + const float mean = stats.mean; + const float rstd = stats.rstd; + for (int col = threadIdx.x; col < width; col += blockDim.x) { + if (col < hidden) { + const float grad = static_cast(dy[row * hidden + col]); + const float scale = static_cast(modulation[batch * 3 * hidden + hidden + col]); + shared[col] = __fmul_rn(grad, __fadd_rn(1.0f, scale)); + } else shared[col] = 0.0f; + } + const float mean_dnorm = __fdiv_rn(fixed_sum(shared, width), hidden); + for (int col = threadIdx.x; col < width; col += blockDim.x) { + if (col < hidden) { + const float centered = __fsub_rn(static_cast(x[row * hidden + col]), mean); + const float norm = __fmul_rn(centered, rstd); + const float grad = static_cast(dy[row * hidden + col]); + const float scale = static_cast(modulation[batch * 3 * hidden + hidden + col]); + shared[col] = __fmul_rn(__fmul_rn(grad, __fadd_rn(1.0f, scale)), norm); + } else shared[col] = 0.0f; + } + const float mean_dnorm_norm = __fdiv_rn(fixed_sum(shared, width), hidden); + for (int col = threadIdx.x; col < hidden; col += blockDim.x) { + const float grad = static_cast(dy[row * hidden + col]); + const float scale = static_cast(modulation[batch * 3 * hidden + hidden + col]); + const float centered = __fsub_rn(static_cast(x[row * hidden + col]), mean); + const float norm = __fmul_rn(centered, rstd); + const float dnorm = __fmul_rn(grad, __fadd_rn(1.0f, scale)); + const float value = __fmul_rn( + __fsub_rn(__fsub_rn(dnorm, mean_dnorm), __fmul_rn(norm, mean_dnorm_norm)), rstd); + dx[row * hidden + col] = static_cast(value); + partial_shift[row * hidden + col] = grad; + partial_scale[row * hidden + col] = __fmul_rn(grad, norm); + } +} + +template +__global__ void backward_reduce_kernel( + const float* partial_shift, const float* partial_scale, const scalar_t* dg, + scalar_t* dm, int seq, int hidden) { + const int col = blockIdx.x * blockDim.x + threadIdx.x; + const int batch = blockIdx.y; + if (col >= hidden) return; + float dshift = 0.0f, dscale = 0.0f; +#pragma unroll 1 + for (int token = 0; token < seq; ++token) { + const int offset = (batch * seq + token) * hidden + col; + dshift = __fadd_rn(dshift, partial_shift[offset]); + dscale = __fadd_rn(dscale, partial_scale[offset]); + } + dm[batch * 3 * hidden + col] = static_cast(dshift); + dm[batch * 3 * hidden + hidden + col] = static_cast(dscale); + dm[batch * 3 * hidden + 2 * hidden + col] = dg[batch * hidden + col]; +} + +void check(const torch::Tensor& x, const torch::Tensor& modulation, double eps, int threads) { + TORCH_CHECK(threads == 128 || threads == 256 || threads == 512, + "AdaLN threads must be 128, 256, or 512"); + TORCH_CHECK(x.is_cuda() && modulation.is_cuda(), "AdaLN inputs must be CUDA tensors"); + TORCH_CHECK(x.is_contiguous() && modulation.is_contiguous(), "AdaLN inputs must be contiguous"); + TORCH_CHECK(x.dim() == 3 && x.size(0) > 0 && x.size(1) > 0 && x.size(2) > 0, + "x must have nonempty shape [B,S,H]"); + TORCH_CHECK(x.size(0) <= 65535, "AdaLN CUDA requires B <= 65535"); + TORCH_CHECK(modulation.dim() == 2 && modulation.size(0) == x.size(0) && + modulation.size(1) == 3 * x.size(2), "modulation must have shape [B,3H]"); + TORCH_CHECK(x.numel() <= std::numeric_limits::max() && + modulation.numel() <= std::numeric_limits::max(), + "AdaLN CUDA requires tensors to fit int32 indexing"); + TORCH_CHECK(x.device() == modulation.device() && x.scalar_type() == modulation.scalar_type(), + "AdaLN inputs must share device and dtype"); + TORCH_CHECK(x.scalar_type() == at::kFloat || x.scalar_type() == at::kBFloat16, + "AdaLN inputs must be FP32 or BF16"); + TORCH_CHECK(x.size(2) <= 4096, "AdaLN CUDA supports H <= 4096"); + TORCH_CHECK(std::isfinite(eps) && eps >= 0.0, "eps must be finite and nonnegative"); + const float eps32 = static_cast(eps); + TORCH_CHECK(std::isfinite(eps32) && (eps == 0.0 || eps32 > 0.0f), + "eps must remain finite FP32 without positive underflow"); +} + +template +__global__ void indexed_forward_kernel( + const scalar_t* x, const scalar_t* m, const int* index, scalar_t* y, + scalar_t* gate, int batch, int seq, int hidden, int width, bool broadcast, float eps) { + extern __shared__ float scratch[]; + const int row = blockIdx.x; + const int bank_row = row / seq + index[broadcast ? row % seq : row] * batch; + const auto stats = row_stats(x, scratch, row, hidden, width, eps); + for (int col = threadIdx.x; col < hidden; col += blockDim.x) { + const float n = __fmul_rn(__fsub_rn(static_cast(x[row * hidden + col]), + stats.mean), stats.rstd); + const int offset = bank_row * 3 * hidden + col; + y[row * hidden + col] = static_cast(__fadd_rn( + __fmul_rn(n, __fadd_rn(1.0f, static_cast(m[offset + hidden]))), + static_cast(m[offset]))); + // Copy original dtype bits, including signaling NaN payloads. + gate[row * hidden + col] = m[offset + 2 * hidden]; + } +} + +template +__global__ void indexed_backward_reduce_kernel( + const float* ps, const float* pc, const scalar_t* dg, const int* index, + scalar_t* dm, int batch, int seq, int hidden, bool broadcast) { + const int sample = blockIdx.y; + const int col = blockIdx.x * blockDim.x + threadIdx.x; + if (col >= hidden) return; + float ds[2] = {0.0f, 0.0f}, dc[2] = {0.0f, 0.0f}, dgate[2] = {0.0f, 0.0f}; +#pragma unroll 1 + for (int token = 0; token < seq; ++token) { + const int row = sample * seq + token; + const int bank = index[broadcast ? token : row]; + const int offset = row * hidden + col; + if (ps != nullptr) { + ds[bank] = __fadd_rn(ds[bank], ps[offset]); + dc[bank] = __fadd_rn(dc[bank], pc[offset]); + } + if (dg != nullptr) dgate[bank] = __fadd_rn(dgate[bank], static_cast(dg[offset])); + } + for (int bank = 0; bank < 2; ++bank) { + const int offset = (sample + bank * batch) * 3 * hidden + col; + dm[offset] = static_cast(ds[bank]); + dm[offset + hidden] = static_cast(dc[bank]); + dm[offset + 2 * hidden] = static_cast(dgate[bank]); + } +} + +void check_indexed( + const torch::Tensor& x, const torch::Tensor& m, const torch::Tensor& index, + double eps, int threads) { + TORCH_CHECK(x.dim() == 3 && x.size(0) > 0 && x.size(1) > 0 && x.size(2) > 0, + "x must have nonempty shape [B,S,H]"); + TORCH_CHECK(m.dim() == 2 && m.size(0) == 2 * x.size(0) && m.size(1) == 3 * x.size(2), + "indexed modulation must have shape [2B,3H]"); + TORCH_CHECK(m.numel() <= std::numeric_limits::max() && + index.numel() <= std::numeric_limits::max(), + "indexed CUDA requires tensors to fit int32 indexing"); + check(x, m.narrow(0, 0, x.size(0)), eps, threads); + TORCH_CHECK(m.is_contiguous(), "indexed modulation must be contiguous"); + TORCH_CHECK(index.is_cuda() && index.device() == x.device() && + index.scalar_type() == at::kInt && index.is_contiguous() && + index.dim() == 2 && (index.size(0) == x.size(0) || index.size(0) == 1) && + index.size(1) == x.size(1), + "indexed CUDA requires contiguous int32 index [B,S] or [1,S] on x device"); + TORCH_CHECK(((index == 0) | (index == 1)).all().item(), + "modulate_index values must be 0 or 1"); +} + +} // namespace + +std::vector adaln_modulation_forward_cuda( + torch::Tensor x, torch::Tensor modulation, double eps, int threads) { + check(x, modulation, eps, threads); + c10::cuda::CUDAGuard guard(x.device()); + const int batch = x.size(0), seq = x.size(1), hidden = x.size(2); + int width = 1; + while (width < hidden) width <<= 1; + auto y = torch::empty_like(x); + auto gate = torch::empty({batch, 1, hidden}, x.options()); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND(at::kBFloat16, x.scalar_type(), "adaln_forward", [&] { + forward_kernel<<>>( + x.data_ptr(), modulation.data_ptr(), y.data_ptr(), + gate.data_ptr(), seq, hidden, width, static_cast(eps)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {y, gate}; +} + +std::vector adaln_modulation_backward_cuda( + torch::Tensor dy, torch::Tensor dg, torch::Tensor x, + torch::Tensor modulation, double eps, int threads) { + check(x, modulation, eps, threads); + TORCH_CHECK(dy.is_cuda() && dy.is_contiguous() && dy.sizes() == x.sizes() && + dy.scalar_type() == x.scalar_type() && dy.device() == x.device(), + "dy must match x"); + TORCH_CHECK(dg.is_cuda() && dg.is_contiguous() && dg.dim() == 3 && + dg.size(0) == x.size(0) && dg.size(1) == 1 && dg.size(2) == x.size(2) && + dg.scalar_type() == x.scalar_type() && dg.device() == x.device(), + "dg must match x device/dtype and have shape [B,1,H]"); + c10::cuda::CUDAGuard guard(x.device()); + const int batch = x.size(0), seq = x.size(1), hidden = x.size(2); + int width = 1; + while (width < hidden) width <<= 1; + auto dx = torch::empty_like(x); + auto dm = torch::empty_like(modulation); + auto partial_shift = torch::empty(x.sizes(), x.options().dtype(torch::kFloat32)); + auto partial_scale = torch::empty_like(partial_shift); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND(at::kBFloat16, x.scalar_type(), "adaln_backward", [&] { + backward_rows_kernel<<>>( + x.data_ptr(), modulation.data_ptr(), dy.data_ptr(), + dx.data_ptr(), partial_shift.data_ptr(), partial_scale.data_ptr(), + seq, hidden, width, static_cast(eps)); + backward_reduce_kernel<<>>( + partial_shift.data_ptr(), partial_scale.data_ptr(), + dg.data_ptr(), dm.data_ptr(), seq, hidden); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dx, dm}; +} + +std::vector adaln_indexed_forward_cuda( + torch::Tensor x, torch::Tensor m, torch::Tensor index, double eps, int threads) { + check_indexed(x, m, index, eps, threads); + c10::cuda::CUDAGuard guard(x.device()); + const int batch = x.size(0), seq = x.size(1), hidden = x.size(2); + int width = 1; + while (width < hidden) width <<= 1; + auto y = torch::empty_like(x), gate = torch::empty_like(x); + auto stream = at::cuda::getCurrentCUDAStream(); + AT_DISPATCH_FLOATING_TYPES_AND(at::kBFloat16, x.scalar_type(), "adaln_indexed_forward", [&] { + indexed_forward_kernel<<>>( + x.data_ptr(), m.data_ptr(), index.data_ptr(), + y.data_ptr(), gate.data_ptr(), batch, seq, hidden, width, + index.size(0) == 1, static_cast(eps)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {y, gate}; +} + +std::vector adaln_indexed_backward_cuda( + c10::optional dy, c10::optional dg, + torch::Tensor x, torch::Tensor m, torch::Tensor index, double eps, int threads) { + check_indexed(x, m, index, eps, threads); + for (const auto& grad : {dy, dg}) { + if (grad.has_value()) { + TORCH_CHECK(grad->device() == x.device() && grad->scalar_type() == x.scalar_type() && + grad->is_contiguous() && grad->sizes() == x.sizes(), + "indexed output gradient must match x device/dtype/shape"); + } + } + c10::cuda::CUDAGuard guard(x.device()); + const int batch = x.size(0), seq = x.size(1), hidden = x.size(2); + int width = 1; + while (width < hidden) width <<= 1; + auto dm = torch::empty_like(m); + torch::Tensor dx, ps, pc; + auto stream = at::cuda::getCurrentCUDAStream(); + if (dy.has_value()) { + dx = torch::empty_like(x); + ps = torch::empty(x.sizes(), x.options().dtype(torch::kFloat32)); + pc = torch::empty_like(ps); + } + AT_DISPATCH_FLOATING_TYPES_AND(at::kBFloat16, x.scalar_type(), "adaln_indexed_backward", [&] { + if (dy.has_value()) { + backward_rows_kernel<<>>( + x.data_ptr(), m.data_ptr(), dy->data_ptr(), + dx.data_ptr(), ps.data_ptr(), pc.data_ptr(), + seq, hidden, width, static_cast(eps), index.data_ptr(), + batch, index.size(0) == 1); + } + // B, not 2B, occupies grid.y; each thread writes both banks. + indexed_backward_reduce_kernel<<>>( + ps.defined() ? ps.data_ptr() : nullptr, + pc.defined() ? pc.data_ptr() : nullptr, + dg.has_value() ? dg->data_ptr() : nullptr, + index.data_ptr(), dm.data_ptr(), batch, seq, hidden, index.size(0) == 1); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dx, dm}; +} diff --git a/docs/.nav.yml b/docs/.nav.yml index 29b5ead37..451697f97 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -17,6 +17,7 @@ nav: - Operators: - operators/README.md - operators/activation.md + - operators/adaln-modulation.md - operators/attention.md - operators/fused-logp.md - operators/linear-logp.md diff --git a/docs/contributing/gtest-usage.md b/docs/contributing/gtest-usage.md index 2e0d9e235..7fb75d88d 100644 --- a/docs/contributing/gtest-usage.md +++ b/docs/contributing/gtest-usage.md @@ -105,7 +105,20 @@ Currently registered (source of truth is the code): ```text rms_norm, qk_norm, attention, logp, linear_logp, embedding, lm_head, -det_gemm, rope, silu, swiglu, batch_invariant_logp, pack +det_gemm, rope, silu, swiglu, batch_invariant_logp, pack, adaln_modulation +``` + +`adaln_modulation` uses the reduction contract and an independent FP32 CPU +LayerNorm/autograd gold. It compares y and gate, and (with `--check-grad`) +x and modulation gradients. Candidates are `pytorch`, `cuda`, and `triton`. +Pass `--normalized-dim 3072` for Qwen-Image; default CLI model dimensions +otherwise describe Qwen3. Indexed select01 and strict byte invariance remain +in the focused [AdaLN tests](../operators/adaln-modulation.md), not C3/C4 adapters. + +```bash +.venv/bin/python tools/validation/operators/check_operator.py --op adaln_modulation \ + --candidate pytorch --device cpu --dtype bf16 --batch 2 --seq 3 \ + --normalized-dim 3072 --check-grad --json ``` `qk_norm` is a first-class `OP_SPECS` key that reuses the RMSNorm kernels on diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md new file mode 100644 index 000000000..e39af1f04 --- /dev/null +++ b/docs/operators/adaln-modulation.md @@ -0,0 +1,284 @@ +# AdaLN Modulation + +## Summary + +AdaLN modulates Qwen-Image attention or MLP inputs with LayerNorm, shift and scale. +It returns a separate gate for downstream use; the gate is not multiplied into the +modulated output. + +Image/text streams and attention/MLP calls use separate inputs. The conditioning +projection (`SiLU + Linear`) and gated residual update are separate operators. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +# x: [B, S, H]; modulation: [B, 3H], ordered as shift, scale, gate. +op, trace = kernel_registry.get_adaln_modulation_op(x.device, hidden=x.shape[-1]) +modulated, gate = op(x, modulation, eps=1e-6) +``` + +Each token is normalized over H without learned LayerNorm affine parameters: + +```text +modulated = LayerNorm(x, eps) * (1 + scale) + shift +``` + +## Tensor Contract + +| Value | Shape | Dtype / requirements | +| --- | --- | --- | +| `x` | `[B,S,H]` | FP32 or BF16; B, S and H must be positive | +| Shared `modulation` | `[B,3H]` | Same device and dtype as x | +| `modulated` | `[B,S,H]` | Same device and dtype as x | +| Shared `gate` | `[B,1,H]` | Same device and dtype as x; retains gradients to modulation | +| `eps` | scalar | Defaults to `1e-6`; see Numerical Contract | + +Shared shift, scale and gate apply to every token in the same sample. GPU wrappers +accept strided inputs and copy them to contiguous storage before kernel execution. + +### Indexed modulation + +For image editing, pass `modulate_index` to both resolution and execution: + +```python +op, trace = kernel_registry.get_adaln_modulation_op( + x.device, hidden=x.shape[-1], modulate_index=modulate_index +) +modulated, gate = op(x, modulation, modulate_index=modulate_index) +``` + +| Value | Indexed contract | +| --- | --- | +| `modulation` | `[2B,3H]`; same device and dtype as x | +| `modulate_index` | Boolean or integer `[B,S]` or `[1,S]`; values 0 or 1; same device as x | +| `modulated` | `[B,S,H]`; same device and dtype as x | +| `gate` | `[B,S,H]`; same device and dtype as x | + +Index 0 selects the sample's modulation from the first B rows; index 1 selects it +from the second B rows. A `[1,S]` index broadcasts across the batch. + +## Backends and Dispatch + +| Backend | Implementation | Use | +| --- | --- | --- | +| CUDA | `CudaAdaLNModulationOp` | NVIDIA CUDA extension; H <= 4096 | +| Triton | `TritonAdaLNModulationOp` | NVIDIA CUDA validation target; ROCm is not qualified | +| PyTorch | `NativeAdaLNModulationOp` | CPU fixed-order baseline and framework fallback | + +On NVIDIA CUDA, shared modulation tries CUDA, then Triton, then PyTorch. +H > 4096 skips CUDA. Missing extension symbols also permit fallback; fast-math +builds omit the AdaLN CUDA symbols. + +CPU, ROCm, MUSA and NPU registry entries select PyTorch. Dedicated indexed CUDA +and NVIDIA Triton accept FP32/BF16, H <= 4096, B <= 65535, both route shapes and +bool/uint8/int8/int16/int32/int64 indices. Inputs and upstreams may be strided or +expanded; internal copies preserve gradients to source leaves. Indexed route storage +is converted to contiguous int32 after validation. + +Indexed CUDA/Triton check element counts and accessed element offsets before copies or +kernel output/temporary allocations. The checks include both modulation banks, +route storage, y/gate/dx/dmod and each FP32 partial buffer. The B grid dimension +contains B entries; each reduction thread writes both banks. + +Valid unsupported geometry, addressing or eps executes the existing framework +select01 fallback in ordinary dispatch. A resolved indexed CUDA implementation +records `indexed_cuda_domain_unsupported`; resolution that selected framework or a +shared implementation records `modulate_index_requires_select01`. Indexed Triton +records `indexed_triton_domain_unsupported` for the corresponding indexed +resolution. Illegal inputs +raise and clear the previous successful execution observation. Direct CUDA/Triton calls +reject unsupported domains instead of silently falling back. + +Historical ticket03 AdaLN-only SM86 evidence recorded an indexed matrix of +361 PASS (including CPU reference/metadata checks), small +memcheck/racecheck 107 PASS each, and related regression 417 PASS / 6 opt-in +workload SKIP. Two oversized-upstream tests require >5 GiB free GPU memory and +remain unexecuted on the 4 GiB target. This evidence applies only to its recorded +snapshot; affected current-source GPU checks need refresh. The dedicated indexed +Triton implementation has a separate local SM86 run on RTX 3050 Ti with +Torch 2.12.1+cu126, Triton 3.7.1 and bundled ptxas 12.8.93: indexed matrix +297 PASS, small sanitizer matrix 129 PASS per tool (memcheck/racecheck/initcheck/ +synccheck, no reported errors or hazards), and shared regression 13 PASS. +FP32/BF16 subnormal checks and retained generated-code inspection are included; +this is local evidence, not CUDA/shared evidence inheritance or general target +qualification. Two large-upstream cases and twelve full-image workload nodes +remain unexecuted because the 4 GiB target cannot meet their memory requirements. +Full-extension and H100 qualification remain separate. Indexed strict stays closed. + +### Strict mode and trace + +Pass `strict=True` to enforce the applicable arithmetic policy. CPU Native and +NVIDIA CUDA/Triton retain the shared profile. Indexed strict additionally requires +an explicit `IndexedAdaLNQualification` from `rl_engine.runtime.adaln`, passed as +`indexed_qualification=scope` to `get_adaln_modulation_op`. Its default is `None`: +**indexed strict delivery qualification remains closed**. Implemented policy and +CPU policy tests do not establish hardware qualification. + +The immutable scope names one backend, an explicit CUDA ordinal, GPU model and +compute capability, exact Torch/CUDA runtime versions, and the Triton version for +Triton. It expresses activation after external evidence review; it does not +verify a certificate or infer evidence. CUDA qualification cannot activate Triton. +Resolution checks backend/runtime/hidden eligibility without launching kernels; +execution rechecks actual device, GPU identity, H and the full indexed domain. +Inputs determine shared/indexed/framework-select01 profiles on every call; there +is no caller profile parameter. An indexed resolution may execute shared inputs +and vice versa, but strict never permits framework select01 or unsupported input. +Illegal input remains an input error and every failure clears earlier observations. + +Acceptance uses this same explicit scope and public resolver/callable path, in a +separate source package for real Triton selection. After ticket06 accepts evidence, +activation uses the reviewed target values without modifying arithmetic, backend +selection or launch code. Source/binary/compiler identity, profiler and sanitizer +evidence stay in external manifests. The scope does not attest those properties; +a changed core path requires new affected evidence. No scope is installed globally +or enabled by an environment variable. + +The returned callable is an `AdaLNDispatch` wrapper with an `implementation` +attribute. `selected_backend` records resolution; `actual_backend`, dtype and +device record the latest successful forward. A failed call clears that observation. + +Copy `trace` to retain an observation and use separate resolutions for concurrent +callers. A successful-forward trace is not a profiler measurement or proof of +backward execution. Semantic fingerprints are not compiled-binary hashes. + +## Numerical Contract + +Shared CUDA/Triton arithmetic and the CPU Native baseline use an FP32 hidden +reduction tree, pairing lower and upper halves after padding to the next power +of two. Shared parameter gradients accumulate in ascending logical token order. + +Intermediates and accumulators remain FP32, with one final cast for BF16 results. +CUDA/Triton kernels use no atomic partial sums, Split-K, Stream-K, TF32 or fast math. + +For shared modulation, eps must be finite, nonnegative and round to finite +FP32. Positive values that round to zero are rejected; FP32 subnormals, including +2^-149, are accepted. + +Indexed framework fallback checks only finite, nonnegative eps. Dedicated indexed +CUDA/Triton use the shared FP32 eps domain. Their per-sample/bank/channel dshift, +dscale and dgate start at FP32 +0 and fold selected tokens in ascending order; +each dscale product rounds to FP32 before addition. Gate copies original bits. +This is the first `adaln-indexed-fp32-tree-v1` arithmetic choice, independently +subject to acceptance, not a permanent restriction on future indexed profiles. + +With eps=0, zero variance produces NaN. Nonfinite normalization arithmetic may +propagate NaNs; those NaN payloads are outside byte-equality guarantees. +The shared path preserves gate/dgate payload bits, including nonfinite values. + +### Fixed-order implementation and independent gold + +`NativeAdaLNModulationOp` uses the fixed reduction order for fallback execution and +byte comparisons. Its square root is evaluated in FP64 and rounded to FP32 to match +correctly rounded CUDA sqrt; its reductions remain FP32. + +`GtestAdaLNReference` is the independent accuracy gold. It uses CPU FP32 +`F.layer_norm` and autograd for normalization, plus an independent affine backward +that follows the specified token order. It calls no production reduction helpers. + +The validation harness creates FP32 reference input leaves to retain FP32 gradients. +It checks both outputs and gradients of x and modulation using the shared +`reduction` tolerances for `forward_accuracy` and `gradient_accuracy`. + +Accuracy checks and fixed-reference byte comparisons are separate tests. +`IndexedAdaLNFixedReference` is a separate CPU FP32/BF16 callable for indexed byte +tests, with its own H tree and explicit VJP. It is not the framework fallback. +Indexed reference literals check sample/bank routing, absent versus explicit zero +dy, frozen leaves, cancellation order, separate product/add rounding, empty-bank ++0 isolation with NaN/Inf, gate payload copies and subnormal contributions. +The accuracy gold uses an independent selected-token affine VJP while retaining +CPU LayerNorm/autograd normalization; it is not a fixed-order byte oracle. +The same literals exercise the public CUDA callable; BF16 tests add final-output +rounding, ties-to-even and raw16 gate payloads. +All0 tests compare indexed bank0 dshift/dscale with shared bytes, check bank1 +0 +separately and retain independent gold comparisons; shared/indexed dgate are not +equated. CUDA and public-dispatch tests require source-specific GPU evidence. +The local Triton run includes the new BF16 minimum-subnormal parameter-gradient +cases for one/three tokens and gate-only/both upstreams. Ticket 03 still has two +unexecuted oversized CUDA upstream cases; indexed strict remains closed. +The CLI records trace snapshots without automatically constructing C1 +`BackendProvenance`; its accuracy report alone is not a WS1 system gate. + +## Tests and Benchmark + +```bash +python -m pytest tests/ops/norm/test_adaln_modulation.py \ + tests/reference/test_adaln_indexed.py tests/ops/norm/test_adaln_indexed.py \ + tests/runtime/test_adaln_indexed.py \ + tests/validation/operators/test_adaln_gtest.py tests/runtime/test_adaln_dispatch.py -q -rs + +python tools/validation/operators/check_operator.py --op adaln_modulation \ + --candidate pytorch --device cpu --dtype bf16 --batch 2 --seq 3 \ + --normalized-dim 3072 --check-grad --json + +python benchmarks/operators/norm/benchmark_adaln_modulation.py --real --dtype bf16 --json +``` + +Tests cover forward/gradient accuracy, fixed-order byte equality, dtype boundaries, +strided inputs/upstreams, gate-only gradients, nonfinite values and dispatch. +Set `RLK_ADALN_REAL_SHAPES=1` to include the three full-image workloads. + +The benchmark reports median forward and forward+backward latency and peak allocated +memory. Use `--dtype fp32` for FP32, `--warmup`/`--repeat` for sample counts and +`--seed` for reproducible inputs (default 386). + +Forward+backward timing includes upstream `ones_like` allocations. +The ordinary `F.layer_norm` plus FP32 affine baseline uses the same timing method, +without a fixed-order replay claim. + +See the [H100 validation runbook](../validation/adaln-h100.md) for workload shapes, +fresh extension builds, sanitizer commands and evidence records. + +## Known Limitations + +Shared CUDA and Triton require at most 2^31-1 elements in each input (`x` and +`modulation`) for int32 element addressing. Python rejects larger inputs before +contiguous copies or output/partial allocation. The CUDA extension additionally +requires B <= 65535 and repeats its bounds checks in C++ before launch. + +Direct backend calls reject unsupported domains. Ordinary dispatch rechecks +actual inputs and uses a real Native fallback when shared device/geometry is +unsupported, preserving the shared validation and arithmetic contract. Its trace +records `shared_cuda_domain_unsupported` or `shared_triton_domain_unsupported`. +Strict rejects this fallback. Resolution uses the metadata available at that point. + +Registry dispatch uses default launch settings. Direct constructors accept +CUDA `threads={128,256,512}` (default 256), Triton `num_warps={4,8}` (default 4) +and `reduction_tile={64,128,256}` (default 128). + +These settings alter scheduling, not the declared reduction order. GPU support +and numerical acceptance must be checked on the target hardware; CPU tests +alone do not establish CUDA/Triton acceptance. + +Indexed NVIDIA dispatch tries CUDA symbols, then Triton, then framework. To test +actual Triton selection while CUDA remains the default, use a separate source +package without `_C.so`, disable editable import hooks, and assert the extension +is absent before resolving. Keep the current binary in place. Successful indexed +Triton execution reports `actual_backend=triton`, indexed arithmetic profile, +`fallback=false` and `fallback_reason=None`; failed calls clear old observations. + +Indexed Triton uses a separate row kernel and selected-bank FP32 fold. It reuses +the shared H statistics helpers without changing shared arithmetic or validation. +Both kernels set `enable_fp_fusion=False` and `enable_reflect_ftz=False`. Gate +copies use original-width integer pointers. Generated code and target subnormal +behavior require separate checks; flags alone are not evidence. Local SM86 +inspection found RN sqrt/div lowering uses internal FTZ/FMA with correction paths, +and single-token BF16 gate addition may become FMA(v,1,+0). The latter is an +equivalent final-result specialization, not multi-token low-precision accumulation. +Do not claim the entire SASS is free of FTZ/FMA. No compiler/runtime +beyond the recorded acceptance environment is qualified. Unsupported geometry, +eps, addressing and non-NVIDIA devices reject directly or explicitly fall back +through ordinary dispatch. Indexed strict remains default-closed; its explicit +activation path is implemented for acceptance and later reviewed delivery. + +Indexed B/S/H/eps, broadcast mode, forward/backward, upstream presence and launch +parameters specialize kernels; S is `tl.constexpr`, so new sequence lengths need +separate variants and identical calls may reuse cache entries. No runtime-S +optimization or launch-annotation performance gain is claimed. + +Triton currently specializes all three shared kernels on S. A new S can require +a new compiled variant; repeated calls can reuse the cached variant. Removing +`tl.constexpr` can reduce variants, but runtime integer specialization can still +distinguish values such as 1 and alignment classes. Runtime division/remainder +and loop-bound costs need target-GPU measurements before changing this choice. diff --git a/docs/validation/adaln-h100.md b/docs/validation/adaln-h100.md new file mode 100644 index 000000000..caae28bd0 --- /dev/null +++ b/docs/validation/adaln-h100.md @@ -0,0 +1,187 @@ +# AdaLN H100 Validation + +This runbook validates [AdaLN modulation](../operators/adaln-modulation.md) on +H100/SM90. Run commands from the repository root in an environment with matching +PyTorch/CUDA, a CUDA toolkit and NVIDIA Compute Sanitizer. + +The commands below are validation instructions; the recorded run is summarized +at the end. Retain the source revision, +working-tree status, environment, build log, test output, sanitizer logs and benchmark +JSON together. A skipped backend is not native CUDA acceptance. + +## Workloads + +| Image resolution | Image sequence length S | Hidden size H | +| --- | --- | --- | +| 1024 x 1024 | 4096 | 3072 | +| 1328 x 1328 | 6889 | 3072 | +| 1664 x 928 | 6032 | 3072 | + +Sequence lengths are derived using VAE scale 8 and 2x2 latent packing. +They are model workloads, not tensor shapes explicitly specified by issue #386. +Validate both FP32 and BF16 and both CUDA and Triton. + +## Record Environment and Build + +```bash +git rev-parse HEAD +git status --porcelain=v1 --untracked-files=all +nvidia-smi +.venv/bin/python -c 'import torch; print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(), torch.cuda.get_device_capability())' + +RL_KERNEL_REQUIRE_EXT=1 KERNEL_ALIGN_USE_FAST_MATH=0 TORCH_CUDA_ARCH_LIST=9.0 \ + .venv/bin/python setup.py build_ext --inplace +``` + +Build the current sources and start new test processes so they load the new +extension. Save the build output and verify the loaded extension identity; +passing tests against an older loaded binary do not validate this build. + +## Forward, Backward and Dispatch + +```bash +RL_KERNEL_REQUIRE_EXT=1 RLK_ADALN_REAL_SHAPES=1 \ + .venv/bin/python -m pytest tests/ops/norm/test_adaln_modulation.py \ + tests/validation/operators/test_adaln_gtest.py \ + tests/backends/test_adaln_shared_bounds.py tests/runtime/test_adaln_dispatch.py -q -rs + +for candidate in cuda triton; do + for dtype in bf16 fp32; do + .venv/bin/python tools/validation/operators/check_operator.py --op adaln_modulation \ + --candidate "$candidate" --device cuda --dtype "$dtype" --batch 2 --seq 3 \ + --normalized-dim 3072 --check-grad --json + done +done +``` + +The opt-in full-image tests compare both GPU backends/dtypes against independent +CPU FP32 forward and backward results. Inspect pass/skip output for each backend. +The CLI separately checks accuracy; it does not establish a WS1 system gate. + +To isolate the GPU dispatch cases: + +```bash +.venv/bin/python -m pytest tests/runtime/test_adaln_dispatch.py -k gpu -q -rs +``` + +## Memory and Race Checks + +```bash +RL_KERNEL_REQUIRE_EXT=1 compute-sanitizer --error-exitcode=1 --tool memcheck \ + .venv/bin/python -m pytest tests/ops/norm/test_adaln_modulation.py \ + -k test_cuda_adaln_bit_equality_harness -q + +RL_KERNEL_REQUIRE_EXT=1 compute-sanitizer --error-exitcode=1 --tool racecheck \ + .venv/bin/python -m pytest tests/ops/norm/test_adaln_modulation.py \ + -k test_cuda_adaln_bit_equality_harness -q +``` + +## Performance and Identity Records + +```bash +for dtype in bf16 fp32; do + .venv/bin/python benchmarks/operators/norm/benchmark_adaln_modulation.py \ + --real --dtype "$dtype" --seed 386 --warmup 10 --repeat 50 --json +done +``` + +JSON records raw event samples, medians, absolute peak allocated memory, +allocation before sampling, GPU/driver/runtime versions, working-tree status, +source hashes and the loaded CUDA extension file hash when available. + +The binary hash identifies the loaded file; it does not prove that file was built +from the recorded source tree. Missing observations are null, and unavailable +CUDA/Triton backends are explicitly skipped. + +Backend observation precedes timing. CUDA events measure the bare operator; +forward+backward includes newly allocated upstream `ones_like` tensors. +The ordinary LayerNorm/affine baseline has no fixed-order replay guarantee. + +## Repository Integration + +After focused acceptance, run the repository suite in an appropriately provisioned +environment. It includes hardware and integration tests beyond AdaLN: + +```bash +.venv/bin/python -m pytest tests -q -rs +``` + +Keep failures or collection errors in the validation record; a focused AdaLN pass +does not imply the repository suite passed. + +## Recorded H100 Acceptance (2026-10-22) + +The locally retained `adaln-h100-evidence-20261021-01.tar.gz` records source +`69bd47f72a24d1177580d9b11b6efdfeb6c6bab5` against comparison base +`16b81175874ecf11d5ff05c00de3b69e9457591b` on one H100 PCIe (SM90), with +Python 3.11.17, PyTorch 2.12.1+cu126, Triton 3.7.1, CUDA toolkit 12.6.85 and +driver 595.91.07. These results are focused operator acceptance, not full WS1, +model, trajectory, repository-suite or cross-platform qualification. + +| Stage | Recorded result | +| --- | --- | +| Correctness | 413 passed; no skips or xfails | +| Racecheck subset | 211 passed; 0 hazards, errors or warnings | +| Memcheck subset | 225 passed; 0 errors | +| Benchmark | 54 points completed; 36 samples per point | + +The run retained a forced extension build, source-file hashes, SM90 binary +targets and per-process loaded-binary hash checks. The archived extension SHA256 +matches `identity.json` and `status.json`: + +```text +deee4d78b05f8ef7fff5f288141ac831b02fe24407f97fccdde474953117743a +``` + +The runtime source profile equals the profile retained by image build +[39741435857](https://github.com/hongleng/RL-Kernel/actions/runs/39741435857). +All repository tracked file hashes match that source revision except +`requirements.txt`, deliberately replaced by the image's pinned dependency list; +its hash matches the build artifact. Image infrastructure remains separate from +this operator PR. + +Evidence erratum: the outer `image-reference.txt` and `image-build-url.txt` +incorrectly name the older build 39565761133 (digest prefix `310f692`). The newer +build 39741435857 publishes digest prefix `682fae9`. Profile agreement supports +the source/environment association but does not independently prove which image +digest the host deployed. Preserve the original archive and this correction. +The inner archive checksum was verified; the outer checksum sidecar was not +retrieved and is not claimed as verified. + +The 54-point measurement uses the separate image acceptance harness: FP32/BF16, +the three real sequence lengths, CUDA/Triton/PyTorch, and forward/backward-only/ +forward+backward with preallocated random upstream gradients. It is distinct from +the repository benchmark command above. Sample coefficients of variation range +from about 0.8% to 50.9%; these timings do not establish stable speedup ratios. + +Subsequent documentation-only edits do not change the accepted executable source. +Any arithmetic, dispatch, build or test change needs an explicit reassessment and +affected H100 checks rerun; these results must not be relabelled as a fresh pass. + +## PR #559 Review Follow-up (2026-10-22) + +Commit `b2eabc9` adds a Python preflight rejection for Triton inputs exceeding +2^31-1 elements, before contiguous copies and output/partial allocation. Commit +`c394895` gates the large-hidden Triton resolution assertion on its dependency. +The Triton kernel bodies, S specialization, launch settings, CUDA sources and +numerical gold are unchanged. The historical H100 results remain evidence for +that unchanged kernel arithmetic and its tested shapes; they do not test the new +rejection behavior or constitute a fresh pass of the updated source profile. + +The new rejection and boundary checks use meta tensors, without large input +allocations. An import-blocked process confirms Native fallback without Triton +and that the Triton-specific test skips. These are host checks, not GPU passes. + +A small SM86 experiment with Triton 3.7.1 compared constexpr S with a runtime-S +candidate at B=2, H=128, S=16/17/33 for FP32/BF16. Each original kernel cached +three variants per dtype; each candidate cached two (aligned and unaligned S). +Repeating S=16 reused both caches. Both versions matched independent FP32 +accuracy gold and fixed-reference forward/backward bytes for these small cases. +Whole-call timings were noisy and include host launch/allocation overhead; they +do not isolate division/remainder/loop costs or establish an H100 speedup. + +S remains constexpr pending target H100 evidence. Any future S/kernel change +requires fresh forward/backward accuracy, byte and launch invariance, relevant +memcheck/racecheck subsets, and performance measurements at the real S workloads +above, with source/cache identities and raw timing samples retained. No fresh +H100 run was performed for this follow-up. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index fb3c4379c..8c1953fca 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -256,6 +256,33 @@ def swiglu_backward( gate: torch.Tensor, up: torch.Tensor, ) -> list[torch.Tensor]: ... +def adaln_modulation_forward( + x: torch.Tensor, modulation: torch.Tensor, eps: float, threads: int = 256 +) -> list[torch.Tensor]: ... +def adaln_modulation_backward( + dy: torch.Tensor, + dg: torch.Tensor, + x: torch.Tensor, + modulation: torch.Tensor, + eps: float, + threads: int = 256, +) -> list[torch.Tensor]: ... +def adaln_indexed_forward( + x: torch.Tensor, + modulation: torch.Tensor, + index: torch.Tensor, + eps: float, + threads: int = 256, +) -> list[torch.Tensor]: ... +def adaln_indexed_backward( + dy: torch.Tensor | None, + dg: torch.Tensor | None, + x: torch.Tensor, + modulation: torch.Tensor, + index: torch.Tensor, + eps: float, + threads: int = 256, +) -> list[torch.Tensor | None]: ... def rmsnorm_forward( x: torch.Tensor, weight: torch.Tensor, diff --git a/rl_engine/backends/cuda/norm/adaln_modulation.py b/rl_engine/backends/cuda/norm/adaln_modulation.py new file mode 100644 index 000000000..82fa6b312 --- /dev/null +++ b/rl_engine/backends/cuda/norm/adaln_modulation.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Autograd wrapper for the native CUDA AdaLN modulation kernels.""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.reference.norm.adaln_modulation import ( + _indexed_domain_reason, + _int32_addressable, + _validate, + _validate_select01, + _validate_select01_metadata, +) + + +class _AdaLNIndexedCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, modulation, index, eps, threads): + ctx.set_materialize_grads(False) + y, gate = _C.adaln_indexed_forward(x, modulation, index, eps, threads) + ctx.save_for_backward(x, modulation, index) + ctx.eps, ctx.threads = eps, threads + return y, gate + + @staticmethod + def backward(ctx, grad_y, grad_gate): + x, modulation, index = ctx.saved_tensors + for grad in (grad_y, grad_gate): + if grad is not None and not _int32_addressable(grad): + raise ValueError("indexed CUDA upstream requires int32 indexing") + dx, dm = _C.adaln_indexed_backward( + None if grad_y is None else grad_y.contiguous(), + None if grad_gate is None else grad_gate.contiguous(), + x, + modulation, + index, + ctx.eps, + ctx.threads, + ) + return None if grad_y is None else dx, dm, None, None, None + + +class _AdaLNCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, modulation, eps, threads): + y, gate = _C.adaln_modulation_forward(x, modulation, eps, threads) + ctx.save_for_backward(x, modulation) + ctx.eps = eps + ctx.threads = threads + return y, gate + + @staticmethod + def backward(ctx, grad_y, grad_gate): + x, modulation = ctx.saved_tensors + dx, dm = _C.adaln_modulation_backward( + grad_y.contiguous(), grad_gate.contiguous(), x, modulation, ctx.eps, ctx.threads + ) + return dx, dm, None, None + + +class CudaAdaLNModulationOp: + def __init__(self, *, threads=256): + if threads not in (128, 256, 512): + raise ValueError("AdaLN threads must be 128, 256, or 512") + self.threads = threads + if not _EXT_AVAILABLE or not all( + hasattr(_C, name) for name in ("adaln_modulation_forward", "adaln_modulation_backward") + ): + raise RuntimeError("CUDA AdaLN extension symbols are unavailable") + self.indexed_available = all( + hasattr(_C, name) for name in ("adaln_indexed_forward", "adaln_indexed_backward") + ) + + def __call__(self, x, modulation, *, eps=1e-6, modulate_index=None): + return self.forward(x, modulation, eps=eps, modulate_index=modulate_index) + + def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): + if modulate_index is not None: + _validate_select01_metadata(x, modulation, modulate_index, eps) + reason = _indexed_domain_reason(x, modulation, modulate_index, eps) + if reason in {"int32 indexing", "B <= 65535"}: + raise ValueError(f"indexed CUDA requires {reason}") + if reason is not None: + raise RuntimeError(f"indexed CUDA requires {reason}") + _validate_select01(x, modulation, modulate_index, eps) + if not self.indexed_available: + raise RuntimeError("indexed CUDA extension symbols are unavailable") + return _AdaLNIndexedCuda.apply( + x.contiguous(), + modulation.contiguous(), + modulate_index.to(dtype=torch.int32).contiguous(), + float(eps), + self.threads, + ) + _validate(x, modulation, eps) + if x.shape[0] > 65535: + raise ValueError("CUDA AdaLN requires B <= 65535") + if max(x.numel(), modulation.numel()) > 2**31 - 1: + raise ValueError("CUDA AdaLN requires tensors to fit int32 indexing") + if not x.is_cuda or torch.version.hip is not None or x.shape[-1] > 4096: + raise RuntimeError("CUDA AdaLN requires NVIDIA CUDA tensors with H <= 4096") + return _AdaLNCuda.apply(x.contiguous(), modulation.contiguous(), float(eps), self.threads) diff --git a/rl_engine/backends/shared/triton/norm/adaln_indexed.py b/rl_engine/backends/shared/triton/norm/adaln_indexed.py new file mode 100644 index 000000000..5d54c4a09 --- /dev/null +++ b/rl_engine/backends/shared/triton/norm/adaln_indexed.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Indexed NVIDIA AdaLN: fixed H tree, FP32 selected-token left folds.""" + +import torch +import triton +import triton.language as tl + +from rl_engine.reference.norm.adaln_modulation import ( + _indexed_domain_reason, + _int32_addressable, + _validate_select01, + _validate_select01_metadata, +) + +from .adaln_modulation import _center_and_rstd, _sum_hidden + + +@triton.jit +def _rows( + X, + M, + INDEX, + Y, + G, + DY, + DX, + PS, + PC, + B: tl.constexpr, + S: tl.constexpr, + H: tl.constexpr, + EPS: tl.constexpr, + BROADCAST: tl.constexpr, + BACKWARD: tl.constexpr, + BLOCK: tl.constexpr, + LOG: tl.constexpr, +): + row = tl.program_id(0) + bank = tl.load(INDEX + (row % S if BROADCAST else row)) + bank_row = row // S + bank * B + col = tl.arange(0, BLOCK) + x = tl.load(X + row * H + col, col < H, other=0).to(tl.float32) + centered, rstd = _center_and_rstd(x, H, EPS, BLOCK, LOG) + norm = centered * rstd + offset = bank_row * 3 * H + col + scale = tl.load(M + offset + H, col < H, other=0).to(tl.float32) + if BACKWARD: + dy = tl.load(DY + row * H + col, col < H, other=0).to(tl.float32) + dnorm = dy * (1.0 + scale) + mean = tl.div_rn(_sum_hidden(tl.where(col < H, dnorm, 0.0), BLOCK, LOG), H) + mean_norm = tl.div_rn(_sum_hidden(tl.where(col < H, dnorm * norm, 0.0), BLOCK, LOG), H) + dx = (dnorm - mean - norm * mean_norm) * rstd + tl.store(DX + row * H + col, dx, col < H) + tl.store(PS + row * H + col, dy, col < H) + # The FP32 store fixes the product's rounding before the token fold. + tl.store(PC + row * H + col, dy * norm, col < H) + else: + shift = tl.load(M + offset, col < H, other=0).to(tl.float32) + tl.store(Y + row * H + col, norm * (1.0 + scale) + shift, col < H) + # Integer loads/stores preserve signaling NaNs without a floating conversion. + if M.dtype.element_ty == tl.bfloat16: + raw = tl.load((M + offset + 2 * H).to(tl.pointer_type(tl.uint16)), col < H, other=0) + tl.store((G + row * H + col).to(tl.pointer_type(tl.uint16)), raw, col < H) + else: + raw = tl.load((M + offset + 2 * H).to(tl.pointer_type(tl.uint32)), col < H, other=0) + tl.store((G + row * H + col).to(tl.pointer_type(tl.uint32)), raw, col < H) + + +@triton.jit +def _fold( + PS, + PC, + DG, + INDEX, + DM, + B: tl.constexpr, + S: tl.constexpr, + H: tl.constexpr, + BROADCAST: tl.constexpr, + HAS_Y: tl.constexpr, + HAS_GATE: tl.constexpr, + TILE: tl.constexpr, +): + sample = tl.program_id(0) + col = tl.program_id(1) * TILE + tl.arange(0, TILE) + # Both banks use independent +0 FP32 accumulators; no masked zero additions. + for bank in tl.static_range(2): + shift = tl.full((TILE,), 0.0, tl.float32) + scale = tl.full((TILE,), 0.0, tl.float32) + gate = tl.full((TILE,), 0.0, tl.float32) + for token in range(S): + row = sample * S + token + selected = tl.load(INDEX + (token if BROADCAST else row)) + if selected == bank: + offset = row * H + col + if HAS_Y: + shift = shift + tl.load(PS + offset, col < H, other=0) + scale = scale + tl.load(PC + offset, col < H, other=0) + if HAS_GATE: + gate = gate + tl.load(DG + offset, col < H, other=0).to(tl.float32) + offset = (sample + bank * B) * 3 * H + col + tl.store(DM + offset, shift, col < H) + tl.store(DM + offset + H, scale, col < H) + tl.store(DM + offset + 2 * H, gate, col < H) + + +class _IndexedTriton(torch.autograd.Function): + @staticmethod + def forward(ctx, x, m, index, eps, num_warps, tile): + ctx.set_materialize_grads(False) + b, s, h = x.shape + y, gate = torch.empty_like(x), torch.empty_like(x) + block = triton.next_power_of_2(h) + with torch.cuda.device(x.device): + _rows[(b * s,)]( + x, + m, + index, + y, + gate, + None, + None, + None, + None, + b, + s, + h, + eps, + index.shape[0] == 1, + False, + block, + block.bit_length() - 1, + num_warps=num_warps, + enable_fp_fusion=False, + enable_reflect_ftz=False, + ) + ctx.save_for_backward(x, m, index) + ctx.eps, ctx.num_warps, ctx.tile = eps, num_warps, tile + return y, gate + + @staticmethod + def backward(ctx, dy, dg): + x, m, index = ctx.saved_tensors + # Reject upstream spans before any copies or temporary allocations. + for grad in (dy, dg): + if grad is not None and not _int32_addressable(grad): + raise ValueError("indexed Triton upstream requires int32 indexing") + dy = None if dy is None else dy.contiguous() + dg = None if dg is None else dg.contiguous() + b, s, h = x.shape + dx = ps = pc = None + if dy is not None: + dx = torch.empty_like(x) + ps = torch.empty(x.shape, dtype=torch.float32, device=x.device) + pc = torch.empty_like(ps) + block = triton.next_power_of_2(h) + _rows[(b * s,)]( + x, + m, + index, + None, + None, + dy, + dx, + ps, + pc, + b, + s, + h, + ctx.eps, + index.shape[0] == 1, + True, + block, + block.bit_length() - 1, + num_warps=ctx.num_warps, + enable_fp_fusion=False, + enable_reflect_ftz=False, + ) + dm = torch.empty_like(m) + _fold[(b, triton.cdiv(h, ctx.tile))]( + ps, + pc, + dg, + index, + dm, + b, + s, + h, + index.shape[0] == 1, + dy is not None, + dg is not None, + ctx.tile, + num_warps=ctx.num_warps, + enable_fp_fusion=False, + enable_reflect_ftz=False, + ) + return dx, dm, None, None, None, None + + +def indexed_forward(x, modulation, index, eps, num_warps, tile): + _validate_select01_metadata(x, modulation, index, eps) + reason = _indexed_domain_reason(x, modulation, index, eps) + if reason in {"int32 indexing", "B <= 65535"}: + raise ValueError(f"indexed Triton requires {reason}") + if reason is not None: + raise RuntimeError(f"indexed Triton requires {reason}") + _validate_select01(x, modulation, index, eps) + return _IndexedTriton.apply( + x.contiguous(), + modulation.contiguous(), + index.to(torch.int32).contiguous(), + float(eps), + num_warps, + tile, + ) diff --git a/rl_engine/backends/shared/triton/norm/adaln_modulation.py b/rl_engine/backends/shared/triton/norm/adaln_modulation.py new file mode 100644 index 000000000..089a70fb1 --- /dev/null +++ b/rl_engine/backends/shared/triton/norm/adaln_modulation.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""One row per program AdaLN; shared gradients fold tokens in ascending order.""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl +from triton.language.extra.cuda import libdevice + +from rl_engine.reference.norm.adaln_modulation import _validate + + +@triton.jit +def _sum_hidden(values, BLOCK: tl.constexpr, LOG: tl.constexpr): + col = tl.arange(0, BLOCK) + for level in tl.static_range(LOG - 1, -1, -1): + half = 1 << level + upper = tl.gather(values, tl.minimum(col + half, BLOCK - 1), 0) + values = tl.where(col < half, values + upper, values) + # Extract the tree root without adding +0, which would erase a -0 sum. + return tl.sum(tl.gather(values, tl.full((1,), 0, tl.int32), 0), 0) + + +@triton.jit +def _center_and_rstd(x, H: tl.constexpr, EPS: tl.constexpr, BLOCK: tl.constexpr, LOG: tl.constexpr): + col = tl.arange(0, BLOCK) + mean = tl.div_rn(_sum_hidden(x, BLOCK, LOG), H) + centered = tl.where(col < H, x - mean, 0.0) + variance = tl.div_rn(_sum_hidden(centered * centered, BLOCK, LOG), H) + rstd = tl.div_rn(1.0, libdevice.sqrt_rn(variance + EPS)) + return centered, rstd + + +@triton.jit +def _fwd( + X, + M, + Y, + G, + S: tl.constexpr, + H: tl.constexpr, + EPS: tl.constexpr, + BLOCK: tl.constexpr, + LOG: tl.constexpr, +): + row = tl.program_id(0) + batch = row // S + col = tl.arange(0, BLOCK) + x = tl.load(X + row * H + col, col < H, other=0).to(tl.float32) + centered, rstd = _center_and_rstd(x, H, EPS, BLOCK, LOG) + shift = tl.load(M + batch * 3 * H + col, col < H, other=0).to(tl.float32) + scale = tl.load(M + batch * 3 * H + H + col, col < H, other=0).to(tl.float32) + y = (centered * rstd) * (1.0 + scale) + shift + tl.store(Y + row * H + col, y, col < H) + if row % S == 0: + gate = tl.load(M + batch * 3 * H + 2 * H + col, col < H, other=0) + tl.store(G + batch * H + col, gate, col < H) + + +@triton.jit +def _bwd_rows( + X, + M, + DY, + DX, + PART_SHIFT, + PART_SCALE, + S: tl.constexpr, + H: tl.constexpr, + EPS: tl.constexpr, + BLOCK: tl.constexpr, + LOG: tl.constexpr, +): + row = tl.program_id(0) + batch = row // S + col = tl.arange(0, BLOCK) + x = tl.load(X + row * H + col, col < H, other=0).to(tl.float32) + dy = tl.load(DY + row * H + col, col < H, other=0).to(tl.float32) + centered, rstd = _center_and_rstd(x, H, EPS, BLOCK, LOG) + norm = centered * rstd + scale = tl.load(M + batch * 3 * H + H + col, col < H, other=0).to(tl.float32) + dnorm = dy * (1.0 + scale) + sum_dnorm = tl.div_rn(_sum_hidden(tl.where(col < H, dnorm, 0.0), BLOCK, LOG), H) + sum_dnorm_norm = tl.div_rn(_sum_hidden(tl.where(col < H, dnorm * norm, 0.0), BLOCK, LOG), H) + dx = (dnorm - sum_dnorm - norm * sum_dnorm_norm) * rstd + tl.store(DX + row * H + col, dx, col < H) + tl.store(PART_SHIFT + row * H + col, dy, col < H) + tl.store(PART_SCALE + row * H + col, dy * norm, col < H) + + +@triton.jit +def _bwd_reduce( + PART_SHIFT, PART_SCALE, DG, DM, S: tl.constexpr, H: tl.constexpr, TILE: tl.constexpr +): + batch = tl.program_id(0) + col = tl.program_id(1) * TILE + tl.arange(0, TILE) + shift = tl.full((TILE,), 0.0, tl.float32) + scale = tl.full((TILE,), 0.0, tl.float32) + # ponytail: O(S) left fold; replace with a documented fixed tree if profiling demands it. + for token in range(S): + offset = (batch * S + token) * H + col + shift += tl.load(PART_SHIFT + offset, col < H, other=0) + scale += tl.load(PART_SCALE + offset, col < H, other=0) + gate = tl.load(DG + batch * H + col, col < H, other=0).to(tl.float32) + tl.store(DM + batch * 3 * H + col, shift, col < H) + tl.store(DM + batch * 3 * H + H + col, scale, col < H) + tl.store(DM + batch * 3 * H + 2 * H + col, gate, col < H) + + +class _AdaLNTriton(torch.autograd.Function): + @staticmethod + def forward(ctx, x, modulation, eps, num_warps, reduction_tile): + batch, seq, hidden = x.shape + y = torch.empty_like(x) + gate = torch.empty((batch, 1, hidden), device=x.device, dtype=x.dtype) + with torch.cuda.device(x.device): + _fwd[(batch * seq,)]( + x, + modulation, + y, + gate, + seq, + hidden, + eps, + triton.next_power_of_2(hidden), + triton.next_power_of_2(hidden).bit_length() - 1, + enable_fp_fusion=False, + enable_reflect_ftz=False, + num_warps=num_warps, + ) + ctx.save_for_backward(x, modulation) + ctx.eps = eps + ctx.num_warps = num_warps + ctx.reduction_tile = reduction_tile + return y, gate + + @staticmethod + def backward(ctx, grad_y, grad_gate): + x, modulation = ctx.saved_tensors + batch, seq, hidden = x.shape + dx = torch.empty_like(x) + partial_shift = torch.empty(x.shape, device=x.device, dtype=torch.float32) + partial_scale = torch.empty_like(partial_shift) + dm = torch.empty_like(modulation) + _bwd_rows[(batch * seq,)]( + x, + modulation, + grad_y.contiguous(), + dx, + partial_shift, + partial_scale, + seq, + hidden, + ctx.eps, + triton.next_power_of_2(hidden), + triton.next_power_of_2(hidden).bit_length() - 1, + enable_fp_fusion=False, + enable_reflect_ftz=False, + num_warps=ctx.num_warps, + ) + _bwd_reduce[(batch, triton.cdiv(hidden, ctx.reduction_tile))]( + partial_shift, + partial_scale, + grad_gate.contiguous(), + dm, + seq, + hidden, + ctx.reduction_tile, + enable_fp_fusion=False, + num_warps=ctx.num_warps, + ) + return dx, dm, None, None, None + + +class TritonAdaLNModulationOp: + def __init__(self, *, num_warps=4, reduction_tile=128): + if num_warps not in (4, 8) or reduction_tile not in (64, 128, 256): + raise ValueError( + "AdaLN requires num_warps in {4, 8} and reduction_tile in {64, 128, 256}" + ) + self.num_warps = num_warps + self.reduction_tile = reduction_tile + + indexed_available = True + + def __call__(self, x, modulation, *, eps=1e-6, modulate_index=None): + return self.forward(x, modulation, eps=eps, modulate_index=modulate_index) + + def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): + if modulate_index is not None: + from .adaln_indexed import indexed_forward + + return indexed_forward( + x, modulation, modulate_index, eps, self.num_warps, self.reduction_tile + ) + _validate(x, modulation, eps) + if max(x.numel(), modulation.numel()) > 2**31 - 1: + raise ValueError("Triton AdaLN requires tensors to fit int32 indexing") + if not x.is_cuda: + raise RuntimeError("Triton AdaLN requires CUDA tensors") + return _AdaLNTriton.apply( + x.contiguous(), + modulation.contiguous(), + float(eps), + self.num_warps, + self.reduction_tile, + ) diff --git a/rl_engine/reference/norm/adaln_indexed.py b/rl_engine/reference/norm/adaln_indexed.py new file mode 100644 index 000000000..719884f02 --- /dev/null +++ b/rl_engine/reference/norm/adaln_indexed.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""CPU byte reference for the indexed FP32 tree/ascending-token profile. + +This is separate from both framework select01 fallback and the accuracy gold. +No backend normalization or reduction helper is used. +""" + +import torch + + +def _hidden_tree(values): + width = 1 << (values.shape[-1] - 1).bit_length() + values = torch.nn.functional.pad(values, (0, width - values.shape[-1])) + while width > 1: + width //= 2 + values = values[..., :width] + values[..., width:] + return values[..., 0] + + +class _IndexedFixed(torch.autograd.Function): + @staticmethod + def forward(ctx, x, modulation, index, eps): + batch, seq, hidden = x.shape + x32 = x.float() + mean = _hidden_tree(x32) / hidden + centered = x32 - mean[..., None] + variance = _hidden_tree(centered * centered) / hidden + root = (variance + eps).double().sqrt().float() + rstd = 1.0 / root + norm = centered * rstd[..., None] + rows = torch.arange(batch)[:, None] + index.long() * batch + selected = modulation[rows] + shift, scale, gate = selected.chunk(3, dim=-1) + shift, scale = shift.float(), scale.float() + ctx.set_materialize_grads(False) + ctx.save_for_backward(norm, rstd, scale, index) + ctx.mod_shape = modulation.shape + ctx.dtype = x.dtype + return (norm * (1.0 + scale) + shift).to(x.dtype), gate + + @staticmethod + def backward(ctx, dy, dg): + norm, rstd, scale, index = ctx.saved_tensors + batch, seq, hidden = norm.shape + if dg is not None: + dg = dg.float() + dx = None + dm = torch.zeros(ctx.mod_shape, dtype=torch.float32) + if dy is not None: + dy = dy.float() + dnorm = dy * (1.0 + scale) + dx = ( + dnorm + - (_hidden_tree(dnorm) / hidden)[..., None] + - norm * (_hidden_tree(dnorm * norm) / hidden)[..., None] + ) * rstd[..., None] + for b in range(batch): + for s in range(seq): + bank = int(index[b if index.shape[0] > 1 else 0, s]) + row = b + bank * batch + if dy is not None: + dm[row, :hidden] = dm[row, :hidden] + dy[b, s] + contribution = dy[b, s] * norm[b, s] + dm[row, hidden : 2 * hidden] = dm[row, hidden : 2 * hidden] + contribution + if dg is not None: + dm[row, 2 * hidden :] = dm[row, 2 * hidden :] + dg[b, s] + return None if dx is None else dx.to(ctx.dtype), dm.to(ctx.dtype), None, None + + +class IndexedAdaLNFixedReference: + """FP32/BF16 CPU callable for fixed-order byte tests, not a dispatch fallback.""" + + def __call__(self, x, modulation, *, modulate_index, eps=1e-6): + if x.device.type != "cpu" or modulation.device.type != "cpu": + raise ValueError("indexed byte reference requires CPU tensors") + if x.dtype not in (torch.float32, torch.bfloat16) or modulation.dtype != x.dtype: + raise TypeError("indexed byte reference requires matching FP32 or BF16 inputs") + return _IndexedFixed.apply(x, modulation, modulate_index.cpu(), float(eps)) diff --git a/rl_engine/reference/norm/adaln_modulation.py b/rl_engine/reference/norm/adaln_modulation.py new file mode 100644 index 000000000..e9da50567 --- /dev/null +++ b/rl_engine/reference/norm/adaln_modulation.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""FP32 AdaLN reference for one Qwen-Image attention or MLP modulation.""" + +from __future__ import annotations + +import math +import struct + +import torch + +from rl_engine.ops.autograd.vjp_fp32 import reduce_rows_fp32 + + +def _sum_hidden(values: torch.Tensor) -> torch.Tensor: + """Fixed pairwise tree over H, independent of batch size and row count.""" + hidden = values.shape[-1] + width = 1 << (hidden - 1).bit_length() + if width != hidden: + values = torch.nn.functional.pad(values, (0, width - hidden)) + while width > 1: + width //= 2 + values = values[..., :width] + values[..., width:] + return values[..., 0] + + +def _normalize_fp32(x: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Tensor]: + hidden = x.shape[-1] + x32 = x.float() + mean = _sum_hidden(x32) / hidden + centered = x32 - mean.unsqueeze(-1) + variance_eps = _sum_hidden(centered * centered) / hidden + eps + # CPU torch.sqrt(float32) can miss a correctly rounded result by one ULP. + root = torch.sqrt(variance_eps.double()).float() + rstd = torch.reciprocal(root) + return centered * rstd.unsqueeze(-1), rstd + + +def _eps_supported_fp32(eps: float) -> bool: + try: + eps32 = struct.unpack("f", struct.pack("f", eps))[0] + except OverflowError: + return False + return math.isfinite(eps32) and eps >= 0 and (eps == 0 or eps32 > 0) + + +def _int32_addressable(*tensors): + # Copied kernels use int32 element offsets for x/y/gate/dx and each FP32 + # partial, both modulation banks/dm, and route rows. Source views are checked + # before their framework copies too; expanded (zero-stride) axes cost no span. + for tensor in tensors: + span = sum( + (size - 1) * stride for size, stride in zip(tensor.shape, tensor.stride(), strict=True) + ) + if tensor.numel() > 2**31 - 1 or span > 2**31 - 1: + return False + return True + + +def _indexed_domain_reason(x, modulation, index, eps): + if not _int32_addressable(x, modulation, index): + return "int32 indexing" + if x.shape[0] > 65535: + return "B <= 65535" + if x.shape[-1] > 4096: + return "H <= 4096" + if not _eps_supported_fp32(eps): + return "finite FP32 eps without positive underflow" + if not x.is_cuda or torch.version.hip is not None: + return "NVIDIA CUDA tensors" + return None + + +def _validate(x: torch.Tensor, modulation: torch.Tensor, eps: float) -> None: + if x.ndim != 3 or min(x.shape) == 0: + raise ValueError("x must have nonempty shape [B, S, H]") + if modulation.shape != (x.shape[0], 3 * x.shape[-1]): + raise ValueError("modulation must have shape [B, 3H]") + if x.device != modulation.device or x.dtype != modulation.dtype: + raise ValueError("x and modulation must share device and dtype") + if x.dtype not in (torch.float32, torch.bfloat16): + raise TypeError("x and modulation must be FP32 or BF16") + if not math.isfinite(eps) or eps < 0: + raise ValueError("eps must be finite and nonnegative") + if not _eps_supported_fp32(eps): + raise ValueError("eps must remain finite FP32 without positive underflow") + + +class _AdaLNReference(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, modulation: torch.Tensor, eps: float): + shift, scale, _ = modulation.float().chunk(3, dim=-1) + norm, rstd = _normalize_fp32(x, eps) + y = (norm * (1.0 + scale[:, None, :]) + shift[:, None, :]).to(x.dtype) + ctx.save_for_backward(norm, rstd, scale) + ctx.x_dtype = x.dtype + ctx.mod_dtype = modulation.dtype + # Gate has no arithmetic; preserve its original BF16 payload bits. + return y, modulation[:, None, 2 * x.shape[-1] :] + + @staticmethod + def backward(ctx, grad_y: torch.Tensor, grad_gate: torch.Tensor): + norm, rstd, scale = ctx.saved_tensors + hidden = norm.shape[-1] + dy = grad_y.float() + dnorm = dy * (1.0 + scale[:, None, :]) + mean_dnorm = _sum_hidden(dnorm) / hidden + mean_dnorm_norm = _sum_hidden(dnorm * norm) / hidden + dx = ( + (dnorm - mean_dnorm.unsqueeze(-1) - norm * mean_dnorm_norm.unsqueeze(-1)) + * rstd.unsqueeze(-1) + ).to(ctx.x_dtype) + # S is folded in ascending logical-token order for every sample. + dm = torch.cat( + [ + torch.stack([reduce_rows_fp32(rows) for rows in dy]).to(ctx.mod_dtype), + torch.stack([reduce_rows_fp32(rows) for rows in dy * norm]).to(ctx.mod_dtype), + grad_gate.squeeze(1), + ], + dim=-1, + ) + return dx, dm, None + + +def _validate_select01_metadata( + x: torch.Tensor, + modulation: torch.Tensor, + modulate_index: torch.Tensor, + eps: float, +) -> None: + if x.ndim != 3 or min(x.shape) == 0: + raise ValueError("x must have nonempty shape [B, S, H]") + batch, seq, hidden = x.shape + if modulation.shape != (2 * batch, 3 * hidden): + raise ValueError("indexed modulation must have shape [2B, 3H]") + if modulate_index.shape not in ((batch, seq), (1, seq)): + raise ValueError("modulate_index must have shape [B, S] or [1, S]") + if x.device != modulation.device or x.device != modulate_index.device: + raise ValueError("x, modulation, and modulate_index must share device") + if x.dtype != modulation.dtype or x.dtype not in (torch.float32, torch.bfloat16): + raise TypeError("x and modulation must share FP32 or BF16 dtype") + if modulate_index.dtype not in ( + torch.bool, + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + ): + raise TypeError("modulate_index must have boolean or integer dtype") + if not math.isfinite(eps) or eps < 0: + raise ValueError("eps must be finite and nonnegative") + + +def _validate_select01(x, modulation, modulate_index, eps): + _validate_select01_metadata(x, modulation, modulate_index, eps) + if not torch.all((modulate_index == 0) | (modulate_index == 1)).item(): + raise ValueError("modulate_index values must be 0 or 1") + + +def _select01_reference(x, modulation, modulate_index, eps): + _validate_select01(x, modulation, modulate_index, eps) + batch = x.shape[0] + shift, scale, gate = modulation.float().chunk(3, dim=-1) + index = modulate_index.bool().unsqueeze(-1) + shift = torch.where(index, shift[batch:, None, :], shift[:batch, None, :]) + scale = torch.where(index, scale[batch:, None, :], scale[:batch, None, :]) + gate = torch.where(index, gate[batch:, None, :], gate[:batch, None, :]) + + norm, _ = _normalize_fp32(x, eps) + return (norm * (1.0 + scale) + shift).to(x.dtype), gate.to(modulation.dtype) + + +class NativeAdaLNModulationOp: + def __call__( + self, + x: torch.Tensor, + modulation: torch.Tensor, + *, + modulate_index: torch.Tensor | None = None, + eps: float = 1e-6, + ): + return self.forward(x, modulation, modulate_index=modulate_index, eps=eps) + + def forward( + self, + x: torch.Tensor, + modulation: torch.Tensor, + *, + modulate_index: torch.Tensor | None = None, + eps: float = 1e-6, + ): + if modulate_index is not None: + return _select01_reference(x, modulation, modulate_index, float(eps)) + _validate(x, modulation, eps) + return _AdaLNReference.apply(x, modulation, float(eps)) diff --git a/rl_engine/runtime/adaln.py b/rl_engine/runtime/adaln.py new file mode 100644 index 000000000..2fbd75830 --- /dev/null +++ b/rl_engine/runtime/adaln.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""AdaLN execution policy; arithmetic stays in the existing backend operators.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass +from typing import Any + +import torch + +SHARED_PROFILE = "adaln-shared-fp32-tree-v1" +SELECT01_PROFILE = "adaln-select01-framework-v1" +INDEXED_PROFILE = "adaln-indexed-fp32-tree-v1" +_BACKENDS = { + "rl_engine.reference.norm.adaln_modulation": ("NativeAdaLNModulationOp", "pytorch"), + "rl_engine.backends.cuda.norm.adaln_modulation": ("CudaAdaLNModulationOp", "cuda"), + "rl_engine.backends.shared.triton.norm.adaln_modulation": ("TritonAdaLNModulationOp", "triton"), +} + + +@dataclass(frozen=True) +class IndexedAdaLNQualification: + """Explicit activation after external evidence review, never an attestation. + + Pass the same scope to acceptance and delivery. No scope is installed by + default. Source/binary hashes and evidence stay in the external manifest. + """ + + backend: str + device: str + device_name: str + capability: tuple[int, int] + torch_version: str + cuda_version: str + triton_version: str | None = None + + def __post_init__(self): + device = torch.device(self.device) + if self.backend not in {"cuda", "triton"}: + raise ValueError("indexed qualification requires cuda or triton backend") + if device.type != "cuda" or device.index is None: + raise ValueError("indexed qualification requires an explicit cuda device ordinal") + if not self.device_name or not self.torch_version or not self.cuda_version: + raise ValueError("indexed qualification requires explicit target/runtime identity") + if len(self.capability) != 2 or any(type(v) is not int or v < 0 for v in self.capability): + raise ValueError("indexed qualification requires a compute capability pair") + if self.backend == "triton" and not self.triton_version: + raise ValueError("indexed Triton qualification requires its runtime version") + + def check(self, backend, device, *, execution=False): + """Fail closed outside the externally reviewed backend/target runtime.""" + device = torch.device(device) + target = torch.device(self.device) + compatible = ( + backend == self.backend + and device.type == "cuda" + and (device.index is None or device.index == target.index) + and torch.version.hip is None + and torch.__version__ == self.torch_version + and torch.version.cuda == self.cuda_version + ) + if compatible and backend == "triton": + import triton + + compatible = triton.__version__ == self.triton_version + if compatible and execution: + compatible = ( + device.index == target.index + and torch.cuda.get_device_name(device) == self.device_name + and torch.cuda.get_device_capability(device) == self.capability + ) + if not compatible: + raise RuntimeError( + "strict AdaLN rejects indexed backend/target/runtime outside qualification" + ) + + +def adaln_backend(op: Any) -> str | None: + cls = type(op) + entry = _BACKENDS.get(cls.__module__) + if entry and getattr(sys.modules[cls.__module__], entry[0], None) is cls: + return entry[1] + return None + + +class AdaLNDispatch: + """A resolution-local trace, updated only after successful operator execution. + + Strict enforces arithmetic and an explicit indexed activation scope. Forward + observations never attest hardware, binaries, backward execution or profiling. + """ + + def __init__( + self, + implementation, + trace, + *, + hidden, + strict=False, + indexed_qualification=None, + device=None, + ): + self.implementation = implementation + self.trace = trace + self.hidden = hidden + self.strict = strict + self.indexed_qualification = indexed_qualification + self.device = torch.device(device) if device is not None else None + trace.update( + actual_backend=None, + observation=None, + execution_dtype=None, + output_dtype=None, + execution_device=None, + binary_path=None, + binary_sha256=None, + hardware_qualified=None, + profiler_observed=False, + semantic_version="adaln-v1", + strict=strict, + ) + self._resolution_trace = trace.copy() + + def __call__(self, x, modulation, *, eps=1e-6, modulate_index=None): + return self.forward(x, modulation, eps=eps, modulate_index=modulate_index) + + def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): + self.trace.clear() + self.trace.update(self._resolution_trace) + backend = adaln_backend(self.implementation) + indexed = modulate_index is not None + profile = SELECT01_PROFILE if indexed else SHARED_PROFILE + implementation = self.implementation + shared_unsupported = False + if indexed: + from rl_engine.reference.norm.adaln_modulation import ( + _indexed_domain_reason, + _validate_select01, + _validate_select01_metadata, + ) + + _validate_select01_metadata(x, modulation, modulate_index, eps) + reason = _indexed_domain_reason(x, modulation, modulate_index, eps) + else: + from rl_engine.reference.norm.adaln_modulation import _int32_addressable, _validate + + _validate(x, modulation, eps) + shared_unsupported = backend in {"cuda", "triton"} and ( + x.device.type != "cuda" + or not _int32_addressable(x, modulation) + or ( + backend == "cuda" + and (torch.version.hip is not None or x.shape[0] > 65535 or x.shape[-1] > 4096) + ) + ) + if self.strict: + try: + if shared_unsupported: + raise RuntimeError("strict AdaLN rejects unsupported shared backend domain") + if indexed: + if reason is not None: + raise RuntimeError( + f"strict AdaLN rejects indexed/select01: requires {reason}" + ) + if not getattr(self.implementation, "indexed_available", False): + raise RuntimeError( + "strict AdaLN rejects select01: indexed symbols unavailable" + ) + if self.indexed_qualification is None: + raise RuntimeError( + "strict AdaLN rejects unqualified indexed/select01 execution" + ) + self.indexed_qualification.check(backend, x.device, execution=True) + if backend is None or (backend == "pytorch" and x.device.type != "cpu"): + raise RuntimeError("strict AdaLN backend has no declared replay profile") + if backend in {"cuda", "triton"} and ( + x.device.type != "cuda" or torch.version.hip is not None + ): + raise RuntimeError("strict AdaLN GPU profile requires NVIDIA CUDA") + if x.ndim != 3 or x.shape[-1] != self.hidden: + raise ValueError("strict AdaLN input hidden differs from resolved hidden") + if self.device is not None and ( + self.device.type != x.device.type + or (self.device.index is not None and self.device.index != x.device.index) + ): + raise RuntimeError("strict AdaLN rejects device differing from resolution") + except (ValueError, RuntimeError): + # Preserve input-error precedence on rejection. Successful calls + # delegate value validation to the public backend exactly once. + if indexed and reason not in {"int32 indexing", "B <= 65535"}: + _validate_select01(x, modulation, modulate_index, eps) + raise + if indexed: + if backend in {"cuda", "triton"}: + if getattr(implementation, "indexed_available", False) and reason is None: + profile = INDEXED_PROFILE + else: + from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp + + implementation = NativeAdaLNModulationOp() + backend = "pytorch" + outputs = implementation(x, modulation, eps=eps, modulate_index=modulate_index) + else: + if shared_unsupported: + from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp + + implementation = NativeAdaLNModulationOp() + backend = "pytorch" + outputs = implementation(x, modulation, eps=eps) + self.trace.update( + actual_backend=backend, + observation="successful_forward", + arithmetic_profile=profile, + execution_dtype=str(x.dtype).removeprefix("torch."), + output_dtype=str(outputs[0].dtype).removeprefix("torch."), + execution_device=str(x.device), + ) + if indexed and profile == SELECT01_PROFILE: + self.trace.update( + fallback=True, + fallback_reason=( + f"indexed_{adaln_backend(self.implementation)}_domain_unsupported" + if self._resolution_trace["arithmetic_profile"] == INDEXED_PROFILE + else "modulate_index_requires_select01" + ), + reduction_order="pairwise_lower_upper_H; framework_select01_backward", + kernel_fingerprint="adaln_modulation_select01_native_v1", + ) + elif indexed: + self.trace.update( + fallback=False, + fallback_reason=None, + reduction_order="pairwise_lower_upper_H; ascending_selected_token_S_per_bank", + kernel_fingerprint=f"adaln_indexed_{backend}_v1", + ) + elif self._resolution_trace["arithmetic_profile"] in {INDEXED_PROFILE, SELECT01_PROFILE}: + self.trace.update( + fallback=bool(self.trace["rejected_backends"]), + fallback_reason="; ".join(self.trace["rejected_backends"]) or None, + reduction_order="pairwise_lower_upper_H; ascending_logical_token_S", + kernel_fingerprint=f"adaln_modulation_v1:{self.trace['selected_backend']}", + ) + if shared_unsupported: + self.trace.update( + fallback=True, + fallback_reason=f"shared_{adaln_backend(self.implementation)}_domain_unsupported", + kernel_fingerprint="adaln_modulation_v1:PYTORCH_ADALN_MODULATION", + ) + if backend == "cuda": + from rl_engine.backends.extension import _C + + self.trace["binary_path"] = getattr(_C, "__file__", None) + return outputs diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 7751d5901..13de847d7 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -31,6 +31,7 @@ MaskMode, ) from rl_engine.platforms.device import device_ctx +from rl_engine.runtime.adaln import IndexedAdaLNQualification from rl_engine.runtime.semantic_registry import ( OperatorBackendDescriptor, OperatorFallbackPolicy, @@ -177,6 +178,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): TRITON_SILU = "rl_engine.kernels.ops.triton.activation.swiglu.TritonSiLUOp" TRITON_SWIGLU = "rl_engine.kernels.ops.triton.activation.swiglu.TritonSwiGLUOp" + PYTORCH_ADALN_MODULATION = "rl_engine.reference.norm.adaln_modulation.NativeAdaLNModulationOp" + TRITON_ADALN_MODULATION = ( + "rl_engine.backends.shared.triton.norm.adaln_modulation.TritonAdaLNModulationOp" + ) + CUDA_ADALN_MODULATION = "rl_engine.backends.cuda.norm.adaln_modulation.CudaAdaLNModulationOp" + # WS1 pure-PyTorch ground-truth attention reference (hand-written fp32 softmax). # Distinct from PYTORCH_ATTN above, which is the production SDPA fallback. PYTORCH_NATIVE_ATTENTION = ( @@ -602,6 +609,11 @@ def __init__(self): OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU, ], + "adaln_modulation": [ + OpBackend.CUDA_ADALN_MODULATION, + OpBackend.TRITON_ADALN_MODULATION, + OpBackend.PYTORCH_ADALN_MODULATION, + ], # Default dispatch logic for new operators "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rope": [ @@ -649,6 +661,7 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "adaln_modulation": [OpBackend.PYTORCH_ADALN_MODULATION], }, "musa": { "logp": [ @@ -697,6 +710,7 @@ def __init__(self): ], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "adaln_modulation": [OpBackend.PYTORCH_ADALN_MODULATION], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -722,11 +736,13 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], + "adaln_modulation": [OpBackend.PYTORCH_ADALN_MODULATION], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only # Ascend-accelerated ops are listed. "npu": { + "adaln_modulation": [OpBackend.PYTORCH_ADALN_MODULATION], "batch_invariant_logp": [ OpBackend.ASCEND_BATCH_INVARIANT_LOGP, OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, @@ -1047,6 +1063,141 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + def get_adaln_modulation_op( + self, + device: torch.device | str | None = None, + *, + hidden: int, + modulate_index: torch.Tensor | None = None, + strict: bool = False, + indexed_qualification: IndexedAdaLNQualification | None = None, + ) -> tuple[Any, dict[str, Any]]: + """Resolve AdaLN; indexed strict requires an explicit AdaLN-local scope.""" + + from rl_engine.runtime.adaln import ( + INDEXED_PROFILE, + SELECT01_PROFILE, + SHARED_PROFILE, + AdaLNDispatch, + adaln_backend, + ) + + if hidden <= 0: + raise ValueError("hidden must be positive") + if modulate_index is not None: + if modulate_index.ndim != 2 or min(modulate_index.shape) == 0: + raise ValueError("modulate_index must have nonempty shape [B, S] or [1, S]") + if modulate_index.dtype not in ( + torch.bool, + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + ): + raise TypeError("modulate_index must have boolean or integer dtype") + profile = SELECT01_PROFILE if modulate_index is not None else SHARED_PROFILE + platform = self._platform_for_device(device) + resolved_device = device if device is not None else device_ctx.device + if modulate_index is not None: + backend = OpBackend.PYTORCH_ADALN_MODULATION + # Resolution is provisional; execution checks actual geometry, device and eps. + if platform == "cuda" and hidden <= 4096: + for candidate in ( + OpBackend.CUDA_ADALN_MODULATION, + OpBackend.TRITON_ADALN_MODULATION, + ): + implementation = self._get_or_create_backend(candidate) + if implementation is not None and getattr( + implementation, "indexed_available", False + ): + backend = candidate + break + op = self._get_or_create_backend(backend) + if op is None: + raise RuntimeError("PyTorch select01 AdaLN fallback is unavailable") + accelerated = backend is not OpBackend.PYTORCH_ADALN_MODULATION + actual = adaln_backend(op) + if strict: + if not accelerated or indexed_qualification is None: + raise RuntimeError( + "strict AdaLN rejects unqualified indexed/select01 execution" + ) + indexed_qualification.check(actual, resolved_device) + trace = { + "selected_backend": backend.name, + "fallback": not accelerated, + "fallback_reason": None if accelerated else "modulate_index_requires_select01", + "rejected_backends": [], + "reduction_order": ( + "pairwise_lower_upper_H; ascending_selected_token_S_per_bank" + if accelerated + else "pairwise_lower_upper_H; framework_select01_backward" + ), + "accumulator_dtype": "fp32", + "split_k": False, + "stream_k": False, + "tf32": False, + "kernel_fingerprint": ( + f"adaln_indexed_{actual}_v1" + if accelerated + else "adaln_modulation_select01_native_v1" + ), + "requested_backend": actual, + "arithmetic_profile": INDEXED_PROFILE if accelerated else profile, + } + return ( + AdaLNDispatch( + op, + trace, + hidden=hidden, + strict=strict, + indexed_qualification=indexed_qualification, + device=resolved_device, + ), + trace, + ) + candidates = self._priority_map[platform]["adaln_modulation"] + rejected: list[str] = [] + for backend in candidates: + if backend is OpBackend.CUDA_ADALN_MODULATION and hidden > 4096: + rejected.append("CUDA_ADALN_MODULATION: H > 4096") + continue + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(backend.name) + continue + actual = adaln_backend(op) + if strict and (actual is None or (actual == "pytorch" and platform != "cpu")): + rejected.append(f"{backend.name}: no declared strict replay profile") + continue + trace = { + "selected_backend": backend.name, + "fallback": bool(rejected), + "rejected_backends": rejected, + "reduction_order": "pairwise_lower_upper_H; ascending_logical_token_S", + "accumulator_dtype": "fp32", + "split_k": False, + "stream_k": False, + "tf32": False, + "kernel_fingerprint": f"adaln_modulation_v1:{backend.name}", + "arithmetic_profile": profile, + "requested_backend": candidates[0].name.split("_")[0].lower(), + "fallback_reason": "; ".join(rejected) if rejected else None, + } + return ( + AdaLNDispatch( + op, + trace, + hidden=hidden, + strict=strict, + indexed_qualification=indexed_qualification, + device=resolved_device, + ), + trace, + ) + raise RuntimeError(f"No functional adaln_modulation backend on {platform}: {rejected}") + def _platform_for_device(self, device: torch.device | str | None) -> str: if device is None: return self._platform() diff --git a/rl_engine/validation/operators/op_checks.py b/rl_engine/validation/operators/op_checks.py index 5f3aded34..c8d511545 100644 --- a/rl_engine/validation/operators/op_checks.py +++ b/rl_engine/validation/operators/op_checks.py @@ -29,6 +29,8 @@ class OperatorCase: inputs: Mapping[str, Any] gold_fn: Callable[..., Any] grad_input_names: tuple[str, ...] = () + gold_input_dtype: torch.dtype | None = None + candidate_output_dtype: torch.dtype | None = None @dataclass(frozen=True) @@ -71,6 +73,7 @@ class CaseCheck: op_class: str passed: bool outputs: list[OutputCheck] + backend_trace: dict[str, Any] | None = None @dataclass(frozen=True) @@ -217,7 +220,9 @@ def _run_case_backward( raise ValueError(f"case {case.name!r} does not declare gradient inputs") candidate_inputs = _clone_inputs_for_backward(case.inputs, case.grad_input_names) - gold_inputs = _clone_inputs_for_backward(case.inputs, case.grad_input_names) + gold_inputs = _clone_inputs_for_backward( + case.inputs, case.grad_input_names, dtype=case.gold_input_dtype + ) candidate_outputs = _flatten_tensors(_call_candidate(candidate.fn, candidate_inputs)) gold_outputs = _flatten_tensors(case.gold_fn(**gold_inputs)) # Candidate and gold must use the same upstream gradients; otherwise we @@ -301,6 +306,7 @@ def _run_case_backward( op_class=case.op_class, passed=all(output.passed for output in checks), outputs=checks, + backend_trace=dict(candidate.fn.trace) if hasattr(candidate.fn, "trace") else None, ) @@ -364,6 +370,7 @@ def _compare_case_outputs( atol=atol, rtol=rtol, judgment="forward_accuracy", + expected_dtype=case.candidate_output_dtype, comparison_lhs_role=( forward_spec.comparison_lhs_role if forward_spec is not None else "bf16_candidate" ), @@ -381,6 +388,7 @@ def _compare_case_outputs( op_class=case.op_class, passed=all(output.passed for output in output_checks), outputs=output_checks, + backend_trace=dict(candidate.fn.trace) if hasattr(candidate.fn, "trace") else None, ) @@ -394,12 +402,16 @@ def _call_candidate(candidate: Callable[..., Any] | Any, inputs: Mapping[str, An def _clone_inputs_for_backward( inputs: Mapping[str, Any], grad_input_names: tuple[str, ...], + *, + dtype: torch.dtype | None = None, ) -> dict[str, Any]: grad_names = set(grad_input_names) cloned: dict[str, Any] = {} for name, value in inputs.items(): if isinstance(value, torch.Tensor): tensor = value.detach().clone() + if dtype is not None and tensor.is_floating_point(): + tensor = tensor.to(dtype=dtype) if name in grad_names: if not tensor.is_floating_point(): raise TypeError(f"gradient input {name!r} must be floating point") @@ -550,6 +562,7 @@ def _compare_output( comparison_lhs_role: str = "bf16_candidate", comparison_rhs_role: str = "fp32_reference", message: str = "", + expected_dtype: torch.dtype | None = None, ) -> OutputCheck: if candidate.shape != gold.shape: return OutputCheck( @@ -569,8 +582,12 @@ def _compare_output( message=f"shape mismatch: candidate={tuple(candidate.shape)} gold={tuple(gold.shape)}", ) - candidate_fp32 = candidate.float() - gold_fp32 = gold.float() + candidate_fp32 = candidate.detach().float() + gold_fp32 = gold.detach().float() + # Independent references may live on CPU even when the candidate runs on GPU. + if candidate.device != gold.device: + candidate_fp32 = candidate_fp32.cpu() + gold_fp32 = gold_fp32.cpu() abs_error = (candidate_fp32 - gold_fp32).abs() if abs_error.numel() == 0: max_abs_error = 0.0 @@ -582,6 +599,9 @@ def _compare_output( rel_error = abs_error / gold_fp32.abs().clamp_min(1e-12) max_rel_error = float(rel_error.max().item()) + dtype_matches = expected_dtype is None or candidate.dtype == expected_dtype + if not dtype_matches: + message = f"dtype mismatch: candidate={candidate.dtype} expected={expected_dtype}" return OutputCheck( output_index=output_index, shape=tuple(candidate.shape), @@ -592,7 +612,8 @@ def _compare_output( max_abs_error=max_abs_error, mean_abs_error=mean_abs_error, max_rel_error=max_rel_error, - passed=bool(torch.allclose(candidate_fp32, gold_fp32, atol=atol, rtol=rtol)), + passed=dtype_matches + and bool(torch.allclose(candidate_fp32, gold_fp32, atol=atol, rtol=rtol)), judgment=judgment, comparison_lhs_role=comparison_lhs_role, comparison_rhs_role=comparison_rhs_role, diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index ca3b7c120..a87677c45 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -25,6 +25,7 @@ def make_operator_inputs( device: torch.device, ) -> dict[str, Any]: builders = { + "adaln_modulation": _make_adaln_modulation_inputs, "rms_norm": _make_rms_norm_inputs, "qk_norm": _make_qk_norm_inputs, "pack": _make_pack_inputs, @@ -53,6 +54,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: batch, seq = _batch_seq(args) vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) names = { + "adaln_modulation": f"{batch}x{seq}x{_normalized_dim(args)}", "rms_norm": f"{batch}x{seq}x{_normalized_dim(args)}", "qk_norm": f"{batch}x{seq}x{_arg_int(args, 'n_heads', DEFAULT_N_HEADS)}x" f"{_arg_int(args, 'head_dim', DEFAULT_HEAD_DIM)}", @@ -79,6 +81,18 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: raise ValueError(f"unsupported operator shape: {op_name}") from exc +def _make_adaln_modulation_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + hidden = _normalized_dim(args) + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, offset=0), + "modulation": _floating_tensor((batch, 3 * hidden), args, dtype, device, offset=1), + "eps": _arg_float(args, "eps", DEFAULT_RMS_EPS), + } + + def _make_rms_norm_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..5b2ac615d 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -22,6 +22,8 @@ class OperatorSpec: gold_method: str candidate_paths: dict[str, str] grad_input_names: tuple[str, ...] = () + gold_input_dtype: torch.dtype | None = None + output_matches_input_dtype: bool = False def _load_object(path: str) -> Any: @@ -32,6 +34,22 @@ def _load_object(path: str) -> Any: OP_SPECS = { + "adaln_modulation": OperatorSpec( + name="adaln_modulation", + op_class="reduction", + gold_path="rl_engine.validation.operators.operator_specs.GtestAdaLNReference", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ("rl_engine.reference.norm.adaln_modulation.NativeAdaLNModulationOp"), + "triton": ( + "rl_engine.backends.shared.triton.norm.adaln_modulation.TritonAdaLNModulationOp" + ), + "cuda": "rl_engine.backends.cuda.norm.adaln_modulation.CudaAdaLNModulationOp", + }, + grad_input_names=("x", "modulation"), + gold_input_dtype=torch.float32, + output_matches_input_dtype=True, + ), "rms_norm": OperatorSpec( name="rms_norm", op_class="reduction", @@ -275,6 +293,79 @@ def _load_object(path: str) -> Any: } +class _GtestSharedAdaLNAffine(torch.autograd.Function): + """Independent CPU affine VJP with the specified ascending FP32 S fold.""" + + @staticmethod + def forward(ctx, norm, modulation): + shift, scale, gate = modulation.chunk(3, dim=-1) + ctx.save_for_backward(norm, scale) + return norm * (1 + scale[:, None, :]) + shift[:, None, :], gate[:, None, :] + + @staticmethod + def backward(ctx, dy, dg): + norm, scale = ctx.saved_tensors + shift_grad = torch.zeros_like(scale) + scale_grad = torch.zeros_like(scale) + # CPU torch's broadcast backward may reassociate this reduction. The + # reference follows the declared token order without any kernel helper. + for token in range(norm.shape[1]): + shift_grad = shift_grad + dy[:, token, :] + scale_grad = scale_grad + dy[:, token, :] * norm[:, token, :] + dm = torch.cat((shift_grad, scale_grad, dg[:, 0, :]), dim=-1) + return dy * (1 + scale[:, None, :]), dm + + +class _GtestIndexedAdaLNAffine(torch.autograd.Function): + """Independent selected-token affine VJP; LayerNorm stays in CPU autograd.""" + + @staticmethod + def forward(ctx, norm, modulation, index): + batch = norm.shape[0] + rows = torch.arange(batch)[:, None] + index.long() * batch + shift, scale, gate = modulation[rows].chunk(3, dim=-1) + ctx.set_materialize_grads(False) + ctx.save_for_backward(norm, scale, rows) + ctx.mod_shape = modulation.shape + return norm * (1 + scale) + shift, gate + + @staticmethod + def backward(ctx, dy, dg): + norm, scale, rows = ctx.saved_tensors + hidden = norm.shape[-1] + dm = norm.new_zeros(ctx.mod_shape) + # Select contributions before arithmetic so other banks cannot see NaNs. + for token in range(norm.shape[1]): + for sample in range(norm.shape[0]): + row = int(rows[sample, token]) + if dy is not None: + grad = dy[sample, token] + dm[row, :hidden] += grad + product = grad * norm[sample, token] + dm[row, hidden : 2 * hidden] += product + if dg is not None: + dm[row, 2 * hidden :] += dg[sample, token] + dnorm = None if dy is None else dy * (1 + scale) + return dnorm, dm, None + + +class GtestAdaLNReference: + """Independent CPU FP32 LayerNorm/autograd gold, not the fixed-tree kernel. + + Transfers/casts retain autograd edges to the caller's inputs. Outputs stay + FP32; no production normalization or backward helper is used. Shared + affine gradients use a separate CPU VJP for the contract's FP32 token fold. + """ + + def forward_fp32(self, x, modulation, *, eps=1e-6, modulate_index=None): + x32 = x.to(device="cpu", dtype=torch.float32) + mod32 = modulation.to(device="cpu", dtype=torch.float32) + norm = torch.nn.functional.layer_norm(x32, (x.shape[-1],), eps=eps) + if modulate_index is None: + return _GtestSharedAdaLNAffine.apply(norm, mod32) + return _GtestIndexedAdaLNAffine.apply(norm, mod32, modulate_index.cpu()) + + class GtestPackOp: """gtest view of NativePackOp: compare the packed rows, not cu_seqlens.""" @@ -345,15 +436,30 @@ def make_operator_case( inputs=make_operator_inputs(args.op, args, dtype, device), gold_fn=gold_fn, grad_input_names=spec.grad_input_names, + gold_input_dtype=spec.gold_input_dtype, + candidate_output_dtype=dtype if spec.output_matches_input_dtype else None, ) -def make_candidate(args: argparse.Namespace) -> CandidateSpec: +def make_candidate(args: argparse.Namespace, *, observe_backend: bool = False) -> CandidateSpec: spec = OP_SPECS[args.op] candidate_name = "pytorch" if args.candidate == "native" else args.candidate if candidate_name in spec.candidate_paths: candidate_op = _load_object(spec.candidate_paths[candidate_name])() + if args.op == "adaln_modulation" and observe_backend: + from rl_engine.runtime.adaln import SHARED_PROFILE, AdaLNDispatch + + candidate_op = AdaLNDispatch( + candidate_op, + { + "requested_backend": candidate_name, + "fallback": False, + "fallback_reason": None, + "arithmetic_profile": SHARED_PROFILE, + }, + hidden=args.normalized_dim, + ) if args.op == "logp" and candidate_name == "cuda-sm90": candidate_op = _LogpSM90CandidateAdapter(candidate_op) return CandidateSpec( diff --git a/tests/backends/test_adaln_shared_bounds.py b/tests/backends/test_adaln_shared_bounds.py new file mode 100644 index 000000000..7a90e8bfa --- /dev/null +++ b/tests/backends/test_adaln_shared_bounds.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Public AdaLN preflight checks; enormous shapes use metadata, not allocations.""" + +import os + +import pytest +import torch + +from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("shape", [(1, 524289, 4096), (1, 2**31, 1), (2**30, 1, 1)]) +def test_triton_shared_rejects_int32_address_overflow_without_allocating(shape, dtype): + pytest.importorskip("triton") + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + x = torch.empty(shape, device="meta", dtype=dtype) + m = torch.empty(shape[0], 3 * shape[-1], device="meta", dtype=dtype) + with pytest.raises(ValueError, match="int32 indexing"): + TritonAdaLNModulationOp()(x, m) + + +@pytest.mark.parametrize("shape", [(1, 2**31 - 1, 1), (1, 1, 715827882)]) +def test_triton_shared_int32_boundary_reaches_device_check(shape): + pytest.importorskip("triton") + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + x = torch.empty(shape, device="meta") + m = torch.empty(shape[0], 3 * shape[-1], device="meta") + with pytest.raises(RuntimeError, match="requires CUDA tensors"): + TritonAdaLNModulationOp()(x, m) + + +@pytest.fixture +def cuda_op(request): + try: + return CudaAdaLNModulationOp(threads=getattr(request, "param", 256)) + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + + +@pytest.mark.parametrize("cuda_op", [128, 256, 512], indirect=True) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_cuda_shared_rejects_batch_above_grid_limit(cuda_op, dtype): + device = "cuda" if torch.cuda.is_available() else "meta" + x = torch.zeros(65536, 1, 1, device=device, dtype=dtype) + m = torch.zeros(65536, 3, device=device, dtype=dtype) + with pytest.raises(ValueError, match="B <= 65535"): + cuda_op(x, m) + + +@pytest.mark.parametrize("shape", [(1, 524289, 4096), (1, 2**31, 1)]) +def test_cuda_shared_rejects_int32_address_overflow_without_allocating(cuda_op, shape): + x = torch.empty(shape, device="meta") + m = torch.empty(shape[0], 3 * shape[-1], device="meta") + with pytest.raises(ValueError, match="int32 indexing"): + cuda_op(x, m) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_cuda_shared_rejects_hidden_4097(cuda_op, dtype): + device = "cuda" if torch.cuda.is_available() else "meta" + x = torch.empty(1, 1, 4097, device=device, dtype=dtype) + m = torch.empty(1, 3 * 4097, device=device, dtype=dtype) + with pytest.raises(RuntimeError, match="H <= 4096"): + cuda_op(x, m) + + +@pytest.mark.parametrize("cuda_op", [128, 256, 512], indirect=True) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_cuda_shared_maximum_batch_forward_backward(cuda_op, dtype): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + x = torch.tensor([-3.0, 3.0], device="cuda", dtype=dtype).repeat(65535, 1, 1) + x.requires_grad_() + m = torch.zeros(65535, 6, device="cuda", dtype=dtype) + # Exact small integer tags expose gate batch addressing, including the last row. + tag = (torch.arange(65535, device="cuda") % 17).to(dtype) + m[:, 4], m[:, 5] = tag, -tag + m.requires_grad_() + y, gate = cuda_op(x, m, eps=7.0) + dy = torch.tensor([1.0, 0.0], device="cuda", dtype=dtype).expand_as(y) + dg = torch.tensor([2.0, -3.0], device="cuda", dtype=dtype).expand_as(gate) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + # variance+eps=16, norm=[-3/4,3/4]; the VJP is [7/128,-7/128]. + assert torch.equal(y, torch.tensor([-0.75, 0.75], device="cuda", dtype=dtype).expand_as(y)) + assert torch.equal(gate[:, 0], torch.stack((tag, -tag), dim=1)) + expected_dx = torch.tensor([7 / 128, -7 / 128], device="cuda", dtype=dtype) + assert torch.equal(dx, expected_dx.expand_as(dx)) + expected = torch.zeros_like(m) + expected[:, 0], expected[:, 2] = 1.0, -0.75 + expected[:, 4], expected[:, 5] = 2.0, -3.0 + assert torch.equal(dm, expected) diff --git a/tests/benchmarks/test_benchmark_adaln_modulation.py b/tests/benchmarks/test_benchmark_adaln_modulation.py new file mode 100644 index 000000000..9588692e8 --- /dev/null +++ b/tests/benchmarks/test_benchmark_adaln_modulation.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""CPU checks of the benchmark's CLI contract and ordinary numerical baseline.""" + +import pytest +import torch + +from benchmarks.operators.norm.benchmark_adaln_modulation import layer_norm_baseline, parse_args + + +def test_benchmark_cli_records_seed_json_and_rejects_empty_measurements(): + args = parse_args(["--seed", "386", "--json", "--repeat", "2"]) + assert args.seed == 386 and args.json and args.repeat == 2 + for argv in (["--repeat", "0"], ["--warmup", "-1"]): + with pytest.raises(SystemExit) as exc: + parse_args(argv) + assert exc.value.code == 2 + + +def test_layer_norm_baseline_keeps_affine_fp32_until_the_output_and_gate_vjp(): + x = torch.tensor([[[-3.0, 3.0]]], dtype=torch.bfloat16, requires_grad=True) + m = torch.tensor([[-1.0, -1.0, 1 / 256, 1 / 256, 2.0, -3.0]], dtype=x.dtype, requires_grad=True) + y, gate = layer_norm_baseline(x, m, eps=7.0) + torch.testing.assert_close( + y, torch.tensor([[[-1.75, -0.2470703125]]], dtype=x.dtype), atol=0, rtol=0 + ) + gate.backward(torch.tensor([[[1.0, -2.0]]], dtype=x.dtype)) + torch.testing.assert_close(m.grad, torch.tensor([[0, 0, 0, 0, 1, -2]], dtype=x.dtype)) diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index b813506b9..b0ee46c4e 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -6,6 +6,7 @@ import runpy from typing import Any +import pytest import setuptools import torch from torch.utils import cpp_extension @@ -83,3 +84,15 @@ def test_package_data_real_json_is_not_reinterpreted(tmp_path) -> None: assert _resolve_package_data_source(source, project_root=tmp_path) == source assert "build_py" in get_cmdclass() + + +@pytest.mark.parametrize( + "hip,fast_math,enabled", [(None, "0", True), (None, "1", False), ("test", "0", False)] +) +def test_adaln_build_registers_native_source_only_for_precise_cuda( + monkeypatch, hip, fast_math, enabled +): + monkeypatch.setenv("KERNEL_ALIGN_USE_FAST_MATH", fast_math) + extension = _load_extension_config(monkeypatch, hip=hip) + assert ("csrc/cuda/norm/adaln_modulation.cu" in extension["sources"]) is enabled + assert ("-DRLK_ADALN_CUDA_ENABLED" in extension["extra_compile_args"]["cxx"]) is enabled diff --git a/tests/helpers/adaln_indexed.py b/tests/helpers/adaln_indexed.py new file mode 100644 index 000000000..fc4a17131 --- /dev/null +++ b/tests/helpers/adaln_indexed.py @@ -0,0 +1,690 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Test-only indexed AdaLN fixtures and comparisons shared by L1/L2/L3.""" + +import os + +import pytest +import torch + +from rl_engine.contracts.numerical import load_contract, resolve_tolerance +from rl_engine.reference.norm.adaln_indexed import IndexedAdaLNFixedReference +from rl_engine.validation.operators.operator_specs import GtestAdaLNReference + + +def assert_bytes(actual, expected): + assert actual.shape == expected.shape and actual.dtype == expected.dtype + assert torch.equal( + actual.detach().cpu().contiguous().view(torch.uint8), + expected.detach().cpu().contiguous().view(torch.uint8), + ) + + +def small_case(device): + generator = torch.Generator().manual_seed(386) + # Large positive eps keeps H2 dx sensitive; samples and banks have distinct tags. + x = (torch.randn(2, 8, 2, generator=generator) * 2 + 0.3).to(device).requires_grad_() + m = (torch.randn(4, 6, generator=generator) * 0.5).to(device).requires_grad_() + index = torch.tensor( + [[0, 1, 0, 1, 1, 0, 1, 0], [1, 0, 1, 0, 0, 1, 0, 1]], dtype=torch.int32, device=device + ) + dy = torch.randn(2, 8, 2, generator=generator).to(device) + dg = torch.randn(2, 8, 2, generator=generator).to(device) + return x, m, index, dy, dg + + +def check_small_case(op, device, trace=None): + x, m, index, dy, dg = small_case(device) + y, gate = op(x, m, modulate_index=index, eps=7.0) + if trace is not None: + # Check routing before numerical comparisons: a correct fallback is still red. + assert trace["actual_backend"] == "cuda" + assert trace["arithmetic_profile"] == "adaln-indexed-fp32-tree-v1" + assert trace["kernel_fingerprint"] == "adaln_indexed_cuda_v1" + assert not trace["fallback"] + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + for reference in (IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32): + rx = x.detach().cpu().requires_grad_() + rm = m.detach().cpu().requires_grad_() + ry, rg = reference(rx, rm, modulate_index=index.cpu(), eps=7.0) + rdx, rdm = torch.autograd.grad((ry, rg), (rx, rm), (dy.cpu(), dg.cpu())) + assert torch.isfinite(rdx).all() and torch.count_nonzero(rdx) == rdx.numel() + for actual, expected, judgment in ( + (y, ry, "forward_accuracy"), + (gate, rg, "forward_accuracy"), + (dx, rdx, "gradient_accuracy"), + (dm, rdm, "gradient_accuracy"), + ): + assert actual.dtype == expected.dtype == torch.float32 + if isinstance(reference, IndexedAdaLNFixedReference): + assert_bytes(actual, expected) + else: + tolerance = resolve_tolerance( + load_contract(), judgment=judgment, op_class="reduction", dtype=actual.dtype + ) + torch.testing.assert_close( + actual.cpu(), expected, atol=tolerance.atol, rtol=tolerance.rtol + ) + + +def require_cuda(): + if torch.version.hip is not None: + pytest.skip("NVIDIA CUDA-only test; ROCm is outside this profile") + if not torch.cuda.is_available(): + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + pytest.fail("NVIDIA CUDA is required for indexed acceptance") + pytest.skip("NVIDIA CUDA required; not an indexed GPU PASS") + from rl_engine.backends.extension import _C + + # Once CUDA is available, missing/old symbols must fail rather than hide a red test. + assert hasattr(_C, "adaln_indexed_forward"), "fresh indexed CUDA extension required" + assert hasattr(_C, "adaln_indexed_backward"), "fresh indexed CUDA extension required" + + +def check_route_identity(op, device, route): + """Hand case: norm=(-3/4,3/4), rstd=1/4; distinct sample/bank tags.""" + x = torch.tensor([-3.0, 3.0], device=device).repeat(2, 8, 1).requires_grad_() + m = torch.tensor( + [ + [10, 11, 1, 0, 20, 21], + [30, 31, 2, 0, 40, 41], + [50, 51, 3, 0, 60, 61], + [70, 71, 4, 0, 80, 81], + ], + dtype=torch.float32, + device=device, + requires_grad=True, + ) + routes = {"all0": [0] * 8, "all1": [1] * 8, "alternating": [0, 1] * 4} + index = torch.tensor([routes[route]] * 2, dtype=torch.int32, device=device) + dy = torch.tensor([[[1.0, 0.0]], [[2.0, 0.0]]], device=device).expand_as(x) + dg = torch.tensor([[[3.0, 4.0]], [[6.0, 8.0]]], device=device).expand_as(x) + y, gate = op(x, m, modulate_index=index, eps=7.0) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + # These four rows are the manually evaluated bank/sample results. + y_rows = torch.tensor([[8.5, 11.75], [27.75, 31.75], [47.0, 51.75], [66.25, 71.75]]) + g_rows = torch.tensor([[20.0, 21.0], [40.0, 41.0], [60.0, 61.0], [80.0, 81.0]]) + dx_rows = torch.tensor( + [[7 / 64, -7 / 64], [21 / 64, -21 / 64], [7 / 32, -7 / 32], [35 / 64, -35 / 64]] + ) + rows = torch.arange(2)[:, None] + index.cpu() * 2 + for actual, expected in ((y, y_rows[rows]), (gate, g_rows[rows]), (dx, dx_rows[rows])): + assert_bytes(actual, expected) + counts = {"all0": [8, 8, 0, 0], "all1": [0, 0, 8, 8], "alternating": [4, 4, 4, 4]}[route] + per_token = torch.tensor( + [[1.0, 0.0, -0.75, 0.0, 3.0, 4.0], [2.0, 0.0, -1.5, 0.0, 6.0, 8.0]] * 2 + ) + expected_dm = torch.zeros(4, 6) + for row, count in enumerate(counts): + if count: + expected_dm[row] = per_token[row] * count + assert_bytes(dm, expected_dm) + + +def check_absent_vs_zero_dy(op, device, explicit_zero): + x = torch.full((2, 8, 2), float("nan"), device=device, requires_grad=True) + m = torch.zeros(4, 6, device=device, requires_grad=True) + index = torch.zeros(2, 8, dtype=torch.int32, device=device) + y, gate = op(x, m, modulate_index=index, eps=7.0) + if explicit_zero: + torch.autograd.backward((y, gate), (torch.zeros_like(y), torch.ones_like(gate))) + assert torch.isnan(x.grad).all() + assert torch.isnan(m.grad[:2, 2:4]).all() + assert_bytes(m.grad[:2, :2], torch.zeros(2, 2)) + else: + gate.backward(torch.ones_like(gate)) + assert x.grad is None + assert_bytes(m.grad[:2, :4], torch.zeros(2, 4)) + assert_bytes(m.grad[:2, 4:], torch.full((2, 2), 8.0)) + assert_bytes(m.grad[2:], torch.zeros(2, 6)) + + +def check_cancellation(op, device, order): + x = torch.tensor([-3.0, 3.0], device=device).repeat(2, 8, 1).requires_grad_() + m = torch.zeros(4, 6, device=device, requires_grad=True) + index = torch.tensor([[0, 1] * 4] * 2, dtype=torch.int32, device=device) + y, gate = op(x, m, modulate_index=index, eps=7.0) + # Lost-unit bank0: (((2**26 + 1) - 2**26) + 1) = 1 in FP32. + # Adjacent pairs instead give (2**26 + 1) + (-2**26 + 1) = 0. + middle, last = (1.0, -(2**26)) if order == "lost_unit" else (-(2**26), 1.0) + tail = 1.0 if order == "lost_unit" else 0.0 + u = torch.tensor([2**26, 4.0, middle, 5.0, last, 6.0, tail, 7.0], device=device) + u = u[:, None].expand(8, 2) + _, dm = torch.autograd.grad( + (y, gate), (x, m), (torch.stack((u, 2 * u)), torch.stack((u, 3 * u))) + ) + bank0 = [[1.0, 1.0, -0.75, 0.75, 1.0, 1.0], [2.0, 2.0, -1.5, 1.5, 3.0, 3.0]] + assert_bytes( + dm, + torch.tensor( + [*bank0, [22.0, 22.0, -16.5, 16.5, 22.0, 22.0], [44.0, 44.0, -33.0, 33.0, 66.0, 66.0]] + ), + ) + + +def check_product_rounding(op, device): + x = torch.tensor([-3.0, 3.0], device=device).repeat(2, 8, 1).requires_grad_() + m = torch.zeros(4, 6, device=device, requires_grad=True) + index = torch.zeros(2, 8, dtype=torch.int32, device=device) + y, _ = op(x, m, modulate_index=index, eps=7.0) + dy = torch.zeros_like(x) + dy[:, 0] = -1.0 + dy[:, 1] = 1 + 2**-23 + _, dm = torch.autograd.grad(y, (x, m), dy) + # norm=+/-3/4 exactly. Product (1+2^-23)*3/4 is a tie: + # .75 + 3*2^-25 rounds to .75 + 2^-23 (even mantissa). + # Adding -.75 afterwards gives 2^-23; FMA instead gives 3*2^-25. + assert_bytes(dm[:2], torch.tensor([[2**-23, 2**-23, -(2**-23), 2**-23, 0.0, 0.0]] * 2)) + assert_bytes(dm[2:], torch.zeros(2, 6)) + + +def check_nonfinite_isolation(op, device, bank, nonfinite): + x = torch.tensor([-3.0, 3.0], device=device).repeat(2, 8, 1) + x[0, :, 0] = nonfinite + x.requires_grad_() + m = torch.zeros(4, 6, device=device, requires_grad=True) + index = torch.full((2, 8), bank, dtype=torch.int32, device=device) + dy = torch.tensor([[[nonfinite, 1.0]], [[1.0, 0.0]]], device=device).expand_as(x) + dg = torch.tensor([[[nonfinite, 1.0]], [[1.0, 2.0]]], device=device).expand_as(x) + y, gate = op(x, m, modulate_index=index, eps=7.0) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + assert torch.isnan(y[0]).all() and torch.isnan(dx[0]).all() + assert_bytes(y[1], torch.tensor([[-0.75, 0.75]] * 8)) + assert_bytes(dx[1], torch.tensor([[7 / 128, -7 / 128]] * 8)) + bad = dm[bank * 2] + classify = torch.isnan if nonfinite != nonfinite else torch.isposinf + assert classify(bad[[0, 4]]).all() + assert torch.isnan(bad[2:4]).all() + assert_bytes(bad[[1, 5]], torch.tensor([8.0, 8.0])) + assert_bytes(dm[bank * 2 + 1], torch.tensor([8.0, 0.0, -6.0, 0.0, 8.0, 16.0])) + assert_bytes(dm[(1 - bank) * 2 : (1 - bank) * 2 + 2], torch.zeros(2, 6)) + + +def check_leaf_participation(op, device, frozen, outputs): + x = torch.tensor([-3.0, 3.0], device=device).repeat(2, 8, 1) + x.requires_grad_(frozen != "x") + m = torch.zeros(4, 6, device=device, requires_grad=frozen != "modulation") + index = torch.tensor([[0, 1] * 4] * 2, dtype=torch.int32, device=device) + y, gate = op(x, m, modulate_index=index, eps=7.0) + dy = torch.tensor([1.0, 0.0], device=device).expand_as(x) + dg = torch.tensor([1.0, 2.0], device=device).expand_as(x) + if outputs == "y": + y.backward(dy) + elif outputs == "gate": + gate.backward(dg) + else: + torch.autograd.backward((y, gate), (dy, dg)) + if frozen == "x" or outputs == "gate": + assert x.grad is None + else: + assert_bytes(x.grad, torch.tensor([7 / 128, -7 / 128]).repeat(2, 8, 1)) + if frozen == "modulation": + assert m.grad is None + else: + expected = { + "y": [4.0, 0.0, -3.0, 0.0, 0.0, 0.0], + "gate": [0.0, 0.0, 0.0, 0.0, 4.0, 8.0], + "both": [4.0, 0.0, -3.0, 0.0, 4.0, 8.0], + }[outputs] + assert_bytes(m.grad, torch.tensor([expected] * 4)) + + +def check_gate_payload(op, device): + x = torch.tensor([-3.0, 3.0], device=device).repeat(2, 8, 1).requires_grad_() + words = torch.tensor( + [[0, -(2**31)], [1, 1 - 2**31], [0x7FC01234, 0x7F800000], [0x7FA01234, -0x800000]], + dtype=torch.int32, + ) + m = torch.zeros(4, 6) + m[:, 4:] = words.view(torch.float32) + m = m.to(device).requires_grad_() + index = torch.tensor([[0, 1] * 4, [1, 0] * 4], dtype=torch.int32, device=device) + _, gate = op(x, m, modulate_index=index, eps=7.0) + rows = torch.arange(2)[:, None] + index.cpu() * 2 + assert_bytes(gate, words.view(torch.float32)[rows]) + gate.backward(torch.ones_like(gate)) + assert x.grad is None + # Arithmetic dgate does not inherit payloads from forward gate values. + assert_bytes(m.grad[:, :4], torch.zeros(4, 4)) + assert_bytes(m.grad[:, 4:], torch.full((4, 2), 4.0)) + + +def check_tiny_contributions(op, device, contribution): + x = torch.tensor([-3.0, 3.0], device=device).repeat(2, 8, 1).requires_grad_() + m = torch.zeros(4, 6, device=device, requires_grad=True) + index = torch.tensor([[0, 1] * 4] * 2, dtype=torch.int32, device=device) + y, gate = op(x, m, modulate_index=index, eps=7.0) + value = -0.0 if contribution == "negative_zero" else 2**-149 + u = torch.full_like(x, value) + _, dm = torch.autograd.grad((y, gate), (x, m), (u, u)) + # Nonempty banks also start at +0: +0 + -0 = +0, without payload identity. + words = [0] * 6 if contribution == "negative_zero" else [4, 4, 4 - 2**31, 4, 4, 4] + assert_bytes(dm, torch.tensor([words] * 4, dtype=torch.int32).view(torch.float32)) + + +def check_shared_bank0(indexed_op, shared_op, device): + x, m, index, dy, dg = small_case(device) + index.zero_() + y, gate = indexed_op(x, m, modulate_index=index, eps=7.0) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + sx = x.detach().clone().requires_grad_() + sm = m.detach()[:2].clone().requires_grad_() + sy, sg = shared_op(sx, sm, eps=7.0) + sdx, sdm = torch.autograd.grad((sy, sg), (sx, sm), (dy, dg[:, :1])) + for actual, expected in ((y, sy), (dx, sdx), (dm[:2, :4], sdm[:, :4])): + assert_bytes(actual, expected) + assert_bytes(dm[2:], torch.zeros(2, 6)) + # Shared dgate is a passthrough; indexed dgate is a token sum, not compared. + rx = x.detach().cpu().requires_grad_() + rm = m.detach().cpu().requires_grad_() + ry, rg = GtestAdaLNReference().forward_fp32(rx, rm, modulate_index=index.cpu(), eps=7.0) + rdx, rdm = torch.autograd.grad((ry, rg), (rx, rm), (dy.cpu(), dg.cpu())) + for actual, expected, judgment in ( + (y, ry, "forward_accuracy"), + (dx, rdx, "gradient_accuracy"), + (dm, rdm, "gradient_accuracy"), + ): + tolerance = resolve_tolerance( + load_contract(), judgment=judgment, op_class="reduction", dtype=actual.dtype + ) + torch.testing.assert_close(actual.cpu(), expected, atol=tolerance.atol, rtol=tolerance.rtol) + + +def check_bf16_rounding(op, device): + # norm=(-.75,.75), shift=-.75, scale=2^-8: y[1]=3/1024. + # Rounding the affine product early instead gives 1/256. + x = torch.tensor([-3.0, 3.0], dtype=torch.bfloat16, device=device).repeat(1, 2, 1) + x.requires_grad_() + m = torch.zeros(2, 6, dtype=torch.bfloat16, device=device) + m[:, :2] = -0.75 + m[:, 2:4] = 2**-8 + m.requires_grad_() + index = torch.zeros(1, 2, dtype=torch.int32, device=device) + y, gate = op(x, m, modulate_index=index, eps=7.0) + # dy=(1,0) and (2^-8,0): sum is BF16 midpoint 1+2^-8 -> 1 (even). + dy = torch.tensor([[[1.0, 0.0], [2**-8, 0.0]]], dtype=torch.bfloat16, device=device) + dg = dy.clone() + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + assert_bytes(y, torch.tensor([[[-1.5, 3 / 1024]] * 2], dtype=torch.bfloat16)) + # dx0=(1+2^-8)*7/128 rounds to 225/4096; early dnorm cast gives 224/4096. + assert_bytes( + dx, + torch.tensor( + [[[225 / 4096, -225 / 4096], [225 / 1048576, -225 / 1048576]]], dtype=torch.bfloat16 + ), + ) + assert_bytes( + dm, torch.tensor([[1.0, 0.0, -0.75390625, 0.0, 1.0, 0.0], [0.0] * 6], dtype=torch.bfloat16) + ) + + +def check_bf16_gate_payload(op, device): + # Includes signed zero, subnormals, infinities, quiet/signaling NaN payloads. + words = torch.tensor( + [[0, 0x8000], [1, 0x8001], [0x7F81, 0xFFC3], [0x7F80, 0xFF80]], dtype=torch.uint16 + ) + m = torch.zeros(4, 6, dtype=torch.bfloat16) + m[:, 4:] = words.view(torch.bfloat16) + m = m.to(device).requires_grad_() + x = ( + torch.tensor([-3.0, 3.0], dtype=torch.bfloat16, device=device) + .repeat(2, 2, 1) + .requires_grad_() + ) + index = torch.tensor([[0, 1], [1, 0]], dtype=torch.int64, device=device) + _, gate = op(x, m, modulate_index=index, eps=7.0) + expected = words.view(torch.bfloat16)[torch.tensor([[0, 2], [3, 1]])] + assert_bytes(gate, expected) + gate.backward(torch.ones_like(gate)) + assert x.grad is None + assert_bytes(m.grad, torch.tensor([[0.0, 0.0, 0.0, 0.0, 1.0, 1.0]] * 4, dtype=torch.bfloat16)) + + +def check_domain_routes(op, device, dtype, index_dtype, broadcast): + x = torch.tensor([-3.0, 3.0], dtype=dtype, device=device).repeat(2, 2, 1).requires_grad_() + m = torch.tensor( + [ + [0.0, 0.0, 0.0, 0.0, 10.0, 11.0], + [1.0, 1.0, 0.0, 0.0, 20.0, 21.0], + [2.0, 2.0, 0.0, 0.0, 30.0, 31.0], + [3.0, 3.0, 0.0, 0.0, 40.0, 41.0], + ], + dtype=dtype, + device=device, + requires_grad=True, + ) + index = torch.tensor( + [[0, 1]] if broadcast else [[0, 1], [1, 0]], dtype=index_dtype, device=device + ) + y, gate = op(x, m, modulate_index=index, eps=7.0) + dx, dm = torch.autograd.grad( + (y, gate), + (x, m), + ( + torch.tensor([1.0, 0.0], dtype=dtype, device=device).expand_as(x), + torch.tensor([2.0, 3.0], dtype=dtype, device=device).expand_as(x), + ), + ) + rows = torch.tensor([[0, 2], [1, 3]] if broadcast else [[0, 2], [3, 1]]) + assert_bytes( + y, + torch.tensor([[-0.75, 0.75], [0.25, 1.75], [1.25, 2.75], [2.25, 3.75]], dtype=dtype)[rows], + ) + assert_bytes( + gate, + torch.tensor([[10.0, 11.0], [20.0, 21.0], [30.0, 31.0], [40.0, 41.0]], dtype=dtype)[rows], + ) + assert_bytes(dx, torch.tensor([7 / 128, -7 / 128], dtype=dtype).expand_as(dx.cpu())) + assert_bytes(dm, torch.tensor([[1.0, 0.0, -0.75, 0.0, 2.0, 3.0]] * 4, dtype=dtype)) + + +def check_view_edges(op, device, dtype, layout): + batch, seq = 2, 4 + if layout == "strided": + source_x = ( + torch.tensor([-3.0, 99.0, 3.0, 99.0], dtype=dtype, device=device) + .repeat(batch, seq, 1) + .requires_grad_() + ) + source_m = torch.zeros(8, 12, dtype=dtype, device=device, requires_grad=True) + x, m = source_x[..., ::2], source_m[::2, ::2] + route_storage = torch.tensor( + [[0, 9, 1, 9, 0, 9, 1, 9]] * batch, dtype=torch.int16, device=device + ) + index = route_storage[:, ::2] + dy = torch.tensor([1.0, 99.0, 0.0, 99.0], dtype=dtype, device=device).repeat(batch, seq, 1)[ + ..., ::2 + ] + dg = torch.tensor([2.0, 99.0, 3.0, 99.0], dtype=dtype, device=device).repeat(batch, seq, 1)[ + ..., ::2 + ] + else: + source_x = ( + torch.tensor([-3.0, 3.0], dtype=dtype, device=device).repeat(1, seq, 1).requires_grad_() + ) + source_m = torch.zeros(1, 6, dtype=dtype, device=device, requires_grad=True) + x, m = source_x.expand(batch, -1, -1), source_m.expand(2 * batch, -1) + index = torch.tensor([[0, 1, 0, 1]], dtype=torch.int16, device=device).expand(batch, -1) + dy = torch.tensor([1.0, 0.0], dtype=dtype, device=device).expand_as(x) + dg = torch.tensor([2.0, 3.0], dtype=dtype, device=device).expand_as(x) + y, gate = op(x, m, modulate_index=index, eps=7.0) + torch.autograd.backward((y, gate), (dy, dg)) + assert_bytes(y, torch.tensor([-0.75, 0.75], dtype=dtype).expand(batch, seq, 2)) + assert_bytes(gate, torch.zeros(batch, seq, 2, dtype=dtype)) + if layout == "strided": + expected_x = torch.tensor([7 / 128, 0.0, -7 / 128, 0.0], dtype=dtype).repeat(batch, seq, 1) + expected_m = torch.zeros(8, 12, dtype=dtype) + expected_m[::2, ::2] = torch.tensor([2.0, 0.0, -1.5, 0.0, 4.0, 6.0], dtype=dtype) + else: + expected_x = torch.tensor([7 / 64, -7 / 64], dtype=dtype).repeat(1, seq, 1) + expected_m = torch.tensor([[8.0, 0.0, -6.0, 0.0, 16.0, 24.0]], dtype=dtype) + assert source_x.is_leaf and source_m.is_leaf + assert_bytes(source_x.grad, expected_x) + assert_bytes(source_m.grad, expected_m) + + +def check_eps_arithmetic(op, device, dtype, eps): + x = torch.zeros(1, 1, 2, dtype=dtype, device=device, requires_grad=True) + m = torch.zeros(2, 6, dtype=dtype, device=device, requires_grad=True) + index = torch.zeros(1, 1, dtype=torch.int32, device=device) + y, gate = op(x, m, modulate_index=index, eps=eps) + dx, dm = torch.autograd.grad( + (y, gate), + (x, m), + ( + torch.tensor([[[1.0, 0.0]]], dtype=dtype, device=device), + torch.tensor([[[1.0, 2.0]]], dtype=dtype, device=device), + ), + ) + if eps == 0: + assert torch.isnan(y).all() and torch.isnan(dx).all() + assert torch.isnan(dm[0, 2:4]).all() + else: + # Decimal sqrt -> FP32 sqrt -> reciprocal independently pins 2^-149. + value = ( + torch.tensor([0x643504F3], dtype=torch.int32).view(torch.float32)[0] + if eps == 2**-149 + else 2**62 + ) + expected_dx = torch.stack((torch.as_tensor(value), -torch.as_tensor(value))).to(dtype) + assert_bytes(dx, expected_dx.reshape(1, 1, 2)) + assert_bytes(y, torch.zeros(1, 1, 2, dtype=dtype)) + assert_bytes(dm[0, 2:4], torch.zeros(2, dtype=dtype)) + assert_bytes(dm[0, [0, 1, 4, 5]], torch.tensor([1.0, 0.0, 1.0, 2.0], dtype=dtype)) + assert_bytes(dm[1], torch.zeros(6, dtype=dtype)) + + +def check_geometry_bytes(op, device, dtype, hidden, *, seq=3, route="mixed"): + """Compare fixed/gold and equivalent batch layouts for identical logical tokens.""" + generator = torch.Generator().manual_seed(303) + x = torch.randn(3, seq, hidden, generator=generator).to(dtype) + m = (torch.randn(6, 3 * hidden, generator=generator) / 4).to(dtype) + if seq == 3 and route == "mixed": + index = torch.tensor([[0, 1, 0], [1, 0, 1], [0, 0, 1]], dtype=torch.int32) + else: + token = torch.arange(seq) + first = (token >= seq // 2) if route == "segmented" else (token % 3 == 1) + index = torch.stack((first, ~first, first)).to(torch.int32) + dy = torch.randn(x.shape, generator=generator).to(dtype) + dg = torch.randn(x.shape, generator=generator).to(dtype) + + def execute(callable_op, target, xx, mm, ii, yy, gg): + xx = xx.to(target).detach().requires_grad_() + mm = mm.to(target).detach().requires_grad_() + y, gate = callable_op(xx, mm, modulate_index=ii.to(target), eps=7.0) + dx, dm = torch.autograd.grad((y, gate), (xx, mm), (yy.to(target), gg.to(target))) + return tuple(t.detach().cpu() for t in (y, gate, dx, dm)) + + expected = execute(IndexedAdaLNFixedReference(), "cpu", x, m, index, dy, dg) + # Gold uses FP32 leaves, so its gradient tolerance check never compares two BF16 leaves. + gold = execute( + GtestAdaLNReference().forward_fp32, + "cpu", + x.float(), + m.float(), + index, + dy.float(), + dg.float(), + ) + actual = execute(op, device, x, m, index, dy, dg) + for i, (a, e, g) in enumerate(zip(actual, expected, gold, strict=True)): + assert_bytes(a, e) + tolerance = resolve_tolerance( + load_contract(), + judgment="forward_accuracy" if i < 2 else "gradient_accuracy", + op_class="reduction", + dtype=dtype, + ) + torch.testing.assert_close(a.float(), g, atol=tolerance.atol, rtol=tolerance.rtol) + # Every logical sample moves to a different batch position; both banks move with it. + permutation = torch.tensor([2, 0, 1]) + bank_rows = torch.cat((permutation, permutation + 3)) + inverse = torch.argsort(permutation) + inverse_banks = torch.cat((inverse, inverse + 3)) + + def strided(t): + return torch.stack((t, torch.full_like(t, 99)), dim=-1).to(device)[..., 0] + + moved = execute( + op, + device, + strided(x[permutation]), + strided(m[bank_rows]), + strided(index[permutation]), + strided(dy[permutation]), + strided(dg[permutation]), + ) + for i, (a, e) in enumerate(zip(moved, actual, strict=True)): + assert_bytes(a[inverse_banks if i == 3 else inverse], e) + + +def check_bf16_ties(op, device): + x = ( + torch.tensor([-3.0, 3.0], dtype=torch.bfloat16, device=device) + .repeat(1, 2, 1) + .requires_grad_() + ) + # y: .75 + {65/256,67/256} are opposite-parity midpoints. + m = torch.tensor( + [[65 / 256, 65 / 256, 0.0, 0.0, 0.0, 0.0], [67 / 256, 67 / 256, 0.0, 0.0, 0.0, 0.0]], + dtype=torch.bfloat16, + device=device, + ) + m.requires_grad_() + index = torch.tensor([[0, 1]], dtype=torch.int32, device=device) + y, _ = op(x, m, modulate_index=index, eps=7.0) + assert_bytes(y[..., 1], torch.tensor([[1.0, 65 / 64]], dtype=torch.bfloat16)) + # dx = (1+scale)*7/128: 455/8192 -> 228/4096 (up), + # 469/8192 -> 234/4096 (down), both ties-to-even. + with torch.no_grad(): + m[0, 2] = 1 / 64 + m[1, 2] = 3 / 64 + y, _ = op(x, m, modulate_index=index, eps=7.0) + dx, _ = torch.autograd.grad( + y, (x, m), torch.tensor([1.0, 0.0], dtype=torch.bfloat16, device=device).expand_as(x) + ) + assert_bytes( + dx, + torch.tensor( + [[[228 / 4096, -228 / 4096], [234 / 4096, -234 / 4096]]], dtype=torch.bfloat16 + ), + ) + + +def check_bf16_parameter_fold_final_cast(op, device): + x = ( + torch.tensor([-3.0, 3.0], dtype=torch.bfloat16, device=device) + .repeat(1, 3, 1) + .requires_grad_() + ) + m = torch.zeros(2, 6, dtype=torch.bfloat16, device=device, requires_grad=True) + index = torch.zeros(1, 3, dtype=torch.int32, device=device) + y, gate = op(x, m, modulate_index=index, eps=7.0) + dy = torch.tensor( + [[[1.0, 1.0], [1 / 512, 1 / 256], [1 / 512, 1 / 256]]], dtype=torch.bfloat16, device=device + ) + dg = torch.tensor( + [[[1.0, 1.0], [1 / 256, 1 / 256], [1 / 256, 1 / 256]]], dtype=torch.bfloat16, device=device + ) + _, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + # dshift[1]/dgate: 1+1/256+1/256 ->129/128 (per-add BF16 gives 1). + # dscale[0]: -.75-3/2048-3/2048 ->-193/256 (per-add gives -.75). + assert_bytes( + dm, + torch.tensor( + [[1.0, 129 / 128, -193 / 256, 194 / 256, 129 / 128, 129 / 128], [0.0] * 6], + dtype=torch.bfloat16, + ), + ) + + +def require_triton_gpu(): + if torch.version.hip is not None or not torch.cuda.is_available(): + pytest.skip("NVIDIA Triton GPU required; not an indexed PASS") + pytest.importorskip("triton") + + +def resolve_qualified_indexed(backend, *, hidden=2, indexed=True): + """Acceptance uses the delivery API; this explicit opt-in is not evidence.""" + from rl_engine.runtime.adaln import IndexedAdaLNQualification + from rl_engine.runtime.registry import KernelRegistry + + if backend == "cuda": + require_cuda() + triton_version = None + else: + require_triton_gpu() + import triton + + from rl_engine.backends.extension import _EXT_AVAILABLE + + if os.environ.get("RLK_ADALN_TRITON_ISOLATED") != "1": + pytest.skip("Triton L3 requires a real extension-free source package") + assert not _EXT_AVAILABLE + triton_version = triton.__version__ + device = torch.device("cuda", torch.cuda.current_device()) + scope = IndexedAdaLNQualification( + backend=backend, + device=str(device), + device_name=torch.cuda.get_device_name(device), + capability=torch.cuda.get_device_capability(device), + torch_version=torch.__version__, + cuda_version=torch.version.cuda, + triton_version=triton_version, + ) + op, trace = KernelRegistry().get_adaln_modulation_op( + device, + hidden=hidden, + strict=True, + indexed_qualification=scope, + modulate_index=torch.zeros(1, 8, dtype=torch.int32) if indexed else None, + ) + assert trace["selected_backend"] == f"{backend.upper()}_ADALN_MODULATION" + return op, trace + + +def check_maximum_batch(op, dtype): + b = 65535 + x = torch.tensor([-3.0, 3.0], dtype=dtype, device="cuda").repeat(b, 2, 1).requires_grad_() + m = torch.zeros(2 * b, 6, dtype=dtype, device="cuda") + tags = (torch.arange(b, device="cuda") % 17).to(dtype) + m[:b, 4], m[b:, 4] = tags, tags + 32 + m.requires_grad_() + index = torch.tensor([[0, 1]], dtype=torch.int32, device="cuda") + y, gate = op(x, m, modulate_index=index, eps=7.0) + dx, dm = torch.autograd.grad( + (y, gate), + (x, m), + ( + torch.tensor([1.0, 0.0], dtype=dtype, device="cuda").expand_as(x), + torch.tensor([2.0, 3.0], dtype=dtype, device="cuda").expand_as(x), + ), + ) + assert_bytes(y, torch.tensor([-0.75, 0.75], dtype=dtype).expand(b, 2, 2)) + assert_bytes(gate[:, 0, 0], tags.cpu()) + assert_bytes(gate[:, 1, 0], (tags + 32).cpu()) + assert_bytes(dx, torch.tensor([7 / 128, -7 / 128], dtype=dtype).expand(b, 2, 2)) + assert_bytes(dm, torch.tensor([1.0, 0.0, -0.75, 0.0, 2.0, 3.0], dtype=dtype).expand(2 * b, 6)) + + +def check_bf16_parameter_ties(op, device): + # eps=0 gives exact norm=(-1,1), so all dmod slices see opposite-parity ties. + x = torch.tensor([-1.0, 1.0], dtype=torch.bfloat16, device=device).repeat(1, 4, 1) + x.requires_grad_() + m = torch.zeros(2, 6, dtype=x.dtype, device=device, requires_grad=True) + index = torch.tensor([[0, 1, 0, 1]], dtype=torch.int32, device=device) + y, gate = op(x, m, modulate_index=index, eps=0.0) + upstream = torch.tensor([1.0, 1.0, 1 / 256, 3 / 256], dtype=x.dtype, device=device) + upstream = upstream[None, :, None].expand_as(x) + _, dm = torch.autograd.grad((y, gate), (x, m), (upstream, upstream)) + assert_bytes( + dm, + torch.tensor( + [ + [1.0, 1.0, -1.0, 1.0, 1.0, 1.0], + [65 / 64, 65 / 64, -65 / 64, 65 / 64, 65 / 64, 65 / 64], + ], + dtype=torch.bfloat16, + ), + ) + + +def check_bf16_subnormal_parameter_fold(op, device, seq, outputs): + x = torch.tensor([-3.0, 3.0], dtype=torch.bfloat16, device=device).repeat(1, seq, 1) + x.requires_grad_() + m = torch.zeros(2, 6, dtype=x.dtype, device=device, requires_grad=True) + index = torch.zeros(1, seq, dtype=torch.int32, device=device) + y, gate = op(x, m, modulate_index=index, eps=7.0) + upstream = torch.tensor([1, 0x8001], dtype=torch.uint16).view(torch.bfloat16) + upstream = upstream.to(device).reshape(1, 1, 2).expand_as(x) + if outputs == "gate": + gate.backward(upstream) + assert x.grad is None + words = [0, 0, 0, 0, seq, 0x8000 + seq] + else: + torch.autograd.backward((y, gate), (upstream, upstream)) + # -.75 * {1,3} BF16 minimum-subnormal units rounds to {-1,-2} units. + scale_units = 1 if seq == 1 else 2 + words = [seq, 0x8000 + seq, 0x8000 + scale_units, 0x8000 + scale_units, seq, 0x8000 + seq] + expected = torch.tensor([words, [0] * 6], dtype=torch.uint16).view(torch.bfloat16) + assert_bytes(m.grad, expected) diff --git a/tests/ops/norm/test_adaln_indexed.py b/tests/ops/norm/test_adaln_indexed.py new file mode 100644 index 000000000..fbb07ea1c --- /dev/null +++ b/tests/ops/norm/test_adaln_indexed.py @@ -0,0 +1,331 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""L2: indexed CUDA/Triton accuracy, bytes and public support boundaries.""" + +import pytest +import torch + +from tests.helpers.adaln_indexed import check_small_case, require_cuda + + +@pytest.fixture(params=["cuda", "triton"]) +def indexed_op(request): + if request.param == "triton": + from tests.helpers.adaln_indexed import require_triton_gpu + + require_triton_gpu() + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + return TritonAdaLNModulationOp() + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + return CudaAdaLNModulationOp() + + +@pytest.mark.cuda_only +def test_indexed_small_fp32_direct_backend(indexed_op): + check_small_case(indexed_op, "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("route", ["all0", "all1", "alternating"]) +def test_indexed_route_identity_literal(indexed_op, route): + from tests.helpers.adaln_indexed import check_route_identity + + check_route_identity(indexed_op, "cuda", route) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("explicit_zero", [False, True]) +def test_indexed_absent_vs_zero_dy_nonfinite_literal(indexed_op, explicit_zero): + from tests.helpers.adaln_indexed import check_absent_vs_zero_dy + + check_absent_vs_zero_dy(indexed_op, "cuda", explicit_zero) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("order", ["lost_unit", "retained_unit"]) +def test_indexed_selected_token_cancellation_literal(indexed_op, order): + from tests.helpers.adaln_indexed import check_cancellation + + check_cancellation(indexed_op, "cuda", order) + + +@pytest.mark.cuda_only +def test_indexed_dscale_product_rounds_before_add_literal(indexed_op): + from tests.helpers.adaln_indexed import check_product_rounding + + check_product_rounding(indexed_op, "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("bank", [0, 1]) +@pytest.mark.parametrize("nonfinite", [float("nan"), float("inf")], ids=["nan", "inf"]) +def test_indexed_nonfinite_selected_bank_isolation_literal(indexed_op, bank, nonfinite): + from tests.helpers.adaln_indexed import check_nonfinite_isolation + + check_nonfinite_isolation(indexed_op, "cuda", bank, nonfinite) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize( + "frozen,outputs", + [ + ("none", "y"), + ("none", "gate"), + ("none", "both"), + ("x", "y"), + ("x", "gate"), + ("x", "both"), + ("modulation", "y"), + ("modulation", "both"), + ], +) +def test_indexed_whole_leaf_participation_literal(indexed_op, frozen, outputs): + from tests.helpers.adaln_indexed import check_leaf_participation + + check_leaf_participation(indexed_op, "cuda", frozen, outputs) + + +@pytest.mark.cuda_only +def test_indexed_gate_copies_original_fp32_payload_literal(indexed_op): + from tests.helpers.adaln_indexed import check_gate_payload + + check_gate_payload(indexed_op, "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("contribution", ["negative_zero", "subnormal"]) +def test_indexed_signed_zero_and_subnormal_fold_literal(indexed_op, contribution): + from tests.helpers.adaln_indexed import check_tiny_contributions + + check_tiny_contributions(indexed_op, "cuda", contribution) + + +@pytest.mark.cuda_only +def test_indexed_all0_matches_shared_arithmetic_and_independent_gold(indexed_op): + from tests.helpers.adaln_indexed import check_shared_bank0 + + check_shared_bank0(indexed_op, indexed_op, "cuda") + + +def test_indexed_two_bank_metadata_overflow_precedes_value_reads(metadata_op): + import torch + + op = metadata_op + # x fits int32, but 2*B*3*H does not. Meta .item()/copies must never run. + x = torch.empty(65535, 1, 5462, device="meta") + m = torch.empty(131070, 16386, device="meta") + index = torch.empty(65535, 1, dtype=torch.int32, device="meta") + with pytest.raises(ValueError, match="int32 indexing"): + op(x, m, modulate_index=index) + + +@pytest.mark.cuda_only +def test_indexed_bf16_final_rounding(indexed_op): + from tests.helpers.adaln_indexed import check_bf16_rounding + + check_bf16_rounding(indexed_op, "cuda") + + +@pytest.mark.cuda_only +def test_indexed_bf16_raw16_gate(indexed_op): + from tests.helpers.adaln_indexed import check_bf16_gate_payload + + check_bf16_gate_payload(indexed_op, "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize( + "index_dtype", [torch.bool, torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64] +) +@pytest.mark.parametrize("broadcast", [False, True]) +def test_indexed_full_route_domain(indexed_op, dtype, index_dtype, broadcast): + from tests.helpers.adaln_indexed import check_domain_routes + + check_domain_routes(indexed_op, "cuda", dtype, index_dtype, broadcast) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("layout", ["strided", "expanded"]) +def test_indexed_source_leaf_gradients(indexed_op, dtype, layout): + from tests.helpers.adaln_indexed import check_view_edges + + check_view_edges(indexed_op, "cuda", dtype, layout) + + +@pytest.fixture(params=["cuda", "triton"]) +def metadata_op(request): + if request.param == "triton": + pytest.importorskip("triton") + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + return TritonAdaLNModulationOp() + import os + + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + try: + return CudaAdaLNModulationOp() + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + + +@pytest.mark.parametrize( + "shape,reason", + [ + ((65536, 1, 2), "B <= 65535"), + ((1, 1, 4097), "H <= 4096"), + ((1, 2**31, 1), "int32 indexing"), + ((1, 524289, 4096), "int32 indexing"), + ], +) +def test_indexed_metadata_rejects_before_copy_or_allocation(metadata_op, shape, reason): + b, s, h = shape + x = torch.empty(shape, device="meta") + m = torch.empty(2 * b, 3 * h, device="meta") + index = torch.empty(1, s, dtype=torch.int32, device="meta") + with pytest.raises((ValueError, RuntimeError), match=reason): + metadata_op(x, m, modulate_index=index) + + +@pytest.mark.parametrize("which", ["x", "m", "index"]) +def test_indexed_metadata_rejects_strided_access_span(metadata_op, which): + x = torch.empty(2, 2, 2, device="meta") + m = torch.empty(4, 6, device="meta") + index = torch.empty(2, 2, dtype=torch.int64, device="meta") + if which == "x": + x = torch.empty_strided(x.shape, (2**31, 2, 1), device="meta") + elif which == "m": + m = torch.empty_strided(m.shape, (2**31, 1), device="meta") + else: + index = torch.empty_strided(index.shape, (2**31, 1), dtype=torch.int64, device="meta") + with pytest.raises(ValueError, match="int32 indexing"): + metadata_op(x, m, modulate_index=index) + + +@pytest.mark.parametrize("shape", [(1, 2**31 - 1, 1), (65535, 1, 4096)]) +def test_indexed_positive_metadata_bounds_reach_device_check(metadata_op, shape): + b, s, h = shape + with pytest.raises(RuntimeError, match="NVIDIA CUDA tensors"): + metadata_op( + torch.empty(shape, device="meta"), + torch.empty(2 * b, 3 * h, device="meta"), + modulate_index=torch.empty(1, s, dtype=torch.bool, device="meta"), + ) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("eps", [0.0, 2**-149, 2**-126]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_indexed_eps_arithmetic(indexed_op, eps, dtype): + from tests.helpers.adaln_indexed import check_eps_arithmetic + + check_eps_arithmetic(indexed_op, "cuda", dtype, eps) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("threads", [128, 256, 512]) +def test_cuda_indexed_maximum_batch_both_banks_nonzero_dx(dtype, threads): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_maximum_batch + + check_maximum_batch(CudaAdaLNModulationOp(threads=threads), dtype) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("threads", [128, 256, 512]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("hidden", [1, 2, 3, 127, 128, 129, 255, 256, 257, 4095, 4096]) +def test_cuda_indexed_threads_geometry_batch_layout_bytes(threads, dtype, hidden): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_geometry_bytes + + check_geometry_bytes(CudaAdaLNModulationOp(threads=threads), "cuda", dtype, hidden) + + +@pytest.mark.cuda_only +def test_indexed_bf16_output_and_dx_ties(indexed_op): + from tests.helpers.adaln_indexed import check_bf16_ties + + check_bf16_ties(indexed_op, "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("invalid", ["value", "shape", "dtype", "eps"]) +def test_indexed_direct_rejects_invalid_inputs(indexed_op, invalid): + x = torch.zeros(1, 2, 3, device="cuda") + m = torch.zeros(2, 9, device="cuda") + index = torch.zeros(1, 2, device="cuda", dtype=torch.int32) + eps = 7.0 + error, message = ValueError, "values must be 0 or 1" + if invalid == "value": + index.fill_(2) + elif invalid == "shape": + index = index[:, :1] + message = "must have shape" + elif invalid == "dtype": + index = index.float() + error, message = TypeError, "boolean or integer" + else: + eps, message = float("nan"), "finite and nonnegative" + with pytest.raises(error, match=message): + indexed_op(x, m, modulate_index=index, eps=eps) + + +def test_indexed_int32_maximum_access_offset_metadata(metadata_op): + # The pointer is at the view start; the last accessed element offset is INT_MAX. + x = torch.empty_strided((2, 1, 2), (2**31 - 2, 2, 1), device="meta") + m = torch.empty(4, 6, device="meta") + index = torch.empty(1, 1, dtype=torch.int32, device="meta") + with pytest.raises(RuntimeError, match="NVIDIA CUDA tensors"): + metadata_op(x, m, modulate_index=index) + + +@pytest.mark.cuda_only +def test_indexed_bf16_parameter_fold_final_cast(indexed_op): + from tests.helpers.adaln_indexed import check_bf16_parameter_fold_final_cast + + check_bf16_parameter_fold_final_cast(indexed_op, "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("output", ["y", "gate"]) +def test_indexed_large_upstream_offset_rejected_before_copy(output, metadata_op): + import os + + if os.environ.get("RLK_ADALN_LARGE_OFFSETS") != "1": + pytest.skip( + "opt-in: upstream storage requires over 4 GiB; metadata alone cannot test backward" + ) + # The backend fixture already checks its own dependencies; CUDA must not + # acquire a Triton requirement when sharing this public backward test. + if torch.version.hip is not None or not torch.cuda.is_available(): + pytest.skip("NVIDIA CUDA required; not a large-offset GPU PASS") + free, _ = torch.cuda.mem_get_info() + assert free > 5 * 2**30, "large-offset acceptance needs >5 GiB free GPU memory" + x = torch.tensor([[[-3.0, 3.0]]], device="cuda", dtype=torch.bfloat16, requires_grad=True) + m = torch.zeros(2, 6, device="cuda", dtype=torch.bfloat16, requires_grad=True) + index = torch.zeros(1, 1, device="cuda", dtype=torch.int32) + y, gate = metadata_op(x, m, modulate_index=index, eps=7.0) + # Two logical elements, but final element address exceeds signed int32. + # Initialize only accessed positions; never fill this entire 4+ GiB storage. + upstream = torch.empty_strided((1, 1, 2), (0, 0, 2**31), device="cuda", dtype=x.dtype) + upstream[0, 0, 0], upstream[0, 0, 1] = 1.0, 0.0 + with pytest.raises(ValueError, match="upstream requires int32 indexing"): + torch.autograd.grad(y if output == "y" else gate, (x, m), upstream) + + +@pytest.mark.cuda_only +def test_indexed_bf16_all_parameter_ties(indexed_op): + from tests.helpers.adaln_indexed import check_bf16_parameter_ties + + check_bf16_parameter_ties(indexed_op, "cuda") diff --git a/tests/ops/norm/test_adaln_indexed_integration.py b/tests/ops/norm/test_adaln_indexed_integration.py new file mode 100644 index 000000000..a7dfb8cf7 --- /dev/null +++ b/tests/ops/norm/test_adaln_indexed_integration.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Prepared L2 integration; backend equality supplements independent oracles.""" + +import os + +import pytest +import torch + +from tests.helpers.adaln_indexed import ( + assert_bytes, + check_geometry_bytes, + check_small_case, + require_cuda, + require_triton_gpu, + small_case, +) + +pytestmark = pytest.mark.cuda_only + + +def test_cuda_triton_indexed_small_bytes_and_independent_oracles(): + require_cuda() + require_triton_gpu() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + results = [] + for op in (CudaAdaLNModulationOp(), TritonAdaLNModulationOp()): + # Each backend must pass fixed bytes AND independent FP32 accuracy first. + check_small_case(op, "cuda") + x, m, index, dy, dg = small_case("cuda") + y, gate = op(x, m, modulate_index=index, eps=7.0) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + results.append((y, gate, dx, dm)) + for cuda, triton in zip(*results, strict=True): + assert_bytes(cuda, triton) + + +@pytest.fixture(params=["cuda", "triton"]) +def direct(request): + if request.param == "cuda": + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + return CudaAdaLNModulationOp() + require_triton_gpu() + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + return TritonAdaLNModulationOp() + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_indexed_direct_integration_fixed_bytes_and_accuracy(direct, dtype): + check_geometry_bytes(direct, "cuda", dtype, 3) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("seq", [4096, 6889, 6032]) +@pytest.mark.parametrize("route", ["segmented", "dispersed"]) +def test_indexed_direct_target_h3072(direct, dtype, seq, route): + if os.environ.get("RLK_ADALN_REAL_SHAPES") != "1": + pytest.skip("user-run workload: >=8 GiB free GPU / >=16 GiB available host") + check_geometry_bytes(direct, "cuda", dtype, 3072, seq=seq, route=route) diff --git a/tests/ops/norm/test_adaln_indexed_triton.py b/tests/ops/norm/test_adaln_indexed_triton.py new file mode 100644 index 000000000..a9457b47f --- /dev/null +++ b/tests/ops/norm/test_adaln_indexed_triton.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""L2: indexed NVIDIA Triton semantics, independent accuracy and fixed bytes.""" + +import pytest +import torch + +from tests.helpers.adaln_indexed import require_triton_gpu, small_case + + +@pytest.fixture +def triton_op(): + pytest.importorskip("triton") + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + return TritonAdaLNModulationOp() + + +def test_triton_indexed_cpu_rejected_through_public_callable(triton_op): + x, m, index, _, _ = small_case("cpu") + with pytest.raises(RuntimeError, match="NVIDIA CUDA tensors"): + triton_op(x, m, modulate_index=index, eps=7.0) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("num_warps", [4, 8]) +@pytest.mark.parametrize("reduction_tile", [64, 128, 256]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("hidden", [1, 2, 3, 127, 128, 129, 255, 256, 257, 3072, 4095, 4096]) +def test_triton_indexed_launch_batch_layout_bytes(num_warps, reduction_tile, dtype, hidden): + require_triton_gpu() + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + from tests.helpers.adaln_indexed import check_geometry_bytes + + check_geometry_bytes( + TritonAdaLNModulationOp(num_warps=num_warps, reduction_tile=reduction_tile), + "cuda", + dtype, + hidden, + ) + + +@pytest.mark.parametrize("num_warps,reduction_tile", [(1, 128), (4, 32), (8, 512)]) +def test_triton_rejects_undeclared_launch_settings(num_warps, reduction_tile): + pytest.importorskip("triton") + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + with pytest.raises(ValueError, match="num_warps"): + TritonAdaLNModulationOp(num_warps=num_warps, reduction_tile=reduction_tile) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("num_warps", [4, 8]) +@pytest.mark.parametrize("reduction_tile", [64, 128, 256]) +def test_triton_indexed_maximum_batch_both_banks(num_warps, reduction_tile, dtype): + require_triton_gpu() + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + from tests.helpers.adaln_indexed import check_maximum_batch + + check_maximum_batch( + TritonAdaLNModulationOp(num_warps=num_warps, reduction_tile=reduction_tile), dtype + ) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("seq", [4096, 6889, 6032]) +@pytest.mark.parametrize("route", ["segmented", "dispersed"]) +def test_triton_indexed_target_h3072_workload(triton_op, dtype, seq, route): + import os + + if os.environ.get("RLK_ADALN_REAL_SHAPES") != "1": + pytest.skip("opt-in target workload; requires >=8 GiB free GPU and >=16 GiB host memory") + require_triton_gpu() + from tests.helpers.adaln_indexed import check_geometry_bytes + + free, _ = torch.cuda.mem_get_info() + assert free >= 8 * 2**30, "target comparison requires >=8 GiB free GPU memory" + check_geometry_bytes(triton_op, "cuda", dtype, 3072, seq=seq, route=route) + + +@pytest.mark.parametrize("eps", [1e300, 2**-150]) +def test_triton_indexed_direct_rejects_legacy_valid_eps_before_copy(triton_op, eps): + x, m, index, _, _ = small_case("cpu") + with pytest.raises(RuntimeError, match="finite FP32 eps"): + triton_op(x, m, modulate_index=index, eps=eps) + + +@pytest.mark.parametrize( + "invalid", + [ + "shape", + "dtype", + "device", + "index_shape", + "index_dtype", + "negative_eps", + "inf_eps", + "nan_eps", + ], +) +def test_triton_indexed_metadata_validation_before_device_or_values(triton_op, invalid): + x, m, index, _, _ = small_case("cpu") + eps = 7.0 + error, message = ValueError, "shape" + if invalid == "shape": + m = m[:, :-1] + elif invalid == "dtype": + m = m.double() + error, message = TypeError, "share FP32 or BF16" + elif invalid == "device": + index = index.to("meta") + message = "share device" + elif invalid == "index_shape": + index = index[:, :-1] + elif invalid == "index_dtype": + index = index.float() + error, message = TypeError, "boolean or integer" + else: + eps = {"negative_eps": -1.0, "inf_eps": float("inf"), "nan_eps": float("nan")}[invalid] + message = "finite and nonnegative" + with pytest.raises(error, match=message): + triton_op(x, m, modulate_index=index, eps=eps) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("seq", [1, 3]) +@pytest.mark.parametrize("outputs", ["gate", "both"]) +def test_triton_indexed_bf16_subnormal_parameter_fold(triton_op, seq, outputs): + require_triton_gpu() + from tests.helpers.adaln_indexed import check_bf16_subnormal_parameter_fold + + check_bf16_subnormal_parameter_fold(triton_op, "cuda", seq, outputs) diff --git a/tests/ops/norm/test_adaln_modulation.py b/tests/ops/norm/test_adaln_modulation.py new file mode 100644 index 000000000..b2de38c46 --- /dev/null +++ b/tests/ops/norm/test_adaln_modulation.py @@ -0,0 +1,452 @@ +import os + +import pytest +import torch + +from rl_engine.contracts.numerical import load_contract, resolve_tolerance +from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp +from rl_engine.runtime.registry import OpBackend, kernel_registry +from rl_engine.validation.operators.operator_specs import GtestAdaLNReference + + +def _assert_bytes_equal(actual, expected): + assert actual.shape == expected.shape and actual.dtype == expected.dtype + actual_bytes = actual.detach().cpu().contiguous().view(torch.uint8) + expected_bytes = expected.detach().cpu().contiguous().view(torch.uint8) + assert torch.equal(actual_bytes, expected_bytes) + + +def _assert_accuracy(actual, expected, *, judgment="forward_accuracy"): + tolerance = resolve_tolerance( + load_contract(), judgment=judgment, op_class="reduction", dtype=actual.dtype + ) + torch.testing.assert_close( + actual.detach().cpu().float(), + expected.detach().cpu().float(), + atol=tolerance.atol, + rtol=tolerance.rtol, + ) + + +def _assert_independent_shared(x, modulation, y, gate, dy, dg): + ref_x = x.detach().cpu().float().requires_grad_() + ref_modulation = modulation.detach().cpu().float().requires_grad_() + ref_y, ref_gate = GtestAdaLNReference().forward_fp32(ref_x, ref_modulation) + torch.autograd.backward((ref_y, ref_gate), (dy.cpu().float(), dg.cpu().float())) + _assert_accuracy(y, ref_y) + _assert_bytes_equal(gate, ref_gate.to(gate.dtype)) + _assert_accuracy(x.grad, ref_x.grad, judgment="gradient_accuracy") + _assert_accuracy(modulation.grad, ref_modulation.grad, judgment="gradient_accuracy") + + +def test_adaln_modulation_forward_and_shared_backward(): + x = torch.tensor([[[1.0, 3.0], [2.0, 4.0]]], requires_grad=True) + # shift, scale, gate; the gate has its own downstream gradient. + modulation = torch.tensor([[0.5, -0.5, 1.0, 0.0, 2.0, 3.0]], requires_grad=True) + y, gate = NativeAdaLNModulationOp()(x, modulation, eps=0.0) + torch.testing.assert_close(y, torch.tensor([[[-1.5, 0.5], [-1.5, 0.5]]])) + torch.testing.assert_close(gate, torch.tensor([[[2.0, 3.0]]])) + + (y * torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])).sum().backward(retain_graph=True) + torch.testing.assert_close(x.grad, torch.zeros_like(x)) + torch.testing.assert_close(modulation.grad, torch.tensor([[4.0, 6.0, -4.0, 6.0, 0.0, 0.0]])) + x.grad.zero_() + modulation.grad.zero_() + gate.sum().backward() + torch.testing.assert_close(modulation.grad, torch.tensor([[0.0, 0.0, 0.0, 0.0, 1.0, 1.0]])) + + +def test_adaln_modulation_rejects_wrong_modulation_shape(): + with pytest.raises(ValueError, match="modulation must have shape"): + NativeAdaLNModulationOp()(torch.zeros(1, 2, 4), torch.zeros(1, 4)) + + +def test_adaln_modulation_matches_independent_layer_norm_autograd(): + torch.manual_seed(386) + x = torch.randn(2, 3, 7, requires_grad=True) + modulation = torch.randn(2, 21, requires_grad=True) + upstream = torch.randn_like(x) + gate_upstream = torch.randn(2, 1, 7) + y, gate = NativeAdaLNModulationOp()(x, modulation) + ((y * upstream).sum() + (gate * gate_upstream).sum()).backward() + + x_ref = x.detach().clone().requires_grad_() + modulation_ref = modulation.detach().clone().requires_grad_() + shift, scale, gate_ref = modulation_ref.chunk(3, dim=-1) + y_ref = torch.nn.functional.layer_norm(x_ref, (7,), eps=1e-6) + y_ref = y_ref * (1 + scale[:, None, :]) + shift[:, None, :] + ((y_ref * upstream).sum() + (gate_ref[:, None, :] * gate_upstream).sum()).backward() + _assert_accuracy(y, y_ref) + _assert_accuracy(x.grad, x_ref.grad, judgment="gradient_accuracy") + _assert_accuracy(modulation.grad, modulation_ref.grad, judgment="gradient_accuracy") + + +def test_adaln_modulation_batch_position_does_not_change_bytes(): + torch.manual_seed(387) + row = torch.randn(1, 5, 17) + mod = torch.randn(1, 51) + op = NativeAdaLNModulationOp() + expected = op(row, mod)[0] + for position in range(3): + batch = torch.randn(3, 5, 17) + batch_mod = torch.randn(3, 51) + batch[position] = row[0] + batch_mod[position] = mod[0] + actual = op(batch, batch_mod)[0][position] + _assert_bytes_equal(actual, expected[0]) + + +def test_adaln_modulation_bf16_cast_and_backward(): + torch.manual_seed(389) + x = torch.randn(1, 4, 9, dtype=torch.bfloat16, requires_grad=True) + modulation = torch.randn(1, 27, dtype=torch.bfloat16, requires_grad=True) + y, gate = NativeAdaLNModulationOp()(x, modulation) + assert y.dtype == gate.dtype == torch.bfloat16 + (y.float().sum() + gate.float().sum()).backward() + assert x.grad is not None and torch.isfinite(x.grad).all() + assert modulation.grad is not None and torch.isfinite(modulation.grad).all() + _assert_bytes_equal(modulation.grad[0, 18:], torch.ones(9, dtype=torch.bfloat16)) + + +def test_adaln_modulation_accepts_qwen_six_chunk_view(): + x = torch.randn(2, 3, 8) + full = torch.randn(2, 48, requires_grad=True) + modulation = full.chunk(2, dim=-1)[0] + assert not modulation.is_contiguous() + y, gate = NativeAdaLNModulationOp()(x, modulation) + (y.sum() + gate.sum()).backward() + assert full.grad is not None + _assert_bytes_equal(full.grad[:, 24:], torch.zeros_like(full.grad[:, 24:])) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_triton_adaln_modulation_forward_backward_and_batch_position(dtype): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + torch.manual_seed(388) + x = torch.randn(2, 3, 32, device="cuda", dtype=dtype, requires_grad=True) + modulation = torch.randn(2, 96, device="cuda", dtype=dtype, requires_grad=True) + upstream = torch.randn_like(x) + gate_upstream = torch.randn(2, 1, 32, device="cuda", dtype=dtype) + y, gate = TritonAdaLNModulationOp()(x, modulation) + ((y * upstream).sum() + (gate * gate_upstream).sum()).backward() + x_ref = x.detach().clone().requires_grad_() + mod_ref = modulation.detach().clone().requires_grad_() + y_ref, gate_ref = NativeAdaLNModulationOp()(x_ref, mod_ref) + ((y_ref * upstream).sum() + (gate_ref * gate_upstream).sum()).backward() + _assert_bytes_equal(y, y_ref) + _assert_bytes_equal(x.grad, x_ref.grad) + _assert_bytes_equal(modulation.grad, mod_ref.grad) + + x_new = torch.cat([torch.randn_like(x[:1]), x[:1]], dim=0).detach().requires_grad_() + mod_new = ( + torch.cat([torch.randn_like(modulation[:1]), modulation[:1]], dim=0) + .detach() + .requires_grad_() + ) + y_new, gate_new = TritonAdaLNModulationOp()(x_new, mod_new) + _assert_bytes_equal(y_new[1], y[0]) + _assert_bytes_equal(gate_new[1], gate[0]) + upstream_new = torch.zeros_like(y_new) + upstream_new[1] = upstream[0] + gate_upstream_new = torch.zeros_like(gate_new) + gate_upstream_new[1] = gate_upstream[0] + torch.autograd.backward((y_new, gate_new), (upstream_new, gate_upstream_new)) + _assert_bytes_equal(x_new.grad[1], x.grad[0]) + _assert_bytes_equal(mod_new.grad[1], modulation.grad[0]) + + +def test_triton_irregular_hidden_and_bf16_tie_rounding(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + for hidden, seq, dtype, seed in ( + (7, 3, torch.float32, 386), + (3072, 64, torch.bfloat16, 391), + ): + torch.manual_seed(seed) + x = torch.randn(1, seq, hidden, device="cuda", dtype=dtype, requires_grad=True) + modulation = torch.randn(1, 3 * hidden, device="cuda", dtype=dtype, requires_grad=True) + y, gate = TritonAdaLNModulationOp()(x, modulation) + (y.float().sum() + gate.float().sum()).backward() + + x_cpu = x.detach().cpu().requires_grad_() + modulation_cpu = modulation.detach().cpu().requires_grad_() + y_cpu, gate_cpu = NativeAdaLNModulationOp()(x_cpu, modulation_cpu) + (y_cpu.float().sum() + gate_cpu.float().sum()).backward() + _assert_bytes_equal(y.cpu(), y_cpu) + _assert_bytes_equal(x.grad.cpu(), x_cpu.grad) + _assert_bytes_equal(modulation.grad.cpu(), modulation_cpu.grad) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_adaln_gpu_padding_and_batch_size_do_not_change_valid_bytes(dtype): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + torch.manual_seed(392) + x = torch.randn(1, 3, 32, device="cuda", dtype=dtype) + modulation = torch.randn(1, 96, device="cuda", dtype=dtype) + dy = torch.randn_like(x) + dgate = torch.randn(1, 1, 32, device="cuda", dtype=dtype) + ops = [TritonAdaLNModulationOp()] + try: + ops.append(CudaAdaLNModulationOp()) + except RuntimeError: + pass + + for op in ops: + baseline_x = x.detach().requires_grad_() + baseline_modulation = modulation.detach().requires_grad_() + baseline_y, baseline_gate = op(baseline_x, baseline_modulation) + torch.autograd.backward((baseline_y, baseline_gate), (dy, dgate)) + + for batch, seq in ((1, 5), (3, 3), (3, 5)): + padded_x = torch.randn(batch, seq, 32, device="cuda", dtype=dtype) + padded_modulation = torch.randn(batch, 96, device="cuda", dtype=dtype) + padded_dy = torch.zeros_like(padded_x) + padded_dgate = torch.zeros(batch, 1, 32, device="cuda", dtype=dtype) + padded_x[-1, :3] = x[0] + padded_modulation[-1] = modulation[0] + padded_dy[-1, :3] = dy[0] + padded_dgate[-1] = dgate[0] + padded_x.requires_grad_() + padded_modulation.requires_grad_() + y, gate = op(padded_x, padded_modulation) + torch.autograd.backward((y, gate), (padded_dy, padded_dgate)) + _assert_bytes_equal(y[-1, :3], baseline_y[0]) + _assert_bytes_equal(gate[-1], baseline_gate[0]) + _assert_bytes_equal(padded_x.grad[-1, :3], baseline_x.grad[0]) + _assert_bytes_equal(padded_modulation.grad[-1], baseline_modulation.grad[0]) + + +@pytest.mark.parametrize("hidden", [32, 64, 3072]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_cuda_adaln_bit_equality_harness(hidden, dtype): + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + try: + cuda_op = CudaAdaLNModulationOp() + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + torch.manual_seed(390) + x = torch.randn(1, 3, hidden, dtype=dtype, requires_grad=True) + modulation = torch.randn(1, 3 * hidden, dtype=dtype, requires_grad=True) + upstream = torch.randn_like(x) + gate_upstream = torch.randn(1, 1, hidden, dtype=dtype) + cpu_y, cpu_gate = NativeAdaLNModulationOp()(x, modulation) + ((cpu_y * upstream).sum() + (cpu_gate * gate_upstream).sum()).backward() + + x_gpu = x.detach().cuda().requires_grad_() + mod_gpu = modulation.detach().cuda().requires_grad_() + cuda_y, cuda_gate = cuda_op(x_gpu, mod_gpu) + ((cuda_y * upstream.cuda()).sum() + (cuda_gate * gate_upstream.cuda()).sum()).backward() + _assert_bytes_equal(cuda_y.cpu(), cpu_y) + _assert_bytes_equal(cuda_gate.cpu(), cpu_gate) + _assert_bytes_equal(x_gpu.grad.cpu(), x.grad) + _assert_bytes_equal(mod_gpu.grad.cpu(), modulation.grad) + + +@pytest.mark.parametrize("seq", [4096, 6889, 6032]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.skipif( + os.environ.get("RLK_ADALN_REAL_SHAPES") != "1", + reason="set RLK_ADALN_REAL_SHAPES=1 for Qwen-Image resolution coverage", +) +def test_adaln_real_image_shapes_forward_backward(seq, dtype): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + op, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=3072) + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + assert trace["selected_backend"] == OpBackend.CUDA_ADALN_MODULATION.name + assert trace["selected_backend"] in { + OpBackend.CUDA_ADALN_MODULATION.name, + OpBackend.TRITON_ADALN_MODULATION.name, + } + torch.manual_seed(396 + seq) + x = torch.randn(1, seq, 3072, device="cuda", dtype=dtype, requires_grad=True) + modulation = torch.randn(1, 9216, device="cuda", dtype=dtype, requires_grad=True) + y, gate = op(x, modulation) + assert y.shape == x.shape and gate.shape == (1, 1, 3072) + dy, dg = torch.randn_like(y), torch.randn_like(gate) + torch.autograd.backward((y, gate), (dy, dg)) + assert torch.isfinite(x.grad).all() and torch.isfinite(modulation.grad).all() + + cpu_x = x.detach().cpu().requires_grad_() + cpu_modulation = modulation.detach().cpu().requires_grad_() + cpu_y, cpu_gate = NativeAdaLNModulationOp()(cpu_x, cpu_modulation) + torch.autograd.backward((cpu_y, cpu_gate), (dy.cpu(), dg.cpu())) + _assert_bytes_equal(y.cpu(), cpu_y) + _assert_bytes_equal(gate.cpu(), cpu_gate) + _assert_bytes_equal(x.grad.cpu(), cpu_x.grad) + _assert_bytes_equal(modulation.grad.cpu(), cpu_modulation.grad) + + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + triton_x = x.detach().clone().requires_grad_() + triton_modulation = modulation.detach().clone().requires_grad_() + triton_y, triton_gate = TritonAdaLNModulationOp()(triton_x, triton_modulation) + torch.autograd.backward((triton_y, triton_gate), (dy, dg)) + _assert_bytes_equal(triton_y, y) + _assert_bytes_equal(triton_gate, gate) + _assert_bytes_equal(triton_x.grad, x.grad) + _assert_bytes_equal(triton_modulation.grad, modulation.grad) + # Independent full-image math, in addition to the fixed-reference byte contract. + _assert_independent_shared(x, modulation, y, gate, dy, dg) + _assert_independent_shared(triton_x, triton_modulation, triton_y, triton_gate, dy, dg) + + +def test_adaln_byte_comparison_distinguishes_signed_zero(): + with pytest.raises(AssertionError): + _assert_bytes_equal(torch.tensor([0.0]), torch.tensor([-0.0])) + + +@pytest.mark.parametrize("backend", ["cuda", "triton"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_adaln_launch_geometry_preserves_bytes(backend, dtype): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + if backend == "cuda": + try: + ops = [CudaAdaLNModulationOp(threads=n) for n in (128, 256, 512)] + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + else: + ops = [ + TritonAdaLNModulationOp(num_warps=w, reduction_tile=t) + for w in (4, 8) + for t in (64, 128, 256) + ] + torch.manual_seed(393) + x = torch.randn(2, 5, 3072, dtype=dtype) + modulation = torch.randn(2, 9216, dtype=dtype) + dy = torch.randn_like(x) + dg = torch.randn(2, 1, 3072, dtype=dtype) + cpu_x = x.clone().requires_grad_() + cpu_modulation = modulation.clone().requires_grad_() + expected_y, expected_gate = NativeAdaLNModulationOp()(cpu_x, cpu_modulation) + torch.autograd.backward((expected_y, expected_gate), (dy, dg)) + for op in ops: + gpu_x = x.cuda().requires_grad_() + gpu_modulation = modulation.cuda().requires_grad_() + y, gate = op(gpu_x, gpu_modulation) + torch.autograd.backward((y, gate), (dy.cuda(), dg.cuda())) + for actual, expected in ( + (y, expected_y), + (gate, expected_gate), + (gpu_x.grad, cpu_x.grad), + (gpu_modulation.grad, cpu_modulation.grad), + ): + _assert_bytes_equal(actual, expected) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("broadcast", [False, True]) +def test_adaln_indexed_matches_independent_autograd(device, dtype, broadcast): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + torch.manual_seed(394) + x = torch.randn(2, 4, 7, device=device, dtype=dtype, requires_grad=True) + modulation = torch.randn(4, 21, device=device, dtype=dtype, requires_grad=True) + index = torch.tensor([[0, 1, 1, 0], [1, 0, 1, 0]], device=device) + if broadcast: + index = index[:1] + dy, dg = torch.randn_like(x), torch.randn_like(x) + op, trace = kernel_registry.get_adaln_modulation_op( + device=device, hidden=7, modulate_index=index + ) + y, gate = op(x, modulation, modulate_index=index) + torch.autograd.backward((y, gate), (dy, dg)) + + ref_x = x.detach().cpu().float().requires_grad_() + ref_modulation = modulation.detach().cpu().float().requires_grad_() + # Independent row indexing and CPU LayerNorm/autograd, not split-and-where. + ref_y, ref_gate = GtestAdaLNReference().forward_fp32( + ref_x, ref_modulation, modulate_index=index + ) + torch.autograd.backward((ref_y, ref_gate), (dy.cpu().float(), dg.cpu().float())) + _assert_accuracy(y, ref_y) + _assert_accuracy(x.grad, ref_x.grad, judgment="gradient_accuracy") + _assert_accuracy(modulation.grad, ref_modulation.grad, judgment="gradient_accuracy") + _assert_bytes_equal(gate, ref_gate.to(dtype)) + assert x.grad.abs().max() > 0 + accelerated = ( + device == "cuda" + and torch.version.hip is None + and trace["selected_backend"] in {"CUDA_ADALN_MODULATION", "TRITON_ADALN_MODULATION"} + ) + if accelerated: + assert trace["actual_backend"] == ( + "cuda" if trace["selected_backend"] == "CUDA_ADALN_MODULATION" else "triton" + ) + assert trace["arithmetic_profile"] == "adaln-indexed-fp32-tree-v1" + assert trace["fallback"] is False and trace["fallback_reason"] is None + assert "ascending_selected_token_S_per_bank" in trace["reduction_order"] + else: + assert trace["actual_backend"] == "pytorch" + assert trace["arithmetic_profile"] == "adaln-select01-framework-v1" + assert trace["fallback"] is True + assert trace["fallback_reason"] == "modulate_index_requires_select01" + assert "framework_select01_backward" in trace["reduction_order"] + + +@pytest.mark.parametrize( + "index, error", + [ + (torch.tensor([[0, 2]]), ValueError), + (torch.tensor([[-1, 0]]), ValueError), + (torch.tensor([[0.0, 1.0]]), TypeError), + (torch.tensor([0, 1]), ValueError), + ], +) +def test_adaln_indexed_rejects_invalid_routing(index, error): + with pytest.raises(error): + NativeAdaLNModulationOp()(torch.ones(1, 2, 7), torch.ones(2, 21), modulate_index=index) + + +@pytest.mark.parametrize("backend", ["pytorch", "triton", "cuda"]) +@pytest.mark.parametrize("hidden", [7, 3072]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_adaln_shared_matches_independent_autograd(backend, hidden, dtype): + if backend != "pytorch" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + if backend == "cuda": + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + try: + op = CudaAdaLNModulationOp() + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + elif backend == "triton": + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + op = TritonAdaLNModulationOp() + else: + op = NativeAdaLNModulationOp() + device = "cpu" if backend == "pytorch" else "cuda" + torch.manual_seed(395) + x = torch.randn(2, 3, hidden, device=device, dtype=dtype, requires_grad=True) + modulation = torch.randn(2, 3 * hidden, device=device, dtype=dtype, requires_grad=True) + dy = torch.randn_like(x) + dg = torch.randn(2, 1, hidden, device=device, dtype=dtype) + y, gate = op(x, modulation) + torch.autograd.backward((y, gate), (dy, dg)) + + _assert_independent_shared(x, modulation, y, gate, dy, dg) diff --git a/tests/ops/norm/test_adaln_triton_device.py b/tests/ops/norm/test_adaln_triton_device.py new file mode 100644 index 000000000..4b7c0da07 --- /dev/null +++ b/tests/ops/norm/test_adaln_triton_device.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Public Triton callable regression for caller/target device and stream ownership.""" + +import pytest +import torch + +from rl_engine.contracts.numerical import load_contract, resolve_tolerance +from rl_engine.validation.operators.operator_specs import GtestAdaLNReference +from tests.helpers.adaln_indexed import assert_bytes, small_case + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("indexed", [False, True], ids=["shared-559", "indexed-621"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_triton_forward_on_other_device_preserves_caller_and_target_stream(indexed, dtype): + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + pytest.skip("two NVIDIA CUDA GPUs required; not a multi-GPU PASS") + pytest.importorskip("triton") + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + original_device = torch.cuda.current_device() + original_stream = torch.cuda.current_stream(1) + stream = torch.cuda.Stream(device=1) + try: + torch.cuda.set_device(0) + torch.cuda.set_stream(stream) + torch.cuda.set_device(0) + x, m, index, dy, dg = small_case("cpu") + if not indexed: + m, dg = m[:2].detach().requires_grad_(), dg[:, :1] + rx, rm = ( + x.to(dtype).float().detach().requires_grad_(), + m.to(dtype).float().detach().requires_grad_(), + ) + kwargs = {"modulate_index": index} if indexed else {} + ry, rg = GtestAdaLNReference().forward_fp32(rx.float(), rm.float(), eps=7.0, **kwargs) + rdx, rdm = torch.autograd.grad( + (ry, rg), (rx, rm), (dy.to(dtype).float(), dg.to(dtype).float()) + ) + # Enqueue input copies and kernel work on the target's non-default stream. + x = x.to(device="cuda:1", dtype=dtype).detach().requires_grad_() + m = m.to(device="cuda:1", dtype=dtype).detach().requires_grad_() + kwargs = {"modulate_index": index.to("cuda:1")} if indexed else {} + y, gate = TritonAdaLNModulationOp()(x, m, eps=7.0, **kwargs) + assert torch.cuda.current_device() == 0 + assert torch.cuda.current_stream(1) == stream + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy.to(y), dg.to(gate))) + assert torch.cuda.current_device() == 0 + stream.synchronize() + for actual, expected, judgment in ( + (y, ry, "forward_accuracy"), + (gate, rg, "forward_accuracy"), + (dx, rdx, "gradient_accuracy"), + (dm, rdm, "gradient_accuracy"), + ): + tolerance = resolve_tolerance( + load_contract(), judgment=judgment, op_class="reduction", dtype=dtype + ) + torch.testing.assert_close( + actual.cpu().float(), expected, atol=tolerance.atol, rtol=tolerance.rtol + ) + assert_bytes(gate, rg.to(dtype)) + finally: + torch.cuda.set_stream(original_stream) + torch.cuda.set_device(original_device) diff --git a/tests/reference/test_adaln_indexed.py b/tests/reference/test_adaln_indexed.py new file mode 100644 index 000000000..6010c440b --- /dev/null +++ b/tests/reference/test_adaln_indexed.py @@ -0,0 +1,285 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""L1: independent indexed accuracy gold and fixed-byte reference checks.""" + +import pytest +import torch + +from rl_engine.reference.norm.adaln_indexed import IndexedAdaLNFixedReference +from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp +from rl_engine.validation.operators.operator_specs import GtestAdaLNReference +from tests.helpers.adaln_indexed import assert_bytes, check_small_case + + +@pytest.fixture( + params=[IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32], + ids=["fixed", "gold"], +) +def reference(request): + return request.param + + +def test_indexed_references_match_hand_calculated_nonzero_vjp(): + # Mean=1, variance+eps=16, norm=[-3/4, 3/4]. No H2 zero-dx degeneration. + for op in (IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32): + x = torch.tensor([[[-2.0, 4.0], [-2.0, 4.0]]], requires_grad=True) + m = torch.tensor( + [[0.5, -0.5, 1.0, 0.0, 2.0, 3.0], [-1.0, 1.0, 0.0, 1.0, 4.0, 5.0]], requires_grad=True + ) + index = torch.tensor([[0, 1]], dtype=torch.int32) + y, gate = op(x, m, modulate_index=index, eps=7.0) + dx, dm = torch.autograd.grad( + (y, gate), + (x, m), + (torch.tensor([[[1.0, 0.0], [0.0, 2.0]]]), torch.tensor([[[2.0, 3.0], [4.0, 5.0]]])), + ) + assert_bytes(y, torch.tensor([[[-1.0, 0.25], [-1.75, 2.5]]])) + assert_bytes(gate, torch.tensor([[[2.0, 3.0], [4.0, 5.0]]])) + assert_bytes(dx, torch.tensor([[[7 / 64, -7 / 64], [-7 / 32, 7 / 32]]])) + assert_bytes( + dm, torch.tensor([[1.0, 0.0, -0.75, 0.0, 2.0, 3.0], [0.0, 2.0, 0.0, 1.5, 4.0, 5.0]]) + ) + assert torch.count_nonzero(dx) == dx.numel() + + +def test_indexed_byte_reference_matches_independent_small_fp32_case(): + check_small_case(IndexedAdaLNFixedReference(), "cpu") + + +@pytest.mark.parametrize( + "eps,supported", [(0.0, True), (2**-149, True), (2**-150, False), (1e300, False)] +) +def test_shared_fp32_eps_boundaries_preserve_wider_indexed_fallback(eps, supported): + op = NativeAdaLNModulationOp() + x = torch.tensor([[[-2.0, 4.0]]]) + if supported: + y, _ = op(x, torch.zeros(1, 6), eps=eps) + assert torch.isfinite(y).all() + else: + with pytest.raises(ValueError, match="finite FP32 without positive underflow"): + op(x, torch.zeros(1, 6), eps=eps) + y, _ = op(x, torch.zeros(2, 6), modulate_index=torch.zeros(1, 1, dtype=torch.int32), eps=eps) + assert torch.isfinite(y).all() + + +@pytest.mark.parametrize("broadcast", [False, True]) +def test_indexed_byte_reference_preserves_selected_token_fold_and_sample_routing(broadcast): + x = torch.tensor([-3.0, 3.0]).repeat(2, 6, 1).requires_grad_() + m = torch.zeros(4, 6, requires_grad=True) + index = torch.tensor([[0, 1, 0, 1, 0, 1]], dtype=torch.int32) + if not broadcast: + index = index.expand(2, -1).clone() + y, gate = IndexedAdaLNFixedReference()(x, m, modulate_index=index, eps=7.0) + # In FP32, 2**26 + 1 - 2**26 is 0 in the declared left-fold order. + upstream = torch.tensor([2**26, 4.0, 1.0, 5.0, -(2**26), 6.0])[:, None].expand(-1, 2) + _, dm = torch.autograd.grad( + (y, gate), + (x, m), + (torch.stack((upstream, 2 * upstream)), torch.stack((upstream, 3 * upstream))), + ) + assert_bytes( + dm, + torch.tensor( + [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + [15.0, 15.0, -11.25, 11.25, 15.0, 15.0], + [30.0, 30.0, -22.5, 22.5, 45.0, 45.0], + ] + ), + ) + + +@pytest.mark.parametrize("route", ["all0", "all1", "alternating"]) +def test_indexed_route_identity_literal(reference, route): + from tests.helpers.adaln_indexed import check_route_identity + + check_route_identity(reference, "cpu", route) + + +@pytest.mark.parametrize("explicit_zero", [False, True]) +def test_indexed_absent_vs_zero_dy_nonfinite_literal(reference, explicit_zero): + from tests.helpers.adaln_indexed import check_absent_vs_zero_dy + + check_absent_vs_zero_dy(reference, "cpu", explicit_zero) + + +@pytest.mark.parametrize("order", ["lost_unit", "retained_unit"]) +def test_indexed_selected_token_cancellation_literal(reference, order): + from tests.helpers.adaln_indexed import check_cancellation + + check_cancellation(reference, "cpu", order) + + +def test_indexed_dscale_product_rounds_before_add_literal(reference): + from tests.helpers.adaln_indexed import check_product_rounding + + check_product_rounding(reference, "cpu") + + +@pytest.mark.parametrize("bank", [0, 1]) +@pytest.mark.parametrize("nonfinite", [float("nan"), float("inf")], ids=["nan", "inf"]) +def test_indexed_nonfinite_selected_bank_isolation_literal(reference, bank, nonfinite): + from tests.helpers.adaln_indexed import check_nonfinite_isolation + + check_nonfinite_isolation(reference, "cpu", bank, nonfinite) + + +@pytest.mark.parametrize( + "frozen,outputs", + [ + ("none", "y"), + ("none", "gate"), + ("none", "both"), + ("x", "y"), + ("x", "gate"), + ("x", "both"), + ("modulation", "y"), + ("modulation", "both"), + ], +) +def test_indexed_whole_leaf_participation_literal(reference, frozen, outputs): + from tests.helpers.adaln_indexed import check_leaf_participation + + check_leaf_participation(reference, "cpu", frozen, outputs) + + +def test_indexed_gate_copies_original_fp32_payload_literal(reference): + from tests.helpers.adaln_indexed import check_gate_payload + + check_gate_payload(reference, "cpu") + + +@pytest.mark.parametrize("contribution", ["negative_zero", "subnormal"]) +def test_indexed_signed_zero_and_subnormal_fold_literal(reference, contribution): + from tests.helpers.adaln_indexed import check_tiny_contributions + + check_tiny_contributions(reference, "cpu", contribution) + + +def test_indexed_all0_matches_shared_arithmetic_and_independent_gold(): + from tests.helpers.adaln_indexed import check_shared_bank0 + + check_shared_bank0(IndexedAdaLNFixedReference(), NativeAdaLNModulationOp(), "cpu") + + +def test_indexed_bf16_reference_casts_once_at_outputs_and_gradients(): + from tests.helpers.adaln_indexed import check_bf16_rounding + + check_bf16_rounding(IndexedAdaLNFixedReference(), "cpu") + + +def test_bf16_rounding_literals_reject_specific_early_casts(): + from fractions import Fraction as F + + # Exact rational arithmetic independently checks the discriminating literals. + product = F(3, 4) * (1 + F(1, 256)) + assert product - F(3, 4) == F(3, 1024) + assert F(193, 256) - F(3, 4) == F(1, 256) # early BF16 product + exact_dx = F(257, 256) * F(7, 128) + assert F(449, 8192) < exact_dx < F(451, 8192) # rounds to 225/4096 + assert F(7, 128) != F(225, 4096) # early BF16 dnorm loses scale + assert F(1) + F(1, 256) == (F(1) + F(129, 128)) / 2 # dshift/dgate tie -> even 1 + assert -F(3, 4) * (1 + F(1, 256)) == -F(771, 1024) # final dscale -> -193/256 + + +def test_indexed_bf16_reference_preserves_raw16_gate_payload(): + from tests.helpers.adaln_indexed import check_bf16_gate_payload + + check_bf16_gate_payload(IndexedAdaLNFixedReference(), "cpu") + # A numeric equality test cannot reject canonicalized NaNs; raw16 can. + payload = torch.tensor([0x7F81], dtype=torch.uint16).view(torch.bfloat16) + assert not torch.equal( + payload.view(torch.uint16), payload.float().bfloat16().view(torch.uint16) + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize( + "index_dtype", [torch.bool, torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64] +) +@pytest.mark.parametrize("broadcast", [False, True]) +def test_indexed_reference_full_route_domain(dtype, index_dtype, broadcast): + from tests.helpers.adaln_indexed import check_domain_routes + + check_domain_routes(IndexedAdaLNFixedReference(), "cpu", dtype, index_dtype, broadcast) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("layout", ["strided", "expanded"]) +def test_indexed_reference_source_leaf_gradients(dtype, layout): + from tests.helpers.adaln_indexed import check_view_edges + + check_view_edges(IndexedAdaLNFixedReference(), "cpu", dtype, layout) + + +@pytest.mark.parametrize("eps", [0.0, 2**-149, 2**-126]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_indexed_reference_eps_arithmetic(eps, dtype): + from tests.helpers.adaln_indexed import check_eps_arithmetic + + check_eps_arithmetic(IndexedAdaLNFixedReference(), "cpu", dtype, eps) + + +def test_subnormal_eps_dx_literal_independent_decimal_sqrt(): + import struct + from decimal import Decimal, localcontext + + with localcontext() as context: + context.prec = 100 + root = struct.unpack("f", struct.pack("f", float((Decimal(2) ** -149).sqrt())))[0] + reciprocal = struct.unpack("f", struct.pack("f", 1 / root))[0] + assert struct.unpack("I", struct.pack("f", reciprocal / 2))[0] == 0x643504F3 + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("hidden", [1, 3, 127, 128, 129, 4095, 4096]) +def test_indexed_reference_geometry_and_batch_layout(dtype, hidden): + from tests.helpers.adaln_indexed import check_geometry_bytes + + check_geometry_bytes(IndexedAdaLNFixedReference(), "cpu", dtype, hidden) + + +def test_indexed_reference_bf16_output_and_dx_ties(): + from fractions import Fraction as F + + from tests.helpers.adaln_indexed import check_bf16_ties + + assert F(3, 4) + F(65, 256) == (F(1) + F(129, 128)) / 2 + assert F(3, 4) + F(67, 256) == (F(129, 128) + F(65, 64)) / 2 + assert F(65, 64) * F(7, 128) == (F(227, 4096) + F(228, 4096)) / 2 + assert F(67, 64) * F(7, 128) == (F(234, 4096) + F(235, 4096)) / 2 + check_bf16_ties(IndexedAdaLNFixedReference(), "cpu") + + +def test_indexed_reference_bf16_parameter_fold_cast_counterexample(): + from fractions import Fraction as F + + from tests.helpers.adaln_indexed import check_bf16_parameter_fold_final_cast + + assert F(1) + F(1, 256) + F(1, 256) == F(129, 128) + # First per-add BF16 tie rounds to 1, and the second repeats that tie. + assert F(1) + F(1, 256) == (F(1) + F(129, 128)) / 2 + exact_scale = -F(3, 4) * (1 + F(1, 512) + F(1, 512)) + assert -F(387, 512) < exact_scale < -F(385, 512) # -> -193/256 + # Each individual tiny contribution is < half an ULP of -.75: + assert F(3, 2048) < F(1, 512) # wrong per-add cast remains -.75 + check_bf16_parameter_fold_final_cast(IndexedAdaLNFixedReference(), "cpu") + + +def test_indexed_reference_bf16_all_parameter_ties(): + from fractions import Fraction as F + + from tests.helpers.adaln_indexed import check_bf16_parameter_ties + + assert F(1) + F(1, 256) == (F(1) + F(129, 128)) / 2 + assert F(1) + F(3, 256) == (F(129, 128) + F(65, 64)) / 2 + check_bf16_parameter_ties(IndexedAdaLNFixedReference(), "cpu") + + +@pytest.mark.parametrize("seq", [1, 3]) +@pytest.mark.parametrize("outputs", ["gate", "both"]) +def test_indexed_reference_bf16_subnormal_parameter_fold(seq, outputs): + from tests.helpers.adaln_indexed import check_bf16_subnormal_parameter_fold + + check_bf16_subnormal_parameter_fold(IndexedAdaLNFixedReference(), "cpu", seq, outputs) diff --git a/tests/runtime/test_adaln_dispatch.py b/tests/runtime/test_adaln_dispatch.py new file mode 100644 index 000000000..5c7d6a400 --- /dev/null +++ b/tests/runtime/test_adaln_dispatch.py @@ -0,0 +1,325 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""AdaLN dispatch policy and observations at the public resolver seam.""" + +import os +from argparse import Namespace + +import pytest +import torch + +from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +from rl_engine.contracts.numerical import ( + BackendProvenance, + ContractResolveError, + load_contract, + resolve_tolerance, + validate_backend_provenance, +) +from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp +from rl_engine.runtime.registry import KernelRegistry, OpBackend, kernel_registry +from rl_engine.validation.operators import CandidateSpec, run_operator_suite +from rl_engine.validation.operators.operator_specs import ( + GtestAdaLNReference, + make_candidate, + make_operator_case, +) + + +@pytest.fixture +def cuda_op(request): + try: + return CudaAdaLNModulationOp(threads=getattr(request, "param", 256)) + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + + +def test_adaln_registry_trace_reports_selected_backend(): + op, trace = kernel_registry.get_adaln_modulation_op(device="cpu", hidden=4) + assert isinstance(op.implementation, NativeAdaLNModulationOp) + assert trace["selected_backend"] == OpBackend.PYTORCH_ADALN_MODULATION.name + assert trace["fallback"] is False + assert trace["accumulator_dtype"] == "fp32" + assert trace["split_k"] is False and trace["stream_k"] is False + + +def test_adaln_modulate_index_uses_explicit_native_fallback(): + x = torch.tensor([[[1.0, 3.0], [2.0, 4.0]]], requires_grad=True) + # First batch row is choice 0, second batch row is choice 1. + modulation = torch.tensor( + [ + [0.5, -0.5, 1.0, 0.0, 2.0, 3.0], + [-1.0, 1.0, 0.0, 1.0, 4.0, 5.0], + ], + requires_grad=True, + ) + modulate_index = torch.tensor([[0, 1]]) + + op, trace = kernel_registry.get_adaln_modulation_op( + device=x.device, + hidden=x.shape[-1], + modulate_index=modulate_index, + ) + y, gate = op(x, modulation, modulate_index=modulate_index, eps=0.0) + + torch.testing.assert_close(y, torch.tensor([[[-1.5, 0.5], [-2.0, 3.0]]])) + torch.testing.assert_close(gate, torch.tensor([[[2.0, 3.0], [4.0, 5.0]]])) + (y.sum() + gate.sum()).backward() + torch.testing.assert_close(x.grad, torch.zeros_like(x)) + torch.testing.assert_close( + modulation.grad, + torch.tensor([[1.0, 1.0, -1.0, 1.0, 1.0, 1.0]]).repeat(2, 1), + ) + assert isinstance(op.implementation, NativeAdaLNModulationOp) + assert trace["selected_backend"] == OpBackend.PYTORCH_ADALN_MODULATION.name + assert trace["fallback"] is True + assert trace["fallback_reason"] == "modulate_index_requires_select01" + + +def test_adaln_cuda_registry_reports_real_fallback(): + pytest.importorskip("triton") + op, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=32) + assert type(op.implementation).__name__ in {"CudaAdaLNModulationOp", "TritonAdaLNModulationOp"} + if type(op.implementation).__name__ == "TritonAdaLNModulationOp": + assert trace["fallback"] is True + assert "CUDA_ADALN_MODULATION" in trace["rejected_backends"] + else: + assert trace["fallback"] is False + + +def test_adaln_registry_skips_cuda_for_large_hidden(): + pytest.importorskip("triton") + op, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=5000) + assert type(op.implementation).__name__ == "TritonAdaLNModulationOp" + assert trace["fallback"] is True + assert "CUDA_ADALN_MODULATION: H > 4096" in trace["rejected_backends"] + + +@pytest.mark.parametrize("hidden", [4095, 4096, 4097]) +def test_cuda_shared_registry_respects_hidden_limit(cuda_op, hidden): + _, trace = kernel_registry.get_adaln_modulation_op("cuda", hidden=hidden) + if hidden <= 4096: + assert trace["selected_backend"] == OpBackend.CUDA_ADALN_MODULATION.name + assert not trace["fallback"] + else: + assert trace["selected_backend"] != OpBackend.CUDA_ADALN_MODULATION.name + assert trace["fallback"] + assert "CUDA_ADALN_MODULATION: H > 4096" in trace["rejected_backends"] + + +def test_strict_cpu_call_records_execution_without_claiming_hardware_qualification(): + op, trace = KernelRegistry().get_adaln_modulation_op("cpu", hidden=2, strict=True) + assert trace["actual_backend"] is None + x = torch.tensor([[[-3.0, 3.0]]], requires_grad=True) + m = torch.zeros(1, 6, requires_grad=True) + y, gate = op(x, m, eps=7.0) + dx, dm = torch.autograd.grad( + (y, gate), (x, m), (torch.tensor([[[1.0, 0.0]]]), torch.ones_like(gate)) + ) + torch.testing.assert_close(y, torch.tensor([[[-0.75, 0.75]]])) + torch.testing.assert_close(dx, torch.tensor([[[7 / 128, -7 / 128]]])) + torch.testing.assert_close(dm, torch.tensor([[1.0, 0.0, -0.75, 0.0, 1.0, 1.0]])) + assert trace["actual_backend"] == "pytorch" + assert trace["execution_dtype"] == "float32" + assert trace["arithmetic_profile"] == "adaln-shared-fp32-tree-v1" + assert trace["observation"] == "successful_forward" + assert trace["binary_sha256"] is None + assert trace["hardware_qualified"] is None + + +def test_gtest_records_cpu_calls_without_inventing_a_bf16_hardware_profile(): + args = Namespace( + op="adaln_modulation", + candidate="pytorch", + arch_key=None, + batch=2, + seq=3, + normalized_dim=7, + seed=386, + ) + candidate = make_candidate(args, observe_backend=True) + case = make_operator_case(args, torch.bfloat16, torch.device("cpu")) + report = run_operator_suite( + "adaln_modulation", candidates=[candidate], cases=[case], check_grad=True + ) + assert report.passed + result = report.candidates[0] + assert result.backend_provenance is None + trace = result.cases[0].backend_trace + assert trace["actual_backend"] == "pytorch" + assert trace["execution_dtype"] == "bfloat16" + assert trace["arithmetic_profile"] == "adaln-shared-fp32-tree-v1" + assert trace["profiler_observed"] is False + + +def test_common_provenance_rejects_cuda_to_triton_fallback(): + proof = BackendProvenance( + backend_profile="triton_cuda_bf16", + requested_backend="cuda", + actual_backend="triton", + execution_dtype="bfloat16", + accumulation_dtype="float32", + output_dtype="bfloat16", + reference_dtype="float32", + candidate_tf32_enabled=False, + reference_tf32_enabled=False, + ) + with pytest.raises(ContractResolveError, match="requested_backend"): + validate_backend_provenance(load_contract(), proof) + + +def test_strict_rejects_indexed_backward(): + registry = KernelRegistry() + with pytest.raises(RuntimeError, match="select01"): + registry.get_adaln_modulation_op( + "cpu", hidden=2, strict=True, modulate_index=torch.tensor([[0]]) + ) + op, _ = registry.get_adaln_modulation_op("cpu", hidden=2, strict=True) + with pytest.raises(RuntimeError, match="select01"): + op(torch.zeros(1, 1, 2), torch.zeros(2, 6), modulate_index=torch.tensor([[0]])) + + +def test_strict_call_rechecks_resolved_hidden_and_observations_are_not_stale(): + op, trace = KernelRegistry().get_adaln_modulation_op("cpu", hidden=2, strict=True) + op(torch.zeros(1, 1, 2), torch.zeros(1, 6)) + with pytest.raises(ValueError, match="hidden"): + op(torch.zeros(1, 1, 4), torch.zeros(1, 12)) + assert trace["actual_backend"] is None + + +def test_ordinary_select01_keeps_explicit_fallback_and_separate_profile(): + index = torch.tensor([[0]]) + op, trace = KernelRegistry().get_adaln_modulation_op("cpu", hidden=2, modulate_index=index) + op(torch.zeros(1, 1, 2), torch.zeros(2, 6), modulate_index=index) + assert trace["fallback"] is True + assert trace["fallback_reason"] == "modulate_index_requires_select01" + assert trace["arithmetic_profile"] == "adaln-select01-framework-v1" + assert trace["actual_backend"] == "pytorch" + + +def test_strict_refuses_unqualified_fallback_when_optimized_backends_are_unavailable(monkeypatch): + registry = KernelRegistry() + load = registry._get_or_create_backend + + def unavailable(backend): + if backend in {OpBackend.CUDA_ADALN_MODULATION, OpBackend.TRITON_ADALN_MODULATION}: + return None + return load(backend) + + monkeypatch.setattr(registry, "_get_or_create_backend", unavailable) + op, trace = registry.get_adaln_modulation_op("cuda", hidden=7) + assert trace["fallback"] and trace["actual_backend"] is None + assert isinstance(op.implementation, NativeAdaLNModulationOp) + with pytest.raises(RuntimeError, match="no declared strict replay profile"): + registry.get_adaln_modulation_op("cuda", hidden=7, strict=True) + cpu_op, cpu_trace = registry.get_adaln_modulation_op("cpu", hidden=7, strict=True) + cpu_op(torch.zeros(1, 1, 7), torch.zeros(1, 21)) + assert cpu_trace["fallback"] is False + assert trace["actual_backend"] is None # Independent resolution, no cache trace leakage. + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is not None, + reason="NVIDIA CUDA required; run explicitly on target GPU", +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("target", ["pytorch", "triton"]) +def test_gpu_fallback_forward_backward_accuracy_and_declared_replay(monkeypatch, dtype, target): + registry = KernelRegistry() + load = registry._get_or_create_backend + + def unavailable(backend): + if backend == OpBackend.CUDA_ADALN_MODULATION: + return None + if target == "pytorch" and backend == OpBackend.TRITON_ADALN_MODULATION: + return None + return load(backend) + + monkeypatch.setattr(registry, "_get_or_create_backend", unavailable) + if target == "triton": + pytest.importorskip("triton") + op, trace = registry.get_adaln_modulation_op("cuda", hidden=3072, strict=target == "triton") + torch.manual_seed(386) + x = torch.randn(2, 3, 3072, device="cuda", dtype=dtype, requires_grad=True) + m = torch.randn(2, 9216, device="cuda", dtype=dtype, requires_grad=True) + dy, dg = torch.randn_like(x), torch.randn(2, 1, 3072, device="cuda", dtype=dtype) + y, gate = op(x, m) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + assert trace["actual_backend"] == target and trace["fallback"] is True + assert trace["fallback_reason"] and trace["profiler_observed"] is False + xx, mm = x.detach().cpu().float().requires_grad_(), m.detach().cpu().float().requires_grad_() + gold_y, gold_gate = GtestAdaLNReference().forward_fp32(xx, mm) + gold_dx, gold_dm = torch.autograd.grad( + (gold_y, gold_gate), (xx, mm), (dy.cpu().float(), dg.cpu().float()) + ) + for actual, gold, judgment in ( + (y, gold_y, "forward_accuracy"), + (dx, gold_dx, "gradient_accuracy"), + (dm, gold_dm, "gradient_accuracy"), + ): + tol = resolve_tolerance( + load_contract(), judgment=judgment, op_class="reduction", dtype=dtype + ) + torch.testing.assert_close(actual.cpu().float(), gold, atol=tol.atol, rtol=tol.rtol) + torch.testing.assert_close(gate.cpu(), gold_gate.to(dtype), atol=0, rtol=0) + # Triton shares the frozen replay contract; Native GPU has only an accuracy claim. + if target == "triton": + xx, mm = x.detach().cpu().requires_grad_(), m.detach().cpu().requires_grad_() + fixed_y, fixed_gate = NativeAdaLNModulationOp()(xx, mm) + fixed_dx, fixed_dm = torch.autograd.grad( + (fixed_y, fixed_gate), (xx, mm), (dy.cpu(), dg.cpu()) + ) + for actual, expected in zip( + (y, gate, dx, dm), (fixed_y, fixed_gate, fixed_dx, fixed_dm), strict=True + ): + assert torch.equal( + actual.detach().cpu().contiguous().view(torch.uint8), + expected.detach().contiguous().view(torch.uint8), + ) + else: + with pytest.raises(RuntimeError, match="no declared strict replay profile"): + registry.get_adaln_modulation_op("cuda", hidden=3072, strict=True) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is not None, + reason="NVIDIA CUDA required; run explicitly on target GPU", +) +def test_gpu_gtest_persists_observed_bf16_fallback_trace(monkeypatch): + pytest.importorskip("triton") + args = Namespace( + op="adaln_modulation", + candidate="triton", + arch_key=None, + batch=2, + seq=3, + normalized_dim=7, + seed=386, + ) + case = make_operator_case(args, torch.bfloat16, torch.device("cuda")) + registry = KernelRegistry() + load = registry._get_or_create_backend + monkeypatch.setattr( + registry, + "_get_or_create_backend", + lambda backend: None if backend == OpBackend.CUDA_ADALN_MODULATION else load(backend), + ) + op, _ = registry.get_adaln_modulation_op("cuda", hidden=7, strict=True) + report = run_operator_suite( + "adaln_modulation", + candidates=[CandidateSpec(name="triton-adaln", backend="triton", fn=op)], + cases=[case], + check_grad=True, + ) + assert report.passed + result = report.candidates[0] + assert result.backend_provenance is None + payload = report.to_dict()["candidates"][0]["cases"][0] + assert "backend_provenance" not in payload + trace = payload["backend_trace"] + assert trace["actual_backend"] == "triton" + assert trace["requested_backend"] == "cuda" and trace["fallback"] is True + assert trace["arithmetic_profile"] == "adaln-shared-fp32-tree-v1" + assert trace["observation"] == "successful_forward" and trace["binary_sha256"] is None diff --git a/tests/runtime/test_adaln_indexed.py b/tests/runtime/test_adaln_indexed.py new file mode 100644 index 000000000..9e4cd8087 --- /dev/null +++ b/tests/runtime/test_adaln_indexed.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""L3: public indexed AdaLN resolution, fallback and execution observations.""" + +import pytest +import torch + +from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp +from rl_engine.runtime.registry import KernelRegistry, kernel_registry +from tests.helpers.adaln_indexed import assert_bytes, check_small_case, require_cuda, small_case + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_indexed_resolution_reports_backend_policy_without_execution(device): + if device == "cuda": + if torch.version.hip is not None: + pytest.skip("NVIDIA resolution policy") + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + try: + cuda = CudaAdaLNModulationOp() + except RuntimeError: + pytest.skip("compiled CUDA AdaLN extension required") + if not cuda.indexed_available: + pytest.skip("indexed CUDA extension required") + index = torch.zeros(2, 8, dtype=torch.int32) + _, trace = KernelRegistry().get_adaln_modulation_op(device, hidden=2, modulate_index=index) + accelerated = device == "cuda" + assert trace["selected_backend"] == ( + "CUDA_ADALN_MODULATION" if accelerated else "PYTORCH_ADALN_MODULATION" + ) + assert trace["fallback"] is (not accelerated) + assert trace["fallback_reason"] == (None if accelerated else "modulate_index_requires_select01") + assert trace["arithmetic_profile"] == ( + "adaln-indexed-fp32-tree-v1" if accelerated else "adaln-select01-framework-v1" + ) + assert trace["requested_backend"] == ("cuda" if accelerated else "pytorch") + assert trace["kernel_fingerprint"] == ( + "adaln_indexed_cuda_v1" if accelerated else "adaln_modulation_select01_native_v1" + ) + assert trace["reduction_order"] == ( + "pairwise_lower_upper_H; ascending_selected_token_S_per_bank" + if accelerated + else "pairwise_lower_upper_H; framework_select01_backward" + ) + assert trace["accumulator_dtype"] == "fp32" + assert not any(trace[key] for key in ("split_k", "stream_k", "tf32", "rejected_backends")) + assert trace["actual_backend"] is None and trace["observation"] is None + + +def test_indexed_small_cpu_public_dispatch_retains_explicit_fallback(): + x, m, index, dy, dg = small_case("cpu") + op, trace = kernel_registry.get_adaln_modulation_op("cpu", hidden=2, modulate_index=index) + y, gate = op(x, m, modulate_index=index, eps=7.0) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + assert torch.isfinite(dx).all() and torch.count_nonzero(dm) + assert trace["actual_backend"] == "pytorch" + assert trace["arithmetic_profile"] == "adaln-select01-framework-v1" + assert trace["fallback"] + + +@pytest.mark.cuda_only +def test_cuda_indexed_small_fp32_public_dispatch(): + require_cuda() + index = small_case("cuda")[2] + op, trace = kernel_registry.get_adaln_modulation_op("cuda", hidden=2, modulate_index=index) + check_small_case(op, "cuda", trace) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("unsupported", ["eps", "eps_underflow", "geometry"]) +def test_cuda_indexed_execution_rechecks_support_and_clears_failed_observation(unsupported): + require_cuda() + x, m, index, _, _ = small_case("cuda") + op, trace = kernel_registry.get_adaln_modulation_op("cuda", hidden=2, modulate_index=index) + op(x, m, modulate_index=index, eps=7.0) + assert trace["actual_backend"] == "cuda" + + fallback_x = x[:1].expand(65536, -1, -1) if unsupported == "geometry" else x + fallback_index = index[:1] if unsupported == "geometry" else index + fallback_m = m[:1].expand(131072, -1) if unsupported == "geometry" else m + eps = {"eps": 1e300, "eps_underflow": 2**-150, "geometry": 7.0}[unsupported] + actual = op(fallback_x, fallback_m, modulate_index=fallback_index, eps=eps) + expected = NativeAdaLNModulationOp()( + fallback_x, fallback_m, modulate_index=fallback_index, eps=eps + ) + for a, e in zip(actual, expected, strict=True): + assert_bytes(a, e) + assert trace["actual_backend"] == "pytorch" + assert trace["arithmetic_profile"] == "adaln-select01-framework-v1" + assert trace["fallback_reason"] == "indexed_cuda_domain_unsupported" + + invalid_index = torch.full_like(index, 2) + with pytest.raises(ValueError, match="values must be 0 or 1"): + op(x, m, modulate_index=invalid_index, eps=7.0) + assert trace["actual_backend"] is None + assert trace["observation"] is None + + +@pytest.mark.cuda_only +def test_cuda_indexed_resolution_can_execute_shared_without_stale_profile(): + require_cuda() + x, m, index, dy, dg = small_case("cuda") + op, trace = kernel_registry.get_adaln_modulation_op("cuda", hidden=2, modulate_index=index) + m_shared = m.detach()[:2].requires_grad_() + y, gate = op(x, m_shared, eps=7.0) + dx, dm = torch.autograd.grad((y, gate), (x, m_shared), (dy, dg[:, :1])) + rx = x.detach().cpu().requires_grad_() + rm = m_shared.detach().cpu().requires_grad_() + ry, rg = NativeAdaLNModulationOp()(rx, rm, eps=7.0) + rdx, rdm = torch.autograd.grad((ry, rg), (rx, rm), (dy.cpu(), dg[:, :1].cpu())) + for a, e in ((y, ry), (gate, rg), (dx, rdx), (dm, rdm)): + assert_bytes(a, e) + assert trace["actual_backend"] == "cuda" + assert trace["arithmetic_profile"] == "adaln-shared-fp32-tree-v1" + assert trace["kernel_fingerprint"] == "adaln_modulation_v1:CUDA_ADALN_MODULATION" + check_small_case(op, "cuda", trace) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("explicit_zero", [False, True]) +def test_cuda_indexed_public_dispatch_preserves_absent_vs_zero_dy(explicit_zero): + require_cuda() + from tests.helpers.adaln_indexed import check_absent_vs_zero_dy + + index = torch.zeros(2, 8, dtype=torch.int32, device="cuda") + op, trace = KernelRegistry().get_adaln_modulation_op("cuda", hidden=2, modulate_index=index) + check_absent_vs_zero_dy(op, "cuda", explicit_zero) + assert trace["actual_backend"] == "cuda" + assert trace["arithmetic_profile"] == "adaln-indexed-fp32-tree-v1" + assert trace["kernel_fingerprint"] == "adaln_indexed_cuda_v1" + assert not trace["fallback"] + + +@pytest.mark.parametrize("hidden", [1, 3, 127, 128, 129, 4096]) +@pytest.mark.parametrize( + "dtype", [torch.bool, torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64] +) +def test_cuda_indexed_resolution_full_domain_metadata(hidden, dtype): + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + try: + cuda = CudaAdaLNModulationOp() + except RuntimeError: + pytest.skip("compiled CUDA extension required for resolution") + if not cuda.indexed_available or torch.version.hip is not None: + pytest.skip("NVIDIA indexed symbols required") + index = torch.zeros(1, 3, dtype=dtype) + _, trace = KernelRegistry().get_adaln_modulation_op("cuda", hidden=hidden, modulate_index=index) + assert trace["selected_backend"] == "CUDA_ADALN_MODULATION" + assert trace["arithmetic_profile"] == "adaln-indexed-fp32-tree-v1" + assert trace["fallback_reason"] is None + assert trace["actual_backend"] is None + + +@pytest.mark.parametrize( + "invalid,exception,message", + [ + ("value", ValueError, "values must be 0 or 1"), + ("negative_value", ValueError, "values must be 0 or 1"), + ("shape", ValueError, "must have shape"), + ("float_index", TypeError, "boolean or integer"), + ("index_device", ValueError, "share device"), + ("modulation_device", ValueError, "share device"), + ("modulation_dtype", TypeError, "share FP32 or BF16"), + ("modulation_shape", ValueError, "shape"), + ("empty", ValueError, "nonempty"), + ("negative_eps", ValueError, "finite and nonnegative"), + ("nan_eps", ValueError, "finite and nonnegative"), + ("inf_eps", ValueError, "finite and nonnegative"), + ], +) +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.cuda_only)]) +def test_indexed_invalid_inputs_clear_successful_trace(invalid, exception, message, device): + if device == "cuda": + require_cuda() + x, m, index, _, _ = small_case(device) + op, trace = KernelRegistry().get_adaln_modulation_op(device, hidden=2, modulate_index=index) + op(x, m, modulate_index=index, eps=7.0) + assert trace["actual_backend"] == ("cuda" if device == "cuda" else "pytorch") + eps = 7.0 + if invalid in {"value", "negative_value"}: + index = torch.full_like(index, 2 if invalid == "value" else -1) + elif invalid == "shape": + index = index[:, :-1] + elif invalid == "float_index": + index = index.float() + elif invalid == "index_device": + index = index.to("meta") + elif invalid == "modulation_device": + m = m.to("meta") + elif invalid == "modulation_dtype": + m = m.double() + elif invalid == "modulation_shape": + m = m[:, :-1] + elif invalid == "empty": + x = x[:, :0] + else: + eps = {"negative_eps": -1.0, "nan_eps": float("nan"), "inf_eps": float("inf")}[invalid] + with pytest.raises(exception, match=message): + op(x, m, modulate_index=index, eps=eps) + assert trace["actual_backend"] is None and trace["observation"] is None + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("broadcast", [False, True]) +@pytest.mark.parametrize( + "index_dtype", [torch.bool, torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64] +) +def test_cuda_indexed_full_domain_actual_execution_trace(dtype, broadcast, index_dtype): + require_cuda() + from tests.helpers.adaln_indexed import check_domain_routes + + index = torch.zeros(1 if broadcast else 2, 2, dtype=index_dtype) + # Resolve using incomplete CPU metadata, execute with real CUDA inputs. + op, trace = KernelRegistry().get_adaln_modulation_op("cuda", hidden=2, modulate_index=index) + check_domain_routes(op, "cuda", dtype, index_dtype, broadcast) + assert trace["actual_backend"] == "cuda" + assert trace["arithmetic_profile"] == "adaln-indexed-fp32-tree-v1" + assert trace["fallback_reason"] is None and not trace["fallback"] + + +@pytest.mark.parametrize("eps", [0.0, 2**-149, 2**-150, 1e300]) +def test_indexed_strict_stays_closed_at_resolution_and_execution(eps): + x, m, index, _, _ = small_case("cpu") + with pytest.raises(RuntimeError, match="strict AdaLN rejects"): + KernelRegistry().get_adaln_modulation_op("cpu", hidden=2, modulate_index=index, strict=True) + op, trace = KernelRegistry().get_adaln_modulation_op("cpu", hidden=2, strict=True) + with pytest.raises(RuntimeError, match="strict AdaLN rejects"): + op(x, m, modulate_index=index, eps=eps) + assert trace["actual_backend"] is None + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("unsupported", ["H", "B", "eps_overflow", "eps_underflow"]) +def test_cuda_indexed_execution_fallback_from_shared_resolution(unsupported): + require_cuda() + b, h = (65536, 2) if unsupported == "B" else (1, 4097 if unsupported == "H" else 3) + x = torch.linspace(-1, 1, h, device="cuda").repeat(b, 1, 1).requires_grad_() + m = torch.zeros(2 * b, 3 * h, device="cuda", requires_grad=True) + index = torch.ones(1, 1, dtype=torch.int64, device="cuda") + eps = {"eps_overflow": 1e300, "eps_underflow": 2**-150}.get(unsupported, 7.0) + op, trace = KernelRegistry().get_adaln_modulation_op("cuda", hidden=3) + actual = op(x, m, modulate_index=index, eps=eps) + expected = NativeAdaLNModulationOp()(x, m, modulate_index=index, eps=eps) + for a, e in zip(actual, expected, strict=True): + assert_bytes(a, e) + assert trace["actual_backend"] == "pytorch" + assert trace["arithmetic_profile"] == "adaln-select01-framework-v1" + assert trace["fallback"] + assert trace["fallback_reason"] == "modulate_index_requires_select01" diff --git a/tests/runtime/test_adaln_indexed_activation.py b/tests/runtime/test_adaln_indexed_activation.py new file mode 100644 index 000000000..9974e07a3 --- /dev/null +++ b/tests/runtime/test_adaln_indexed_activation.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Prepared L3 GPU acceptance through the default-closed delivery mechanism.""" + +import pytest +import torch + +from tests.helpers.adaln_indexed import ( + check_absent_vs_zero_dy, + check_domain_routes, + check_leaf_participation, + check_small_case, + check_view_edges, + resolve_qualified_indexed, + small_case, +) + +pytestmark = pytest.mark.cuda_only + + +@pytest.fixture(params=["cuda", "triton"]) +def activated(request): + return resolve_qualified_indexed(request.param) + + +def test_activated_indexed_small_forward_backward_and_trace(activated): + op, trace = activated + check_small_case(op, "cuda") + assert trace["actual_backend"] in {"cuda", "triton"} + assert trace["arithmetic_profile"] == "adaln-indexed-fp32-tree-v1" + assert trace["fallback"] is False and trace["fallback_reason"] is None + assert trace["observation"] == "successful_forward" + assert trace["hardware_qualified"] is None and trace["binary_sha256"] is None + assert trace["profiler_observed"] is False + + +@pytest.mark.parametrize("backend", ["cuda", "triton"]) +@pytest.mark.parametrize("strict", [False, True]) +def test_indexed_dispatch_adds_no_route_scan_or_scalar_read(backend, strict): + from rl_engine.runtime.registry import KernelRegistry + from tests.helpers.adaln_indexed import assert_bytes + + op, trace = resolve_qualified_indexed(backend) + x, m, index, _, _ = small_case("cuda") + if not strict: + op, trace = KernelRegistry().get_adaln_modulation_op( + x.device, hidden=2, modulate_index=index + ) + # Warm compilation separately; count only the public direct/dispatch calls. + check_small_case(op, "cuda") + op.implementation(x, m, modulate_index=index, eps=7.0) + observed, outputs = [], [] + for callable_op in (op.implementation, op): + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU], acc_events=True + ) as profile: + outputs.append(callable_op(x, m, modulate_index=index, eps=7.0)) + observed.append({event.key: event.count for event in profile.key_averages()}) + for actual, expected in zip(outputs[1], outputs[0], strict=True): + assert_bytes(actual, expected) + for event in ("aten::all", "aten::item"): + # CUDA retains its native ABI check in addition to the Python check. + assert observed[0][event] > 0 + assert observed[1][event] == observed[0][event] + assert trace["actual_backend"] == backend and trace["fallback"] is False + assert trace["profiler_observed"] is False + + +@pytest.mark.parametrize("explicit_zero", [False, True]) +def test_activated_indexed_absent_vs_zero(activated, explicit_zero): + check_absent_vs_zero_dy(activated[0], "cuda", explicit_zero) + + +@pytest.mark.parametrize( + "frozen,outputs", [("x", "both"), ("modulation", "both"), (None, "gate"), (None, "y")] +) +def test_activated_indexed_leaf_participation(activated, frozen, outputs): + check_leaf_participation(activated[0], "cuda", frozen, outputs) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("layout", ["strided", "expanded"]) +def test_activated_indexed_source_leaf_gradients(activated, dtype, layout): + check_view_edges(activated[0], "cuda", dtype, layout) + + +@pytest.mark.parametrize("broadcast", [False, True]) +def test_activated_indexed_routes(activated, broadcast): + check_domain_routes(activated[0], "cuda", torch.bfloat16, torch.int16, broadcast) + + +@pytest.mark.parametrize("backend", ["cuda", "triton"]) +@pytest.mark.parametrize("resolution", ["indexed", "shared"]) +def test_activated_strict_derives_mode_on_every_call(backend, resolution): + op, trace = resolve_qualified_indexed(backend, indexed=resolution == "indexed") + x, m, index, _, _ = small_case("cuda") + for indexed in (True, False, True): + op(x, m if indexed else m[:2], eps=7.0, modulate_index=index if indexed else None) + assert trace["arithmetic_profile"] == ( + "adaln-indexed-fp32-tree-v1" if indexed else "adaln-shared-fp32-tree-v1" + ) + assert trace["actual_backend"] == backend + with pytest.raises(ValueError, match="values must be 0 or 1"): + op(x, m, modulate_index=index + 2) + assert trace["actual_backend"] is None and trace["observation"] is None + + +@pytest.mark.parametrize( + "unsupported", ["H", "changed_hidden", "B", "device", "eps_overflow", "eps_underflow"] +) +def test_activated_strict_rejects_changed_execution_domain(activated, unsupported): + op, trace = activated + check_small_case(op, "cuda") + h = {"H": 4097, "changed_hidden": 3}.get(unsupported, 2) + b = 65536 if unsupported == "B" else 2 + device = "cpu" if unsupported == "device" else "cuda" + x = torch.zeros(b, 1, h, device=device) + m = torch.zeros(2 * b, 3 * h, device=device) + index = torch.zeros(1, 1, dtype=torch.int32, device=device) + eps = {"eps_overflow": 1e300, "eps_underflow": 2**-150}.get(unsupported, 7.0) + with pytest.raises((ValueError, RuntimeError), match="strict AdaLN"): + op(x, m, modulate_index=index, eps=eps) + assert trace["actual_backend"] is None and trace["observation"] is None + + +@pytest.mark.parametrize("backend", ["cuda", "triton"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("seq", [4096, 6889, 6032]) +@pytest.mark.parametrize("route", ["segmented", "dispersed"]) +def test_activated_indexed_target_h3072(backend, dtype, seq, route): + import os + + from tests.helpers.adaln_indexed import check_geometry_bytes + + if os.environ.get("RLK_ADALN_REAL_SHAPES") != "1": + pytest.skip("user-run workload: >=8 GiB free GPU / >=16 GiB available host") + op, trace = resolve_qualified_indexed(backend, hidden=3072) + check_geometry_bytes(op, "cuda", dtype, 3072, seq=seq, route=route) + assert trace["actual_backend"] == backend and not trace["fallback"] diff --git a/tests/runtime/test_adaln_indexed_triton.py b/tests/runtime/test_adaln_indexed_triton.py new file mode 100644 index 000000000..73bdea9d3 --- /dev/null +++ b/tests/runtime/test_adaln_indexed_triton.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""L3: real missing-extension resolution and indexed Triton execution traces.""" + +import hashlib +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +from rl_engine.runtime.registry import KernelRegistry +from tests.helpers.adaln_indexed import assert_bytes, check_small_case, small_case + + +@pytest.mark.parametrize("triton_state", ["available", "absent", "incompatible"]) +@pytest.mark.parametrize( + "missing_symbol", + [ + None, + "adaln_indexed_forward", + "adaln_indexed_backward", + "adaln_modulation_forward", + "adaln_modulation_backward", + ], +) +def test_indexed_missing_extension_real_resolution(tmp_path, triton_state, missing_symbol): + triton_available = triton_state == "available" + if triton_available: + pytest.importorskip("triton") + if torch.version.hip is not None: + pytest.skip("NVIDIA resolution policy") + root = Path(__file__).resolve().parents[2] + # A distinct real package has no extension; the current _C.so is untouched. + shutil.copytree( + root / "rl_engine", + tmp_path / "rl_engine", + ignore=shutil.ignore_patterns("*.so", "*.pyc", "__pycache__"), + ) + binary = None + original_digest = None + fixture_digest = None + if missing_symbol is not None: + from rl_engine.backends.extension import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE: + pytest.skip("a real native binary is required for a missing-symbol ABI fixture") + binary = Path(_C.__file__).resolve() + original = binary.read_bytes() + original_digest = hashlib.sha256(original).hexdigest() + # Rename only one NUL-delimited pybind registration in an independent + # copy. Native kernel instructions and the current binary are untouched. + pattern = b"\0" + missing_symbol.encode() + b"\0" + assert original.count(pattern) == 1, "ABI fixture needs one exact registration string" + fixture = original.replace(pattern, b"\0x" + missing_symbol[1:].encode() + b"\0") + assert sum(a != b for a, b in zip(original, fixture, strict=True)) == 1 + (tmp_path / "rl_engine/_C.so").write_bytes(fixture) + fixture_digest = hashlib.sha256(fixture).hexdigest() + shared_cuda = missing_symbol is not None and missing_symbol.startswith("adaln_indexed") + script = ( + f"TRITON_AVAILABLE = {triton_available!r}\n" + f"MISSING_SYMBOL = {missing_symbol!r}\nSHARED_CUDA = {shared_cuda!r}\n" + f"FIXTURE_SHA256 = {fixture_digest!r}\n" + + """ +import pathlib, torch, rl_engine, importlib.util +assert (importlib.util.find_spec("triton") is not None) == TRITON_INSTALLED +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.runtime.registry import KernelRegistry +assert pathlib.Path(rl_engine.__file__).resolve().parent.parent == pathlib.Path.cwd() +assert _EXT_AVAILABLE == (MISSING_SYMBOL is not None) +if MISSING_SYMBOL is not None: + import hashlib + assert pathlib.Path(_C.__file__).resolve() == pathlib.Path.cwd()/"rl_engine/_C.so" + assert hashlib.sha256(pathlib.Path(_C.__file__).read_bytes()).hexdigest() == FIXTURE_SHA256 + for name in ("adaln_modulation_forward","adaln_modulation_backward", + "adaln_indexed_forward","adaln_indexed_backward"): + assert hasattr(_C,name) == (name != MISSING_SYMBOL) + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + if SHARED_CUDA: + assert not CudaAdaLNModulationOp().indexed_available + else: + try: + CudaAdaLNModulationOp() + except RuntimeError as exc: + assert "symbols" in str(exc) + else: + raise AssertionError("incomplete shared ABI accepted") +op, trace = KernelRegistry().get_adaln_modulation_op( + "cuda", hidden=2, modulate_index=torch.zeros(1, 8, dtype=torch.int32)) +assert type(op.implementation).__name__ == ( + "TritonAdaLNModulationOp" if TRITON_AVAILABLE else "NativeAdaLNModulationOp"), trace +assert trace["requested_backend"] == ("triton" if TRITON_AVAILABLE else "pytorch") +assert trace["arithmetic_profile"] == ( + "adaln-indexed-fp32-tree-v1" if TRITON_AVAILABLE else "adaln-select01-framework-v1") +assert trace["fallback"] is (not TRITON_AVAILABLE) +assert trace["fallback_reason"] == ( + None if TRITON_AVAILABLE else "modulate_index_requires_select01") +assert trace["actual_backend"] is None +if TRITON_AVAILABLE: + import triton + from rl_engine.runtime.adaln import IndexedAdaLNQualification + cuda_scope = IndexedAdaLNQualification( + backend="cuda", device="cuda:0", device_name="candidate target", + capability=(8,6), torch_version=torch.__version__, cuda_version=torch.version.cuda) + try: + KernelRegistry().get_adaln_modulation_op( + "cuda:0",hidden=2,modulate_index=torch.zeros(1,8,dtype=torch.int32), + strict=True,indexed_qualification=cuda_scope) + except RuntimeError: + pass + else: + raise AssertionError("CUDA scope qualified the real Triton fallback") + for version in ("unqualified-runtime", triton.__version__): + scope = IndexedAdaLNQualification( + backend="triton", device="cuda:0", device_name="candidate target", + capability=(8,6), torch_version=torch.__version__, + cuda_version=torch.version.cuda, triton_version=version) + try: + activated, activation_trace = KernelRegistry().get_adaln_modulation_op( + "cuda:0",hidden=2,modulate_index=torch.zeros(1,8,dtype=torch.int32), + strict=True,indexed_qualification=scope) + except RuntimeError: + assert version == "unqualified-runtime" + else: + assert version == triton.__version__ + assert activation_trace["arithmetic_profile"] == "adaln-indexed-fp32-tree-v1" + assert activation_trace["actual_backend"] is None +x = torch.tensor([[[-3.,3.]]], requires_grad=True) +m = torch.zeros(2,6, requires_grad=True) +index = torch.zeros(1,1,dtype=torch.int32) +y,g = op(x,m,modulate_index=index,eps=7.) +dx,dm = torch.autograd.grad((y,g),(x,m),(torch.tensor([[[1.,0.]]]),torch.ones_like(g))) +assert torch.equal(dx,torch.tensor([[[7/128,-7/128]]])) +assert torch.equal(dm,torch.tensor([[1.,0.,-.75,0.,1.,1.],[0.]*6])) +assert trace["actual_backend"] == "pytorch" +assert trace["arithmetic_profile"] == "adaln-select01-framework-v1" +assert trace["fallback"] is True +assert trace["fallback_reason"] == ( + "indexed_triton_domain_unsupported" if TRITON_AVAILABLE else "modulate_index_requires_select01") +try: + op(x,m,modulate_index=index+2,eps=7.) +except ValueError: + pass +else: + raise AssertionError("illegal route accepted") +assert trace["actual_backend"] is None and trace["observation"] is None +shared, shared_trace = KernelRegistry().get_adaln_modulation_op("cuda",hidden=2) +assert shared_trace["selected_backend"] == ( + "CUDA_ADALN_MODULATION" if SHARED_CUDA else ( + "TRITON_ADALN_MODULATION" if TRITON_AVAILABLE else "PYTORCH_ADALN_MODULATION")) +if SHARED_CUDA: + shared_strict, st = KernelRegistry().get_adaln_modulation_op("cuda",hidden=2,strict=True) + assert st["arithmetic_profile"] == "adaln-shared-fp32-tree-v1" + assert st["selected_backend"] == "CUDA_ADALN_MODULATION" +elif not TRITON_AVAILABLE: + try: + KernelRegistry().get_adaln_modulation_op("cuda",hidden=2,strict=True) + except RuntimeError: + pass + else: + raise AssertionError("shared strict accepted unqualified GPU framework") +try: + KernelRegistry().get_adaln_modulation_op("cuda",hidden=2,modulate_index=index,strict=True) +except RuntimeError: + pass +else: + raise AssertionError("indexed strict opened") +""" + ) + env = dict(os.environ, PYTHONPATH=str(tmp_path), CUDA_VISIBLE_DEVICES="") + # -S avoids the installed editable finder resolving absent submodules back + # into the original checkout. Add dependency directories without .pth hooks. + dependencies = [p for p in sys.path if "site-packages" in p] + if not triton_available: + dependency_view = tmp_path / "dependencies" + dependency_view.mkdir() + for directory in dependencies: + for entry in Path(directory).iterdir(): + if ( + entry.name.startswith(("triton", "__editable__")) + or entry.suffix == ".pth" + or (dependency_view / entry.name).exists() + ): + continue + (dependency_view / entry.name).symlink_to(entry) + dependencies = [str(dependency_view)] + if triton_state == "incompatible": + # An installed dependency missing the required language/CUDA API. + # Registry and all production code remain real and unpatched. + incompatible = tmp_path / "triton" + incompatible.mkdir() + (incompatible / "__init__.py").write_text('__version__ = "incompatible"\n') + script = f"TRITON_INSTALLED = {triton_state != 'absent'!r}\n" + script + script = "import sys; sys.path.extend(" + json.dumps(dependencies) + ")\n" + script + result = subprocess.run( + [sys.executable, "-S", "-c", script], + cwd=tmp_path, + env=env, + text=True, + capture_output=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + if binary is not None: + assert hashlib.sha256(binary.read_bytes()).hexdigest() == original_digest + + +@pytest.fixture +def resolved_triton(): + from tests.helpers.adaln_indexed import require_triton_gpu + + require_triton_gpu() + if os.environ.get("RLK_ADALN_TRITON_ISOLATED") != "1": + pytest.skip("run in the prepared extension-free package for real Triton L3 selection") + from rl_engine.backends.extension import _EXT_AVAILABLE + + assert not _EXT_AVAILABLE, "isolation must actually exclude the extension" + index = torch.zeros(1, 8, dtype=torch.int32) + op, trace = KernelRegistry().get_adaln_modulation_op("cuda", hidden=2, modulate_index=index) + assert trace["selected_backend"] == "TRITON_ADALN_MODULATION" + return op, trace + + +@pytest.mark.cuda_only +def test_triton_indexed_small_public_dispatch(resolved_triton): + op, trace = resolved_triton + check_small_case(op, "cuda") + assert trace["actual_backend"] == "triton" + assert trace["arithmetic_profile"] == "adaln-indexed-fp32-tree-v1" + assert trace["kernel_fingerprint"] == "adaln_indexed_triton_v1" + assert trace["fallback"] is False and trace["fallback_reason"] is None + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("explicit_zero", [False, True]) +def test_triton_dispatch_absent_vs_zero(resolved_triton, explicit_zero): + from tests.helpers.adaln_indexed import check_absent_vs_zero_dy + + op, trace = resolved_triton + check_absent_vs_zero_dy(op, "cuda", explicit_zero) + assert trace["actual_backend"] == "triton" + assert trace["fallback"] is False and trace["fallback_reason"] is None + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("broadcast", [False, True]) +@pytest.mark.parametrize( + "index_dtype", [torch.bool, torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64] +) +def test_triton_dispatch_route_domain(resolved_triton, dtype, broadcast, index_dtype): + from tests.helpers.adaln_indexed import check_domain_routes + + op, trace = resolved_triton + check_domain_routes(op, "cuda", dtype, index_dtype, broadcast) + assert trace["actual_backend"] == "triton" + assert trace["arithmetic_profile"] == "adaln-indexed-fp32-tree-v1" + assert trace["fallback"] is False and trace["fallback_reason"] is None + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("unsupported", ["H", "B", "eps_overflow", "eps_underflow", "cpu"]) +@pytest.mark.parametrize("resolution", ["indexed", "shared"]) +def test_triton_execution_fallback_and_failed_trace(resolved_triton, unsupported, resolution): + from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp + + op, trace = resolved_triton + if resolution == "shared": + op, trace = KernelRegistry().get_adaln_modulation_op("cuda", hidden=2) + x, m, index, _, _ = small_case("cuda") + op(x, m, modulate_index=index, eps=7.0) + assert trace["actual_backend"] == "triton" + if unsupported == "cpu": + x, m, index = x.cpu(), m.cpu(), index.cpu() + else: + b, h = (65536, 2) if unsupported == "B" else (1, 4097 if unsupported == "H" else 2) + x = torch.tensor([-3.0, 3.0], device="cuda").repeat(b, 1, (h + 1) // 2)[..., :h] + m = torch.zeros(2 * b, 3 * h, device="cuda") + index = torch.ones(1, 1, dtype=torch.int32, device="cuda") + eps = {"eps_overflow": 1e300, "eps_underflow": 2**-150}.get(unsupported, 7.0) + actual = op(x, m, modulate_index=index, eps=eps) + expected = NativeAdaLNModulationOp()(x, m, modulate_index=index, eps=eps) + for a, e in zip(actual, expected, strict=True): + assert_bytes(a, e) + assert trace["actual_backend"] == "pytorch" + assert trace["arithmetic_profile"] == "adaln-select01-framework-v1" + assert trace["fallback"] is True + assert trace["fallback_reason"] == ( + "indexed_triton_domain_unsupported" + if resolution == "indexed" + else "modulate_index_requires_select01" + ) + # A failed call after successful fallback must clear the old observation too. + with pytest.raises(ValueError, match="values must be 0 or 1"): + op(x, m, modulate_index=torch.full_like(index, 2), eps=eps) + assert trace["actual_backend"] is None and trace["observation"] is None + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("eps", [0.0, 2**-149, 2**-150, 1e300]) +def test_triton_indexed_strict_closed(resolved_triton, eps): + with pytest.raises(RuntimeError, match="strict AdaLN rejects"): + KernelRegistry().get_adaln_modulation_op( + "cuda", hidden=2, modulate_index=torch.zeros(1, 8, dtype=torch.int32), strict=True + ) + op, trace = KernelRegistry().get_adaln_modulation_op("cuda", hidden=2, strict=True) + x, m, index, _, _ = small_case("cuda") + with pytest.raises(RuntimeError, match="strict AdaLN rejects"): + op(x, m, modulate_index=index, eps=eps) + assert trace["actual_backend"] is None diff --git a/tests/runtime/test_adaln_strict_policy.py b/tests/runtime/test_adaln_strict_policy.py new file mode 100644 index 000000000..f9e430c10 --- /dev/null +++ b/tests/runtime/test_adaln_strict_policy.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""AdaLN-local strict activation at the public resolver/callable seam.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +from rl_engine.runtime.registry import KernelRegistry +from tests.helpers.adaln_indexed import small_case + + +def test_rocm_default_device_resolution_preserves_shared_and_indexed_fallback(): + root = Path(__file__).resolve().parents[2] + # Simulate only the external platform APIs before real device discovery. + script = """ +import torch +torch.version.hip = "test-ROCm-runtime" +torch.cuda.is_available = lambda: True +from rl_engine.runtime.registry import KernelRegistry +from tests.helpers.adaln_indexed import check_domain_routes, assert_bytes + +for device in (None, "cuda"): + for indexed in (False, True): + op, trace = KernelRegistry().get_adaln_modulation_op( + device, hidden=2, + modulate_index=torch.zeros(1, 2, dtype=torch.int32) if indexed else None) + assert trace["selected_backend"] == "PYTORCH_ADALN_MODULATION" + if indexed: + check_domain_routes(op, "cpu", torch.float32, torch.int32, False) + assert trace["arithmetic_profile"] == "adaln-select01-framework-v1" + assert trace["fallback"] is True + else: + x = torch.tensor([[[-3.,3.]]],requires_grad=True) + m = torch.zeros(1,6,requires_grad=True) + y,g = op(x,m,eps=7.) + dx,dm = torch.autograd.grad((y,g),(x,m), + (torch.tensor([[[1.,0.]]]),torch.ones_like(g))) + assert_bytes(y,torch.tensor([[[-.75,.75]]])) + assert_bytes(dx,torch.tensor([[[7/128,-7/128]]])) + assert_bytes(dm,torch.tensor([[1.,0.,-.75,0.,1.,1.]])) + assert trace["arithmetic_profile"] == "adaln-shared-fp32-tree-v1" + assert trace["actual_backend"] == "pytorch" + try: + KernelRegistry().get_adaln_modulation_op( + device, hidden=2, strict=True, + modulate_index=torch.zeros(1,2,dtype=torch.int32) if indexed else None) + except RuntimeError: + pass + else: + raise AssertionError("ROCm framework accepted as strict GPU replay") +""" + result = subprocess.run( + [sys.executable, "-c", script], + cwd=root, + env=dict(os.environ, PYTHONPATH=str(root), CUDA_VISIBLE_DEVICES=""), + capture_output=True, + text=True, + timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def test_ordinary_indexed_dispatch_adds_no_route_scan_or_scalar_read(): + from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp + from tests.helpers.adaln_indexed import assert_bytes + + x, m, index, _, _ = small_case("cpu") + op, trace = KernelRegistry().get_adaln_modulation_op("cpu", hidden=2, modulate_index=index) + observed = [] + outputs = [] + for callable_op in (NativeAdaLNModulationOp(), op): + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CPU], acc_events=True + ) as profile: + outputs.append(callable_op(x, m, modulate_index=index, eps=7.0)) + observed.append({event.key: event.count for event in profile.key_averages()}) + for actual, expected in zip(outputs[1], outputs[0], strict=True): + assert_bytes(actual, expected) + for event in ("aten::all", "aten::item"): + assert observed[0][event] == 1 + assert observed[1][event] == observed[0][event] + assert trace["actual_backend"] == "pytorch" and trace["fallback"] is True + assert trace["profiler_observed"] is False + + +def test_default_indexed_strict_rejection_clears_previous_shared_success(): + x, m, index, _, _ = small_case("cpu") + op, trace = KernelRegistry().get_adaln_modulation_op("cpu", hidden=2, strict=True) + op(x, m[:2], eps=7.0) + assert trace["actual_backend"] == "pytorch" + with pytest.raises(RuntimeError, match="strict AdaLN rejects"): + op(x, m, modulate_index=index, eps=7.0) + for key in ( + "actual_backend", + "observation", + "execution_dtype", + "output_dtype", + "execution_device", + "binary_path", + "binary_sha256", + "hardware_qualified", + ): + assert trace[key] is None + assert trace["profiler_observed"] is False + + +@pytest.fixture +def indexed_cuda_symbols(): + # Resolution needs a CUDA runtime identity and symbols, but no physical GPU. + if torch.version.cuda is None or torch.version.hip is not None: + pytest.skip("NVIDIA CUDA runtime identity required") + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + try: + op = CudaAdaLNModulationOp() + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + if not op.indexed_available: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + pytest.fail("indexed CUDA extension symbols required") + pytest.skip("indexed CUDA extension symbols required") + + +def test_explicit_indexed_activation_resolves_without_enabling_cpu_select01(indexed_cuda_symbols): + from rl_engine.runtime.adaln import IndexedAdaLNQualification + + scope = IndexedAdaLNQualification( + backend="cuda", + device="cuda:0", + device_name="reviewed target", + capability=(8, 6), + torch_version=torch.__version__, + cuda_version=torch.version.cuda, + ) + x, m, index, _, _ = small_case("cpu") + # Resolution is metadata-only; hardware is checked again before execution. + op, trace = KernelRegistry().get_adaln_modulation_op( + "cuda:0", + hidden=2, + modulate_index=index, + strict=True, + indexed_qualification=scope, + ) + assert trace["arithmetic_profile"] == "adaln-indexed-fp32-tree-v1" + with pytest.raises(RuntimeError, match="strict AdaLN rejects"): + op(x, m, modulate_index=index, eps=7.0) + assert trace["actual_backend"] is None + + +@pytest.mark.parametrize("eps", [7.0, 1e300, 2**-150]) +def test_strict_invalid_index_is_an_input_error_before_qualification_rejection(eps): + x, m, index, _, _ = small_case("cpu") + op, trace = KernelRegistry().get_adaln_modulation_op("cpu", hidden=2, strict=True) + op(x, m[:2], eps=7.0) + with pytest.raises(ValueError, match="values must be 0 or 1"): + op(x, m, modulate_index=index + 2, eps=eps) + assert trace["actual_backend"] is None and trace["observation"] is None + + +def test_indexed_resolution_then_shared_execution_reports_actual_shared_contract(): + x, m, index, _, _ = small_case("cpu") + op, trace = KernelRegistry().get_adaln_modulation_op("cpu", hidden=2, modulate_index=index) + op(x, m, modulate_index=index, eps=7.0) + assert trace["fallback"] is True + y, gate = op(x, m[:2], eps=7.0) + assert y.shape == x.shape and gate.shape == (2, 1, 2) + assert trace["actual_backend"] == "pytorch" + assert trace["arithmetic_profile"] == "adaln-shared-fp32-tree-v1" + assert trace["fallback"] is False and trace["fallback_reason"] is None + assert trace["kernel_fingerprint"] == "adaln_modulation_v1:PYTORCH_ADALN_MODULATION" + + +@pytest.mark.parametrize( + "eps,reason", + [(7.0, "NVIDIA CUDA tensors"), (1e300, "finite FP32 eps"), (2**-150, "finite FP32 eps")], +) +def test_activated_strict_rechecks_actual_domain_before_execution( + indexed_cuda_symbols, eps, reason +): + from rl_engine.runtime.adaln import IndexedAdaLNQualification + + scope = IndexedAdaLNQualification( + backend="cuda", + device="cuda:0", + device_name="reviewed target", + capability=(8, 6), + torch_version=torch.__version__, + cuda_version=torch.version.cuda, + ) + x, m, index, _, _ = small_case("cpu") + op, trace = KernelRegistry().get_adaln_modulation_op( + "cuda:0", + hidden=2, + modulate_index=index, + strict=True, + indexed_qualification=scope, + ) + with pytest.raises(RuntimeError, match=reason): + op(x, m, modulate_index=index, eps=eps) + assert trace["actual_backend"] is None + + +def test_ordinary_gpu_resolution_rechecks_shared_cpu_device_with_real_fallback(): + from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp + from tests.helpers.adaln_indexed import assert_bytes + + x, m, _, dy, dg = small_case("cpu") + op, trace = KernelRegistry().get_adaln_modulation_op("cuda", hidden=2) + y, gate = op(x, m[:2], eps=7.0) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg[:, :1])) + rx, rm = x.detach().requires_grad_(), m.detach().requires_grad_() + ry, rg = NativeAdaLNModulationOp()(rx, rm[:2], eps=7.0) + rdx, rdm = torch.autograd.grad((ry, rg), (rx, rm), (dy, dg[:, :1])) + for a, e in zip((y, gate, dx, dm), (ry, rg, rdx, rdm), strict=True): + assert_bytes(a, e) + assert trace["actual_backend"] == "pytorch" + assert trace["arithmetic_profile"] == "adaln-shared-fp32-tree-v1" + assert trace["fallback"] and trace["fallback_reason"] == "shared_cuda_domain_unsupported" + + +@pytest.mark.parametrize( + "index,error", + [ + (torch.zeros(8, dtype=torch.int32), ValueError), + (torch.zeros(1, 8), TypeError), + (torch.zeros(1, 0, dtype=torch.int32), ValueError), + ], +) +@pytest.mark.parametrize("strict", [False, True]) +def test_resolver_rejects_illegal_route_metadata_before_policy(index, error, strict): + with pytest.raises(error, match="modulate_index"): + KernelRegistry().get_adaln_modulation_op( + "cpu", hidden=2, modulate_index=index, strict=strict + ) + + +@pytest.mark.parametrize("mismatch", ["backend", "torch", "cuda", "device", "cpu", "hidden"]) +def test_activation_scope_does_not_qualify_another_backend_or_runtime(monkeypatch, mismatch): + from rl_engine.runtime.adaln import IndexedAdaLNQualification + + # Scope checks are metadata-only. Simulate a consistent runtime even on CPU/HIP. + monkeypatch.setattr(torch.version, "cuda", "reviewed-CUDA-runtime") + monkeypatch.setattr(torch.version, "hip", None) + baseline = IndexedAdaLNQualification( + backend="cuda", + device="cuda:0", + device_name="reviewed target", + capability=(8, 6), + torch_version=torch.__version__, + cuda_version=torch.version.cuda, + ) + baseline.check("cuda", "cuda:0") + scope = IndexedAdaLNQualification( + backend="triton" if mismatch == "backend" else "cuda", + device="cuda:1" if mismatch == "device" else "cuda:0", + device_name="reviewed target", + capability=(8, 6), + torch_version="unqualified" if mismatch == "torch" else torch.__version__, + cuda_version="unqualified" if mismatch == "cuda" else torch.version.cuda, + triton_version="unqualified" if mismatch == "backend" else None, + ) + if mismatch != "hidden": + with pytest.raises(RuntimeError, match="strict AdaLN rejects"): + scope.check("cuda", "cpu" if mismatch == "cpu" else "cuda:0") + with pytest.raises(RuntimeError, match="strict AdaLN rejects"): + KernelRegistry().get_adaln_modulation_op( + "cpu" if mismatch == "cpu" else "cuda:0", + hidden=4097 if mismatch == "hidden" else 2, + modulate_index=torch.zeros(1, 8, dtype=torch.int32), + strict=True, + indexed_qualification=scope, + ) + + +def test_cuda_device_string_on_hip_cannot_open_indexed_qualification(monkeypatch): + # torch.version is the external platform boundary, never a registry collaborator. + monkeypatch.setattr(torch.version, "hip", "test-ROCm-runtime") + index = torch.zeros(1, 8, dtype=torch.int32) + op, trace = KernelRegistry().get_adaln_modulation_op("cuda", hidden=2, modulate_index=index) + assert trace["selected_backend"] == "PYTORCH_ADALN_MODULATION" + assert trace["arithmetic_profile"] == "adaln-select01-framework-v1" + with pytest.raises(RuntimeError, match="strict AdaLN rejects"): + KernelRegistry().get_adaln_modulation_op( + "cuda", hidden=2, modulate_index=index, strict=True + ) + x, m, index, _, _ = small_case("cpu") + op(x, m, modulate_index=index, eps=7.0) + assert trace["actual_backend"] == "pytorch" and trace["fallback"] + + +@pytest.mark.parametrize("resolution", ["indexed", "shared"]) +@pytest.mark.parametrize("hidden", [2, 3]) +def test_ordinary_cpu_rechecks_mode_and_hidden_without_stale_trace(resolution, hidden): + x, m, index, _, _ = small_case("cpu") + op, trace = KernelRegistry().get_adaln_modulation_op( + "cpu", + hidden=hidden, + modulate_index=index if resolution == "indexed" else None, + ) + for indexed in (True, False, True): + op(x, m if indexed else m[:2], modulate_index=index if indexed else None, eps=7.0) + assert trace["arithmetic_profile"] == ( + "adaln-select01-framework-v1" if indexed else "adaln-shared-fp32-tree-v1" + ) + assert trace["fallback"] is indexed + assert trace["fallback_reason"] == ("modulate_index_requires_select01" if indexed else None) diff --git a/tests/runtime/test_adaln_test_environments.py b/tests/runtime/test_adaln_test_environments.py new file mode 100644 index 000000000..86306de16 --- /dev/null +++ b/tests/runtime/test_adaln_test_environments.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Exercise pytest environment behavior with controlled external runtime identity.""" + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[2] + + +def run_pytest(script, *nodes, require_ext="0"): + result = subprocess.run( + [ + sys.executable, + "-c", + script + "\nimport pytest, sys\nsys.exit(pytest.main(sys.argv[1:]))", + "-q", + "-r", + "s", + *nodes, + ], + cwd=ROOT, + env=dict( + os.environ, + PYTHONPATH=str(ROOT), + CUDA_VISIBLE_DEVICES="", + RL_KERNEL_REQUIRE_EXT=require_ext, + RLK_ADALN_REAL_SHAPES="1", + ), + capture_output=True, + text=True, + timeout=45, + ) + return result + + +CPU_WHEEL = """ +import torch +torch.version.cuda = None +torch.version.hip = None +torch.cuda.is_available = lambda: False +def no_cuda_init(): + raise AssertionError("Torch not compiled with CUDA enabled") +torch.cuda.init = no_cuda_init +""" + + +def test_cuda_execution_tests_skip_with_cpu_wheel_identity(): + result = run_pytest( + CPU_WHEEL, + "tests/ops/norm/test_adaln_modulation.py::test_triton_adaln_modulation_forward_backward_and_batch_position", + "tests/ops/norm/test_adaln_modulation.py::test_triton_irregular_hidden_and_bf16_tie_rounding", + "tests/ops/norm/test_adaln_modulation.py::test_adaln_gpu_padding_and_batch_size_do_not_change_valid_bytes", + "tests/ops/norm/test_adaln_modulation.py::test_cuda_adaln_bit_equality_harness", + "tests/ops/norm/test_adaln_modulation.py::test_adaln_real_image_shapes_forward_backward", + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "17 skipped" in result.stdout + + +def test_resolution_with_cpu_wheel_identity_needs_no_physical_gpu(): + pytest.importorskip("triton") + result = run_pytest( + CPU_WHEEL, + "tests/runtime/test_adaln_dispatch.py::test_adaln_cuda_registry_reports_real_fallback", + "tests/runtime/test_adaln_dispatch.py::test_adaln_registry_skips_cuda_for_large_hidden", + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "2 passed" in result.stdout + + +def test_missing_extension_still_fails_when_required(): + result = run_pytest( + CPU_WHEEL + '\nimport sys\nsys.modules["rl_engine._C"] = None\n', + "tests/ops/norm/test_adaln_modulation.py::test_cuda_adaln_bit_equality_harness", + "tests/runtime/test_adaln_dispatch.py::test_cuda_shared_registry_respects_hidden_limit", + require_ext="1", + ) + assert result.returncode == 1, result.stdout + result.stderr + assert "CUDA AdaLN extension symbols are unavailable" in result.stdout + assert "6 failed, 3 errors" in result.stdout + + +@pytest.mark.parametrize("hip", [None, "test-HIP-runtime"]) +def test_indexed_qualification_tests_with_cpu_wheel_identity(hip): + result = run_pytest( + CPU_WHEEL + f"\ntorch.version.hip = {hip!r}\n", + "tests/runtime/test_adaln_strict_policy.py::test_explicit_indexed_activation_resolves_without_enabling_cpu_select01", + "tests/runtime/test_adaln_strict_policy.py::test_activated_strict_rechecks_actual_domain_before_execution", + "tests/runtime/test_adaln_strict_policy.py::test_activation_scope_does_not_qualify_another_backend_or_runtime", + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "6 passed, 4 skipped" in result.stdout + + +@pytest.mark.parametrize("require_ext", ["0", "1"]) +def test_cli_tests_without_cuda_extension_only_skip_optional_cuda_candidate(require_ext): + result = run_pytest( + 'import sys\nsys.modules["rl_engine._C"] = None\n', + "tests/validation/operators/test_check_operator_cli.py", + require_ext=require_ext, + ) + if require_ext == "1": + assert result.returncode == 1, result.stdout + result.stderr + assert "1 failed, 2 passed" in result.stdout + assert "CUDA AdaLN extension symbols are unavailable" in result.stdout + else: + assert result.returncode == 0, result.stdout + result.stderr + assert "2 passed, 1 skipped" in result.stdout diff --git a/tests/runtime/test_layout_compatibility.py b/tests/runtime/test_layout_compatibility.py index a3ab9f969..1f7d6fa02 100644 --- a/tests/runtime/test_layout_compatibility.py +++ b/tests/runtime/test_layout_compatibility.py @@ -123,6 +123,9 @@ def test_packaged_contract_bytes_match_main(package, resource, digest): [ "rl_engine/repro.py", "scripts/ws1_reference.py", + "scripts/check_operator.py", + "tools/validation/operators/check_operator.py", + "benchmarks/operators/norm/benchmark_adaln_modulation.py", "rl_engine/testing/logprob_comparison.py", "examples/vime_rocm_attention_ablation/run.py", ], diff --git a/tests/validation/operators/test_adaln_gtest.py b/tests/validation/operators/test_adaln_gtest.py new file mode 100644 index 000000000..a5882e25d --- /dev/null +++ b/tests/validation/operators/test_adaln_gtest.py @@ -0,0 +1,695 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""CPU checks of AdaLN gtest wiring, independent gold and gradient sensitivity.""" + +import os +from argparse import Namespace +from dataclasses import replace + +import pytest +import torch + +from rl_engine.contracts.numerical import load_contract, resolve_tolerance +from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp +from rl_engine.runtime.registry import OpBackend, kernel_registry +from rl_engine.validation.operators import run_operator_suite +from rl_engine.validation.operators.operator_inputs import operator_shape_name +from rl_engine.validation.operators.operator_specs import ( + GtestAdaLNReference, + make_candidate, + make_operator_case, +) + + +def _args(hidden=7): + return Namespace( + op="adaln_modulation", + candidate="pytorch", + arch_key=None, + batch=2, + seq=3, + normalized_dim=hidden, + seed=386, + ) + + +@pytest.fixture( + params=["pytorch", "triton"] + + [pytest.param(("cuda", threads), id=f"cuda-threads{threads}") for threads in (128, 256, 512)] +) +def shared_op(request): + backend, threads = request.param if isinstance(request.param, tuple) else (request.param, None) + if backend != "pytorch" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + args = _args() + args.candidate = backend + try: + op = make_candidate(args).fn + if threads is not None: + op = type(op)(threads=threads) + except RuntimeError: + if backend != "cuda": + raise + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + return op, "cpu" if backend == "pytorch" else "cuda" + + +def _bytes_equal(actual, expected): + assert actual.shape == expected.shape and actual.dtype == expected.dtype + assert torch.equal( + actual.detach().cpu().contiguous().view(torch.uint8), + expected.detach().cpu().contiguous().view(torch.uint8), + ) + + +@pytest.mark.parametrize("nan_upstream", [False, True]) +def test_adaln_gate_preserves_bf16_nan_payload_and_identity_vjp(shared_op, nan_upstream): + op, device = shared_op + x = torch.tensor([[[-1.0, 1.0]]], device=device, dtype=torch.bfloat16) + m = torch.zeros(1, 6, device=device, dtype=torch.bfloat16) + # Gate is a passthrough: converting these payloads through FP32 is lossy. + payload = torch.tensor([0x7F81, -63], dtype=torch.int16).view(torch.bfloat16) + m[:, 4:] = payload.to(device) + m.requires_grad_() + _, gate = op(x, m) + _bytes_equal(gate, payload.reshape(1, 1, 2)) + dg = ( + torch.tensor([0x7F82, 0x7FFF], dtype=torch.int16).view(m.dtype).reshape(1, 1, 2) + if nan_upstream + else torch.tensor([[[2.0, -3.0]]], dtype=m.dtype) + ) + (dm,) = torch.autograd.grad(gate, (m,), dg.to(device)) + expected = torch.zeros(1, 6, dtype=m.dtype) + expected[:, 4:] = dg[:, 0] + _bytes_equal(dm, expected) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_adaln_signed_zero_follows_hidden_tree(shared_op, dtype): + op, device = shared_op + x = torch.full((1, 1, 2), -0.0, device=device, dtype=dtype) + m = torch.zeros(1, 6, device=device, dtype=dtype) + m[:, :2] = -0.0 + y, _ = op(x, m) + # (-0 + -0)/2 = -0, so centered = -0 - -0 = +0. An extra +0 in + # the final reduction erases mean's sign and incorrectly makes y=-0. + _bytes_equal(y, torch.zeros(1, 1, 2, dtype=dtype)) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize( + "eps,rstd", + [ + (1e-40, 1.000002658868784e20), + (2**-149, 2.671373844909537e22), + (3.4028234663852886e38, 5.421011508662376e-20), + ], +) +def test_adaln_subnormal_eps_is_not_flushed_to_zero(shared_op, dtype, eps, rstd): + op, device = shared_op + x = torch.zeros(1, 1, 2, device=device, dtype=dtype, requires_grad=True) + m = torch.zeros(1, 6, device=device, dtype=dtype, requires_grad=True) + y, gate = op(x, m, eps=eps) + _bytes_equal(y, torch.zeros(1, 1, 2, dtype=dtype)) + dy = torch.tensor([[[1.0, -1.0]]], device=device, dtype=dtype) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, torch.zeros_like(gate))) + # Independent scalar rounding literals: FP32(1 / FP32(sqrt(FP32(eps)))). + _bytes_equal(dx, torch.tensor([[[rstd, -rstd]]], dtype=dtype)) + _bytes_equal(dm, torch.tensor([[1.0, -1.0, 0.0, 0.0, 0.0, 0.0]], dtype=dtype)) + + +@pytest.mark.parametrize("scale, positive", [(0.0, -(2**-21)), (1 / 256, 1 / 256)]) +def test_adaln_bf16_forward_casts_only_at_output(shared_op, scale, positive): + op, device = shared_op + x = torch.tensor([[[-1.0, 1.0]]], device=device, dtype=torch.bfloat16) + m = torch.tensor([[-1.0, -1.0, scale, scale, -0.0, 2.0]], device=device, dtype=torch.bfloat16) + y, gate = op(x, m) + # FP32 sqrt(1 + 1e-6) = 1 + 2^-21, reciprocal = 1 - 2^-21. + # Rounding norm or (1 + 1/256) to BF16 early erases the positive residual. + _bytes_equal(y, torch.tensor([[[-2.0, positive]]], dtype=torch.bfloat16)) + _bytes_equal(gate, torch.tensor([[[-0.0, 2.0]]], dtype=torch.bfloat16)) + + +def test_adaln_bf16_backward_keeps_partials_and_accumulators_fp32(shared_op): + op, device = shared_op + x = torch.tensor([[[-3.0, 3.0]] * 3], device=device, dtype=torch.bfloat16, requires_grad=True) + m = torch.tensor( + [[0.0, 0.0, 1 / 256, 1 / 256, 2.0, 3.0]], + device=device, + dtype=torch.bfloat16, + requires_grad=True, + ) + dy = torch.tensor( + [[[256.0, 0.0], [1.0, 0.0], [-256.0, 0.0]]], device=device, dtype=torch.bfloat16 + ) + dg = torch.tensor([[[-2.0, 0.5]]], device=device, dtype=torch.bfloat16) + y, gate = op(x, m, eps=7.0) + torch.autograd.backward((y, gate), (dy, dg)) + # variance+eps=16, rstd=1/4, norm=[-3/4,+3/4]. The exact VJP coefficient + # is (1+1/256)*(1-9/16)/8 = 1799/32768, independent of kernel helpers. + coefficient = 1799 / 32768 + expected_dx = torch.tensor( + [ + [ + [256 * coefficient, -256 * coefficient], + [coefficient, -coefficient], + [-256 * coefficient, 256 * coefficient], + ] + ], + dtype=torch.bfloat16, + ) + _bytes_equal(x.grad, expected_dx) + # FP32 left folds: 256+1-256=1; -192-3/4+192=-3/4. + _bytes_equal(m.grad, torch.tensor([[1.0, 0.0, -0.75, 0.0, -2.0, 0.5]], dtype=torch.bfloat16)) + + +@pytest.mark.parametrize("hidden", [7, 3072]) +def test_adaln_bf16_vjp_matches_fp32_with_one_final_cast(shared_op, hidden): + op, device = shared_op + torch.manual_seed(389) + x = torch.randn(2, 3, hidden, device=device, dtype=torch.bfloat16) + m = torch.randn(2, 3 * hidden, device=device, dtype=torch.bfloat16) + dy, dg = torch.randn_like(x), torch.randn(2, 1, hidden, device=device, dtype=x.dtype) + observations = [] + for dtype in (torch.bfloat16, torch.float32): + xx, mm = ( + x.to(dtype).detach().requires_grad_(), + m.to(dtype).detach().requires_grad_(), + ) + y, gate = op(xx, mm) + dx, dm = torch.autograd.grad((y, gate), (xx, mm), (dy.to(dtype), dg.to(dtype))) + observations.append((y, gate, dx, dm)) + # Identical quantized inputs also expose BF16 saved norms/partials used only + # in backward. Independent LayerNorm math is checked in a separate test. + for actual, fp32 in zip(*observations, strict=True): + _bytes_equal(actual, fp32.to(torch.bfloat16)) + + +@pytest.mark.parametrize("fraction, rounded", [(1 / 256, 1.0), (3 / 256, 1 + 4 / 256)]) +def test_adaln_bf16_boundary_rounds_halfway_values_to_even(shared_op, fraction, rounded): + op, device = shared_op + x = torch.tensor([[[-1.0, 1.0]] * 2], device=device, dtype=torch.bfloat16) + m = torch.zeros(1, 6, device=device, dtype=x.dtype) + m[:, :2] = fraction + m.requires_grad_() + y, _ = op(x, m, eps=0.0) + _bytes_equal(y, torch.tensor([[[-1 + fraction, rounded]] * 2], dtype=x.dtype)) + dy = torch.tensor([[[0.0, 1.0], [0.0, fraction]]], device=device, dtype=x.dtype) + (dm,) = torch.autograd.grad(y, (m,), dy) + _bytes_equal(dm, torch.tensor([[0.0, rounded, 0.0, rounded, 0.0, 0.0]], dtype=x.dtype)) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_adaln_h3072_sample_bytes_survive_batch_positions_and_padding(shared_op, dtype): + _assert_h3072_relocation(shared_op, dtype) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize( + "backend,options", + [("triton", {"num_warps": w, "reduction_tile": t}) for w in (4, 8) for t in (64, 128, 256)] + + [("cuda", {"threads": n}) for n in (128, 256, 512)], +) +def test_adaln_h3072_relocation_with_all_launch_configurations(dtype, backend, options): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + args = _args() + args.candidate = backend + try: + default_op = make_candidate(args).fn + except RuntimeError: + if backend != "cuda" or os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + op_class = type(default_op) + op = op_class(**options) + _assert_h3072_relocation((op, "cuda"), dtype, reference_op=default_op) + + +def _assert_h3072_relocation(shared_op, dtype, *, reference_op=None): + op, device = shared_op + torch.manual_seed(386) + hidden, valid = 3072, 3 + x0 = torch.randn(1, valid, hidden, device=device, dtype=dtype) + full0 = torch.randn(1, 6 * hidden, device=device, dtype=dtype) + dy0 = torch.randn_like(x0) + dg0 = torch.randn(1, 1, hidden, device=device, dtype=dtype) + + def observe(x, full, dy, dg, *, candidate=op): + x = x.detach().requires_grad_() + full = full.detach().requires_grad_() + y, gate = candidate(x, full[:, : 3 * hidden]) + dx, dfull = torch.autograd.grad((y, gate), (x, full), (dy, dg)) + _bytes_equal(dfull[:, 3 * hidden :], torch.zeros_like(dfull[:, 3 * hidden :])) + return y, gate, dx, dfull[:, : 3 * hidden] + + baseline = observe(x0, full0, dy0, dg0, candidate=reference_op or op) + for batch, seq in ((1, 5), (3, 3), (3, 5)): + for position in range(batch): + # Strides and six-chunk views match model callers, alongside topology changes. + x = torch.randn(batch, 2 * seq, hidden, device=device, dtype=dtype)[:, ::2] + full = torch.randn(batch, 6 * hidden, device=device, dtype=dtype) + dy = torch.randn_like(x) + dg = torch.randn(batch, 1, hidden, device=device, dtype=dtype) + x[position, :valid], full[position] = x0[0], full0[0] + dy[position, :valid], dg[position] = dy0[0], dg0[0] + dy[position, valid:] = 0 + actual = observe(x, full, dy, dg) + for index, (result, expected) in enumerate(zip(actual, baseline, strict=True)): + result = result[position, :valid] if index in (0, 2) else result[position] + _bytes_equal(result, expected[0]) + # Padding has zero VJP; IEEE arithmetic may produce either zero sign. + assert torch.equal( + actual[2][position, valid:], + torch.zeros_like(actual[2][position, valid:]), + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_adaln_backward_accepts_transposed_and_broadcast_upstreams(shared_op, dtype): + op, device = shared_op + torch.manual_seed(391) + x = torch.randn(2, 3, 7, device=device, dtype=dtype, requires_grad=True) + m = torch.randn(2, 21, device=device, dtype=dtype, requires_grad=True) + dy = torch.randn(2, 7, 3, device=device, dtype=dtype).transpose(1, 2) + dg_storage = torch.randn(2, 1, 7, device=device, dtype=dtype) + dg = dg_storage[:1].expand(2, 1, 7) + assert not dy.is_contiguous() and dg.stride(0) == 0 + actual = torch.autograd.grad(op(x, m), (x, m), (dy, dg)) + contiguous = torch.autograd.grad(op(x, m), (x, m), (dy.contiguous(), dg.contiguous())) + for result, expected in zip(actual, contiguous, strict=True): + _bytes_equal(result, expected) + xx = x.detach().cpu().float().requires_grad_() + mm = m.detach().cpu().float().requires_grad_() + gold = torch.autograd.grad( + GtestAdaLNReference().forward_fp32(xx, mm), + (xx, mm), + (dy.cpu().float(), dg.cpu().float()), + ) + tol = resolve_tolerance( + load_contract(), judgment="gradient_accuracy", op_class="reduction", dtype=dtype + ) + for result, expected in zip(actual, gold, strict=True): + torch.testing.assert_close(result.cpu().float(), expected, atol=tol.atol, rtol=tol.rtol) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("trainable", ["x", "modulation", "gate_only"]) +def test_adaln_single_trainable_input_and_unused_output(shared_op, dtype, trainable): + op, device = shared_op + x_trainable = trainable != "modulation" + x = torch.tensor( + [[[-3.0, 3.0] if trainable == "x" else [-1.0, 1.0]]], + device=device, + dtype=dtype, + requires_grad=x_trainable, + ) + m = torch.zeros(1, 6, device=device, dtype=dtype, requires_grad=trainable != "x") + y, gate = op(x, m, eps=7.0 if trainable == "x" else 0.0) + if trainable == "x": + y.backward(torch.tensor([[[1.0, 0.0]]], device=device, dtype=dtype)) + _bytes_equal(x.grad, torch.tensor([[[7 / 128, -7 / 128]]], dtype=dtype)) + assert m.grad is None + elif trainable == "modulation": + y.sum().backward() # Gate is genuinely unused, rather than given a zero VJP. + assert x.grad is None + _bytes_equal(m.grad, torch.tensor([[1.0, 1.0, -1.0, 1.0, 0.0, 0.0]], dtype=dtype)) + else: + gate.backward(torch.tensor([[[2.0, -3.0]]], device=device, dtype=dtype)) + assert torch.equal(x.grad, torch.zeros_like(x)) + _bytes_equal(m.grad, torch.tensor([[0.0, 0.0, 0.0, 0.0, 2.0, -3.0]], dtype=dtype)) + + +def test_adaln_hidden_reduction_uses_declared_pairwise_order(shared_op): + op, device = shared_op + x = torch.tensor([[[2**24, 1.0, -(2**24), 1.0]]], device=device) + m = torch.zeros(1, 12, device=device) + y, _ = op(x, m, eps=0.0) + # Lower+upper sum is 2 (left fold is 1), mean=1/2, FP32 variance=2^47. + _bytes_equal( + y, + torch.tensor( + [ + [ + [ + 1.4142135381698608, + 4.214684778958144e-8, + -1.4142135381698608, + 4.214684778958144e-8, + ] + ] + ] + ), + ) + + +def test_adaln_variance_reduction_uses_declared_pairwise_order(shared_op): + op, device = shared_op + x = torch.tensor([[[4096.0, 1.25, -4096.0, -1.25]]], device=device) + y, _ = op(x, torch.zeros(1, 12, device=device), eps=0.0) + # Mean is zero under either tree. Pairwise variance=8388609, left fold=8388608. + _bytes_equal( + y, + torch.tensor( + [ + [ + [ + 1.4142134189605713, + 0.0004315836704336107, + -1.4142134189605713, + -0.0004315836704336107, + ] + ] + ] + ), + ) + + +@pytest.mark.parametrize("sign", [1, -1]) +def test_adaln_backward_hidden_sums_use_declared_pairwise_order(shared_op, sign): + op, device = shared_op + x = torch.tensor([[[-1.0, 1.0, -1.0, 1.0]]], device=device, requires_grad=True) + m = torch.zeros(1, 12, device=device) + y, _ = op(x, m, eps=0.0) + dy = torch.tensor([[[sign * 2**24, 1.0, -sign * 2**24, 1.0]]], device=device) + (dx,) = torch.autograd.grad(y, (x,), dy) + # Both hidden means are 1/2. A left fold in either reduction leaves 1/4 + # in channels 1 and 3; the two signs discriminate the reductions separately. + _bytes_equal(dx, torch.tensor([[[sign * 2**24, 0.0, -sign * 2**24, 0.0]]])) + + +def test_adaln_modulation_gradient_uses_ascending_token_order(shared_op): + op, device = shared_op + x = torch.tensor([[[-1.0, 1.0]] * 4], device=device, requires_grad=True) + m = torch.zeros(1, 6, device=device, requires_grad=True) + y, gate = op(x, m, eps=0.0) + dy = torch.tensor([[[2**24, 0.0], [1.0, 0.0], [-(2**24), 0.0], [1.0, 0.0]]], device=device) + (dm,) = torch.autograd.grad((y, gate), (m,), (dy, torch.zeros_like(gate))) + # Ascending FP32 fold is 1; reordering to lower+upper pairs would yield 2. + _bytes_equal(dm, torch.tensor([[1.0, 0.0, -1.0, 0.0, 0.0, 0.0]])) + + +def test_adaln_independent_gold_uses_contract_token_order(): + x = torch.tensor([[[-1.0, 1.0]] * 8]) + m = torch.zeros(1, 6, requires_grad=True) + dy = torch.tensor([[[2**24, 0.0], [1.0, 0.0], [-(2**24), 0.0], [1.0, 0.0]]]) + y, gate = GtestAdaLNReference().forward_fp32(x, m, eps=0.0) + (dm,) = torch.autograd.grad((y, gate), (m,), (dy.repeat(1, 2, 1), torch.zeros_like(gate))) + # The carry from the first block is lost at 1+2^24; each ascending fold ends at 1. + # A reassociated CPU broadcast backward returns +/-2 for two blocks. + _bytes_equal(dm, torch.tensor([[1.0, 0.0, -1.0, 0.0, 0.0, 0.0]])) + + +@pytest.mark.parametrize( + "invalid, error, message", + [ + ("rank", ValueError, "nonempty shape"), + ("empty_batch", ValueError, "nonempty shape"), + ("empty_sequence", ValueError, "nonempty shape"), + ("empty_hidden", ValueError, "nonempty shape"), + ("modulation_shape", ValueError, "modulation must have shape"), + ("mixed_dtype", ValueError, "share device and dtype"), + ("fp16", TypeError, "FP32 or BF16"), + ("negative_eps", ValueError, "finite and nonnegative"), + ("nan_eps", ValueError, "finite and nonnegative"), + ("infinite_eps", ValueError, "finite and nonnegative"), + ], +) +def test_adaln_shared_rejects_unsupported_inputs(shared_op, invalid, error, message): + op, device = shared_op + shape = { + "rank": (1, 2), + "empty_batch": (0, 2, 2), + "empty_sequence": (1, 0, 2), + "empty_hidden": (1, 2, 0), + }.get(invalid, (1, 2, 2)) + x = torch.zeros(shape, device=device) + m = torch.zeros(shape[0], 3 * shape[-1], device=device) + if invalid == "modulation_shape": + m = m[:, :-1] + if invalid == "mixed_dtype": + m = m.bfloat16() + if invalid == "fp16": + x, m = x.half(), m.half() + eps = { + "negative_eps": -1.0, + "nan_eps": float("nan"), + "infinite_eps": float("inf"), + }.get(invalid, 1e-6) + with pytest.raises(error, match=message): + op(x, m, eps=eps) + + +def test_adaln_shared_rejects_mixed_devices(shared_op): + op, device = shared_op + with pytest.raises(ValueError, match="share device and dtype"): + op(torch.zeros(1, 2, 2, device=device), torch.empty(1, 6, device="meta")) + + +@pytest.mark.parametrize("eps", [1e39, 1e-46]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_adaln_shared_rejects_eps_overflow_and_positive_underflow(shared_op, dtype, eps): + op, device = shared_op + x = torch.zeros(1, 1, 2, device=device, dtype=dtype) + m = torch.zeros(1, 6, device=device, dtype=dtype) + with pytest.raises(ValueError, match="finite FP32.*underflow"): + op(x, m, eps=eps) + + +@pytest.mark.parametrize("value", [float("nan"), float("inf")]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_adaln_nonfinite_activation_propagates_nan_and_preserves_gate_vjp(shared_op, dtype, value): + op, device = shared_op + x = torch.tensor([[[value, 1.0]]], device=device, dtype=dtype, requires_grad=True) + m = torch.zeros(1, 6, device=device, dtype=dtype) + m[:, 4:] = torch.tensor([2.0, -0.0], device=device, dtype=dtype) + m.requires_grad_() + y, gate = op(x, m) + assert torch.isnan(y).all() + _bytes_equal(gate, torch.tensor([[[2.0, -0.0]]], dtype=dtype)) + dy = torch.tensor([[[1.0, -1.0]]], device=device, dtype=dtype) + dg = torch.tensor([[[2.0, -3.0]]], device=device, dtype=dtype) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy, dg)) + assert torch.isnan(dx).all() and torch.isnan(dm[:, 2:4]).all() + _bytes_equal(dm[:, :2], torch.tensor([[1.0, -1.0]], dtype=dtype)) + _bytes_equal(dm[:, 4:], dg[:, 0]) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("hidden", [7, 3072]) +def test_adaln_gtest_outputs_and_gradients_use_shared_contract(dtype, hidden): + args = _args(hidden) + case = make_operator_case(args, dtype, torch.device("cpu")) + assert operator_shape_name(args.op, args) == f"2x3x{hidden}" + assert case.inputs["modulation"].shape == (2, 3 * hidden) + gold = case.gold_fn(**case.inputs) + assert all(t.device.type == "cpu" and t.dtype == torch.float32 for t in gold) + assert gold[0].shape == (2, 3, hidden) and gold[1].shape == (2, 1, hidden) + report = run_operator_suite( + args.op, candidates=[make_candidate(args)], cases=[case], check_grad=True + ) + assert report.passed, report.to_dict() + checks = report.candidates[0].cases[0].outputs + assert len(checks) == 4 + assert all(c.gold_dtype == "torch.float32" for c in checks) + assert [c.message for c in checks[2:]] == ["gradient:x", "gradient:modulation"] + for check in checks: + threshold = resolve_tolerance( + load_contract(), judgment=check.judgment, op_class="reduction", dtype=dtype + ) + assert (check.atol, check.rtol) == (threshold.atol, threshold.rtol) + + +@pytest.mark.parametrize("broken", ["gate", "x", "modulation"]) +def test_adaln_gtest_detects_wrong_output_and_backward(broken): + args = _args() + case = make_operator_case(args, torch.float32, torch.device("cpu")) + good = make_candidate(args) + + def corrupt(x, modulation, eps): + # Zero-valued terms preserve connectivity while deliberately erasing a VJP. + xx = x.detach() + x * 0 if broken == "x" else x + mm = modulation.detach() + modulation * 0 if broken == "modulation" else modulation + y, gate = good.fn(xx, mm, eps=eps) + return y, gate + 1 if broken == "gate" else gate + + report = run_operator_suite( + args.op, candidates=[replace(good, fn=corrupt)], cases=[case], check_grad=True + ) + assert not report.passed + checks = report.candidates[0].cases[0].outputs + check = ( + checks[1] + if broken == "gate" + else next(c for c in checks if c.message == f"gradient:{broken}") + ) + assert not check.passed + + +@pytest.mark.parametrize("check_grad", [False, True]) +@pytest.mark.parametrize("wrong_output", [0, 1]) +def test_adaln_gtest_rejects_fp32_outputs_for_bf16_inputs(check_grad, wrong_output): + args = _args() + case = make_operator_case(args, torch.bfloat16, torch.device("cpu")) + good = make_candidate(args) + + def wrong_dtype(**inputs): + outputs = list(good.fn(**inputs)) + outputs[wrong_output] = outputs[wrong_output].float() + return tuple(outputs) + + report = run_operator_suite( + args.op, + candidates=[replace(good, fn=wrong_dtype)], + cases=[case], + check_grad=check_grad, + ) + assert not report.passed, report.to_dict() + checks = report.candidates[0].cases[0].outputs[:2] + assert not checks[wrong_output].passed + assert "dtype mismatch" in checks[wrong_output].message + assert checks[1 - wrong_output].passed + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("triple", [0, 1]) +def test_adaln_gate_only_and_separate_stream_chunks(shared_op, dtype, triple): + op, device = shared_op + torch.manual_seed(386) + hidden = 7 + # Separate image/text projections, each exposing a strided triple of six chunks. + image = torch.randn(2, 5, hidden, device=device, dtype=dtype, requires_grad=True) + text = torch.randn(2, 3, hidden, device=device, dtype=dtype, requires_grad=True) + image_full = torch.randn(2, 6 * hidden, device=device, dtype=dtype, requires_grad=True) + text_full = torch.randn(2, 6 * hidden, device=device, dtype=dtype, requires_grad=True) + start, stop = triple * 3 * hidden, (triple + 1) * 3 * hidden + unused = slice(3 * hidden, None) if triple == 0 else slice(None, 3 * hidden) + for x, full, other in ( + (image, image_full, text_full), + (text, text_full, image_full), + ): + m = full[:, start:stop] + other_before = None if other.grad is None else other.grad.clone() + assert not m.is_contiguous() + y, gate = op(x, m) + dg = torch.randn_like(gate) + torch.autograd.backward((y, gate), (torch.zeros_like(y), dg)) + assert torch.equal(x.grad, torch.zeros_like(x)) + expected = torch.zeros_like(full) + expected[:, start + 2 * hidden : stop] = dg[:, 0] + assert torch.equal(full.grad, expected) + if other_before is None: + assert other.grad is None + else: + _bytes_equal(other.grad, other_before) + # Independent CPU gold with nonuniform y/gate upstreams on each stream. + ref_x = x.detach().cpu().float().requires_grad_() + ref_m = m.detach().cpu().float().requires_grad_() + ref_y, ref_gate = GtestAdaLNReference().forward_fp32(ref_x, ref_m) + forward_tol = resolve_tolerance( + load_contract(), + judgment="forward_accuracy", + op_class="reduction", + dtype=dtype, + ) + torch.testing.assert_close( + y.detach().cpu().float(), + ref_y.detach(), + atol=forward_tol.atol, + rtol=forward_tol.rtol, + ) + _bytes_equal(gate, m[:, None, 2 * hidden :]) + dy = torch.randn_like(y) + actual = torch.autograd.grad(op(x, m), (x, full), (dy, dg)) + gold = torch.autograd.grad( + (ref_y, ref_gate), (ref_x, ref_m), (dy.cpu().float(), dg.cpu().float()) + ) + tol = resolve_tolerance( + load_contract(), + judgment="gradient_accuracy", + op_class="reduction", + dtype=dtype, + ) + for result, expected_grad in ( + (actual[0], gold[0]), + (actual[1][:, start:stop], gold[1]), + ): + torch.testing.assert_close( + result.cpu().float(), expected_grad, atol=tol.atol, rtol=tol.rtol + ) + assert torch.equal(actual[1][:, unused], torch.zeros_like(actual[1][:, unused])) + + +@pytest.mark.parametrize("hidden", [4095, 4096]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_adaln_hidden_boundary_matches_independent_math_and_fixed_bytes(shared_op, dtype, hidden): + op, device = shared_op + torch.manual_seed(395) + x_cpu = torch.randn(2, 3, hidden, dtype=dtype) + m_cpu = torch.randn(2, 3 * hidden, dtype=dtype) + dy, dg = torch.randn_like(x_cpu), torch.randn(2, 1, hidden, dtype=dtype) + x = x_cpu.to(device).detach().requires_grad_() + m = m_cpu.to(device).detach().requires_grad_() + y, gate = op(x, m) + dx, dm = torch.autograd.grad((y, gate), (x, m), (dy.to(device), dg.to(device))) + xx, mm = x_cpu.float().requires_grad_(), m_cpu.float().requires_grad_() + gold_y, gold_gate = GtestAdaLNReference().forward_fp32(xx, mm) + gold_dx, gold_dm = torch.autograd.grad((gold_y, gold_gate), (xx, mm), (dy.float(), dg.float())) + for actual, expected, judgment in ( + (y, gold_y, "forward_accuracy"), + (gate, gold_gate, "forward_accuracy"), + (dx, gold_dx, "gradient_accuracy"), + (dm, gold_dm, "gradient_accuracy"), + ): + tol = resolve_tolerance( + load_contract(), judgment=judgment, op_class="reduction", dtype=dtype + ) + torch.testing.assert_close(actual.cpu().float(), expected, atol=tol.atol, rtol=tol.rtol) + fx, fm = x_cpu.detach().requires_grad_(), m_cpu.detach().requires_grad_() + fixed_y, fixed_gate = NativeAdaLNModulationOp()(fx, fm) + fixed_dx, fixed_dm = torch.autograd.grad((fixed_y, fixed_gate), (fx, fm), (dy, dg)) + for actual, expected in zip( + (y, gate, dx, dm), (fixed_y, fixed_gate, fixed_dx, fixed_dm), strict=True + ): + _bytes_equal(actual, expected) + + +def test_adaln_declared_unavailable_fallbacks_and_exhaustion(monkeypatch): + original = kernel_registry._get_or_create_backend + missing = set() + + def resolve(backend): + if backend in missing: + return None + if backend is OpBackend.TRITON_ADALN_MODULATION: + return object() # Resolver-only sentinel: CPU CI does not install Triton. + return original(backend) + + monkeypatch.setattr(kernel_registry, "_get_or_create_backend", resolve) + missing.add(OpBackend.CUDA_ADALN_MODULATION) + # No GPU execution: test resolution and trace, then exercise the CPU fallback. + _, trace = kernel_registry.get_adaln_modulation_op("cuda", hidden=7) + assert trace["selected_backend"] == OpBackend.TRITON_ADALN_MODULATION.name + assert trace["fallback"] is True + missing.add(OpBackend.TRITON_ADALN_MODULATION) + op, trace = kernel_registry.get_adaln_modulation_op("cuda", hidden=7) + assert isinstance(op.implementation, NativeAdaLNModulationOp) + assert trace["rejected_backends"] == [ + OpBackend.CUDA_ADALN_MODULATION.name, + OpBackend.TRITON_ADALN_MODULATION.name, + ] + x, m = torch.randn(1, 2, 7), torch.randn(1, 21) + y, _ = op(x, m) + gold, _ = GtestAdaLNReference().forward_fp32(x, m) + tol = resolve_tolerance( + load_contract(), + judgment="forward_accuracy", + op_class="reduction", + dtype=torch.float32, + ) + torch.testing.assert_close(y, gold, atol=tol.atol, rtol=tol.rtol) + missing.add(OpBackend.PYTORCH_ADALN_MODULATION) + with pytest.raises(RuntimeError, match="No functional adaln_modulation"): + kernel_registry.get_adaln_modulation_op("cuda", hidden=7) diff --git a/tests/validation/operators/test_check_operator_cli.py b/tests/validation/operators/test_check_operator_cli.py new file mode 100644 index 000000000..427eead4d --- /dev/null +++ b/tests/validation/operators/test_check_operator_cli.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""User-visible backend observations in the validation CLI.""" + +import json +import os +import subprocess +import sys +from argparse import Namespace +from dataclasses import replace +from pathlib import Path + +import pytest +import torch + +from rl_engine.validation.operators import run_operator_suite +from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case +from tools.validation.operators.check_operator import _summarize + +ROOT = Path(__file__).resolve().parents[3] + + +@pytest.mark.parametrize( + "op,candidate,actual", [("adaln_modulation", "cuda", "pytorch"), ("logp", "pytorch", "unknown")] +) +def test_cli_text_and_json_report_observed_backend(op, candidate, actual): + if candidate == "cuda": + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + try: + CudaAdaLNModulationOp() + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + command = [ + sys.executable, + "tools/validation/operators/check_operator.py", + "--op", + op, + "--candidate", + candidate, + "--device", + "cpu", + "--batch", + "1", + "--seq", + "2", + "--normalized-dim", + "2", + "--vocab", + "7", + ] + results = [ + subprocess.run( + command + flags, + cwd=ROOT, + capture_output=True, + text=True, + env=dict(os.environ, PYTHONPATH=str(ROOT)), + timeout=30, + ) + for flags in ([], ["--json"]) + ] + for result in results: + assert result.returncode == 0, result.stdout + result.stderr + output = results[0].stdout + payload = json.loads(results[1].stdout) + assert payload["passed"] is True + assert payload["candidates"][0]["backend"] == candidate + assert f"requested_backend={candidate}" in output + assert f"actual_backend={actual}" in output + assert "passed=True pass_rate=1.0000" in output + trace = payload["candidates"][0]["cases"][0]["backend_trace"] + if op == "adaln_modulation": + assert trace["actual_backend"] == actual and trace["fallback"] is True + assert "fallback=True" in output + assert f"fallback_reason={trace['fallback_reason']}" in output + else: + assert trace is None + assert "fallback=unknown" in output + + +def test_text_reports_each_case_trace_independently(capsys): + args = Namespace( + op="adaln_modulation", + candidate="pytorch", + arch_key=None, + batch=1, + seq=2, + normalized_dim=2, + seed=386, + ) + candidate = make_candidate(args, observe_backend=True) + case = make_operator_case(args, torch.float32, torch.device("cpu")) + report = run_operator_suite("adaln_modulation", candidates=[candidate], cases=[case]) + result = report.candidates[0] + # Formatting consumes report snapshots, including observations from other devices. + traces = [ + dict(actual_backend="cuda", fallback=False, fallback_reason=None), + dict(actual_backend="pytorch", fallback=True, fallback_reason="case-domain"), + None, + dict(actual_backend=None, fallback=None), + ] + cases = [ + replace(result.cases[0], case_name=f"case-{i}", backend_trace=trace) + for i, trace in enumerate(traces) + ] + report = replace(report, candidates=[replace(result, backend="cuda", cases=cases)]) + _summarize(report) + output = capsys.readouterr().out + for i, expected in enumerate( + ( + "actual_backend=cuda fallback=False", + "actual_backend=pytorch fallback=True", + "actual_backend=unknown fallback=unknown", + "actual_backend=unknown fallback=unknown", + ) + ): + assert f"case=case-{i} {expected}" in output + assert "fallback_reason=case-domain" in output diff --git a/tests/validation/operators/test_op_checks.py b/tests/validation/operators/test_op_checks.py index 9c9b6999b..311a7042d 100644 --- a/tests/validation/operators/test_op_checks.py +++ b/tests/validation/operators/test_op_checks.py @@ -112,6 +112,68 @@ def test_logp_native_candidate_suite_passes(): assert all(case.passed for case in report.candidates[0].cases) +@pytest.mark.parametrize("check_grad", [False, True]) +def test_report_snapshots_generic_trace_without_interpreting_backend_policy(check_grad): + class TracedDouble: + def __init__(self): + self.trace = {} + + def __call__(self, x): + self.trace["sample"] = float(x[0].detach()) + return x * 2 + + op = TracedDouble() + cases = [ + OperatorCase( + name=f"sample-{value}", + op_class="elementwise", + dtype=torch.float32, + inputs={"x": torch.tensor([value])}, + gold_fn=lambda x: x * 2, + grad_input_names=("x",), + ) + for value in (1.0, 3.0) + ] + report = run_operator_suite( + "trace", + candidates=[CandidateSpec(name="double", fn=op)], + cases=cases, + check_grad=check_grad, + ) + assert report.passed + op.trace.clear() + result = report.to_dict()["candidates"][0] + assert result["backend_provenance"] is None + assert [case["backend_trace"] for case in result["cases"]] == [ + {"sample": 1.0}, + {"sample": 3.0}, + ] + assert all("backend_provenance" not in case for case in result["cases"]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_gpu_candidate_compares_with_cpu_gold_for_outputs_and_gradients(dtype): + x = torch.tensor([1.0, -2.0], device="cuda", dtype=dtype) + case = OperatorCase( + name="gpu-candidate-cpu-gold", + op_class="elementwise", + dtype=dtype, + inputs={"x": x}, + gold_fn=lambda x: x.cpu().float() * 2, + grad_input_names=("x",), + gold_input_dtype=torch.float32, + ) + report = run_operator_suite( + "cross-device", + candidates=[CandidateSpec(name="double", fn=lambda x: x * 2)], + cases=[case], + check_grad=True, + ) + assert report.passed + assert all(output.max_abs_error == 0 for output in report.candidates[0].cases[0].outputs) + + def test_embedding_native_candidate_suite_passes_issue_108_helper(): report = run_operator_suite( "embedding", diff --git a/tools/validation/operators/check_operator.py b/tools/validation/operators/check_operator.py index 288549235..ca13bd506 100644 --- a/tools/validation/operators/check_operator.py +++ b/tools/validation/operators/check_operator.py @@ -74,10 +74,20 @@ def _summarize(report: Any) -> None: print(f"suite={report.suite_name} passed={report.passed} pass_rate={report.pass_rate:.4f}") for candidate in report.candidates: print( - f"candidate={candidate.candidate_name} backend={candidate.backend} " + f"candidate={candidate.candidate_name} requested_backend={candidate.backend} " f"passed={candidate.passed} pass_rate={candidate.pass_rate:.4f}" ) for case in candidate.cases: + trace = case.backend_trace or {} + actual = trace.get("actual_backend") or "unknown" + fallback = trace.get("fallback") + if fallback is None: + fallback = "unknown" + reason = trace.get("fallback_reason") + print( + f" case={case.case_name} actual_backend={actual} fallback={fallback}" + + (f" fallback_reason={reason}" if reason else "") + ) for output in case.outputs: label = f" {output.message}" if output.message else "" print( @@ -99,8 +109,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--candidate", default="pytorch", - help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton, " - "ascend.", + help="Candidate backend to validate, for example pytorch, cuda, cuda-sm90, triton, ascend.", ) parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="fp32") parser.add_argument("--device", default="auto") @@ -144,7 +153,7 @@ def main() -> None: args = parse_args() dtype = _parse_dtype(args.dtype) device = _select_device(args.device) - candidate = make_candidate(args) + candidate = make_candidate(args, observe_backend=True) case = make_operator_case(args, dtype, device) report = run_operator_suite( args.op,