From d666d3116f84a8a8964732363bd2731ba9628303 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Tue, 22 Sep 2026 23:56:16 +0800 Subject: [PATCH 01/37] WIP #386: add AdaLN modulation WS1 implementations and tests Signed-off-by: hongleng-github --- benchmarks/benchmark_adaln_modulation.py | 83 ++++++ build_tools/extensions.py | 5 + csrc/bindings/ops.cpp | 14 + csrc/cuda/adaln_modulation.cu | 198 ++++++++++++++ docs/operators/adaln-modulation.md | 62 +++++ rl_engine/_C.pyi | 6 + .../kernels/ops/cuda/norm/adaln_modulation.py | 45 ++++ .../ops/pytorch/norm/adaln_modulation.py | 88 +++++++ .../ops/triton/norm/adaln_modulation.py | 134 ++++++++++ rl_engine/runtime/registry.py | 171 +++++++----- tests/test_adaln_modulation.py | 244 ++++++++++++++++++ 11 files changed, 985 insertions(+), 65 deletions(-) create mode 100644 benchmarks/benchmark_adaln_modulation.py create mode 100644 csrc/cuda/adaln_modulation.cu create mode 100644 docs/operators/adaln-modulation.md create mode 100644 rl_engine/kernels/ops/cuda/norm/adaln_modulation.py create mode 100644 rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py create mode 100644 rl_engine/kernels/ops/triton/norm/adaln_modulation.py create mode 100644 tests/test_adaln_modulation.py diff --git a/benchmarks/benchmark_adaln_modulation.py b/benchmarks/benchmark_adaln_modulation.py new file mode 100644 index 000000000..b09f06bed --- /dev/null +++ b/benchmarks/benchmark_adaln_modulation.py @@ -0,0 +1,83 @@ +# 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 platform +import statistics + +import torch + +from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp + + +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() + 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 statistics.median(samples), torch.cuda.max_memory_allocated() + + +def main(): + 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") + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("CUDA device required") + from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + backends = [("pytorch", NativeAdaLNModulationOp()), ("triton", TritonAdaLNModulationOp())] + try: + backends.append(("cuda", CudaAdaLNModulationOp())) + except RuntimeError: + print("cuda extension unavailable: explicit skip") + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 + print( + f"GPU={torch.cuda.get_device_name()} cuda_runtime={torch.version.cuda} " + f"torch={torch.__version__} python={platform.python_version()} " + f"dtype={args.dtype} warmup={args.warmup} repeat={args.repeat}" + ) + sizes = (4096, 6889, 6032) if args.real else (64,) + for seq in sizes: + 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) + for name, op in backends: + for backward in (False, True): + ms, peak = _bench( + op, x, modulation, backward=backward, warmup=args.warmup, + repeat=args.repeat, + ) + print( + f"S={seq} backend={name} backward={backward} " + f"median_ms={ms:.3f} peak_MB={peak / 2**20:.1f}" + ) + + +if __name__ == "__main__": + main() diff --git a/build_tools/extensions.py b/build_tools/extensions.py index db079e828..d1e554611 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/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/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index 0c7c3e7f4..6e2019bfe 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -256,6 +256,14 @@ 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); +std::vector adaln_modulation_backward_cuda( + torch::Tensor dy, torch::Tensor dg, torch::Tensor x, + torch::Tensor modulation, double eps); +#endif + // RMSNorm Declarations & Wrappers void rmsnorm_forward_cuda( @@ -737,6 +745,12 @@ 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"); + m.def("adaln_modulation_backward", &adaln_modulation_backward_cuda, + "AdaLN modulation backward CUDA"); +#endif // Deterministic standard-softmax attention (issue #147) m.def( diff --git a/csrc/cuda/adaln_modulation.cu b/csrc/cuda/adaln_modulation.cu new file mode 100644 index 000000000..c22703e36 --- /dev/null +++ b/csrc/cuda/adaln_modulation.cu @@ -0,0 +1,198 @@ +// 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 + +#if defined(__CUDA_FAST_MATH__) +#error "adaln_modulation requires CUDA compilation without --use_fast_math" +#endif + +namespace { + +constexpr int THREADS = 256; + +__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(); + } + return shared[0]; +} + +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; + 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))); + 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) { + extern __shared__ float shared[]; + const int row = blockIdx.x; + const int batch = row / seq; + 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))); + 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 shift = 0.0f, scale = 0.0f; +#pragma unroll 1 + for (int token = 0; token < seq; ++token) { + const int offset = (batch * seq + token) * hidden + col; + shift = __fadd_rn(shift, partial_shift[offset]); + scale = __fadd_rn(scale, partial_scale[offset]); + } + dm[batch * 3 * hidden + col] = static_cast(shift); + dm[batch * 3 * hidden + hidden + col] = static_cast(scale); + dm[batch * 3 * hidden + 2 * hidden + col] = dg[batch * hidden + col]; +} + +void check(const torch::Tensor& x, const torch::Tensor& modulation, double eps) { + 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(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.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"); +} + +} // namespace + +std::vector adaln_modulation_forward_cuda( + torch::Tensor x, torch::Tensor modulation, double eps) { + check(x, modulation, eps); + 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) { + check(x, modulation, eps); + 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}; +} diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md new file mode 100644 index 000000000..4a79127af --- /dev/null +++ b/docs/operators/adaln-modulation.md @@ -0,0 +1,62 @@ +# AdaLN modulation (Qwen-Image WS1, #386) + +`kernel_registry.get_adaln_modulation_op(device, hidden=H)` returns `(op, trace)`. Call +`modulated, gate = op(x, modulation, eps=1e-6)` with `x: [B,S,H]` +and shared `modulation: [B,3H]`. Both inputs must share device and dtype (FP32 +or BF16). The chunks are **shift, scale, gate**. Each row uses LayerNorm over +H with no affine parameters, then `norm(x) * (1 + scale) + shift`. The gate is +`[B,1,H]` and remains connected to `modulation` for downstream gradients. +The Qwen-Image image/text streams and attention/MLP calls use separate inputs. +The `SiLU + Linear` projection and `adaln_gate_residual` are separate operators. + +The FP32 hidden reductions use a lower-plus-upper pairwise tree padded to the +next power of two. Each sample's `dshift` and `dscale` contributions are folded +in ascending logical token order. No atomic partial sums, Split-K, Stream-K, +TF32, or fast math are used by these kernels. BF16 outputs are cast once after +FP32 arithmetic. The CPU reference computes square root in FP64 and rounds it +to FP32 because CPU `torch.sqrt(float32)` can differ by one ULP from CUDA's +correctly rounded square root. Hidden and token accumulators remain FP32. +Inputs with S=0 or H=0 and mixed dtype/device are rejected. +Strided model chunks are copied to contiguous storage inside +GPU wrappers. The CUDA kernel supports H<=4096; shape-aware dispatch selects +Triton for larger H. CUDA symbols are omitted from builds with +`KERNEL_ALIGN_USE_FAST_MATH=1`, making fallback explicit. Indexed +`modulate_index` is outside this shared-parameter +API pending the maintainer's answer on #386. + +The resolution matrix `{1024², 1328², 1664×928}` implies image sequence +lengths `{4096, 6889, 6032}` for VAE scale 8 and 2x2 latent packing. H=3072. +These are derived model shapes, not tensor shapes stated by the issue. + +## Validation + +Run `python -m pytest tests/test_adaln_modulation.py -q` for focused tests. +Set `RLK_ADALN_REAL_SHAPES=1` to include the three full image shapes. +Run `python benchmarks/benchmark_adaln_modulation.py --real` for the three +image shapes, with `--dtype fp32` for the second dtype. The benchmark prints +median forward and forward+backward latency plus peak allocated memory; +`--warmup` and `--repeat` control the sample count. It explicitly skips the +CUDA extension if its symbols are absent. The registry trace records selected +backend, fallback, reduction order, accumulator dtype, disabled algorithm +flags, and a semantic kernel version. The semantic version is not a hash of the +compiled binary. + +Local smoke environment (2026-09-22): RTX 3050 Ti (SM86, 4 GiB), driver +610.60, PyTorch 2.12.1+cu126, Triton 3.7.1, Python 3.12.12. The focused suite +passed 12 tests with 4 skips (three opt-in real-shape cases and the unavailable +CUDA extension). The three real-shape Triton BF16 forward/backward tests passed +with byte equality against the CPU reference when enabled. A smoke benchmark +at B=1, S=64, H=3072, BF16, 2 warmups and 3 +repeats measured medians (ms): PyTorch 6.723 forward / 26.161 forward+backward; +Triton 0.364 / 2.631. The CUDA extension was explicitly skipped because no +matching `nvcc` is installed. These numbers are a script smoke check, not the +required full resolution benchmark. + +**Unconfirmed acceptance:** #386 requires CPU-to-CUDA byte equality and CUDA +as the bit-level comparison backend. Existing repository tests for other ops +use a tolerance for CPU-to-GPU accuracy, but no such exception has been granted +for this operator. The test file checks FP32 accuracy against the independent +PyTorch LayerNorm expression, CPU-to-Triton byte equality for FP32/BF16 and the +three real image shapes, and exact batch-position invariance. +Do not claim cross-backend byte equality or full #386 acceptance until the CUDA +extension is built and compared on the required shapes and dtypes. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index fb3c4379c..8453a87da 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -256,6 +256,12 @@ def swiglu_backward( gate: torch.Tensor, up: torch.Tensor, ) -> list[torch.Tensor]: ... +def adaln_modulation_forward( + x: torch.Tensor, modulation: torch.Tensor, eps: float +) -> list[torch.Tensor]: ... +def adaln_modulation_backward( + dy: torch.Tensor, dg: torch.Tensor, x: torch.Tensor, modulation: torch.Tensor, eps: float +) -> list[torch.Tensor]: ... def rmsnorm_forward( x: torch.Tensor, weight: torch.Tensor, diff --git a/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py b/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py new file mode 100644 index 000000000..ffacc8f51 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py @@ -0,0 +1,45 @@ +# 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.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import _validate + + +class _AdaLNCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, modulation, eps): + y, gate = _C.adaln_modulation_forward(x, modulation, eps) + ctx.save_for_backward(x, modulation) + ctx.eps = eps + 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 + ) + return dx, dm, None + + +class CudaAdaLNModulationOp: + def __init__(self): + 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") + + def __call__(self, x, modulation, *, eps=1e-6): + return self.forward(x, modulation, eps=eps) + + def forward(self, x, modulation, *, eps=1e-6): + _validate(x, modulation, eps) + 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)) diff --git a/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py b/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py new file mode 100644 index 000000000..adf6159e8 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py @@ -0,0 +1,88 @@ +# 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 torch + +from rl_engine.kernels.ops.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 _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") + + +class _AdaLNReference(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, modulation: torch.Tensor, eps: float): + hidden = x.shape[-1] + shift, scale, gate = modulation.float().chunk(3, dim=-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) + norm = centered * rstd.unsqueeze(-1) + 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 + return y, gate[:, None, :].to(modulation.dtype) + + @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]), + torch.stack([reduce_rows_fp32(rows) for rows in dy * norm]), + grad_gate.float().squeeze(1), + ], + dim=-1, + ).to(ctx.mod_dtype) + return dx, dm, None + + +class NativeAdaLNModulationOp: + def __call__(self, x: torch.Tensor, modulation: torch.Tensor, *, eps: float = 1e-6): + return self.forward(x, modulation, eps=eps) + + def forward(self, x: torch.Tensor, modulation: torch.Tensor, *, eps: float = 1e-6): + _validate(x, modulation, eps) + return _AdaLNReference.apply(x, modulation, float(eps)) diff --git a/rl_engine/kernels/ops/triton/norm/adaln_modulation.py b/rl_engine/kernels/ops/triton/norm/adaln_modulation.py new file mode 100644 index 000000000..22129f0f8 --- /dev/null +++ b/rl_engine/kernels/ops/triton/norm/adaln_modulation.py @@ -0,0 +1,134 @@ +# 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.kernels.ops.pytorch.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) + return tl.sum(tl.where(col == 0, values, 0.0), 0) + + +@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) + 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)) + 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) + 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)) + 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): + batch, seq, hidden = x.shape + y = torch.empty_like(x) + gate = torch.empty((batch, 1, hidden), device=x.device, dtype=x.dtype) + _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, + ) + ctx.save_for_backward(x, modulation) + ctx.eps = eps + 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, + ) + _bwd_reduce[(batch, triton.cdiv(hidden, 128))]( + partial_shift, partial_scale, grad_gate.contiguous(), dm, seq, hidden, 128, + enable_fp_fusion=False, + ) + return dx, dm, None + + +class TritonAdaLNModulationOp: + def __call__(self, x, modulation, *, eps=1e-6): + return self.forward(x, modulation, eps=eps) + + def forward(self, x, modulation, *, eps=1e-6): + _validate(x, modulation, eps) + if not x.is_cuda: + raise RuntimeError("Triton AdaLN requires CUDA tensors") + return _AdaLNTriton.apply(x.contiguous(), modulation.contiguous(), float(eps)) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 7751d5901..3799d9543 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -67,138 +67,141 @@ def __getitem__(cls, name: str): class OpBackend(Enum, metaclass=_KernelEnumMeta): # NVIDIA optimized stack - FLASH_ATTN = "rl_engine.kernels.ops.cuda.attention.flash_attn.FlashAttentionOp" + FLASH_ATTN = "rl_engine.backends.cuda.attention.flash_attn.FlashAttentionOp" FLASHINFER = "rl_engine.kernels.ops.cuda.flashinfer.FlashInferOp" # TMA-accelerated LogP for SM90+ (Warp Specialization) - CUDA_FUSED_LOGP_SM90 = "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpSM90Op" - CUDA_FUSED_LOGP_GENERIC = "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp" + CUDA_FUSED_LOGP_SM90 = "rl_engine.backends.cuda.logprob.logp.FusedLogpSM90Op" + CUDA_FUSED_LOGP_GENERIC = "rl_engine.backends.cuda.logprob.logp.FusedLogpGenericOp" MUSA_FUSED_LOGP_GENERIC = "rl_engine.backends.musa.logprob.logp.MusaFusedLogpOp" - CUDA_DETERMINISTIC_LOGP = "rl_engine.kernels.ops.cuda.loss.logp.DeterministicLogpCUDAOp" + CUDA_DETERMINISTIC_LOGP = "rl_engine.backends.cuda.logprob.logp.DeterministicLogpCUDAOp" # Deterministic standard-softmax attention (issue #147); not FlashAttention. CUDA_DETERMINISTIC_ATTENTION = ( - "rl_engine.kernels.ops.cuda.attention.deterministic_attn.DeterministicAttentionOp" + "rl_engine.backends.cuda.attention.deterministic_attn.DeterministicAttentionOp" ) # AMD ROCm optimized stack ROCM_AITER = "rl_engine.kernels.ops.rocm.aiter.AiterOp" ROCM_CK = "rl_engine.kernels.ops.rocm.composable_kernel.CKOp" - ROCM_FLASH_ATTN = "rl_engine.kernels.ops.rocm.attention.flash_attn.RocmFlashAttentionOp" + ROCM_FLASH_ATTN = "rl_engine.backends.rocm.attention.flash_attn.RocmFlashAttentionOp" # WS2 strict ROCm attention core (AITER/CK dense MHA, Split-KV disabled). # Reachable only through ``get_attention_op``: it is not a drop-in for the # SDPA-shaped ``attn`` wrappers and must never be a silent fallback. ROCM_STRICT_ATTENTION = ( - "rl_engine.kernels.ops.rocm.attention.flash_attn.StrictRocmAiterCKAttentionCore" + "rl_engine.backends.rocm.attention.flash_attn.StrictRocmAiterCKAttentionCore" ) # GRPO loss (group reward normalization + clipped surrogate + KL) - TRITON_GRPO_LOSS = "rl_engine.kernels.ops.triton.loss.grpo_loss.TritonGRPOLossOp" - PYTORCH_GRPO_LOSS = "rl_engine.kernels.ops.pytorch.loss.grpo_loss.NativeGRPOLossOp" + TRITON_GRPO_LOSS = "rl_engine.backends.shared.triton.loss.grpo_loss.TritonGRPOLossOp" + PYTORCH_GRPO_LOSS = "rl_engine.reference.loss.grpo_loss.NativeGRPOLossOp" # Fused linear log-prob (hidden @ W^T -> selected-token logp, no [N, V] logits) CUDA_FUSED_LINEAR_LOGP_SM90 = ( - "rl_engine.kernels.ops.cuda.loss.linear_logp.FusedLinearLogpSM90Op" + "rl_engine.backends.cuda.logprob.linear_logp.FusedLinearLogpSM90Op" ) - TRITON_LINEAR_LOGP = "rl_engine.kernels.ops.triton.loss.linear_logp.TritonLinearLogpOp" - PYTORCH_LINEAR_LOGP = "rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp" + TRITON_LINEAR_LOGP = "rl_engine.backends.shared.triton.logprob.linear_logp.TritonLinearLogpOp" + PYTORCH_LINEAR_LOGP = "rl_engine.reference.logprob.linear_logp.NativeLinearLogpOp" # Fused policy-ratio + KL-penalty front-end (PPO/GRPO), logits -> (ratio, kl) - TRITON_RATIO_KL = "rl_engine.kernels.ops.triton.loss.ratio_kl.TritonRatioKLOp" - PYTORCH_RATIO_KL = "rl_engine.kernels.ops.pytorch.loss.ratio_kl.NativeRatioKLOp" + TRITON_RATIO_KL = "rl_engine.backends.shared.triton.loss.ratio_kl.TritonRatioKLOp" + PYTORCH_RATIO_KL = "rl_engine.reference.loss.ratio_kl.NativeRatioKLOp" # Variable-length packing (pack-and-pad), [B,S,...] -> [Total_Active,...] - PYTORCH_PACK = "rl_engine.kernels.ops.pytorch.packing.pack.NativePackOp" + PYTORCH_PACK = "rl_engine.reference.packing.pack.NativePackOp" # Batch-invariant deterministic GEMM (WS1 #146) - CUDA_DET_GEMM = "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp" - TRITON_DET_GEMM = "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp" - ASCEND_DET_GEMM = "rl_engine.kernels.ops.ascend.matmul.det_gemm.DetGemmAscendOp" + CUDA_DET_GEMM = "rl_engine.backends.cuda.gemm.det_gemm.DetGemmOp" + TRITON_DET_GEMM = "rl_engine.backends.shared.triton.gemm.det_gemm.TritonDetGemmOp" + ASCEND_DET_GEMM = "rl_engine.backends.ascend.gemm.det_gemm.DetGemmAscendOp" # NON-deterministic reference (torch.matmul); reference/benchmark ONLY, # intentionally excluded from det_gemm dispatch (cuBLAS breaks invariance). - PYTORCH_GEMM = "rl_engine.kernels.ops.pytorch.matmul.det_gemm.NativeGemmOp" + PYTORCH_GEMM = "rl_engine.reference.gemm.det_gemm.NativeGemmOp" # Batch-invariant selected-logprob (WS1 #148: locked reduction order) TRITON_BATCH_INVARIANT_LOGP = ( - "rl_engine.kernels.ops.triton.loss.batch_invariant_logp.TritonBatchInvariantLogpOp" + "rl_engine.backends.shared.triton.logprob.batch_invariant_logp.TritonBatchInvariantLogpOp" ) PYTORCH_BATCH_INVARIANT_LOGP = ( - "rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp.NativeBatchInvariantLogpOp" + "rl_engine.reference.logprob.batch_invariant_logp.NativeBatchInvariantLogpOp" ) CUDA_BATCH_INVARIANT_LOGP_SM90 = ( - "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90Op" + "rl_engine.backends.cuda.logprob.batch_invariant_logp.BatchInvariantLogpSM90Op" ) ASCEND_BATCH_INVARIANT_LOGP = ( - "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp" + "rl_engine.backends.ascend.logprob.batch_invariant_logp.BatchInvariantLogpAscendOp" ) - ASCEND_RMS_NORM = "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp" - ASCEND_EMBEDDING = "rl_engine.kernels.ops.ascend.linear.embedding.AscendEmbeddingOp" - ASCEND_FUSED_LOGP = "rl_engine.kernels.ops.ascend.loss.logp.FusedLogpAscendOp" - ASCEND_LM_HEAD = "rl_engine.kernels.ops.ascend.linear.lm_head.AscendLMHeadOp" + ASCEND_RMS_NORM = "rl_engine.backends.ascend.norm.rmsnorm.RMSNormAscendOp" + ASCEND_EMBEDDING = "rl_engine.backends.ascend.embedding.embedding.AscendEmbeddingOp" + ASCEND_FUSED_LOGP = "rl_engine.backends.ascend.logprob.logp.FusedLogpAscendOp" + ASCEND_LM_HEAD = "rl_engine.backends.ascend.embedding.lm_head.AscendLMHeadOp" ASCEND_FUSED_LINEAR_LOGP = ( - "rl_engine.kernels.ops.ascend.loss.linear_logp.FusedLinearLogpAscendOp" + "rl_engine.backends.ascend.logprob.linear_logp.FusedLinearLogpAscendOp" ) # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) PYTORCH_VOCAB_PARALLEL_LOGP = ( - "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" + "rl_engine.reference.logprob.vocab_parallel_logp.VocabParallelLogprobOp" ) ROCM_VOCAB_PARALLEL_LOGP = ( - "rl_engine.kernels.ops.rocm.loss.vocab_parallel_logp.RocmVocabParallelLogprobOp" + "rl_engine.backends.rocm.logprob.vocab_parallel_logp.RocmVocabParallelLogprobOp" ) TRITON_VOCAB_PARALLEL_LOGP = ( - "rl_engine.kernels.ops.triton.loss.vocab_parallel_logp.TritonVocabParallelLogprobOp" + "rl_engine.backends.shared.triton.logprob.vocab_parallel_logp.TritonVocabParallelLogprobOp" ) # Deterministic GRPO loss on the TP-aware logprob path (WS2 #241 PR5) PYTORCH_DISTRIBUTED_GRPO_LOSS = ( - "rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss.DistributedGRPOLossOp" + "rl_engine.reference.loss.distributed_grpo_loss.DistributedGRPOLossOp" ) # Ascend NPU deterministic (batch-invariant, no split-K) standard-softmax # attention (issue #147); Ascend C forward + reference backward. ASCEND_DETERMINISTIC_ATTENTION = ( - "rl_engine.kernels.ops.ascend.attention.deterministic_attn.DeterministicAttentionAscendOp" + "rl_engine.backends.ascend.attention.deterministic_attn.DeterministicAttentionAscendOp" ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) - TRITON_RMS_NORM = "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp" - PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" + TRITON_RMS_NORM = "rl_engine.backends.shared.triton.norm.rmsnorm.RMSNormTritonOp" + PYTORCH_NATIVE_RMS_NORM = "rl_engine.reference.norm.rms_norm.NativeRMSNormOp" # Generic fallback TRITON_GENERIC = "rl_engine.kernels.ops.triton.generic.TritonOp" PYTORCH_ATTN = "rl_engine.kernels.ops.pytorch.attention.NativeAttentionOp" - TRITON_LOGP = "rl_engine.kernels.ops.triton.loss.logp.TritonLogpOp" - PYTORCH_NATIVE = "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp" - PYTORCH_NATIVE_MATMUL = "rl_engine.kernels.ops.pytorch.linear.matmul.NativeMatmulOp" - PYTORCH_NATIVE_ROPE = "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp" - TRITON_ROPE = "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp" - CUDA_ROPE_SM90 = "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op" - ASCEND_ROPE = "rl_engine.kernels.ops.ascend.rotary_embedding.rope.RoPEAscendOp" - PYTORCH_NATIVE_SILU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp" - PYTORCH_NATIVE_SWIGLU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp" - CUDA_SILU = "rl_engine.kernels.ops.cuda.activation.swiglu.SiLUCudaOp" - CUDA_SWIGLU = "rl_engine.kernels.ops.cuda.activation.swiglu.SwiGLUCudaOp" - ASCEND_SWIGLU = "rl_engine.kernels.ops.ascend.activation.swiglu.SwiGLUAscendOp" - ASCEND_SILU = "rl_engine.kernels.ops.ascend.activation.silu.SiLUAscendOp" - TRITON_SILU = "rl_engine.kernels.ops.triton.activation.swiglu.TritonSiLUOp" - TRITON_SWIGLU = "rl_engine.kernels.ops.triton.activation.swiglu.TritonSwiGLUOp" + TRITON_LOGP = "rl_engine.backends.shared.triton.logprob.logp.TritonLogpOp" + PYTORCH_NATIVE = "rl_engine.reference.logprob.logp.NativeLogpOp" + PYTORCH_NATIVE_MATMUL = "rl_engine.reference.embedding.matmul.NativeMatmulOp" + PYTORCH_NATIVE_ROPE = "rl_engine.reference.rope.rope.NativeRoPEOp" + TRITON_ROPE = "rl_engine.backends.shared.triton.rope.rope.TritonRoPEOp" + CUDA_ROPE_SM90 = "rl_engine.backends.cuda.rope.rope.RoPESM90Op" + ASCEND_ROPE = "rl_engine.backends.ascend.rope.rope.RoPEAscendOp" + PYTORCH_NATIVE_SILU = "rl_engine.reference.activation.swiglu.NativeSiLUOp" + PYTORCH_NATIVE_SWIGLU = "rl_engine.reference.activation.swiglu.NativeSwiGLUOp" + CUDA_SILU = "rl_engine.backends.cuda.activation.swiglu.SiLUCudaOp" + CUDA_SWIGLU = "rl_engine.backends.cuda.activation.swiglu.SwiGLUCudaOp" + ASCEND_SWIGLU = "rl_engine.backends.ascend.activation.swiglu.SwiGLUAscendOp" + ASCEND_SILU = "rl_engine.backends.ascend.activation.silu.SiLUAscendOp" + TRITON_SILU = "rl_engine.backends.shared.triton.activation.swiglu.TritonSiLUOp" + TRITON_SWIGLU = "rl_engine.backends.shared.triton.activation.swiglu.TritonSwiGLUOp" + PYTORCH_ADALN_MODULATION = ( + "rl_engine.kernels.ops.pytorch.norm.adaln_modulation.NativeAdaLNModulationOp" + ) + TRITON_ADALN_MODULATION = ( + "rl_engine.kernels.ops.triton.norm.adaln_modulation.TritonAdaLNModulationOp" + ) + CUDA_ADALN_MODULATION = "rl_engine.kernels.ops.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 = ( - "rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp" - ) + PYTORCH_NATIVE_ATTENTION = "rl_engine.reference.attention.standard_attn.NativeAttentionOp" # WS1 pure-PyTorch ground-truth KV-cache (decode/incremental) attention # reference; concats cache+new then reuses the standard attention reduction. - PYTORCH_NATIVE_KV_CACHE_ATTN = ( - "rl_engine.kernels.ops.pytorch.attention.kv_cache.NativeKVCacheAttnOp" - ) + PYTORCH_NATIVE_KV_CACHE_ATTN = "rl_engine.reference.attention.kv_cache.NativeKVCacheAttnOp" # WS2 correctness-first context-parallel attention reference. It emulates # CP prefill/chunked-prefill with fp32 attention-domain LSE merges. PYTORCH_CP_ATTENTION = ( - "rl_engine.kernels.ops.pytorch.attention.cp_attention.DeterministicCPAttentionReferenceOp" + "rl_engine.reference.attention.cp_attention.DeterministicCPAttentionReferenceOp" ) # WS1 pure-PyTorch ground-truth linear ops - PYTORCH_NATIVE_LM_HEAD = "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp" + PYTORCH_NATIVE_LM_HEAD = "rl_engine.reference.embedding.lm_head.NativeLMHeadOp" # WS1 pure-PyTorch ground-truth embedding ops - TRITON_EMBEDDING = "rl_engine.kernels.ops.triton.linear.embedding.TritonEmbeddingOp" - PYTORCH_NATIVE_EMBEDDING = "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp" - CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp" - CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" + TRITON_EMBEDDING = "rl_engine.backends.shared.triton.embedding.embedding.TritonEmbeddingOp" + PYTORCH_NATIVE_EMBEDDING = "rl_engine.reference.embedding.embedding.NativeEmbeddingOp" + CUDA_SM90_LM_HEAD = "rl_engine.backends.cuda.embedding.lm_head.SM90LMHeadOp" + CUDA_SM90_EMBEDDING = "rl_engine.backends.cuda.embedding.embedding.SM90EmbeddingOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -227,9 +230,7 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "strict_observable": True, }, lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, - implementation_class_or_factory=( - "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp" - ), + implementation_class_or_factory=("rl_engine.reference.logprob.logp.NativeLogpOp"), fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint="NativeLogpOp-selected-logprob-v1", ), @@ -602,6 +603,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 +655,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 +704,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 +730,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 +1057,37 @@ 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 + ) -> tuple[Any, dict[str, Any]]: + """Resolve AdaLN with a trace that makes backend fallback visible.""" + if hidden <= 0: + raise ValueError("hidden must be positive") + platform = self._platform_for_device(device) + 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 + 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}", + } + return op, 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/tests/test_adaln_modulation.py b/tests/test_adaln_modulation.py new file mode 100644 index 000000000..143718382 --- /dev/null +++ b/tests/test_adaln_modulation.py @@ -0,0 +1,244 @@ +import os + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp +from rl_engine.kernels.registry import OpBackend, kernel_registry + + +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_registry_trace_reports_selected_backend(): + op, trace = kernel_registry.get_adaln_modulation_op(device="cpu", hidden=4) + assert isinstance(op, 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_cuda_registry_reports_real_fallback(): + try: + torch.cuda.init() + except RuntimeError: + pytest.skip("CUDA required") + op, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=32) + assert type(op).__name__ in {"CudaAdaLNModulationOp", "TritonAdaLNModulationOp"} + if type(op).__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(): + op, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=5000) + assert type(op).__name__ == "TritonAdaLNModulationOp" + assert trace["fallback"] is True + assert "CUDA_ADALN_MODULATION: H > 4096" in trace["rejected_backends"] + + +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() + torch.testing.assert_close(y, y_ref, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(x.grad, x_ref.grad, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(modulation.grad, modulation_ref.grad, atol=1e-6, rtol=1e-6) + + +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 torch.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 torch.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 torch.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): + try: + torch.cuda.init() + except RuntimeError: + pytest.skip("CUDA required") + from rl_engine.kernels.ops.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 torch.equal(y, y_ref) + assert torch.equal(x.grad, x_ref.grad) + assert torch.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 torch.equal(y_new[1], y[0]) + assert torch.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 torch.equal(x_new.grad[1], x.grad[0]) + assert torch.equal(mod_new.grad[1], modulation.grad[0]) + + +def test_triton_irregular_hidden_and_bf16_tie_rounding(): + from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + try: + torch.cuda.init() + except RuntimeError: + pytest.skip("CUDA required") + 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 torch.equal(y.cpu(), y_cpu) + assert torch.equal(x.grad.cpu(), x_cpu.grad) + assert torch.equal(modulation.grad.cpu(), modulation_cpu.grad) + + +def test_cuda_adaln_bit_equality_harness(): + from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + try: + torch.cuda.init() + cuda_op = CudaAdaLNModulationOp() + except RuntimeError: + pytest.skip("compiled CUDA AdaLN extension required") + torch.manual_seed(390) + x = torch.randn(1, 3, 32, requires_grad=True) + modulation = torch.randn(1, 96, requires_grad=True) + upstream = torch.randn_like(x) + gate_upstream = torch.randn(1, 1, 32) + 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 torch.equal(cuda_y.cpu(), cpu_y) + assert torch.equal(cuda_gate.cpu(), cpu_gate) + assert torch.equal(x_gpu.grad.cpu(), x.grad) + assert torch.equal(mod_gpu.grad.cpu(), modulation.grad) + + +@pytest.mark.parametrize("seq", [4096, 6889, 6032]) +@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): + try: + torch.cuda.init() + except RuntimeError: + pytest.skip("CUDA required") + op, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=3072) + assert trace["selected_backend"] in { + OpBackend.CUDA_ADALN_MODULATION.name, + OpBackend.TRITON_ADALN_MODULATION.name, + } + x = torch.randn(1, seq, 3072, device="cuda", dtype=torch.bfloat16, requires_grad=True) + modulation = torch.randn(1, 9216, device="cuda", dtype=torch.bfloat16, requires_grad=True) + y, gate = op(x, modulation) + assert y.shape == x.shape and gate.shape == (1, 1, 3072) + (y.float().sum() + gate.float().sum()).backward() + 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) + (cpu_y.float().sum() + cpu_gate.float().sum()).backward() + assert torch.equal(y.cpu(), cpu_y) + assert torch.equal(gate.cpu(), cpu_gate) + assert torch.equal(x.grad.cpu(), cpu_x.grad) + assert torch.equal(modulation.grad.cpu(), cpu_modulation.grad) From da269ee30a3c560486fd0bb1272883feccd36948 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Tue, 22 Sep 2026 23:59:47 +0800 Subject: [PATCH 02/37] test: cover AdaLN padding and batch invariance Signed-off-by: hongleng-github --- tests/test_adaln_modulation.py | 39 ++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tests/test_adaln_modulation.py b/tests/test_adaln_modulation.py index 143718382..54048f483 100644 --- a/tests/test_adaln_modulation.py +++ b/tests/test_adaln_modulation.py @@ -186,6 +186,45 @@ def test_triton_irregular_hidden_and_bf16_tie_rounding(): assert torch.equal(modulation.grad.cpu(), modulation_cpu.grad) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_triton_adaln_padding_and_batch_size_do_not_change_valid_bytes(dtype): + from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + try: + torch.cuda.init() + except RuntimeError: + pytest.skip("CUDA required") + 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) + op = TritonAdaLNModulationOp() + + 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 torch.equal(y[-1, :3], baseline_y[0]) + assert torch.equal(gate[-1], baseline_gate[0]) + assert torch.equal(padded_x.grad[-1, :3], baseline_x.grad[0]) + assert torch.equal(padded_modulation.grad[-1], baseline_modulation.grad[0]) + + def test_cuda_adaln_bit_equality_harness(): from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp From 28ea0577e5fc7be13886183dcb0cab47cfbacb51 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 23 Sep 2026 15:24:57 +0800 Subject: [PATCH 03/37] test: validate AdaLN native CUDA acceptance Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 51 +++++++++----- .../01_install_cuda_12_6_wsl.sh | 19 +++++ .../02_verify_cuda_12_6.sh | 17 +++++ .../03_build_extension.sh | 23 +++++++ .../04_test_adaln_cuda.sh | 19 +++++ .../05_benchmark_adaln_cuda.sh | 15 ++++ tests/test_adaln_modulation.py | 69 ++++++++++++------- 7 files changed, 169 insertions(+), 44 deletions(-) create mode 100755 scripts/adaln_cuda_acceptance/01_install_cuda_12_6_wsl.sh create mode 100755 scripts/adaln_cuda_acceptance/02_verify_cuda_12_6.sh create mode 100755 scripts/adaln_cuda_acceptance/03_build_extension.sh create mode 100755 scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh create mode 100755 scripts/adaln_cuda_acceptance/05_benchmark_adaln_cuda.sh diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 4a79127af..ae7ee6826 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -41,22 +41,35 @@ backend, fallback, reduction order, accumulator dtype, disabled algorithm flags, and a semantic kernel version. The semantic version is not a hash of the compiled binary. -Local smoke environment (2026-09-22): RTX 3050 Ti (SM86, 4 GiB), driver -610.60, PyTorch 2.12.1+cu126, Triton 3.7.1, Python 3.12.12. The focused suite -passed 12 tests with 4 skips (three opt-in real-shape cases and the unavailable -CUDA extension). The three real-shape Triton BF16 forward/backward tests passed -with byte equality against the CPU reference when enabled. A smoke benchmark -at B=1, S=64, H=3072, BF16, 2 warmups and 3 -repeats measured medians (ms): PyTorch 6.723 forward / 26.161 forward+backward; -Triton 0.364 / 2.631. The CUDA extension was explicitly skipped because no -matching `nvcc` is installed. These numbers are a script smoke check, not the -required full resolution benchmark. - -**Unconfirmed acceptance:** #386 requires CPU-to-CUDA byte equality and CUDA -as the bit-level comparison backend. Existing repository tests for other ops -use a tolerance for CPU-to-GPU accuracy, but no such exception has been granted -for this operator. The test file checks FP32 accuracy against the independent -PyTorch LayerNorm expression, CPU-to-Triton byte equality for FP32/BF16 and the -three real image shapes, and exact batch-position invariance. -Do not claim cross-backend byte equality or full #386 acceptance until the CUDA -extension is built and compared on the required shapes and dtypes. +Local CUDA validation environment (2026-09-23): RTX 3050 Ti (SM86, 4 GiB), +driver 610.60, CUDA toolkit/runtime 12.6, PyTorch 2.12.1+cu126, Triton 3.7.1, +Python 3.12.12. The extension was built with fast math disabled for SM86. +The focused suite plus all three real shapes passed 18 tests with no skips. +FP32 small-shape and BF16 real-shape CUDA forward/backward outputs were byte +equal to the CPU reference. CUDA and Triton were also byte equal directly at +all three real shapes, and both passed FP32/BF16 batch-size and token-padding +byte-invariance checks. + +Real-shape medians below use one warmup and three samples, so they are smoke +measurements rather than stable performance claims: + +| dtype | S | Triton fwd / fwd+bwd (ms) | CUDA fwd / fwd+bwd (ms) | +| --- | ---: | ---: | ---: | +| BF16 | 4096 | 3.069 / 20.709 | 2.998 / 23.262 | +| BF16 | 6889 | 5.125 / 34.551 | 5.150 / 39.963 | +| BF16 | 6032 | 4.689 / 30.302 | 4.357 / 34.718 | +| FP32 | 4096 | 5.944 / 29.269 | 7.229 / 38.693 | +| FP32 | 6889 | 9.840 / 48.579 | 11.948 / 65.271 | +| FP32 | 6032 | 8.685 / 81.903 | 10.639 / 57.270 | + +Forward and forward+backward peak allocation matched between CUDA and Triton +at each shape. The CUDA backward path was generally slower in this short SM86 +run; rerun with the default warmup/repeat counts on target hardware before +making a performance claim. + +**Remaining acceptance boundaries:** the maintainer has not answered whether +the indexed modulate_index path may stay on fallback. The trace fingerprint +is a semantic kernel version rather than a compiled-binary hash. The exactness +evidence above covers FP32 at the focused shape and BF16 at the three real image +shapes; a broader CUDA shape matrix remains optional evidence, not a result +claimed here. diff --git a/scripts/adaln_cuda_acceptance/01_install_cuda_12_6_wsl.sh b/scripts/adaln_cuda_acceptance/01_install_cuda_12_6_wsl.sh new file mode 100755 index 000000000..5a492deff --- /dev/null +++ b/scripts/adaln_cuda_acceptance/01_install_cuda_12_6_wsl.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +set -euo pipefail + +grep -qi microsoft /proc/version || { echo "WSL2 is required" >&2; exit 1; } +if [[ -x /usr/local/cuda-12.6/bin/nvcc ]] && + /usr/local/cuda-12.6/bin/nvcc --version | grep -q 'release 12\.6'; then + echo "CUDA Toolkit 12.6 is already installed" + exit 0 +fi + +read -r -p "Install the large cuda-toolkit-12-6 package (no Linux driver)? [y/N] " reply +[[ "$reply" =~ ^[Yy]$ ]] || exit 0 + +deb=$(mktemp --suffix=.deb) +trap 'rm -f "$deb"' EXIT +curl -fL https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-keyring_1.1-1_all.deb -o "$deb" +sudo dpkg -i "$deb" +sudo apt-get update +sudo apt-get install -y cuda-toolkit-12-6 diff --git a/scripts/adaln_cuda_acceptance/02_verify_cuda_12_6.sh b/scripts/adaln_cuda_acceptance/02_verify_cuda_12_6.sh new file mode 100755 index 000000000..8b18e3994 --- /dev/null +++ b/scripts/adaln_cuda_acceptance/02_verify_cuda_12_6.sh @@ -0,0 +1,17 @@ +#!/usr/bin/env bash +set -euo pipefail + +root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) +cuda_home=${CUDA_HOME:-/usr/local/cuda-12.6} +nvcc="$cuda_home/bin/nvcc" + +[[ -x "$nvcc" ]] || { echo "Missing $nvcc" >&2; exit 1; } +"$nvcc" --version | grep 'release 12\.6' +[[ -e /usr/lib/wsl/lib/libcuda.so ]] || { echo "Missing WSL libcuda.so" >&2; exit 1; } +"$root/.venv/bin/python" - <<'PY' +import torch +assert torch.version.cuda == "12.6", torch.version.cuda +assert torch.cuda.is_available() +assert torch.cuda.get_device_capability() == (8, 6) +print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name()) +PY diff --git a/scripts/adaln_cuda_acceptance/03_build_extension.sh b/scripts/adaln_cuda_acceptance/03_build_extension.sh new file mode 100755 index 000000000..e60258c17 --- /dev/null +++ b/scripts/adaln_cuda_acceptance/03_build_extension.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +set -euo pipefail + +root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) +cd "$root" +export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.6} +export PATH="$CUDA_HOME/bin:$PATH" +export LIBRARY_PATH="/usr/lib/wsl/lib:${LIBRARY_PATH:-}" +export LD_LIBRARY_PATH="/usr/lib/wsl/lib:$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}" +export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST:-8.6} +export MAX_JOBS=${MAX_JOBS:-2} +export FORCE_CUDA=1 +export RL_KERNEL_REQUIRE_EXT=1 +export KERNEL_ALIGN_USE_FAST_MATH=0 + +uv pip install --python .venv/bin/python ninja +.venv/bin/python setup.py build_ext --inplace +.venv/bin/python - <<'PY' +import rl_engine._C as ext +for name in ("adaln_modulation_forward", "adaln_modulation_backward"): + assert hasattr(ext, name), name +print("AdaLN CUDA extension symbols verified") +PY diff --git a/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh b/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh new file mode 100755 index 000000000..4fb604081 --- /dev/null +++ b/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +set -euo pipefail + +root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) +cd "$root" +export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.6} +export LD_LIBRARY_PATH="/usr/lib/wsl/lib:$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}" +export RL_KERNEL_REQUIRE_EXT=1 +export RLK_ADALN_REAL_SHAPES=1 +export PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 + +.venv/bin/python - <<'PY' +from rl_engine.kernels.registry import OpBackend, kernel_registry +_, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=3072) +assert trace["selected_backend"] == OpBackend.CUDA_ADALN_MODULATION.name, trace +assert trace["fallback"] is False, trace +print(trace) +PY +.venv/bin/python -m pytest tests/test_extension_smoke.py tests/test_adaln_modulation.py -q -rs diff --git a/scripts/adaln_cuda_acceptance/05_benchmark_adaln_cuda.sh b/scripts/adaln_cuda_acceptance/05_benchmark_adaln_cuda.sh new file mode 100755 index 000000000..f46294cbf --- /dev/null +++ b/scripts/adaln_cuda_acceptance/05_benchmark_adaln_cuda.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +set -euo pipefail + +root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) +cd "$root" +.venv/bin/python - <<'PYCODE' +from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +CudaAdaLNModulationOp() +print("AdaLN CUDA extension verified") +PYCODE +warmup=${WARMUP:-1} +repeat=${REPEAT:-3} + +.venv/bin/python benchmarks/benchmark_adaln_modulation.py --real --dtype bf16 --warmup "$warmup" --repeat "$repeat" +.venv/bin/python benchmarks/benchmark_adaln_modulation.py --real --dtype fp32 --warmup "$warmup" --repeat "$repeat" diff --git a/tests/test_adaln_modulation.py b/tests/test_adaln_modulation.py index 54048f483..f3caa0ef3 100644 --- a/tests/test_adaln_modulation.py +++ b/tests/test_adaln_modulation.py @@ -187,7 +187,8 @@ def test_triton_irregular_hidden_and_bf16_tie_rounding(): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_triton_adaln_padding_and_batch_size_do_not_change_valid_bytes(dtype): +def test_adaln_gpu_padding_and_batch_size_do_not_change_valid_bytes(dtype): + from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp try: @@ -199,30 +200,35 @@ def test_triton_adaln_padding_and_batch_size_do_not_change_valid_bytes(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) - op = TritonAdaLNModulationOp() - - 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 torch.equal(y[-1, :3], baseline_y[0]) - assert torch.equal(gate[-1], baseline_gate[0]) - assert torch.equal(padded_x.grad[-1, :3], baseline_x.grad[0]) - assert torch.equal(padded_modulation.grad[-1], baseline_modulation.grad[0]) + 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 torch.equal(y[-1, :3], baseline_y[0]) + assert torch.equal(gate[-1], baseline_gate[0]) + assert torch.equal(padded_x.grad[-1, :3], baseline_x.grad[0]) + assert torch.equal(padded_modulation.grad[-1], baseline_modulation.grad[0]) def test_cuda_adaln_bit_equality_harness(): @@ -262,6 +268,8 @@ def test_adaln_real_image_shapes_forward_backward(seq): except RuntimeError: 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, @@ -281,3 +289,14 @@ def test_adaln_real_image_shapes_forward_backward(seq): assert torch.equal(gate.cpu(), cpu_gate) assert torch.equal(x.grad.cpu(), cpu_x.grad) assert torch.equal(modulation.grad.cpu(), cpu_modulation.grad) + + from rl_engine.kernels.ops.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) + (triton_y.float().sum() + triton_gate.float().sum()).backward() + assert torch.equal(triton_y, y) + assert torch.equal(triton_gate, gate) + assert torch.equal(triton_x.grad, x.grad) + assert torch.equal(triton_modulation.grad, modulation.grad) From 649dcd7f73e65e44143934a6b8680e7f1308247b Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 23 Sep 2026 16:00:48 +0800 Subject: [PATCH 04/37] feat: add explicit AdaLN indexed fallback Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 24 ++++-- .../ops/pytorch/norm/adaln_modulation.py | 85 ++++++++++++++++--- rl_engine/runtime/registry.py | 23 ++++- tests/test_adaln_modulation.py | 33 +++++++ 4 files changed, 147 insertions(+), 18 deletions(-) diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index ae7ee6826..6570ad743 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -20,9 +20,22 @@ Inputs with S=0 or H=0 and mixed dtype/device are rejected. Strided model chunks are copied to contiguous storage inside GPU wrappers. The CUDA kernel supports H<=4096; shape-aware dispatch selects Triton for larger H. CUDA symbols are omitted from builds with -`KERNEL_ALIGN_USE_FAST_MATH=1`, making fallback explicit. Indexed -`modulate_index` is outside this shared-parameter -API pending the maintainer's answer on #386. +`KERNEL_ALIGN_USE_FAST_MATH=1`, making fallback explicit. + +For Qwen-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) +``` + +This path accepts modulation `[2B,3H]` and an index `[B,S]` or `[1,S]` whose +values are 0 or 1. It uses the explicit PyTorch select01 reference and returns +gate `[B,S,H]`. The trace sets `fallback=true` and +`fallback_reason=modulate_index_requires_select01`; the indexed path does not +claim the fixed-order CUDA/Triton backward contract. The resolution matrix `{1024², 1328², 1664×928}` implies image sequence lengths `{4096, 6889, 6032}` for VAE scale 8 and 2x2 latent packing. H=3072. @@ -68,8 +81,9 @@ run; rerun with the default warmup/repeat counts on target hardware before making a performance claim. **Remaining acceptance boundaries:** the maintainer has not answered whether -the indexed modulate_index path may stay on fallback. The trace fingerprint -is a semantic kernel version rather than a compiled-binary hash. The exactness +the indexed `modulate_index` path must be optimized beyond its explicit +PyTorch fallback. The trace fingerprint is a semantic kernel version rather +than a compiled-binary hash. The exactness evidence above covers FP32 at the focused shape and BF16 at the three real image shapes; a broader CUDA shape matrix remains optional evidence, not a result claimed here. diff --git a/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py b/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py index adf6159e8..6d4fe5463 100644 --- a/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py +++ b/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py @@ -23,6 +23,18 @@ def _sum_hidden(values: torch.Tensor) -> torch.Tensor: 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 _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]") @@ -39,16 +51,8 @@ def _validate(x: torch.Tensor, modulation: torch.Tensor, eps: float) -> None: class _AdaLNReference(torch.autograd.Function): @staticmethod def forward(ctx, x: torch.Tensor, modulation: torch.Tensor, eps: float): - hidden = x.shape[-1] shift, scale, gate = modulation.float().chunk(3, dim=-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) - norm = centered * rstd.unsqueeze(-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 @@ -79,10 +83,67 @@ def backward(ctx, grad_y: torch.Tensor, grad_gate: torch.Tensor): return dx, dm, None +def _select01_reference( + x: torch.Tensor, + modulation: torch.Tensor, + modulate_index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + 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 torch.all((modulate_index == 0) | (modulate_index == 1)).item(): + raise ValueError("modulate_index values must be 0 or 1") + if not math.isfinite(eps) or eps < 0: + raise ValueError("eps must be finite and nonnegative") + + 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, *, eps: float = 1e-6): - return self.forward(x, modulation, eps=eps) + 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, *, eps: float = 1e-6): + 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/registry.py b/rl_engine/runtime/registry.py index 3799d9543..669f22405 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -1058,12 +1058,33 @@ 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 + self, + device: torch.device | str | None = None, + *, + hidden: int, + modulate_index: torch.Tensor | None = None, ) -> tuple[Any, dict[str, Any]]: """Resolve AdaLN with a trace that makes backend fallback visible.""" if hidden <= 0: raise ValueError("hidden must be positive") platform = self._platform_for_device(device) + if modulate_index is not None: + backend = OpBackend.PYTORCH_ADALN_MODULATION + op = self._get_or_create_backend(backend) + if op is None: + raise RuntimeError("PyTorch select01 AdaLN fallback is unavailable") + return op, { + "selected_backend": backend.name, + "fallback": True, + "fallback_reason": "modulate_index_requires_select01", + "rejected_backends": [], + "reduction_order": "pairwise_lower_upper_H; framework_select01_backward", + "accumulator_dtype": "fp32", + "split_k": False, + "stream_k": False, + "tf32": False, + "kernel_fingerprint": "adaln_modulation_select01_native_v1", + } candidates = self._priority_map[platform]["adaln_modulation"] rejected: list[str] = [] for backend in candidates: diff --git a/tests/test_adaln_modulation.py b/tests/test_adaln_modulation.py index f3caa0ef3..36607336b 100644 --- a/tests/test_adaln_modulation.py +++ b/tests/test_adaln_modulation.py @@ -40,6 +40,39 @@ def test_adaln_registry_trace_reports_selected_backend(): 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, 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(): try: torch.cuda.init() From f38ba6f27b2e495291250cc401d979eb0c8fc724 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Fri, 25 Sep 2026 01:21:04 +0800 Subject: [PATCH 05/37] Enhance AdaLN Modulation with Thread Configuration and Synchronization Fixes - Updated CUDA kernel functions to accept a configurable number of threads (128, 256, or 512) for improved flexibility. - Added synchronization after reading the reduction result in the fixed_sum function to prevent shared memory hazards. - Modified the Python bindings and C++ interface to accommodate the new threads parameter in forward and backward functions. - Enhanced documentation to include the new thread configuration options. - Implemented additional tests to validate the behavior of the AdaLN modulation with different thread configurations and ensure numerical stability. - Introduced a race condition check script to validate shared memory access in CUDA kernels. Signed-off-by: hongleng-github --- csrc/bindings/ops.cpp | 11 +- csrc/cuda/adaln_modulation.cu | 27 +- docs/.nav.yml | 1 + docs/operators/adaln-modulation.md | 96 ++++++- rl_engine/_C.pyi | 5 +- .../kernels/ops/cuda/norm/adaln_modulation.py | 18 +- .../ops/triton/norm/adaln_modulation.py | 29 ++- .../03_build_extension.sh | 3 +- .../04_test_adaln_cuda.sh | 8 +- .../06_racecheck_adaln_cuda.sh | 23 ++ tests/test_adaln_modulation.py | 243 +++++++++++++++--- 11 files changed, 382 insertions(+), 82 deletions(-) create mode 100644 scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index 6e2019bfe..1dbf203ac 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -258,10 +258,10 @@ std::vector swiglu_packed_backward_cuda( #if defined(RLK_ADALN_CUDA_ENABLED) std::vector adaln_modulation_forward_cuda( - torch::Tensor x, torch::Tensor modulation, double eps); + 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); + torch::Tensor modulation, double eps, int threads); #endif // RMSNorm Declarations & Wrappers @@ -747,9 +747,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "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"); + "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"); + "AdaLN modulation backward CUDA", + py::arg("dy"), py::arg("dg"), py::arg("x"), py::arg("modulation"), + py::arg("eps"), py::arg("threads") = 256); #endif // Deterministic standard-softmax attention (issue #147) diff --git a/csrc/cuda/adaln_modulation.cu b/csrc/cuda/adaln_modulation.cu index c22703e36..0ff7eb7bb 100644 --- a/csrc/cuda/adaln_modulation.cu +++ b/csrc/cuda/adaln_modulation.cu @@ -15,8 +15,6 @@ namespace { -constexpr int THREADS = 256; - __device__ float fixed_sum(float* shared, int width) { __syncthreads(); for (int stride = width / 2; stride > 0; stride /= 2) { @@ -25,7 +23,10 @@ __device__ float fixed_sum(float* shared, int width) { } __syncthreads(); } - return shared[0]; + const float result = shared[0]; + // Every warp must read the result before callers reuse the shared buffer. + __syncthreads(); + return result; } template @@ -127,7 +128,9 @@ __global__ void backward_reduce_kernel( dm[batch * 3 * hidden + 2 * hidden + col] = dg[batch * hidden + col]; } -void check(const torch::Tensor& x, const torch::Tensor& modulation, double eps) { +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, @@ -145,8 +148,8 @@ void check(const torch::Tensor& x, const torch::Tensor& modulation, double eps) } // namespace std::vector adaln_modulation_forward_cuda( - torch::Tensor x, torch::Tensor modulation, double eps) { - check(x, modulation, eps); + 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; @@ -155,7 +158,7 @@ std::vector adaln_modulation_forward_cuda( 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<<>>( + forward_kernel<<>>( x.data_ptr(), modulation.data_ptr(), y.data_ptr(), gate.data_ptr(), seq, hidden, width, static_cast(eps)); }); @@ -165,8 +168,8 @@ std::vector adaln_modulation_forward_cuda( std::vector adaln_modulation_backward_cuda( torch::Tensor dy, torch::Tensor dg, torch::Tensor x, - torch::Tensor modulation, double eps) { - check(x, modulation, eps); + 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"); @@ -184,12 +187,12 @@ std::vector adaln_modulation_backward_cuda( 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<<>>( + 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<<>>( + backward_reduce_kernel<<>>( partial_shift.data_ptr(), partial_scale.data_ptr(), dg.data_ptr(), dm.data_ptr(), seq, hidden); }); 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/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 6570ad743..694a0997c 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -11,7 +11,8 @@ The `SiLU + Linear` projection and `adaln_gate_residual` are separate operators. The FP32 hidden reductions use a lower-plus-upper pairwise tree padded to the next power of two. Each sample's `dshift` and `dscale` contributions are folded -in ascending logical token order. No atomic partial sums, Split-K, Stream-K, +in ascending logical token order. The CUDA shared reduction reads its result +into each thread before a block barrier permits buffer reuse. No atomic partial sums, Split-K, Stream-K, TF32, or fast math are used by these kernels. BF16 outputs are cast once after FP32 arithmetic. The CPU reference computes square root in FP64 and rounds it to FP32 because CPU `torch.sqrt(float32)` can differ by one ULP from CUDA's @@ -52,12 +53,16 @@ median forward and forward+backward latency plus peak allocated memory; CUDA extension if its symbols are absent. The registry trace records selected backend, fallback, reduction order, accumulator dtype, disabled algorithm flags, and a semantic kernel version. The semantic version is not a hash of the -compiled binary. +compiled binary. Direct backend constructors accept validated launch settings: +CUDA `threads={128,256,512}` (default 256); Triton `num_warps={4,8}` +(default 4), `reduction_tile={64,128,256}` (default 128). Registry dispatch +uses those defaults. These settings change scheduling, not the reduction order. Local CUDA validation environment (2026-09-23): RTX 3050 Ti (SM86, 4 GiB), driver 610.60, CUDA toolkit/runtime 12.6, PyTorch 2.12.1+cu126, Triton 3.7.1, Python 3.12.12. The extension was built with fast math disabled for SM86. -The focused suite plus all three real shapes passed 18 tests with no skips. +The pre-review focused suite plus all three real shapes passed 18 tests with no skips. +These historical results are not acceptance evidence for the synchronization fix. FP32 small-shape and BF16 real-shape CUDA forward/backward outputs were byte equal to the CPU reference. CUDA and Triton were also byte equal directly at all three real shapes, and both passed FP32/BF16 batch-size and token-padding @@ -80,10 +85,81 @@ at each shape. The CUDA backward path was generally slower in this short SM86 run; rerun with the default warmup/repeat counts on target hardware before making a performance claim. -**Remaining acceptance boundaries:** the maintainer has not answered whether -the indexed `modulate_index` path must be optimized beyond its explicit -PyTorch fallback. The trace fingerprint is a semantic kernel version rather -than a compiled-binary hash. The exactness -evidence above covers FP32 at the focused shape and BF16 at the three real image -shapes; a broader CUDA shape matrix remains optional evidence, not a result -claimed here. +## Review regression checks (2026-09-24) + +The pre-fix current-HEAD suite passed 20 cases (including extension smoke), +but compute-sanitizer found one forward and three backward shared-memory hazards. +The missing synchronization was after reading the reduction result, before +reusing its buffer. The minimal fixed-seed H=64 regression returned exit code 7 +even though its numerical assertion passed. + +The non-native review selection passed 30 cases with 26 deselected using +`-k 'not cuda and not real_image and not gpu_padding'`. This does not validate +the rebuilt native extension. + +Post-fix validation of the rebuilt extension passed **57 tests in 43.73s** with +no skips, including native extension smoke and the complete focused matrix. +The native registry selected `CUDA_ADALN_MODULATION` with `fallback=False`. +The expanded sanitizer regression passed **8 tests** with **0 hazards, +0 errors, and 0 warnings** (exit 0). The original seed-386 FP32 +`[1,2,3072]` forward/backward reproducer also reported zero hazards after +the fix. Environment remained PyTorch 2.12.1+cu126, Triton 3.7.1, Python +3.12.12, RTX 3050 Ti SM86. + +Observed commands (the sanitizer ran separately before the full pytest matrix): + +```bash +bash scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh +PYTHONDONTWRITEBYTECODE=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ +RLK_ADALN_REAL_SHAPES=1 RL_KERNEL_REQUIRE_EXT=1 \ +.venv/bin/python -m pytest tests/test_extension_smoke.py \ + tests/test_adaln_modulation.py -q -rs -p no:cacheprovider +``` + +Local logs are `/tmp/rlk-adaln-race-fix-build.log`, +`/tmp/rlk-adaln-race-fix-racecheck.log`, and +`/tmp/rlk-adaln-race-fix-tests.log`. They describe the uncommitted review +fixes on top of `d14e4ed`; that commit alone predates the fix. + +Run the combined fail-closed checks using the cu126 repository environment: + +```bash +bash scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh +# Or only the sanitizer regression: +bash scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh +``` + +Script 04 runs the full focused suite with real shapes and required native CUDA, +then script 06. Missing extension symbols or sanitizer failures must not count +as successful native acceptance. Do not use bare `uv run` or `uv sync` to +run these checks: invoke `.venv/bin/python` directly to preserve cu126. + +The statuses below apply to the measured inputs and SM86 configurations, not +all shapes or hardware. + +| Requirement | Implementation | Test | Observed evidence | Remaining gap | Status | +| --- | --- | --- | --- | --- | --- | +| Shared forward/backward math at H=7/3072 | Fixed FP32 normalization and VJP | Independent LayerNorm/autograd, both dtypes, random dy/dgate | All 12 backend/dtype/H cases passed | Independent full-image autograd not covered | PROVEN | +| Race-free CUDA shared reduction | Read-result/barrier/reuse | Native harness and geometry under racecheck | 8 tests, zero hazards; original reproducer also green | Other architectures not exercised | PROVEN | +| Launch geometry / tiling | CUDA threads 128/256/512; Triton 4/8 warps and 64/128/256 tiles | Public constructors, CPU-reference byte checks | Both dtypes passed all 3 CUDA and 6 Triton configurations | Coverage is H=3072, B=2, S=5 | PROVEN | +| Strict byte comparator | Contiguous uint8 views | Signed-zero regression and backend comparisons | Comparator rejects +0 versus -0; full matrix passed | Exceptional floating-point payloads not covered | PROVEN | +| Three real image sizes | S=4096/6889/6032, H=3072 | Both dtypes, random dy/dgate, CPU/native/Triton byte checks | All 6 real-shape cases passed | Full-size oracle is the fixed CPU implementation, not independent LayerNorm/autograd | PARTIAL | +| Indexed select01 | Explicit PyTorch fallback with [2B,3H] modulation | B=2, both index shapes/devices/dtypes, independent row-index and autograd gold; invalid routes | 12 indexed cases passed | No fixed-order indexed backward claim | PROVEN | +| FP32 accumulation / final BF16 cast | FP32 reductions, CUDA RN operations, Triton fusion disabled | Cast, byte-reference and independent accuracy tests | All focused checks passed | Exhaustive rounding boundaries and fast-math-on build exclusion not exercised | PARTIAL | +| Batch / padding invariance | Row-local hidden reduction, ascending token fold | Batch position, B=1/3, S=3/5, zero-gradient padding; raw bytes | FP32/BF16 CUDA and Triton cases passed | Tested padding matrix uses H=32 | PROVEN | +| Registry / trace | Explicit backend selection, H limit, indexed fallback reason | Registry tests and native preflight | Native selected without fallback; indexed fallback and H limit passed | Semantic fingerprint is not a binary hash; unavailable-backend combinations are not exhaustive | PARTIAL | +| Native SM86 build and binding | Conditional sources/symbols, defaulted threads argument | Symbol verification, native/reference tests and geometry launches | User build completed; new four-argument binding and full suite passed | Other build configurations untested | PROVEN | +| Indexed optimized kernels | Agreed explicit PyTorch fallback | Fallback and trace checks | Within agreed scope | No optimized indexed kernel required | OUT OF SCOPE | +| Full #386 WS1 exit criteria | All 19 operators and assembled model | MMDiT/LoRA, 60-layer curve, reproducible rollout logp | Not exercised by this operator suite | Separate system-level work | UNPROVEN | + + +The independent autograd checks are accuracy checks, not a replacement tolerance +profile for the fixed-reference byte contract. Shared checks use FP32 +`atol=rtol=2e-5`, BF16 `atol=0.02, rtol=0.016`; indexed checks use +`atol=rtol=2e-6`. The full-size byte checks still use the fixed CPU reference +and do not independently validate LayerNorm/autograd at every image size. + +The semantic fingerprint does not identify a compiled binary. This suite does +not claim coverage across other GPU architectures, every floating-point input, +or fast-math-enabled builds. Benchmark numbers above predate the synchronization +fix; they are historical smoke measurements, not performance claims for the fix. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 8453a87da..d80c16b5c 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -257,10 +257,11 @@ def swiglu_backward( up: torch.Tensor, ) -> list[torch.Tensor]: ... def adaln_modulation_forward( - x: torch.Tensor, modulation: torch.Tensor, eps: float + 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 + dy: torch.Tensor, dg: torch.Tensor, x: torch.Tensor, modulation: torch.Tensor, + eps: float, threads: int = 256 ) -> list[torch.Tensor]: ... def rmsnorm_forward( x: torch.Tensor, diff --git a/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py b/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py index ffacc8f51..0fd5cbe68 100644 --- a/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py +++ b/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py @@ -12,23 +12,27 @@ class _AdaLNCuda(torch.autograd.Function): @staticmethod - def forward(ctx, x, modulation, eps): - y, gate = _C.adaln_modulation_forward(x, modulation, eps) + 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 + grad_y.contiguous(), grad_gate.contiguous(), x, modulation, ctx.eps, ctx.threads ) - return dx, dm, None + return dx, dm, None, None class CudaAdaLNModulationOp: - def __init__(self): + 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") @@ -42,4 +46,6 @@ def forward(self, x, modulation, *, eps=1e-6): _validate(x, modulation, eps) 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)) + return _AdaLNCuda.apply( + x.contiguous(), modulation.contiguous(), float(eps), self.threads + ) diff --git a/rl_engine/kernels/ops/triton/norm/adaln_modulation.py b/rl_engine/kernels/ops/triton/norm/adaln_modulation.py index 22129f0f8..9843a238e 100644 --- a/rl_engine/kernels/ops/triton/norm/adaln_modulation.py +++ b/rl_engine/kernels/ops/triton/norm/adaln_modulation.py @@ -90,17 +90,19 @@ def _bwd_reduce(PART_SHIFT, PART_SCALE, DG, DM, S: tl.constexpr, H: tl.constexpr class _AdaLNTriton(torch.autograd.Function): @staticmethod - def forward(ctx, x, modulation, eps): + 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) _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_fp_fusion=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 @@ -115,15 +117,25 @@ def backward(ctx, grad_y, grad_gate): 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, + num_warps=ctx.num_warps, ) - _bwd_reduce[(batch, triton.cdiv(hidden, 128))]( - partial_shift, partial_scale, grad_gate.contiguous(), dm, seq, hidden, 128, - enable_fp_fusion=False, + _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 + 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 + def __call__(self, x, modulation, *, eps=1e-6): return self.forward(x, modulation, eps=eps) @@ -131,4 +143,7 @@ def forward(self, x, modulation, *, eps=1e-6): _validate(x, modulation, eps) if not x.is_cuda: raise RuntimeError("Triton AdaLN requires CUDA tensors") - return _AdaLNTriton.apply(x.contiguous(), modulation.contiguous(), float(eps)) + return _AdaLNTriton.apply( + x.contiguous(), modulation.contiguous(), float(eps), + self.num_warps, self.reduction_tile, + ) diff --git a/scripts/adaln_cuda_acceptance/03_build_extension.sh b/scripts/adaln_cuda_acceptance/03_build_extension.sh index e60258c17..6ad635421 100755 --- a/scripts/adaln_cuda_acceptance/03_build_extension.sh +++ b/scripts/adaln_cuda_acceptance/03_build_extension.sh @@ -4,7 +4,7 @@ set -euo pipefail root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) cd "$root" export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.6} -export PATH="$CUDA_HOME/bin:$PATH" +export PATH="$root/.venv/bin:$CUDA_HOME/bin:$PATH" export LIBRARY_PATH="/usr/lib/wsl/lib:${LIBRARY_PATH:-}" export LD_LIBRARY_PATH="/usr/lib/wsl/lib:$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}" export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST:-8.6} @@ -14,6 +14,7 @@ export RL_KERNEL_REQUIRE_EXT=1 export KERNEL_ALIGN_USE_FAST_MATH=0 uv pip install --python .venv/bin/python ninja +ninja --version .venv/bin/python setup.py build_ext --inplace .venv/bin/python - <<'PY' import rl_engine._C as ext diff --git a/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh b/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh index 4fb604081..9ae064d56 100755 --- a/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh +++ b/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh @@ -10,10 +10,16 @@ export RLK_ADALN_REAL_SHAPES=1 export PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 .venv/bin/python - <<'PY' +import sys +import torch +import triton +assert torch.version.cuda == "12.6", torch.__version__ +print(sys.version.split()[0], torch.__version__, torch.version.cuda, triton.__version__) from rl_engine.kernels.registry import OpBackend, kernel_registry _, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=3072) assert trace["selected_backend"] == OpBackend.CUDA_ADALN_MODULATION.name, trace assert trace["fallback"] is False, trace print(trace) PY -.venv/bin/python -m pytest tests/test_extension_smoke.py tests/test_adaln_modulation.py -q -rs +.venv/bin/python -m pytest tests/test_extension_smoke.py tests/test_adaln_modulation.py -q -rs -p no:cacheprovider +bash scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh diff --git a/scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh b/scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh new file mode 100644 index 000000000..e7827fdf4 --- /dev/null +++ b/scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +set -euo pipefail + +root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) +cd "$root" +export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.6} +export LD_LIBRARY_PATH="/usr/lib/wsl/lib:$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}" +export RL_KERNEL_REQUIRE_EXT=1 +export PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 +export PYTHONDONTWRITEBYTECODE=1 + +.venv/bin/python - <<'PY' +import torch +assert torch.version.cuda == "12.6", torch.__version__ +print(torch.__version__, torch.cuda.get_device_name(), flush=True) +PY + +# Numerical equality alone does not detect shared-memory read/write hazards. +exec "$CUDA_HOME/bin/compute-sanitizer" --tool racecheck \ + --racecheck-report analysis --error-exitcode 7 \ + .venv/bin/python -m pytest -q -rs -p no:cacheprovider \ + tests/test_adaln_modulation.py::test_cuda_adaln_bit_equality_harness \ + tests/test_adaln_modulation.py::test_adaln_launch_geometry_preserves_bytes -k cuda "$@" diff --git a/tests/test_adaln_modulation.py b/tests/test_adaln_modulation.py index 36607336b..22f39662f 100644 --- a/tests/test_adaln_modulation.py +++ b/tests/test_adaln_modulation.py @@ -7,6 +7,13 @@ from rl_engine.kernels.registry import OpBackend, kernel_registry +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 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. @@ -126,7 +133,7 @@ def test_adaln_modulation_batch_position_does_not_change_bytes(): batch[position] = row[0] batch_mod[position] = mod[0] actual = op(batch, batch_mod)[0][position] - assert torch.equal(actual, expected[0]) + _assert_bytes_equal(actual, expected[0]) def test_adaln_modulation_bf16_cast_and_backward(): @@ -138,7 +145,7 @@ def test_adaln_modulation_bf16_cast_and_backward(): (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 torch.equal(modulation.grad[0, 18:], torch.ones(9, dtype=torch.bfloat16)) + _assert_bytes_equal(modulation.grad[0, 18:], torch.ones(9, dtype=torch.bfloat16)) def test_adaln_modulation_accepts_qwen_six_chunk_view(): @@ -149,7 +156,7 @@ def test_adaln_modulation_accepts_qwen_six_chunk_view(): y, gate = NativeAdaLNModulationOp()(x, modulation) (y.sum() + gate.sum()).backward() assert full.grad is not None - assert torch.equal(full.grad[:, 24:], torch.zeros_like(full.grad[:, 24:])) + _assert_bytes_equal(full.grad[:, 24:], torch.zeros_like(full.grad[:, 24:])) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -171,9 +178,9 @@ def test_triton_adaln_modulation_forward_backward_and_batch_position(dtype): 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 torch.equal(y, y_ref) - assert torch.equal(x.grad, x_ref.grad) - assert torch.equal(modulation.grad, mod_ref.grad) + _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 = ( @@ -182,15 +189,15 @@ def test_triton_adaln_modulation_forward_backward_and_batch_position(dtype): .requires_grad_() ) y_new, gate_new = TritonAdaLNModulationOp()(x_new, mod_new) - assert torch.equal(y_new[1], y[0]) - assert torch.equal(gate_new[1], gate[0]) + _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 torch.equal(x_new.grad[1], x.grad[0]) - assert torch.equal(mod_new.grad[1], modulation.grad[0]) + _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(): @@ -214,9 +221,9 @@ def test_triton_irregular_hidden_and_bf16_tie_rounding(): 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 torch.equal(y.cpu(), y_cpu) - assert torch.equal(x.grad.cpu(), x_cpu.grad) - assert torch.equal(modulation.grad.cpu(), modulation_cpu.grad) + _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]) @@ -258,25 +265,29 @@ def test_adaln_gpu_padding_and_batch_size_do_not_change_valid_bytes(dtype): padded_modulation.requires_grad_() y, gate = op(padded_x, padded_modulation) torch.autograd.backward((y, gate), (padded_dy, padded_dgate)) - assert torch.equal(y[-1, :3], baseline_y[0]) - assert torch.equal(gate[-1], baseline_gate[0]) - assert torch.equal(padded_x.grad[-1, :3], baseline_x.grad[0]) - assert torch.equal(padded_modulation.grad[-1], baseline_modulation.grad[0]) + _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]) -def test_cuda_adaln_bit_equality_harness(): +@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.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp try: torch.cuda.init() cuda_op = CudaAdaLNModulationOp() except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise pytest.skip("compiled CUDA AdaLN extension required") torch.manual_seed(390) - x = torch.randn(1, 3, 32, requires_grad=True) - modulation = torch.randn(1, 96, requires_grad=True) + 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, 32) + 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() @@ -284,18 +295,19 @@ def test_cuda_adaln_bit_equality_harness(): 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 torch.equal(cuda_y.cpu(), cpu_y) - assert torch.equal(cuda_gate.cpu(), cpu_gate) - assert torch.equal(x_gpu.grad.cpu(), x.grad) - assert torch.equal(mod_gpu.grad.cpu(), modulation.grad) + _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): +def test_adaln_real_image_shapes_forward_backward(seq, dtype): try: torch.cuda.init() except RuntimeError: @@ -307,29 +319,182 @@ def test_adaln_real_image_shapes_forward_backward(seq): OpBackend.CUDA_ADALN_MODULATION.name, OpBackend.TRITON_ADALN_MODULATION.name, } - x = torch.randn(1, seq, 3072, device="cuda", dtype=torch.bfloat16, requires_grad=True) - modulation = torch.randn(1, 9216, device="cuda", dtype=torch.bfloat16, requires_grad=True) + 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) - (y.float().sum() + gate.float().sum()).backward() + 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) - (cpu_y.float().sum() + cpu_gate.float().sum()).backward() - assert torch.equal(y.cpu(), cpu_y) - assert torch.equal(gate.cpu(), cpu_gate) - assert torch.equal(x.grad.cpu(), cpu_x.grad) - assert torch.equal(modulation.grad.cpu(), cpu_modulation.grad) + 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.kernels.ops.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) - (triton_y.float().sum() + triton_gate.float().sum()).backward() - assert torch.equal(triton_y, y) - assert torch.equal(triton_gate, gate) - assert torch.equal(triton_x.grad, x.grad) - assert torch.equal(triton_modulation.grad, modulation.grad) + 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) + + +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.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from rl_engine.kernels.ops.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().clone().requires_grad_() + ref_modulation = modulation.detach().clone().requires_grad_() + # Independent row indexing, not the implementation's split-and-where path. + selected = ref_modulation.float()[torch.arange(2, device=device)[:, None] + index * 2] + shift, scale, ref_gate = selected.chunk(3, dim=-1) + ref_y = ( + torch.nn.functional.layer_norm(ref_x.float(), (7,), eps=1e-6) * (1 + scale) + shift + ).to(dtype) + ref_gate = ref_gate.to(dtype) + torch.autograd.backward((ref_y, ref_gate), (dy, dg)) + for actual, expected in ( + (y, ref_y), (x.grad, ref_x.grad), (modulation.grad, ref_modulation.grad) + ): + torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-6) + _assert_bytes_equal(gate, ref_gate) + assert x.grad.abs().max() > 0 + 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") + from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + ops = {"pytorch": NativeAdaLNModulationOp, "triton": TritonAdaLNModulationOp, + "cuda": CudaAdaLNModulationOp} + if backend == "cuda": + try: + op = ops[backend]() + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + else: + op = ops[backend]() + 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)) + + ref_x = x.detach().cpu().requires_grad_() + ref_modulation = modulation.detach().cpu().requires_grad_() + shift, scale, ref_gate = ref_modulation.float().chunk(3, dim=-1) + ref_y = ( + torch.nn.functional.layer_norm(ref_x.float(), (hidden,), eps=1e-6) + * (1 + scale[:, None, :]) + shift[:, None, :] + ).to(dtype) + ref_gate = ref_gate[:, None, :].to(dtype) + torch.autograd.backward((ref_y, ref_gate), (dy.cpu(), dg.cpu())) + for actual, expected in ( + (y, ref_y), (x.grad, ref_x.grad), (modulation.grad, ref_modulation.grad) + ): + torch.testing.assert_close( + actual.cpu(), expected, + atol=2e-5 if dtype == torch.float32 else 0.02, + rtol=2e-5 if dtype == torch.float32 else 0.016, + ) + _assert_bytes_equal(gate, ref_gate) From a566b16fbe5b3e9587447c0ed243420fd6ee5750 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Fri, 9 Oct 2026 23:29:03 +0800 Subject: [PATCH 06/37] fix(adaln): close shared precision and validation gaps Signed-off-by: hongleng-github --- .github/workflows/ci.yml | 10 +- .github/workflows/gpu-ci.yml | 2 +- .github/workflows/ws1-chain-gpu.yml | 4 +- .github/workflows/ws1-chain-npu.yml | 4 +- .github/workflows/ws1-gtest-gpu.yml | 16 +- ci/scripts/run_ws1_gtest.sh | 11 + csrc/cuda/adaln_modulation.cu | 8 + docs/contributing/gtest-usage.md | 15 +- docs/operators/adaln-modulation.md | 90 ++- docs/reviews/adaln-shared-b31fac9-followup.md | 172 +++++ docs/reviews/adaln-shared-b31fac9.md | 351 +++++++++ .../kernels/ops/cuda/norm/adaln_modulation.py | 11 +- .../ops/pytorch/norm/adaln_modulation.py | 21 +- .../ops/triton/norm/adaln_modulation.py | 6 +- rl_engine/validation/operators/op_checks.py | 21 +- .../validation/operators/operator_inputs.py | 14 + .../validation/operators/operator_specs.py | 207 +++--- .../04_test_adaln_cuda.sh | 3 +- .../06_racecheck_adaln_cuda.sh | 3 +- tests/test_adaln_gtest.py | 665 ++++++++++++++++++ tests/test_adaln_modulation.py | 119 ++-- tests/test_adaln_shared_bounds.py | 74 ++ 22 files changed, 1638 insertions(+), 189 deletions(-) create mode 100644 docs/reviews/adaln-shared-b31fac9-followup.md create mode 100644 docs/reviews/adaln-shared-b31fac9.md create mode 100644 tests/test_adaln_gtest.py create mode 100644 tests/test_adaln_shared_bounds.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4634569cd..a678e3015 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,9 +5,9 @@ name: CI-Pipeline on: push: - branches: [ main, test ] + branches: [ main, test, test-qwenimage ] pull_request: - branches: [ main, test ] + branches: [ main, test, test-qwenimage ] permissions: contents: read @@ -71,6 +71,12 @@ jobs: - name: Check refactor packaging and compatibility run: python ci/run.py layout -q + - name: Run AdaLN CPU contract and reference tests + env: + CUDA_VISIBLE_DEVICES: "" + PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1" + run: python -m pytest tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py -q -rs + - name: Run Mocked Hardware Discovery Tests run: | python -m pytest tests/runtime/test_dispatch.py -v diff --git a/.github/workflows/gpu-ci.yml b/.github/workflows/gpu-ci.yml index dbf822d4d..c73b791a3 100644 --- a/.github/workflows/gpu-ci.yml +++ b/.github/workflows/gpu-ci.yml @@ -2,7 +2,7 @@ name: GPU CI on: pull_request_target: - branches: [ main, test ] + branches: [ main, test, test-qwenimage ] paths: - 'csrc/**' - 'rl_engine/**' diff --git a/.github/workflows/ws1-chain-gpu.yml b/.github/workflows/ws1-chain-gpu.yml index 43d63fc13..9849ffe4a 100644 --- a/.github/workflows/ws1-chain-gpu.yml +++ b/.github/workflows/ws1-chain-gpu.yml @@ -8,9 +8,9 @@ name: WS1-chain-GPU on: pull_request: - branches: [ main, test ] + branches: [ main, test, test-qwenimage ] push: - branches: [ main, test ] + branches: [ main, test, test-qwenimage ] workflow_dispatch: inputs: source_repository: diff --git a/.github/workflows/ws1-chain-npu.yml b/.github/workflows/ws1-chain-npu.yml index 1d14094ed..45fb70c7f 100644 --- a/.github/workflows/ws1-chain-npu.yml +++ b/.github/workflows/ws1-chain-npu.yml @@ -15,9 +15,9 @@ name: WS1-chain-NPU on: pull_request: - branches: [ main, test ] + branches: [ main, test, test-qwenimage ] push: - branches: [ main, test ] + branches: [ main, test, test-qwenimage ] workflow_dispatch: inputs: source_repository: diff --git a/.github/workflows/ws1-gtest-gpu.yml b/.github/workflows/ws1-gtest-gpu.yml index a0d8ac0c4..56bdaa9a4 100644 --- a/.github/workflows/ws1-gtest-gpu.yml +++ b/.github/workflows/ws1-gtest-gpu.yml @@ -10,7 +10,7 @@ name: WS1-gtest-GPU on: pull_request: - branches: [ main ] + branches: [ main, test-qwenimage ] paths: - "rl_engine/validation/**" - "rl_engine/backends/**" @@ -22,6 +22,12 @@ on: - "rl_engine/reference/**" - "build_tools/**" - "csrc/**" + - "setup.py" + - "tests/test_adaln*.py" + - "tools/validation/operators/check_operator.py" + - "rl_engine/kernels/gtest/**" + - "rl_engine/kernels/ops/**" + - "rl_engine/testing/**" - "tools/validation/operators/sweep_ws1_four_judgments.py" - "tools/validation/operators/check_forward_invariance.py" - "tools/validation/operators/check_gradient_invariance.py" @@ -36,7 +42,7 @@ on: - "ci/scripts/run_gpu_ci.sh" - ".github/workflows/ws1-gtest-gpu.yml" push: - branches: [ main ] + branches: [ main, test-qwenimage ] paths: - "rl_engine/validation/**" - "rl_engine/backends/**" @@ -48,6 +54,12 @@ on: - "rl_engine/reference/**" - "build_tools/**" - "csrc/**" + - "setup.py" + - "tests/test_adaln*.py" + - "tools/validation/operators/check_operator.py" + - "rl_engine/kernels/gtest/**" + - "rl_engine/kernels/ops/**" + - "rl_engine/testing/**" - "tools/validation/operators/sweep_ws1_four_judgments.py" - "tools/validation/operators/check_forward_invariance.py" - "tools/validation/operators/check_gradient_invariance.py" diff --git a/ci/scripts/run_ws1_gtest.sh b/ci/scripts/run_ws1_gtest.sh index cc47dd678..2c866fb8a 100755 --- a/ci/scripts/run_ws1_gtest.sh +++ b/ci/scripts/run_ws1_gtest.sh @@ -23,8 +23,19 @@ 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/test_adaln_gtest.py \ + tests/test_adaln_shared_bounds.py \ + tests/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/cuda/adaln_modulation.cu b/csrc/cuda/adaln_modulation.cu index 0ff7eb7bb..826a61d59 100644 --- a/csrc/cuda/adaln_modulation.cu +++ b/csrc/cuda/adaln_modulation.cu @@ -8,6 +8,7 @@ #include #include #include +#include #if defined(__CUDA_FAST_MATH__) #error "adaln_modulation requires CUDA compilation without --use_fast_math" @@ -135,14 +136,21 @@ void check(const torch::Tensor& x, const torch::Tensor& modulation, double eps, 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"); } } // namespace diff --git a/docs/contributing/gtest-usage.md b/docs/contributing/gtest-usage.md index 2e0d9e235..95331d316 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 scripts/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 index 694a0997c..de1c96bf4 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -19,10 +19,22 @@ to FP32 because CPU `torch.sqrt(float32)` can differ by one ULP from CUDA's correctly rounded square root. Hidden and token accumulators remain FP32. Inputs with S=0 or H=0 and mixed dtype/device are rejected. Strided model chunks are copied to contiguous storage inside -GPU wrappers. The CUDA kernel supports H<=4096; shape-aware dispatch selects -Triton for larger H. CUDA symbols are omitted from builds with +GPU wrappers. The CUDA kernel supports H<=4096, B<=65535, and at most +2^31-1 elements in each input tensor. Python checks these bounds before copying +strided inputs; C++ checks them again before launch. Oversized B/address ranges +raise an explicit error; the registry only resolves using H and does not silently +fallback for these other limits. Shape-aware dispatch selects Triton for larger H. +CUDA symbols are omitted from builds with `KERNEL_ALIGN_USE_FAST_MATH=1`, making fallback explicit. +Shared eps must be finite and nonnegative and round to finite FP32. A positive +eps that rounds to zero is rejected; representable subnormals, including 2^-149, +remain supported. eps=0 remains supported; zero variance then produces IEEE NaN. +Nonfinite LayerNorm arithmetic follows the FP32 formula and may propagate NaNs; +its NaN payloads are outside the strict byte-equality qualification. Gate and +dgate use a bit-preserving passthrough, including nonfinite payloads. This shared +eps policy does not change the separate indexed/select01 fallback contract. + For Qwen-Image editing, pass `modulate_index` to both resolution and execution: ```python @@ -44,6 +56,48 @@ These are derived model shapes, not tensor shapes stated by the issue. ## Validation +The [2026-10-09 shared-workspace review](../reviews/adaln-shared-b31fac9.md) +compares `b31fac9` with the actual worktree, including untracked tests. It records +current CPU/SM86 Triton evidence separately from pending native CUDA/H100, +full-image, geometry, sanitizer and benchmark acceptance. New public-interface +canaries check BF16 cast boundaries, each reduction, signed zero, subnormal eps, +gate/dgate payloads, strided upstreams and H=3072 batch positions with padding +across all declared launch configurations. AdaLN gtest rejects a candidate's wrong +output dtype and explicitly promotes gold leaves to FP32 +before autograd; an FP32 cast inside the gold function alone would still return +BF16 gradients to BF16 leaves. The review includes the red/green evidence and +H100 commands; historical results below are not current-worktree evidence. + +AdaLN is registered in `gtest/operator_specs.py` and `operator_inputs.py` as +`op_class="reduction"`, with PyTorch, CUDA and Triton candidates. The independent +gold uses FP32 CPU `torch.nn.functional.layer_norm` and autograd, without the +production fixed-reduction or custom backward helpers. It checks both outputs +(y and gate) and gradients of x and modulation. Accuracy thresholds come from +`resolve_tolerance`: `forward_accuracy` for outputs and `gradient_accuracy` +for input gradients; no private AdaLN thresholds or SM90 override are added. +Existing fixed-reference byte comparisons remain separate, unchanged gates. + +```bash +.venv/bin/python scripts/check_operator.py --op adaln_modulation \ + --candidate pytorch --device cpu --dtype bf16 --batch 2 --seq 3 \ + --normalized-dim 3072 --check-grad --json +.venv/bin/python -m pytest tests/test_adaln_gtest.py -q -rs +``` + +Use CUDA/Triton candidates on the target GPU for native accuracy evidence. +The three full-image tests additionally check independent CPU FP32 forward and +backward accuracy for both GPU backends/dtypes. These additions are **pending +GPU execution**; old SM86 results do not validate the new assertions. +CPU tests also cover gate-only gradients, distinct image/text modulation +chunks, corrupted-output/VJP rejection and unavailable-backend routing. +The CLI accuracy report alone is not a provenance-checked WS1 system gate. + +CI branch filters now include `test-qwenimage`; CPU CI runs the new contract +tests and the existing WS1 GPU script includes AdaLN accuracy checks. Fork +isolation and existing GPU authorization conditions remain. A missing/skipped +workflow is not a passing check; no CI run is claimed for these local changes. + + Run `python -m pytest tests/test_adaln_modulation.py -q` for focused tests. Set `RLK_ADALN_REAL_SHAPES=1` to include the three full image shapes. Run `python benchmarks/benchmark_adaln_modulation.py --real` for the three @@ -58,6 +112,15 @@ CUDA `threads={128,256,512}` (default 256); Triton `num_warps={4,8}` (default 4), `reduction_tile={64,128,256}` (default 128). Registry dispatch uses those defaults. These settings change scheduling, not the reduction order. +### Historical evidence limitation (verified 2026-10-02) + +The following SM86 results and timings are historical observations. The two +original `/tmp` handoffs and raw build/test/racecheck logs are no longer +available. Surviving excerpts and recorded hashes cannot reconstruct or +substitute for complete logs, binary provenance or newly executed evidence. +Re-run acceptance and publish fresh complete artifacts before using it as a +reviewable native acceptance result. H100 validation is **pending**. + Local CUDA validation environment (2026-09-23): RTX 3050 Ti (SM86, 4 GiB), driver 610.60, CUDA toolkit/runtime 12.6, PyTorch 2.12.1+cu126, Triton 3.7.1, Python 3.12.12. The extension was built with fast math disabled for SM86. @@ -116,7 +179,7 @@ RLK_ADALN_REAL_SHAPES=1 RL_KERNEL_REQUIRE_EXT=1 \ tests/test_adaln_modulation.py -q -rs -p no:cacheprovider ``` -Local logs are `/tmp/rlk-adaln-race-fix-build.log`, +Original, now missing logs were `/tmp/rlk-adaln-race-fix-build.log`, `/tmp/rlk-adaln-race-fix-racecheck.log`, and `/tmp/rlk-adaln-race-fix-tests.log`. They describe the uncommitted review fixes on top of `d14e4ed`; that commit alone predates the fix. @@ -134,30 +197,29 @@ then script 06. Missing extension symbols or sanitizer failures must not count as successful native acceptance. Do not use bare `uv run` or `uv sync` to run these checks: invoke `.venv/bin/python` directly to preserve cu126. -The statuses below apply to the measured inputs and SM86 configurations, not -all shapes or hardware. +The statuses below describe historical assertions within measured SM86 cases, +not retained complete execution evidence or current H100 acceptance. | Requirement | Implementation | Test | Observed evidence | Remaining gap | Status | | --- | --- | --- | --- | --- | --- | -| Shared forward/backward math at H=7/3072 | Fixed FP32 normalization and VJP | Independent LayerNorm/autograd, both dtypes, random dy/dgate | All 12 backend/dtype/H cases passed | Independent full-image autograd not covered | PROVEN | +| Shared forward/backward math at H=7/3072 | Fixed FP32 normalization and VJP | Independent LayerNorm/autograd, both dtypes, random dy/dgate | All 12 backend/dtype/H cases passed | Independent full-image assertions added; GPU execution pending | PROVEN | | Race-free CUDA shared reduction | Read-result/barrier/reuse | Native harness and geometry under racecheck | 8 tests, zero hazards; original reproducer also green | Other architectures not exercised | PROVEN | | Launch geometry / tiling | CUDA threads 128/256/512; Triton 4/8 warps and 64/128/256 tiles | Public constructors, CPU-reference byte checks | Both dtypes passed all 3 CUDA and 6 Triton configurations | Coverage is H=3072, B=2, S=5 | PROVEN | | Strict byte comparator | Contiguous uint8 views | Signed-zero regression and backend comparisons | Comparator rejects +0 versus -0; full matrix passed | Exceptional floating-point payloads not covered | PROVEN | -| Three real image sizes | S=4096/6889/6032, H=3072 | Both dtypes, random dy/dgate, CPU/native/Triton byte checks | All 6 real-shape cases passed | Full-size oracle is the fixed CPU implementation, not independent LayerNorm/autograd | PARTIAL | +| Three real image sizes | S=4096/6889/6032, H=3072 | Both dtypes, random dy/dgate, CPU/native/Triton byte checks | All 6 real-shape cases passed | Independent FP32 CPU LayerNorm/autograd added; GPU execution pending | PARTIAL | | Indexed select01 | Explicit PyTorch fallback with [2B,3H] modulation | B=2, both index shapes/devices/dtypes, independent row-index and autograd gold; invalid routes | 12 indexed cases passed | No fixed-order indexed backward claim | PROVEN | | FP32 accumulation / final BF16 cast | FP32 reductions, CUDA RN operations, Triton fusion disabled | Cast, byte-reference and independent accuracy tests | All focused checks passed | Exhaustive rounding boundaries and fast-math-on build exclusion not exercised | PARTIAL | | Batch / padding invariance | Row-local hidden reduction, ascending token fold | Batch position, B=1/3, S=3/5, zero-gradient padding; raw bytes | FP32/BF16 CUDA and Triton cases passed | Tested padding matrix uses H=32 | PROVEN | -| Registry / trace | Explicit backend selection, H limit, indexed fallback reason | Registry tests and native preflight | Native selected without fallback; indexed fallback and H limit passed | Semantic fingerprint is not a binary hash; unavailable-backend combinations are not exhaustive | PARTIAL | +| Registry / trace | Explicit backend selection, H limit, indexed fallback reason | Registry tests and native preflight | Native selected without fallback; indexed fallback and H limit passed | Semantic fingerprint is not a binary hash; CPU unavailable-backend routing tests added; actual GPU fallback execution pending | PARTIAL | | Native SM86 build and binding | Conditional sources/symbols, defaulted threads argument | Symbol verification, native/reference tests and geometry launches | User build completed; new four-argument binding and full suite passed | Other build configurations untested | PROVEN | | Indexed optimized kernels | Agreed explicit PyTorch fallback | Fallback and trace checks | Within agreed scope | No optimized indexed kernel required | OUT OF SCOPE | -| Full #386 WS1 exit criteria | All 19 operators and assembled model | MMDiT/LoRA, 60-layer curve, reproducible rollout logp | Not exercised by this operator suite | Separate system-level work | UNPROVEN | +| Full #386 WS1 exit criteria | 22 WS1 rows plus Foundation, WS2 and integration | MMDiT/LoRA, 60-layer curve, reproducible rollout logp | Not exercised by this operator suite | Separate system-level work | UNPROVEN | -The independent autograd checks are accuracy checks, not a replacement tolerance -profile for the fixed-reference byte contract. Shared checks use FP32 -`atol=rtol=2e-5`, BF16 `atol=0.02, rtol=0.016`; indexed checks use -`atol=rtol=2e-6`. The full-size byte checks still use the fixed CPU reference -and do not independently validate LayerNorm/autograd at every image size. +The independent autograd checks use the shared reduction accuracy contract. +They do not replace strict fixed-reference and batch/topology byte checks. +Full-image independent assertions are implemented but await GPU evidence; +exhaustive rounding boundaries and alternative fast-math builds remain gaps. The semantic fingerprint does not identify a compiled binary. This suite does not claim coverage across other GPU architectures, every floating-point input, diff --git a/docs/reviews/adaln-shared-b31fac9-followup.md b/docs/reviews/adaln-shared-b31fac9-followup.md new file mode 100644 index 000000000..f65ab8e32 --- /dev/null +++ b/docs/reviews/adaln-shared-b31fac9-followup.md @@ -0,0 +1,172 @@ +# AdaLN shared：新增缺口与当前版本证据（2026-10-09） + +> 本文是第二轮中间源码快照。用户随后授权继续改实现、补测试；最新状态与证据已统一记录在 +> [指定主报告](adaln-shared-b31fac9.md)。本文的 hash、运行表和未关闭项状态不代表最新工作区。 + +仍以 [issue #386](https://github.com/RL-Align/RL-Kernel/issues/386) 为规格，固定比较点 +`b31fac988faae9de71ab1f683e4a5c9cc8578b84`，HEAD +`120f4de7274798c6310d87799ebb6538f2f8f531`。审查范围包括 `git diff b31fac9` +及所有未跟踪文件；原有用户修改保留，没有提交或推送。 +[首轮报告](adaln-shared-b31fac9.md) 的源码 hash 与运行表属于此前快照，不能用作当前版本运行证据。 + +本次沿用 `/code-review` 的独立 Standards/Spec 审查及 `/tdd` 公开接口测试流程。 +测试只通过公开 op 调用、`torch.autograd` 和 `run_operator_suite` 观察结果。 +确实失败的生产行为逐项先 red 再修复；原本正确但缺断言的行为补 canary，另用公共 +callable 错误候选验证断言敏感性，没有为了制造 red 修改生产实现。 + +PROVEN 仅指表中指定有限反例在当前 CPU/SM86 Triton 版本运行通过。 +PARTIAL 表示有局部证据而未覆盖完整域。UNPROVEN 表示尚无当前版本运行证据。 +这些状态不表示 CUDA/H100 全平台验收完成。 + +## Standards：验证器确实放过错误 dtype + +**P1，已修复。** AdaLN BF16 输入必须输出 BF16，但无 provenance 的默认 gtest candidate +会将输出转 FP32 后比较，不检查声明的输出边界。公开错误候选 +`tuple(output.float() for output in good.fn(**inputs))` 在数值与梯度比较中原本可以通过。 +新增 forward-only 和 check-grad 两个反例均先失败。 + +最小修复：`OperatorCase.candidate_output_dtype` 和 spec 的 `output_matches_input_dtype` +显式启用 AdaLN 输出 dtype 检查；失败报告包含 `dtype mismatch`。 +其他算子默认不启用，避免将本就应输出 FP32 的算子误判。gold leaves 的 FP32 策略继续保留。 +修复后检查器相关 39 项通过，最终 CPU suite 89 项通过。 + +## Spec：三项实际实现错误与其他证明缺口 + +**P1，已修复：Triton 默认 FTZ 改变合法 subnormal eps。** 本机 Triton 3.7.1 默认 +`enable_reflect_ftz=True`；本轮 probe 的 PTX 中出现 `sqrt.rn.ftz.f32`。 +公开 `eps=1e-40`、常量零输入返回 NaN，而 CPU 按 FP32 eps 应输出零。 +forward 与 backward rows 现均显式关闭 reflect FTZ；两 dtype 的 y/dX/dMod literal 验证通过。 +不能仅凭 `sqrt_rn` 名称或 `enable_fp_fusion=False` 断言没有 FTZ。 + +**P2，已修复：Triton tree root 提取额外加正零。** `tl.sum(where(col==0,root,+0))` +把 tree 的负零变成正零，改变 mean 和最后输出符号。现直接 gather 单元素 root 再归约， +移除额外正零。不是放宽 byte assertion。 + +**P2，已修复：CPU BF16 gate FP32 往返转换损坏 NaN payload。** gate 没有算术， +原实现仍先转 FP32 再转回 BF16；现从原调制张量直接取 gate view。继续检查发现 CPU +backward 的 dgate 同样损坏上游 payload，第二轮 red→green 后改为原 dtype 直接拼入 dMod， +只在 dshift/dscale 各自输出边界 cast。finite 与 nonfinite identity VJP 都验证。 +issue 未显式定义 LayerNorm 的 nonfinite 域，本项只证明 gate 位透传，不声称整个 nonfinite +输入域已符合规格。 + +以下新增测试均在 [test_adaln_gtest.py](../../tests/test_adaln_gtest.py)。 +独立 literal 来自标出的标量手算；数学 accuracy 使用独立 CPU FP32 LayerNorm/autograd； +拓扑 differential 只用于 bytes 不变，不充当独立数学参考。 + +| 要求 / 当前实现与独立参考 | 具体错误实现 | 最小反例 | 应失败的断言 | 当前结论 | +| --- | --- | --- | --- | --- | +| BF16 kernel 输出边界；checker 现显式检查 dtype | 正确 BF16 y/gate 再 `.float()` 返回；旧 checker 数值全通过 | B2,S3,H7 BF16，公开错误 candidate | `gtest_rejects_fp32_outputs_for_bf16_inputs`：forward/check-grad report 都失败，两输出有 dtype mismatch | **PROVEN** CPU checker，实际 red→green | +| 固定 hidden tree 的零符号；CPU 独立 IEEE literal | Triton tree root 提取时再加 +0，擦掉负零 mean | H2，x=[-0,-0]，shift=-0，scale=0 | `signed_zero_follows_hidden_tree`：y 必须是 +0 bytes，错误实现是 -0 | **PROVEN** CPU/SM86 Triton，两 dtype；CUDA **UNPROVEN** | +| FP32 scalar/平方根不能无声明 FTZ；独立标量舍入 | `sqrt.rn.ftz` 把 FP32(1e-40) 当作0 | x=[0,0]，m=0，eps=1e-40；dy=[1,-1] | `subnormal_eps_is_not_flushed_to_zero`：y=0；rstd=1.000002658868784e20，dX=±rstd 最终 cast；dMod=[1,-1,0,0,0,0] | **PROVEN** CPU/SM86 Triton，两 dtype，实际 red→green；CUDA **UNPROVEN** | +| gate/dgate 按原 dtype 位透传，独立 raw literal | BF16 gate或dgate→FP32→BF16，NaN canonicalize | gate raw16=[0x7f81,0xffc1]；dg分别取[2,-3]或raw16=[0x7f82,0x7fff]；x、shift、scale有限 | `gate_preserves_bf16_nan_payload_and_identity_vjp`：gate原uint8 bytes不变，dg原位进入dMod | **PROVEN** CPU/SM86 Triton；CPU fwd/bwd各实际red→green；nonfinite LN合同 **PARTIAL** | +| variance 的 hidden tree 必须独立于 mean 的 canary 检验 | mean 正确，只把 squared-centered 的树换成左折叠 | H4 x=[4096,1.25,-4096,-1.25]，eps0；mean 两算法都0；variance 应8388609，错误为8388608 | `variance_reduction_uses_declared_pairwise_order`：y literal ±1.4142134189605713、±0.0004315836704336107 的 bytes | **PROVEN** CPU/SM86 Triton；对应 mutant 被拒；CUDA **UNPROVEN** | +| backward sum(dnorm) 的固定树 | 只把第一个 backward H sum 改左折叠 | x=[-1,1,-1,1]，eps0，dy=[2^24,1,-2^24,1] | `backward_hidden_sums_use_declared_pairwise_order(sign=1)`：dx=[2^24,0,-2^24,0]；错误中间通道为.25 | **PROVEN** CPU/SM86 Triton；单独 mutant 被拒；CUDA **UNPROVEN** | +| backward sum(dnorm*norm) 的固定树 | 只把第二个 backward H sum 改左折叠 | 同上，dy=[-2^24,1,2^24,1] | 同一测试 sign=-1；dx=[-2^24,0,2^24,0]；错误中间通道为.25 | **PROVEN** CPU/SM86 Triton；单独 mutant 被拒；CUDA **UNPROVEN** | +| BF16 最终边界 RN ties-to-even，forward/backward 都须检查 | 中点向上舍入或截断；随机用例很难撞中点 | x=[-1,1]，eps0，正通道 shift=f；S2 正通道 dy=[1,f]；f=1/256、3/256 | `boundary_rounds_halfway_values_to_even`：1+f 应分别 round 到1和1+4/256；y、dshift、dscale bytes | **PROVEN** CPU/SM86 Triton；四个 forward/backward 舍入 mutant 被拒；CUDA **UNPROVEN** | +| upstream stride 必须按公开 autograd 语义处理 | 按紧密连续地址读 dy，或忽略 dg 的 stride0 | B2,S3,H7；dy 从[B,H,S] transpose，dg 首样本 expand 至B2 | `backward_accepts_transposed_and_broadcast_upstreams`：dX/dMod 等于 contiguous upstream 的 bytes，且对独立 CPU FP32 gold 在声明容差内 | **PROVEN** CPU/SM86 Triton，两 dtype；raw-read mutant 被拒；CUDA **UNPROVEN** | +| 输入可分别冻结，输出可真正不参与 loss | 强制两个输入都求导；backward 假设两个输出一定都有传入 VJP | 仅m训练，y.sum且gate不用；仅x训练、dy=[1,0]、x=[-3,3]eps7；仅gate loss | `single_trainable_input_and_unused_output`：冻结输入.grad is None；dm=[1,1,-1,1,0,0]；dx=±7/128；gate-only dx=0/dm gate identity | **PROVEN** CPU/SM86 Triton，两 dtype；不扩展为二阶梯度或 LoRA 模型验收 | +| H3072 的 batch/padding 必须与 launch 同时改变 | warps8/tile256 路径只在 position>0 读取错 mod；单独 topology、单独 geometry 可漏 | default4/128的B1,S3基准；改4/64或8/256，同时B3的每个position与S3/5，strided输入与6H调制切片 | `h3072_relocation_with_alternate_triton_launch`：y/gate/dX/dMod bytes 等于默认 launch 的同一样本；unused半区0，padding dx数值0 | **PROVEN** SM86 Triton 两额外配置×两 dtype；全部六配置/native三threads交叉 **PARTIAL** | +| unsupported batch 几何必须在入口清楚处理 | C++ backward `grid.y=batch`，既未限B也未展平，B65536进入非法grid | B65536,S1,H1，gate-only backward；不到几MiB | 允许的输入应成功且dx0/dm gate1；若明确不支持，应在公开入口清楚拒绝；不能以 invalid configuration launch失败 | **PARTIAL** 源码定位在 CUDA L194；当前 native runtime **UNPROVEN**，未冒充已修复 | +| 地址乘积不能静默 int32 overflow | `row*hidden`、`(batch*seq+token)*hidden`、`batch*3*hidden` 全以int计算，无 numel 范围检查 | x寻址 B1,S524289,H4096 BF16（约4GiB）；mod寻址 forward B174763,S1,H4096 | 有界支持则公开入口拒绝；支持则 int64寻址并memcheck零错误；不是普通数值 allclose | **PARTIAL** 静态类型/边界证据；实际巨大输入与当前 native **UNPROVEN**；未分配/运行 | +| eps 公共值域与FP32 scalar转换策略须明确 | finite double eps 通过校验，但cast后为Inf或0；数值行为未定义 | eps1e39→FP32 Inf，eps1e-46→0；常量x的后者产生NaN | 按选定政策验证round-to-FP32结果，或清楚拒绝超范围值；不能只断言finite(eps) | **PARTIAL** CPU/SM86 probe 已观察转换结果；是否允许此域规格未明确，未擅自更改API | +| native H上限与非二次幂边界必须交叉验收 | off-by-one允许H4097或错拒H4096；尾部padding树错误 | H4095/4096/4097，两dtype，direct native及registry | 4095/4096正确，4097明确拒绝或有声明fallback；独立数学与bytes都检查 | **UNPROVEN** 当前 native边界；旧大H/irregularH测试不替代这组边界 | + +CUDA grid.y 上限65535见 +[NVIDIA compute-capability limits](https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/compute-capabilities.html)。 +第12/13项是未覆盖的实现范围错误候选，不是模型常用 B 的性能问题。 +若将B限制到65535且H≤4096，modulation乘积自然小于int32上限，但x的B*S*H仍需单独检查。 + +## 当前版本运行证据 + +环境仍为 Python3.12.12、PyTorch2.12.1+cu126、Triton3.7.1、RTX3050Ti/SM86。 +没有构建 native；现存 `.so` 未绑定当前源码，因此没有执行其数值测试。 +下面最终 suite 与 mutation 运行对应文末 hash；red/局部green是本轮 TDD过程的中间源码证据。 + +| 运行 | 结果 / artifact | +| --- | --- | +| signed-zero red→green | 2 failed/2 passed→4 passed；[/tmp red](/tmp/rlk-adaln-more-signed-zero-red.log)、[green](/tmp/rlk-adaln-more-signed-zero-green.log) | +| subnormal eps red→green | 2 failed/2 passed→4 passed；[red](/tmp/rlk-adaln-more-subnormal-red.log)、[green](/tmp/rlk-adaln-more-subnormal-green.log) | +| CPU gate payload red→green | 1 failed→1 passed；[red](/tmp/rlk-adaln-more-gate-red.log)、[green](/tmp/rlk-adaln-more-gate-green.log) | +| CPU gate VJP payload red→green | 1 failed/1 passed→2 passed；[red](/tmp/rlk-adaln-more-gate-vjp-red.log)、[green](/tmp/rlk-adaln-more-gate-vjp-green.log)；[Triton局部检查](/tmp/rlk-adaln-more-gate-vjp-triton.log)2 passed | +| gtest dtype red→green | 2 failed→相关39 passed；[red](/tmp/rlk-adaln-more-dtype-red.log)、[green](/tmp/rlk-adaln-more-dtype-green.log) | +| 最终CPU/checker | **89 passed、83 deselected，7.29s，无skip**；[log](/tmp/rlk-adaln-more-final-cpu.log) | +| 最终CPU/SM86 Triton，包括额外launch交叉 | **124 passed、67 deselected，14.07s，无skip**；[log](/tmp/rlk-adaln-more-final-gpu.log)；未运行native、大图、旧完整geometry sweep | +| 最终新增错误候选 | **8/8拒绝**；[script](/tmp/rlk_adaln_more_mutations.py)、[log](/tmp/rlk-adaln-more-mutations-final.log) | +| 最终首轮错误候选重跑 | **6/6拒绝**；[script](/tmp/rlk_adaln_review_mutations.py)、[log](/tmp/rlk-adaln-more-old-mutants-final.log) | + +```bash +# 最终 CPU 命令 +PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 .venv/bin/python -m pytest \ + tests/test_adaln_gtest.py tests/test_op_checks.py tests/test_operator_inputs.py \ + -q -rs -k 'not triton and not cuda' + +# 最终 SM86 短测试命令；明确排除未经当前源码构建的native及长检查 +PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 TRITON_CACHE_DIR=/tmp/rlk-adaln-review-triton \ + .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_adaln_modulation.py \ + -q -rs -k 'not cuda and not real_image and not gpu_padding and not launch_geometry' +``` + +## H100：新增验证清单与命令，未启动 + +先执行[首轮报告的源码快照、pinned构建和完整测试命令](adaln-shared-b31fac9.md)。 +需要将本 followup 与新的 untracked test 一起带过去，不能用旧 hash/旧 binary。 + +在已经完成当前源码构建的 H100 环境中，新增 canaries 可单独检查： + +```bash +RL_KERNEL_REQUIRE_EXT=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ + .venv/bin/python -m pytest tests/test_adaln_gtest.py -q -rs \ + -k 'signed_zero or subnormal_eps or gate_preserves or rejects_fp32_outputs or variance_reduction or backward_hidden_sums or halfway_values or broadcast_upstreams or single_trainable or all_launch_configurations' +``` + +小内存 native batch 边界公开接口复现命令（这里未运行）： + +```bash +RL_KERNEL_REQUIRE_EXT=1 .venv/bin/python - <<'PY' +import torch +from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +op = CudaAdaLNModulationOp() +for batch in (65535, 65536): + x = torch.zeros(batch, 1, 1, device="cuda", requires_grad=True) + m = torch.zeros(batch, 3, device="cuda", requires_grad=True) + y, gate = op(x, m) + dx, dm = torch.autograd.grad(gate, (x, m), torch.ones_like(gate)) + torch.cuda.synchronize() + assert torch.equal(y, torch.zeros_like(y)) + assert torch.equal(dx, torch.zeros_like(dx)) + expected = torch.zeros_like(m) + expected[:, 2] = 1 + assert torch.equal(dm, expected) + print("PASS batch", batch) +PY +``` + +按当前未限定B的API，第二个batch是预期暴露错误的点。若修复选择限定支持域， +需同步声明限制并将边界断言改为入口明确拒绝，不能将 invalid launch当正常拒绝。 +int32 巨大输入反例只准备参数与memcheck清单；没有分配约4GiB输入，更没有启动该测试。 + +- 全部新增 canary 在当前 native/SM90 和Triton两dtype运行，记录零skip/xfail与完整输出。 +- 把全部3种native threads、全部6种Triton warps/tile与H3072 relocation/padding交叉， + 直接比较同一默认B1基准的 y/gate/dX/dMod bytes；本机只新增两组Triton配置。 +- native B65535/65536、H4095/4096/4097正负边界不能遗漏。 +- 寻址选择64位支持或显式有界拒绝；若支持巨大输入，再用compute-sanitizer memcheck检查 + B1,S524289,H4096及modulation地址范围，保留0 errors证据。 +- 对新FTZ修复保存当前Triton版本/compile options/PTX，确认fwd和bwd_rows均无sqrt FTZ； + 其他Triton版本/ROCm配置须独立qualification。 +- 保留首轮要求的三真实image S、benchmark、racecheck、源码/二进制hash及pinnedruntime记录。 + +## 对应最终测试的源码 SHA256 + +| 文件 | SHA256 | +| --- | --- | +| 未跟踪 `tests/test_adaln_gtest.py` | `85de4e0760d6ead947b0af1e8bd4eb896cfdea6bccf4434ca46d5a3adf67e165` | +| `gtest/op_checks.py` | `70f9c648b5da5ace001d295569c4b05a15d2e630e84e1870b5734a5377d76b39` | +| `gtest/operator_specs.py` | `b78d15322199b4986f19697c7987dea879238b7d8b14fe2f1b05fc3fe0268e79` | +| PyTorch AdaLN | `1b573827279e89c97e4f293bf867fe5a3a77102c89cf9427ff0e1f429a351c57` | +| Triton AdaLN | `ef2b13c8e13da8d324b160489efad4b2080d3a1058cf3b58d566c2bc92a3b4a2` | +| CUDA source（未修改） | `27074706332b1db32236b6cbb225ee99514938f38d1996ef4812e221d857e80d` | + +静态检查：修改的五个Python文件ruff/isort均通过；新增测试与checker的Black check通过， +`git diff --check`通过，验收shell的`bash -n`通过。 + +`/tmp`日志与mutation脚本是本次artifact，应另行归档。新增文档不会改变上述测试源码hash。 diff --git a/docs/reviews/adaln-shared-b31fac9.md b/docs/reviews/adaln-shared-b31fac9.md new file mode 100644 index 000000000..d9ff01587 --- /dev/null +++ b/docs/reviews/adaln-shared-b31fac9.md @@ -0,0 +1,351 @@ +# AdaLN shared：工作区审查与 TDD 证据（2026-10-09） + +## 剩余缺口与本次执行范围 + +用户已授权继续按 `/tdd` 改实现、补测试；沿用公开 op 调用、autograd、registry、gtest +接口，不要求新的接口确认。不提交/推送，不启动 H100 长测试、巨大输入或无法估计耗时的构建。 +下面是实施前待办；其最新完成状态和本次运行证据记录在下一节。 +本文件下方第一轮审查及 followup 文档都属于历史源码快照，旧运行表不能算当前版本证据。 + +| 缺口 | 修补动作 | 实施前结论 | +| --- | --- | --- | +| CUDA B65536 超出 backward grid.y | Python/C++ 声明并检查 B≤65535，补65535/65536边界测试，launch前明确拒绝 | 源码PARTIAL，当前native运行UNPROVEN | +| CUDA int32寻址范围无检查 | 在复制/launch前检查索引乘积，Python/C++双层保护；巨大形状使用meta输入测试，不分配几GiB张量 | 源码PARTIAL，运行UNPROVEN | +| shared eps和nonfinite支持域未明确 | 拒绝FP32转换后Inf或正eps下溢到0，保留合法subnormal与eps=0；补公开反例并记录nonfinite LN范围 | PARTIAL | +| geometry×batch/padding交叉不完整 | 补全6种Triton配置及3种CUDA threads，比较默认B1同样本y/gate/dX/dMod bytes | PARTIAL | +| GPU两stream及第二triple覆盖不足 | 扩展已有小形状公开测试到各后端，独立FP32梯度、unused半区0、另一stream不受影响 | PARTIAL | +| H4095/4096/4097 native边界 | 补direct op与registry正负边界，两dtype | UNPROVEN | +| 当前源码/native binary闭环缺失 | 保存源/产物hash；仅在能够可靠估计快速完成时构建，否则提供构建与验收命令 | UNPROVEN | +| H100真实图像、安全性、性能 | 保留三真实S、racecheck/memcheck、benchmark和pinned环境命令与清单，不自行启动 | UNPROVEN | + +## 本次 TDD 完成状态与当前源码证据 + +以下均通过公开 op、autograd、registry 或 gtest 观察，没有新增生产 helper mock。 +已复用旧检查器、固定字节参考、独立 FP32 CPU LayerNorm/autograd 和模型六块切片用例。 +比较点仍为 `b31fac9`，HEAD仍为 `120f4de7274798c6310d87799ebb6538f2f8f531`； +范围包含当前 tracked diff 与未跟踪的两个测试文件和审查文档。未提交或推送。 + +实现采用明确的有界拒绝：native 支持 B≤65535、H≤4096、每个输入 numel≤2^31-1。 +Python 在 `.contiguous()` 前校验，C++ 在 narrowing/launch 前重新校验；没有扩大支持域， +因此不需要实际分配约4GiB的溢出反例。registry只按H选择后端，B/numel超限在公开调用 +入口清楚拒绝，不算隐含fallback。meta测试验证的是公开Python形状preflight，不是native内核。 + +shared eps 检查 round-to-FP32 后仍有限，且正值不变成0；合法subnormal和eps=0保留。 +nonfinite LN按FP32公式传播，eps=0且方差0可以产生NaN；不承诺LN运算的NaN payload +或非有限运算域的跨后端字节相等。gate/dgate无算术，仍按原dtype位透传。 +这不是给整个nonfinite输入域宣称PROVEN，也没有改变独立indexed/select01合同。 + +| 缺口 / 可能错误的实现 | 最小反例与关键断言 | 已完成的修改 / 测试 | 当前结论 | +| --- | --- | --- | --- | +| B65536直接进入backward grid.y | B65536,S1,H1；公开调用必须在launch前明确拒绝；B65535的y/gate/dX=0，dMod gate=1 | Python/C++新增B≤65535校验；`shared_bounds::rejects_batch_above_grid_limit`及`maximum_batch_forward_backward` | Python拒绝 **PROVEN**（meta及SM86 CUDA tensor）；C++产物与合法最大B数值 **UNPROVEN**，测试已准备 | +| x地址或seq narrowing以int溢出，先连续化会尝试几GiB分配 | meta x形状[1,524289,4096]、[1,2^31,1]；必须报int32 indexing错误，无实数据分配 | Python/C++检查x和modulation numel≤2^31-1；`rejects_int32_address_overflow_without_allocating` | Python metadata拒绝 **PROVEN**；C++重新构建后的运行 **UNPROVEN** | +| double eps有限但FP32为Inf/0 | 1e39、1e-46必须明确ValueError；1e-40、2^-149、FP32最大有限值必须保持y0、dX=±rstd、dMod=[1,-1,0,0,0,0] | 共享Python校验用stdlib struct检查FP32舍入，C++也检查；公开`eps_overflow`及扩展`subnormal_eps` literal测试，两dtype | 三Python公开后端拒绝 **PROVEN**；合法边界CPU/SM86 Triton **PROVEN**；native合法eps数值 **UNPROVEN** | +| 把nonfinite LN误当finite验收；gate/VJP因NaN LN串值 | x=[NaN,1]或[Inf,1]，y/dX/dscale应NaN；dshift=[1,-1]、gate=[2,-0]、dgate=[2,-3]独立保持 | 声明nonfinite传播和qualification范围，新增`nonfinite_activation_propagates_nan_and_preserves_gate_vjp` | CPU/SM86 Triton指定反例 **PROVEN**；一般NaN payload与strict非有限域不在资格声明中 | +| 只单独比较geometry或relocation，错过组合路径 | H3072：默认B1,S3基准→B3每个position、S3/5、strided x和6H调制；y/gate/dX/dMod bytes一致、unused半区0 | `h3072_relocation_with_all_launch_configurations`覆盖全部6种Triton warps/tile及3种CUDA threads，两dtype | 全部Triton组合 **PROVEN**（12参数例）；native组合测试已写、运行 **UNPROVEN** | +| GPU第二triple/两stream不独立；只验梯度会漏错shift | image S5/text S3，各取前/后3H；另一个stream梯度不变、unused半区0；错误候选y+1仍有正确VJP | 原`gate_only_and_separate_stream_chunks`扩展到三后端，独立CPU gold的forward+backward、gate bytes及另一stream隔离；错误shift候选旧断言实际放过，新断言拒绝 | CPU/SM86 Triton **PROVEN**（8参数例，均遍历两stream）；native测试已写、运行 **UNPROVEN** | +| H上限或零填充off-by-one | H4095/4096：两dtype y/gate/dX/dMod对独立CPU数学及固定参考bytes；H4097直接拒绝、registry明确fallback | 新`hidden_boundary_matches_independent_math_and_fixed_bytes`及`shared_bounds` direct/registry三个H边界测试 | CPU/SM86 Triton合法H数学 **PROVEN**；Python direct拒绝及registry **PROVEN**；native合法H数值 **UNPROVEN** | +| 用旧_C或语义fingerprint证明当前C++ | 当前CUDA源码已有校验变更，现存_C未重建，不能因symbols/constructor可用就算通过 | 记录新source hash与旧binary hash，准备fresh build/验收命令，未调用未绑定当前源码的native数值路径 | 当前native内核资格 **UNPROVEN** | +| 三真实图像S、shared race、OOB、claimed平台性能缺当前证据 | S4096/6889/6032,H3072、两dtype、两个GPU实现；racecheck/memcheck必须0 errors；benchmark完整数据 | 更新04验收脚本、06racecheck脚本及WS1/CPU CI入口接入新测试；准备下面H100命令，未启动CI或长测试 | H100/真实图像/sanitizer/benchmark **UNPROVEN** | + +当前合法eps的手算rstd literals由独立stdlib math+IEEE FP32舍入得到，分别为 +`1.000002658868784e20`、`2.671373844909537e22`、`5.421011508662376e-20`。 +隐藏边界的数学参考使用独立CPU FP32 leaves；固定CPU参考只作为字节/归约合同参考， +不将其当独立数学gold。三种比较不能互相替代。 + +### 当前版本运行表 + +下面最终suite和错误候选重跑对应本节末尾hash；red/局部green是本次TDD的中间源码证据。 +旧 `.ninja_log` 只用于评估构建耗时：full build历史上超过17分钟，不作为正确性运行证据。 +没有自行启动full build，也没有为了构建安装/下载依赖或索取新权限。 +本机GPU检查沿用已有Python执行权限，NVIDIA RTX3050Ti/SM86、PyTorch2.12.1+cu126、 +Python3.12.12、Triton3.7.1；未声明SM90或当前native binary资格。 + +| 本次检查 | 结果 / evidence | +| --- | --- | +| B限制red→green | meta公开入口1 failed→1 passed；[red](/tmp/rlk-adaln-close-batch-red.log)、[green](/tmp/rlk-adaln-close-batch-green.log)；[SM86公开拒绝](/tmp/rlk-adaln-close-batch-gpu-green.log)1 passed。没有故意启动非法native launch | +| int32范围red→green | 2 failed→bounds 3 passed；[red](/tmp/rlk-adaln-close-int32-red.log)、[green](/tmp/rlk-adaln-close-int32-green.log)，全是小输入/meta，无巨大分配 | +| eps范围red→green | CPU4 failed→4 passed；[red](/tmp/rlk-adaln-close-eps-red.log)、[green](/tmp/rlk-adaln-close-eps-green.log)；[合法/非法eps小形状CPU+Triton](/tmp/rlk-adaln-close-eps-boundaries.log)20 passed | +| 第二triple错shift断言red→green | [错误候选旧断言SURVIVED，exit1](/tmp/rlk-adaln-close-shift-mutant-red.log)→[新断言REJECTED，exit0](/tmp/rlk-adaln-close-shift-mutant-green.log)，不是生产算子数值错误的声明 | +| 全Triton geometry×topology | [log](/tmp/rlk-adaln-close-all-triton-launches.log)12 passed/9.68s；全部6配置×2dtype，每例含B3的每个position与S3/5 | +| 两stream/两triple | [log](/tmp/rlk-adaln-close-streams-green.log)8 passed/5.43s；含独立forward及backward | +| H4095/4096合法边界数学与bytes | [log](/tmp/rlk-adaln-close-hidden-math.log)8 passed/25.65s | +| B/H/寻址公开preflight及registry | [log](/tmp/rlk-adaln-close-hidden-preflight.log)8 passed/8.73s；未运行maximum_batch数值例 | +| nonfinite传播/独立gate VJP | [log](/tmp/rlk-adaln-close-nonfinite.log)8 passed/5.15s；只断言NaN传播，未将NaN算术payload计为byte资格 | +| 最终CPU/checker/preflight | [log](/tmp/rlk-adaln-close-final-cpu.log)：**113 passed，139 deselected，10.84s，0 skip/xfail** | +| 最终CPU+SM86 Triton+CUDA公开拒绝 | [log](/tmp/rlk-adaln-close-final-gpu.log)：**180 passed，91 deselected，11.00s，0 skip/xfail**；排除native数值、真实图像、旧完整geometry/sanitizer | +| 最终错误候选重跑 | [6个原候选](/tmp/rlk-adaln-close-mutants-six-final.log)、[8个后续候选](/tmp/rlk-adaln-close-mutants-eight-final.log)、[新shift候选](/tmp/rlk-adaln-close-shift-mutant-final.log)：**15/15拒绝**；[新候选脚本](/tmp/rlk_adaln_close_shift_mutant.py) | + +静态检查通过：7个相关Python文件ruff，新增/本次修改文件isort及Black检查， +三个修改shell的`bash -n`、CPU CI YAML解析、`git diff --check`。 + +```bash +# 当前CPU命令：bounds是metadata/preflight，不运行native内核 +PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 .venv/bin/python -m pytest \ + tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py \ + tests/test_op_checks.py tests/test_operator_inputs.py -q -rs \ + -k '(not triton and not cuda) or (shared_bounds and not maximum_batch)' + +# 当前SM86短检查。eps_overflow和shared_bounds仅公开拒绝/registry，不调用native数值路径 +PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 RL_KERNEL_REQUIRE_EXT=1 \ +TRITON_CACHE_DIR=/tmp/rlk-adaln-review-triton .venv/bin/python -m pytest \ + tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py tests/test_adaln_modulation.py \ + -q -rs -k '(not cuda and not real_image and not gpu_padding and not launch_geometry) or eps_overflow or (shared_bounds and not maximum_batch)' +``` + +### 当前源码 SHA256 + +| 文件 | SHA256 | +| --- | --- | +| 未跟踪 `tests/test_adaln_gtest.py` | `db9617770d5637a4aa35e4d0688bb3e3e3abdb7c127f5d04288ab12ce6237119` | +| 未跟踪 `tests/test_adaln_shared_bounds.py` | `97c5a99ae6f73545304114786e935d30f21af8a4bbe73984c296972ae997b0ba` | +| CUDA wrapper | `e5a52b8e29c6be131176e0fdad64aa259cc780b5a4c7529bad98901b7ec53797` | +| PyTorch AdaLN | `96707e412a15790962a15a939ecc835eaf363829c515304a2467db2a294fb5c6` | +| Triton AdaLN | `ef2b13c8e13da8d324b160489efad4b2080d3a1058cf3b58d566c2bc92a3b4a2` | +| 当前 CUDA source | `aae87fb84f9b41954a727aa09e926ddfc4deb7c82130aca508607248d0670c05` | +| gtest checker | `70f9c648b5da5ace001d295569c4b05a15d2e630e84e1870b5734a5377d76b39` | +| gtest specs | `b78d15322199b4986f19697c7987dea879238b7d8b14fe2f1b05fc3fe0268e79` | +| 现存旧_C（没有本次构建来源证明） | `d7b91b28b81c7f53db47e757eab9305f617dbc5649666daa3237e9946eb350b6` | + +### 待用户执行的native/H100检查，未启动 + +下方历史报告的H100构建命令仍是待执行方案；必须先从当前dirty/untracked源码重新构建， +不能复用上表旧_C。源码快照要包含两个未跟踪测试及本文件。建议保存build log、完整 +compiler flags、源码/产物hash、Python/PyTorch/Triton/CUDA/driver/GPU信息，再执行： + +```bash +# 仅在当前源码已fresh build的准备好环境中运行；这里未执行 +RL_KERNEL_REQUIRE_EXT=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ + .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py \ + -q -rs -k cuda + +# 完整真实形状与数值入口；04已包含新bounds，06包含完整native launch×topology +bash scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh + +# 支持域已限定，巨大输入应在入口拒绝；memcheck检查真实支持边界，不分配4GiB反例 +RL_KERNEL_REQUIRE_EXT=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ + "${CUDA_HOME:?set CUDA_HOME}/bin/compute-sanitizer" --tool memcheck --error-exitcode 7 \ + .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py \ + -q -rs -p no:cacheprovider \ + -k 'cuda and (hidden_boundary or maximum_batch or all_launch_configurations)' + +bash scripts/adaln_cuda_acceptance/05_benchmark_adaln_cuda.sh +``` + +验收必须覆盖：native全部canaries、最大B65535与超限65536、H4095/4096/4097、三threads、 +两dtype、两stream/两triple；Triton全六配置；S4096/6889/6032真实图像；racecheck/memcheck +0 hazards/errors；benchmark有每个dtype/S/backend的forward及forward+backward数据。 +要求零skip/xfail、native无fallback，source/binary在验收期间不变；配置和CLI报告不是运行证据。 +一般NaN算术payload、ROCm或其他未固定Triton环境不在本机资格声明中。 + +## 第一轮审查及运行记录（历史源码快照） + +> 下方内容及[第二轮报告](adaln-shared-b31fac9-followup.md)保留中间源码证据。 +> 它们的运行表和hash不代表当前工作区；最新状态以上面“本次TDD”章节为准。 +> H100/native构建与长验收没有在本次执行。 + +规格：[RL-Align/RL-Kernel #386](https://github.com/RL-Align/RL-Kernel/issues/386)。 +比较点 `b31fac988faae9de71ab1f683e4a5c9cc8578b84`，HEAD +`120f4de7274798c6310d87799ebb6538f2f8f531`。范围是 **比较点到实际工作区**， +使用 `git diff b31fac9` 加 `git ls-files --others --exclude-standard`，包括索引、 +未提交修改和未跟踪 `tests/test_adaln_gtest.py`。三点 HEAD diff 不代表本次范围。 + +开始时工作区有 11 个已修改文件和上述未跟踪文件；其既有修改均保留。 +已读比较点之后五个提交、CUDA/PyTorch/Triton 实现、bindings、registry、 +gtest reference/inputs/checker、两个测试文件、benchmark、acceptance scripts、CI 和文档。 +indexed select01、projection、gate residual、最终 AdaLN、完整 MMDiT/WS2 不属于本次 shared 实现验收。 + +这里 **PROVEN** 只表示明确列出的当前版本、后端与有限反例已经运行通过； +**PARTIAL** 表示有实现或部分证据，仍有未运行的平台/几何/验收项； +**UNPROVEN** 表示没有当前工作区的有效运行证据。没有用旧日志填补当前证据。 + +## Standards + +1. **P1，已修复：BF16 gtest 参考梯度被再次舍入成 BF16。** + [gtest 规范](../contributing/gtest-usage.md)要求独立 FP32 reference。 + `GtestAdaLNReference.forward_fp32` 虽在内部转 FP32，但旧 checker 给 gold 克隆的 + 求导叶子仍为 BF16；autograd 返回 BF16 dX/dMod，报告却标 `fp32_reference`。 + 新增公开 `run_operator_suite` 断言所有 output/gradient checks 的 `gold_dtype` + 都为 `torch.float32`；BF16 H=7 的运行实际失败（1 failed、1 passed，exit 1）。 + 修复为 `OperatorSpec/OperatorCase.gold_input_dtype` 显式启用 FP32 gold leaves； + 只转换浮点张量，candidate 和其他算子的默认输入策略保留。随后相关 47 项检查通过。 + 这修复的是验证参考，未假称发现 CUDA 算术错误。 +2. **判断性异味：registry 的两条 trace 分支重复若干精度/算法字段。** + 可能导致未来元数据漂移；目前无实际错误证据，没有为此增加抽象或重构。 + +另修复 focused shared CPU 测试无条件导入 Triton 的问题，并将两个 GPU 测试的 +Triton import 移到 CUDA 可用检查之后。CPU 原本不需要这个可选依赖。 + +Standards：1 个已修复规范问题、1 个判断性异味；最严重项是参考梯度精度声明错误。 + +## Spec + +原有 dtype/finite/容差断言不足以证明 BF16 输出边界；H=17/32 的 relocation/padding +也不能覆盖 H=3072 的批位置。现有 H=3072、B=2 geometry 测试比较的是各自 CPU +reference，没有把同一个样本从 B=1 移到 B=3 的每个位置。两项均已补公开入口反例。 +另有明确的独立 FP32 gold 梯度错误，已按上述 red→green 修复。 +shared 非法输入的反例也已补齐；当前 native binary、真实图像尺寸、racecheck、H100 +和当前 benchmark 仍没有本轮运行证据。 + +Spec：3 个核心精度/拓扑缺口、1 组非法输入覆盖缺口、1 组平台验收证据缺口； +最严重项是原 gtest 的 FP32 gold 梯度并非 FP32。没有合并或重排 Standards 结论。 + +## 要求、独立参考、反例与关键断言 + +位置简称：`native` = `ops/pytorch/norm/adaln_modulation.py`；`cuda` = +`csrc/cuda/adaln_modulation.cu`；`triton` = `ops/triton/norm/adaln_modulation.py`。 +新测试均在 [test_adaln_gtest.py](../../tests/test_adaln_gtest.py)，旧断言保留在 +[test_adaln_modulation.py](../../tests/test_adaln_modulation.py)。 + +| 要求 / 当前实现与参考 | 具体可能错误的实现 | 最小反例 | 应失败的断言 / 测试 | 当前结论 | +| --- | --- | --- | --- | --- | +| 无 affine LayerNorm，shift/scale/gate 顺序，默认 eps=1e-6；三后端均按此实现；独立 CPU `F.layer_norm` + autograd | 用 RMSNorm，交换 shift/scale，gate 混入 y | x=[1,3]、shift=[.5,-.5]、scale=[1,0]、gate=[2,3]、eps=0；y=[-1.5,.5] | 旧定值 forward/backward；独立 shared H=7/3072、非均匀 dy/dg 的 accuracy 断言 | **PROVEN** CPU/SM86 Triton 的已跑小形状；CUDA **UNPROVEN** | +| y、gate 及 dX/dMod 独立 FP32 CPU reference；现已启用 FP32 gold leaves | 内部转 FP32但向 BF16 叶子求导，或求导后再 `.float()` 掩盖已舍入的参考 | BF16 B=2,S=3,H=7 gtest case | 所有四项 check 的 `gold_dtype == torch.float32`；独立 FP32 leaves 的 accuracy；旧错误实际 red、新 checker green | **PROVEN** 当前 CPU checker；native 平台验收另计 | +| BF16 normalization 中间量不得提前舍入 | 把 norm 转 BF16 再做 scale/shift | BF16 x=[-1,1]、shift=-1、scale=0、默认 eps；正通道应为 -2^-21，早舍入为 0 | `bf16_forward_casts_only_at_output` 对手算 literal 做 uint8 bytes 比较 | **PROVEN** CPU/SM86 Triton 该反例；完整内部精度政策 **PARTIAL** | +| `1+scale` 和 affine 保持 FP32 到输出边界 | BF16 的 `1+1/256` 先舍入成 1 | 同上，scale=1/256；正通道 BF16 应为 1/256，错误结果约为 -2^-21 | 同一 forward canary 的第二个参数例；不能用容差替代 | **PROVEN** CPU/SM86 Triton 该反例；CUDA **UNPROVEN** | +| backward dnorm、partial、token accumulators FP32 | BF16 dnorm、BF16 partial，或逐 token BF16 fold | B=1,S=3,H=2，x 每行[-3,3]，eps=7，scale=1/256，dy 首通道[256,1,-256] | `backward_keeps_partials_and_accumulators_fp32`：dshift=[1,0]，dscale=[-.75,0]；dX 按独立手算系数1799/32768；gate VJP=[-2,.5] | **PROVEN** CPU/SM86 Triton 该反例；CUDA **UNPROVEN** | +| BF16 backward 保存的 norm/rstd/partials 不能成为隐含额外 cast | 只在 BF16 backward 保存量转 BF16，forward 完全正确 | 相同 BF16 量化输入、dy/dg，H=7/3072，分别以 BF16 和提升后的 FP32 调用公开 op | `bf16_vjp_matches_fp32_with_one_final_cast`：y/gate/dX/dMod 必须等于 FP32 调用结果最终 cast 的 bytes；独立 math 另测，避免把 differential 当独立数学 gold | **PROVEN** CPU/SM86 Triton 的固定 seed；更广输入域 **PARTIAL** | +| H 固定 lower+upper pairwise tree，非二次幂零填充 | 改成 ascending left fold、按线程分组求和 | FP32 H=4 x=[2^24,1,-2^24,1]，pairwise mean=.5，left-fold mean=.25 | `hidden_reduction_uses_declared_pairwise_order` 的手算 FP32 literal bytes；旧 H=7/3072 accuracy/byte tests | **PROVEN** CPU/SM86 Triton 该反例；其他重排与 CUDA **PARTIAL** | +| dMod 按逻辑 token 递增 FP32 fold；无 atomic/Split-K/Stream-K | 用 pairwise token tree、调换 token/分块重关联 | S=4,H=2 x=[-1,1]，dy 首通道[2^24,1,-2^24,1] | `modulation_gradient_uses_ascending_token_order`：dshift=1、dscale=-1；改树可能得到2/-2 | **PROVEN** CPU/SM86 Triton 该反例；禁止算法同时有源码证据，native **UNPROVEN** | +| batch size、batch position、padding 下有效样本 bytes 不变 | H=3072 路径始终读 mod[0]、错误 batch stride、跨样本共享 dMod | B=1,S=3,H=3072 基线→B=3 全部位置；S=3/5；其他样本不同 x/m/dy/dg，padding dy=0 | `h3072_sample_bytes_survive_batch_positions_and_padding`：y/gate/dX/dMod uint8 相等；包含非连续 x 和六块 view；未使用调制块梯度为0 | **PROVEN** CPU/SM86 Triton BF16/FP32 本次矩阵；更大 B/S 与 CUDA **PARTIAL** | +| launch geometry / tiling 不改 bytes | H tree 随 CUDA threads、Triton warps/tile 改变；跨 warp 共享内存未同步 | H=3072,B=2,S=5，同一输入换 threads128/256/512，warps4/8、tile64/128/256 | 旧 `launch_geometry_preserves_bytes` 比较 y/gate/dX/dMod 对固定 CPU reference | 测试已存在，本轮未运行 geometry sweep；**UNPROVEN** 当前完整 sweep | +| 两 stream 独立调制、attention/MLP 两 triples、gate downstream gradient、strided 模型 view | 串用 image/text modulation；用错第二 triple；gate detached；漏掉 contiguous 复制 | image S=5、text S=3，各自 B=2,H=7、6H 调制；分别取前/后3H，非均匀 dy/dg 和 gate-only | `gate_only_and_separate_stream_chunks` 已扩展到两 triples，独立 CPU accuracy、未使用半区0、另一 stream 无梯度；H3072 topology 的 strided view | **PROVEN** CPU 两 stream/两 triples；SM86 Triton H3072 strided 第一 triple **PROVEN**；GPU 第二 triple **PARTIAL** | +| unsupported shared 输入清楚拒绝；CUDA H<=4096，几何参数有界 | empty/mixed shape 被 launch；FP16 悄悄进入；NaN eps 未拒绝；H>4096 native误选 | 任一空 B/S/H、rank2、错3H、mixed dtype/device、FP16、eps<0/NaN/Inf | 新 `shared_rejects_unsupported_inputs` 和 `shared_rejects_mixed_devices`；旧 large-H registry test；constructors 源码验证 launch 参数 | **PROVEN** CPU/SM86 Triton 输入矩阵；CUDA H/geometry 的当前 runtime **PARTIAL** | +| 实际 backend/fallback、算术 policy 与 kernel identity 可追踪 | native缺失但谎称CUDA；将语义version当编译binary hash；fallback算作strict成功 | native symbol missing、Triton missing、所有后端missing | 旧/新 registry resolution tests；当前 preflight selected native但仅说明构造/符号可用；actual execution须结合direct candidate+source/binary manifest | **PARTIAL**；语义 fingerprint 不绑定二进制，不能凭它证明当前 native | +| CUDA shared reduction race-free | 读shared[0]后其他warp立即重用buffer | H=64 或3072、多warp fwd/bwd | `compute-sanitizer --tool racecheck --error-exitcode 7`；数值 bytes 相等不是 race assertion | **UNPROVEN** 当前 binary；旧 zero-hazard 日志不计 | +| 三参考图像尺寸与短 synthetic，forward/backward、两 dtype、两个GPU实现 | S=6889 tail越界、dMod 漏 token、仅注册优先backend正确 | S={4096,6889,6032},H=3072，随机 dy/dg | 旧 real-image tests 同时独立 FP32 CPU accuracy、fixed-reference bytes、CUDA↔Triton bytes | synthetic 已有当前证据；三 full-image 当前 **UNPROVEN**，H100命令如下 | +| 当前 claimed CUDA platform benchmark、pinned环境 | 只用旧binary/旧timings，CUDA skip却当benchmark完成 | 当前源码构建→两dtype三S，fwd及fwd+bwd | benchmark backend标签、无native skip、完整timings/peak memory、build/runtime/source manifest | 实现存在，当前运行 **UNPROVEN** | +| 无TF32、fast math、compiler-dependent重关联；unsupported严格配置拒绝 | 编译开启fast math仍注册native，或Triton融合乘加 | fast-math-on build / 不同编译配置 | CUDA RN intrinsics、宏#error、setup条件排除symbols；Triton enable_fp_fusion=False；必须验独立build/provenance，不用普通数值测试冒充compiler proof | **PARTIAL** 静态证据；替代build/H100 profile **UNPROVEN** | + +没有把原本正确的算子改成迎合测试的实现。补测时两个期望问题已纠正：大值 +hidden-tree 反例的标准化值应为 ±sqrt(2)，不是 ±1/sqrt(2)；padding 零 VJP 可为 +负零，规格要求有效行 bytes 不变，并未要求 padding 的零必须为正零。 +有效行、gate 和有效 dMod 的原始字节比较没有放宽。 + +## 本轮运行证据 + +所有结果来自本次工作区。源文件与测试的 SHA256 见下，未提交或推送。 + +| 运行 | 命令 / artifact | 结果 | +| --- | --- | --- | +| 初始 red | `pytest tests/test_adaln_gtest.py -q -x -k outputs_and_gradients` | BF16 gold_dtype 断言失败，exit 1；1 failed、1 passed、6 deselected | +| 修复后 checker 回归 | `pytest tests/test_adaln_gtest.py tests/test_op_checks.py tests/test_operator_inputs.py -q -rs`，当时尚未追加其余 canaries | 47 passed / 19.19s | +| 最终 CPU/checker | `PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_op_checks.py tests/test_operator_inputs.py -q -rs -k 'not triton and not cuda'`；[/tmp CPU log](/tmp/rlk-adaln-review-cpu-20261009.log) | 68 passed、41 deselected / 16.94s；无 skip | +| 最终 CPU+SM86 Triton | `PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 TRITON_CACHE_DIR=/tmp/rlk-adaln-review-triton .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_adaln_modulation.py -q -rs -k 'not cuda and not real_image and not gpu_padding and not launch_geometry'`；[/tmp GPU log](/tmp/rlk-adaln-review-gpu-20261009.log) | 80 passed、48 deselected / 32.06s；无 skip。此选择不包含 native CUDA 或 geometry sweep | +| 断言敏感性 | [/tmp mutation probe](/tmp/rlk_adaln_review_mutations.py)，[/tmp log](/tmp/rlk-adaln-review-mutations-20261009.log) | 6/6 公共接口模拟错误被拒绝：early norm、early affine、BF16 token fold、token reorder、hidden left fold、错误 batch 调制 | +| CPU 可选依赖隔离 | [/tmp probe](/tmp/rlk_adaln_review_cpu_only.py)，[/tmp log](/tmp/rlk-adaln-review-cpu-portability-20261009.log)；拒绝所有 Triton imports，再运行 focused shared PyTorch independent cases | 4 passed、52 deselected / 3.44s | +| 静态检查 | ruff、isort、Black check：新增测试与 checker/spec；`git diff --check`；两个修改 shell 的 `bash -n` | 全部通过；Black 按 py310、line-length100 逐文件检查 | + +Mutation probe 用错误候选替换测试的公共 callable，不修改生产 helper，也不是硬件证据。 +独立 LayerNorm/autograd 检验数学正确性；同值 FP32/BF16 differential 检验 cast 边界； +literal 与 uint8 比较检验特定舍入/归约错误。三者分工，不能互相替代。 + +本机环境：Python 3.12.12、PyTorch 2.12.1+cu126、CUDA runtime 12.6、Triton 3.7.1， +RTX 3050 Ti Laptop / SM86。没有本轮 driver/toolkit build qualification 声明。 +native `_C` 已存在但没有与当前源码绑定的 build artifact，本轮没有执行其数值测试。 + +| 文件 | SHA256 | +| --- | --- | +| `tests/test_adaln_gtest.py`（未跟踪文件，已纳入） | `9fea7d7ae184b74641ef8368734a446169c05d7063259956eb42d7a5d28e77f2` | +| `gtest/op_checks.py` | `8b1ff607983d40d60c3a21e03302884228c75ec1beb56b20dccf25215f9c18c1` | +| `gtest/operator_specs.py` | `93c44ff298f2896de40c0b43d9fc5314ef377f0d7c2294cb2420826116ac3917` | +| PyTorch AdaLN | `682a32ffad1ecbbdb28b7bdbb25413acabdde51c71c784656e2d272e16b19003` | +| Triton AdaLN | `7b08d9b4bc047a342875040372bd128f2cbc66eb02779e1c9f9520160b28f781` | +| CUDA source | `27074706332b1db32236b6cbb225ee99514938f38d1996ef4812e221d857e80d` | +| 现存 `_C.cpython-312-x86_64-linux-gnu.so` | `d7b91b28b81c7f53db47e757eab9305f617dbc5649666daa3237e9946eb350b6`；仅记录身份，非当前 native 验收 | + +tracked diff 快照：[patch](/tmp/rlk-adaln-review-tracked-b31fac9.patch)； +未跟踪测试快照：[test source](/tmp/rlk-adaln-review-test_adaln_gtest.py)； +规格快照：[issue JSON](/tmp/rlk-adaln-review-issue386-20261009.json)。 +这些 `/tmp` 文件应随验收 artifact 保存，不能假设长期存在。 + +## H100：只准备命令与验证清单，未执行 + +在用户准备好的、依赖已经安装且版本已固定的 H100 checkout 中执行。以下包含可能长时间 +运行的构建、真实大图、racecheck 与 benchmark；本轮均没有启动。 +H100 native 必须重新构建，不能沿用本机 SM86 `.so` 或仅按符号可用判定成功。 + +```bash +set -euo pipefail +export ADALN_EVIDENCE=/tmp/adaln-h100-current-worktree +mkdir -p "$ADALN_EVIDENCE" +git rev-parse HEAD > "$ADALN_EVIDENCE/head.txt" +git status --short --untracked-files=all > "$ADALN_EVIDENCE/status.txt" +git diff b31fac9 --binary > "$ADALN_EVIDENCE/tracked.patch" +git ls-files --others --exclude-standard > "$ADALN_EVIDENCE/untracked-files.txt" +cp tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py "$ADALN_EVIDENCE/" +cp docs/reviews/adaln-shared-b31fac9*.md "$ADALN_EVIDENCE/" +sha256sum csrc/cuda/adaln_modulation.cu \ + rl_engine/kernels/ops/{pytorch,triton,cuda}/norm/adaln_modulation.py \ + rl_engine/kernels/gtest/{op_checks,operator_specs,operator_inputs}.py \ + rl_engine/kernels/registry.py tests/test_adaln*.py setup.py csrc/ops.cpp \ + > "$ADALN_EVIDENCE/sources.sha256" +nvidia-smi > "$ADALN_EVIDENCE/nvidia-smi.txt" +nvcc --version > "$ADALN_EVIDENCE/nvcc.txt" +.venv/bin/python -m pip freeze > "$ADALN_EVIDENCE/packages.txt" + +# 建议使用与项目兼容且已固定的 CUDA toolkit；不在这里安装/下载依赖。 +FORCE_CUDA=1 KERNEL_ALIGN_USE_FAST_MATH=0 KERNEL_ALIGN_FORCE_SM90=1 \ +TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=4 \ + .venv/bin/python setup.py build_ext --inplace --force \ + 2>&1 | tee "$ADALN_EVIDENCE/build.log" +sha256sum rl_engine/_C*.so > "$ADALN_EVIDENCE/binary.sha256" + +export RL_KERNEL_REQUIRE_EXT=1 +export PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 +export RLK_ADALN_REAL_SHAPES=1 +.venv/bin/python - <<'PY' 2>&1 | tee "$ADALN_EVIDENCE/preflight.log" +import sys, torch, triton +from rl_engine.kernels.registry import kernel_registry +assert torch.cuda.is_available() +assert torch.cuda.get_device_capability()[0] == 9 +assert "H100" in torch.cuda.get_device_name(), torch.cuda.get_device_name() +print(sys.version, torch.__version__, torch.version.cuda, triton.__version__) +print(torch.cuda.get_device_name(), torch.cuda.get_device_capability()) +_, trace = kernel_registry.get_adaln_modulation_op("cuda", hidden=3072) +assert trace["selected_backend"] == "CUDA_ADALN_MODULATION" and not trace["fallback"], trace +print(trace) +PY +.venv/bin/python -m pytest tests/test_extension_smoke.py \ + tests/test_adaln_modulation.py tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py \ + -q -rs -p no:cacheprovider 2>&1 | tee "$ADALN_EVIDENCE/tests.log" + +# 使用与 CUDA_HOME 相同且兼容的 sanitizer；检查退出码及完整 hazard/error summary。 +"${CUDA_HOME:?set CUDA_HOME}/bin/compute-sanitizer" --tool racecheck \ + --racecheck-report analysis --error-exitcode 7 \ + .venv/bin/python -m pytest tests/test_adaln_modulation.py \ + tests/test_adaln_gtest.py -q -rs -p no:cacheprovider \ + -k 'cuda and (bit_equality or launch_geometry or sample_bytes or backward_keeps)' \ + 2>&1 | tee "$ADALN_EVIDENCE/racecheck.log" + +for backend in cuda triton; do + for dtype in bf16 fp32; do + .venv/bin/python scripts/check_operator.py --op adaln_modulation \ + --candidate "$backend" --device cuda --dtype "$dtype" \ + --batch 3 --seq 5 --normalized-dim 3072 --check-grad --json \ + > "$ADALN_EVIDENCE/gtest-$backend-$dtype.json" + done +done +for dtype in bf16 fp32; do + .venv/bin/python benchmarks/benchmark_adaln_modulation.py \ + --real --dtype "$dtype" --warmup 3 --repeat 10 \ + 2>&1 | tee "$ADALN_EVIDENCE/benchmark-$dtype.log" +done +sha256sum -c "$ADALN_EVIDENCE/sources.sha256" +sha256sum -c "$ADALN_EVIDENCE/binary.sha256" +``` + +验收清单： + +- 保存 worktree patch、所有相关 untracked 源文件、源码 hash、binary hash、完整 build/test logs。 +- 源码/二进制在整个验收期间不变；确认 fast math disabled、SM90 构建、H100 实际执行。 +- focused + gtest + extension smoke 全部通过，0 skip/xfail；三 S、两 dtype、两个 GPU backend 不能漏。 +- BF16 forward 两个 literal、backward cancellation、FP32 single-cast VJP differential 全部通过。 +- H3072 B=1→B=3 的每个位置、S3/5、非连续视图、y/gate/dX/dMod bytes 全部通过。 +- CUDA 三 threads 与 Triton 六 warps/tile 配置的 geometry bytes 全部通过。 +- real-image independent CPU accuracy 用共享 forward/gradient tolerance；fixed-reference 与 backend bytes 另验。 +- gold 的两项 gradient dtype 为 FP32；不能只见 `report.passed` 或 dtype/finite 就宣称精度通过。 +- sanitizer exit 0，0 hazards/errors；有 race 报告即失败，即使数值断言通过。 +- native 没有 fallback/skip；benchmark 每个 dtype/S 都有 CUDA/Triton fwd/fwd+bwd 与 peak-memory 数据。 +- CLI JSON 是 operator debugging evidence；缺少受验证 provenance 时不能当整个 WS1 系统 EXIT。 + +GPU CI 与 script 04 现在也会运行新 canaries。CI 未启动，workflow 配置不是运行证据。 diff --git a/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py b/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py index 0fd5cbe68..ce827e340 100644 --- a/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py +++ b/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py @@ -34,8 +34,7 @@ def __init__(self, *, threads=256): 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") + hasattr(_C, name) for name in ("adaln_modulation_forward", "adaln_modulation_backward") ): raise RuntimeError("CUDA AdaLN extension symbols are unavailable") @@ -44,8 +43,10 @@ def __call__(self, x, modulation, *, eps=1e-6): def forward(self, x, modulation, *, eps=1e-6): _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 - ) + return _AdaLNCuda.apply(x.contiguous(), modulation.contiguous(), float(eps), self.threads) diff --git a/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py b/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py index 6d4fe5463..3f28d1224 100644 --- a/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py +++ b/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py @@ -5,6 +5,7 @@ from __future__ import annotations import math +import struct import torch @@ -46,18 +47,26 @@ def _validate(x: torch.Tensor, modulation: torch.Tensor, eps: float) -> None: 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") + message = "eps must remain finite FP32 without positive underflow" + try: + eps32 = struct.unpack("f", struct.pack("f", eps))[0] + except OverflowError: + raise ValueError(message) from None + if not math.isfinite(eps32) or (eps > 0 and eps32 == 0): + raise ValueError(message) class _AdaLNReference(torch.autograd.Function): @staticmethod def forward(ctx, x: torch.Tensor, modulation: torch.Tensor, eps: float): - shift, scale, gate = modulation.float().chunk(3, dim=-1) + 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 - return y, gate[:, None, :].to(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): @@ -74,12 +83,12 @@ def backward(ctx, grad_y: torch.Tensor, grad_gate: torch.Tensor): # S is folded in ascending logical-token order for every sample. dm = torch.cat( [ - torch.stack([reduce_rows_fp32(rows) for rows in dy]), - torch.stack([reduce_rows_fp32(rows) for rows in dy * norm]), - grad_gate.float().squeeze(1), + 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, - ).to(ctx.mod_dtype) + ) return dx, dm, None diff --git a/rl_engine/kernels/ops/triton/norm/adaln_modulation.py b/rl_engine/kernels/ops/triton/norm/adaln_modulation.py index 9843a238e..fdb66dc80 100644 --- a/rl_engine/kernels/ops/triton/norm/adaln_modulation.py +++ b/rl_engine/kernels/ops/triton/norm/adaln_modulation.py @@ -19,7 +19,8 @@ def _sum_hidden(values, BLOCK: tl.constexpr, LOG: tl.constexpr): half = 1 << level upper = tl.gather(values, tl.minimum(col + half, BLOCK - 1), 0) values = tl.where(col < half, values + upper, values) - return tl.sum(tl.where(col == 0, values, 0.0), 0) + # 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 @@ -97,7 +98,7 @@ def forward(ctx, x, modulation, eps, num_warps, reduction_tile): _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, num_warps=num_warps, + enable_fp_fusion=False, enable_reflect_ftz=False, num_warps=num_warps, ) ctx.save_for_backward(x, modulation) ctx.eps = eps @@ -117,6 +118,7 @@ def backward(ctx, grad_y, grad_gate): 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))]( diff --git a/rl_engine/validation/operators/op_checks.py b/rl_engine/validation/operators/op_checks.py index 5f3aded34..7f46f4174 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) @@ -217,7 +219,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 @@ -364,6 +368,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" ), @@ -394,12 +399,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") @@ -520,8 +529,7 @@ def _resolve_tolerance( # They only mirror forward accuracy thresholds; never apply them to grads. if judgment != "forward_accuracy": raise ContractResolveError( - f"legacy accuracy contracts only support judgment='forward_accuracy'; " - f"got {judgment!r}" + f"legacy accuracy contracts only support judgment='forward_accuracy'; got {judgment!r}" ) dtype_name = normalize_dtype_name(dtype) if arch_key is not None: @@ -550,6 +558,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( @@ -582,6 +591,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 +604,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..6eb03d9f2 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -11,7 +11,10 @@ import torch from rl_engine.validation.operators.op_checks import CandidateSpec, OperatorCase -from rl_engine.validation.operators.operator_inputs import make_operator_inputs, operator_shape_name +from rl_engine.validation.operators.operator_inputs import ( + make_operator_inputs, + operator_shape_name, +) @dataclass(frozen=True) @@ -22,6 +25,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,55 +37,67 @@ 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.kernels.ops.pytorch.norm.adaln_modulation.NativeAdaLNModulationOp" + ), + "triton": "rl_engine.kernels.ops.triton.norm.adaln_modulation.TritonAdaLNModulationOp", + "cuda": "rl_engine.kernels.ops.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", - gold_path="rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp", + gold_path="rl_engine.reference.norm.rms_norm.NativeRMSNormOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp", - "triton": "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp", - "cuda": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", - "cuda-sm90": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", - "ascend": "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp", + "pytorch": "rl_engine.reference.norm.rms_norm.NativeRMSNormOp", + "triton": "rl_engine.backends.shared.triton.norm.rmsnorm.RMSNormTritonOp", + "cuda": "rl_engine.backends.cuda.norm.rmsnorm.RMSNormCudaOp", + "cuda-sm90": "rl_engine.backends.cuda.norm.rmsnorm.RMSNormCudaOp", + "ascend": "rl_engine.backends.ascend.norm.rmsnorm.RMSNormAscendOp", }, grad_input_names=("x", "weight"), ), "qk_norm": OperatorSpec( name="qk_norm", op_class="reduction", - gold_path="rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp", + gold_path="rl_engine.reference.norm.rms_norm.NativeRMSNormOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp", - "triton": "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp", - "cuda": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", - "cuda-sm90": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", - "ascend": "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp", + "pytorch": "rl_engine.reference.norm.rms_norm.NativeRMSNormOp", + "triton": "rl_engine.backends.shared.triton.norm.rmsnorm.RMSNormTritonOp", + "cuda": "rl_engine.backends.cuda.norm.rmsnorm.RMSNormCudaOp", + "cuda-sm90": "rl_engine.backends.cuda.norm.rmsnorm.RMSNormCudaOp", + "ascend": "rl_engine.backends.ascend.norm.rmsnorm.RMSNormAscendOp", }, grad_input_names=("x", "weight"), ), "attention": OperatorSpec( name="attention", op_class="attention", - gold_path="rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp", + gold_path="rl_engine.reference.attention.standard_attn.NativeAttentionOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp", + "pytorch": "rl_engine.reference.attention.standard_attn.NativeAttentionOp", "triton": ( - ( - "rl_engine.kernels.ops.triton.attention.standard_attn.TritonBatchInvariantAt" - "tentionOp" - ) + "rl_engine.backends.shared.triton.attention.standard_attn.TritonBatchInvariantAt" + "tentionOp" ), "cuda": ( - "rl_engine.kernels.ops.cuda.attention.deterministic_attn.DeterministicAttentionOp" + "rl_engine.backends.cuda.attention.deterministic_attn.DeterministicAttentionOp" ), "ascend": ( - ( - "rl_engine.kernels.ops.ascend.attention.deterministic_attn.DeterministicAtte" - "ntionAscendOp" - ) + "rl_engine.backends.ascend.attention.deterministic_attn.DeterministicAtte" + "ntionAscendOp" ), }, grad_input_names=("q", "k", "v"), @@ -96,13 +113,11 @@ def _load_object(path: str) -> Any: candidate_paths={ "pytorch": "rl_engine.validation.operators.operator_specs.GtestPrefixSharedAttentionOp", "cuda": ( - "rl_engine.kernels.ops.cuda.attention.prefix_shared_attn.PrefixSharedAttentionOp" + "rl_engine.backends.cuda.attention.prefix_shared_attn.PrefixSharedAttentionOp" ), "ascend": ( - ( - "rl_engine.kernels.ops.ascend.attention.prefix_shared_attn.PrefixSharedAtten" - "tionAscendOp" - ) + "rl_engine.backends.ascend.attention.prefix_shared_attn.PrefixSharedAtten" + "tionAscendOp" ), }, ), @@ -110,18 +125,12 @@ def _load_object(path: str) -> Any: name="cp_attention", op_class="attention", gold_path=( - ( - "rl_engine.kernels.ops.pytorch.attention.cp_attention.DeterministicCPAttenti" - "onReferenceOp" - ) + "rl_engine.reference.attention.cp_attention.DeterministicCPAttentionReferenceOp" ), gold_method="forward_fp32", candidate_paths={ "pytorch": ( - ( - "rl_engine.kernels.ops.pytorch.attention.cp_attention.DeterministicCPAttenti" - "onReferenceOp" - ) + "rl_engine.reference.attention.cp_attention.DeterministicCPAttentionReferenceOp" ), }, grad_input_names=("q", "k", "v"), @@ -129,54 +138,54 @@ def _load_object(path: str) -> Any: "logp": OperatorSpec( name="logp", op_class="logprob", - gold_path="rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp", + gold_path="rl_engine.reference.logprob.logp.NativeLogpOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp", - "triton": "rl_engine.kernels.ops.triton.loss.logp.TritonLogpOp", - "cuda": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp", - "cuda-generic": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp", - "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpSM90Op", - "ascend": "rl_engine.kernels.ops.ascend.loss.logp.FusedLogpAscendOp", + "pytorch": "rl_engine.reference.logprob.logp.NativeLogpOp", + "triton": "rl_engine.backends.shared.triton.logprob.logp.TritonLogpOp", + "cuda": "rl_engine.backends.cuda.logprob.logp.FusedLogpGenericOp", + "cuda-generic": "rl_engine.backends.cuda.logprob.logp.FusedLogpGenericOp", + "cuda-sm90": "rl_engine.backends.cuda.logprob.logp.FusedLogpSM90Op", + "ascend": "rl_engine.backends.ascend.logprob.logp.FusedLogpAscendOp", }, grad_input_names=("logits",), ), "linear_logp": OperatorSpec( name="linear_logp", op_class="logprob", - gold_path="rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp", + gold_path="rl_engine.reference.logprob.linear_logp.NativeLinearLogpOp", gold_method="apply", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp", - "triton": "rl_engine.kernels.ops.triton.loss.linear_logp.TritonLinearLogpOp", - "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.linear_logp.FusedLinearLogpSM90Op", - "ascend": "rl_engine.kernels.ops.ascend.loss.linear_logp.FusedLinearLogpAscendOp", + "pytorch": "rl_engine.reference.logprob.linear_logp.NativeLinearLogpOp", + "triton": "rl_engine.backends.shared.triton.logprob.linear_logp.TritonLinearLogpOp", + "cuda-sm90": "rl_engine.backends.cuda.logprob.linear_logp.FusedLinearLogpSM90Op", + "ascend": "rl_engine.backends.ascend.logprob.linear_logp.FusedLinearLogpAscendOp", }, grad_input_names=("hidden", "lm_head_weight"), ), "embedding": OperatorSpec( name="embedding", op_class="elementwise", - gold_path="rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", + gold_path="rl_engine.reference.embedding.embedding.NativeEmbeddingOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", - "triton": "rl_engine.kernels.ops.triton.linear.embedding.TritonEmbeddingOp", - "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp", - "ascend": "rl_engine.kernels.ops.ascend.linear.embedding.AscendEmbeddingOp", + "pytorch": "rl_engine.reference.embedding.embedding.NativeEmbeddingOp", + "triton": "rl_engine.backends.shared.triton.embedding.embedding.TritonEmbeddingOp", + "cuda-sm90": "rl_engine.backends.cuda.embedding.embedding.SM90EmbeddingOp", + "ascend": "rl_engine.backends.ascend.embedding.embedding.AscendEmbeddingOp", }, grad_input_names=("weight",), ), "lm_head": OperatorSpec( name="lm_head", op_class="reduction", - gold_path="rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", + gold_path="rl_engine.reference.embedding.lm_head.NativeLMHeadOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", - "triton": "rl_engine.kernels.ops.triton.linear.lm_head.TritonLMHeadOp", - "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp", - "ascend": "rl_engine.kernels.ops.ascend.linear.lm_head.AscendLMHeadOp", + "pytorch": "rl_engine.reference.embedding.lm_head.NativeLMHeadOp", + "triton": "rl_engine.backends.shared.triton.embedding.lm_head.TritonLMHeadOp", + "cuda-sm90": "rl_engine.backends.cuda.embedding.lm_head.SM90LMHeadOp", + "ascend": "rl_engine.backends.ascend.embedding.lm_head.AscendLMHeadOp", }, grad_input_names=("hidden", "weight"), ), @@ -187,77 +196,73 @@ def _load_object(path: str) -> Any: # the accuracy gold must be the same leaf-space tree, not the # single-rounding torch.matmul (which fails structurally at # near-cancellation outputs on random inputs). - gold_path="rl_engine.kernels.ops.pytorch.matmul.det_gemm.DetGemmTreeReferenceOp", + gold_path="rl_engine.reference.gemm.det_gemm.DetGemmTreeReferenceOp", gold_method="__call__", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.matmul.det_gemm.NativeGemmOp", - "cuda": "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp", - "triton": "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp", - "ascend": "rl_engine.kernels.ops.ascend.matmul.det_gemm.DetGemmAscendOp", + "pytorch": "rl_engine.reference.gemm.det_gemm.NativeGemmOp", + "cuda": "rl_engine.backends.cuda.gemm.det_gemm.DetGemmOp", + "triton": "rl_engine.backends.shared.triton.gemm.det_gemm.TritonDetGemmOp", + "ascend": "rl_engine.backends.ascend.gemm.det_gemm.DetGemmAscendOp", }, grad_input_names=("a", "b"), ), "rope": OperatorSpec( name="rope", op_class="elementwise", - gold_path="rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp", + gold_path="rl_engine.reference.rope.rope.NativeRoPEOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp", - "triton": "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp", - "cuda-sm90": "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op", - "ascend": "rl_engine.kernels.ops.ascend.rotary_embedding.rope.RoPEAscendOp", + "pytorch": "rl_engine.reference.rope.rope.NativeRoPEOp", + "triton": "rl_engine.backends.shared.triton.rope.rope.TritonRoPEOp", + "cuda-sm90": "rl_engine.backends.cuda.rope.rope.RoPESM90Op", + "ascend": "rl_engine.backends.ascend.rope.rope.RoPEAscendOp", }, grad_input_names=("x",), ), "silu": OperatorSpec( name="silu", op_class="elementwise", - gold_path="rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp", + gold_path="rl_engine.reference.activation.swiglu.NativeSiLUOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp", - "triton": "rl_engine.kernels.ops.triton.activation.swiglu.TritonSiLUOp", - "cuda": "rl_engine.kernels.ops.cuda.activation.swiglu.SiLUCudaOp", - "ascend": "rl_engine.kernels.ops.ascend.activation.silu.SiLUAscendOp", + "pytorch": "rl_engine.reference.activation.swiglu.NativeSiLUOp", + "triton": "rl_engine.backends.shared.triton.activation.swiglu.TritonSiLUOp", + "cuda": "rl_engine.backends.cuda.activation.swiglu.SiLUCudaOp", + "ascend": "rl_engine.backends.ascend.activation.silu.SiLUAscendOp", }, grad_input_names=("x",), ), "swiglu": OperatorSpec( name="swiglu", op_class="elementwise", - gold_path="rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp", + gold_path="rl_engine.reference.activation.swiglu.NativeSwiGLUOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp", - "triton": "rl_engine.kernels.ops.triton.activation.swiglu.TritonSwiGLUOp", - "cuda": "rl_engine.kernels.ops.cuda.activation.swiglu.SwiGLUCudaOp", - "ascend": "rl_engine.kernels.ops.ascend.activation.swiglu.SwiGLUAscendOp", + "pytorch": "rl_engine.reference.activation.swiglu.NativeSwiGLUOp", + "triton": "rl_engine.backends.shared.triton.activation.swiglu.TritonSwiGLUOp", + "cuda": "rl_engine.backends.cuda.activation.swiglu.SwiGLUCudaOp", + "ascend": "rl_engine.backends.ascend.activation.swiglu.SwiGLUAscendOp", }, grad_input_names=("gate", "up"), ), "batch_invariant_logp": OperatorSpec( name="batch_invariant_logp", op_class="logprob", - gold_path=( - "rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp.NativeBatchInvarian" "tLogpOp" - ), + gold_path=("rl_engine.reference.logprob.batch_invariant_logp.NativeBatchInvariantLogpOp"), gold_method="apply", candidate_paths={ "pytorch": ( - "rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp.NativeBatchInvarian" - "tLogpOp" + "rl_engine.reference.logprob.batch_invariant_logp.NativeBatchInvariantLogpOp" ), "triton": ( - "rl_engine.kernels.ops.triton.loss.batch_invariant_logp.TritonBatchInvariant" + "rl_engine.backends.shared.triton.logprob.batch_invariant_logp.TritonBatchInvariant" "LogpOp" ), "cuda-sm90": ( - "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90" "Op" + "rl_engine.backends.cuda.logprob.batch_invariant_logp.BatchInvariantLogpSM90Op" ), "ascend": ( - "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAs" - "cendOp" + "rl_engine.backends.ascend.logprob.batch_invariant_logp.BatchInvariantLogpAscendOp" ), }, grad_input_names=("logits",), @@ -275,6 +280,26 @@ def _load_object(path: str) -> Any: } +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 custom backward helper is used. + """ + + 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) + if modulate_index is None: + shift, scale, gate = mod32[:, None, :].chunk(3, dim=-1) + else: + batch = x.shape[0] + rows = torch.arange(batch)[:, None] + modulate_index.cpu().long() * batch + shift, scale, gate = mod32[rows].chunk(3, dim=-1) + norm = torch.nn.functional.layer_norm(x32, (x.shape[-1],), eps=eps) + return norm * (1 + scale) + shift, gate + + class GtestPackOp: """gtest view of NativePackOp: compare the packed rows, not cu_seqlens.""" @@ -306,7 +331,9 @@ class GtestPrefixSharedAttentionOp: op_class = "attention" def __init__(self) -> None: - from rl_engine.reference.attention.standard_attn import NativeAttentionOp + from rl_engine.reference.attention.standard_attn import ( + NativeAttentionOp, + ) self._op = NativeAttentionOp() @@ -345,6 +372,8 @@ 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, ) diff --git a/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh b/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh index 9ae064d56..684fb31d3 100755 --- a/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh +++ b/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh @@ -21,5 +21,6 @@ assert trace["selected_backend"] == OpBackend.CUDA_ADALN_MODULATION.name, trace assert trace["fallback"] is False, trace print(trace) PY -.venv/bin/python -m pytest tests/test_extension_smoke.py tests/test_adaln_modulation.py -q -rs -p no:cacheprovider +.venv/bin/python -m pytest tests/test_extension_smoke.py tests/test_adaln_modulation.py \ + tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py -q -rs -p no:cacheprovider bash scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh diff --git a/scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh b/scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh index e7827fdf4..9018ea9cc 100644 --- a/scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh +++ b/scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh @@ -20,4 +20,5 @@ exec "$CUDA_HOME/bin/compute-sanitizer" --tool racecheck \ --racecheck-report analysis --error-exitcode 7 \ .venv/bin/python -m pytest -q -rs -p no:cacheprovider \ tests/test_adaln_modulation.py::test_cuda_adaln_bit_equality_harness \ - tests/test_adaln_modulation.py::test_adaln_launch_geometry_preserves_bytes -k cuda "$@" + tests/test_adaln_modulation.py::test_adaln_launch_geometry_preserves_bytes \ + tests/test_adaln_gtest.py::test_adaln_h3072_relocation_with_all_launch_configurations -k cuda "$@" diff --git a/tests/test_adaln_gtest.py b/tests/test_adaln_gtest.py new file mode 100644 index 000000000..bab77a8be --- /dev/null +++ b/tests/test_adaln_gtest.py @@ -0,0 +1,665 @@ +# 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.kernels.gtest import run_operator_suite +from rl_engine.kernels.gtest.operator_inputs import operator_shape_name +from rl_engine.kernels.gtest.operator_specs import ( + GtestAdaLNReference, + make_candidate, + make_operator_case, +) +from rl_engine.kernels.gtest.tolerance import load_contract, resolve_tolerance +from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp +from rl_engine.kernels.registry import OpBackend, kernel_registry + + +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", "cuda"]) +def shared_op(request): + backend = request.param + if backend != "pytorch" and not torch.cuda.is_available(): + pytest.skip("CUDA required") + args = _args() + args.candidate = backend + try: + op = make_candidate(args).fn + 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]])) + + +@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]) +def test_adaln_gtest_rejects_fp32_outputs_for_bf16_inputs(check_grad): + args = _args() + case = make_operator_case(args, torch.bfloat16, torch.device("cpu")) + good = make_candidate(args) + + def wrong_dtype(**inputs): + return tuple(output.float() for output in good.fn(**inputs)) + + 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 all(not check.passed and "dtype mismatch" in check.message for check in checks) + + +@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, 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/test_adaln_modulation.py b/tests/test_adaln_modulation.py index 22f39662f..c3658520c 100644 --- a/tests/test_adaln_modulation.py +++ b/tests/test_adaln_modulation.py @@ -3,6 +3,8 @@ import pytest import torch +from rl_engine.kernels.gtest.operator_specs import GtestAdaLNReference +from rl_engine.kernels.gtest.tolerance import load_contract, resolve_tolerance from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp from rl_engine.kernels.registry import OpBackend, kernel_registry @@ -14,6 +16,29 @@ def _assert_bytes_equal(actual, expected): 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. @@ -24,9 +49,7 @@ def test_adaln_modulation_forward_and_shared_backward(): (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]]) - ) + 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() @@ -116,9 +139,9 @@ def test_adaln_modulation_matches_independent_layer_norm_autograd(): 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() - torch.testing.assert_close(y, y_ref, atol=1e-6, rtol=1e-6) - torch.testing.assert_close(x.grad, x_ref.grad, atol=1e-6, rtol=1e-6) - torch.testing.assert_close(modulation.grad, modulation_ref.grad, atol=1e-6, rtol=1e-6) + _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(): @@ -201,12 +224,11 @@ def test_triton_adaln_modulation_forward_backward_and_batch_position(dtype): def test_triton_irregular_hidden_and_bf16_tie_rounding(): - from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp - try: torch.cuda.init() except RuntimeError: pytest.skip("CUDA required") + from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp for hidden, seq, dtype, seed in ( (7, 3, torch.float32, 386), (3072, 64, torch.bfloat16, 391), @@ -228,13 +250,12 @@ def test_triton_irregular_hidden_and_bf16_tie_rounding(): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_adaln_gpu_padding_and_batch_size_do_not_change_valid_bytes(dtype): - from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp - from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp - try: torch.cuda.init() except RuntimeError: pytest.skip("CUDA required") + from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from rl_engine.kernels.ops.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) @@ -347,6 +368,9 @@ def test_adaln_real_image_shapes_forward_backward(seq, dtype): _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(): @@ -372,7 +396,8 @@ def test_adaln_launch_geometry_preserves_bytes(backend, dtype): else: ops = [ TritonAdaLNModulationOp(num_warps=w, reduction_tile=t) - for w in (4, 8) for t in (64, 128, 256) + for w in (4, 8) + for t in (64, 128, 256) ] torch.manual_seed(393) x = torch.randn(2, 5, 3072, dtype=dtype) @@ -389,8 +414,10 @@ def test_adaln_launch_geometry_preserves_bytes(backend, dtype): 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), + (y, expected_y), + (gate, expected_gate), + (gpu_x.grad, cpu_x.grad), + (gpu_modulation.grad, cpu_modulation.grad), ): _assert_bytes_equal(actual, expected) @@ -414,21 +441,17 @@ def test_adaln_indexed_matches_independent_autograd(device, dtype, broadcast): y, gate = op(x, modulation, modulate_index=index) torch.autograd.backward((y, gate), (dy, dg)) - ref_x = x.detach().clone().requires_grad_() - ref_modulation = modulation.detach().clone().requires_grad_() - # Independent row indexing, not the implementation's split-and-where path. - selected = ref_modulation.float()[torch.arange(2, device=device)[:, None] + index * 2] - shift, scale, ref_gate = selected.chunk(3, dim=-1) - ref_y = ( - torch.nn.functional.layer_norm(ref_x.float(), (7,), eps=1e-6) * (1 + scale) + shift - ).to(dtype) - ref_gate = ref_gate.to(dtype) - torch.autograd.backward((ref_y, ref_gate), (dy, dg)) - for actual, expected in ( - (y, ref_y), (x.grad, ref_x.grad), (modulation.grad, ref_modulation.grad) - ): - torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-6) - _assert_bytes_equal(gate, ref_gate) + 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 assert trace["fallback"] is True assert trace["fallback_reason"] == "modulate_index_requires_select01" @@ -446,9 +469,7 @@ def test_adaln_indexed_matches_independent_autograd(device, dtype, broadcast): ) 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 - ) + NativeAdaLNModulationOp()(torch.ones(1, 2, 7), torch.ones(2, 21), modulate_index=index) @pytest.mark.parametrize("backend", ["pytorch", "triton", "cuda"]) @@ -457,20 +478,21 @@ def test_adaln_indexed_rejects_invalid_routing(index, error): def test_adaln_shared_matches_independent_autograd(backend, hidden, dtype): if backend != "pytorch" and not torch.cuda.is_available(): pytest.skip("CUDA required") - from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp - from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp - - ops = {"pytorch": NativeAdaLNModulationOp, "triton": TritonAdaLNModulationOp, - "cuda": CudaAdaLNModulationOp} if backend == "cuda": + from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + try: - op = ops[backend]() + 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.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + + op = TritonAdaLNModulationOp() else: - op = ops[backend]() + 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) @@ -480,21 +502,4 @@ def test_adaln_shared_matches_independent_autograd(backend, hidden, dtype): y, gate = op(x, modulation) torch.autograd.backward((y, gate), (dy, dg)) - ref_x = x.detach().cpu().requires_grad_() - ref_modulation = modulation.detach().cpu().requires_grad_() - shift, scale, ref_gate = ref_modulation.float().chunk(3, dim=-1) - ref_y = ( - torch.nn.functional.layer_norm(ref_x.float(), (hidden,), eps=1e-6) - * (1 + scale[:, None, :]) + shift[:, None, :] - ).to(dtype) - ref_gate = ref_gate[:, None, :].to(dtype) - torch.autograd.backward((ref_y, ref_gate), (dy.cpu(), dg.cpu())) - for actual, expected in ( - (y, ref_y), (x.grad, ref_x.grad), (modulation.grad, ref_modulation.grad) - ): - torch.testing.assert_close( - actual.cpu(), expected, - atol=2e-5 if dtype == torch.float32 else 0.02, - rtol=2e-5 if dtype == torch.float32 else 0.016, - ) - _assert_bytes_equal(gate, ref_gate) + _assert_independent_shared(x, modulation, y, gate, dy, dg) diff --git a/tests/test_adaln_shared_bounds.py b/tests/test_adaln_shared_bounds.py new file mode 100644 index 000000000..00e5e2b21 --- /dev/null +++ b/tests/test_adaln_shared_bounds.py @@ -0,0 +1,74 @@ +# 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.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +from rl_engine.kernels.registry import OpBackend, kernel_registry + + +@pytest.fixture +def cuda_op(): + try: + return CudaAdaLNModulationOp() + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + pytest.skip("compiled CUDA AdaLN extension required") + + +def test_cuda_shared_rejects_batch_above_grid_limit(cuda_op): + device = "cuda" if torch.cuda.is_available() else "meta" + x = torch.zeros(65536, 1, 1, device=device) + m = torch.zeros(65536, 3, device=device) + 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("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"] + + +@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.zeros(65535, 1, 1, device="cuda", dtype=dtype, requires_grad=True) + m = torch.zeros(65535, 3, device="cuda", dtype=dtype, requires_grad=True) + y, gate = cuda_op(x, m) + dx, dm = torch.autograd.grad(gate, (x, m), torch.ones_like(gate)) + assert torch.equal(y, torch.zeros_like(y)) + assert torch.equal(gate, torch.zeros_like(gate)) + assert torch.equal(dx, torch.zeros_like(dx)) + expected = torch.zeros_like(m) + expected[:, 2] = 1 + assert torch.equal(dm, expected) From 46774190e2f920b0c2e16f1bfc0adbf8966d623f Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Sat, 10 Oct 2026 01:44:40 +0800 Subject: [PATCH 07/37] feat(test): enhance AdaLN tests with CUDA thread configurations and dtype support Signed-off-by: hongleng-github --- docs/reviews/adaln-shared-b31fac9.md | 86 +++++++++++++++++- tests/test_adaln_gtest.py | 129 +++++++++++++++++++++------ tests/test_adaln_shared_bounds.py | 41 ++++++--- 3 files changed, 210 insertions(+), 46 deletions(-) diff --git a/docs/reviews/adaln-shared-b31fac9.md b/docs/reviews/adaln-shared-b31fac9.md index d9ff01587..98dce9edd 100644 --- a/docs/reviews/adaln-shared-b31fac9.md +++ b/docs/reviews/adaln-shared-b31fac9.md @@ -1,10 +1,88 @@ -# AdaLN shared：工作区审查与 TDD 证据（2026-10-09） +# AdaLN shared：工作区审查与 TDD 证据（更新于 2026-10-10） + +## 当前 native SM86 短验收（2026-10-10） + +用户授权先本地提交、重新构建 native，再执行全部短测试。输入提交为 +`f21007cccc216d214ec2589cc175bc288a5363e0`，固定比较点仍为 `b31fac9`。 +该提交包含此前全部未提交改动、两个未跟踪测试和两个审查文档，未推送。 +本节取代下方构建前记录中的 native UNPROVEN 结论；仅对明确列出的本机反例升级资格。 + +全量构建使用独立 objects/lib 目录和 `build_ext --inplace --force`，旧 `.so` 已隔离。 +环境固定为 Python 3.12.12、PyTorch 2.12.1+cu126、CUDA toolkit 12.6.85、GCC 11.4、 +RTX3050Ti/SM86、MAX_JOBS=2、fast math=0，未安装或下载依赖。 +构建成功，实际 import 路径、构建目录产物和安装产物 hash 一致；构建期间源码不变。 +[构建参数及工具 hash](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/build-manifest.json)、 +[全部源码 hash](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/source.sha256)、 +[源码快照](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/source.tar.gz)、 +[完整构建日志](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/build.log)、 +[实际编译参数](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/objects/build.ninja)、 +[import 验证](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/import-verification.log)。 + +当前 `.so`：`rl_engine/_C.cpython-312-x86_64-linux-gnu.so`；SHA256： +`94b86feb52cb621794598b9ba6361ae3d521a370b9e6c5df74942112f2f35e61`。 +旧二进制 `d7b91b28...` 不参与本次运行，也不算当前证据。 + +先运行提交中现有用例：**89 passed，182 deselected，5.38s**。 +检查后补齐 B65536 的 BF16 和全部 threads，强化 B65535 的非零 y backward，得到 +**98 passed，182 deselected，5.68s**。最终进一步将 CUDA fixture 扩展到三种 threads， +让每个 canary、两 stream/两个 triple 和合法 H 边界均交叉执行： +**216 passed，182 deselected，7.23s；零 skip、xfail、xpass、失败或 teardown 错误**。 +没有暴露需要修改 native 实现的错误；本轮只修改两个测试文件和本报告，未新增提交。 + +| 可能错误 / 要求 | 本次反例与必须成立的断言 | 当前结论 | +| --- | --- | --- | +| 提前 BF16 舍入 norm/affine、保存量、partial 或 token accumulator | 早 cast 残差、1799/32768 VJP、H7/3072 BF16 对提升后 FP32 单次最终 cast、两 halfway tie；比较 y/gate/dX/dMod bytes；各自遍历 threads128/256/512 | 指定 native canary **PROVEN** | +| hidden mean/variance 或两个 backward sum 改树，token fold 被重关联 | `[2^24,1,-2^24,1]`、variance-only 向量、正负 dy、S4 cancellation；独立手算 literal 的输出/VJP bytes；三 threads | 指定 native 归约反例 **PROVEN** | +| signed zero 被额外 +0 擦除、合法 eps 被 FTZ、gate/dgate NaN payload 被转换 | signed-zero bytes、eps1e-40/2^-149/FP32 max 的 y0/dX±rstd/dMod literal、BF16 原始 payload 透传；三 threads，适用的两 dtype | 指定 native 反例 **PROVEN**；一般 LN NaN payload 不在资格声明中 | +| B65535 的最后 batch 寻址或 backward grid.y 错误 | B65535,S1,H2,x=[-3,3],eps7；y=[-.75,.75]，dX=[7/128,-7/128]，dshift=[1,0]，dscale=[-.75,0]，dgate=[2,-3]；每个 batch 的 gate 使用 mod17 标签；两 dtype×三 threads共6例，检查所有行 | 合法最大 B native forward/backward **PROVEN** | +| B65536 被非法 launch | B65536,S1,H1；公开 op 必须 ValueError `B <= 65535`，两 dtype×三 threads共6例 | 公开 fail-closed **PROVEN**；负例在 Python preflight 拒绝，不冒充独立 C++ guard 运行证据 | +| H4095/4096 的补零或上限 off-by-one | B2,S3，两 dtype×三 threads；y/gate/dX/dMod 对独立 CPU FP32 math accuracy 和固定参考 bytes，共12例 | 指定合法 native H 边界 **PROVEN** | +| H4097 静默进入 native / 未声明 fallback | 两 dtype公开 op 明确拒绝；H4095/4096 registry 必须选 CUDA 且 fallback=False；4097必须明确 rejected reason及fallback=True | direct 拒绝及 registry policy **PROVEN**；未声称 H4097 fallback 数值已验收 | +| threads×H3072 batch position/padding 组合路径出错 | 128/256/512×两 dtype，B1/S3 baseline→B3每个位置、S3/5、strided x和6H view；y/gate/dX/dMod bytes、unused半区0、padding dX0 | 当前 native 三 geometry×topology **PROVEN** | +| 第二 triple 或两 stream 错 shift / VJP 串值 | image S5/text S3、两 triple×两 dtype×三 threads共12例；独立 FP32 CPU forward/backward、gate bytes、unused半区0、另一个stream梯度不变 | 指定两 stream/两个 triple **PROVEN** | +| 上游 stride、单个 trainable input、unused output 失效 | transposed dy、stride0 dg、x-only/modulation-only/gate-only；独立 math 或 literal bytes；三 threads×两 dtype | 指定 native autograd 用例 **PROVEN** | +| int32 地址溢出、eps 输入域或非有限传播未验证 | 巨大 meta shape明确拒绝，不分配实数据；非法eps公开拒绝；NaN/Inf x传播及独立gate VJP | 公开 preflight及指定 nonfinite native 反例 **PROVEN**；C++负例守卫未绕过Python单独测试，完整非有限域 **PARTIAL** | +| 用旧二进制 / 隐含 fallback 冒充 native | 加载路径和 .so SHA256与fresh build一致；运行前后所有生产源码及binary hash不变；可支持 H 的 registry必须 native/fallback=False；数值用例直接使用 CudaAdaLNModulationOp | 当前 SM86 源码/产物/运行闭环 **PROVEN** | +| H100、三真实图像形状、racecheck/memcheck、benchmark | SM90固定环境 fresh build；S4096/6889/6032；sanitizer零hazards/errors；完整时间/峰值内存 | **UNPROVEN**，本轮未启动 | + +最终证据：[pytest 完整日志](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/pytest.log)、 +[JUnit](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/junit.xml)、 +[逐阶段结果及零 skip/xfail 检查](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/results.json)、 +[实际选中用例](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/selected-tests.json)、 +[固定运行环境及 registry traces](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/manifest.json)。 +182 deselected 是未选择的 CPU/Triton、indexed、真实图像及非 native checker 用例，不是 skip。 +测试中的 H4097/H5000 registry fallback 是声明过的 unsupported-H policy；支持域的数值用例无 fallback。 + +最终测试 SHA256： + +- `tests/test_adaln_gtest.py`：`dbb499193632e23ccaf415d78a246a4ac1d3f3a340809cda8ad4d1e7717efa54`。 +- `tests/test_adaln_shared_bounds.py`：`48c4432a6456bb9b25e06ffbe2d9e7457cdff7e9f25b5208fcbe310e7cc98810`。 + +运行前相对于 build source manifest 只有两个测试及本报告改变，所有生产源码保持一致；完整 +[当前源码 hash](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/source-sha256-current.json)、 +[测试 patch](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/test-changes.patch) +和两个完整测试快照均已保存。运行期间源码/产物不变，报告在运行后更新。 +ruff、Black py312、isort black profile及 `git diff --check` 检查通过；没有重新编译或更改生产源码。 + +复跑短矩阵使用如下选择，并保留固定环境、import hash、registry和零skip/xfail校验； +这些额外校验的完整可审查脚本保存在证据目录 `runner.py`。 + +```bash +OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 RLK_ADALN_REAL_SHAPES=0 \ +RL_KERNEL_REQUIRE_EXT=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ + .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py \ + tests/test_adaln_modulation.py -vv -ra -p no:cacheprovider \ + -k 'cuda and not real_image and not indexed' +``` + +下一步仍是单独授权/由用户执行的 H100真实形状、sanitizer及benchmark； +本机 SM86 短验收不能升级为 SM90 或一般输入域的证明。 -## 剩余缺口与本次执行范围 +## 构建前 TDD 范围与记录（2026-10-09 历史快照） 用户已授权继续按 `/tdd` 改实现、补测试；沿用公开 op 调用、autograd、registry、gtest 接口，不要求新的接口确认。不提交/推送，不启动 H100 长测试、巨大输入或无法估计耗时的构建。 -下面是实施前待办；其最新完成状态和本次运行证据记录在下一节。 +下面保留构建前实施待办和中间运行记录；最新 native 状态以上一节为准。 本文件下方第一轮审查及 followup 文档都属于历史源码快照，旧运行表不能算当前版本证据。 | 缺口 | 修补动作 | 实施前结论 | @@ -18,7 +96,7 @@ | 当前源码/native binary闭环缺失 | 保存源/产物hash；仅在能够可靠估计快速完成时构建，否则提供构建与验收命令 | UNPROVEN | | H100真实图像、安全性、性能 | 保留三真实S、racecheck/memcheck、benchmark和pinned环境命令与清单，不自行启动 | UNPROVEN | -## 本次 TDD 完成状态与当前源码证据 +## 构建前 TDD 完成状态与源码证据（历史快照） 以下均通过公开 op、autograd、registry 或 gtest 观察，没有新增生产 helper mock。 已复用旧检查器、固定字节参考、独立 FP32 CPU LayerNorm/autograd 和模型六块切片用例。 diff --git a/tests/test_adaln_gtest.py b/tests/test_adaln_gtest.py index bab77a8be..0a5e54818 100644 --- a/tests/test_adaln_gtest.py +++ b/tests/test_adaln_gtest.py @@ -33,15 +33,25 @@ def _args(hidden=7): ) -@pytest.fixture(params=["pytorch", "triton", "cuda"]) +@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 = request.param + 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 @@ -60,7 +70,9 @@ def _bytes_equal(actual, expected): @pytest.mark.parametrize("nan_upstream", [False, True]) -def test_adaln_gate_preserves_bf16_nan_payload_and_identity_vjp(shared_op, nan_upstream): +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) @@ -119,7 +131,9 @@ def test_adaln_subnormal_eps_is_not_flushed_to_zero(shared_op, dtype, eps, rstd) 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) + 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. @@ -129,7 +143,9 @@ def test_adaln_bf16_forward_casts_only_at_output(shared_op, scale, positive): 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) + 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, @@ -157,7 +173,9 @@ def test_adaln_bf16_backward_keeps_partials_and_accumulators_fp32(shared_op): ) _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)) + _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]) @@ -166,10 +184,15 @@ def test_adaln_bf16_vjp_matches_fp32_with_one_final_cast(shared_op, hidden): 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) + 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_() + 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)) @@ -180,7 +203,9 @@ def test_adaln_bf16_vjp_matches_fp32_with_one_final_cast(shared_op, hidden): @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): +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) @@ -190,7 +215,9 @@ def test_adaln_bf16_boundary_rounds_halfway_values_to_even(shared_op, fraction, _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)) + _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]) @@ -201,7 +228,11 @@ def test_adaln_h3072_sample_bytes_survive_batch_positions_and_padding(shared_op, @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)] + [ + ("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): @@ -249,8 +280,12 @@ def observe(x, full, dy, dg, *, candidate=op): 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] + 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( @@ -270,7 +305,9 @@ def test_adaln_backward_accepts_transposed_and_broadcast_upstreams(shared_op, dt 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())) + 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_() @@ -284,7 +321,9 @@ def test_adaln_backward_accepts_transposed_and_broadcast_upstreams(shared_op, dt 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) + torch.testing.assert_close( + result.cpu().float(), expected, atol=tol.atol, rtol=tol.rtol + ) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -307,11 +346,15 @@ def test_adaln_single_trainable_input_and_unused_output(shared_op, dtype, traina 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)) + _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)) + _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): @@ -377,7 +420,9 @@ def test_adaln_modulation_gradient_uses_ascending_token_order(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) + 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]])) @@ -431,7 +476,9 @@ def test_adaln_shared_rejects_mixed_devices(shared_op): @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): +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) @@ -441,7 +488,9 @@ def test_adaln_shared_rejects_eps_overflow_and_positive_underflow(shared_op, dty @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): +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) @@ -492,7 +541,11 @@ def test_adaln_gtest_detects_wrong_output_and_backward(broken): 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 + 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 @@ -519,11 +572,16 @@ def wrong_dtype(**inputs): return tuple(output.float() for output in good.fn(**inputs)) report = run_operator_suite( - args.op, candidates=[replace(good, fn=wrong_dtype)], cases=[case], check_grad=check_grad + 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 all(not check.passed and "dtype mismatch" in check.message for check in checks) + assert all( + not check.passed and "dtype mismatch" in check.message for check in checks + ) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -535,8 +593,12 @@ def test_adaln_gate_only_and_separate_stream_chunks(shared_op, dtype, triple): # 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) + 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 ( @@ -562,7 +624,10 @@ def test_adaln_gate_only_and_separate_stream_chunks(shared_op, dtype, triple): 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 + load_contract(), + judgment="forward_accuracy", + op_class="reduction", + dtype=dtype, ) torch.testing.assert_close( y.detach().cpu().float(), @@ -594,7 +659,9 @@ def test_adaln_gate_only_and_separate_stream_chunks(shared_op, dtype, triple): @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): +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) @@ -606,7 +673,9 @@ def test_adaln_hidden_boundary_matches_independent_math_and_fixed_bytes(shared_o 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())) + 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"), @@ -616,7 +685,9 @@ def test_adaln_hidden_boundary_matches_independent_math_and_fixed_bytes(shared_o 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) + 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)) diff --git a/tests/test_adaln_shared_bounds.py b/tests/test_adaln_shared_bounds.py index 00e5e2b21..854559b15 100644 --- a/tests/test_adaln_shared_bounds.py +++ b/tests/test_adaln_shared_bounds.py @@ -12,19 +12,21 @@ @pytest.fixture -def cuda_op(): +def cuda_op(request): try: - return CudaAdaLNModulationOp() + 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_cuda_shared_rejects_batch_above_grid_limit(cuda_op): +@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) - m = torch.zeros(65536, 3, device=device) + 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) @@ -58,17 +60,30 @@ def test_cuda_shared_registry_respects_hidden_limit(cuda_op, hidden): assert "CUDA_ADALN_MODULATION: H > 4096" in trace["rejected_backends"] +@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.zeros(65535, 1, 1, device="cuda", dtype=dtype, requires_grad=True) - m = torch.zeros(65535, 3, device="cuda", dtype=dtype, requires_grad=True) - y, gate = cuda_op(x, m) - dx, dm = torch.autograd.grad(gate, (x, m), torch.ones_like(gate)) - assert torch.equal(y, torch.zeros_like(y)) - assert torch.equal(gate, torch.zeros_like(gate)) - assert torch.equal(dx, torch.zeros_like(dx)) + 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[:, 2] = 1 + expected[:, 0], expected[:, 2] = 1.0, -0.75 + expected[:, 4], expected[:, 5] = 2.0, -3.0 assert torch.equal(dm, expected) From 8e1e99ec3c5dfe5f8b7f82987f863230cbd1dc56 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Sat, 10 Oct 2026 02:32:58 +0800 Subject: [PATCH 08/37] feat(adaln): implement independent CPU affine VJP and enhance gradient sensitivity tests Signed-off-by: hongleng-github --- docs/reviews/adaln-shared-b31fac9.md | 42 +++++++++++++++++++ .../validation/operators/operator_specs.py | 29 ++++++++++++- tests/test_adaln_gtest.py | 13 ++++++ 3 files changed, 82 insertions(+), 2 deletions(-) diff --git a/docs/reviews/adaln-shared-b31fac9.md b/docs/reviews/adaln-shared-b31fac9.md index 98dce9edd..27929d32a 100644 --- a/docs/reviews/adaln-shared-b31fac9.md +++ b/docs/reviews/adaln-shared-b31fac9.md @@ -427,3 +427,45 @@ sha256sum -c "$ADALN_EVIDENCE/binary.sha256" - CLI JSON 是 operator debugging evidence；缺少受验证 provenance 时不能当整个 WS1 系统 EXIT。 GPU CI 与 script 04 现在也会运行新 canaries。CI 未启动，workflow 配置不是运行证据。 + + +## 本地完整验收流水线（2026-10-10，SM86） + +当前流水线状态：**failed**；H100/SM90仍为 **UNPROVEN**。 + +[完整状态与逐阶段命令](/tmp/rlk-adaln-local-acceptance-20261009T174946Z/status.json)、[固定环境及 binary hash](/tmp/rlk-adaln-local-acceptance-20261009T174946Z/manifest.json)、[源码 hash](/tmp/rlk-adaln-local-acceptance-20261009T174946Z/source.sha256)。 + +- correctness：exit=1；passed=n/a；[完整日志](/tmp/rlk-adaln-local-acceptance-20261009T174946Z/correctness.log)。 + +失败原因见[failure.log](/tmp/rlk-adaln-local-acceptance-20261009T174946Z/failure.log)，后续阶段未启动，不升级对应资格。 + + +## 真实形状暴露的独立参考错误（2026-10-10） + +首次完整运行381 passed、3 failed、14 deselected；三个FP32真实S均在dMod accuracy断言失败，fixed CPU与CUDA/Triton的bytes比较已经通过。旧独立参考使用CPU broadcast autograd默认S归约，该归约不遵守合同的token递增FP32 left-fold。 + +独立最小反例：B1,S8,H2，x每行[-1,1]、eps0、shift/scale0，dy首通道[2^24,1,-2^24,1]重复两次；合同dshift=1、dscale=-1，旧gold得到2、-2。关键断言是公开GtestAdaLNReference VJP对手算literal的bytes，而非仅互相比对两个算子。 + +[test-first red](/tmp/rlk-adaln-gold-token-order-red.log)1 failed→[green](/tmp/rlk-adaln-gold-token-order-green.log)1 passed。修复只在独立gold的shared affine使用独立CPU FP32 token fold，未调用任何生产normalization/backward helper；norm和dX仍由独立CPU F.layer_norm/autograd提供。indexed参考路径及所有原容差不变，native源码/.so未变。 + +当前最小gold归约反例 **PROVEN**；完整真实形状复跑和后续sanitizer/benchmark状态以后续流水线证据为准。 + + +## 本地完整验收流水线（2026-10-10，SM86） + +当前流水线状态：**succeeded**；H100/SM90仍为 **UNPROVEN**。 + +[完整状态与逐阶段命令](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/status.json)、[固定环境及 binary hash](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/manifest.json)、[源码 hash](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/source.sha256)。 + +- correctness：exit=0；passed=385；[完整日志](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/correctness.log)。 +- racecheck：exit=0；passed=210；[完整日志](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/racecheck.log)。 + ========= RACECHECK SUMMARY: 0 hazards displayed (0 errors, 0 warnings) +- memcheck：exit=0；passed=224；[完整日志](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/memcheck.log)。 + ========= ERROR SUMMARY: 0 errors +- benchmark：exit=0；passed=n/a；[完整日志](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/benchmark.log)。 + +54项 benchmark 覆盖2 dtype×3 S×3 backend×3 mode；[JSON含每次采样](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/benchmark/benchmark.json)、[CSV](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/benchmark/benchmark.csv)。 + +mode分别为forward、backward_only（forward graph在计时区域外准备）、forward_backward。warmup1/repeat3，时间为GPU event median；保存allocated/reserved峰值及相对计时前的extra峰值。属于本机短采样，H100需重新采样。 + +本次可将三真实S、完整launch×batch/padding、所选sanitizer矩阵及benchmark的本机指定配置记为PROVEN；不代表一般输入域或SM90资格。源码和二进制在各阶段前后保持不变。 diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 6eb03d9f2..023da17c7 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -280,18 +280,43 @@ 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 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 custom backward helper is used. + 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) if modulate_index is None: - shift, scale, gate = mod32[:, None, :].chunk(3, dim=-1) + norm = torch.nn.functional.layer_norm(x32, (x.shape[-1],), eps=eps) + return _GtestSharedAdaLNAffine.apply(norm, mod32) else: batch = x.shape[0] rows = torch.arange(batch)[:, None] + modulate_index.cpu().long() * batch diff --git a/tests/test_adaln_gtest.py b/tests/test_adaln_gtest.py index 0a5e54818..c505399ac 100644 --- a/tests/test_adaln_gtest.py +++ b/tests/test_adaln_gtest.py @@ -428,6 +428,19 @@ def test_adaln_modulation_gradient_uses_ascending_token_order(shared_op): _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", [ From e2c36d300914d6fc2de1b682218f88b62cf0688f Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Sat, 10 Oct 2026 12:56:05 +0800 Subject: [PATCH 09/37] Remove CUDA 12.6 installation and verification scripts - Deleted installation script for CUDA 12.6 on WSL. - Removed verification script for CUDA 12.6 installation. - Eliminated build extension script for CUDA-related extensions. - Removed testing scripts for AdaLN CUDA extension. - Deleted benchmarking script for AdaLN CUDA operations. - Removed race checking script for AdaLN CUDA extension. Signed-off-by: hongleng-github --- .github/workflows/ci.yml | 4 +- .github/workflows/gpu-ci.yml | 2 +- .github/workflows/ws1-chain-gpu.yml | 4 +- .github/workflows/ws1-chain-npu.yml | 4 +- .github/workflows/ws1-gtest-gpu.yml | 4 +- docs/operators/adaln-modulation.md | 146 +----- docs/reviews/adaln-shared-b31fac9-followup.md | 172 ------- docs/reviews/adaln-shared-b31fac9.md | 471 ------------------ .../01_install_cuda_12_6_wsl.sh | 19 - .../02_verify_cuda_12_6.sh | 17 - .../03_build_extension.sh | 24 - .../04_test_adaln_cuda.sh | 26 - .../05_benchmark_adaln_cuda.sh | 15 - .../06_racecheck_adaln_cuda.sh | 24 - 14 files changed, 21 insertions(+), 911 deletions(-) delete mode 100644 docs/reviews/adaln-shared-b31fac9-followup.md delete mode 100644 docs/reviews/adaln-shared-b31fac9.md delete mode 100755 scripts/adaln_cuda_acceptance/01_install_cuda_12_6_wsl.sh delete mode 100755 scripts/adaln_cuda_acceptance/02_verify_cuda_12_6.sh delete mode 100755 scripts/adaln_cuda_acceptance/03_build_extension.sh delete mode 100755 scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh delete mode 100755 scripts/adaln_cuda_acceptance/05_benchmark_adaln_cuda.sh delete mode 100644 scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a678e3015..edcef1e36 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,9 +5,9 @@ name: CI-Pipeline on: push: - branches: [ main, test, test-qwenimage ] + branches: [ main, test ] pull_request: - branches: [ main, test, test-qwenimage ] + branches: [ main, test ] permissions: contents: read diff --git a/.github/workflows/gpu-ci.yml b/.github/workflows/gpu-ci.yml index c73b791a3..dbf822d4d 100644 --- a/.github/workflows/gpu-ci.yml +++ b/.github/workflows/gpu-ci.yml @@ -2,7 +2,7 @@ name: GPU CI on: pull_request_target: - branches: [ main, test, test-qwenimage ] + branches: [ main, test ] paths: - 'csrc/**' - 'rl_engine/**' diff --git a/.github/workflows/ws1-chain-gpu.yml b/.github/workflows/ws1-chain-gpu.yml index 9849ffe4a..43d63fc13 100644 --- a/.github/workflows/ws1-chain-gpu.yml +++ b/.github/workflows/ws1-chain-gpu.yml @@ -8,9 +8,9 @@ name: WS1-chain-GPU on: pull_request: - branches: [ main, test, test-qwenimage ] + branches: [ main, test ] push: - branches: [ main, test, test-qwenimage ] + branches: [ main, test ] workflow_dispatch: inputs: source_repository: diff --git a/.github/workflows/ws1-chain-npu.yml b/.github/workflows/ws1-chain-npu.yml index 45fb70c7f..1d14094ed 100644 --- a/.github/workflows/ws1-chain-npu.yml +++ b/.github/workflows/ws1-chain-npu.yml @@ -15,9 +15,9 @@ name: WS1-chain-NPU on: pull_request: - branches: [ main, test, test-qwenimage ] + branches: [ main, test ] push: - branches: [ main, test, test-qwenimage ] + branches: [ main, test ] workflow_dispatch: inputs: source_repository: diff --git a/.github/workflows/ws1-gtest-gpu.yml b/.github/workflows/ws1-gtest-gpu.yml index 56bdaa9a4..51b1bb471 100644 --- a/.github/workflows/ws1-gtest-gpu.yml +++ b/.github/workflows/ws1-gtest-gpu.yml @@ -10,7 +10,7 @@ name: WS1-gtest-GPU on: pull_request: - branches: [ main, test-qwenimage ] + branches: [ main ] paths: - "rl_engine/validation/**" - "rl_engine/backends/**" @@ -42,7 +42,7 @@ on: - "ci/scripts/run_gpu_ci.sh" - ".github/workflows/ws1-gtest-gpu.yml" push: - branches: [ main, test-qwenimage ] + branches: [ main ] paths: - "rl_engine/validation/**" - "rl_engine/backends/**" diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index de1c96bf4..65584218c 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -56,22 +56,17 @@ These are derived model shapes, not tensor shapes stated by the issue. ## Validation -The [2026-10-09 shared-workspace review](../reviews/adaln-shared-b31fac9.md) -compares `b31fac9` with the actual worktree, including untracked tests. It records -current CPU/SM86 Triton evidence separately from pending native CUDA/H100, -full-image, geometry, sanitizer and benchmark acceptance. New public-interface -canaries check BF16 cast boundaries, each reduction, signed zero, subnormal eps, -gate/dgate payloads, strided upstreams and H=3072 batch positions with padding -across all declared launch configurations. AdaLN gtest rejects a candidate's wrong -output dtype and explicitly promotes gold leaves to FP32 -before autograd; an FP32 cast inside the gold function alone would still return -BF16 gradients to BF16 leaves. The review includes the red/green evidence and -H100 commands; historical results below are not current-worktree evidence. +Public-interface tests check BF16 cast boundaries, each reduction, signed zero, +subnormal eps, gate/dgate payloads, strided upstreams and H=3072 batch positions +with padding across all declared launch configurations. AdaLN gtest checks +candidate output dtype and promotes gold leaves to FP32 before autograd; +casting inside the gold function alone would return BF16 gradients to BF16 leaves. AdaLN is registered in `gtest/operator_specs.py` and `operator_inputs.py` as `op_class="reduction"`, with PyTorch, CUDA and Triton candidates. The independent -gold uses FP32 CPU `torch.nn.functional.layer_norm` and autograd, without the -production fixed-reduction or custom backward helpers. It checks both outputs +gold uses FP32 CPU `torch.nn.functional.layer_norm` and autograd for normalization. +Its independent affine VJP folds token contributions in ascending FP32 order; +it does not call production reduction or backward helpers. It checks both outputs (y and gate) and gradients of x and modulation. Accuracy thresholds come from `resolve_tolerance`: `forward_accuracy` for outputs and `gradient_accuracy` for input gradients; no private AdaLN thresholds or SM90 override are added. @@ -86,18 +81,12 @@ Existing fixed-reference byte comparisons remain separate, unchanged gates. Use CUDA/Triton candidates on the target GPU for native accuracy evidence. The three full-image tests additionally check independent CPU FP32 forward and -backward accuracy for both GPU backends/dtypes. These additions are **pending -GPU execution**; old SM86 results do not validate the new assertions. +backward accuracy for both GPU backends/dtypes. These assertions require GPU +execution; CPU results alone do not validate them. CPU tests also cover gate-only gradients, distinct image/text modulation chunks, corrupted-output/VJP rejection and unavailable-backend routing. The CLI accuracy report alone is not a provenance-checked WS1 system gate. -CI branch filters now include `test-qwenimage`; CPU CI runs the new contract -tests and the existing WS1 GPU script includes AdaLN accuracy checks. Fork -isolation and existing GPU authorization conditions remain. A missing/skipped -workflow is not a passing check; no CI run is claimed for these local changes. - - Run `python -m pytest tests/test_adaln_modulation.py -q` for focused tests. Set `RLK_ADALN_REAL_SHAPES=1` to include the three full image shapes. Run `python benchmarks/benchmark_adaln_modulation.py --real` for the three @@ -112,116 +101,5 @@ CUDA `threads={128,256,512}` (default 256); Triton `num_warps={4,8}` (default 4), `reduction_tile={64,128,256}` (default 128). Registry dispatch uses those defaults. These settings change scheduling, not the reduction order. -### Historical evidence limitation (verified 2026-10-02) - -The following SM86 results and timings are historical observations. The two -original `/tmp` handoffs and raw build/test/racecheck logs are no longer -available. Surviving excerpts and recorded hashes cannot reconstruct or -substitute for complete logs, binary provenance or newly executed evidence. -Re-run acceptance and publish fresh complete artifacts before using it as a -reviewable native acceptance result. H100 validation is **pending**. - -Local CUDA validation environment (2026-09-23): RTX 3050 Ti (SM86, 4 GiB), -driver 610.60, CUDA toolkit/runtime 12.6, PyTorch 2.12.1+cu126, Triton 3.7.1, -Python 3.12.12. The extension was built with fast math disabled for SM86. -The pre-review focused suite plus all three real shapes passed 18 tests with no skips. -These historical results are not acceptance evidence for the synchronization fix. -FP32 small-shape and BF16 real-shape CUDA forward/backward outputs were byte -equal to the CPU reference. CUDA and Triton were also byte equal directly at -all three real shapes, and both passed FP32/BF16 batch-size and token-padding -byte-invariance checks. - -Real-shape medians below use one warmup and three samples, so they are smoke -measurements rather than stable performance claims: - -| dtype | S | Triton fwd / fwd+bwd (ms) | CUDA fwd / fwd+bwd (ms) | -| --- | ---: | ---: | ---: | -| BF16 | 4096 | 3.069 / 20.709 | 2.998 / 23.262 | -| BF16 | 6889 | 5.125 / 34.551 | 5.150 / 39.963 | -| BF16 | 6032 | 4.689 / 30.302 | 4.357 / 34.718 | -| FP32 | 4096 | 5.944 / 29.269 | 7.229 / 38.693 | -| FP32 | 6889 | 9.840 / 48.579 | 11.948 / 65.271 | -| FP32 | 6032 | 8.685 / 81.903 | 10.639 / 57.270 | - -Forward and forward+backward peak allocation matched between CUDA and Triton -at each shape. The CUDA backward path was generally slower in this short SM86 -run; rerun with the default warmup/repeat counts on target hardware before -making a performance claim. - -## Review regression checks (2026-09-24) - -The pre-fix current-HEAD suite passed 20 cases (including extension smoke), -but compute-sanitizer found one forward and three backward shared-memory hazards. -The missing synchronization was after reading the reduction result, before -reusing its buffer. The minimal fixed-seed H=64 regression returned exit code 7 -even though its numerical assertion passed. - -The non-native review selection passed 30 cases with 26 deselected using -`-k 'not cuda and not real_image and not gpu_padding'`. This does not validate -the rebuilt native extension. - -Post-fix validation of the rebuilt extension passed **57 tests in 43.73s** with -no skips, including native extension smoke and the complete focused matrix. -The native registry selected `CUDA_ADALN_MODULATION` with `fallback=False`. -The expanded sanitizer regression passed **8 tests** with **0 hazards, -0 errors, and 0 warnings** (exit 0). The original seed-386 FP32 -`[1,2,3072]` forward/backward reproducer also reported zero hazards after -the fix. Environment remained PyTorch 2.12.1+cu126, Triton 3.7.1, Python -3.12.12, RTX 3050 Ti SM86. - -Observed commands (the sanitizer ran separately before the full pytest matrix): - -```bash -bash scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh -PYTHONDONTWRITEBYTECODE=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ -RLK_ADALN_REAL_SHAPES=1 RL_KERNEL_REQUIRE_EXT=1 \ -.venv/bin/python -m pytest tests/test_extension_smoke.py \ - tests/test_adaln_modulation.py -q -rs -p no:cacheprovider -``` - -Original, now missing logs were `/tmp/rlk-adaln-race-fix-build.log`, -`/tmp/rlk-adaln-race-fix-racecheck.log`, and -`/tmp/rlk-adaln-race-fix-tests.log`. They describe the uncommitted review -fixes on top of `d14e4ed`; that commit alone predates the fix. - -Run the combined fail-closed checks using the cu126 repository environment: - -```bash -bash scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh -# Or only the sanitizer regression: -bash scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh -``` - -Script 04 runs the full focused suite with real shapes and required native CUDA, -then script 06. Missing extension symbols or sanitizer failures must not count -as successful native acceptance. Do not use bare `uv run` or `uv sync` to -run these checks: invoke `.venv/bin/python` directly to preserve cu126. - -The statuses below describe historical assertions within measured SM86 cases, -not retained complete execution evidence or current H100 acceptance. - -| Requirement | Implementation | Test | Observed evidence | Remaining gap | Status | -| --- | --- | --- | --- | --- | --- | -| Shared forward/backward math at H=7/3072 | Fixed FP32 normalization and VJP | Independent LayerNorm/autograd, both dtypes, random dy/dgate | All 12 backend/dtype/H cases passed | Independent full-image assertions added; GPU execution pending | PROVEN | -| Race-free CUDA shared reduction | Read-result/barrier/reuse | Native harness and geometry under racecheck | 8 tests, zero hazards; original reproducer also green | Other architectures not exercised | PROVEN | -| Launch geometry / tiling | CUDA threads 128/256/512; Triton 4/8 warps and 64/128/256 tiles | Public constructors, CPU-reference byte checks | Both dtypes passed all 3 CUDA and 6 Triton configurations | Coverage is H=3072, B=2, S=5 | PROVEN | -| Strict byte comparator | Contiguous uint8 views | Signed-zero regression and backend comparisons | Comparator rejects +0 versus -0; full matrix passed | Exceptional floating-point payloads not covered | PROVEN | -| Three real image sizes | S=4096/6889/6032, H=3072 | Both dtypes, random dy/dgate, CPU/native/Triton byte checks | All 6 real-shape cases passed | Independent FP32 CPU LayerNorm/autograd added; GPU execution pending | PARTIAL | -| Indexed select01 | Explicit PyTorch fallback with [2B,3H] modulation | B=2, both index shapes/devices/dtypes, independent row-index and autograd gold; invalid routes | 12 indexed cases passed | No fixed-order indexed backward claim | PROVEN | -| FP32 accumulation / final BF16 cast | FP32 reductions, CUDA RN operations, Triton fusion disabled | Cast, byte-reference and independent accuracy tests | All focused checks passed | Exhaustive rounding boundaries and fast-math-on build exclusion not exercised | PARTIAL | -| Batch / padding invariance | Row-local hidden reduction, ascending token fold | Batch position, B=1/3, S=3/5, zero-gradient padding; raw bytes | FP32/BF16 CUDA and Triton cases passed | Tested padding matrix uses H=32 | PROVEN | -| Registry / trace | Explicit backend selection, H limit, indexed fallback reason | Registry tests and native preflight | Native selected without fallback; indexed fallback and H limit passed | Semantic fingerprint is not a binary hash; CPU unavailable-backend routing tests added; actual GPU fallback execution pending | PARTIAL | -| Native SM86 build and binding | Conditional sources/symbols, defaulted threads argument | Symbol verification, native/reference tests and geometry launches | User build completed; new four-argument binding and full suite passed | Other build configurations untested | PROVEN | -| Indexed optimized kernels | Agreed explicit PyTorch fallback | Fallback and trace checks | Within agreed scope | No optimized indexed kernel required | OUT OF SCOPE | -| Full #386 WS1 exit criteria | 22 WS1 rows plus Foundation, WS2 and integration | MMDiT/LoRA, 60-layer curve, reproducible rollout logp | Not exercised by this operator suite | Separate system-level work | UNPROVEN | - - -The independent autograd checks use the shared reduction accuracy contract. -They do not replace strict fixed-reference and batch/topology byte checks. -Full-image independent assertions are implemented but await GPU evidence; -exhaustive rounding boundaries and alternative fast-math builds remain gaps. - -The semantic fingerprint does not identify a compiled binary. This suite does -not claim coverage across other GPU architectures, every floating-point input, -or fast-math-enabled builds. Benchmark numbers above predate the synchronization -fix; they are historical smoke measurements, not performance claims for the fix. +Native acceptance requires a fresh extension build, source/binary identity records, +and complete test and sanitizer logs. diff --git a/docs/reviews/adaln-shared-b31fac9-followup.md b/docs/reviews/adaln-shared-b31fac9-followup.md deleted file mode 100644 index f65ab8e32..000000000 --- a/docs/reviews/adaln-shared-b31fac9-followup.md +++ /dev/null @@ -1,172 +0,0 @@ -# AdaLN shared：新增缺口与当前版本证据（2026-10-09） - -> 本文是第二轮中间源码快照。用户随后授权继续改实现、补测试；最新状态与证据已统一记录在 -> [指定主报告](adaln-shared-b31fac9.md)。本文的 hash、运行表和未关闭项状态不代表最新工作区。 - -仍以 [issue #386](https://github.com/RL-Align/RL-Kernel/issues/386) 为规格，固定比较点 -`b31fac988faae9de71ab1f683e4a5c9cc8578b84`，HEAD -`120f4de7274798c6310d87799ebb6538f2f8f531`。审查范围包括 `git diff b31fac9` -及所有未跟踪文件；原有用户修改保留，没有提交或推送。 -[首轮报告](adaln-shared-b31fac9.md) 的源码 hash 与运行表属于此前快照，不能用作当前版本运行证据。 - -本次沿用 `/code-review` 的独立 Standards/Spec 审查及 `/tdd` 公开接口测试流程。 -测试只通过公开 op 调用、`torch.autograd` 和 `run_operator_suite` 观察结果。 -确实失败的生产行为逐项先 red 再修复；原本正确但缺断言的行为补 canary，另用公共 -callable 错误候选验证断言敏感性，没有为了制造 red 修改生产实现。 - -PROVEN 仅指表中指定有限反例在当前 CPU/SM86 Triton 版本运行通过。 -PARTIAL 表示有局部证据而未覆盖完整域。UNPROVEN 表示尚无当前版本运行证据。 -这些状态不表示 CUDA/H100 全平台验收完成。 - -## Standards：验证器确实放过错误 dtype - -**P1，已修复。** AdaLN BF16 输入必须输出 BF16，但无 provenance 的默认 gtest candidate -会将输出转 FP32 后比较，不检查声明的输出边界。公开错误候选 -`tuple(output.float() for output in good.fn(**inputs))` 在数值与梯度比较中原本可以通过。 -新增 forward-only 和 check-grad 两个反例均先失败。 - -最小修复：`OperatorCase.candidate_output_dtype` 和 spec 的 `output_matches_input_dtype` -显式启用 AdaLN 输出 dtype 检查；失败报告包含 `dtype mismatch`。 -其他算子默认不启用，避免将本就应输出 FP32 的算子误判。gold leaves 的 FP32 策略继续保留。 -修复后检查器相关 39 项通过，最终 CPU suite 89 项通过。 - -## Spec：三项实际实现错误与其他证明缺口 - -**P1，已修复：Triton 默认 FTZ 改变合法 subnormal eps。** 本机 Triton 3.7.1 默认 -`enable_reflect_ftz=True`；本轮 probe 的 PTX 中出现 `sqrt.rn.ftz.f32`。 -公开 `eps=1e-40`、常量零输入返回 NaN，而 CPU 按 FP32 eps 应输出零。 -forward 与 backward rows 现均显式关闭 reflect FTZ；两 dtype 的 y/dX/dMod literal 验证通过。 -不能仅凭 `sqrt_rn` 名称或 `enable_fp_fusion=False` 断言没有 FTZ。 - -**P2，已修复：Triton tree root 提取额外加正零。** `tl.sum(where(col==0,root,+0))` -把 tree 的负零变成正零，改变 mean 和最后输出符号。现直接 gather 单元素 root 再归约， -移除额外正零。不是放宽 byte assertion。 - -**P2，已修复：CPU BF16 gate FP32 往返转换损坏 NaN payload。** gate 没有算术， -原实现仍先转 FP32 再转回 BF16；现从原调制张量直接取 gate view。继续检查发现 CPU -backward 的 dgate 同样损坏上游 payload，第二轮 red→green 后改为原 dtype 直接拼入 dMod， -只在 dshift/dscale 各自输出边界 cast。finite 与 nonfinite identity VJP 都验证。 -issue 未显式定义 LayerNorm 的 nonfinite 域，本项只证明 gate 位透传，不声称整个 nonfinite -输入域已符合规格。 - -以下新增测试均在 [test_adaln_gtest.py](../../tests/test_adaln_gtest.py)。 -独立 literal 来自标出的标量手算；数学 accuracy 使用独立 CPU FP32 LayerNorm/autograd； -拓扑 differential 只用于 bytes 不变，不充当独立数学参考。 - -| 要求 / 当前实现与独立参考 | 具体错误实现 | 最小反例 | 应失败的断言 | 当前结论 | -| --- | --- | --- | --- | --- | -| BF16 kernel 输出边界；checker 现显式检查 dtype | 正确 BF16 y/gate 再 `.float()` 返回；旧 checker 数值全通过 | B2,S3,H7 BF16，公开错误 candidate | `gtest_rejects_fp32_outputs_for_bf16_inputs`：forward/check-grad report 都失败，两输出有 dtype mismatch | **PROVEN** CPU checker，实际 red→green | -| 固定 hidden tree 的零符号；CPU 独立 IEEE literal | Triton tree root 提取时再加 +0，擦掉负零 mean | H2，x=[-0,-0]，shift=-0，scale=0 | `signed_zero_follows_hidden_tree`：y 必须是 +0 bytes，错误实现是 -0 | **PROVEN** CPU/SM86 Triton，两 dtype；CUDA **UNPROVEN** | -| FP32 scalar/平方根不能无声明 FTZ；独立标量舍入 | `sqrt.rn.ftz` 把 FP32(1e-40) 当作0 | x=[0,0]，m=0，eps=1e-40；dy=[1,-1] | `subnormal_eps_is_not_flushed_to_zero`：y=0；rstd=1.000002658868784e20，dX=±rstd 最终 cast；dMod=[1,-1,0,0,0,0] | **PROVEN** CPU/SM86 Triton，两 dtype，实际 red→green；CUDA **UNPROVEN** | -| gate/dgate 按原 dtype 位透传，独立 raw literal | BF16 gate或dgate→FP32→BF16，NaN canonicalize | gate raw16=[0x7f81,0xffc1]；dg分别取[2,-3]或raw16=[0x7f82,0x7fff]；x、shift、scale有限 | `gate_preserves_bf16_nan_payload_and_identity_vjp`：gate原uint8 bytes不变，dg原位进入dMod | **PROVEN** CPU/SM86 Triton；CPU fwd/bwd各实际red→green；nonfinite LN合同 **PARTIAL** | -| variance 的 hidden tree 必须独立于 mean 的 canary 检验 | mean 正确，只把 squared-centered 的树换成左折叠 | H4 x=[4096,1.25,-4096,-1.25]，eps0；mean 两算法都0；variance 应8388609，错误为8388608 | `variance_reduction_uses_declared_pairwise_order`：y literal ±1.4142134189605713、±0.0004315836704336107 的 bytes | **PROVEN** CPU/SM86 Triton；对应 mutant 被拒；CUDA **UNPROVEN** | -| backward sum(dnorm) 的固定树 | 只把第一个 backward H sum 改左折叠 | x=[-1,1,-1,1]，eps0，dy=[2^24,1,-2^24,1] | `backward_hidden_sums_use_declared_pairwise_order(sign=1)`：dx=[2^24,0,-2^24,0]；错误中间通道为.25 | **PROVEN** CPU/SM86 Triton；单独 mutant 被拒；CUDA **UNPROVEN** | -| backward sum(dnorm*norm) 的固定树 | 只把第二个 backward H sum 改左折叠 | 同上，dy=[-2^24,1,2^24,1] | 同一测试 sign=-1；dx=[-2^24,0,2^24,0]；错误中间通道为.25 | **PROVEN** CPU/SM86 Triton；单独 mutant 被拒；CUDA **UNPROVEN** | -| BF16 最终边界 RN ties-to-even，forward/backward 都须检查 | 中点向上舍入或截断；随机用例很难撞中点 | x=[-1,1]，eps0，正通道 shift=f；S2 正通道 dy=[1,f]；f=1/256、3/256 | `boundary_rounds_halfway_values_to_even`：1+f 应分别 round 到1和1+4/256；y、dshift、dscale bytes | **PROVEN** CPU/SM86 Triton；四个 forward/backward 舍入 mutant 被拒；CUDA **UNPROVEN** | -| upstream stride 必须按公开 autograd 语义处理 | 按紧密连续地址读 dy，或忽略 dg 的 stride0 | B2,S3,H7；dy 从[B,H,S] transpose，dg 首样本 expand 至B2 | `backward_accepts_transposed_and_broadcast_upstreams`：dX/dMod 等于 contiguous upstream 的 bytes，且对独立 CPU FP32 gold 在声明容差内 | **PROVEN** CPU/SM86 Triton，两 dtype；raw-read mutant 被拒；CUDA **UNPROVEN** | -| 输入可分别冻结，输出可真正不参与 loss | 强制两个输入都求导；backward 假设两个输出一定都有传入 VJP | 仅m训练，y.sum且gate不用；仅x训练、dy=[1,0]、x=[-3,3]eps7；仅gate loss | `single_trainable_input_and_unused_output`：冻结输入.grad is None；dm=[1,1,-1,1,0,0]；dx=±7/128；gate-only dx=0/dm gate identity | **PROVEN** CPU/SM86 Triton，两 dtype；不扩展为二阶梯度或 LoRA 模型验收 | -| H3072 的 batch/padding 必须与 launch 同时改变 | warps8/tile256 路径只在 position>0 读取错 mod；单独 topology、单独 geometry 可漏 | default4/128的B1,S3基准；改4/64或8/256，同时B3的每个position与S3/5，strided输入与6H调制切片 | `h3072_relocation_with_alternate_triton_launch`：y/gate/dX/dMod bytes 等于默认 launch 的同一样本；unused半区0，padding dx数值0 | **PROVEN** SM86 Triton 两额外配置×两 dtype；全部六配置/native三threads交叉 **PARTIAL** | -| unsupported batch 几何必须在入口清楚处理 | C++ backward `grid.y=batch`，既未限B也未展平，B65536进入非法grid | B65536,S1,H1，gate-only backward；不到几MiB | 允许的输入应成功且dx0/dm gate1；若明确不支持，应在公开入口清楚拒绝；不能以 invalid configuration launch失败 | **PARTIAL** 源码定位在 CUDA L194；当前 native runtime **UNPROVEN**，未冒充已修复 | -| 地址乘积不能静默 int32 overflow | `row*hidden`、`(batch*seq+token)*hidden`、`batch*3*hidden` 全以int计算，无 numel 范围检查 | x寻址 B1,S524289,H4096 BF16（约4GiB）；mod寻址 forward B174763,S1,H4096 | 有界支持则公开入口拒绝；支持则 int64寻址并memcheck零错误；不是普通数值 allclose | **PARTIAL** 静态类型/边界证据；实际巨大输入与当前 native **UNPROVEN**；未分配/运行 | -| eps 公共值域与FP32 scalar转换策略须明确 | finite double eps 通过校验，但cast后为Inf或0；数值行为未定义 | eps1e39→FP32 Inf，eps1e-46→0；常量x的后者产生NaN | 按选定政策验证round-to-FP32结果，或清楚拒绝超范围值；不能只断言finite(eps) | **PARTIAL** CPU/SM86 probe 已观察转换结果；是否允许此域规格未明确，未擅自更改API | -| native H上限与非二次幂边界必须交叉验收 | off-by-one允许H4097或错拒H4096；尾部padding树错误 | H4095/4096/4097，两dtype，direct native及registry | 4095/4096正确，4097明确拒绝或有声明fallback；独立数学与bytes都检查 | **UNPROVEN** 当前 native边界；旧大H/irregularH测试不替代这组边界 | - -CUDA grid.y 上限65535见 -[NVIDIA compute-capability limits](https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/compute-capabilities.html)。 -第12/13项是未覆盖的实现范围错误候选，不是模型常用 B 的性能问题。 -若将B限制到65535且H≤4096，modulation乘积自然小于int32上限，但x的B*S*H仍需单独检查。 - -## 当前版本运行证据 - -环境仍为 Python3.12.12、PyTorch2.12.1+cu126、Triton3.7.1、RTX3050Ti/SM86。 -没有构建 native；现存 `.so` 未绑定当前源码，因此没有执行其数值测试。 -下面最终 suite 与 mutation 运行对应文末 hash；red/局部green是本轮 TDD过程的中间源码证据。 - -| 运行 | 结果 / artifact | -| --- | --- | -| signed-zero red→green | 2 failed/2 passed→4 passed；[/tmp red](/tmp/rlk-adaln-more-signed-zero-red.log)、[green](/tmp/rlk-adaln-more-signed-zero-green.log) | -| subnormal eps red→green | 2 failed/2 passed→4 passed；[red](/tmp/rlk-adaln-more-subnormal-red.log)、[green](/tmp/rlk-adaln-more-subnormal-green.log) | -| CPU gate payload red→green | 1 failed→1 passed；[red](/tmp/rlk-adaln-more-gate-red.log)、[green](/tmp/rlk-adaln-more-gate-green.log) | -| CPU gate VJP payload red→green | 1 failed/1 passed→2 passed；[red](/tmp/rlk-adaln-more-gate-vjp-red.log)、[green](/tmp/rlk-adaln-more-gate-vjp-green.log)；[Triton局部检查](/tmp/rlk-adaln-more-gate-vjp-triton.log)2 passed | -| gtest dtype red→green | 2 failed→相关39 passed；[red](/tmp/rlk-adaln-more-dtype-red.log)、[green](/tmp/rlk-adaln-more-dtype-green.log) | -| 最终CPU/checker | **89 passed、83 deselected，7.29s，无skip**；[log](/tmp/rlk-adaln-more-final-cpu.log) | -| 最终CPU/SM86 Triton，包括额外launch交叉 | **124 passed、67 deselected，14.07s，无skip**；[log](/tmp/rlk-adaln-more-final-gpu.log)；未运行native、大图、旧完整geometry sweep | -| 最终新增错误候选 | **8/8拒绝**；[script](/tmp/rlk_adaln_more_mutations.py)、[log](/tmp/rlk-adaln-more-mutations-final.log) | -| 最终首轮错误候选重跑 | **6/6拒绝**；[script](/tmp/rlk_adaln_review_mutations.py)、[log](/tmp/rlk-adaln-more-old-mutants-final.log) | - -```bash -# 最终 CPU 命令 -PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 .venv/bin/python -m pytest \ - tests/test_adaln_gtest.py tests/test_op_checks.py tests/test_operator_inputs.py \ - -q -rs -k 'not triton and not cuda' - -# 最终 SM86 短测试命令；明确排除未经当前源码构建的native及长检查 -PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 TRITON_CACHE_DIR=/tmp/rlk-adaln-review-triton \ - .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_adaln_modulation.py \ - -q -rs -k 'not cuda and not real_image and not gpu_padding and not launch_geometry' -``` - -## H100：新增验证清单与命令，未启动 - -先执行[首轮报告的源码快照、pinned构建和完整测试命令](adaln-shared-b31fac9.md)。 -需要将本 followup 与新的 untracked test 一起带过去，不能用旧 hash/旧 binary。 - -在已经完成当前源码构建的 H100 环境中，新增 canaries 可单独检查： - -```bash -RL_KERNEL_REQUIRE_EXT=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ - .venv/bin/python -m pytest tests/test_adaln_gtest.py -q -rs \ - -k 'signed_zero or subnormal_eps or gate_preserves or rejects_fp32_outputs or variance_reduction or backward_hidden_sums or halfway_values or broadcast_upstreams or single_trainable or all_launch_configurations' -``` - -小内存 native batch 边界公开接口复现命令（这里未运行）： - -```bash -RL_KERNEL_REQUIRE_EXT=1 .venv/bin/python - <<'PY' -import torch -from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp -op = CudaAdaLNModulationOp() -for batch in (65535, 65536): - x = torch.zeros(batch, 1, 1, device="cuda", requires_grad=True) - m = torch.zeros(batch, 3, device="cuda", requires_grad=True) - y, gate = op(x, m) - dx, dm = torch.autograd.grad(gate, (x, m), torch.ones_like(gate)) - torch.cuda.synchronize() - assert torch.equal(y, torch.zeros_like(y)) - assert torch.equal(dx, torch.zeros_like(dx)) - expected = torch.zeros_like(m) - expected[:, 2] = 1 - assert torch.equal(dm, expected) - print("PASS batch", batch) -PY -``` - -按当前未限定B的API，第二个batch是预期暴露错误的点。若修复选择限定支持域， -需同步声明限制并将边界断言改为入口明确拒绝，不能将 invalid launch当正常拒绝。 -int32 巨大输入反例只准备参数与memcheck清单；没有分配约4GiB输入，更没有启动该测试。 - -- 全部新增 canary 在当前 native/SM90 和Triton两dtype运行，记录零skip/xfail与完整输出。 -- 把全部3种native threads、全部6种Triton warps/tile与H3072 relocation/padding交叉， - 直接比较同一默认B1基准的 y/gate/dX/dMod bytes；本机只新增两组Triton配置。 -- native B65535/65536、H4095/4096/4097正负边界不能遗漏。 -- 寻址选择64位支持或显式有界拒绝；若支持巨大输入，再用compute-sanitizer memcheck检查 - B1,S524289,H4096及modulation地址范围，保留0 errors证据。 -- 对新FTZ修复保存当前Triton版本/compile options/PTX，确认fwd和bwd_rows均无sqrt FTZ； - 其他Triton版本/ROCm配置须独立qualification。 -- 保留首轮要求的三真实image S、benchmark、racecheck、源码/二进制hash及pinnedruntime记录。 - -## 对应最终测试的源码 SHA256 - -| 文件 | SHA256 | -| --- | --- | -| 未跟踪 `tests/test_adaln_gtest.py` | `85de4e0760d6ead947b0af1e8bd4eb896cfdea6bccf4434ca46d5a3adf67e165` | -| `gtest/op_checks.py` | `70f9c648b5da5ace001d295569c4b05a15d2e630e84e1870b5734a5377d76b39` | -| `gtest/operator_specs.py` | `b78d15322199b4986f19697c7987dea879238b7d8b14fe2f1b05fc3fe0268e79` | -| PyTorch AdaLN | `1b573827279e89c97e4f293bf867fe5a3a77102c89cf9427ff0e1f429a351c57` | -| Triton AdaLN | `ef2b13c8e13da8d324b160489efad4b2080d3a1058cf3b58d566c2bc92a3b4a2` | -| CUDA source（未修改） | `27074706332b1db32236b6cbb225ee99514938f38d1996ef4812e221d857e80d` | - -静态检查：修改的五个Python文件ruff/isort均通过；新增测试与checker的Black check通过， -`git diff --check`通过，验收shell的`bash -n`通过。 - -`/tmp`日志与mutation脚本是本次artifact，应另行归档。新增文档不会改变上述测试源码hash。 diff --git a/docs/reviews/adaln-shared-b31fac9.md b/docs/reviews/adaln-shared-b31fac9.md deleted file mode 100644 index 27929d32a..000000000 --- a/docs/reviews/adaln-shared-b31fac9.md +++ /dev/null @@ -1,471 +0,0 @@ -# AdaLN shared：工作区审查与 TDD 证据（更新于 2026-10-10） - -## 当前 native SM86 短验收（2026-10-10） - -用户授权先本地提交、重新构建 native，再执行全部短测试。输入提交为 -`f21007cccc216d214ec2589cc175bc288a5363e0`，固定比较点仍为 `b31fac9`。 -该提交包含此前全部未提交改动、两个未跟踪测试和两个审查文档，未推送。 -本节取代下方构建前记录中的 native UNPROVEN 结论；仅对明确列出的本机反例升级资格。 - -全量构建使用独立 objects/lib 目录和 `build_ext --inplace --force`，旧 `.so` 已隔离。 -环境固定为 Python 3.12.12、PyTorch 2.12.1+cu126、CUDA toolkit 12.6.85、GCC 11.4、 -RTX3050Ti/SM86、MAX_JOBS=2、fast math=0，未安装或下载依赖。 -构建成功，实际 import 路径、构建目录产物和安装产物 hash 一致；构建期间源码不变。 -[构建参数及工具 hash](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/build-manifest.json)、 -[全部源码 hash](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/source.sha256)、 -[源码快照](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/source.tar.gz)、 -[完整构建日志](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/build.log)、 -[实际编译参数](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/objects/build.ninja)、 -[import 验证](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/import-verification.log)。 - -当前 `.so`：`rl_engine/_C.cpython-312-x86_64-linux-gnu.so`；SHA256： -`94b86feb52cb621794598b9ba6361ae3d521a370b9e6c5df74942112f2f35e61`。 -旧二进制 `d7b91b28...` 不参与本次运行，也不算当前证据。 - -先运行提交中现有用例：**89 passed，182 deselected，5.38s**。 -检查后补齐 B65536 的 BF16 和全部 threads，强化 B65535 的非零 y backward，得到 -**98 passed，182 deselected，5.68s**。最终进一步将 CUDA fixture 扩展到三种 threads， -让每个 canary、两 stream/两个 triple 和合法 H 边界均交叉执行： -**216 passed，182 deselected，7.23s；零 skip、xfail、xpass、失败或 teardown 错误**。 -没有暴露需要修改 native 实现的错误；本轮只修改两个测试文件和本报告，未新增提交。 - -| 可能错误 / 要求 | 本次反例与必须成立的断言 | 当前结论 | -| --- | --- | --- | -| 提前 BF16 舍入 norm/affine、保存量、partial 或 token accumulator | 早 cast 残差、1799/32768 VJP、H7/3072 BF16 对提升后 FP32 单次最终 cast、两 halfway tie；比较 y/gate/dX/dMod bytes；各自遍历 threads128/256/512 | 指定 native canary **PROVEN** | -| hidden mean/variance 或两个 backward sum 改树，token fold 被重关联 | `[2^24,1,-2^24,1]`、variance-only 向量、正负 dy、S4 cancellation；独立手算 literal 的输出/VJP bytes；三 threads | 指定 native 归约反例 **PROVEN** | -| signed zero 被额外 +0 擦除、合法 eps 被 FTZ、gate/dgate NaN payload 被转换 | signed-zero bytes、eps1e-40/2^-149/FP32 max 的 y0/dX±rstd/dMod literal、BF16 原始 payload 透传；三 threads，适用的两 dtype | 指定 native 反例 **PROVEN**；一般 LN NaN payload 不在资格声明中 | -| B65535 的最后 batch 寻址或 backward grid.y 错误 | B65535,S1,H2,x=[-3,3],eps7；y=[-.75,.75]，dX=[7/128,-7/128]，dshift=[1,0]，dscale=[-.75,0]，dgate=[2,-3]；每个 batch 的 gate 使用 mod17 标签；两 dtype×三 threads共6例，检查所有行 | 合法最大 B native forward/backward **PROVEN** | -| B65536 被非法 launch | B65536,S1,H1；公开 op 必须 ValueError `B <= 65535`，两 dtype×三 threads共6例 | 公开 fail-closed **PROVEN**；负例在 Python preflight 拒绝，不冒充独立 C++ guard 运行证据 | -| H4095/4096 的补零或上限 off-by-one | B2,S3，两 dtype×三 threads；y/gate/dX/dMod 对独立 CPU FP32 math accuracy 和固定参考 bytes，共12例 | 指定合法 native H 边界 **PROVEN** | -| H4097 静默进入 native / 未声明 fallback | 两 dtype公开 op 明确拒绝；H4095/4096 registry 必须选 CUDA 且 fallback=False；4097必须明确 rejected reason及fallback=True | direct 拒绝及 registry policy **PROVEN**；未声称 H4097 fallback 数值已验收 | -| threads×H3072 batch position/padding 组合路径出错 | 128/256/512×两 dtype，B1/S3 baseline→B3每个位置、S3/5、strided x和6H view；y/gate/dX/dMod bytes、unused半区0、padding dX0 | 当前 native 三 geometry×topology **PROVEN** | -| 第二 triple 或两 stream 错 shift / VJP 串值 | image S5/text S3、两 triple×两 dtype×三 threads共12例；独立 FP32 CPU forward/backward、gate bytes、unused半区0、另一个stream梯度不变 | 指定两 stream/两个 triple **PROVEN** | -| 上游 stride、单个 trainable input、unused output 失效 | transposed dy、stride0 dg、x-only/modulation-only/gate-only；独立 math 或 literal bytes；三 threads×两 dtype | 指定 native autograd 用例 **PROVEN** | -| int32 地址溢出、eps 输入域或非有限传播未验证 | 巨大 meta shape明确拒绝，不分配实数据；非法eps公开拒绝；NaN/Inf x传播及独立gate VJP | 公开 preflight及指定 nonfinite native 反例 **PROVEN**；C++负例守卫未绕过Python单独测试，完整非有限域 **PARTIAL** | -| 用旧二进制 / 隐含 fallback 冒充 native | 加载路径和 .so SHA256与fresh build一致；运行前后所有生产源码及binary hash不变；可支持 H 的 registry必须 native/fallback=False；数值用例直接使用 CudaAdaLNModulationOp | 当前 SM86 源码/产物/运行闭环 **PROVEN** | -| H100、三真实图像形状、racecheck/memcheck、benchmark | SM90固定环境 fresh build；S4096/6889/6032；sanitizer零hazards/errors；完整时间/峰值内存 | **UNPROVEN**，本轮未启动 | - -最终证据：[pytest 完整日志](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/pytest.log)、 -[JUnit](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/junit.xml)、 -[逐阶段结果及零 skip/xfail 检查](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/results.json)、 -[实际选中用例](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/selected-tests.json)、 -[固定运行环境及 registry traces](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/manifest.json)。 -182 deselected 是未选择的 CPU/Triton、indexed、真实图像及非 native checker 用例，不是 skip。 -测试中的 H4097/H5000 registry fallback 是声明过的 unsupported-H policy；支持域的数值用例无 fallback。 - -最终测试 SHA256： - -- `tests/test_adaln_gtest.py`：`dbb499193632e23ccaf415d78a246a4ac1d3f3a340809cda8ad4d1e7717efa54`。 -- `tests/test_adaln_shared_bounds.py`：`48c4432a6456bb9b25e06ffbe2d9e7457cdff7e9f25b5208fcbe310e7cc98810`。 - -运行前相对于 build source manifest 只有两个测试及本报告改变，所有生产源码保持一致；完整 -[当前源码 hash](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/source-sha256-current.json)、 -[测试 patch](/tmp/rlk-adaln-native-f21007c-20261009T153114Z/native-short-20261009T174217Z/test-changes.patch) -和两个完整测试快照均已保存。运行期间源码/产物不变，报告在运行后更新。 -ruff、Black py312、isort black profile及 `git diff --check` 检查通过；没有重新编译或更改生产源码。 - -复跑短矩阵使用如下选择，并保留固定环境、import hash、registry和零skip/xfail校验； -这些额外校验的完整可审查脚本保存在证据目录 `runner.py`。 - -```bash -OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 RLK_ADALN_REAL_SHAPES=0 \ -RL_KERNEL_REQUIRE_EXT=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ - .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py \ - tests/test_adaln_modulation.py -vv -ra -p no:cacheprovider \ - -k 'cuda and not real_image and not indexed' -``` - -下一步仍是单独授权/由用户执行的 H100真实形状、sanitizer及benchmark； -本机 SM86 短验收不能升级为 SM90 或一般输入域的证明。 - -## 构建前 TDD 范围与记录（2026-10-09 历史快照） - -用户已授权继续按 `/tdd` 改实现、补测试；沿用公开 op 调用、autograd、registry、gtest -接口，不要求新的接口确认。不提交/推送，不启动 H100 长测试、巨大输入或无法估计耗时的构建。 -下面保留构建前实施待办和中间运行记录；最新 native 状态以上一节为准。 -本文件下方第一轮审查及 followup 文档都属于历史源码快照，旧运行表不能算当前版本证据。 - -| 缺口 | 修补动作 | 实施前结论 | -| --- | --- | --- | -| CUDA B65536 超出 backward grid.y | Python/C++ 声明并检查 B≤65535，补65535/65536边界测试，launch前明确拒绝 | 源码PARTIAL，当前native运行UNPROVEN | -| CUDA int32寻址范围无检查 | 在复制/launch前检查索引乘积，Python/C++双层保护；巨大形状使用meta输入测试，不分配几GiB张量 | 源码PARTIAL，运行UNPROVEN | -| shared eps和nonfinite支持域未明确 | 拒绝FP32转换后Inf或正eps下溢到0，保留合法subnormal与eps=0；补公开反例并记录nonfinite LN范围 | PARTIAL | -| geometry×batch/padding交叉不完整 | 补全6种Triton配置及3种CUDA threads，比较默认B1同样本y/gate/dX/dMod bytes | PARTIAL | -| GPU两stream及第二triple覆盖不足 | 扩展已有小形状公开测试到各后端，独立FP32梯度、unused半区0、另一stream不受影响 | PARTIAL | -| H4095/4096/4097 native边界 | 补direct op与registry正负边界，两dtype | UNPROVEN | -| 当前源码/native binary闭环缺失 | 保存源/产物hash；仅在能够可靠估计快速完成时构建，否则提供构建与验收命令 | UNPROVEN | -| H100真实图像、安全性、性能 | 保留三真实S、racecheck/memcheck、benchmark和pinned环境命令与清单，不自行启动 | UNPROVEN | - -## 构建前 TDD 完成状态与源码证据（历史快照） - -以下均通过公开 op、autograd、registry 或 gtest 观察，没有新增生产 helper mock。 -已复用旧检查器、固定字节参考、独立 FP32 CPU LayerNorm/autograd 和模型六块切片用例。 -比较点仍为 `b31fac9`，HEAD仍为 `120f4de7274798c6310d87799ebb6538f2f8f531`； -范围包含当前 tracked diff 与未跟踪的两个测试文件和审查文档。未提交或推送。 - -实现采用明确的有界拒绝：native 支持 B≤65535、H≤4096、每个输入 numel≤2^31-1。 -Python 在 `.contiguous()` 前校验，C++ 在 narrowing/launch 前重新校验；没有扩大支持域， -因此不需要实际分配约4GiB的溢出反例。registry只按H选择后端，B/numel超限在公开调用 -入口清楚拒绝，不算隐含fallback。meta测试验证的是公开Python形状preflight，不是native内核。 - -shared eps 检查 round-to-FP32 后仍有限，且正值不变成0；合法subnormal和eps=0保留。 -nonfinite LN按FP32公式传播，eps=0且方差0可以产生NaN；不承诺LN运算的NaN payload -或非有限运算域的跨后端字节相等。gate/dgate无算术，仍按原dtype位透传。 -这不是给整个nonfinite输入域宣称PROVEN，也没有改变独立indexed/select01合同。 - -| 缺口 / 可能错误的实现 | 最小反例与关键断言 | 已完成的修改 / 测试 | 当前结论 | -| --- | --- | --- | --- | -| B65536直接进入backward grid.y | B65536,S1,H1；公开调用必须在launch前明确拒绝；B65535的y/gate/dX=0，dMod gate=1 | Python/C++新增B≤65535校验；`shared_bounds::rejects_batch_above_grid_limit`及`maximum_batch_forward_backward` | Python拒绝 **PROVEN**（meta及SM86 CUDA tensor）；C++产物与合法最大B数值 **UNPROVEN**，测试已准备 | -| x地址或seq narrowing以int溢出，先连续化会尝试几GiB分配 | meta x形状[1,524289,4096]、[1,2^31,1]；必须报int32 indexing错误，无实数据分配 | Python/C++检查x和modulation numel≤2^31-1；`rejects_int32_address_overflow_without_allocating` | Python metadata拒绝 **PROVEN**；C++重新构建后的运行 **UNPROVEN** | -| double eps有限但FP32为Inf/0 | 1e39、1e-46必须明确ValueError；1e-40、2^-149、FP32最大有限值必须保持y0、dX=±rstd、dMod=[1,-1,0,0,0,0] | 共享Python校验用stdlib struct检查FP32舍入，C++也检查；公开`eps_overflow`及扩展`subnormal_eps` literal测试，两dtype | 三Python公开后端拒绝 **PROVEN**；合法边界CPU/SM86 Triton **PROVEN**；native合法eps数值 **UNPROVEN** | -| 把nonfinite LN误当finite验收；gate/VJP因NaN LN串值 | x=[NaN,1]或[Inf,1]，y/dX/dscale应NaN；dshift=[1,-1]、gate=[2,-0]、dgate=[2,-3]独立保持 | 声明nonfinite传播和qualification范围，新增`nonfinite_activation_propagates_nan_and_preserves_gate_vjp` | CPU/SM86 Triton指定反例 **PROVEN**；一般NaN payload与strict非有限域不在资格声明中 | -| 只单独比较geometry或relocation，错过组合路径 | H3072：默认B1,S3基准→B3每个position、S3/5、strided x和6H调制；y/gate/dX/dMod bytes一致、unused半区0 | `h3072_relocation_with_all_launch_configurations`覆盖全部6种Triton warps/tile及3种CUDA threads，两dtype | 全部Triton组合 **PROVEN**（12参数例）；native组合测试已写、运行 **UNPROVEN** | -| GPU第二triple/两stream不独立；只验梯度会漏错shift | image S5/text S3，各取前/后3H；另一个stream梯度不变、unused半区0；错误候选y+1仍有正确VJP | 原`gate_only_and_separate_stream_chunks`扩展到三后端，独立CPU gold的forward+backward、gate bytes及另一stream隔离；错误shift候选旧断言实际放过，新断言拒绝 | CPU/SM86 Triton **PROVEN**（8参数例，均遍历两stream）；native测试已写、运行 **UNPROVEN** | -| H上限或零填充off-by-one | H4095/4096：两dtype y/gate/dX/dMod对独立CPU数学及固定参考bytes；H4097直接拒绝、registry明确fallback | 新`hidden_boundary_matches_independent_math_and_fixed_bytes`及`shared_bounds` direct/registry三个H边界测试 | CPU/SM86 Triton合法H数学 **PROVEN**；Python direct拒绝及registry **PROVEN**；native合法H数值 **UNPROVEN** | -| 用旧_C或语义fingerprint证明当前C++ | 当前CUDA源码已有校验变更，现存_C未重建，不能因symbols/constructor可用就算通过 | 记录新source hash与旧binary hash，准备fresh build/验收命令，未调用未绑定当前源码的native数值路径 | 当前native内核资格 **UNPROVEN** | -| 三真实图像S、shared race、OOB、claimed平台性能缺当前证据 | S4096/6889/6032,H3072、两dtype、两个GPU实现；racecheck/memcheck必须0 errors；benchmark完整数据 | 更新04验收脚本、06racecheck脚本及WS1/CPU CI入口接入新测试；准备下面H100命令，未启动CI或长测试 | H100/真实图像/sanitizer/benchmark **UNPROVEN** | - -当前合法eps的手算rstd literals由独立stdlib math+IEEE FP32舍入得到，分别为 -`1.000002658868784e20`、`2.671373844909537e22`、`5.421011508662376e-20`。 -隐藏边界的数学参考使用独立CPU FP32 leaves；固定CPU参考只作为字节/归约合同参考， -不将其当独立数学gold。三种比较不能互相替代。 - -### 当前版本运行表 - -下面最终suite和错误候选重跑对应本节末尾hash；red/局部green是本次TDD的中间源码证据。 -旧 `.ninja_log` 只用于评估构建耗时：full build历史上超过17分钟，不作为正确性运行证据。 -没有自行启动full build，也没有为了构建安装/下载依赖或索取新权限。 -本机GPU检查沿用已有Python执行权限，NVIDIA RTX3050Ti/SM86、PyTorch2.12.1+cu126、 -Python3.12.12、Triton3.7.1；未声明SM90或当前native binary资格。 - -| 本次检查 | 结果 / evidence | -| --- | --- | -| B限制red→green | meta公开入口1 failed→1 passed；[red](/tmp/rlk-adaln-close-batch-red.log)、[green](/tmp/rlk-adaln-close-batch-green.log)；[SM86公开拒绝](/tmp/rlk-adaln-close-batch-gpu-green.log)1 passed。没有故意启动非法native launch | -| int32范围red→green | 2 failed→bounds 3 passed；[red](/tmp/rlk-adaln-close-int32-red.log)、[green](/tmp/rlk-adaln-close-int32-green.log)，全是小输入/meta，无巨大分配 | -| eps范围red→green | CPU4 failed→4 passed；[red](/tmp/rlk-adaln-close-eps-red.log)、[green](/tmp/rlk-adaln-close-eps-green.log)；[合法/非法eps小形状CPU+Triton](/tmp/rlk-adaln-close-eps-boundaries.log)20 passed | -| 第二triple错shift断言red→green | [错误候选旧断言SURVIVED，exit1](/tmp/rlk-adaln-close-shift-mutant-red.log)→[新断言REJECTED，exit0](/tmp/rlk-adaln-close-shift-mutant-green.log)，不是生产算子数值错误的声明 | -| 全Triton geometry×topology | [log](/tmp/rlk-adaln-close-all-triton-launches.log)12 passed/9.68s；全部6配置×2dtype，每例含B3的每个position与S3/5 | -| 两stream/两triple | [log](/tmp/rlk-adaln-close-streams-green.log)8 passed/5.43s；含独立forward及backward | -| H4095/4096合法边界数学与bytes | [log](/tmp/rlk-adaln-close-hidden-math.log)8 passed/25.65s | -| B/H/寻址公开preflight及registry | [log](/tmp/rlk-adaln-close-hidden-preflight.log)8 passed/8.73s；未运行maximum_batch数值例 | -| nonfinite传播/独立gate VJP | [log](/tmp/rlk-adaln-close-nonfinite.log)8 passed/5.15s；只断言NaN传播，未将NaN算术payload计为byte资格 | -| 最终CPU/checker/preflight | [log](/tmp/rlk-adaln-close-final-cpu.log)：**113 passed，139 deselected，10.84s，0 skip/xfail** | -| 最终CPU+SM86 Triton+CUDA公开拒绝 | [log](/tmp/rlk-adaln-close-final-gpu.log)：**180 passed，91 deselected，11.00s，0 skip/xfail**；排除native数值、真实图像、旧完整geometry/sanitizer | -| 最终错误候选重跑 | [6个原候选](/tmp/rlk-adaln-close-mutants-six-final.log)、[8个后续候选](/tmp/rlk-adaln-close-mutants-eight-final.log)、[新shift候选](/tmp/rlk-adaln-close-shift-mutant-final.log)：**15/15拒绝**；[新候选脚本](/tmp/rlk_adaln_close_shift_mutant.py) | - -静态检查通过：7个相关Python文件ruff，新增/本次修改文件isort及Black检查， -三个修改shell的`bash -n`、CPU CI YAML解析、`git diff --check`。 - -```bash -# 当前CPU命令：bounds是metadata/preflight，不运行native内核 -PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 .venv/bin/python -m pytest \ - tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py \ - tests/test_op_checks.py tests/test_operator_inputs.py -q -rs \ - -k '(not triton and not cuda) or (shared_bounds and not maximum_batch)' - -# 当前SM86短检查。eps_overflow和shared_bounds仅公开拒绝/registry，不调用native数值路径 -PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 RL_KERNEL_REQUIRE_EXT=1 \ -TRITON_CACHE_DIR=/tmp/rlk-adaln-review-triton .venv/bin/python -m pytest \ - tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py tests/test_adaln_modulation.py \ - -q -rs -k '(not cuda and not real_image and not gpu_padding and not launch_geometry) or eps_overflow or (shared_bounds and not maximum_batch)' -``` - -### 当前源码 SHA256 - -| 文件 | SHA256 | -| --- | --- | -| 未跟踪 `tests/test_adaln_gtest.py` | `db9617770d5637a4aa35e4d0688bb3e3e3abdb7c127f5d04288ab12ce6237119` | -| 未跟踪 `tests/test_adaln_shared_bounds.py` | `97c5a99ae6f73545304114786e935d30f21af8a4bbe73984c296972ae997b0ba` | -| CUDA wrapper | `e5a52b8e29c6be131176e0fdad64aa259cc780b5a4c7529bad98901b7ec53797` | -| PyTorch AdaLN | `96707e412a15790962a15a939ecc835eaf363829c515304a2467db2a294fb5c6` | -| Triton AdaLN | `ef2b13c8e13da8d324b160489efad4b2080d3a1058cf3b58d566c2bc92a3b4a2` | -| 当前 CUDA source | `aae87fb84f9b41954a727aa09e926ddfc4deb7c82130aca508607248d0670c05` | -| gtest checker | `70f9c648b5da5ace001d295569c4b05a15d2e630e84e1870b5734a5377d76b39` | -| gtest specs | `b78d15322199b4986f19697c7987dea879238b7d8b14fe2f1b05fc3fe0268e79` | -| 现存旧_C（没有本次构建来源证明） | `d7b91b28b81c7f53db47e757eab9305f617dbc5649666daa3237e9946eb350b6` | - -### 待用户执行的native/H100检查，未启动 - -下方历史报告的H100构建命令仍是待执行方案；必须先从当前dirty/untracked源码重新构建， -不能复用上表旧_C。源码快照要包含两个未跟踪测试及本文件。建议保存build log、完整 -compiler flags、源码/产物hash、Python/PyTorch/Triton/CUDA/driver/GPU信息，再执行： - -```bash -# 仅在当前源码已fresh build的准备好环境中运行；这里未执行 -RL_KERNEL_REQUIRE_EXT=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ - .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py \ - -q -rs -k cuda - -# 完整真实形状与数值入口；04已包含新bounds，06包含完整native launch×topology -bash scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh - -# 支持域已限定，巨大输入应在入口拒绝；memcheck检查真实支持边界，不分配4GiB反例 -RL_KERNEL_REQUIRE_EXT=1 PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ - "${CUDA_HOME:?set CUDA_HOME}/bin/compute-sanitizer" --tool memcheck --error-exitcode 7 \ - .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py \ - -q -rs -p no:cacheprovider \ - -k 'cuda and (hidden_boundary or maximum_batch or all_launch_configurations)' - -bash scripts/adaln_cuda_acceptance/05_benchmark_adaln_cuda.sh -``` - -验收必须覆盖：native全部canaries、最大B65535与超限65536、H4095/4096/4097、三threads、 -两dtype、两stream/两triple；Triton全六配置；S4096/6889/6032真实图像；racecheck/memcheck -0 hazards/errors；benchmark有每个dtype/S/backend的forward及forward+backward数据。 -要求零skip/xfail、native无fallback，source/binary在验收期间不变；配置和CLI报告不是运行证据。 -一般NaN算术payload、ROCm或其他未固定Triton环境不在本机资格声明中。 - -## 第一轮审查及运行记录（历史源码快照） - -> 下方内容及[第二轮报告](adaln-shared-b31fac9-followup.md)保留中间源码证据。 -> 它们的运行表和hash不代表当前工作区；最新状态以上面“本次TDD”章节为准。 -> H100/native构建与长验收没有在本次执行。 - -规格：[RL-Align/RL-Kernel #386](https://github.com/RL-Align/RL-Kernel/issues/386)。 -比较点 `b31fac988faae9de71ab1f683e4a5c9cc8578b84`，HEAD -`120f4de7274798c6310d87799ebb6538f2f8f531`。范围是 **比较点到实际工作区**， -使用 `git diff b31fac9` 加 `git ls-files --others --exclude-standard`，包括索引、 -未提交修改和未跟踪 `tests/test_adaln_gtest.py`。三点 HEAD diff 不代表本次范围。 - -开始时工作区有 11 个已修改文件和上述未跟踪文件；其既有修改均保留。 -已读比较点之后五个提交、CUDA/PyTorch/Triton 实现、bindings、registry、 -gtest reference/inputs/checker、两个测试文件、benchmark、acceptance scripts、CI 和文档。 -indexed select01、projection、gate residual、最终 AdaLN、完整 MMDiT/WS2 不属于本次 shared 实现验收。 - -这里 **PROVEN** 只表示明确列出的当前版本、后端与有限反例已经运行通过； -**PARTIAL** 表示有实现或部分证据，仍有未运行的平台/几何/验收项； -**UNPROVEN** 表示没有当前工作区的有效运行证据。没有用旧日志填补当前证据。 - -## Standards - -1. **P1，已修复：BF16 gtest 参考梯度被再次舍入成 BF16。** - [gtest 规范](../contributing/gtest-usage.md)要求独立 FP32 reference。 - `GtestAdaLNReference.forward_fp32` 虽在内部转 FP32，但旧 checker 给 gold 克隆的 - 求导叶子仍为 BF16；autograd 返回 BF16 dX/dMod，报告却标 `fp32_reference`。 - 新增公开 `run_operator_suite` 断言所有 output/gradient checks 的 `gold_dtype` - 都为 `torch.float32`；BF16 H=7 的运行实际失败（1 failed、1 passed，exit 1）。 - 修复为 `OperatorSpec/OperatorCase.gold_input_dtype` 显式启用 FP32 gold leaves； - 只转换浮点张量，candidate 和其他算子的默认输入策略保留。随后相关 47 项检查通过。 - 这修复的是验证参考，未假称发现 CUDA 算术错误。 -2. **判断性异味：registry 的两条 trace 分支重复若干精度/算法字段。** - 可能导致未来元数据漂移；目前无实际错误证据，没有为此增加抽象或重构。 - -另修复 focused shared CPU 测试无条件导入 Triton 的问题，并将两个 GPU 测试的 -Triton import 移到 CUDA 可用检查之后。CPU 原本不需要这个可选依赖。 - -Standards：1 个已修复规范问题、1 个判断性异味；最严重项是参考梯度精度声明错误。 - -## Spec - -原有 dtype/finite/容差断言不足以证明 BF16 输出边界；H=17/32 的 relocation/padding -也不能覆盖 H=3072 的批位置。现有 H=3072、B=2 geometry 测试比较的是各自 CPU -reference，没有把同一个样本从 B=1 移到 B=3 的每个位置。两项均已补公开入口反例。 -另有明确的独立 FP32 gold 梯度错误，已按上述 red→green 修复。 -shared 非法输入的反例也已补齐；当前 native binary、真实图像尺寸、racecheck、H100 -和当前 benchmark 仍没有本轮运行证据。 - -Spec：3 个核心精度/拓扑缺口、1 组非法输入覆盖缺口、1 组平台验收证据缺口； -最严重项是原 gtest 的 FP32 gold 梯度并非 FP32。没有合并或重排 Standards 结论。 - -## 要求、独立参考、反例与关键断言 - -位置简称：`native` = `ops/pytorch/norm/adaln_modulation.py`；`cuda` = -`csrc/cuda/adaln_modulation.cu`；`triton` = `ops/triton/norm/adaln_modulation.py`。 -新测试均在 [test_adaln_gtest.py](../../tests/test_adaln_gtest.py)，旧断言保留在 -[test_adaln_modulation.py](../../tests/test_adaln_modulation.py)。 - -| 要求 / 当前实现与参考 | 具体可能错误的实现 | 最小反例 | 应失败的断言 / 测试 | 当前结论 | -| --- | --- | --- | --- | --- | -| 无 affine LayerNorm，shift/scale/gate 顺序，默认 eps=1e-6；三后端均按此实现；独立 CPU `F.layer_norm` + autograd | 用 RMSNorm，交换 shift/scale，gate 混入 y | x=[1,3]、shift=[.5,-.5]、scale=[1,0]、gate=[2,3]、eps=0；y=[-1.5,.5] | 旧定值 forward/backward；独立 shared H=7/3072、非均匀 dy/dg 的 accuracy 断言 | **PROVEN** CPU/SM86 Triton 的已跑小形状；CUDA **UNPROVEN** | -| y、gate 及 dX/dMod 独立 FP32 CPU reference；现已启用 FP32 gold leaves | 内部转 FP32但向 BF16 叶子求导，或求导后再 `.float()` 掩盖已舍入的参考 | BF16 B=2,S=3,H=7 gtest case | 所有四项 check 的 `gold_dtype == torch.float32`；独立 FP32 leaves 的 accuracy；旧错误实际 red、新 checker green | **PROVEN** 当前 CPU checker；native 平台验收另计 | -| BF16 normalization 中间量不得提前舍入 | 把 norm 转 BF16 再做 scale/shift | BF16 x=[-1,1]、shift=-1、scale=0、默认 eps；正通道应为 -2^-21，早舍入为 0 | `bf16_forward_casts_only_at_output` 对手算 literal 做 uint8 bytes 比较 | **PROVEN** CPU/SM86 Triton 该反例；完整内部精度政策 **PARTIAL** | -| `1+scale` 和 affine 保持 FP32 到输出边界 | BF16 的 `1+1/256` 先舍入成 1 | 同上，scale=1/256；正通道 BF16 应为 1/256，错误结果约为 -2^-21 | 同一 forward canary 的第二个参数例；不能用容差替代 | **PROVEN** CPU/SM86 Triton 该反例；CUDA **UNPROVEN** | -| backward dnorm、partial、token accumulators FP32 | BF16 dnorm、BF16 partial，或逐 token BF16 fold | B=1,S=3,H=2，x 每行[-3,3]，eps=7，scale=1/256，dy 首通道[256,1,-256] | `backward_keeps_partials_and_accumulators_fp32`：dshift=[1,0]，dscale=[-.75,0]；dX 按独立手算系数1799/32768；gate VJP=[-2,.5] | **PROVEN** CPU/SM86 Triton 该反例；CUDA **UNPROVEN** | -| BF16 backward 保存的 norm/rstd/partials 不能成为隐含额外 cast | 只在 BF16 backward 保存量转 BF16，forward 完全正确 | 相同 BF16 量化输入、dy/dg，H=7/3072，分别以 BF16 和提升后的 FP32 调用公开 op | `bf16_vjp_matches_fp32_with_one_final_cast`：y/gate/dX/dMod 必须等于 FP32 调用结果最终 cast 的 bytes；独立 math 另测，避免把 differential 当独立数学 gold | **PROVEN** CPU/SM86 Triton 的固定 seed；更广输入域 **PARTIAL** | -| H 固定 lower+upper pairwise tree，非二次幂零填充 | 改成 ascending left fold、按线程分组求和 | FP32 H=4 x=[2^24,1,-2^24,1]，pairwise mean=.5，left-fold mean=.25 | `hidden_reduction_uses_declared_pairwise_order` 的手算 FP32 literal bytes；旧 H=7/3072 accuracy/byte tests | **PROVEN** CPU/SM86 Triton 该反例；其他重排与 CUDA **PARTIAL** | -| dMod 按逻辑 token 递增 FP32 fold；无 atomic/Split-K/Stream-K | 用 pairwise token tree、调换 token/分块重关联 | S=4,H=2 x=[-1,1]，dy 首通道[2^24,1,-2^24,1] | `modulation_gradient_uses_ascending_token_order`：dshift=1、dscale=-1；改树可能得到2/-2 | **PROVEN** CPU/SM86 Triton 该反例；禁止算法同时有源码证据，native **UNPROVEN** | -| batch size、batch position、padding 下有效样本 bytes 不变 | H=3072 路径始终读 mod[0]、错误 batch stride、跨样本共享 dMod | B=1,S=3,H=3072 基线→B=3 全部位置；S=3/5；其他样本不同 x/m/dy/dg，padding dy=0 | `h3072_sample_bytes_survive_batch_positions_and_padding`：y/gate/dX/dMod uint8 相等；包含非连续 x 和六块 view；未使用调制块梯度为0 | **PROVEN** CPU/SM86 Triton BF16/FP32 本次矩阵；更大 B/S 与 CUDA **PARTIAL** | -| launch geometry / tiling 不改 bytes | H tree 随 CUDA threads、Triton warps/tile 改变；跨 warp 共享内存未同步 | H=3072,B=2,S=5，同一输入换 threads128/256/512，warps4/8、tile64/128/256 | 旧 `launch_geometry_preserves_bytes` 比较 y/gate/dX/dMod 对固定 CPU reference | 测试已存在，本轮未运行 geometry sweep；**UNPROVEN** 当前完整 sweep | -| 两 stream 独立调制、attention/MLP 两 triples、gate downstream gradient、strided 模型 view | 串用 image/text modulation；用错第二 triple；gate detached；漏掉 contiguous 复制 | image S=5、text S=3，各自 B=2,H=7、6H 调制；分别取前/后3H，非均匀 dy/dg 和 gate-only | `gate_only_and_separate_stream_chunks` 已扩展到两 triples，独立 CPU accuracy、未使用半区0、另一 stream 无梯度；H3072 topology 的 strided view | **PROVEN** CPU 两 stream/两 triples；SM86 Triton H3072 strided 第一 triple **PROVEN**；GPU 第二 triple **PARTIAL** | -| unsupported shared 输入清楚拒绝；CUDA H<=4096，几何参数有界 | empty/mixed shape 被 launch；FP16 悄悄进入；NaN eps 未拒绝；H>4096 native误选 | 任一空 B/S/H、rank2、错3H、mixed dtype/device、FP16、eps<0/NaN/Inf | 新 `shared_rejects_unsupported_inputs` 和 `shared_rejects_mixed_devices`；旧 large-H registry test；constructors 源码验证 launch 参数 | **PROVEN** CPU/SM86 Triton 输入矩阵；CUDA H/geometry 的当前 runtime **PARTIAL** | -| 实际 backend/fallback、算术 policy 与 kernel identity 可追踪 | native缺失但谎称CUDA；将语义version当编译binary hash；fallback算作strict成功 | native symbol missing、Triton missing、所有后端missing | 旧/新 registry resolution tests；当前 preflight selected native但仅说明构造/符号可用；actual execution须结合direct candidate+source/binary manifest | **PARTIAL**；语义 fingerprint 不绑定二进制，不能凭它证明当前 native | -| CUDA shared reduction race-free | 读shared[0]后其他warp立即重用buffer | H=64 或3072、多warp fwd/bwd | `compute-sanitizer --tool racecheck --error-exitcode 7`；数值 bytes 相等不是 race assertion | **UNPROVEN** 当前 binary；旧 zero-hazard 日志不计 | -| 三参考图像尺寸与短 synthetic，forward/backward、两 dtype、两个GPU实现 | S=6889 tail越界、dMod 漏 token、仅注册优先backend正确 | S={4096,6889,6032},H=3072，随机 dy/dg | 旧 real-image tests 同时独立 FP32 CPU accuracy、fixed-reference bytes、CUDA↔Triton bytes | synthetic 已有当前证据；三 full-image 当前 **UNPROVEN**，H100命令如下 | -| 当前 claimed CUDA platform benchmark、pinned环境 | 只用旧binary/旧timings，CUDA skip却当benchmark完成 | 当前源码构建→两dtype三S，fwd及fwd+bwd | benchmark backend标签、无native skip、完整timings/peak memory、build/runtime/source manifest | 实现存在，当前运行 **UNPROVEN** | -| 无TF32、fast math、compiler-dependent重关联；unsupported严格配置拒绝 | 编译开启fast math仍注册native，或Triton融合乘加 | fast-math-on build / 不同编译配置 | CUDA RN intrinsics、宏#error、setup条件排除symbols；Triton enable_fp_fusion=False；必须验独立build/provenance，不用普通数值测试冒充compiler proof | **PARTIAL** 静态证据；替代build/H100 profile **UNPROVEN** | - -没有把原本正确的算子改成迎合测试的实现。补测时两个期望问题已纠正：大值 -hidden-tree 反例的标准化值应为 ±sqrt(2)，不是 ±1/sqrt(2)；padding 零 VJP 可为 -负零，规格要求有效行 bytes 不变，并未要求 padding 的零必须为正零。 -有效行、gate 和有效 dMod 的原始字节比较没有放宽。 - -## 本轮运行证据 - -所有结果来自本次工作区。源文件与测试的 SHA256 见下，未提交或推送。 - -| 运行 | 命令 / artifact | 结果 | -| --- | --- | --- | -| 初始 red | `pytest tests/test_adaln_gtest.py -q -x -k outputs_and_gradients` | BF16 gold_dtype 断言失败，exit 1；1 failed、1 passed、6 deselected | -| 修复后 checker 回归 | `pytest tests/test_adaln_gtest.py tests/test_op_checks.py tests/test_operator_inputs.py -q -rs`，当时尚未追加其余 canaries | 47 passed / 19.19s | -| 最终 CPU/checker | `PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_op_checks.py tests/test_operator_inputs.py -q -rs -k 'not triton and not cuda'`；[/tmp CPU log](/tmp/rlk-adaln-review-cpu-20261009.log) | 68 passed、41 deselected / 16.94s；无 skip | -| 最终 CPU+SM86 Triton | `PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 TRITON_CACHE_DIR=/tmp/rlk-adaln-review-triton .venv/bin/python -m pytest tests/test_adaln_gtest.py tests/test_adaln_modulation.py -q -rs -k 'not cuda and not real_image and not gpu_padding and not launch_geometry'`；[/tmp GPU log](/tmp/rlk-adaln-review-gpu-20261009.log) | 80 passed、48 deselected / 32.06s；无 skip。此选择不包含 native CUDA 或 geometry sweep | -| 断言敏感性 | [/tmp mutation probe](/tmp/rlk_adaln_review_mutations.py)，[/tmp log](/tmp/rlk-adaln-review-mutations-20261009.log) | 6/6 公共接口模拟错误被拒绝：early norm、early affine、BF16 token fold、token reorder、hidden left fold、错误 batch 调制 | -| CPU 可选依赖隔离 | [/tmp probe](/tmp/rlk_adaln_review_cpu_only.py)，[/tmp log](/tmp/rlk-adaln-review-cpu-portability-20261009.log)；拒绝所有 Triton imports，再运行 focused shared PyTorch independent cases | 4 passed、52 deselected / 3.44s | -| 静态检查 | ruff、isort、Black check：新增测试与 checker/spec；`git diff --check`；两个修改 shell 的 `bash -n` | 全部通过；Black 按 py310、line-length100 逐文件检查 | - -Mutation probe 用错误候选替换测试的公共 callable，不修改生产 helper，也不是硬件证据。 -独立 LayerNorm/autograd 检验数学正确性；同值 FP32/BF16 differential 检验 cast 边界； -literal 与 uint8 比较检验特定舍入/归约错误。三者分工，不能互相替代。 - -本机环境：Python 3.12.12、PyTorch 2.12.1+cu126、CUDA runtime 12.6、Triton 3.7.1， -RTX 3050 Ti Laptop / SM86。没有本轮 driver/toolkit build qualification 声明。 -native `_C` 已存在但没有与当前源码绑定的 build artifact，本轮没有执行其数值测试。 - -| 文件 | SHA256 | -| --- | --- | -| `tests/test_adaln_gtest.py`（未跟踪文件，已纳入） | `9fea7d7ae184b74641ef8368734a446169c05d7063259956eb42d7a5d28e77f2` | -| `gtest/op_checks.py` | `8b1ff607983d40d60c3a21e03302884228c75ec1beb56b20dccf25215f9c18c1` | -| `gtest/operator_specs.py` | `93c44ff298f2896de40c0b43d9fc5314ef377f0d7c2294cb2420826116ac3917` | -| PyTorch AdaLN | `682a32ffad1ecbbdb28b7bdbb25413acabdde51c71c784656e2d272e16b19003` | -| Triton AdaLN | `7b08d9b4bc047a342875040372bd128f2cbc66eb02779e1c9f9520160b28f781` | -| CUDA source | `27074706332b1db32236b6cbb225ee99514938f38d1996ef4812e221d857e80d` | -| 现存 `_C.cpython-312-x86_64-linux-gnu.so` | `d7b91b28b81c7f53db47e757eab9305f617dbc5649666daa3237e9946eb350b6`；仅记录身份，非当前 native 验收 | - -tracked diff 快照：[patch](/tmp/rlk-adaln-review-tracked-b31fac9.patch)； -未跟踪测试快照：[test source](/tmp/rlk-adaln-review-test_adaln_gtest.py)； -规格快照：[issue JSON](/tmp/rlk-adaln-review-issue386-20261009.json)。 -这些 `/tmp` 文件应随验收 artifact 保存，不能假设长期存在。 - -## H100：只准备命令与验证清单，未执行 - -在用户准备好的、依赖已经安装且版本已固定的 H100 checkout 中执行。以下包含可能长时间 -运行的构建、真实大图、racecheck 与 benchmark；本轮均没有启动。 -H100 native 必须重新构建，不能沿用本机 SM86 `.so` 或仅按符号可用判定成功。 - -```bash -set -euo pipefail -export ADALN_EVIDENCE=/tmp/adaln-h100-current-worktree -mkdir -p "$ADALN_EVIDENCE" -git rev-parse HEAD > "$ADALN_EVIDENCE/head.txt" -git status --short --untracked-files=all > "$ADALN_EVIDENCE/status.txt" -git diff b31fac9 --binary > "$ADALN_EVIDENCE/tracked.patch" -git ls-files --others --exclude-standard > "$ADALN_EVIDENCE/untracked-files.txt" -cp tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py "$ADALN_EVIDENCE/" -cp docs/reviews/adaln-shared-b31fac9*.md "$ADALN_EVIDENCE/" -sha256sum csrc/cuda/adaln_modulation.cu \ - rl_engine/kernels/ops/{pytorch,triton,cuda}/norm/adaln_modulation.py \ - rl_engine/kernels/gtest/{op_checks,operator_specs,operator_inputs}.py \ - rl_engine/kernels/registry.py tests/test_adaln*.py setup.py csrc/ops.cpp \ - > "$ADALN_EVIDENCE/sources.sha256" -nvidia-smi > "$ADALN_EVIDENCE/nvidia-smi.txt" -nvcc --version > "$ADALN_EVIDENCE/nvcc.txt" -.venv/bin/python -m pip freeze > "$ADALN_EVIDENCE/packages.txt" - -# 建议使用与项目兼容且已固定的 CUDA toolkit；不在这里安装/下载依赖。 -FORCE_CUDA=1 KERNEL_ALIGN_USE_FAST_MATH=0 KERNEL_ALIGN_FORCE_SM90=1 \ -TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=4 \ - .venv/bin/python setup.py build_ext --inplace --force \ - 2>&1 | tee "$ADALN_EVIDENCE/build.log" -sha256sum rl_engine/_C*.so > "$ADALN_EVIDENCE/binary.sha256" - -export RL_KERNEL_REQUIRE_EXT=1 -export PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 -export RLK_ADALN_REAL_SHAPES=1 -.venv/bin/python - <<'PY' 2>&1 | tee "$ADALN_EVIDENCE/preflight.log" -import sys, torch, triton -from rl_engine.kernels.registry import kernel_registry -assert torch.cuda.is_available() -assert torch.cuda.get_device_capability()[0] == 9 -assert "H100" in torch.cuda.get_device_name(), torch.cuda.get_device_name() -print(sys.version, torch.__version__, torch.version.cuda, triton.__version__) -print(torch.cuda.get_device_name(), torch.cuda.get_device_capability()) -_, trace = kernel_registry.get_adaln_modulation_op("cuda", hidden=3072) -assert trace["selected_backend"] == "CUDA_ADALN_MODULATION" and not trace["fallback"], trace -print(trace) -PY -.venv/bin/python -m pytest tests/test_extension_smoke.py \ - tests/test_adaln_modulation.py tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py \ - -q -rs -p no:cacheprovider 2>&1 | tee "$ADALN_EVIDENCE/tests.log" - -# 使用与 CUDA_HOME 相同且兼容的 sanitizer；检查退出码及完整 hazard/error summary。 -"${CUDA_HOME:?set CUDA_HOME}/bin/compute-sanitizer" --tool racecheck \ - --racecheck-report analysis --error-exitcode 7 \ - .venv/bin/python -m pytest tests/test_adaln_modulation.py \ - tests/test_adaln_gtest.py -q -rs -p no:cacheprovider \ - -k 'cuda and (bit_equality or launch_geometry or sample_bytes or backward_keeps)' \ - 2>&1 | tee "$ADALN_EVIDENCE/racecheck.log" - -for backend in cuda triton; do - for dtype in bf16 fp32; do - .venv/bin/python scripts/check_operator.py --op adaln_modulation \ - --candidate "$backend" --device cuda --dtype "$dtype" \ - --batch 3 --seq 5 --normalized-dim 3072 --check-grad --json \ - > "$ADALN_EVIDENCE/gtest-$backend-$dtype.json" - done -done -for dtype in bf16 fp32; do - .venv/bin/python benchmarks/benchmark_adaln_modulation.py \ - --real --dtype "$dtype" --warmup 3 --repeat 10 \ - 2>&1 | tee "$ADALN_EVIDENCE/benchmark-$dtype.log" -done -sha256sum -c "$ADALN_EVIDENCE/sources.sha256" -sha256sum -c "$ADALN_EVIDENCE/binary.sha256" -``` - -验收清单： - -- 保存 worktree patch、所有相关 untracked 源文件、源码 hash、binary hash、完整 build/test logs。 -- 源码/二进制在整个验收期间不变；确认 fast math disabled、SM90 构建、H100 实际执行。 -- focused + gtest + extension smoke 全部通过，0 skip/xfail；三 S、两 dtype、两个 GPU backend 不能漏。 -- BF16 forward 两个 literal、backward cancellation、FP32 single-cast VJP differential 全部通过。 -- H3072 B=1→B=3 的每个位置、S3/5、非连续视图、y/gate/dX/dMod bytes 全部通过。 -- CUDA 三 threads 与 Triton 六 warps/tile 配置的 geometry bytes 全部通过。 -- real-image independent CPU accuracy 用共享 forward/gradient tolerance；fixed-reference 与 backend bytes 另验。 -- gold 的两项 gradient dtype 为 FP32；不能只见 `report.passed` 或 dtype/finite 就宣称精度通过。 -- sanitizer exit 0，0 hazards/errors；有 race 报告即失败，即使数值断言通过。 -- native 没有 fallback/skip；benchmark 每个 dtype/S 都有 CUDA/Triton fwd/fwd+bwd 与 peak-memory 数据。 -- CLI JSON 是 operator debugging evidence；缺少受验证 provenance 时不能当整个 WS1 系统 EXIT。 - -GPU CI 与 script 04 现在也会运行新 canaries。CI 未启动，workflow 配置不是运行证据。 - - -## 本地完整验收流水线（2026-10-10，SM86） - -当前流水线状态：**failed**；H100/SM90仍为 **UNPROVEN**。 - -[完整状态与逐阶段命令](/tmp/rlk-adaln-local-acceptance-20261009T174946Z/status.json)、[固定环境及 binary hash](/tmp/rlk-adaln-local-acceptance-20261009T174946Z/manifest.json)、[源码 hash](/tmp/rlk-adaln-local-acceptance-20261009T174946Z/source.sha256)。 - -- correctness：exit=1；passed=n/a；[完整日志](/tmp/rlk-adaln-local-acceptance-20261009T174946Z/correctness.log)。 - -失败原因见[failure.log](/tmp/rlk-adaln-local-acceptance-20261009T174946Z/failure.log)，后续阶段未启动，不升级对应资格。 - - -## 真实形状暴露的独立参考错误（2026-10-10） - -首次完整运行381 passed、3 failed、14 deselected；三个FP32真实S均在dMod accuracy断言失败，fixed CPU与CUDA/Triton的bytes比较已经通过。旧独立参考使用CPU broadcast autograd默认S归约，该归约不遵守合同的token递增FP32 left-fold。 - -独立最小反例：B1,S8,H2，x每行[-1,1]、eps0、shift/scale0，dy首通道[2^24,1,-2^24,1]重复两次；合同dshift=1、dscale=-1，旧gold得到2、-2。关键断言是公开GtestAdaLNReference VJP对手算literal的bytes，而非仅互相比对两个算子。 - -[test-first red](/tmp/rlk-adaln-gold-token-order-red.log)1 failed→[green](/tmp/rlk-adaln-gold-token-order-green.log)1 passed。修复只在独立gold的shared affine使用独立CPU FP32 token fold，未调用任何生产normalization/backward helper；norm和dX仍由独立CPU F.layer_norm/autograd提供。indexed参考路径及所有原容差不变，native源码/.so未变。 - -当前最小gold归约反例 **PROVEN**；完整真实形状复跑和后续sanitizer/benchmark状态以后续流水线证据为准。 - - -## 本地完整验收流水线（2026-10-10，SM86） - -当前流水线状态：**succeeded**；H100/SM90仍为 **UNPROVEN**。 - -[完整状态与逐阶段命令](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/status.json)、[固定环境及 binary hash](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/manifest.json)、[源码 hash](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/source.sha256)。 - -- correctness：exit=0；passed=385；[完整日志](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/correctness.log)。 -- racecheck：exit=0；passed=210；[完整日志](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/racecheck.log)。 - ========= RACECHECK SUMMARY: 0 hazards displayed (0 errors, 0 warnings) -- memcheck：exit=0；passed=224；[完整日志](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/memcheck.log)。 - ========= ERROR SUMMARY: 0 errors -- benchmark：exit=0；passed=n/a；[完整日志](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/benchmark.log)。 - -54项 benchmark 覆盖2 dtype×3 S×3 backend×3 mode；[JSON含每次采样](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/benchmark/benchmark.json)、[CSV](/tmp/rlk-adaln-local-acceptance-20261009T175945Z/benchmark/benchmark.csv)。 - -mode分别为forward、backward_only（forward graph在计时区域外准备）、forward_backward。warmup1/repeat3，时间为GPU event median；保存allocated/reserved峰值及相对计时前的extra峰值。属于本机短采样，H100需重新采样。 - -本次可将三真实S、完整launch×batch/padding、所选sanitizer矩阵及benchmark的本机指定配置记为PROVEN；不代表一般输入域或SM90资格。源码和二进制在各阶段前后保持不变。 diff --git a/scripts/adaln_cuda_acceptance/01_install_cuda_12_6_wsl.sh b/scripts/adaln_cuda_acceptance/01_install_cuda_12_6_wsl.sh deleted file mode 100755 index 5a492deff..000000000 --- a/scripts/adaln_cuda_acceptance/01_install_cuda_12_6_wsl.sh +++ /dev/null @@ -1,19 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -grep -qi microsoft /proc/version || { echo "WSL2 is required" >&2; exit 1; } -if [[ -x /usr/local/cuda-12.6/bin/nvcc ]] && - /usr/local/cuda-12.6/bin/nvcc --version | grep -q 'release 12\.6'; then - echo "CUDA Toolkit 12.6 is already installed" - exit 0 -fi - -read -r -p "Install the large cuda-toolkit-12-6 package (no Linux driver)? [y/N] " reply -[[ "$reply" =~ ^[Yy]$ ]] || exit 0 - -deb=$(mktemp --suffix=.deb) -trap 'rm -f "$deb"' EXIT -curl -fL https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-keyring_1.1-1_all.deb -o "$deb" -sudo dpkg -i "$deb" -sudo apt-get update -sudo apt-get install -y cuda-toolkit-12-6 diff --git a/scripts/adaln_cuda_acceptance/02_verify_cuda_12_6.sh b/scripts/adaln_cuda_acceptance/02_verify_cuda_12_6.sh deleted file mode 100755 index 8b18e3994..000000000 --- a/scripts/adaln_cuda_acceptance/02_verify_cuda_12_6.sh +++ /dev/null @@ -1,17 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) -cuda_home=${CUDA_HOME:-/usr/local/cuda-12.6} -nvcc="$cuda_home/bin/nvcc" - -[[ -x "$nvcc" ]] || { echo "Missing $nvcc" >&2; exit 1; } -"$nvcc" --version | grep 'release 12\.6' -[[ -e /usr/lib/wsl/lib/libcuda.so ]] || { echo "Missing WSL libcuda.so" >&2; exit 1; } -"$root/.venv/bin/python" - <<'PY' -import torch -assert torch.version.cuda == "12.6", torch.version.cuda -assert torch.cuda.is_available() -assert torch.cuda.get_device_capability() == (8, 6) -print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name()) -PY diff --git a/scripts/adaln_cuda_acceptance/03_build_extension.sh b/scripts/adaln_cuda_acceptance/03_build_extension.sh deleted file mode 100755 index 6ad635421..000000000 --- a/scripts/adaln_cuda_acceptance/03_build_extension.sh +++ /dev/null @@ -1,24 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) -cd "$root" -export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.6} -export PATH="$root/.venv/bin:$CUDA_HOME/bin:$PATH" -export LIBRARY_PATH="/usr/lib/wsl/lib:${LIBRARY_PATH:-}" -export LD_LIBRARY_PATH="/usr/lib/wsl/lib:$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}" -export TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST:-8.6} -export MAX_JOBS=${MAX_JOBS:-2} -export FORCE_CUDA=1 -export RL_KERNEL_REQUIRE_EXT=1 -export KERNEL_ALIGN_USE_FAST_MATH=0 - -uv pip install --python .venv/bin/python ninja -ninja --version -.venv/bin/python setup.py build_ext --inplace -.venv/bin/python - <<'PY' -import rl_engine._C as ext -for name in ("adaln_modulation_forward", "adaln_modulation_backward"): - assert hasattr(ext, name), name -print("AdaLN CUDA extension symbols verified") -PY diff --git a/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh b/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh deleted file mode 100755 index 684fb31d3..000000000 --- a/scripts/adaln_cuda_acceptance/04_test_adaln_cuda.sh +++ /dev/null @@ -1,26 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) -cd "$root" -export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.6} -export LD_LIBRARY_PATH="/usr/lib/wsl/lib:$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}" -export RL_KERNEL_REQUIRE_EXT=1 -export RLK_ADALN_REAL_SHAPES=1 -export PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 - -.venv/bin/python - <<'PY' -import sys -import torch -import triton -assert torch.version.cuda == "12.6", torch.__version__ -print(sys.version.split()[0], torch.__version__, torch.version.cuda, triton.__version__) -from rl_engine.kernels.registry import OpBackend, kernel_registry -_, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=3072) -assert trace["selected_backend"] == OpBackend.CUDA_ADALN_MODULATION.name, trace -assert trace["fallback"] is False, trace -print(trace) -PY -.venv/bin/python -m pytest tests/test_extension_smoke.py tests/test_adaln_modulation.py \ - tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py -q -rs -p no:cacheprovider -bash scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh diff --git a/scripts/adaln_cuda_acceptance/05_benchmark_adaln_cuda.sh b/scripts/adaln_cuda_acceptance/05_benchmark_adaln_cuda.sh deleted file mode 100755 index f46294cbf..000000000 --- a/scripts/adaln_cuda_acceptance/05_benchmark_adaln_cuda.sh +++ /dev/null @@ -1,15 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) -cd "$root" -.venv/bin/python - <<'PYCODE' -from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp -CudaAdaLNModulationOp() -print("AdaLN CUDA extension verified") -PYCODE -warmup=${WARMUP:-1} -repeat=${REPEAT:-3} - -.venv/bin/python benchmarks/benchmark_adaln_modulation.py --real --dtype bf16 --warmup "$warmup" --repeat "$repeat" -.venv/bin/python benchmarks/benchmark_adaln_modulation.py --real --dtype fp32 --warmup "$warmup" --repeat "$repeat" diff --git a/scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh b/scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh deleted file mode 100644 index 9018ea9cc..000000000 --- a/scripts/adaln_cuda_acceptance/06_racecheck_adaln_cuda.sh +++ /dev/null @@ -1,24 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) -cd "$root" -export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-12.6} -export LD_LIBRARY_PATH="/usr/lib/wsl/lib:$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}" -export RL_KERNEL_REQUIRE_EXT=1 -export PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 -export PYTHONDONTWRITEBYTECODE=1 - -.venv/bin/python - <<'PY' -import torch -assert torch.version.cuda == "12.6", torch.__version__ -print(torch.__version__, torch.cuda.get_device_name(), flush=True) -PY - -# Numerical equality alone does not detect shared-memory read/write hazards. -exec "$CUDA_HOME/bin/compute-sanitizer" --tool racecheck \ - --racecheck-report analysis --error-exitcode 7 \ - .venv/bin/python -m pytest -q -rs -p no:cacheprovider \ - tests/test_adaln_modulation.py::test_cuda_adaln_bit_equality_harness \ - tests/test_adaln_modulation.py::test_adaln_launch_geometry_preserves_bytes \ - tests/test_adaln_gtest.py::test_adaln_h3072_relocation_with_all_launch_configurations -k cuda "$@" From 0988fab3730352f0a624550d036fac1f377659c9 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Mon, 12 Oct 2026 17:07:25 +0800 Subject: [PATCH 10/37] Implement strict dispatch for AdaLN modulation and enhance backend observation - Introduced strict dispatch and execution observations for AdaLN modulation, enforcing arithmetic profiles and backend qualifications. - Added `AdaLNDispatch` class to manage execution traces and backend information. - Updated `KernelRegistry` to support strict mode and handle arithmetic profiles. - Enhanced backend provenance tracking in GTest to validate execution against declared profiles. - Created tests for strict dispatch behavior, including validation of backend observations and error handling for unsupported profiles. - Added benchmark tests for CPU checks and CLI contract validation. - Refactored existing tests to ensure compatibility with new dispatch and backend observation mechanisms. Signed-off-by: hongleng-github --- benchmarks/benchmark_adaln_modulation.py | 220 ++++++++++++-- docs/operators/adaln-modulation.md | 62 ++++ rl_engine/contracts/numerical.py | 27 +- rl_engine/kernels/adaln_contract.py | 101 +++++++ rl_engine/runtime/registry.py | 30 +- rl_engine/validation/operators/op_checks.py | 68 ++++- .../validation/operators/operator_specs.py | 15 +- tests/test_adaln_dispatch.py | 274 ++++++++++++++++++ tests/test_adaln_gtest.py | 2 +- tests/test_adaln_modulation.py | 10 +- tests/test_benchmark_adaln_modulation.py | 28 ++ tests/validation/operators/test_op_checks.py | 17 ++ tools/validation/operators/check_operator.py | 5 +- 13 files changed, 809 insertions(+), 50 deletions(-) create mode 100644 rl_engine/kernels/adaln_contract.py create mode 100644 tests/test_adaln_dispatch.py create mode 100644 tests/test_benchmark_adaln_modulation.py diff --git a/benchmarks/benchmark_adaln_modulation.py b/benchmarks/benchmark_adaln_modulation.py index b09f06bed..1efd2a505 100644 --- a/benchmarks/benchmark_adaln_modulation.py +++ b/benchmarks/benchmark_adaln_modulation.py @@ -7,12 +7,30 @@ """ import argparse +import contextlib +import hashlib +import json +from pathlib import Path import platform import statistics +import subprocess +import sys import torch -from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp +REPO_ROOT = Path(__file__).resolve().parents[1] +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): @@ -29,6 +47,7 @@ def run(): 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) @@ -37,46 +56,187 @@ def run(): end.record() end.synchronize() samples.append(start.elapsed_time(end)) - return statistics.median(samples), torch.cuda.max_memory_allocated() + return { + "median_ms": statistics.median(samples), + "samples_ms": samples, + "peak_allocated_bytes": torch.cuda.max_memory_allocated(), + "allocated_before_samples_bytes": allocated_before, + } -def main(): +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") - args = parser.parse_args() + 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") - from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp - from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + # Imports and extension diagnostics must not corrupt the JSON stdout stream. + with contextlib.redirect_stdout(sys.stderr): + from rl_engine.kernels.adaln_contract import SHARED_PROFILE, AdaLNDispatch + from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp - backends = [("pytorch", NativeAdaLNModulationOp()), ("triton", TritonAdaLNModulationOp())] - try: - backends.append(("cuda", CudaAdaLNModulationOp())) - except RuntimeError: - print("cuda extension unavailable: explicit skip") - dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 - print( - f"GPU={torch.cuda.get_device_name()} cuda_runtime={torch.version.cuda} " - f"torch={torch.__version__} python={platform.python_version()} " - f"dtype={args.dtype} warmup={args.warmup} repeat={args.repeat}" - ) - sizes = (4096, 6889, 6032) if args.real else (64,) - for seq in sizes: - 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) - for name, op in backends: - for backward in (False, True): - ms, peak = _bench( - op, x, modulation, backward=backward, warmup=args.warmup, - repeat=args.repeat, - ) - print( - f"S={seq} backend={name} backward={backward} " - f"median_ms={ms:.3f} peak_MB={peak / 2**20:.1f}" + backends = [("pytorch", NativeAdaLNModulationOp())] + skipped = [] + try: + from rl_engine.kernels.ops.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.kernels.ops.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/benchmark_adaln_modulation.py", + "rl_engine/kernels/adaln_contract.py", + "rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py", + "rl_engine/kernels/ops/cuda/norm/adaln_modulation.py", + "rl_engine/kernels/ops/triton/norm/adaln_modulation.py", + "rl_engine/kernels/ops/vjp_fp32.py", + "csrc/cuda/adaln_modulation.cu", + "csrc/ops.cpp", + "setup.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__": diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 65584218c..534626a48 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -103,3 +103,65 @@ uses those defaults. These settings change scheduling, not the reduction order. Native acceptance requires a fresh extension build, source/binary identity records, and complete test and sanitizer logs. + +## Strict dispatch and execution observations + +```python +op, trace = kernel_registry.get_adaln_modulation_op( + x.device, hidden=x.shape[-1], strict=True, + arithmetic_profile="adaln-shared-fp32-tree-v1", +) +y, gate = op(x, modulation) +``` + +Strict dispatch pins the shared FP32 hidden tree, ascending token fold and final +cast. NVIDIA CUDA may fall back to Triton when the extension is unavailable or +H exceeds 4096: both implement this same arithmetic profile, accept FP32/BF16, +and retain the existing launch-configuration replay tests. Strict rejects unknown +profiles/backends, the indexed/select01 framework backward and PyTorch execution +on GPU, whose cross-backend replay is not qualified here. It checks the actual +input device and hidden size again at execution. CPU Native is permitted for the +fixed reference. This policy does not itself certify any hardware/runtime profile. +Ordinary dispatch retains the prior CUDA/Triton/PyTorch and indexed fallbacks. + +The resolver still returns a callable and a trace; the callable is now a +resolution-local `AdaLNDispatch`, exposing its backend as `implementation`. +`selected_backend` is a resolution decision; `actual_backend`, execution dtype, +device and output dtype are filled only after a successful forward call. A failed +call clears the previous observation. `fallback_reason` and `arithmetic_profile` +identify routing and arithmetic separately. The trace describes the last call; +copy it when retaining an observation, and use separate resolutions for concurrent +callers. `observation="successful_forward"` is evidence of the invoked operator +path, not a profiler measurement or independent proof of backward execution. +`profiler_observed=false` and `hardware_qualified=null` make that limit explicit. +Semantic versions/fingerprints are not binary hashes. + +`check_operator.py` observes AdaLN candidates and includes per-case backend traces. +Successful BF16 CUDA/Triton calls also populate the existing C1 +`BackendProvenance`; CPU and FP32 have no declared C1 backend profile and remain +trace-only. The shared tolerance contract and independent reference are unchanged. +Library callers can request this with `make_candidate(args, observe_backend=True)`. + +The benchmark accepts `--seed` (default 386) and `--json`. JSON includes raw event +samples, median, absolute peak allocated memory and 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 that file; +it does not assert that it was built from the recorded working tree. Missing +observations are null. Extension/Triton skips are explicit. +Backend observations happen before timing; the original event interval still +measures the bare operator, including newly allocated `ones_like` upstreams for +forward+backward. The added ordinary `F.layer_norm` + FP32 affine baseline uses +the same inputs and timing, but does not claim fixed-order replay. + +Run the new GPU cases explicitly on the target environment: + +```bash +.venv/bin/python -m pytest tests/test_adaln_dispatch.py -k gpu -q -rs +.venv/bin/python benchmarks/benchmark_adaln_modulation.py \ + --real --dtype bf16 --seed 386 --warmup 10 --repeat 50 --json +``` + +Directory migration remains a separate integration step: the upstream refactor +must first be integrated, then move the reference/backends, registry, gtest, +bindings/build and tests/benchmark into its existing layout and rerun their +relevant checks. This change does not copy or partially recreate that refactor. diff --git a/rl_engine/contracts/numerical.py b/rl_engine/contracts/numerical.py index 3d2a10c3e..5a1e75785 100644 --- a/rl_engine/contracts/numerical.py +++ b/rl_engine/contracts/numerical.py @@ -90,6 +90,13 @@ class BackendProvenance: reference_dtype: str candidate_tf32_enabled: bool reference_tf32_enabled: bool + arithmetic_profile: str | None = None + fallback: bool = False + fallback_reason: str | None = None + semantic_version: str | None = None + binary_path: str | None = None + binary_sha256: str | None = None + observation: str | None = None def to_dict(self) -> dict[str, Any]: return asdict(self) @@ -267,11 +274,25 @@ def validate_backend_provenance( ) profile_contract = profile_contracts[provenance.backend_profile] expected_backend = str(profile_contract["backend_family"]) + declared_fallback = False + if provenance.fallback: + from rl_engine.kernels.adaln_contract import SHARED_PROFILE + + # Only AdaLN's observed same-arithmetic CUDA -> Triton route is declared. + declared_fallback = ( + provenance.requested_backend == "cuda" + and provenance.actual_backend == "triton" + and provenance.arithmetic_profile == SHARED_PROFILE + and provenance.observation == "successful_forward" + and bool(provenance.fallback_reason) + ) for field_name, actual in ( ("requested_backend", provenance.requested_backend), ("actual_backend", provenance.actual_backend), ): - if actual != expected_backend: + if actual != expected_backend and not ( + field_name == "requested_backend" and declared_fallback + ): raise ContractResolveError( f"backend provenance mismatch for {field_name}: expected " f"{expected_backend!r}, got {actual!r}" @@ -844,7 +865,7 @@ def _validate_chain_aggregates(root: Mapping[str, Any]) -> None: sole = list(root["sole_chain_level_logprob_metrics"]) if set(sole) != set(CHAIN_AGGREGATE_METRICS) or len(sole) != 3: raise ContractSchemaError( - "sole_chain_level_logprob_metrics must be exactly " f"{list(CHAIN_AGGREGATE_METRICS)}" + f"sole_chain_level_logprob_metrics must be exactly {list(CHAIN_AGGREGATE_METRICS)}" ) interval = root["default_clip_interval"] if len(interval) != 2 or float(interval[0]) >= float(interval[1]): @@ -897,7 +918,7 @@ def _validate_compat_views(contract: Mapping[str, Any]) -> None: cell["rtol"] ): raise ContractSchemaError( - f"compat accuracy mismatch vs forward_accuracy for " f"{op_class}/{dtype_name}" + f"compat accuracy mismatch vs forward_accuracy for {op_class}/{dtype_name}" ) diff --git a/rl_engine/kernels/adaln_contract.py b/rl_engine/kernels/adaln_contract.py new file mode 100644 index 000000000..1192ff0cb --- /dev/null +++ b/rl_engine/kernels/adaln_contract.py @@ -0,0 +1,101 @@ +# 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 typing import Any + +import torch + +SHARED_PROFILE = "adaln-shared-fp32-tree-v1" +SELECT01_PROFILE = "adaln-select01-framework-v1" +_BACKENDS = { + "rl_engine.kernels.ops.pytorch.norm.adaln_modulation": ("NativeAdaLNModulationOp", "pytorch"), + "rl_engine.kernels.ops.cuda.norm.adaln_modulation": ("CudaAdaLNModulationOp", "cuda"), + "rl_engine.kernels.ops.triton.norm.adaln_modulation": ("TritonAdaLNModulationOp", "triton"), +} + + +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 the declared arithmetic, not hardware qualification. CUDA and + Triton share the profile; Native on GPU is excluded pending its replay audit. + """ + + def __init__(self, implementation, trace, *, hidden, strict=False): + self.implementation = implementation + self.trace = trace + self.hidden = hidden + self.strict = strict + self.arithmetic_profile = trace["arithmetic_profile"] + 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 + if self.strict: + if indexed: + raise RuntimeError("strict AdaLN rejects select01 framework backward") + 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 self.arithmetic_profile != profile: + raise RuntimeError("strict AdaLN arithmetic profile mismatch") + if x.ndim != 3 or x.shape[-1] != self.hidden: + raise ValueError("strict AdaLN input hidden differs from resolved hidden") + if indexed: + outputs = self.implementation(x, modulation, eps=eps, modulate_index=modulate_index) + else: + outputs = self.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: + self.trace.update( + fallback=True, + fallback_reason="modulate_index_requires_select01", + reduction_order="pairwise_lower_upper_H; framework_select01_backward", + ) + if backend == "cuda": + from rl_engine.kernels.ops.base 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 669f22405..1a7fa2aa4 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -1063,17 +1063,31 @@ def get_adaln_modulation_op( *, hidden: int, modulate_index: torch.Tensor | None = None, + strict: bool = False, + arithmetic_profile: str | None = None, ) -> tuple[Any, dict[str, Any]]: - """Resolve AdaLN with a trace that makes backend fallback visible.""" + """Resolve AdaLN; strict permits only the shared declared replay profile.""" + from rl_engine.kernels.adaln_contract import ( + SELECT01_PROFILE, + SHARED_PROFILE, + AdaLNDispatch, + adaln_backend, + ) + if hidden <= 0: raise ValueError("hidden must be positive") + profile = SELECT01_PROFILE if modulate_index is not None else SHARED_PROFILE + if arithmetic_profile is not None and arithmetic_profile != profile: + raise ValueError(f"unsupported AdaLN arithmetic profile {arithmetic_profile!r}") + if strict and modulate_index is not None: + raise RuntimeError("strict AdaLN rejects select01 framework backward") platform = self._platform_for_device(device) if modulate_index is not None: backend = OpBackend.PYTORCH_ADALN_MODULATION op = self._get_or_create_backend(backend) if op is None: raise RuntimeError("PyTorch select01 AdaLN fallback is unavailable") - return op, { + trace = { "selected_backend": backend.name, "fallback": True, "fallback_reason": "modulate_index_requires_select01", @@ -1084,7 +1098,10 @@ def get_adaln_modulation_op( "stream_k": False, "tf32": False, "kernel_fingerprint": "adaln_modulation_select01_native_v1", + "requested_backend": "pytorch", + "arithmetic_profile": profile, } + return AdaLNDispatch(op, trace, hidden=hidden, strict=strict), trace candidates = self._priority_map[platform]["adaln_modulation"] rejected: list[str] = [] for backend in candidates: @@ -1095,6 +1112,10 @@ def get_adaln_modulation_op( 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), @@ -1105,8 +1126,11 @@ def get_adaln_modulation_op( "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 op, trace + return AdaLNDispatch(op, trace, hidden=hidden, strict=strict), trace raise RuntimeError(f"No functional adaln_modulation backend on {platform}: {rejected}") def _platform_for_device(self, device: torch.device | str | None) -> str: diff --git a/rl_engine/validation/operators/op_checks.py b/rl_engine/validation/operators/op_checks.py index 7f46f4174..370672f17 100644 --- a/rl_engine/validation/operators/op_checks.py +++ b/rl_engine/validation/operators/op_checks.py @@ -4,7 +4,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping, Sequence -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, replace from typing import Any import torch @@ -73,6 +73,8 @@ class CaseCheck: op_class: str passed: bool outputs: list[OutputCheck] + backend_trace: dict[str, Any] | None = None + backend_provenance: BackendProvenance | None = None @dataclass(frozen=True) @@ -185,6 +187,9 @@ def _run_candidate( total_outputs = sum(len(case.outputs) for case in case_checks) passed_outputs = sum(1 for case in case_checks for output in case.outputs if output.passed) pass_rate = float(passed_outputs / total_outputs) if total_outputs else 0.0 + observed_provenance = case_checks[0].backend_provenance if case_checks else None + if any(case.backend_provenance != observed_provenance for case in case_checks): + observed_provenance = None return CandidateReport( candidate_name=candidate.name, backend=candidate.backend, @@ -193,7 +198,7 @@ def _run_candidate( pass_rate=pass_rate, passed=passed_outputs == total_outputs, cases=case_checks, - backend_provenance=candidate.provenance, + backend_provenance=candidate.provenance or observed_provenance, ) @@ -204,9 +209,58 @@ def _run_case( ) -> CaseCheck: candidate_outputs = _flatten_tensors(_call_candidate(candidate.fn, case.inputs)) gold_outputs = _flatten_tensors(case.gold_fn(**case.inputs)) + candidate = _observe_adaln_candidate(candidate, case, gold_outputs, contract) return _compare_case_outputs(candidate, case, contract, candidate_outputs, gold_outputs) +def _observe_adaln_candidate( + candidate: CandidateSpec, + case: OperatorCase, + gold_outputs: list[torch.Tensor], + contract: Mapping[str, Any], +) -> CandidateSpec: + from rl_engine.kernels.adaln_contract import AdaLNDispatch + + if not isinstance(candidate.fn, AdaLNDispatch): + return candidate + trace = candidate.fn.trace + actual = trace["actual_backend"] + if actual != candidate.backend: + raise ContractResolveError( + f"candidate backend {candidate.backend!r} differs from {actual!r}" + ) + if normalize_dtype_name(case.dtype) != trace["execution_dtype"]: + raise ContractResolveError("case dtype differs from observed AdaLN execution dtype") + reference_dtypes = {normalize_dtype_name(output.dtype) for output in gold_outputs} + if reference_dtypes != {"float32"}: + raise ContractResolveError("AdaLN reference dtype must be observed FP32 for every output") + # C1 declares BF16 GPU profiles only; CPU/FP32 observations remain case traces. + if actual not in {"cuda", "triton"} or trace["execution_dtype"] != "bfloat16": + return candidate + provenance = BackendProvenance( + backend_profile="cuda_bf16" if actual == "cuda" else "triton_cuda_bf16", + requested_backend=trace["requested_backend"], + actual_backend=actual, + execution_dtype=trace["execution_dtype"], + accumulation_dtype="float32", + output_dtype=trace["output_dtype"], + reference_dtype=normalize_dtype_name(gold_outputs[0].dtype), + candidate_tf32_enabled=False, + reference_tf32_enabled=False, + arithmetic_profile=trace["arithmetic_profile"], + fallback=trace["fallback"], + fallback_reason=trace["fallback_reason"], + semantic_version=trace["semantic_version"], + binary_path=trace["binary_path"], + binary_sha256=trace["binary_sha256"], + observation=trace["observation"], + ) + validate_backend_provenance(contract, provenance) + if candidate.provenance is not None and candidate.provenance != provenance: + raise ContractResolveError("declared AdaLN provenance differs from successful execution") + return replace(candidate, provenance=provenance) + + def _run_case_backward( candidate: CandidateSpec, case: OperatorCase, @@ -245,6 +299,7 @@ def _run_case_backward( case.grad_input_names, grad_outputs=_match_grad_outputs(shared_upstreams, gold_outputs), ) + candidate = _observe_adaln_candidate(candidate, case, gold_outputs, contract) output_checks = _compare_case_outputs( candidate, case, @@ -305,6 +360,8 @@ 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, + backend_provenance=candidate.provenance, ) @@ -386,6 +443,8 @@ 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, + backend_provenance=candidate.provenance, ) @@ -578,8 +637,9 @@ def _compare_output( message=f"shape mismatch: candidate={tuple(candidate.shape)} gold={tuple(gold.shape)}", ) - candidate_fp32 = candidate.float() - gold_fp32 = gold.float() + # Independent references may live on CPU even when the candidate runs on GPU. + candidate_fp32 = candidate.detach().to(device="cpu", dtype=torch.float32) + gold_fp32 = gold.detach().to(device="cpu", dtype=torch.float32) abs_error = (candidate_fp32 - gold_fp32).abs() if abs_error.numel() == 0: max_abs_error = 0.0 diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 023da17c7..c2a5f53a1 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -402,12 +402,25 @@ def make_operator_case( ) -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.kernels.adaln_contract 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/test_adaln_dispatch.py b/tests/test_adaln_dispatch.py new file mode 100644 index 000000000..619a40803 --- /dev/null +++ b/tests/test_adaln_dispatch.py @@ -0,0 +1,274 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""AdaLN dispatch policy and observations at the public resolver seam.""" + +import pytest +import torch + +from argparse import Namespace +from dataclasses import replace + +from rl_engine.kernels.gtest import CandidateSpec, run_operator_suite +from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case +from rl_engine.kernels.gtest.operator_specs import GtestAdaLNReference +from rl_engine.kernels.gtest.tolerance import ( + BackendProvenance, + ContractResolveError, + load_contract, + resolve_tolerance, + validate_backend_provenance, +) +from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp +from rl_engine.kernels.registry import OpBackend +from rl_engine.kernels.registry import KernelRegistry + + +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 + + +@pytest.mark.parametrize("invalid", ["reference_dtype", "case_dtype"]) +def test_observed_gtest_rejects_false_reference_and_case_dtype_claims(invalid): + args = Namespace( + op="adaln_modulation", + candidate="pytorch", + arch_key=None, + batch=2, + seq=3, + normalized_dim=7, + seed=386, + ) + case = make_operator_case(args, torch.bfloat16, torch.device("cpu")) + case = ( + replace(case, gold_fn=NativeAdaLNModulationOp()) + if invalid == "reference_dtype" + else replace(case, dtype=torch.float32) + ) + with pytest.raises(ContractResolveError, match="reference dtype|case dtype"): + run_operator_suite( + "adaln_modulation", + candidates=[make_candidate(args, observe_backend=True)], + cases=[case], + ) + + +def test_provenance_accepts_only_the_declared_observed_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, + arithmetic_profile="adaln-shared-fp32-tree-v1", + fallback=True, + fallback_reason="CUDA extension unavailable", + observation="successful_forward", + ) + assert validate_backend_provenance(load_contract(), proof) == proof + for invalid in ( + replace(proof, fallback=False), + replace(proof, observation=None), + replace(proof, arithmetic_profile="unknown"), + replace(proof, fallback_reason=None), + replace(proof, requested_backend="pytorch"), + ): + with pytest.raises(ContractResolveError): + validate_backend_provenance(load_contract(), invalid) + + +def test_strict_rejects_unknown_profile_and_indexed_backward(): + registry = KernelRegistry() + with pytest.raises(ValueError, match="arithmetic profile"): + registry.get_adaln_modulation_op("cpu", hidden=2, strict=True, arithmetic_profile="unknown") + 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_provenance(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 + proof = report.candidates[0].backend_provenance + assert proof.actual_backend == "triton" + assert proof.requested_backend == "cuda" and proof.fallback is True + assert proof.arithmetic_profile == "adaln-shared-fp32-tree-v1" + assert proof.observation == "successful_forward" and proof.binary_sha256 is None diff --git a/tests/test_adaln_gtest.py b/tests/test_adaln_gtest.py index c505399ac..d43b17070 100644 --- a/tests/test_adaln_gtest.py +++ b/tests/test_adaln_gtest.py @@ -729,7 +729,7 @@ def resolve(backend): 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, NativeAdaLNModulationOp) + assert isinstance(op.implementation, NativeAdaLNModulationOp) assert trace["rejected_backends"] == [ OpBackend.CUDA_ADALN_MODULATION.name, OpBackend.TRITON_ADALN_MODULATION.name, diff --git a/tests/test_adaln_modulation.py b/tests/test_adaln_modulation.py index c3658520c..6a70fbe6a 100644 --- a/tests/test_adaln_modulation.py +++ b/tests/test_adaln_modulation.py @@ -63,7 +63,7 @@ def test_adaln_modulation_rejects_wrong_modulation_shape(): def test_adaln_registry_trace_reports_selected_backend(): op, trace = kernel_registry.get_adaln_modulation_op(device="cpu", hidden=4) - assert isinstance(op, NativeAdaLNModulationOp) + assert isinstance(op.implementation, NativeAdaLNModulationOp) assert trace["selected_backend"] == OpBackend.PYTORCH_ADALN_MODULATION.name assert trace["fallback"] is False assert trace["accumulator_dtype"] == "fp32" @@ -97,7 +97,7 @@ def test_adaln_modulate_index_uses_explicit_native_fallback(): modulation.grad, torch.tensor([[1.0, 1.0, -1.0, 1.0, 1.0, 1.0]]).repeat(2, 1), ) - assert isinstance(op, NativeAdaLNModulationOp) + 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" @@ -109,8 +109,8 @@ def test_adaln_cuda_registry_reports_real_fallback(): except RuntimeError: pytest.skip("CUDA required") op, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=32) - assert type(op).__name__ in {"CudaAdaLNModulationOp", "TritonAdaLNModulationOp"} - if type(op).__name__ == "TritonAdaLNModulationOp": + 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: @@ -119,7 +119,7 @@ def test_adaln_cuda_registry_reports_real_fallback(): def test_adaln_registry_skips_cuda_for_large_hidden(): op, trace = kernel_registry.get_adaln_modulation_op(device="cuda", hidden=5000) - assert type(op).__name__ == "TritonAdaLNModulationOp" + assert type(op.implementation).__name__ == "TritonAdaLNModulationOp" assert trace["fallback"] is True assert "CUDA_ADALN_MODULATION: H > 4096" in trace["rejected_backends"] diff --git a/tests/test_benchmark_adaln_modulation.py b/tests/test_benchmark_adaln_modulation.py new file mode 100644 index 000000000..c70b742e2 --- /dev/null +++ b/tests/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.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/validation/operators/test_op_checks.py b/tests/validation/operators/test_op_checks.py index 9c9b6999b..26a84e35b 100644 --- a/tests/validation/operators/test_op_checks.py +++ b/tests/validation/operators/test_op_checks.py @@ -112,6 +112,23 @@ def test_logp_native_candidate_suite_passes(): assert all(case.passed for case in report.candidates[0].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..d6ae1e2fa 100644 --- a/tools/validation/operators/check_operator.py +++ b/tools/validation/operators/check_operator.py @@ -99,8 +99,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 +143,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, From b31a814e884829a3065ca8c00e10f2ff94f70a35 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Mon, 12 Oct 2026 17:48:51 +0800 Subject: [PATCH 11/37] refactor(adaln): migrate WS1 operator to upstream layout Rebase the existing AdaLN work onto fixed test-qwenimage revision 16b81175874ecf11d5ff05c00de3b69e9457591b. Move reference, backends, dispatch, tests and benchmark into their canonical owners; preserve legacy module and stable operator identity, arithmetic, independent gold and shared precision thresholds. Validation: 229 CPU-focused tests passed, 302 hardware/opt-in tests skipped; focused mypy, Ruff and BF16 H3072 forward/backward CLI passed. GPU builds, H100 tests, sanitizer, benchmarks and the repository-wide suite remain manual; updated commands are documented. Signed-off-by: hongleng-github --- .github/workflows/ci.yml | 2 +- .github/workflows/ws1-gtest-gpu.yml | 14 +- benchmarks/benchmark_adaln_modulation.py | 247 +----------------- .../norm/benchmark_adaln_modulation.py | 246 +++++++++++++++++ build_tools/extensions.py | 2 +- ci/scripts/run_ws1_gtest.sh | 6 +- csrc/cuda/{ => norm}/adaln_modulation.cu | 0 docs/contributing/gtest-usage.md | 2 +- docs/operators/adaln-modulation.md | 78 +++++- .../backends/cuda/norm/adaln_modulation.py | 52 ++++ .../shared/triton/norm/adaln_modulation.py | 194 ++++++++++++++ rl_engine/contracts/numerical.py | 2 +- rl_engine/kernels/adaln_contract.py | 105 +------- .../kernels/ops/cuda/norm/adaln_modulation.py | 58 +--- .../ops/pytorch/norm/adaln_modulation.py | 164 +----------- .../ops/triton/norm/adaln_modulation.py | 157 +---------- rl_engine/reference/norm/adaln_modulation.py | 158 +++++++++++ rl_engine/runtime/adaln.py | 101 +++++++ .../runtime/provenance/module_names.json | 3 + rl_engine/runtime/registry.py | 8 +- rl_engine/validation/operators/op_checks.py | 2 +- .../validation/operators/operator_specs.py | 10 +- .../test_adaln_shared_bounds.py | 8 +- .../test_benchmark_adaln_modulation.py | 2 +- .../build/test_build_platform_collectives.py | 13 + tests/{ => ops/norm}/test_adaln_modulation.py | 30 ++- tests/{ => runtime}/test_adaln_dispatch.py | 14 +- tests/runtime/test_layout_compatibility.py | 28 ++ .../operators}/test_adaln_gtest.py | 125 +++------ tools/migration/layout.json | 13 +- 30 files changed, 1013 insertions(+), 831 deletions(-) create mode 100644 benchmarks/operators/norm/benchmark_adaln_modulation.py rename csrc/cuda/{ => norm}/adaln_modulation.cu (100%) create mode 100644 rl_engine/backends/cuda/norm/adaln_modulation.py create mode 100644 rl_engine/backends/shared/triton/norm/adaln_modulation.py create mode 100644 rl_engine/reference/norm/adaln_modulation.py create mode 100644 rl_engine/runtime/adaln.py rename tests/{ => backends}/test_adaln_shared_bounds.py (93%) rename tests/{ => benchmarks}/test_benchmark_adaln_modulation.py (92%) rename tests/{ => ops/norm}/test_adaln_modulation.py (94%) rename tests/{ => runtime}/test_adaln_dispatch.py (95%) rename tests/{ => validation/operators}/test_adaln_gtest.py (89%) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index edcef1e36..24a5b1ad6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -75,7 +75,7 @@ jobs: env: CUDA_VISIBLE_DEVICES: "" PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1" - run: python -m pytest tests/test_adaln_gtest.py tests/test_adaln_shared_bounds.py -q -rs + run: python -m pytest tests/validation/operators/test_adaln_gtest.py tests/backends/test_adaln_shared_bounds.py -q -rs - name: Run Mocked Hardware Discovery Tests run: | diff --git a/.github/workflows/ws1-gtest-gpu.yml b/.github/workflows/ws1-gtest-gpu.yml index 51b1bb471..ee255a0d3 100644 --- a/.github/workflows/ws1-gtest-gpu.yml +++ b/.github/workflows/ws1-gtest-gpu.yml @@ -23,11 +23,10 @@ on: - "build_tools/**" - "csrc/**" - "setup.py" - - "tests/test_adaln*.py" + - "tests/ops/norm/test_adaln*.py" + - "tests/backends/test_adaln*.py" + - "tests/runtime/test_adaln*.py" - "tools/validation/operators/check_operator.py" - - "rl_engine/kernels/gtest/**" - - "rl_engine/kernels/ops/**" - - "rl_engine/testing/**" - "tools/validation/operators/sweep_ws1_four_judgments.py" - "tools/validation/operators/check_forward_invariance.py" - "tools/validation/operators/check_gradient_invariance.py" @@ -55,11 +54,10 @@ on: - "build_tools/**" - "csrc/**" - "setup.py" - - "tests/test_adaln*.py" + - "tests/ops/norm/test_adaln*.py" + - "tests/backends/test_adaln*.py" + - "tests/runtime/test_adaln*.py" - "tools/validation/operators/check_operator.py" - - "rl_engine/kernels/gtest/**" - - "rl_engine/kernels/ops/**" - - "rl_engine/testing/**" - "tools/validation/operators/sweep_ws1_four_judgments.py" - "tools/validation/operators/check_forward_invariance.py" - "tools/validation/operators/check_gradient_invariance.py" diff --git a/benchmarks/benchmark_adaln_modulation.py b/benchmarks/benchmark_adaln_modulation.py index 1efd2a505..8962638f9 100644 --- a/benchmarks/benchmark_adaln_modulation.py +++ b/benchmarks/benchmark_adaln_modulation.py @@ -1,243 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors -"""Compare AdaLN backends on Qwen-Image image-token lengths. +"""Compatibility import; the canonical module is loaded below.""" -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 -from pathlib import Path -import platform -import statistics -import subprocess +import importlib import sys +from pathlib import Path -import torch - -REPO_ROOT = Path(__file__).resolve().parents[1] -if str(REPO_ROOT) not in sys.path: - sys.path.insert(0, str(REPO_ROOT)) +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +if __name__ == "__main__": + import runpy -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] :], + runpy.run_module("benchmarks.operators.norm.benchmark_adaln_modulation", run_name="__main__") +else: + sys.modules[__name__] = importlib.import_module( + "benchmarks.operators.norm.benchmark_adaln_modulation" ) - - -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.kernels.adaln_contract import SHARED_PROFILE, AdaLNDispatch - from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp - - backends = [("pytorch", NativeAdaLNModulationOp())] - skipped = [] - try: - from rl_engine.kernels.ops.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.kernels.ops.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/benchmark_adaln_modulation.py", - "rl_engine/kernels/adaln_contract.py", - "rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py", - "rl_engine/kernels/ops/cuda/norm/adaln_modulation.py", - "rl_engine/kernels/ops/triton/norm/adaln_modulation.py", - "rl_engine/kernels/ops/vjp_fp32.py", - "csrc/cuda/adaln_modulation.cu", - "csrc/ops.cpp", - "setup.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/benchmarks/operators/norm/benchmark_adaln_modulation.py b/benchmarks/operators/norm/benchmark_adaln_modulation.py new file mode 100644 index 000000000..605f02cae --- /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 +from pathlib import Path +import platform +import statistics +import subprocess +import sys + +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.runtime.adaln import SHARED_PROFILE, AdaLNDispatch + from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp + + 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 d1e554611..8de45a676 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -237,7 +237,7 @@ def get_extensions(): # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. if adaln_cuda_enabled: - cuda_sources.append("csrc/cuda/adaln_modulation.cu") + 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. diff --git a/ci/scripts/run_ws1_gtest.sh b/ci/scripts/run_ws1_gtest.sh index 2c866fb8a..ecb2951b6 100755 --- a/ci/scripts/run_ws1_gtest.sh +++ b/ci/scripts/run_ws1_gtest.sh @@ -23,9 +23,9 @@ 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/test_adaln_gtest.py \ - tests/test_adaln_shared_bounds.py \ - tests/test_adaln_modulation.py \ + tests/validation/operators/test_adaln_gtest.py \ + tests/backends/test_adaln_shared_bounds.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" diff --git a/csrc/cuda/adaln_modulation.cu b/csrc/cuda/norm/adaln_modulation.cu similarity index 100% rename from csrc/cuda/adaln_modulation.cu rename to csrc/cuda/norm/adaln_modulation.cu diff --git a/docs/contributing/gtest-usage.md b/docs/contributing/gtest-usage.md index 95331d316..7fb75d88d 100644 --- a/docs/contributing/gtest-usage.md +++ b/docs/contributing/gtest-usage.md @@ -116,7 +116,7 @@ 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 scripts/check_operator.py --op adaln_modulation \ +.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 ``` diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 534626a48..98b4050b3 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -62,7 +62,7 @@ with padding across all declared launch configurations. AdaLN gtest checks candidate output dtype and promotes gold leaves to FP32 before autograd; casting inside the gold function alone would return BF16 gradients to BF16 leaves. -AdaLN is registered in `gtest/operator_specs.py` and `operator_inputs.py` as +AdaLN is registered in `rl_engine/validation/operators/operator_specs.py` and `operator_inputs.py` as `op_class="reduction"`, with PyTorch, CUDA and Triton candidates. The independent gold uses FP32 CPU `torch.nn.functional.layer_norm` and autograd for normalization. Its independent affine VJP folds token contributions in ascending FP32 order; @@ -73,10 +73,10 @@ for input gradients; no private AdaLN thresholds or SM90 override are added. Existing fixed-reference byte comparisons remain separate, unchanged gates. ```bash -.venv/bin/python scripts/check_operator.py --op adaln_modulation \ +.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 -.venv/bin/python -m pytest tests/test_adaln_gtest.py -q -rs +.venv/bin/python -m pytest tests/validation/operators/test_adaln_gtest.py -q -rs ``` Use CUDA/Triton candidates on the target GPU for native accuracy evidence. @@ -87,9 +87,9 @@ CPU tests also cover gate-only gradients, distinct image/text modulation chunks, corrupted-output/VJP rejection and unavailable-backend routing. The CLI accuracy report alone is not a provenance-checked WS1 system gate. -Run `python -m pytest tests/test_adaln_modulation.py -q` for focused tests. +Run `python -m pytest tests/ops/norm/test_adaln_modulation.py -q` for focused tests. Set `RLK_ADALN_REAL_SHAPES=1` to include the three full image shapes. -Run `python benchmarks/benchmark_adaln_modulation.py --real` for the three +Run `python benchmarks/operators/norm/benchmark_adaln_modulation.py --real` for the three image shapes, with `--dtype fp32` for the second dtype. The benchmark prints median forward and forward+backward latency plus peak allocated memory; `--warmup` and `--repeat` control the sample count. It explicitly skips the @@ -156,12 +156,68 @@ the same inputs and timing, but does not claim fixed-order replay. Run the new GPU cases explicitly on the target environment: ```bash -.venv/bin/python -m pytest tests/test_adaln_dispatch.py -k gpu -q -rs -.venv/bin/python benchmarks/benchmark_adaln_modulation.py \ +.venv/bin/python -m pytest tests/runtime/test_adaln_dispatch.py -k gpu -q -rs +.venv/bin/python benchmarks/operators/norm/benchmark_adaln_modulation.py \ --real --dtype bf16 --seed 386 --warmup 10 --repeat 50 --json ``` -Directory migration remains a separate integration step: the upstream refactor -must first be integrated, then move the reference/backends, registry, gtest, -bindings/build and tests/benchmark into its existing layout and rerun their -relevant checks. This change does not copy or partially recreate that refactor. +## Directory integration + +The implementation follows the upstream layout: reference in `rl_engine/reference/norm/`, +CUDA in `rl_engine/backends/cuda/norm/`, Triton in `rl_engine/backends/shared/triton/norm/`, +and dispatch policy in `rl_engine/runtime/adaln.py`. The shared Triton directory does +not claim ROCm qualification; strict GPU execution requires NVIDIA CUDA. +Legacy imports alias the canonical modules and retain class and dispatch state identity. +Native bindings live in `csrc/bindings/ops.cpp`; `build_tools/extensions.py` selects +`csrc/cuda/norm/adaln_modulation.cu` only for CUDA builds without fast math. +This directory migration preserves the arithmetic profile and precision thresholds. +GPU acceptance must be rerun against the migrated sources and a fresh extension. + +## Migrated-source H100 validation commands + +Run these commands yourself from the repository root on H100/SM90. They are not +migration validation results. Keep the source revision, working-tree status, build +log and test/sanitizer logs together; the benchmark JSON also records source and +loaded extension hashes. A skip is not native CUDA acceptance. + +```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 + +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 + +# Repeat each candidate for bf16 and fp32. +.venv/bin/python tools/validation/operators/check_operator.py --op adaln_modulation \ + --candidate cuda --device cuda --dtype bf16 --batch 2 --seq 3 \ + --normalized-dim 3072 --check-grad --json +.venv/bin/python tools/validation/operators/check_operator.py --op adaln_modulation \ + --candidate triton --device cuda --dtype bf16 --batch 2 --seq 3 \ + --normalized-dim 3072 --check-grad --json + +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 + +.venv/bin/python benchmarks/operators/norm/benchmark_adaln_modulation.py \ + --real --dtype bf16 --seed 386 --warmup 10 --repeat 50 --json +.venv/bin/python benchmarks/operators/norm/benchmark_adaln_modulation.py \ + --real --dtype fp32 --seed 386 --warmup 10 --repeat 50 --json +``` + +The repository-wide suite is also left for manual execution because it includes +hardware and integration scopes whose duration is not bounded by this migration: + +```bash +.venv/bin/python -m pytest tests -q -rs +``` 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..a26a0adf3 --- /dev/null +++ b/rl_engine/backends/cuda/norm/adaln_modulation.py @@ -0,0 +1,52 @@ +# 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 _validate + + +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") + + def __call__(self, x, modulation, *, eps=1e-6): + return self.forward(x, modulation, eps=eps) + + def forward(self, x, modulation, *, eps=1e-6): + _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_modulation.py b/rl_engine/backends/shared/triton/norm/adaln_modulation.py new file mode 100644 index 000000000..95c4f333f --- /dev/null +++ b/rl_engine/backends/shared/triton/norm/adaln_modulation.py @@ -0,0 +1,194 @@ +# 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 _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) + 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)) + 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) + 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)) + 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) + _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 + + def __call__(self, x, modulation, *, eps=1e-6): + return self.forward(x, modulation, eps=eps) + + def forward(self, x, modulation, *, eps=1e-6): + _validate(x, modulation, eps) + 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/contracts/numerical.py b/rl_engine/contracts/numerical.py index 5a1e75785..cba5e6701 100644 --- a/rl_engine/contracts/numerical.py +++ b/rl_engine/contracts/numerical.py @@ -276,7 +276,7 @@ def validate_backend_provenance( expected_backend = str(profile_contract["backend_family"]) declared_fallback = False if provenance.fallback: - from rl_engine.kernels.adaln_contract import SHARED_PROFILE + from rl_engine.runtime.adaln import SHARED_PROFILE # Only AdaLN's observed same-arithmetic CUDA -> Triton route is declared. declared_fallback = ( diff --git a/rl_engine/kernels/adaln_contract.py b/rl_engine/kernels/adaln_contract.py index 1192ff0cb..06f902e3f 100644 --- a/rl_engine/kernels/adaln_contract.py +++ b/rl_engine/kernels/adaln_contract.py @@ -1,101 +1,16 @@ # 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 +"""Compatibility import; the canonical module is loaded below.""" +import importlib import sys -from typing import Any - -import torch - -SHARED_PROFILE = "adaln-shared-fp32-tree-v1" -SELECT01_PROFILE = "adaln-select01-framework-v1" -_BACKENDS = { - "rl_engine.kernels.ops.pytorch.norm.adaln_modulation": ("NativeAdaLNModulationOp", "pytorch"), - "rl_engine.kernels.ops.cuda.norm.adaln_modulation": ("CudaAdaLNModulationOp", "cuda"), - "rl_engine.kernels.ops.triton.norm.adaln_modulation": ("TritonAdaLNModulationOp", "triton"), -} - - -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 the declared arithmetic, not hardware qualification. CUDA and - Triton share the profile; Native on GPU is excluded pending its replay audit. - """ - - def __init__(self, implementation, trace, *, hidden, strict=False): - self.implementation = implementation - self.trace = trace - self.hidden = hidden - self.strict = strict - self.arithmetic_profile = trace["arithmetic_profile"] - 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() +from pathlib import Path - def __call__(self, x, modulation, *, eps=1e-6, modulate_index=None): - return self.forward(x, modulation, eps=eps, modulate_index=modulate_index) +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - 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 - if self.strict: - if indexed: - raise RuntimeError("strict AdaLN rejects select01 framework backward") - 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 self.arithmetic_profile != profile: - raise RuntimeError("strict AdaLN arithmetic profile mismatch") - if x.ndim != 3 or x.shape[-1] != self.hidden: - raise ValueError("strict AdaLN input hidden differs from resolved hidden") - if indexed: - outputs = self.implementation(x, modulation, eps=eps, modulate_index=modulate_index) - else: - outputs = self.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: - self.trace.update( - fallback=True, - fallback_reason="modulate_index_requires_select01", - reduction_order="pairwise_lower_upper_H; framework_select01_backward", - ) - if backend == "cuda": - from rl_engine.kernels.ops.base import _C +if __name__ == "__main__": + import runpy - self.trace["binary_path"] = getattr(_C, "__file__", None) - return outputs + runpy.run_module("rl_engine.runtime.adaln", run_name="__main__") +else: + sys.modules[__name__] = importlib.import_module("rl_engine.runtime.adaln") diff --git a/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py b/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py index ce827e340..714557035 100644 --- a/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py +++ b/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py @@ -1,52 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors -"""Autograd wrapper for the native CUDA AdaLN modulation kernels.""" +"""Compatibility import; the canonical module is loaded below.""" -from __future__ import annotations +import importlib +import sys +from pathlib import Path -import torch +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[5])) -from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE -from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import _validate +if __name__ == "__main__": + import runpy - -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") - - def __call__(self, x, modulation, *, eps=1e-6): - return self.forward(x, modulation, eps=eps) - - def forward(self, x, modulation, *, eps=1e-6): - _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) + runpy.run_module("rl_engine.backends.cuda.norm.adaln_modulation", run_name="__main__") +else: + sys.modules[__name__] = importlib.import_module("rl_engine.backends.cuda.norm.adaln_modulation") diff --git a/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py b/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py index 3f28d1224..fa6a59b15 100644 --- a/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py +++ b/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py @@ -1,158 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors -"""FP32 AdaLN reference for one Qwen-Image attention or MLP modulation.""" +"""Compatibility import; the canonical module is loaded below.""" -from __future__ import annotations +import importlib +import sys +from pathlib import Path -import math -import struct +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[5])) -import torch +if __name__ == "__main__": + import runpy -from rl_engine.kernels.ops.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 _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") - message = "eps must remain finite FP32 without positive underflow" - try: - eps32 = struct.unpack("f", struct.pack("f", eps))[0] - except OverflowError: - raise ValueError(message) from None - if not math.isfinite(eps32) or (eps > 0 and eps32 == 0): - raise ValueError(message) - - -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 _select01_reference( - x: torch.Tensor, - modulation: torch.Tensor, - modulate_index: torch.Tensor, - eps: float, -) -> tuple[torch.Tensor, torch.Tensor]: - 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 torch.all((modulate_index == 0) | (modulate_index == 1)).item(): - raise ValueError("modulate_index values must be 0 or 1") - if not math.isfinite(eps) or eps < 0: - raise ValueError("eps must be finite and nonnegative") - - 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)) + runpy.run_module("rl_engine.reference.norm.adaln_modulation", run_name="__main__") +else: + sys.modules[__name__] = importlib.import_module("rl_engine.reference.norm.adaln_modulation") diff --git a/rl_engine/kernels/ops/triton/norm/adaln_modulation.py b/rl_engine/kernels/ops/triton/norm/adaln_modulation.py index fdb66dc80..a4296f5e7 100644 --- a/rl_engine/kernels/ops/triton/norm/adaln_modulation.py +++ b/rl_engine/kernels/ops/triton/norm/adaln_modulation.py @@ -1,151 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors -"""One row per program AdaLN; shared gradients fold tokens in ascending order.""" +"""Compatibility import; the canonical module is loaded below.""" -from __future__ import annotations +import importlib +import sys +from pathlib import Path -import torch -import triton -import triton.language as tl -from triton.language.extra.cuda import libdevice +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[5])) -from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import _validate +if __name__ == "__main__": + import runpy - -@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 _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) - 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)) - 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) - 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)) - 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 + runpy.run_module("rl_engine.backends.shared.triton.norm.adaln_modulation", run_name="__main__") +else: + sys.modules[__name__] = importlib.import_module( + "rl_engine.backends.shared.triton.norm.adaln_modulation" ) - 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) - _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 - - def __call__(self, x, modulation, *, eps=1e-6): - return self.forward(x, modulation, eps=eps) - - def forward(self, x, modulation, *, eps=1e-6): - _validate(x, modulation, eps) - 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_modulation.py b/rl_engine/reference/norm/adaln_modulation.py new file mode 100644 index 000000000..e2a35968f --- /dev/null +++ b/rl_engine/reference/norm/adaln_modulation.py @@ -0,0 +1,158 @@ +# 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 _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") + message = "eps must remain finite FP32 without positive underflow" + try: + eps32 = struct.unpack("f", struct.pack("f", eps))[0] + except OverflowError: + raise ValueError(message) from None + if not math.isfinite(eps32) or (eps > 0 and eps32 == 0): + raise ValueError(message) + + +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 _select01_reference( + x: torch.Tensor, + modulation: torch.Tensor, + modulate_index: torch.Tensor, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + 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 torch.all((modulate_index == 0) | (modulate_index == 1)).item(): + raise ValueError("modulate_index values must be 0 or 1") + if not math.isfinite(eps) or eps < 0: + raise ValueError("eps must be finite and nonnegative") + + 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..b1b3df7b8 --- /dev/null +++ b/rl_engine/runtime/adaln.py @@ -0,0 +1,101 @@ +# 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 typing import Any + +import torch + +SHARED_PROFILE = "adaln-shared-fp32-tree-v1" +SELECT01_PROFILE = "adaln-select01-framework-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"), +} + + +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 the declared arithmetic, not hardware qualification. CUDA and + Triton share the profile; Native on GPU is excluded pending its replay audit. + """ + + def __init__(self, implementation, trace, *, hidden, strict=False): + self.implementation = implementation + self.trace = trace + self.hidden = hidden + self.strict = strict + self.arithmetic_profile = trace["arithmetic_profile"] + 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 + if self.strict: + if indexed: + raise RuntimeError("strict AdaLN rejects select01 framework backward") + 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 self.arithmetic_profile != profile: + raise RuntimeError("strict AdaLN arithmetic profile mismatch") + if x.ndim != 3 or x.shape[-1] != self.hidden: + raise ValueError("strict AdaLN input hidden differs from resolved hidden") + if indexed: + outputs = self.implementation(x, modulation, eps=eps, modulate_index=modulate_index) + else: + outputs = self.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: + self.trace.update( + fallback=True, + fallback_reason="modulate_index_requires_select01", + reduction_order="pairwise_lower_upper_H; framework_select01_backward", + ) + 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/provenance/module_names.json b/rl_engine/runtime/provenance/module_names.json index dc3dbb23e..4afad7ccc 100644 --- a/rl_engine/runtime/provenance/module_names.json +++ b/rl_engine/runtime/provenance/module_names.json @@ -26,6 +26,7 @@ "rl_engine.backends.cuda.logprob.linear_logp": "rl_engine.kernels.ops.cuda.loss.linear_logp", "rl_engine.backends.cuda.logprob.logp": "rl_engine.kernels.ops.cuda.loss.logp", "rl_engine.backends.cuda.logprob.vocab_parallel_logp": "rl_engine.kernels.ops.cuda.loss.vocab_parallel_logp", + "rl_engine.backends.cuda.norm.adaln_modulation": "rl_engine.kernels.ops.cuda.norm.adaln_modulation", "rl_engine.backends.cuda.norm.rmsnorm": "rl_engine.kernels.ops.cuda.norm.rmsnorm", "rl_engine.backends.cuda.rope.rope": "rl_engine.kernels.ops.cuda.rotary_embedding.rope", "rl_engine.backends.extension": "rl_engine.kernels.ops.base", @@ -58,6 +59,7 @@ "rl_engine.backends.shared.triton.logprob.vocab_parallel_logp": "rl_engine.kernels.ops.triton.loss.vocab_parallel_logp", "rl_engine.backends.shared.triton.loss.grpo_loss": "rl_engine.kernels.ops.triton.loss.grpo_loss", "rl_engine.backends.shared.triton.loss.ratio_kl": "rl_engine.kernels.ops.triton.loss.ratio_kl", + "rl_engine.backends.shared.triton.norm.adaln_modulation": "rl_engine.kernels.ops.triton.norm.adaln_modulation", "rl_engine.backends.shared.triton.norm.rmsnorm": "rl_engine.kernels.ops.triton.rmsnorm_triton", "rl_engine.backends.shared.triton.rope.rope": "rl_engine.kernels.ops.triton.rotary_embedding.rope", "rl_engine.contracts.diagnostics.attention_axes": "rl_engine.kernels.ops.pytorch.attention.debug_taxonomy", @@ -89,6 +91,7 @@ "rl_engine.reference.loss.distributed_grpo_loss": "rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss", "rl_engine.reference.loss.grpo_loss": "rl_engine.kernels.ops.pytorch.loss.grpo_loss", "rl_engine.reference.loss.ratio_kl": "rl_engine.kernels.ops.pytorch.loss.ratio_kl", + "rl_engine.reference.norm.adaln_modulation": "rl_engine.kernels.ops.pytorch.norm.adaln_modulation", "rl_engine.reference.norm.rms_norm": "rl_engine.kernels.ops.pytorch.norm.rms_norm", "rl_engine.reference.packing.pack": "rl_engine.kernels.ops.pytorch.packing.pack", "rl_engine.reference.rope.rope": "rl_engine.kernels.ops.pytorch.rotary_embedding.rope" diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 1a7fa2aa4..05d82e71a 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -177,12 +177,12 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): TRITON_SILU = "rl_engine.backends.shared.triton.activation.swiglu.TritonSiLUOp" TRITON_SWIGLU = "rl_engine.backends.shared.triton.activation.swiglu.TritonSwiGLUOp" PYTORCH_ADALN_MODULATION = ( - "rl_engine.kernels.ops.pytorch.norm.adaln_modulation.NativeAdaLNModulationOp" + "rl_engine.reference.norm.adaln_modulation.NativeAdaLNModulationOp" ) TRITON_ADALN_MODULATION = ( - "rl_engine.kernels.ops.triton.norm.adaln_modulation.TritonAdaLNModulationOp" + "rl_engine.backends.shared.triton.norm.adaln_modulation.TritonAdaLNModulationOp" ) - CUDA_ADALN_MODULATION = "rl_engine.kernels.ops.cuda.norm.adaln_modulation.CudaAdaLNModulationOp" + 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. @@ -1067,7 +1067,7 @@ def get_adaln_modulation_op( arithmetic_profile: str | None = None, ) -> tuple[Any, dict[str, Any]]: """Resolve AdaLN; strict permits only the shared declared replay profile.""" - from rl_engine.kernels.adaln_contract import ( + from rl_engine.runtime.adaln import ( SELECT01_PROFILE, SHARED_PROFILE, AdaLNDispatch, diff --git a/rl_engine/validation/operators/op_checks.py b/rl_engine/validation/operators/op_checks.py index 370672f17..8f6a69ac0 100644 --- a/rl_engine/validation/operators/op_checks.py +++ b/rl_engine/validation/operators/op_checks.py @@ -219,7 +219,7 @@ def _observe_adaln_candidate( gold_outputs: list[torch.Tensor], contract: Mapping[str, Any], ) -> CandidateSpec: - from rl_engine.kernels.adaln_contract import AdaLNDispatch + from rl_engine.runtime.adaln import AdaLNDispatch if not isinstance(candidate.fn, AdaLNDispatch): return candidate diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index c2a5f53a1..503b19454 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -44,10 +44,12 @@ def _load_object(path: str) -> Any: gold_method="forward_fp32", candidate_paths={ "pytorch": ( - "rl_engine.kernels.ops.pytorch.norm.adaln_modulation.NativeAdaLNModulationOp" + "rl_engine.reference.norm.adaln_modulation.NativeAdaLNModulationOp" ), - "triton": "rl_engine.kernels.ops.triton.norm.adaln_modulation.TritonAdaLNModulationOp", - "cuda": "rl_engine.kernels.ops.cuda.norm.adaln_modulation.CudaAdaLNModulationOp", + "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, @@ -409,7 +411,7 @@ def make_candidate(args: argparse.Namespace, *, observe_backend: bool = False) - 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.kernels.adaln_contract import SHARED_PROFILE, AdaLNDispatch + from rl_engine.runtime.adaln import SHARED_PROFILE, AdaLNDispatch candidate_op = AdaLNDispatch( candidate_op, diff --git a/tests/test_adaln_shared_bounds.py b/tests/backends/test_adaln_shared_bounds.py similarity index 93% rename from tests/test_adaln_shared_bounds.py rename to tests/backends/test_adaln_shared_bounds.py index 854559b15..79927d892 100644 --- a/tests/test_adaln_shared_bounds.py +++ b/tests/backends/test_adaln_shared_bounds.py @@ -7,8 +7,8 @@ import pytest import torch -from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp -from rl_engine.kernels.registry import OpBackend, kernel_registry +from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +from rl_engine.runtime.registry import OpBackend, kernel_registry @pytest.fixture @@ -77,9 +77,7 @@ def test_cuda_shared_maximum_batch_forward_backward(cuda_op, dtype): 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(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)) diff --git a/tests/test_benchmark_adaln_modulation.py b/tests/benchmarks/test_benchmark_adaln_modulation.py similarity index 92% rename from tests/test_benchmark_adaln_modulation.py rename to tests/benchmarks/test_benchmark_adaln_modulation.py index c70b742e2..9588692e8 100644 --- a/tests/test_benchmark_adaln_modulation.py +++ b/tests/benchmarks/test_benchmark_adaln_modulation.py @@ -5,7 +5,7 @@ import pytest import torch -from benchmarks.benchmark_adaln_modulation import layer_norm_baseline, parse_args +from benchmarks.operators.norm.benchmark_adaln_modulation import layer_norm_baseline, parse_args def test_benchmark_cli_records_seed_json_and_rejects_empty_measurements(): 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/test_adaln_modulation.py b/tests/ops/norm/test_adaln_modulation.py similarity index 94% rename from tests/test_adaln_modulation.py rename to tests/ops/norm/test_adaln_modulation.py index 6a70fbe6a..c4c37c9e5 100644 --- a/tests/test_adaln_modulation.py +++ b/tests/ops/norm/test_adaln_modulation.py @@ -3,10 +3,10 @@ import pytest import torch -from rl_engine.kernels.gtest.operator_specs import GtestAdaLNReference -from rl_engine.kernels.gtest.tolerance import load_contract, resolve_tolerance -from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp -from rl_engine.kernels.registry import OpBackend, kernel_registry +from rl_engine.validation.operators.operator_specs import GtestAdaLNReference +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 def _assert_bytes_equal(actual, expected): @@ -188,7 +188,7 @@ def test_triton_adaln_modulation_forward_backward_and_batch_position(dtype): torch.cuda.init() except RuntimeError: pytest.skip("CUDA required") - from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + 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) @@ -228,7 +228,8 @@ def test_triton_irregular_hidden_and_bf16_tie_rounding(): torch.cuda.init() except RuntimeError: pytest.skip("CUDA required") - from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + 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), @@ -254,8 +255,9 @@ def test_adaln_gpu_padding_and_batch_size_do_not_change_valid_bytes(dtype): torch.cuda.init() except RuntimeError: pytest.skip("CUDA required") - from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp - from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + 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) @@ -295,7 +297,7 @@ def test_adaln_gpu_padding_and_batch_size_do_not_change_valid_bytes(dtype): @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.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp try: torch.cuda.init() @@ -358,7 +360,7 @@ def test_adaln_real_image_shapes_forward_backward(seq, dtype): _assert_bytes_equal(x.grad.cpu(), cpu_x.grad) _assert_bytes_equal(modulation.grad.cpu(), cpu_modulation.grad) - from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + 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_() @@ -383,8 +385,8 @@ def test_adaln_byte_comparison_distinguishes_signed_zero(): def test_adaln_launch_geometry_preserves_bytes(backend, dtype): if not torch.cuda.is_available(): pytest.skip("CUDA required") - from rl_engine.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp - from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + 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: @@ -479,7 +481,7 @@ 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.kernels.ops.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp try: op = CudaAdaLNModulationOp() @@ -488,7 +490,7 @@ def test_adaln_shared_matches_independent_autograd(backend, hidden, dtype): raise pytest.skip("compiled CUDA AdaLN extension required") elif backend == "triton": - from rl_engine.kernels.ops.triton.norm.adaln_modulation import TritonAdaLNModulationOp + from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp op = TritonAdaLNModulationOp() else: diff --git a/tests/test_adaln_dispatch.py b/tests/runtime/test_adaln_dispatch.py similarity index 95% rename from tests/test_adaln_dispatch.py rename to tests/runtime/test_adaln_dispatch.py index 619a40803..7245fe7ab 100644 --- a/tests/test_adaln_dispatch.py +++ b/tests/runtime/test_adaln_dispatch.py @@ -8,19 +8,19 @@ from argparse import Namespace from dataclasses import replace -from rl_engine.kernels.gtest import CandidateSpec, run_operator_suite -from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case -from rl_engine.kernels.gtest.operator_specs import GtestAdaLNReference -from rl_engine.kernels.gtest.tolerance import ( +from rl_engine.validation.operators import CandidateSpec, run_operator_suite +from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case +from rl_engine.validation.operators.operator_specs import GtestAdaLNReference +from rl_engine.contracts.numerical import ( BackendProvenance, ContractResolveError, load_contract, resolve_tolerance, validate_backend_provenance, ) -from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp -from rl_engine.kernels.registry import OpBackend -from rl_engine.kernels.registry import KernelRegistry +from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp +from rl_engine.runtime.registry import OpBackend +from rl_engine.runtime.registry import KernelRegistry def test_strict_cpu_call_records_execution_without_claiming_hardware_qualification(): diff --git a/tests/runtime/test_layout_compatibility.py b/tests/runtime/test_layout_compatibility.py index a3ab9f969..c45d8d505 100644 --- a/tests/runtime/test_layout_compatibility.py +++ b/tests/runtime/test_layout_compatibility.py @@ -24,6 +24,22 @@ ("legacy", "canonical", "shared_object"), [ ("rl_engine.kernels.registry", "rl_engine.runtime.registry", "kernel_registry"), + ("rl_engine.kernels.adaln_contract", "rl_engine.runtime.adaln", "AdaLNDispatch"), + ( + "rl_engine.kernels.ops.pytorch.norm.adaln_modulation", + "rl_engine.reference.norm.adaln_modulation", + "NativeAdaLNModulationOp", + ), + ( + "rl_engine.kernels.ops.cuda.norm.adaln_modulation", + "rl_engine.backends.cuda.norm.adaln_modulation", + "CudaAdaLNModulationOp", + ), + ( + "rl_engine.kernels.ops.triton.norm.adaln_modulation", + "rl_engine.backends.shared.triton.norm.adaln_modulation", + "TritonAdaLNModulationOp", + ), ("rl_engine.runtime_mode", "rl_engine.runtime.policy", "strict_contract_enabled"), ("rl_engine.integrations.state", "rl_engine.integrations.common.state", "_ACTIVE"), ("rl_engine.kernels.ops.vjp_fp32", "rl_engine.ops.autograd.vjp_fp32", "reduce_rows_fp32"), @@ -63,6 +79,14 @@ def test_ffn_compatibility_imports_share_platform_modules(legacy, canonical, sha assert getattr(old, shared_object) is getattr(new, shared_object) +def test_adaln_stable_operator_id_survives_directory_migration(): + from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp + + assert stable_operator_path(NativeAdaLNModulationOp()) == ( + "rl_engine.kernels.ops.pytorch.norm.adaln_modulation.NativeAdaLNModulationOp" + ) + + def test_stable_operator_id_does_not_hide_an_unknown_implementation(): from rl_engine.reference.gemm.det_gemm import NativeGemmOp @@ -123,6 +147,10 @@ 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/benchmark_adaln_modulation.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/test_adaln_gtest.py b/tests/validation/operators/test_adaln_gtest.py similarity index 89% rename from tests/test_adaln_gtest.py rename to tests/validation/operators/test_adaln_gtest.py index d43b17070..dc2719764 100644 --- a/tests/test_adaln_gtest.py +++ b/tests/validation/operators/test_adaln_gtest.py @@ -9,16 +9,16 @@ import pytest import torch -from rl_engine.kernels.gtest import run_operator_suite -from rl_engine.kernels.gtest.operator_inputs import operator_shape_name -from rl_engine.kernels.gtest.operator_specs import ( +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, ) -from rl_engine.kernels.gtest.tolerance import load_contract, resolve_tolerance -from rl_engine.kernels.ops.pytorch.norm.adaln_modulation import NativeAdaLNModulationOp -from rl_engine.kernels.registry import OpBackend, kernel_registry +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 def _args(hidden=7): @@ -35,15 +35,10 @@ def _args(hidden=7): @pytest.fixture( params=["pytorch", "triton"] - + [ - pytest.param(("cuda", threads), id=f"cuda-threads{threads}") - for threads in (128, 256, 512) - ] + + [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) - ) + 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() @@ -70,9 +65,7 @@ def _bytes_equal(actual, expected): @pytest.mark.parametrize("nan_upstream", [False, True]) -def test_adaln_gate_preserves_bf16_nan_payload_and_identity_vjp( - shared_op, nan_upstream -): +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) @@ -131,9 +124,7 @@ def test_adaln_subnormal_eps_is_not_flushed_to_zero(shared_op, dtype, eps, rstd) 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 - ) + 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. @@ -143,9 +134,7 @@ def test_adaln_bf16_forward_casts_only_at_output(shared_op, scale, positive): 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 - ) + 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, @@ -173,9 +162,7 @@ def test_adaln_bf16_backward_keeps_partials_and_accumulators_fp32(shared_op): ) _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) - ) + _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]) @@ -184,9 +171,7 @@ def test_adaln_bf16_vjp_matches_fp32_with_one_final_cast(shared_op, hidden): 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 - ) + 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 = ( @@ -203,9 +188,7 @@ def test_adaln_bf16_vjp_matches_fp32_with_one_final_cast(shared_op, hidden): @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 -): +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) @@ -215,9 +198,7 @@ def test_adaln_bf16_boundary_rounds_halfway_values_to_even( _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) - ) + _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]) @@ -228,11 +209,7 @@ def test_adaln_h3072_sample_bytes_survive_batch_positions_and_padding(shared_op, @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) - ] + [("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): @@ -280,12 +257,8 @@ def observe(x, full, dy, dg, *, candidate=op): 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] - ) + 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( @@ -305,9 +278,7 @@ def test_adaln_backward_accepts_transposed_and_broadcast_upstreams(shared_op, dt 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()) - ) + 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_() @@ -321,9 +292,7 @@ def test_adaln_backward_accepts_transposed_and_broadcast_upstreams(shared_op, dt 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 - ) + torch.testing.assert_close(result.cpu().float(), expected, atol=tol.atol, rtol=tol.rtol) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -346,15 +315,11 @@ def test_adaln_single_trainable_input_and_unused_output(shared_op, dtype, traina 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) - ) + _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) - ) + _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): @@ -420,9 +385,7 @@ def test_adaln_modulation_gradient_uses_ascending_token_order(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 - ) + 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]])) @@ -433,9 +396,7 @@ def test_adaln_independent_gold_uses_contract_token_order(): 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)) - ) + (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]])) @@ -489,9 +450,7 @@ def test_adaln_shared_rejects_mixed_devices(shared_op): @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 -): +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) @@ -501,9 +460,7 @@ def test_adaln_shared_rejects_eps_overflow_and_positive_underflow( @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 -): +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) @@ -554,11 +511,7 @@ def test_adaln_gtest_detects_wrong_output_and_backward(broken): 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 - ) + 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 @@ -592,9 +545,7 @@ def wrong_dtype(**inputs): ) assert not report.passed, report.to_dict() checks = report.candidates[0].cases[0].outputs[:2] - assert all( - not check.passed and "dtype mismatch" in check.message for check in checks - ) + assert all(not check.passed and "dtype mismatch" in check.message for check in checks) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -606,12 +557,8 @@ def test_adaln_gate_only_and_separate_stream_chunks(shared_op, dtype, triple): # 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 - ) + 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 ( @@ -672,9 +619,7 @@ def test_adaln_gate_only_and_separate_stream_chunks(shared_op, dtype, triple): @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 -): +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) @@ -686,9 +631,7 @@ def test_adaln_hidden_boundary_matches_independent_math_and_fixed_bytes( 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()) - ) + 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"), @@ -698,9 +641,7 @@ def test_adaln_hidden_boundary_matches_independent_math_and_fixed_bytes( 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 - ) + 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)) diff --git a/tools/migration/layout.json b/tools/migration/layout.json index cb8fe469a..6f52bf111 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -968,7 +968,18 @@ "tests/test_ws1_workload.py": "tests/validation/operators/test_ws1_workload.py", "tests/test_ws2_attention_gpu_acceptance.py": "tests/validation/distributed/test_ws2_attention_gpu_acceptance.py", "tests/test_ws2_cp_attention_drift_benchmark.py": "tests/validation/distributed/test_ws2_cp_attention_drift_benchmark.py", - "tests/test_ws2_megatron_te_cp_compare.py": "tests/validation/distributed/test_ws2_megatron_te_cp_compare.py" + "tests/test_ws2_megatron_te_cp_compare.py": "tests/validation/distributed/test_ws2_megatron_te_cp_compare.py", + "rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py": "rl_engine/reference/norm/adaln_modulation.py", + "rl_engine/kernels/ops/cuda/norm/adaln_modulation.py": "rl_engine/backends/cuda/norm/adaln_modulation.py", + "rl_engine/kernels/ops/triton/norm/adaln_modulation.py": "rl_engine/backends/shared/triton/norm/adaln_modulation.py", + "rl_engine/kernels/adaln_contract.py": "rl_engine/runtime/adaln.py", + "csrc/cuda/adaln_modulation.cu": "csrc/cuda/norm/adaln_modulation.cu", + "benchmarks/benchmark_adaln_modulation.py": "benchmarks/operators/norm/benchmark_adaln_modulation.py", + "tests/test_adaln_modulation.py": "tests/ops/norm/test_adaln_modulation.py", + "tests/test_adaln_gtest.py": "tests/validation/operators/test_adaln_gtest.py", + "tests/test_adaln_shared_bounds.py": "tests/backends/test_adaln_shared_bounds.py", + "tests/test_adaln_dispatch.py": "tests/runtime/test_adaln_dispatch.py", + "tests/test_benchmark_adaln_modulation.py": "tests/benchmarks/test_benchmark_adaln_modulation.py" }, "splits": { "rl_engine/integrations/framework_operators.py": [ From 42e965b0004408cd0b9d6b59010b31076979c3d4 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Tue, 13 Oct 2026 01:57:09 +0800 Subject: [PATCH 12/37] refactor(adaln): remove local legacy compatibility entries Delete the five AdaLN import/benchmark shims and their compatibility checks. Remove AdaLN-only local migration and identity aliases while keeping upstream compatibility and canonical implementations unchanged. Validation: focused CPU regression 223 passed, 302 skipped; BF16 H3072 CPU forward/backward CLI, Ruff and diff checks passed. Migrated ten external local scripts, including a CPU-only probe (4 passed). GPU/H100 and full-repository validation were not run. Signed-off-by: hongleng-github --- benchmarks/benchmark_adaln_modulation.py | 18 ------------- docs/operators/adaln-modulation.md | 1 - rl_engine/kernels/adaln_contract.py | 16 ------------ .../kernels/ops/cuda/norm/adaln_modulation.py | 16 ------------ .../ops/pytorch/norm/adaln_modulation.py | 16 ------------ .../ops/triton/norm/adaln_modulation.py | 18 ------------- .../runtime/provenance/module_names.json | 3 --- tests/runtime/test_layout_compatibility.py | 25 ------------------- tools/migration/layout.json | 13 +--------- 9 files changed, 1 insertion(+), 125 deletions(-) delete mode 100644 benchmarks/benchmark_adaln_modulation.py delete mode 100644 rl_engine/kernels/adaln_contract.py delete mode 100644 rl_engine/kernels/ops/cuda/norm/adaln_modulation.py delete mode 100644 rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py delete mode 100644 rl_engine/kernels/ops/triton/norm/adaln_modulation.py diff --git a/benchmarks/benchmark_adaln_modulation.py b/benchmarks/benchmark_adaln_modulation.py deleted file mode 100644 index 8962638f9..000000000 --- a/benchmarks/benchmark_adaln_modulation.py +++ /dev/null @@ -1,18 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Compatibility import; the canonical module is loaded below.""" - -import importlib -import sys -from pathlib import Path - -if __package__ in (None, ""): - sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -if __name__ == "__main__": - import runpy - - runpy.run_module("benchmarks.operators.norm.benchmark_adaln_modulation", run_name="__main__") -else: - sys.modules[__name__] = importlib.import_module( - "benchmarks.operators.norm.benchmark_adaln_modulation" - ) diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 98b4050b3..5f6c90829 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -167,7 +167,6 @@ The implementation follows the upstream layout: reference in `rl_engine/referenc CUDA in `rl_engine/backends/cuda/norm/`, Triton in `rl_engine/backends/shared/triton/norm/`, and dispatch policy in `rl_engine/runtime/adaln.py`. The shared Triton directory does not claim ROCm qualification; strict GPU execution requires NVIDIA CUDA. -Legacy imports alias the canonical modules and retain class and dispatch state identity. Native bindings live in `csrc/bindings/ops.cpp`; `build_tools/extensions.py` selects `csrc/cuda/norm/adaln_modulation.cu` only for CUDA builds without fast math. This directory migration preserves the arithmetic profile and precision thresholds. diff --git a/rl_engine/kernels/adaln_contract.py b/rl_engine/kernels/adaln_contract.py deleted file mode 100644 index 06f902e3f..000000000 --- a/rl_engine/kernels/adaln_contract.py +++ /dev/null @@ -1,16 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Compatibility import; the canonical module is loaded below.""" - -import importlib -import sys -from pathlib import Path - -if __package__ in (None, ""): - sys.path.insert(0, str(Path(__file__).resolve().parents[2])) - -if __name__ == "__main__": - import runpy - - runpy.run_module("rl_engine.runtime.adaln", run_name="__main__") -else: - sys.modules[__name__] = importlib.import_module("rl_engine.runtime.adaln") diff --git a/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py b/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py deleted file mode 100644 index 714557035..000000000 --- a/rl_engine/kernels/ops/cuda/norm/adaln_modulation.py +++ /dev/null @@ -1,16 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Compatibility import; the canonical module is loaded below.""" - -import importlib -import sys -from pathlib import Path - -if __package__ in (None, ""): - sys.path.insert(0, str(Path(__file__).resolve().parents[5])) - -if __name__ == "__main__": - import runpy - - runpy.run_module("rl_engine.backends.cuda.norm.adaln_modulation", run_name="__main__") -else: - sys.modules[__name__] = importlib.import_module("rl_engine.backends.cuda.norm.adaln_modulation") diff --git a/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py b/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py deleted file mode 100644 index fa6a59b15..000000000 --- a/rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py +++ /dev/null @@ -1,16 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Compatibility import; the canonical module is loaded below.""" - -import importlib -import sys -from pathlib import Path - -if __package__ in (None, ""): - sys.path.insert(0, str(Path(__file__).resolve().parents[5])) - -if __name__ == "__main__": - import runpy - - runpy.run_module("rl_engine.reference.norm.adaln_modulation", run_name="__main__") -else: - sys.modules[__name__] = importlib.import_module("rl_engine.reference.norm.adaln_modulation") diff --git a/rl_engine/kernels/ops/triton/norm/adaln_modulation.py b/rl_engine/kernels/ops/triton/norm/adaln_modulation.py deleted file mode 100644 index a4296f5e7..000000000 --- a/rl_engine/kernels/ops/triton/norm/adaln_modulation.py +++ /dev/null @@ -1,18 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -"""Compatibility import; the canonical module is loaded below.""" - -import importlib -import sys -from pathlib import Path - -if __package__ in (None, ""): - sys.path.insert(0, str(Path(__file__).resolve().parents[5])) - -if __name__ == "__main__": - import runpy - - runpy.run_module("rl_engine.backends.shared.triton.norm.adaln_modulation", run_name="__main__") -else: - sys.modules[__name__] = importlib.import_module( - "rl_engine.backends.shared.triton.norm.adaln_modulation" - ) diff --git a/rl_engine/runtime/provenance/module_names.json b/rl_engine/runtime/provenance/module_names.json index 4afad7ccc..dc3dbb23e 100644 --- a/rl_engine/runtime/provenance/module_names.json +++ b/rl_engine/runtime/provenance/module_names.json @@ -26,7 +26,6 @@ "rl_engine.backends.cuda.logprob.linear_logp": "rl_engine.kernels.ops.cuda.loss.linear_logp", "rl_engine.backends.cuda.logprob.logp": "rl_engine.kernels.ops.cuda.loss.logp", "rl_engine.backends.cuda.logprob.vocab_parallel_logp": "rl_engine.kernels.ops.cuda.loss.vocab_parallel_logp", - "rl_engine.backends.cuda.norm.adaln_modulation": "rl_engine.kernels.ops.cuda.norm.adaln_modulation", "rl_engine.backends.cuda.norm.rmsnorm": "rl_engine.kernels.ops.cuda.norm.rmsnorm", "rl_engine.backends.cuda.rope.rope": "rl_engine.kernels.ops.cuda.rotary_embedding.rope", "rl_engine.backends.extension": "rl_engine.kernels.ops.base", @@ -59,7 +58,6 @@ "rl_engine.backends.shared.triton.logprob.vocab_parallel_logp": "rl_engine.kernels.ops.triton.loss.vocab_parallel_logp", "rl_engine.backends.shared.triton.loss.grpo_loss": "rl_engine.kernels.ops.triton.loss.grpo_loss", "rl_engine.backends.shared.triton.loss.ratio_kl": "rl_engine.kernels.ops.triton.loss.ratio_kl", - "rl_engine.backends.shared.triton.norm.adaln_modulation": "rl_engine.kernels.ops.triton.norm.adaln_modulation", "rl_engine.backends.shared.triton.norm.rmsnorm": "rl_engine.kernels.ops.triton.rmsnorm_triton", "rl_engine.backends.shared.triton.rope.rope": "rl_engine.kernels.ops.triton.rotary_embedding.rope", "rl_engine.contracts.diagnostics.attention_axes": "rl_engine.kernels.ops.pytorch.attention.debug_taxonomy", @@ -91,7 +89,6 @@ "rl_engine.reference.loss.distributed_grpo_loss": "rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss", "rl_engine.reference.loss.grpo_loss": "rl_engine.kernels.ops.pytorch.loss.grpo_loss", "rl_engine.reference.loss.ratio_kl": "rl_engine.kernels.ops.pytorch.loss.ratio_kl", - "rl_engine.reference.norm.adaln_modulation": "rl_engine.kernels.ops.pytorch.norm.adaln_modulation", "rl_engine.reference.norm.rms_norm": "rl_engine.kernels.ops.pytorch.norm.rms_norm", "rl_engine.reference.packing.pack": "rl_engine.kernels.ops.pytorch.packing.pack", "rl_engine.reference.rope.rope": "rl_engine.kernels.ops.pytorch.rotary_embedding.rope" diff --git a/tests/runtime/test_layout_compatibility.py b/tests/runtime/test_layout_compatibility.py index c45d8d505..1f7d6fa02 100644 --- a/tests/runtime/test_layout_compatibility.py +++ b/tests/runtime/test_layout_compatibility.py @@ -24,22 +24,6 @@ ("legacy", "canonical", "shared_object"), [ ("rl_engine.kernels.registry", "rl_engine.runtime.registry", "kernel_registry"), - ("rl_engine.kernels.adaln_contract", "rl_engine.runtime.adaln", "AdaLNDispatch"), - ( - "rl_engine.kernels.ops.pytorch.norm.adaln_modulation", - "rl_engine.reference.norm.adaln_modulation", - "NativeAdaLNModulationOp", - ), - ( - "rl_engine.kernels.ops.cuda.norm.adaln_modulation", - "rl_engine.backends.cuda.norm.adaln_modulation", - "CudaAdaLNModulationOp", - ), - ( - "rl_engine.kernels.ops.triton.norm.adaln_modulation", - "rl_engine.backends.shared.triton.norm.adaln_modulation", - "TritonAdaLNModulationOp", - ), ("rl_engine.runtime_mode", "rl_engine.runtime.policy", "strict_contract_enabled"), ("rl_engine.integrations.state", "rl_engine.integrations.common.state", "_ACTIVE"), ("rl_engine.kernels.ops.vjp_fp32", "rl_engine.ops.autograd.vjp_fp32", "reduce_rows_fp32"), @@ -79,14 +63,6 @@ def test_ffn_compatibility_imports_share_platform_modules(legacy, canonical, sha assert getattr(old, shared_object) is getattr(new, shared_object) -def test_adaln_stable_operator_id_survives_directory_migration(): - from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp - - assert stable_operator_path(NativeAdaLNModulationOp()) == ( - "rl_engine.kernels.ops.pytorch.norm.adaln_modulation.NativeAdaLNModulationOp" - ) - - def test_stable_operator_id_does_not_hide_an_unknown_implementation(): from rl_engine.reference.gemm.det_gemm import NativeGemmOp @@ -149,7 +125,6 @@ def test_packaged_contract_bytes_match_main(package, resource, digest): "scripts/ws1_reference.py", "scripts/check_operator.py", "tools/validation/operators/check_operator.py", - "benchmarks/benchmark_adaln_modulation.py", "benchmarks/operators/norm/benchmark_adaln_modulation.py", "rl_engine/testing/logprob_comparison.py", "examples/vime_rocm_attention_ablation/run.py", diff --git a/tools/migration/layout.json b/tools/migration/layout.json index 6f52bf111..cb8fe469a 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -968,18 +968,7 @@ "tests/test_ws1_workload.py": "tests/validation/operators/test_ws1_workload.py", "tests/test_ws2_attention_gpu_acceptance.py": "tests/validation/distributed/test_ws2_attention_gpu_acceptance.py", "tests/test_ws2_cp_attention_drift_benchmark.py": "tests/validation/distributed/test_ws2_cp_attention_drift_benchmark.py", - "tests/test_ws2_megatron_te_cp_compare.py": "tests/validation/distributed/test_ws2_megatron_te_cp_compare.py", - "rl_engine/kernels/ops/pytorch/norm/adaln_modulation.py": "rl_engine/reference/norm/adaln_modulation.py", - "rl_engine/kernels/ops/cuda/norm/adaln_modulation.py": "rl_engine/backends/cuda/norm/adaln_modulation.py", - "rl_engine/kernels/ops/triton/norm/adaln_modulation.py": "rl_engine/backends/shared/triton/norm/adaln_modulation.py", - "rl_engine/kernels/adaln_contract.py": "rl_engine/runtime/adaln.py", - "csrc/cuda/adaln_modulation.cu": "csrc/cuda/norm/adaln_modulation.cu", - "benchmarks/benchmark_adaln_modulation.py": "benchmarks/operators/norm/benchmark_adaln_modulation.py", - "tests/test_adaln_modulation.py": "tests/ops/norm/test_adaln_modulation.py", - "tests/test_adaln_gtest.py": "tests/validation/operators/test_adaln_gtest.py", - "tests/test_adaln_shared_bounds.py": "tests/backends/test_adaln_shared_bounds.py", - "tests/test_adaln_dispatch.py": "tests/runtime/test_adaln_dispatch.py", - "tests/test_benchmark_adaln_modulation.py": "tests/benchmarks/test_benchmark_adaln_modulation.py" + "tests/test_ws2_megatron_te_cp_compare.py": "tests/validation/distributed/test_ws2_megatron_te_cp_compare.py" }, "splits": { "rl_engine/integrations/framework_operators.py": [ From 8fb3180b6e85401c2c2a0ee4d54b691e7ce02f9f Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Tue, 13 Oct 2026 16:59:41 +0800 Subject: [PATCH 13/37] test(adaln): include runtime dispatch tests in CI pipeline Signed-off-by: hongleng-github --- .github/workflows/ci.yml | 2 +- ci/scripts/run_ws1_gtest.sh | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 24a5b1ad6..57c0dd7b4 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -75,7 +75,7 @@ jobs: env: CUDA_VISIBLE_DEVICES: "" PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1" - run: python -m pytest tests/validation/operators/test_adaln_gtest.py tests/backends/test_adaln_shared_bounds.py -q -rs + run: python -m pytest tests/validation/operators/test_adaln_gtest.py tests/backends/test_adaln_shared_bounds.py tests/runtime/test_adaln_dispatch.py -q -rs - name: Run Mocked Hardware Discovery Tests run: | diff --git a/ci/scripts/run_ws1_gtest.sh b/ci/scripts/run_ws1_gtest.sh index ecb2951b6..c0d27c277 100755 --- a/ci/scripts/run_ws1_gtest.sh +++ b/ci/scripts/run_ws1_gtest.sh @@ -25,6 +25,7 @@ echo "[ws1-gtest] interpreter=$PY out=$OUT" 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 From d84d12fe02182a8a9f21b0e0d6e4d9e0fe7802e4 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Fri, 16 Oct 2026 16:35:21 +0800 Subject: [PATCH 14/37] refactor(ci): remove AdaLN CPU contract and reference tests from CI pipeline Signed-off-by: hongleng-github --- .github/workflows/ci.yml | 6 ------ 1 file changed, 6 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 57c0dd7b4..4634569cd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -71,12 +71,6 @@ jobs: - name: Check refactor packaging and compatibility run: python ci/run.py layout -q - - name: Run AdaLN CPU contract and reference tests - env: - CUDA_VISIBLE_DEVICES: "" - PYTEST_DISABLE_PLUGIN_AUTOLOAD: "1" - run: python -m pytest tests/validation/operators/test_adaln_gtest.py tests/backends/test_adaln_shared_bounds.py tests/runtime/test_adaln_dispatch.py -q -rs - - name: Run Mocked Hardware Discovery Tests run: | python -m pytest tests/runtime/test_dispatch.py -v From 538f4f42037670542fbd0718daf8756ee7db5c55 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Fri, 16 Oct 2026 16:38:25 +0800 Subject: [PATCH 15/37] refactor(ci): remove obsolete test paths from WS1 GPU workflow Signed-off-by: hongleng-github --- .github/workflows/ws1-gtest-gpu.yml | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/.github/workflows/ws1-gtest-gpu.yml b/.github/workflows/ws1-gtest-gpu.yml index ee255a0d3..a0d8ac0c4 100644 --- a/.github/workflows/ws1-gtest-gpu.yml +++ b/.github/workflows/ws1-gtest-gpu.yml @@ -22,11 +22,6 @@ on: - "rl_engine/reference/**" - "build_tools/**" - "csrc/**" - - "setup.py" - - "tests/ops/norm/test_adaln*.py" - - "tests/backends/test_adaln*.py" - - "tests/runtime/test_adaln*.py" - - "tools/validation/operators/check_operator.py" - "tools/validation/operators/sweep_ws1_four_judgments.py" - "tools/validation/operators/check_forward_invariance.py" - "tools/validation/operators/check_gradient_invariance.py" @@ -53,11 +48,6 @@ on: - "rl_engine/reference/**" - "build_tools/**" - "csrc/**" - - "setup.py" - - "tests/ops/norm/test_adaln*.py" - - "tests/backends/test_adaln*.py" - - "tests/runtime/test_adaln*.py" - - "tools/validation/operators/check_operator.py" - "tools/validation/operators/sweep_ws1_four_judgments.py" - "tools/validation/operators/check_forward_invariance.py" - "tools/validation/operators/check_gradient_invariance.py" From 11fb8f849ce0592818299b90fe065e90cbee2d72 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 21 Oct 2026 00:55:50 +0800 Subject: [PATCH 16/37] fix(adaln): rename shift and scale variables to dshift and dscale in backward_reduce_kernel Signed-off-by: hongleng-github --- csrc/cuda/norm/adaln_modulation.cu | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/csrc/cuda/norm/adaln_modulation.cu b/csrc/cuda/norm/adaln_modulation.cu index 826a61d59..9c1745eed 100644 --- a/csrc/cuda/norm/adaln_modulation.cu +++ b/csrc/cuda/norm/adaln_modulation.cu @@ -117,15 +117,15 @@ __global__ void backward_reduce_kernel( const int col = blockIdx.x * blockDim.x + threadIdx.x; const int batch = blockIdx.y; if (col >= hidden) return; - float shift = 0.0f, scale = 0.0f; + 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; - shift = __fadd_rn(shift, partial_shift[offset]); - scale = __fadd_rn(scale, partial_scale[offset]); + dshift = __fadd_rn(dshift, partial_shift[offset]); + dscale = __fadd_rn(dscale, partial_scale[offset]); } - dm[batch * 3 * hidden + col] = static_cast(shift); - dm[batch * 3 * hidden + hidden + col] = static_cast(scale); + 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]; } From 90b467ac5b6aed7b7da2113696996e26b3ebdfa8 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 21 Oct 2026 14:01:34 +0800 Subject: [PATCH 17/37] feat(adaln): implement _center_and_rstd function for mean and rstd calculation Signed-off-by: hongleng-github --- csrc/cuda/norm/adaln_modulation.cu | 40 +++++++++++-------- .../shared/triton/norm/adaln_modulation.py | 20 ++++++---- 2 files changed, 35 insertions(+), 25 deletions(-) diff --git a/csrc/cuda/norm/adaln_modulation.cu b/csrc/cuda/norm/adaln_modulation.cu index 9c1745eed..ad311d483 100644 --- a/csrc/cuda/norm/adaln_modulation.cu +++ b/csrc/cuda/norm/adaln_modulation.cu @@ -30,13 +30,14 @@ __device__ float fixed_sum(float* shared, int width) { return result; } +struct RowStats { + float mean; + float 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; +__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; } @@ -48,6 +49,19 @@ __global__ void forward_kernel( } 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); @@ -67,17 +81,9 @@ __global__ void backward_rows_kernel( extern __shared__ float shared[]; const int row = blockIdx.x; const int batch = row / seq; - 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))); + 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]); diff --git a/rl_engine/backends/shared/triton/norm/adaln_modulation.py b/rl_engine/backends/shared/triton/norm/adaln_modulation.py index 95c4f333f..a0bd4fd79 100644 --- a/rl_engine/backends/shared/triton/norm/adaln_modulation.py +++ b/rl_engine/backends/shared/triton/norm/adaln_modulation.py @@ -23,6 +23,16 @@ def _sum_hidden(values, BLOCK: tl.constexpr, LOG: tl.constexpr): 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, @@ -39,10 +49,7 @@ def _fwd( batch = row // S col = tl.arange(0, BLOCK) x = tl.load(X + row * H + col, col < H, other=0).to(tl.float32) - 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)) + 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 @@ -71,10 +78,7 @@ def _bwd_rows( 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) - 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)) + 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) From 903943a56277de4e0b9b70aa4c4b59dd4a49fb97 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 21 Oct 2026 14:40:31 +0800 Subject: [PATCH 18/37] fix(numerical): correct string formatting in error messages for contract validation Signed-off-by: hongleng-github --- rl_engine/contracts/numerical.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rl_engine/contracts/numerical.py b/rl_engine/contracts/numerical.py index cba5e6701..d0ca04d41 100644 --- a/rl_engine/contracts/numerical.py +++ b/rl_engine/contracts/numerical.py @@ -865,7 +865,7 @@ def _validate_chain_aggregates(root: Mapping[str, Any]) -> None: sole = list(root["sole_chain_level_logprob_metrics"]) if set(sole) != set(CHAIN_AGGREGATE_METRICS) or len(sole) != 3: raise ContractSchemaError( - f"sole_chain_level_logprob_metrics must be exactly {list(CHAIN_AGGREGATE_METRICS)}" + "sole_chain_level_logprob_metrics must be exactly " f"{list(CHAIN_AGGREGATE_METRICS)}" ) interval = root["default_clip_interval"] if len(interval) != 2 or float(interval[0]) >= float(interval[1]): @@ -918,7 +918,7 @@ def _validate_compat_views(contract: Mapping[str, Any]) -> None: cell["rtol"] ): raise ContractSchemaError( - f"compat accuracy mismatch vs forward_accuracy for {op_class}/{dtype_name}" + f"compat accuracy mismatch vs forward_accuracy for " f"{op_class}/{dtype_name}" ) From c139bbf3ee14d0a518e3d01f43cb308dafb96ce5 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 21 Oct 2026 14:52:04 +0800 Subject: [PATCH 19/37] refactor(adaln): streamline backend provenance handling and validation logic Signed-off-by: hongleng-github --- rl_engine/contracts/numerical.py | 23 +----- rl_engine/validation/operators/op_checks.py | 26 ++++--- tests/runtime/test_adaln_dispatch.py | 79 +++++++++++++++------ 3 files changed, 77 insertions(+), 51 deletions(-) diff --git a/rl_engine/contracts/numerical.py b/rl_engine/contracts/numerical.py index d0ca04d41..3d2a10c3e 100644 --- a/rl_engine/contracts/numerical.py +++ b/rl_engine/contracts/numerical.py @@ -90,13 +90,6 @@ class BackendProvenance: reference_dtype: str candidate_tf32_enabled: bool reference_tf32_enabled: bool - arithmetic_profile: str | None = None - fallback: bool = False - fallback_reason: str | None = None - semantic_version: str | None = None - binary_path: str | None = None - binary_sha256: str | None = None - observation: str | None = None def to_dict(self) -> dict[str, Any]: return asdict(self) @@ -274,25 +267,11 @@ def validate_backend_provenance( ) profile_contract = profile_contracts[provenance.backend_profile] expected_backend = str(profile_contract["backend_family"]) - declared_fallback = False - if provenance.fallback: - from rl_engine.runtime.adaln import SHARED_PROFILE - - # Only AdaLN's observed same-arithmetic CUDA -> Triton route is declared. - declared_fallback = ( - provenance.requested_backend == "cuda" - and provenance.actual_backend == "triton" - and provenance.arithmetic_profile == SHARED_PROFILE - and provenance.observation == "successful_forward" - and bool(provenance.fallback_reason) - ) for field_name, actual in ( ("requested_backend", provenance.requested_backend), ("actual_backend", provenance.actual_backend), ): - if actual != expected_backend and not ( - field_name == "requested_backend" and declared_fallback - ): + if actual != expected_backend: raise ContractResolveError( f"backend provenance mismatch for {field_name}: expected " f"{expected_backend!r}, got {actual!r}" diff --git a/rl_engine/validation/operators/op_checks.py b/rl_engine/validation/operators/op_checks.py index 8f6a69ac0..a6f889616 100644 --- a/rl_engine/validation/operators/op_checks.py +++ b/rl_engine/validation/operators/op_checks.py @@ -219,7 +219,7 @@ def _observe_adaln_candidate( gold_outputs: list[torch.Tensor], contract: Mapping[str, Any], ) -> CandidateSpec: - from rl_engine.runtime.adaln import AdaLNDispatch + from rl_engine.runtime.adaln import SHARED_PROFILE, AdaLNDispatch if not isinstance(candidate.fn, AdaLNDispatch): return candidate @@ -237,9 +237,20 @@ def _observe_adaln_candidate( # C1 declares BF16 GPU profiles only; CPU/FP32 observations remain case traces. if actual not in {"cuda", "triton"} or trace["execution_dtype"] != "bfloat16": return candidate + fallback = trace["requested_backend"] != actual + if fallback and not ( + trace["fallback"] + and trace["requested_backend"] == "cuda" + and actual == "triton" + and trace["arithmetic_profile"] == SHARED_PROFILE + and trace["observation"] == "successful_forward" + and trace["fallback_reason"] + ): + raise ContractResolveError("undeclared AdaLN backend fallback") + # Validate the executed backend's profile; request history stays in backend_trace. provenance = BackendProvenance( backend_profile="cuda_bf16" if actual == "cuda" else "triton_cuda_bf16", - requested_backend=trace["requested_backend"], + requested_backend=actual, actual_backend=actual, execution_dtype=trace["execution_dtype"], accumulation_dtype="float32", @@ -247,15 +258,12 @@ def _observe_adaln_candidate( reference_dtype=normalize_dtype_name(gold_outputs[0].dtype), candidate_tf32_enabled=False, reference_tf32_enabled=False, - arithmetic_profile=trace["arithmetic_profile"], - fallback=trace["fallback"], - fallback_reason=trace["fallback_reason"], - semantic_version=trace["semantic_version"], - binary_path=trace["binary_path"], - binary_sha256=trace["binary_sha256"], - observation=trace["observation"], ) validate_backend_provenance(contract, provenance) + if fallback: + if candidate.provenance is not None: + raise ContractResolveError("AdaLN fallback must be reported through backend_trace") + return candidate if candidate.provenance is not None and candidate.provenance != provenance: raise ContractResolveError("declared AdaLN provenance differs from successful execution") return replace(candidate, provenance=provenance) diff --git a/tests/runtime/test_adaln_dispatch.py b/tests/runtime/test_adaln_dispatch.py index 7245fe7ab..9fe87655d 100644 --- a/tests/runtime/test_adaln_dispatch.py +++ b/tests/runtime/test_adaln_dispatch.py @@ -93,7 +93,7 @@ def test_observed_gtest_rejects_false_reference_and_case_dtype_claims(invalid): ) -def test_provenance_accepts_only_the_declared_observed_cuda_to_triton_fallback(): +def test_common_provenance_rejects_cuda_to_triton_fallback(): proof = BackendProvenance( backend_profile="triton_cuda_bf16", requested_backend="cuda", @@ -104,21 +104,56 @@ def test_provenance_accepts_only_the_declared_observed_cuda_to_triton_fallback() reference_dtype="float32", candidate_tf32_enabled=False, reference_tf32_enabled=False, - arithmetic_profile="adaln-shared-fp32-tree-v1", - fallback=True, - fallback_reason="CUDA extension unavailable", - observation="successful_forward", ) - assert validate_backend_provenance(load_contract(), proof) == proof - for invalid in ( - replace(proof, fallback=False), - replace(proof, observation=None), - replace(proof, arithmetic_profile="unknown"), - replace(proof, fallback_reason=None), - replace(proof, requested_backend="pytorch"), - ): - with pytest.raises(ContractResolveError): - validate_backend_provenance(load_contract(), invalid) + with pytest.raises(ContractResolveError, match="requested_backend"): + validate_backend_provenance(load_contract(), proof) + + +@pytest.mark.parametrize( + "override", + [ + {}, + {"fallback": False}, + {"observation": None}, + {"arithmetic_profile": "unknown"}, + {"fallback_reason": None}, + {"requested_backend": "pytorch"}, + ], +) +def test_gtest_validates_and_serializes_fallback_trace_on_cpu(monkeypatch, override): + from rl_engine.runtime import adaln + + # Simulate backend observations to exercise report policy without a GPU kernel. + monkeypatch.setattr(adaln, "adaln_backend", lambda op: "triton") + + class ObservedDispatch(adaln.AdaLNDispatch): + def forward(self, *args, **kwargs): + outputs = super().forward(*args, **kwargs) + self.trace.update(override) + return outputs + + trace = { + "requested_backend": "cuda", + "fallback": True, + "fallback_reason": "CUDA extension unavailable", + "arithmetic_profile": adaln.SHARED_PROFILE, + } + op = ObservedDispatch(NativeAdaLNModulationOp(), trace, hidden=7) + args = Namespace(op="adaln_modulation", batch=2, seq=3, normalized_dim=7, seed=386) + case = make_operator_case(args, torch.bfloat16, torch.device("cpu")) + candidate = CandidateSpec(name="observed-adaln", backend="triton", fn=op) + if override: + with pytest.raises(ContractResolveError, match="undeclared AdaLN backend fallback"): + run_operator_suite("adaln_modulation", candidates=[candidate], cases=[case]) + else: + report = run_operator_suite("adaln_modulation", candidates=[candidate], cases=[case]) + assert report.passed + result = report.to_dict()["candidates"][0] + assert result["backend_provenance"] is None + observed = result["cases"][0]["backend_trace"] + assert observed["requested_backend"] == "cuda" + assert observed["actual_backend"] == "triton" + assert observed["fallback"] is True and observed["fallback_reason"] def test_strict_rejects_unknown_profile_and_indexed_backward(): @@ -267,8 +302,12 @@ def test_gpu_gtest_persists_observed_bf16_fallback_provenance(monkeypatch): check_grad=True, ) assert report.passed - proof = report.candidates[0].backend_provenance - assert proof.actual_backend == "triton" - assert proof.requested_backend == "cuda" and proof.fallback is True - assert proof.arithmetic_profile == "adaln-shared-fp32-tree-v1" - assert proof.observation == "successful_forward" and proof.binary_sha256 is None + result = report.candidates[0] + assert result.backend_provenance is None + payload = report.to_dict()["candidates"][0]["cases"][0] + assert payload["backend_provenance"] is None + 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 From 50f3397ce4bb54ed6b45a28d73ac5cc1abc82809 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 21 Oct 2026 14:59:43 +0800 Subject: [PATCH 20/37] refactor(adaln): update arithmetic profile handling and simplify test for strict rejection Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 6 +- rl_engine/runtime/registry.py | 128 ++++++++++++++------------- tests/runtime/test_adaln_dispatch.py | 4 +- 3 files changed, 70 insertions(+), 68 deletions(-) diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 5f6c90829..5224ae588 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -109,7 +109,6 @@ and complete test and sanitizer logs. ```python op, trace = kernel_registry.get_adaln_modulation_op( x.device, hidden=x.shape[-1], strict=True, - arithmetic_profile="adaln-shared-fp32-tree-v1", ) y, gate = op(x, modulation) ``` @@ -117,8 +116,9 @@ y, gate = op(x, modulation) Strict dispatch pins the shared FP32 hidden tree, ascending token fold and final cast. NVIDIA CUDA may fall back to Triton when the extension is unavailable or H exceeds 4096: both implement this same arithmetic profile, accept FP32/BF16, -and retain the existing launch-configuration replay tests. Strict rejects unknown -profiles/backends, the indexed/select01 framework backward and PyTorch execution +and retain the existing launch-configuration replay tests. The arithmetic profile +is inferred from the requested operation and recorded in the trace. Strict rejects +unknown backends, the indexed/select01 framework backward and PyTorch execution on GPU, whose cross-backend replay is not qualified here. It checks the actual input device and hidden size again at execution. CPU Native is permitted for the fixed reference. This policy does not itself certify any hardware/runtime profile. diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 05d82e71a..dc29f0e91 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -67,115 +67,116 @@ def __getitem__(cls, name: str): class OpBackend(Enum, metaclass=_KernelEnumMeta): # NVIDIA optimized stack - FLASH_ATTN = "rl_engine.backends.cuda.attention.flash_attn.FlashAttentionOp" + FLASH_ATTN = "rl_engine.kernels.ops.cuda.attention.flash_attn.FlashAttentionOp" FLASHINFER = "rl_engine.kernels.ops.cuda.flashinfer.FlashInferOp" # TMA-accelerated LogP for SM90+ (Warp Specialization) - CUDA_FUSED_LOGP_SM90 = "rl_engine.backends.cuda.logprob.logp.FusedLogpSM90Op" - CUDA_FUSED_LOGP_GENERIC = "rl_engine.backends.cuda.logprob.logp.FusedLogpGenericOp" + CUDA_FUSED_LOGP_SM90 = "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpSM90Op" + CUDA_FUSED_LOGP_GENERIC = "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp" MUSA_FUSED_LOGP_GENERIC = "rl_engine.backends.musa.logprob.logp.MusaFusedLogpOp" - CUDA_DETERMINISTIC_LOGP = "rl_engine.backends.cuda.logprob.logp.DeterministicLogpCUDAOp" + CUDA_DETERMINISTIC_LOGP = "rl_engine.kernels.ops.cuda.loss.logp.DeterministicLogpCUDAOp" # Deterministic standard-softmax attention (issue #147); not FlashAttention. CUDA_DETERMINISTIC_ATTENTION = ( - "rl_engine.backends.cuda.attention.deterministic_attn.DeterministicAttentionOp" + "rl_engine.kernels.ops.cuda.attention.deterministic_attn.DeterministicAttentionOp" ) # AMD ROCm optimized stack ROCM_AITER = "rl_engine.kernels.ops.rocm.aiter.AiterOp" ROCM_CK = "rl_engine.kernels.ops.rocm.composable_kernel.CKOp" - ROCM_FLASH_ATTN = "rl_engine.backends.rocm.attention.flash_attn.RocmFlashAttentionOp" + ROCM_FLASH_ATTN = "rl_engine.kernels.ops.rocm.attention.flash_attn.RocmFlashAttentionOp" # WS2 strict ROCm attention core (AITER/CK dense MHA, Split-KV disabled). # Reachable only through ``get_attention_op``: it is not a drop-in for the # SDPA-shaped ``attn`` wrappers and must never be a silent fallback. ROCM_STRICT_ATTENTION = ( - "rl_engine.backends.rocm.attention.flash_attn.StrictRocmAiterCKAttentionCore" + "rl_engine.kernels.ops.rocm.attention.flash_attn.StrictRocmAiterCKAttentionCore" ) # GRPO loss (group reward normalization + clipped surrogate + KL) - TRITON_GRPO_LOSS = "rl_engine.backends.shared.triton.loss.grpo_loss.TritonGRPOLossOp" - PYTORCH_GRPO_LOSS = "rl_engine.reference.loss.grpo_loss.NativeGRPOLossOp" + TRITON_GRPO_LOSS = "rl_engine.kernels.ops.triton.loss.grpo_loss.TritonGRPOLossOp" + PYTORCH_GRPO_LOSS = "rl_engine.kernels.ops.pytorch.loss.grpo_loss.NativeGRPOLossOp" # Fused linear log-prob (hidden @ W^T -> selected-token logp, no [N, V] logits) CUDA_FUSED_LINEAR_LOGP_SM90 = ( - "rl_engine.backends.cuda.logprob.linear_logp.FusedLinearLogpSM90Op" + "rl_engine.kernels.ops.cuda.loss.linear_logp.FusedLinearLogpSM90Op" ) - TRITON_LINEAR_LOGP = "rl_engine.backends.shared.triton.logprob.linear_logp.TritonLinearLogpOp" - PYTORCH_LINEAR_LOGP = "rl_engine.reference.logprob.linear_logp.NativeLinearLogpOp" + TRITON_LINEAR_LOGP = "rl_engine.kernels.ops.triton.loss.linear_logp.TritonLinearLogpOp" + PYTORCH_LINEAR_LOGP = "rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp" # Fused policy-ratio + KL-penalty front-end (PPO/GRPO), logits -> (ratio, kl) - TRITON_RATIO_KL = "rl_engine.backends.shared.triton.loss.ratio_kl.TritonRatioKLOp" - PYTORCH_RATIO_KL = "rl_engine.reference.loss.ratio_kl.NativeRatioKLOp" + TRITON_RATIO_KL = "rl_engine.kernels.ops.triton.loss.ratio_kl.TritonRatioKLOp" + PYTORCH_RATIO_KL = "rl_engine.kernels.ops.pytorch.loss.ratio_kl.NativeRatioKLOp" # Variable-length packing (pack-and-pad), [B,S,...] -> [Total_Active,...] - PYTORCH_PACK = "rl_engine.reference.packing.pack.NativePackOp" + PYTORCH_PACK = "rl_engine.kernels.ops.pytorch.packing.pack.NativePackOp" # Batch-invariant deterministic GEMM (WS1 #146) - CUDA_DET_GEMM = "rl_engine.backends.cuda.gemm.det_gemm.DetGemmOp" - TRITON_DET_GEMM = "rl_engine.backends.shared.triton.gemm.det_gemm.TritonDetGemmOp" - ASCEND_DET_GEMM = "rl_engine.backends.ascend.gemm.det_gemm.DetGemmAscendOp" + CUDA_DET_GEMM = "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp" + TRITON_DET_GEMM = "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp" + ASCEND_DET_GEMM = "rl_engine.kernels.ops.ascend.matmul.det_gemm.DetGemmAscendOp" # NON-deterministic reference (torch.matmul); reference/benchmark ONLY, # intentionally excluded from det_gemm dispatch (cuBLAS breaks invariance). - PYTORCH_GEMM = "rl_engine.reference.gemm.det_gemm.NativeGemmOp" + PYTORCH_GEMM = "rl_engine.kernels.ops.pytorch.matmul.det_gemm.NativeGemmOp" # Batch-invariant selected-logprob (WS1 #148: locked reduction order) TRITON_BATCH_INVARIANT_LOGP = ( - "rl_engine.backends.shared.triton.logprob.batch_invariant_logp.TritonBatchInvariantLogpOp" + "rl_engine.kernels.ops.triton.loss.batch_invariant_logp.TritonBatchInvariantLogpOp" ) PYTORCH_BATCH_INVARIANT_LOGP = ( - "rl_engine.reference.logprob.batch_invariant_logp.NativeBatchInvariantLogpOp" + "rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp.NativeBatchInvariantLogpOp" ) CUDA_BATCH_INVARIANT_LOGP_SM90 = ( - "rl_engine.backends.cuda.logprob.batch_invariant_logp.BatchInvariantLogpSM90Op" + "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90Op" ) ASCEND_BATCH_INVARIANT_LOGP = ( - "rl_engine.backends.ascend.logprob.batch_invariant_logp.BatchInvariantLogpAscendOp" + "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAscendOp" ) - ASCEND_RMS_NORM = "rl_engine.backends.ascend.norm.rmsnorm.RMSNormAscendOp" - ASCEND_EMBEDDING = "rl_engine.backends.ascend.embedding.embedding.AscendEmbeddingOp" - ASCEND_FUSED_LOGP = "rl_engine.backends.ascend.logprob.logp.FusedLogpAscendOp" - ASCEND_LM_HEAD = "rl_engine.backends.ascend.embedding.lm_head.AscendLMHeadOp" + ASCEND_RMS_NORM = "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp" + ASCEND_EMBEDDING = "rl_engine.kernels.ops.ascend.linear.embedding.AscendEmbeddingOp" + ASCEND_FUSED_LOGP = "rl_engine.kernels.ops.ascend.loss.logp.FusedLogpAscendOp" + ASCEND_LM_HEAD = "rl_engine.kernels.ops.ascend.linear.lm_head.AscendLMHeadOp" ASCEND_FUSED_LINEAR_LOGP = ( - "rl_engine.backends.ascend.logprob.linear_logp.FusedLinearLogpAscendOp" + "rl_engine.kernels.ops.ascend.loss.linear_logp.FusedLinearLogpAscendOp" ) # Deterministic vocab-parallel TP logprob reference (WS2 #241 PR3) PYTORCH_VOCAB_PARALLEL_LOGP = ( - "rl_engine.reference.logprob.vocab_parallel_logp.VocabParallelLogprobOp" + "rl_engine.kernels.ops.pytorch.loss.vocab_parallel_logp.VocabParallelLogprobOp" ) ROCM_VOCAB_PARALLEL_LOGP = ( - "rl_engine.backends.rocm.logprob.vocab_parallel_logp.RocmVocabParallelLogprobOp" + "rl_engine.kernels.ops.rocm.loss.vocab_parallel_logp.RocmVocabParallelLogprobOp" ) TRITON_VOCAB_PARALLEL_LOGP = ( - "rl_engine.backends.shared.triton.logprob.vocab_parallel_logp.TritonVocabParallelLogprobOp" + "rl_engine.kernels.ops.triton.loss.vocab_parallel_logp.TritonVocabParallelLogprobOp" ) # Deterministic GRPO loss on the TP-aware logprob path (WS2 #241 PR5) PYTORCH_DISTRIBUTED_GRPO_LOSS = ( - "rl_engine.reference.loss.distributed_grpo_loss.DistributedGRPOLossOp" + "rl_engine.kernels.ops.pytorch.loss.distributed_grpo_loss.DistributedGRPOLossOp" ) # Ascend NPU deterministic (batch-invariant, no split-K) standard-softmax # attention (issue #147); Ascend C forward + reference backward. ASCEND_DETERMINISTIC_ATTENTION = ( - "rl_engine.backends.ascend.attention.deterministic_attn.DeterministicAttentionAscendOp" + "rl_engine.kernels.ops.ascend.attention.deterministic_attn.DeterministicAttentionAscendOp" ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) - TRITON_RMS_NORM = "rl_engine.backends.shared.triton.norm.rmsnorm.RMSNormTritonOp" - PYTORCH_NATIVE_RMS_NORM = "rl_engine.reference.norm.rms_norm.NativeRMSNormOp" + TRITON_RMS_NORM = "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp" + PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" # Generic fallback TRITON_GENERIC = "rl_engine.kernels.ops.triton.generic.TritonOp" PYTORCH_ATTN = "rl_engine.kernels.ops.pytorch.attention.NativeAttentionOp" - TRITON_LOGP = "rl_engine.backends.shared.triton.logprob.logp.TritonLogpOp" - PYTORCH_NATIVE = "rl_engine.reference.logprob.logp.NativeLogpOp" - PYTORCH_NATIVE_MATMUL = "rl_engine.reference.embedding.matmul.NativeMatmulOp" - PYTORCH_NATIVE_ROPE = "rl_engine.reference.rope.rope.NativeRoPEOp" - TRITON_ROPE = "rl_engine.backends.shared.triton.rope.rope.TritonRoPEOp" - CUDA_ROPE_SM90 = "rl_engine.backends.cuda.rope.rope.RoPESM90Op" - ASCEND_ROPE = "rl_engine.backends.ascend.rope.rope.RoPEAscendOp" - PYTORCH_NATIVE_SILU = "rl_engine.reference.activation.swiglu.NativeSiLUOp" - PYTORCH_NATIVE_SWIGLU = "rl_engine.reference.activation.swiglu.NativeSwiGLUOp" - CUDA_SILU = "rl_engine.backends.cuda.activation.swiglu.SiLUCudaOp" - CUDA_SWIGLU = "rl_engine.backends.cuda.activation.swiglu.SwiGLUCudaOp" - ASCEND_SWIGLU = "rl_engine.backends.ascend.activation.swiglu.SwiGLUAscendOp" - ASCEND_SILU = "rl_engine.backends.ascend.activation.silu.SiLUAscendOp" - TRITON_SILU = "rl_engine.backends.shared.triton.activation.swiglu.TritonSiLUOp" - TRITON_SWIGLU = "rl_engine.backends.shared.triton.activation.swiglu.TritonSwiGLUOp" + TRITON_LOGP = "rl_engine.kernels.ops.triton.loss.logp.TritonLogpOp" + PYTORCH_NATIVE = "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp" + PYTORCH_NATIVE_MATMUL = "rl_engine.kernels.ops.pytorch.linear.matmul.NativeMatmulOp" + PYTORCH_NATIVE_ROPE = "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp" + TRITON_ROPE = "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp" + CUDA_ROPE_SM90 = "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op" + ASCEND_ROPE = "rl_engine.kernels.ops.ascend.rotary_embedding.rope.RoPEAscendOp" + PYTORCH_NATIVE_SILU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp" + PYTORCH_NATIVE_SWIGLU = "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp" + CUDA_SILU = "rl_engine.kernels.ops.cuda.activation.swiglu.SiLUCudaOp" + CUDA_SWIGLU = "rl_engine.kernels.ops.cuda.activation.swiglu.SwiGLUCudaOp" + ASCEND_SWIGLU = "rl_engine.kernels.ops.ascend.activation.swiglu.SwiGLUAscendOp" + ASCEND_SILU = "rl_engine.kernels.ops.ascend.activation.silu.SiLUAscendOp" + 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" ) @@ -186,22 +187,26 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): # 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 = "rl_engine.reference.attention.standard_attn.NativeAttentionOp" + PYTORCH_NATIVE_ATTENTION = ( + "rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp" + ) # WS1 pure-PyTorch ground-truth KV-cache (decode/incremental) attention # reference; concats cache+new then reuses the standard attention reduction. - PYTORCH_NATIVE_KV_CACHE_ATTN = "rl_engine.reference.attention.kv_cache.NativeKVCacheAttnOp" + PYTORCH_NATIVE_KV_CACHE_ATTN = ( + "rl_engine.kernels.ops.pytorch.attention.kv_cache.NativeKVCacheAttnOp" + ) # WS2 correctness-first context-parallel attention reference. It emulates # CP prefill/chunked-prefill with fp32 attention-domain LSE merges. PYTORCH_CP_ATTENTION = ( - "rl_engine.reference.attention.cp_attention.DeterministicCPAttentionReferenceOp" + "rl_engine.kernels.ops.pytorch.attention.cp_attention.DeterministicCPAttentionReferenceOp" ) # WS1 pure-PyTorch ground-truth linear ops - PYTORCH_NATIVE_LM_HEAD = "rl_engine.reference.embedding.lm_head.NativeLMHeadOp" + PYTORCH_NATIVE_LM_HEAD = "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp" # WS1 pure-PyTorch ground-truth embedding ops - TRITON_EMBEDDING = "rl_engine.backends.shared.triton.embedding.embedding.TritonEmbeddingOp" - PYTORCH_NATIVE_EMBEDDING = "rl_engine.reference.embedding.embedding.NativeEmbeddingOp" - CUDA_SM90_LM_HEAD = "rl_engine.backends.cuda.embedding.lm_head.SM90LMHeadOp" - CUDA_SM90_EMBEDDING = "rl_engine.backends.cuda.embedding.embedding.SM90EmbeddingOp" + TRITON_EMBEDDING = "rl_engine.kernels.ops.triton.linear.embedding.TritonEmbeddingOp" + PYTORCH_NATIVE_EMBEDDING = "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp" + CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp" + CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -230,7 +235,9 @@ def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: "strict_observable": True, }, lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, - implementation_class_or_factory=("rl_engine.reference.logprob.logp.NativeLogpOp"), + implementation_class_or_factory=( + "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp" + ), fallback_policy=OperatorFallbackPolicy.ERROR, version_or_build_fingerprint="NativeLogpOp-selected-logprob-v1", ), @@ -1064,7 +1071,6 @@ def get_adaln_modulation_op( hidden: int, modulate_index: torch.Tensor | None = None, strict: bool = False, - arithmetic_profile: str | None = None, ) -> tuple[Any, dict[str, Any]]: """Resolve AdaLN; strict permits only the shared declared replay profile.""" from rl_engine.runtime.adaln import ( @@ -1077,8 +1083,6 @@ def get_adaln_modulation_op( if hidden <= 0: raise ValueError("hidden must be positive") profile = SELECT01_PROFILE if modulate_index is not None else SHARED_PROFILE - if arithmetic_profile is not None and arithmetic_profile != profile: - raise ValueError(f"unsupported AdaLN arithmetic profile {arithmetic_profile!r}") if strict and modulate_index is not None: raise RuntimeError("strict AdaLN rejects select01 framework backward") platform = self._platform_for_device(device) diff --git a/tests/runtime/test_adaln_dispatch.py b/tests/runtime/test_adaln_dispatch.py index 9fe87655d..1ded3631e 100644 --- a/tests/runtime/test_adaln_dispatch.py +++ b/tests/runtime/test_adaln_dispatch.py @@ -156,10 +156,8 @@ def forward(self, *args, **kwargs): assert observed["fallback"] is True and observed["fallback_reason"] -def test_strict_rejects_unknown_profile_and_indexed_backward(): +def test_strict_rejects_indexed_backward(): registry = KernelRegistry() - with pytest.raises(ValueError, match="arithmetic profile"): - registry.get_adaln_modulation_op("cpu", hidden=2, strict=True, arithmetic_profile="unknown") with pytest.raises(RuntimeError, match="select01"): registry.get_adaln_modulation_op( "cpu", hidden=2, strict=True, modulate_index=torch.tensor([[0]]) From 3837c4fa76b5265f660926f67385c3f099ba7a34 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 21 Oct 2026 15:15:23 +0800 Subject: [PATCH 21/37] refactor(validation): keep AdaLN policy out of generic operator checks Remove automatic AdaLN provenance observation and retain generic trace snapshots, FP32 reference inputs and output dtype checks. Preserve comparison on the original device when candidate and gold share it. Validation: op_checks 19 passed; AdaLN subset 333 passed; targeted mypy and Ruff passed. Full pytest stopped during collection on an unchanged ROCm import of _AITER_FWD_REQUIRED_KEYWORDS. Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 9 ++- rl_engine/validation/operators/op_checks.py | 78 +++---------------- tests/runtime/test_adaln_dispatch.py | 77 +----------------- .../validation/operators/test_adaln_gtest.py | 11 ++- tests/validation/operators/test_op_checks.py | 33 ++++++++ 5 files changed, 57 insertions(+), 151 deletions(-) diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 5224ae588..6d7ad0958 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -136,10 +136,11 @@ path, not a profiler measurement or independent proof of backward execution. `profiler_observed=false` and `hardware_qualified=null` make that limit explicit. Semantic versions/fingerprints are not binary hashes. -`check_operator.py` observes AdaLN candidates and includes per-case backend traces. -Successful BF16 CUDA/Triton calls also populate the existing C1 -`BackendProvenance`; CPU and FP32 have no declared C1 backend profile and remain -trace-only. The shared tolerance contract and independent reference are unchanged. +`check_operator.py` includes per-case AdaLN backend trace snapshots without +interpreting their dispatch policy or automatically constructing C1 +`BackendProvenance`. Explicit candidate provenance remains supported by the common +validator. AdaLN dispatch policy is covered by its dedicated runtime tests; +the shared tolerance contract and independent reference are unchanged. Library callers can request this with `make_candidate(args, observe_backend=True)`. The benchmark accepts `--seed` (default 386) and `--json`. JSON includes raw event diff --git a/rl_engine/validation/operators/op_checks.py b/rl_engine/validation/operators/op_checks.py index a6f889616..c8d511545 100644 --- a/rl_engine/validation/operators/op_checks.py +++ b/rl_engine/validation/operators/op_checks.py @@ -4,7 +4,7 @@ from __future__ import annotations from collections.abc import Callable, Mapping, Sequence -from dataclasses import asdict, dataclass, replace +from dataclasses import asdict, dataclass from typing import Any import torch @@ -74,7 +74,6 @@ class CaseCheck: passed: bool outputs: list[OutputCheck] backend_trace: dict[str, Any] | None = None - backend_provenance: BackendProvenance | None = None @dataclass(frozen=True) @@ -187,9 +186,6 @@ def _run_candidate( total_outputs = sum(len(case.outputs) for case in case_checks) passed_outputs = sum(1 for case in case_checks for output in case.outputs if output.passed) pass_rate = float(passed_outputs / total_outputs) if total_outputs else 0.0 - observed_provenance = case_checks[0].backend_provenance if case_checks else None - if any(case.backend_provenance != observed_provenance for case in case_checks): - observed_provenance = None return CandidateReport( candidate_name=candidate.name, backend=candidate.backend, @@ -198,7 +194,7 @@ def _run_candidate( pass_rate=pass_rate, passed=passed_outputs == total_outputs, cases=case_checks, - backend_provenance=candidate.provenance or observed_provenance, + backend_provenance=candidate.provenance, ) @@ -209,66 +205,9 @@ def _run_case( ) -> CaseCheck: candidate_outputs = _flatten_tensors(_call_candidate(candidate.fn, case.inputs)) gold_outputs = _flatten_tensors(case.gold_fn(**case.inputs)) - candidate = _observe_adaln_candidate(candidate, case, gold_outputs, contract) return _compare_case_outputs(candidate, case, contract, candidate_outputs, gold_outputs) -def _observe_adaln_candidate( - candidate: CandidateSpec, - case: OperatorCase, - gold_outputs: list[torch.Tensor], - contract: Mapping[str, Any], -) -> CandidateSpec: - from rl_engine.runtime.adaln import SHARED_PROFILE, AdaLNDispatch - - if not isinstance(candidate.fn, AdaLNDispatch): - return candidate - trace = candidate.fn.trace - actual = trace["actual_backend"] - if actual != candidate.backend: - raise ContractResolveError( - f"candidate backend {candidate.backend!r} differs from {actual!r}" - ) - if normalize_dtype_name(case.dtype) != trace["execution_dtype"]: - raise ContractResolveError("case dtype differs from observed AdaLN execution dtype") - reference_dtypes = {normalize_dtype_name(output.dtype) for output in gold_outputs} - if reference_dtypes != {"float32"}: - raise ContractResolveError("AdaLN reference dtype must be observed FP32 for every output") - # C1 declares BF16 GPU profiles only; CPU/FP32 observations remain case traces. - if actual not in {"cuda", "triton"} or trace["execution_dtype"] != "bfloat16": - return candidate - fallback = trace["requested_backend"] != actual - if fallback and not ( - trace["fallback"] - and trace["requested_backend"] == "cuda" - and actual == "triton" - and trace["arithmetic_profile"] == SHARED_PROFILE - and trace["observation"] == "successful_forward" - and trace["fallback_reason"] - ): - raise ContractResolveError("undeclared AdaLN backend fallback") - # Validate the executed backend's profile; request history stays in backend_trace. - provenance = BackendProvenance( - backend_profile="cuda_bf16" if actual == "cuda" else "triton_cuda_bf16", - requested_backend=actual, - actual_backend=actual, - execution_dtype=trace["execution_dtype"], - accumulation_dtype="float32", - output_dtype=trace["output_dtype"], - reference_dtype=normalize_dtype_name(gold_outputs[0].dtype), - candidate_tf32_enabled=False, - reference_tf32_enabled=False, - ) - validate_backend_provenance(contract, provenance) - if fallback: - if candidate.provenance is not None: - raise ContractResolveError("AdaLN fallback must be reported through backend_trace") - return candidate - if candidate.provenance is not None and candidate.provenance != provenance: - raise ContractResolveError("declared AdaLN provenance differs from successful execution") - return replace(candidate, provenance=provenance) - - def _run_case_backward( candidate: CandidateSpec, case: OperatorCase, @@ -307,7 +246,6 @@ def _run_case_backward( case.grad_input_names, grad_outputs=_match_grad_outputs(shared_upstreams, gold_outputs), ) - candidate = _observe_adaln_candidate(candidate, case, gold_outputs, contract) output_checks = _compare_case_outputs( candidate, case, @@ -369,7 +307,6 @@ def _run_case_backward( passed=all(output.passed for output in checks), outputs=checks, backend_trace=dict(candidate.fn.trace) if hasattr(candidate.fn, "trace") else None, - backend_provenance=candidate.provenance, ) @@ -452,7 +389,6 @@ def _compare_case_outputs( 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, - backend_provenance=candidate.provenance, ) @@ -596,7 +532,8 @@ def _resolve_tolerance( # They only mirror forward accuracy thresholds; never apply them to grads. if judgment != "forward_accuracy": raise ContractResolveError( - f"legacy accuracy contracts only support judgment='forward_accuracy'; got {judgment!r}" + f"legacy accuracy contracts only support judgment='forward_accuracy'; " + f"got {judgment!r}" ) dtype_name = normalize_dtype_name(dtype) if arch_key is not None: @@ -645,9 +582,12 @@ def _compare_output( message=f"shape mismatch: candidate={tuple(candidate.shape)} gold={tuple(gold.shape)}", ) + candidate_fp32 = candidate.detach().float() + gold_fp32 = gold.detach().float() # Independent references may live on CPU even when the candidate runs on GPU. - candidate_fp32 = candidate.detach().to(device="cpu", dtype=torch.float32) - gold_fp32 = gold.detach().to(device="cpu", dtype=torch.float32) + 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 diff --git a/tests/runtime/test_adaln_dispatch.py b/tests/runtime/test_adaln_dispatch.py index 1ded3631e..6ab700e4e 100644 --- a/tests/runtime/test_adaln_dispatch.py +++ b/tests/runtime/test_adaln_dispatch.py @@ -6,7 +6,6 @@ import torch from argparse import Namespace -from dataclasses import replace from rl_engine.validation.operators import CandidateSpec, run_operator_suite from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case @@ -68,31 +67,6 @@ def test_gtest_records_cpu_calls_without_inventing_a_bf16_hardware_profile(): assert trace["profiler_observed"] is False -@pytest.mark.parametrize("invalid", ["reference_dtype", "case_dtype"]) -def test_observed_gtest_rejects_false_reference_and_case_dtype_claims(invalid): - args = Namespace( - op="adaln_modulation", - candidate="pytorch", - arch_key=None, - batch=2, - seq=3, - normalized_dim=7, - seed=386, - ) - case = make_operator_case(args, torch.bfloat16, torch.device("cpu")) - case = ( - replace(case, gold_fn=NativeAdaLNModulationOp()) - if invalid == "reference_dtype" - else replace(case, dtype=torch.float32) - ) - with pytest.raises(ContractResolveError, match="reference dtype|case dtype"): - run_operator_suite( - "adaln_modulation", - candidates=[make_candidate(args, observe_backend=True)], - cases=[case], - ) - - def test_common_provenance_rejects_cuda_to_triton_fallback(): proof = BackendProvenance( backend_profile="triton_cuda_bf16", @@ -109,53 +83,6 @@ def test_common_provenance_rejects_cuda_to_triton_fallback(): validate_backend_provenance(load_contract(), proof) -@pytest.mark.parametrize( - "override", - [ - {}, - {"fallback": False}, - {"observation": None}, - {"arithmetic_profile": "unknown"}, - {"fallback_reason": None}, - {"requested_backend": "pytorch"}, - ], -) -def test_gtest_validates_and_serializes_fallback_trace_on_cpu(monkeypatch, override): - from rl_engine.runtime import adaln - - # Simulate backend observations to exercise report policy without a GPU kernel. - monkeypatch.setattr(adaln, "adaln_backend", lambda op: "triton") - - class ObservedDispatch(adaln.AdaLNDispatch): - def forward(self, *args, **kwargs): - outputs = super().forward(*args, **kwargs) - self.trace.update(override) - return outputs - - trace = { - "requested_backend": "cuda", - "fallback": True, - "fallback_reason": "CUDA extension unavailable", - "arithmetic_profile": adaln.SHARED_PROFILE, - } - op = ObservedDispatch(NativeAdaLNModulationOp(), trace, hidden=7) - args = Namespace(op="adaln_modulation", batch=2, seq=3, normalized_dim=7, seed=386) - case = make_operator_case(args, torch.bfloat16, torch.device("cpu")) - candidate = CandidateSpec(name="observed-adaln", backend="triton", fn=op) - if override: - with pytest.raises(ContractResolveError, match="undeclared AdaLN backend fallback"): - run_operator_suite("adaln_modulation", candidates=[candidate], cases=[case]) - else: - report = run_operator_suite("adaln_modulation", candidates=[candidate], cases=[case]) - assert report.passed - result = report.to_dict()["candidates"][0] - assert result["backend_provenance"] is None - observed = result["cases"][0]["backend_trace"] - assert observed["requested_backend"] == "cuda" - assert observed["actual_backend"] == "triton" - assert observed["fallback"] is True and observed["fallback_reason"] - - def test_strict_rejects_indexed_backward(): registry = KernelRegistry() with pytest.raises(RuntimeError, match="select01"): @@ -273,7 +200,7 @@ def unavailable(backend): 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_provenance(monkeypatch): +def test_gpu_gtest_persists_observed_bf16_fallback_trace(monkeypatch): pytest.importorskip("triton") args = Namespace( op="adaln_modulation", @@ -303,7 +230,7 @@ def test_gpu_gtest_persists_observed_bf16_fallback_provenance(monkeypatch): result = report.candidates[0] assert result.backend_provenance is None payload = report.to_dict()["candidates"][0]["cases"][0] - assert payload["backend_provenance"] is None + 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 diff --git a/tests/validation/operators/test_adaln_gtest.py b/tests/validation/operators/test_adaln_gtest.py index dc2719764..0fcfa5af7 100644 --- a/tests/validation/operators/test_adaln_gtest.py +++ b/tests/validation/operators/test_adaln_gtest.py @@ -529,13 +529,16 @@ def corrupt(x, modulation, eps): @pytest.mark.parametrize("check_grad", [False, True]) -def test_adaln_gtest_rejects_fp32_outputs_for_bf16_inputs(check_grad): +@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): - return tuple(output.float() for output in good.fn(**inputs)) + outputs = list(good.fn(**inputs)) + outputs[wrong_output] = outputs[wrong_output].float() + return tuple(outputs) report = run_operator_suite( args.op, @@ -545,7 +548,9 @@ def wrong_dtype(**inputs): ) assert not report.passed, report.to_dict() checks = report.candidates[0].cases[0].outputs[:2] - assert all(not check.passed and "dtype mismatch" in check.message for check in checks) + 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]) diff --git a/tests/validation/operators/test_op_checks.py b/tests/validation/operators/test_op_checks.py index 26a84e35b..b72e82959 100644 --- a/tests/validation/operators/test_op_checks.py +++ b/tests/validation/operators/test_op_checks.py @@ -112,6 +112,39 @@ 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): From ebd5e6fccc5c6aa7cb6904fb7caf138a024526cc Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 21 Oct 2026 15:22:04 +0800 Subject: [PATCH 22/37] refactor(validation): restore unrelated operator specs to baseline Restore non-AdaLN gold and candidate paths and unrelated import formatting to the fixed integration baseline. Preserve AdaLN registration, precision settings, independent reference and optional backend observation. Validation: 65 focused tests passed; mypy, Ruff and diff checks passed. Both review axes reported no findings. Full pytest stopped during collection on the existing ROCm _AITER_FWD_REQUIRED_KEYWORDS import error. Signed-off-by: hongleng-github --- .../validation/operators/operator_specs.py | 167 ++++++++++-------- 1 file changed, 89 insertions(+), 78 deletions(-) diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 503b19454..901d14800 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -11,10 +11,7 @@ import torch from rl_engine.validation.operators.op_checks import CandidateSpec, OperatorCase -from rl_engine.validation.operators.operator_inputs import ( - make_operator_inputs, - operator_shape_name, -) +from rl_engine.validation.operators.operator_inputs import make_operator_inputs, operator_shape_name @dataclass(frozen=True) @@ -58,48 +55,52 @@ def _load_object(path: str) -> Any: "rms_norm": OperatorSpec( name="rms_norm", op_class="reduction", - gold_path="rl_engine.reference.norm.rms_norm.NativeRMSNormOp", + gold_path="rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.reference.norm.rms_norm.NativeRMSNormOp", - "triton": "rl_engine.backends.shared.triton.norm.rmsnorm.RMSNormTritonOp", - "cuda": "rl_engine.backends.cuda.norm.rmsnorm.RMSNormCudaOp", - "cuda-sm90": "rl_engine.backends.cuda.norm.rmsnorm.RMSNormCudaOp", - "ascend": "rl_engine.backends.ascend.norm.rmsnorm.RMSNormAscendOp", + "pytorch": "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp", + "triton": "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp", + "cuda": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", + "ascend": "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp", }, grad_input_names=("x", "weight"), ), "qk_norm": OperatorSpec( name="qk_norm", op_class="reduction", - gold_path="rl_engine.reference.norm.rms_norm.NativeRMSNormOp", + gold_path="rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.reference.norm.rms_norm.NativeRMSNormOp", - "triton": "rl_engine.backends.shared.triton.norm.rmsnorm.RMSNormTritonOp", - "cuda": "rl_engine.backends.cuda.norm.rmsnorm.RMSNormCudaOp", - "cuda-sm90": "rl_engine.backends.cuda.norm.rmsnorm.RMSNormCudaOp", - "ascend": "rl_engine.backends.ascend.norm.rmsnorm.RMSNormAscendOp", + "pytorch": "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp", + "triton": "rl_engine.kernels.ops.triton.rmsnorm_triton.RMSNormTritonOp", + "cuda": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.norm.rmsnorm.RMSNormCudaOp", + "ascend": "rl_engine.kernels.ops.ascend.norm.rmsnorm.RMSNormAscendOp", }, grad_input_names=("x", "weight"), ), "attention": OperatorSpec( name="attention", op_class="attention", - gold_path="rl_engine.reference.attention.standard_attn.NativeAttentionOp", + gold_path="rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.reference.attention.standard_attn.NativeAttentionOp", + "pytorch": "rl_engine.kernels.ops.pytorch.attention.standard_attn.NativeAttentionOp", "triton": ( - "rl_engine.backends.shared.triton.attention.standard_attn.TritonBatchInvariantAt" - "tentionOp" + ( + "rl_engine.kernels.ops.triton.attention.standard_attn.TritonBatchInvariantAt" + "tentionOp" + ) ), "cuda": ( - "rl_engine.backends.cuda.attention.deterministic_attn.DeterministicAttentionOp" + "rl_engine.kernels.ops.cuda.attention.deterministic_attn.DeterministicAttentionOp" ), "ascend": ( - "rl_engine.backends.ascend.attention.deterministic_attn.DeterministicAtte" - "ntionAscendOp" + ( + "rl_engine.kernels.ops.ascend.attention.deterministic_attn.DeterministicAtte" + "ntionAscendOp" + ) ), }, grad_input_names=("q", "k", "v"), @@ -115,11 +116,13 @@ def _load_object(path: str) -> Any: candidate_paths={ "pytorch": "rl_engine.validation.operators.operator_specs.GtestPrefixSharedAttentionOp", "cuda": ( - "rl_engine.backends.cuda.attention.prefix_shared_attn.PrefixSharedAttentionOp" + "rl_engine.kernels.ops.cuda.attention.prefix_shared_attn.PrefixSharedAttentionOp" ), "ascend": ( - "rl_engine.backends.ascend.attention.prefix_shared_attn.PrefixSharedAtten" - "tionAscendOp" + ( + "rl_engine.kernels.ops.ascend.attention.prefix_shared_attn.PrefixSharedAtten" + "tionAscendOp" + ) ), }, ), @@ -127,12 +130,18 @@ def _load_object(path: str) -> Any: name="cp_attention", op_class="attention", gold_path=( - "rl_engine.reference.attention.cp_attention.DeterministicCPAttentionReferenceOp" + ( + "rl_engine.kernels.ops.pytorch.attention.cp_attention.DeterministicCPAttenti" + "onReferenceOp" + ) ), gold_method="forward_fp32", candidate_paths={ "pytorch": ( - "rl_engine.reference.attention.cp_attention.DeterministicCPAttentionReferenceOp" + ( + "rl_engine.kernels.ops.pytorch.attention.cp_attention.DeterministicCPAttenti" + "onReferenceOp" + ) ), }, grad_input_names=("q", "k", "v"), @@ -140,54 +149,54 @@ def _load_object(path: str) -> Any: "logp": OperatorSpec( name="logp", op_class="logprob", - gold_path="rl_engine.reference.logprob.logp.NativeLogpOp", + gold_path="rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.reference.logprob.logp.NativeLogpOp", - "triton": "rl_engine.backends.shared.triton.logprob.logp.TritonLogpOp", - "cuda": "rl_engine.backends.cuda.logprob.logp.FusedLogpGenericOp", - "cuda-generic": "rl_engine.backends.cuda.logprob.logp.FusedLogpGenericOp", - "cuda-sm90": "rl_engine.backends.cuda.logprob.logp.FusedLogpSM90Op", - "ascend": "rl_engine.backends.ascend.logprob.logp.FusedLogpAscendOp", + "pytorch": "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp", + "triton": "rl_engine.kernels.ops.triton.loss.logp.TritonLogpOp", + "cuda": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp", + "cuda-generic": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpGenericOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.logp.FusedLogpSM90Op", + "ascend": "rl_engine.kernels.ops.ascend.loss.logp.FusedLogpAscendOp", }, grad_input_names=("logits",), ), "linear_logp": OperatorSpec( name="linear_logp", op_class="logprob", - gold_path="rl_engine.reference.logprob.linear_logp.NativeLinearLogpOp", + gold_path="rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp", gold_method="apply", candidate_paths={ - "pytorch": "rl_engine.reference.logprob.linear_logp.NativeLinearLogpOp", - "triton": "rl_engine.backends.shared.triton.logprob.linear_logp.TritonLinearLogpOp", - "cuda-sm90": "rl_engine.backends.cuda.logprob.linear_logp.FusedLinearLogpSM90Op", - "ascend": "rl_engine.backends.ascend.logprob.linear_logp.FusedLinearLogpAscendOp", + "pytorch": "rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp", + "triton": "rl_engine.kernels.ops.triton.loss.linear_logp.TritonLinearLogpOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.loss.linear_logp.FusedLinearLogpSM90Op", + "ascend": "rl_engine.kernels.ops.ascend.loss.linear_logp.FusedLinearLogpAscendOp", }, grad_input_names=("hidden", "lm_head_weight"), ), "embedding": OperatorSpec( name="embedding", op_class="elementwise", - gold_path="rl_engine.reference.embedding.embedding.NativeEmbeddingOp", + gold_path="rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.reference.embedding.embedding.NativeEmbeddingOp", - "triton": "rl_engine.backends.shared.triton.embedding.embedding.TritonEmbeddingOp", - "cuda-sm90": "rl_engine.backends.cuda.embedding.embedding.SM90EmbeddingOp", - "ascend": "rl_engine.backends.ascend.embedding.embedding.AscendEmbeddingOp", + "pytorch": "rl_engine.kernels.ops.pytorch.linear.embedding.NativeEmbeddingOp", + "triton": "rl_engine.kernels.ops.triton.linear.embedding.TritonEmbeddingOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp", + "ascend": "rl_engine.kernels.ops.ascend.linear.embedding.AscendEmbeddingOp", }, grad_input_names=("weight",), ), "lm_head": OperatorSpec( name="lm_head", op_class="reduction", - gold_path="rl_engine.reference.embedding.lm_head.NativeLMHeadOp", + gold_path="rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.reference.embedding.lm_head.NativeLMHeadOp", - "triton": "rl_engine.backends.shared.triton.embedding.lm_head.TritonLMHeadOp", - "cuda-sm90": "rl_engine.backends.cuda.embedding.lm_head.SM90LMHeadOp", - "ascend": "rl_engine.backends.ascend.embedding.lm_head.AscendLMHeadOp", + "pytorch": "rl_engine.kernels.ops.pytorch.linear.lm_head.NativeLMHeadOp", + "triton": "rl_engine.kernels.ops.triton.linear.lm_head.TritonLMHeadOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp", + "ascend": "rl_engine.kernels.ops.ascend.linear.lm_head.AscendLMHeadOp", }, grad_input_names=("hidden", "weight"), ), @@ -198,73 +207,77 @@ def _load_object(path: str) -> Any: # the accuracy gold must be the same leaf-space tree, not the # single-rounding torch.matmul (which fails structurally at # near-cancellation outputs on random inputs). - gold_path="rl_engine.reference.gemm.det_gemm.DetGemmTreeReferenceOp", + gold_path="rl_engine.kernels.ops.pytorch.matmul.det_gemm.DetGemmTreeReferenceOp", gold_method="__call__", candidate_paths={ - "pytorch": "rl_engine.reference.gemm.det_gemm.NativeGemmOp", - "cuda": "rl_engine.backends.cuda.gemm.det_gemm.DetGemmOp", - "triton": "rl_engine.backends.shared.triton.gemm.det_gemm.TritonDetGemmOp", - "ascend": "rl_engine.backends.ascend.gemm.det_gemm.DetGemmAscendOp", + "pytorch": "rl_engine.kernels.ops.pytorch.matmul.det_gemm.NativeGemmOp", + "cuda": "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp", + "triton": "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp", + "ascend": "rl_engine.kernels.ops.ascend.matmul.det_gemm.DetGemmAscendOp", }, grad_input_names=("a", "b"), ), "rope": OperatorSpec( name="rope", op_class="elementwise", - gold_path="rl_engine.reference.rope.rope.NativeRoPEOp", + gold_path="rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.reference.rope.rope.NativeRoPEOp", - "triton": "rl_engine.backends.shared.triton.rope.rope.TritonRoPEOp", - "cuda-sm90": "rl_engine.backends.cuda.rope.rope.RoPESM90Op", - "ascend": "rl_engine.backends.ascend.rope.rope.RoPEAscendOp", + "pytorch": "rl_engine.kernels.ops.pytorch.rotary_embedding.rope.NativeRoPEOp", + "triton": "rl_engine.kernels.ops.triton.rotary_embedding.rope.TritonRoPEOp", + "cuda-sm90": "rl_engine.kernels.ops.cuda.rotary_embedding.rope.RoPESM90Op", + "ascend": "rl_engine.kernels.ops.ascend.rotary_embedding.rope.RoPEAscendOp", }, grad_input_names=("x",), ), "silu": OperatorSpec( name="silu", op_class="elementwise", - gold_path="rl_engine.reference.activation.swiglu.NativeSiLUOp", + gold_path="rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.reference.activation.swiglu.NativeSiLUOp", - "triton": "rl_engine.backends.shared.triton.activation.swiglu.TritonSiLUOp", - "cuda": "rl_engine.backends.cuda.activation.swiglu.SiLUCudaOp", - "ascend": "rl_engine.backends.ascend.activation.silu.SiLUAscendOp", + "pytorch": "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSiLUOp", + "triton": "rl_engine.kernels.ops.triton.activation.swiglu.TritonSiLUOp", + "cuda": "rl_engine.kernels.ops.cuda.activation.swiglu.SiLUCudaOp", + "ascend": "rl_engine.kernels.ops.ascend.activation.silu.SiLUAscendOp", }, grad_input_names=("x",), ), "swiglu": OperatorSpec( name="swiglu", op_class="elementwise", - gold_path="rl_engine.reference.activation.swiglu.NativeSwiGLUOp", + gold_path="rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp", gold_method="forward_fp32", candidate_paths={ - "pytorch": "rl_engine.reference.activation.swiglu.NativeSwiGLUOp", - "triton": "rl_engine.backends.shared.triton.activation.swiglu.TritonSwiGLUOp", - "cuda": "rl_engine.backends.cuda.activation.swiglu.SwiGLUCudaOp", - "ascend": "rl_engine.backends.ascend.activation.swiglu.SwiGLUAscendOp", + "pytorch": "rl_engine.kernels.ops.pytorch.activation.swiglu.NativeSwiGLUOp", + "triton": "rl_engine.kernels.ops.triton.activation.swiglu.TritonSwiGLUOp", + "cuda": "rl_engine.kernels.ops.cuda.activation.swiglu.SwiGLUCudaOp", + "ascend": "rl_engine.kernels.ops.ascend.activation.swiglu.SwiGLUAscendOp", }, grad_input_names=("gate", "up"), ), "batch_invariant_logp": OperatorSpec( name="batch_invariant_logp", op_class="logprob", - gold_path=("rl_engine.reference.logprob.batch_invariant_logp.NativeBatchInvariantLogpOp"), + gold_path=( + "rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp.NativeBatchInvarian" "tLogpOp" + ), gold_method="apply", candidate_paths={ "pytorch": ( - "rl_engine.reference.logprob.batch_invariant_logp.NativeBatchInvariantLogpOp" + "rl_engine.kernels.ops.pytorch.loss.batch_invariant_logp.NativeBatchInvarian" + "tLogpOp" ), "triton": ( - "rl_engine.backends.shared.triton.logprob.batch_invariant_logp.TritonBatchInvariant" + "rl_engine.kernels.ops.triton.loss.batch_invariant_logp.TritonBatchInvariant" "LogpOp" ), "cuda-sm90": ( - "rl_engine.backends.cuda.logprob.batch_invariant_logp.BatchInvariantLogpSM90Op" + "rl_engine.kernels.ops.cuda.loss.batch_invariant_logp.BatchInvariantLogpSM90" "Op" ), "ascend": ( - "rl_engine.backends.ascend.logprob.batch_invariant_logp.BatchInvariantLogpAscendOp" + "rl_engine.kernels.ops.ascend.loss.batch_invariant_logp.BatchInvariantLogpAs" + "cendOp" ), }, grad_input_names=("logits",), @@ -358,9 +371,7 @@ class GtestPrefixSharedAttentionOp: op_class = "attention" def __init__(self) -> None: - from rl_engine.reference.attention.standard_attn import ( - NativeAttentionOp, - ) + from rl_engine.reference.attention.standard_attn import NativeAttentionOp self._op = NativeAttentionOp() From 111aa446d74374b11f9e9d3963b6a8048a675a4b Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 21 Oct 2026 15:35:27 +0800 Subject: [PATCH 23/37] docs(validation): add H100 validation runbook for AdaLN modulation Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 347 +++++++++++++---------------- docs/validation/adaln-h100.md | 109 +++++++++ 2 files changed, 258 insertions(+), 198 deletions(-) create mode 100644 docs/validation/adaln-h100.md diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 6d7ad0958..79867afd7 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -1,223 +1,174 @@ -# AdaLN modulation (Qwen-Image WS1, #386) - -`kernel_registry.get_adaln_modulation_op(device, hidden=H)` returns `(op, trace)`. Call -`modulated, gate = op(x, modulation, eps=1e-6)` with `x: [B,S,H]` -and shared `modulation: [B,3H]`. Both inputs must share device and dtype (FP32 -or BF16). The chunks are **shift, scale, gate**. Each row uses LayerNorm over -H with no affine parameters, then `norm(x) * (1 + scale) + shift`. The gate is -`[B,1,H]` and remains connected to `modulation` for downstream gradients. -The Qwen-Image image/text streams and attention/MLP calls use separate inputs. -The `SiLU + Linear` projection and `adaln_gate_residual` are separate operators. - -The FP32 hidden reductions use a lower-plus-upper pairwise tree padded to the -next power of two. Each sample's `dshift` and `dscale` contributions are folded -in ascending logical token order. The CUDA shared reduction reads its result -into each thread before a block barrier permits buffer reuse. No atomic partial sums, Split-K, Stream-K, -TF32, or fast math are used by these kernels. BF16 outputs are cast once after -FP32 arithmetic. The CPU reference computes square root in FP64 and rounds it -to FP32 because CPU `torch.sqrt(float32)` can differ by one ULP from CUDA's -correctly rounded square root. Hidden and token accumulators remain FP32. -Inputs with S=0 or H=0 and mixed dtype/device are rejected. -Strided model chunks are copied to contiguous storage inside -GPU wrappers. The CUDA kernel supports H<=4096, B<=65535, and at most -2^31-1 elements in each input tensor. Python checks these bounds before copying -strided inputs; C++ checks them again before launch. Oversized B/address ranges -raise an explicit error; the registry only resolves using H and does not silently -fallback for these other limits. Shape-aware dispatch selects Triton for larger H. -CUDA symbols are omitted from builds with -`KERNEL_ALIGN_USE_FAST_MATH=1`, making fallback explicit. - -Shared eps must be finite and nonnegative and round to finite FP32. A positive -eps that rounds to zero is rejected; representable subnormals, including 2^-149, -remain supported. eps=0 remains supported; zero variance then produces IEEE NaN. -Nonfinite LayerNorm arithmetic follows the FP32 formula and may propagate NaNs; -its NaN payloads are outside the strict byte-equality qualification. Gate and -dgate use a bit-preserving passthrough, including nonfinite payloads. This shared -eps policy does not change the separate indexed/select01 fallback contract. - -For Qwen-Image editing, pass `modulate_index` to both resolution and execution: +# 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 -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) +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) ``` -This path accepts modulation `[2B,3H]` and an index `[B,S]` or `[1,S]` whose -values are 0 or 1. It uses the explicit PyTorch select01 reference and returns -gate `[B,S,H]`. The trace sets `fallback=true` and -`fallback_reason=modulate_index_requires_select01`; the indexed path does not -claim the fixed-order CUDA/Triton backward contract. - -The resolution matrix `{1024², 1328², 1664×928}` implies image sequence -lengths `{4096, 6889, 6032}` for VAE scale 8 and 2x2 latent packing. H=3072. -These are derived model shapes, not tensor shapes stated by the issue. - -## Validation - -Public-interface tests check BF16 cast boundaries, each reduction, signed zero, -subnormal eps, gate/dgate payloads, strided upstreams and H=3072 batch positions -with padding across all declared launch configurations. AdaLN gtest checks -candidate output dtype and promotes gold leaves to FP32 before autograd; -casting inside the gold function alone would return BF16 gradients to BF16 leaves. - -AdaLN is registered in `rl_engine/validation/operators/operator_specs.py` and `operator_inputs.py` as -`op_class="reduction"`, with PyTorch, CUDA and Triton candidates. The independent -gold uses FP32 CPU `torch.nn.functional.layer_norm` and autograd for normalization. -Its independent affine VJP folds token contributions in ascending FP32 order; -it does not call production reduction or backward helpers. It checks both outputs -(y and gate) and gradients of x and modulation. Accuracy thresholds come from -`resolve_tolerance`: `forward_accuracy` for outputs and `gradient_accuracy` -for input gradients; no private AdaLN thresholds or SM90 override are added. -Existing fixed-reference byte comparisons remain separate, unchanged gates. +Each token is normalized over H without learned LayerNorm affine parameters: -```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 -.venv/bin/python -m pytest tests/validation/operators/test_adaln_gtest.py -q -rs +```text +modulated = LayerNorm(x, eps) * (1 + scale) + shift ``` -Use CUDA/Triton candidates on the target GPU for native accuracy evidence. -The three full-image tests additionally check independent CPU FP32 forward and -backward accuracy for both GPU backends/dtypes. These assertions require GPU -execution; CPU results alone do not validate them. -CPU tests also cover gate-only gradients, distinct image/text modulation -chunks, corrupted-output/VJP rejection and unavailable-backend routing. -The CLI accuracy report alone is not a provenance-checked WS1 system gate. - -Run `python -m pytest tests/ops/norm/test_adaln_modulation.py -q` for focused tests. -Set `RLK_ADALN_REAL_SHAPES=1` to include the three full image shapes. -Run `python benchmarks/operators/norm/benchmark_adaln_modulation.py --real` for the three -image shapes, with `--dtype fp32` for the second dtype. The benchmark prints -median forward and forward+backward latency plus peak allocated memory; -`--warmup` and `--repeat` control the sample count. It explicitly skips the -CUDA extension if its symbols are absent. The registry trace records selected -backend, fallback, reduction order, accumulator dtype, disabled algorithm -flags, and a semantic kernel version. The semantic version is not a hash of the -compiled binary. Direct backend constructors accept validated launch settings: -CUDA `threads={128,256,512}` (default 256); Triton `num_warps={4,8}` -(default 4), `reduction_tile={64,128,256}` (default 128). Registry dispatch -uses those defaults. These settings change scheduling, not the reduction order. - -Native acceptance requires a fresh extension build, source/binary identity records, -and complete test and sanitizer logs. - -## Strict dispatch and execution observations +## 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], strict=True, + x.device, hidden=x.shape[-1], modulate_index=modulate_index ) -y, gate = op(x, modulation) +modulated, gate = op(x, modulation, modulate_index=modulate_index) ``` -Strict dispatch pins the shared FP32 hidden tree, ascending token fold and final -cast. NVIDIA CUDA may fall back to Triton when the extension is unavailable or -H exceeds 4096: both implement this same arithmetic profile, accept FP32/BF16, -and retain the existing launch-configuration replay tests. The arithmetic profile -is inferred from the requested operation and recorded in the trace. Strict rejects -unknown backends, the indexed/select01 framework backward and PyTorch execution -on GPU, whose cross-backend replay is not qualified here. It checks the actual -input device and hidden size again at execution. CPU Native is permitted for the -fixed reference. This policy does not itself certify any hardware/runtime profile. -Ordinary dispatch retains the prior CUDA/Triton/PyTorch and indexed fallbacks. - -The resolver still returns a callable and a trace; the callable is now a -resolution-local `AdaLNDispatch`, exposing its backend as `implementation`. -`selected_backend` is a resolution decision; `actual_backend`, execution dtype, -device and output dtype are filled only after a successful forward call. A failed -call clears the previous observation. `fallback_reason` and `arithmetic_profile` -identify routing and arithmetic separately. The trace describes the last call; -copy it when retaining an observation, and use separate resolutions for concurrent -callers. `observation="successful_forward"` is evidence of the invoked operator -path, not a profiler measurement or independent proof of backward execution. -`profiler_observed=false` and `hardware_qualified=null` make that limit explicit. -Semantic versions/fingerprints are not binary hashes. - -`check_operator.py` includes per-case AdaLN backend trace snapshots without -interpreting their dispatch policy or automatically constructing C1 -`BackendProvenance`. Explicit candidate provenance remains supported by the common -validator. AdaLN dispatch policy is covered by its dedicated runtime tests; -the shared tolerance contract and independent reference are unchanged. -Library callers can request this with `make_candidate(args, observe_backend=True)`. - -The benchmark accepts `--seed` (default 386) and `--json`. JSON includes raw event -samples, median, absolute peak allocated memory and 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 that file; -it does not assert that it was built from the recorded working tree. Missing -observations are null. Extension/Triton skips are explicit. -Backend observations happen before timing; the original event interval still -measures the bare operator, including newly allocated `ones_like` upstreams for -forward+backward. The added ordinary `F.layer_norm` + FP32 affine baseline uses -the same inputs and timing, but does not claim fixed-order replay. - -Run the new GPU cases explicitly on the target environment: +| 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 | -```bash -.venv/bin/python -m pytest tests/runtime/test_adaln_dispatch.py -k gpu -q -rs -.venv/bin/python benchmarks/operators/norm/benchmark_adaln_modulation.py \ - --real --dtype bf16 --seed 386 --warmup 10 --repeat 50 --json -``` +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. Indexed modulation always +selects the PyTorch select01 path, with `fallback=true` and reason +`modulate_index_requires_select01`. + +### Strict mode and trace -## Directory integration +Pass `strict=True` when resolving the shared path to require its declared arithmetic +profile. CPU Native and NVIDIA CUDA/Triton are permitted; indexed modulation, +unknown backends and PyTorch execution on GPU are rejected. -The implementation follows the upstream layout: reference in `rl_engine/reference/norm/`, -CUDA in `rl_engine/backends/cuda/norm/`, Triton in `rl_engine/backends/shared/triton/norm/`, -and dispatch policy in `rl_engine/runtime/adaln.py`. The shared Triton directory does -not claim ROCm qualification; strict GPU execution requires NVIDIA CUDA. -Native bindings live in `csrc/bindings/ops.cpp`; `build_tools/extensions.py` selects -`csrc/cuda/norm/adaln_modulation.cu` only for CUDA builds without fast math. -This directory migration preserves the arithmetic profile and precision thresholds. -GPU acceptance must be rerun against the migrated sources and a fresh extension. +Strict mode rechecks the input device and hidden size at execution. It enforces +the declared calculation policy, not hardware qualification. The indexed path +does not claim fixed-order CUDA/Triton backward equivalence. -## Migrated-source H100 validation commands +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. -Run these commands yourself from the repository root on H100/SM90. They are not -migration validation results. Keep the source revision, working-tree status, build -log and test/sanitizer logs together; the benchmark JSON also records source and -loaded extension hashes. A skip is not native CUDA acceptance. +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 modulation checks only finite, nonnegative eps. + +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. +The CLI records trace snapshots without automatically constructing C1 +`BackendProvenance`; its accuracy report alone is not a WS1 system gate. + +## Tests and Benchmark ```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 - -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 - -# Repeat each candidate for bf16 and fp32. -.venv/bin/python tools/validation/operators/check_operator.py --op adaln_modulation \ - --candidate cuda --device cuda --dtype bf16 --batch 2 --seq 3 \ - --normalized-dim 3072 --check-grad --json -.venv/bin/python tools/validation/operators/check_operator.py --op adaln_modulation \ - --candidate triton --device cuda --dtype bf16 --batch 2 --seq 3 \ +python -m pytest tests/ops/norm/test_adaln_modulation.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 -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 - -.venv/bin/python benchmarks/operators/norm/benchmark_adaln_modulation.py \ - --real --dtype bf16 --seed 386 --warmup 10 --repeat 50 --json -.venv/bin/python benchmarks/operators/norm/benchmark_adaln_modulation.py \ - --real --dtype fp32 --seed 386 --warmup 10 --repeat 50 --json +python benchmarks/operators/norm/benchmark_adaln_modulation.py --real --dtype bf16 --json ``` -The repository-wide suite is also left for manual execution because it includes -hardware and integration scopes whose duration is not bounded by this migration: +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. -```bash -.venv/bin/python -m pytest tests -q -rs -``` +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 + +The CUDA extension requires B <= 65535 and at most 2^31-1 elements in each input. +Python checks these bounds before contiguous copies; C++ checks them before launch. + +Exceeding these limits raises an error; resolution only uses H to select fallback. + +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. diff --git a/docs/validation/adaln-h100.md b/docs/validation/adaln-h100.md new file mode 100644 index 000000000..9b6a22341 --- /dev/null +++ b/docs/validation/adaln-h100.md @@ -0,0 +1,109 @@ +# 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. + +These are validation instructions, not recorded results. 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. From aefd69a528bc149f06e0510600cb892625676adf Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Thu, 22 Oct 2026 16:43:59 +0800 Subject: [PATCH 24/37] docs(adaln): record H100 acceptance and indexed follow-up Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 3 ++ docs/validation/adaln-h100.md | 52 +++++++++++++++++++++++++++++- 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 79867afd7..340f7f5fd 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -75,6 +75,9 @@ CPU, ROCm, MUSA and NPU registry entries select PyTorch. Indexed modulation alwa selects the PyTorch select01 path, with `fallback=true` and reason `modulate_index_requires_select01`. +Dedicated CUDA/Triton indexed kernels are deferred to a follow-up PR. This PR +does not qualify indexed framework backward for the shared replay profile. + ### Strict mode and trace Pass `strict=True` when resolving the shared path to require its declared arithmetic diff --git a/docs/validation/adaln-h100.md b/docs/validation/adaln-h100.md index 9b6a22341..7a8949c81 100644 --- a/docs/validation/adaln-h100.md +++ b/docs/validation/adaln-h100.md @@ -4,7 +4,8 @@ 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. -These are validation instructions, not recorded results. Retain the source revision, +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. @@ -107,3 +108,52 @@ environment. It includes hardware and integration tests beyond AdaLN: 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. From b2eabc9bc3db6bed6a83794ad5e44316055f2801 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Thu, 22 Oct 2026 17:26:02 +0800 Subject: [PATCH 25/37] fix(adaln): reject Triton inputs beyond int32 addressing Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 6 +++-- .../shared/triton/norm/adaln_modulation.py | 2 ++ tests/backends/test_adaln_shared_bounds.py | 23 +++++++++++++++++++ 3 files changed, 29 insertions(+), 2 deletions(-) diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 340f7f5fd..2e1339b5c 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -163,8 +163,10 @@ fresh extension builds, sanitizer commands and evidence records. ## Known Limitations -The CUDA extension requires B <= 65535 and at most 2^31-1 elements in each input. -Python checks these bounds before contiguous copies; C++ checks them before launch. +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. Exceeding these limits raises an error; resolution only uses H to select fallback. diff --git a/rl_engine/backends/shared/triton/norm/adaln_modulation.py b/rl_engine/backends/shared/triton/norm/adaln_modulation.py index a0bd4fd79..2c8654cf5 100644 --- a/rl_engine/backends/shared/triton/norm/adaln_modulation.py +++ b/rl_engine/backends/shared/triton/norm/adaln_modulation.py @@ -187,6 +187,8 @@ def __call__(self, x, modulation, *, eps=1e-6): def forward(self, x, modulation, *, eps=1e-6): _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( diff --git a/tests/backends/test_adaln_shared_bounds.py b/tests/backends/test_adaln_shared_bounds.py index 79927d892..df5d5d7fd 100644 --- a/tests/backends/test_adaln_shared_bounds.py +++ b/tests/backends/test_adaln_shared_bounds.py @@ -11,6 +11,29 @@ from rl_engine.runtime.registry import OpBackend, kernel_registry +@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: From c39489566b171534522fd8ab73098e0dc2477938 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Thu, 22 Oct 2026 17:26:03 +0800 Subject: [PATCH 26/37] test(adaln): require Triton for large-hidden resolution assertion Signed-off-by: hongleng-github --- tests/ops/norm/test_adaln_modulation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/ops/norm/test_adaln_modulation.py b/tests/ops/norm/test_adaln_modulation.py index c4c37c9e5..16416244a 100644 --- a/tests/ops/norm/test_adaln_modulation.py +++ b/tests/ops/norm/test_adaln_modulation.py @@ -118,6 +118,7 @@ def test_adaln_cuda_registry_reports_real_fallback(): 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 From beb5bff59d2a9611b3d12eeeec775ac6d757321b Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Thu, 22 Oct 2026 17:27:22 +0800 Subject: [PATCH 27/37] docs(adaln): qualify S specialization and review validation evidence Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 6 ++++++ docs/validation/adaln-h100.md | 28 ++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 2e1339b5c..4668900bc 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -177,3 +177,9 @@ 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. + +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 index 7a8949c81..caae28bd0 100644 --- a/docs/validation/adaln-h100.md +++ b/docs/validation/adaln-h100.md @@ -157,3 +157,31 @@ 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. From d18c4c27f0c033552ac66b847910766634ddc3ab Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Fri, 23 Oct 2026 01:03:34 +0800 Subject: [PATCH 28/37] test(adaln): align dispatch test ownership with parent Apply the test-only layout update from parent PR #559 commit 38352b007cd748e2e1080a098b8ee5f20224b807. Move resolver and fallback policy cases into tests/runtime while retaining operator accuracy and invariance cases in tests/ops. Preserve existing test names, assertions and decorators. Signed-off-by: hongleng-github --- tests/backends/test_adaln_shared_bounds.py | 13 ---- tests/ops/norm/test_adaln_modulation.py | 64 --------------- tests/runtime/test_adaln_dispatch.py | 91 +++++++++++++++++++++- 3 files changed, 90 insertions(+), 78 deletions(-) diff --git a/tests/backends/test_adaln_shared_bounds.py b/tests/backends/test_adaln_shared_bounds.py index df5d5d7fd..7a90e8bfa 100644 --- a/tests/backends/test_adaln_shared_bounds.py +++ b/tests/backends/test_adaln_shared_bounds.py @@ -8,7 +8,6 @@ import torch from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp -from rl_engine.runtime.registry import OpBackend, kernel_registry @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) @@ -71,18 +70,6 @@ def test_cuda_shared_rejects_hidden_4097(cuda_op, dtype): cuda_op(x, m) -@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"] - - @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): diff --git a/tests/ops/norm/test_adaln_modulation.py b/tests/ops/norm/test_adaln_modulation.py index 16416244a..d5a717b07 100644 --- a/tests/ops/norm/test_adaln_modulation.py +++ b/tests/ops/norm/test_adaln_modulation.py @@ -61,70 +61,6 @@ def test_adaln_modulation_rejects_wrong_modulation_shape(): NativeAdaLNModulationOp()(torch.zeros(1, 2, 4), torch.zeros(1, 4)) -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(): - try: - torch.cuda.init() - except RuntimeError: - pytest.skip("CUDA required") - 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"] - - def test_adaln_modulation_matches_independent_layer_norm_autograd(): torch.manual_seed(386) x = torch.randn(2, 3, 7, requires_grad=True) diff --git a/tests/runtime/test_adaln_dispatch.py b/tests/runtime/test_adaln_dispatch.py index 6ab700e4e..2030a31aa 100644 --- a/tests/runtime/test_adaln_dispatch.py +++ b/tests/runtime/test_adaln_dispatch.py @@ -2,6 +2,8 @@ # Copyright (c) 2026 RL-Kernel Contributors """AdaLN dispatch policy and observations at the public resolver seam.""" +import os + import pytest import torch @@ -17,9 +19,96 @@ resolve_tolerance, validate_backend_provenance, ) +from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp from rl_engine.runtime.registry import OpBackend -from rl_engine.runtime.registry import KernelRegistry +from rl_engine.runtime.registry import KernelRegistry, kernel_registry + + +@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(): + try: + torch.cuda.init() + except RuntimeError: + pytest.skip("CUDA required") + 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(): From dd550ecce2e0df2bf003856192cdff013915ea29 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Fri, 23 Oct 2026 01:03:45 +0800 Subject: [PATCH 29/37] feat(adaln): add small FP32 indexed CUDA forward and backward Route supported B2/S8/H2 FP32 inputs with int32 indices through dedicated CUDA kernels. Preserve explicit framework fallback outside the slice and keep indexed strict execution rejected. Add an independent CPU byte reference, literal VJP and token-order checks, and direct backend plus public dispatch coverage. Share FP32 eps eligibility and indexed trace construction, and isolate NVIDIA-only tests from ROCm. Validation: bounded CPU-only regression passed with 135 passed and 305 skipped; targeted Ruff and diff checks passed. Historical SM86 evidence predates the Python refactors; current-source GPU acceptance, full-extension linkage, H100 and strict qualification remain pending. Signed-off-by: hongleng-github --- csrc/bindings/ops.cpp | 13 ++ csrc/cuda/norm/adaln_modulation.cu | 131 ++++++++++++++++++ docs/operators/adaln-modulation.md | 33 ++++- rl_engine/_C.pyi | 16 +++ .../backends/cuda/norm/adaln_modulation.py | 66 ++++++++- rl_engine/reference/norm/adaln_indexed.py | 73 ++++++++++ rl_engine/reference/norm/adaln_modulation.py | 25 ++-- rl_engine/runtime/adaln.py | 39 +++++- rl_engine/runtime/registry.py | 33 ++++- tests/helpers/adaln_indexed.py | 81 +++++++++++ tests/ops/norm/test_adaln_indexed.py | 15 ++ tests/reference/test_adaln_indexed.py | 80 +++++++++++ tests/runtime/test_adaln_indexed.py | 115 +++++++++++++++ 13 files changed, 691 insertions(+), 29 deletions(-) create mode 100644 rl_engine/reference/norm/adaln_indexed.py create mode 100644 tests/helpers/adaln_indexed.py create mode 100644 tests/ops/norm/test_adaln_indexed.py create mode 100644 tests/reference/test_adaln_indexed.py create mode 100644 tests/runtime/test_adaln_indexed.py diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index 1dbf203ac..63ed671d8 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -262,6 +262,11 @@ std::vector adaln_modulation_forward_cuda( 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 @@ -753,6 +758,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "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) diff --git a/csrc/cuda/norm/adaln_modulation.cu b/csrc/cuda/norm/adaln_modulation.cu index ad311d483..09ad2f482 100644 --- a/csrc/cuda/norm/adaln_modulation.cu +++ b/csrc/cuda/norm/adaln_modulation.cu @@ -159,6 +159,94 @@ void check(const torch::Tensor& x, const torch::Tensor& modulation, double eps, "eps must remain finite FP32 without positive underflow"); } +// First indexed slice: FP32 B2/S8/H2. Expand only with the domain tests. +__global__ void indexed_forward_kernel( + const float* x, const float* m, const int* index, float* y, float* gate, float eps) { + extern __shared__ float scratch[]; + const int row = blockIdx.x; + const int sample = row / 8; + const int bank_row = sample + index[row] * 2; + const auto stats = row_stats(x, scratch, row, 2, 2, eps); + if (threadIdx.x < 2) { + const int col = threadIdx.x; + const float n = __fmul_rn(__fsub_rn(x[row * 2 + col], stats.mean), stats.rstd); + y[row * 2 + col] = __fadd_rn( + __fmul_rn(n, __fadd_rn(1.0f, m[bank_row * 6 + 2 + col])), m[bank_row * 6 + col]); + gate[row * 2 + col] = m[bank_row * 6 + 4 + col]; + } +} + +__global__ void indexed_backward_rows_kernel( + const float* x, const float* m, const int* index, const float* dy, + float* dx, float* partial_shift, float* partial_scale, float eps) { + extern __shared__ float scratch[]; + const int row = blockIdx.x; + const int bank_row = row / 8 + index[row] * 2; + const auto stats = row_stats(x, scratch, row, 2, 2, eps); + if (threadIdx.x < 2) { + const int col = threadIdx.x; + const float n0 = __fmul_rn(__fsub_rn(x[row * 2], stats.mean), stats.rstd); + const float n1 = __fmul_rn(__fsub_rn(x[row * 2 + 1], stats.mean), stats.rstd); + const float dn0 = __fmul_rn(dy[row * 2], __fadd_rn(1.0f, m[bank_row * 6 + 2])); + const float dn1 = __fmul_rn(dy[row * 2 + 1], __fadd_rn(1.0f, m[bank_row * 6 + 3])); + const float mean_dn = __fdiv_rn(__fadd_rn(dn0, dn1), 2.0f); + const float mean_dnn = __fdiv_rn( + __fadd_rn(__fmul_rn(dn0, n0), __fmul_rn(dn1, n1)), 2.0f); + const float n = col == 0 ? n0 : n1; + const float dn = col == 0 ? dn0 : dn1; + dx[row * 2 + col] = __fmul_rn( + __fsub_rn(__fsub_rn(dn, mean_dn), __fmul_rn(n, mean_dnn)), stats.rstd); + partial_shift[row * 2 + col] = dy[row * 2 + col]; + partial_scale[row * 2 + col] = __fmul_rn(dy[row * 2 + col], n); + } +} + +__global__ void indexed_backward_reduce_kernel( + const float* ps, const float* pc, const float* dg, const int* index, float* dm) { + const int sample = blockIdx.x; + const int col = threadIdx.x; + if (col >= 2) 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 < 8; ++token) { + const int row = sample * 8 + token; + const int bank = index[row]; + const int offset = row * 2 + col; + // The unselected bank performs no add and reads no contribution. + 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], dg[offset]); + } + for (int bank = 0; bank < 2; ++bank) { + const int offset = (sample + bank * 2) * 6 + col; + dm[offset] = ds[bank]; + dm[offset + 2] = dc[bank]; + dm[offset + 4] = 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) == 2 && x.size(1) == 8 && x.size(2) == 2, + "indexed CUDA slice requires x [2,8,2]"); + TORCH_CHECK(m.dim() == 2 && m.size(0) == 4 && m.size(1) == 6, + "indexed CUDA slice requires modulation [4,6]"); + TORCH_CHECK(x.scalar_type() == at::kFloat && m.scalar_type() == at::kFloat, + "indexed CUDA slice requires FP32"); + // The shared preflight also checks CUDA, dtype/device, eps and launch settings. + check(x, m.narrow(0, 0, 2), 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) == 2 && index.size(1) == 8, + "indexed CUDA slice requires contiguous int32 index [2,8] 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( @@ -213,3 +301,46 @@ std::vector adaln_modulation_backward_cuda( 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()); + auto y = torch::empty_like(x), gate = torch::empty_like(x); + indexed_forward_kernel<<<16, threads, 2 * sizeof(float), at::cuda::getCurrentCUDAStream()>>>( + x.data_ptr(), m.data_ptr(), index.data_ptr(), + y.data_ptr(), gate.data_ptr(), 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() == at::kFloat && + grad->is_contiguous() && grad->sizes() == x.sizes(), + "indexed output gradient must match x device/dtype/shape"); + } + } + c10::cuda::CUDAGuard guard(x.device()); + 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_like(x); + pc = torch::empty_like(x); + indexed_backward_rows_kernel<<<16, threads, 2 * sizeof(float), stream>>>( + x.data_ptr(), m.data_ptr(), index.data_ptr(), dy->data_ptr(), + dx.data_ptr(), ps.data_ptr(), pc.data_ptr(), static_cast(eps)); + } + indexed_backward_reduce_kernel<<<2, threads, 0, stream>>>( + 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()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dx, dm}; +} diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 4668900bc..53a2ee4a9 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -71,12 +71,18 @@ 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. Indexed modulation always -selects the PyTorch select01 path, with `fallback=true` and reason -`modulate_index_requires_select01`. - -Dedicated CUDA/Triton indexed kernels are deferred to a follow-up PR. This PR -does not qualify indexed framework backward for the shared replay profile. +CPU, ROCm, MUSA and NPU registry entries select PyTorch. The initial dedicated +indexed CUDA slice accepts only FP32 x `[2,8,2]` and int32 index `[2,8]`, with +the usual `[4,6]` modulation. It requires fresh indexed extension symbols. +The small slice has local SM86 coverage using an AdaLN-only extension build; +full-extension and target H100 qualification remain separate. It does not enable strict. + +Other indexed inputs select PyTorch select01 with `fallback=true` and reason +`modulate_index_requires_select01`. If a resolved indexed CUDA slice receives +unsupported geometry, dtype or eps at execution, it explicitly executes the +framework path with reason `indexed_cuda_slice_unsupported`. +Expanded CUDA coverage and dedicated indexed Triton remain follow-up work. +Indexed framework backward has no shared replay-profile qualification. ### Strict mode and trace @@ -109,7 +115,12 @@ 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 modulation checks only finite, nonnegative eps. +Indexed framework fallback checks only finite, nonnegative eps. The small indexed +CUDA slice uses the shared FP32 eps domain. Its 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. @@ -130,6 +141,12 @@ 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 callable for indexed byte +tests, with its own H tree and explicit VJP. It is not the framework fallback. +The first indexed tests independently check a hand-calculated nonzero VJP and +a small mixed-bank case; broader indexed reference/autograd boundaries remain +part of the follow-up. The accuracy gold's indexed branch still uses framework +autograd and is not a fixed-order byte oracle. The CLI records trace snapshots without automatically constructing C1 `BackendProvenance`; its accuracy report alone is not a WS1 system gate. @@ -137,6 +154,8 @@ The CLI records trace snapshots without automatically constructing C1 ```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 \ diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index d80c16b5c..bb194f8fb 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -263,6 +263,22 @@ 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 index a26a0adf3..001fec3d4 100644 --- a/rl_engine/backends/cuda/norm/adaln_modulation.py +++ b/rl_engine/backends/cuda/norm/adaln_modulation.py @@ -7,7 +7,49 @@ import torch from rl_engine.backends.extension import _C, _EXT_AVAILABLE -from rl_engine.reference.norm.adaln_modulation import _validate +from rl_engine.reference.norm.adaln_modulation import ( + _eps_supported_fp32, + _validate, + _validate_select01, +) + + +def _indexed_slice_supported(x, index, eps): + # shortcut: only the first FP32 tracer bullet; expand with ticket 03 tests. + if not ( + x.is_cuda + and torch.version.hip is None + and x.shape == (2, 8, 2) + and x.dtype == torch.float32 + and index.shape == (2, 8) + and index.dtype == torch.int32 + ): + return False + return _eps_supported_fp32(eps) + + +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 + 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): @@ -37,11 +79,27 @@ def __init__(self, *, threads=256): 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): - return self.forward(x, modulation, eps=eps) + 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): + def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): + if modulate_index is not None: + _validate_select01(x, modulation, modulate_index, eps) + if not self.indexed_available or not _indexed_slice_supported(x, modulate_index, eps): + raise RuntimeError( + "indexed CUDA slice requires FP32 x [2,8,2] and int32 index [2,8]" + ) + return _AdaLNIndexedCuda.apply( + x.contiguous(), + modulation.contiguous(), + modulate_index.contiguous(), + float(eps), + self.threads, + ) _validate(x, modulation, eps) if x.shape[0] > 65535: raise ValueError("CUDA AdaLN requires B <= 65535") diff --git a/rl_engine/reference/norm/adaln_indexed.py b/rl_engine/reference/norm/adaln_indexed.py new file mode 100644 index 000000000..64e06c28a --- /dev/null +++ b/rl_engine/reference/norm/adaln_indexed.py @@ -0,0 +1,73 @@ +# 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 + mean = _hidden_tree(x.float()) / hidden + centered = x.float() - 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) + ctx.set_materialize_grads(False) + ctx.save_for_backward(norm, rstd, scale, index) + ctx.mod_shape = modulation.shape + return norm * (1.0 + scale) + shift, gate + + @staticmethod + def backward(ctx, dy, dg): + norm, rstd, scale, index = ctx.saved_tensors + batch, seq, hidden = norm.shape + dx = None + dm = torch.zeros(ctx.mod_shape, dtype=torch.float32) + if dy is not None: + 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 dx, dm, None, None + + +class IndexedAdaLNFixedReference: + """FP32 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 != torch.float32 or modulation.dtype != torch.float32: + raise TypeError("indexed byte reference requires FP32 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 index e2a35968f..120d57da0 100644 --- a/rl_engine/reference/norm/adaln_modulation.py +++ b/rl_engine/reference/norm/adaln_modulation.py @@ -36,6 +36,14 @@ def _normalize_fp32(x: torch.Tensor, eps: float) -> tuple[torch.Tensor, torch.Te 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 _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]") @@ -47,13 +55,8 @@ def _validate(x: torch.Tensor, modulation: torch.Tensor, eps: float) -> None: 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") - message = "eps must remain finite FP32 without positive underflow" - try: - eps32 = struct.unpack("f", struct.pack("f", eps))[0] - except OverflowError: - raise ValueError(message) from None - if not math.isfinite(eps32) or (eps > 0 and eps32 == 0): - raise ValueError(message) + if not _eps_supported_fp32(eps): + raise ValueError("eps must remain finite FP32 without positive underflow") class _AdaLNReference(torch.autograd.Function): @@ -92,12 +95,12 @@ def backward(ctx, grad_y: torch.Tensor, grad_gate: torch.Tensor): return dx, dm, None -def _select01_reference( +def _validate_select01( x: torch.Tensor, modulation: torch.Tensor, modulate_index: torch.Tensor, eps: float, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> 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 @@ -123,6 +126,10 @@ def _select01_reference( if not math.isfinite(eps) or eps < 0: raise ValueError("eps must be finite and nonnegative") + +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, :]) diff --git a/rl_engine/runtime/adaln.py b/rl_engine/runtime/adaln.py index b1b3df7b8..e091a8143 100644 --- a/rl_engine/runtime/adaln.py +++ b/rl_engine/runtime/adaln.py @@ -11,6 +11,7 @@ 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"), @@ -77,7 +78,22 @@ def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): if x.ndim != 3 or x.shape[-1] != self.hidden: raise ValueError("strict AdaLN input hidden differs from resolved hidden") if indexed: - outputs = self.implementation(x, modulation, eps=eps, modulate_index=modulate_index) + implementation = self.implementation + if backend in {"cuda", "triton"}: + from rl_engine.backends.cuda.norm.adaln_modulation import _indexed_slice_supported + + if ( + backend == "cuda" + and getattr(implementation, "indexed_available", False) + and _indexed_slice_supported(x, modulate_index, eps) + ): + 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: outputs = self.implementation(x, modulation, eps=eps) self.trace.update( @@ -88,11 +104,28 @@ def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): output_dtype=str(outputs[0].dtype).removeprefix("torch."), execution_device=str(x.device), ) - if indexed: + if indexed and profile == SELECT01_PROFILE: self.trace.update( fallback=True, - fallback_reason="modulate_index_requires_select01", + fallback_reason=( + "indexed_cuda_slice_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="adaln_indexed_cuda_v1", + ) + elif self._resolution_trace["arithmetic_profile"] == INDEXED_PROFILE: + self.trace.update( + reduction_order="pairwise_lower_upper_H; ascending_logical_token_S", + kernel_fingerprint=f"adaln_modulation_v1:{self.trace['selected_backend']}", ) if backend == "cuda": from rl_engine.backends.extension import _C diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index dc29f0e91..b5e8a0b73 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -1073,7 +1073,9 @@ def get_adaln_modulation_op( strict: bool = False, ) -> tuple[Any, dict[str, Any]]: """Resolve AdaLN; strict permits only the shared declared replay profile.""" + from rl_engine.runtime.adaln import ( + INDEXED_PROFILE, SELECT01_PROFILE, SHARED_PROFILE, AdaLNDispatch, @@ -1088,22 +1090,41 @@ def get_adaln_modulation_op( platform = self._platform_for_device(device) if modulate_index is not None: backend = OpBackend.PYTORCH_ADALN_MODULATION + # Only the first indexed CUDA slice is accelerated; execution rechecks x/eps. + if ( + platform == "cuda" + and hidden == 2 + and modulate_index.shape == (2, 8) + and modulate_index.dtype == torch.int32 + ): + cuda = self._get_or_create_backend(OpBackend.CUDA_ADALN_MODULATION) + if cuda is not None and getattr(cuda, "indexed_available", False): + backend = OpBackend.CUDA_ADALN_MODULATION op = self._get_or_create_backend(backend) if op is None: raise RuntimeError("PyTorch select01 AdaLN fallback is unavailable") + accelerated = backend is OpBackend.CUDA_ADALN_MODULATION trace = { "selected_backend": backend.name, - "fallback": True, - "fallback_reason": "modulate_index_requires_select01", + "fallback": not accelerated, + "fallback_reason": None if accelerated else "modulate_index_requires_select01", "rejected_backends": [], - "reduction_order": "pairwise_lower_upper_H; framework_select01_backward", + "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": "adaln_modulation_select01_native_v1", - "requested_backend": "pytorch", - "arithmetic_profile": profile, + "kernel_fingerprint": ( + "adaln_indexed_cuda_v1" + if accelerated + else "adaln_modulation_select01_native_v1" + ), + "requested_backend": "cuda" if accelerated else "pytorch", + "arithmetic_profile": INDEXED_PROFILE if accelerated else profile, } return AdaLNDispatch(op, trace, hidden=hidden, strict=strict), trace candidates = self._priority_map[platform]["adaln_modulation"] diff --git a/tests/helpers/adaln_indexed.py b/tests/helpers/adaln_indexed.py new file mode 100644 index 000000000..8e08cff2a --- /dev/null +++ b/tests/helpers/adaln_indexed.py @@ -0,0 +1,81 @@ +# 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" diff --git a/tests/ops/norm/test_adaln_indexed.py b/tests/ops/norm/test_adaln_indexed.py new file mode 100644 index 000000000..46dce8ae2 --- /dev/null +++ b/tests/ops/norm/test_adaln_indexed.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""L2: indexed AdaLN forward/backward accuracy and byte comparisons on CUDA.""" + +import pytest + +from tests.helpers.adaln_indexed import check_small_case, require_cuda + + +@pytest.mark.cuda_only +def test_cuda_indexed_small_fp32_direct_backend(): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + check_small_case(CudaAdaLNModulationOp(), "cuda") diff --git a/tests/reference/test_adaln_indexed.py b/tests/reference/test_adaln_indexed.py new file mode 100644 index 000000000..030932383 --- /dev/null +++ b/tests/reference/test_adaln_indexed.py @@ -0,0 +1,80 @@ +# 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 + + +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], + ]), + ) diff --git a/tests/runtime/test_adaln_indexed.py b/tests/runtime/test_adaln_indexed.py new file mode 100644 index 000000000..643f3f964 --- /dev/null +++ b/tests/runtime/test_adaln_indexed.py @@ -0,0 +1,115 @@ +# 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[:, :7] if unsupported == "geometry" else x + fallback_index = index[:, :7] if unsupported == "geometry" else index + eps = {"eps": 1e300, "eps_underflow": 2**-150, "geometry": 7.0}[unsupported] + actual = op(fallback_x, m, modulate_index=fallback_index, eps=eps) + expected = NativeAdaLNModulationOp()(fallback_x, 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_slice_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) From d079f7c156dd9d2ddc7742227f5e226bfed321f5 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Fri, 23 Oct 2026 01:19:21 +0800 Subject: [PATCH 30/37] test(adaln): cover indexed CUDA gradient semantics Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 15 +- .../validation/operators/operator_specs.py | 39 +++- tests/helpers/adaln_indexed.py | 205 ++++++++++++++++++ tests/ops/norm/test_adaln_indexed.py | 101 +++++++++ tests/reference/test_adaln_indexed.py | 96 ++++++++ tests/runtime/test_adaln_indexed.py | 14 ++ 6 files changed, 461 insertions(+), 9 deletions(-) diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 53a2ee4a9..f4983dc05 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -143,10 +143,17 @@ It checks both outputs and gradients of x and modulation using the shared Accuracy checks and fixed-reference byte comparisons are separate tests. `IndexedAdaLNFixedReference` is a separate CPU FP32 callable for indexed byte tests, with its own H tree and explicit VJP. It is not the framework fallback. -The first indexed tests independently check a hand-calculated nonzero VJP and -a small mixed-bank case; broader indexed reference/autograd boundaries remain -part of the follow-up. The accuracy gold's indexed branch still uses framework -autograd and is not a fixed-order byte oracle. +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 in the FP32 [2,8,2] slice. +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 fresh current-source GPU evidence. +Ticket 02 adds no wider input-domain or strict qualification claim; indexed strict +remains closed. The CLI records trace snapshots without automatically constructing C1 `BackendProvenance`; its accuracy report alone is not a WS1 system gate. diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 901d14800..13c076d34 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -318,6 +318,39 @@ def backward(ctx, dy, dg): 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. @@ -332,12 +365,8 @@ def forward_fp32(self, x, modulation, *, eps=1e-6, modulate_index=None): if modulate_index is None: norm = torch.nn.functional.layer_norm(x32, (x.shape[-1],), eps=eps) return _GtestSharedAdaLNAffine.apply(norm, mod32) - else: - batch = x.shape[0] - rows = torch.arange(batch)[:, None] + modulate_index.cpu().long() * batch - shift, scale, gate = mod32[rows].chunk(3, dim=-1) norm = torch.nn.functional.layer_norm(x32, (x.shape[-1],), eps=eps) - return norm * (1 + scale) + shift, gate + return _GtestIndexedAdaLNAffine.apply(norm, mod32, modulate_index.cpu()) class GtestPackOp: diff --git a/tests/helpers/adaln_indexed.py b/tests/helpers/adaln_indexed.py index 8e08cff2a..dff728d4c 100644 --- a/tests/helpers/adaln_indexed.py +++ b/tests/helpers/adaln_indexed.py @@ -79,3 +79,208 @@ def require_cuda(): # 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) + # Selected bank0: FP32 fold (2**26 + 1) - 2**26 = 0. + # Reordering to cancel the large terms first leaves 1 (scale: +/-3/4). + middle, last = (1.0, -(2**26)) if order == "lost_unit" else (-(2**26), 1.0) + u = torch.tensor([2**26, 4.0, middle, 5.0, last, 6.0, 0.0, 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 = ( + [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]] * 2 + if order == "lost_unit" + else [[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 in ((y, ry), (dx, rdx), (dm, rdm)): + torch.testing.assert_close(actual.cpu(), expected, rtol=1e-5, atol=1e-6) diff --git a/tests/ops/norm/test_adaln_indexed.py b/tests/ops/norm/test_adaln_indexed.py index 46dce8ae2..31b7becb0 100644 --- a/tests/ops/norm/test_adaln_indexed.py +++ b/tests/ops/norm/test_adaln_indexed.py @@ -13,3 +13,104 @@ def test_cuda_indexed_small_fp32_direct_backend(): from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp check_small_case(CudaAdaLNModulationOp(), "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("route", ["all0", "all1", "alternating"]) +def test_cuda_indexed_route_identity_literal(route): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_route_identity + + check_route_identity(CudaAdaLNModulationOp(), "cuda", route) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("explicit_zero", [False, True]) +def test_cuda_indexed_absent_vs_zero_dy_nonfinite_literal(explicit_zero): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_absent_vs_zero_dy + + check_absent_vs_zero_dy(CudaAdaLNModulationOp(), "cuda", explicit_zero) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("order", ["lost_unit", "retained_unit"]) +def test_cuda_indexed_selected_token_cancellation_literal(order): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_cancellation + + check_cancellation(CudaAdaLNModulationOp(), "cuda", order) + + +@pytest.mark.cuda_only +def test_cuda_indexed_dscale_product_rounds_before_add_literal(): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_product_rounding + + check_product_rounding(CudaAdaLNModulationOp(), "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("bank", [0, 1]) +@pytest.mark.parametrize("nonfinite", [float("nan"), float("inf")], ids=["nan", "inf"]) +def test_cuda_indexed_nonfinite_selected_bank_isolation_literal(bank, nonfinite): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_nonfinite_isolation + + check_nonfinite_isolation(CudaAdaLNModulationOp(), "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_cuda_indexed_whole_leaf_participation_literal(frozen, outputs): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_leaf_participation + + check_leaf_participation(CudaAdaLNModulationOp(), "cuda", frozen, outputs) + + +@pytest.mark.cuda_only +def test_cuda_indexed_gate_copies_original_fp32_payload_literal(): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_gate_payload + + check_gate_payload(CudaAdaLNModulationOp(), "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("contribution", ["negative_zero", "subnormal"]) +def test_cuda_indexed_signed_zero_and_subnormal_fold_literal(contribution): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_tiny_contributions + + check_tiny_contributions(CudaAdaLNModulationOp(), "cuda", contribution) + + +@pytest.mark.cuda_only +def test_cuda_indexed_all0_matches_shared_arithmetic_and_independent_gold(): + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + from tests.helpers.adaln_indexed import check_shared_bank0 + + op = CudaAdaLNModulationOp() + check_shared_bank0(op, op, "cuda") diff --git a/tests/reference/test_adaln_indexed.py b/tests/reference/test_adaln_indexed.py index 030932383..cdcf41b5d 100644 --- a/tests/reference/test_adaln_indexed.py +++ b/tests/reference/test_adaln_indexed.py @@ -78,3 +78,99 @@ def test_indexed_byte_reference_preserves_selected_token_fold_and_sample_routing [30.0, 30.0, -22.5, 22.5, 45.0, 45.0], ]), ) + +@pytest.mark.parametrize("route", ["all0", "all1", "alternating"]) +@pytest.mark.parametrize( + "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] +) +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]) +@pytest.mark.parametrize( + "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] +) +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( + "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] +) +@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) + + +@pytest.mark.parametrize( + "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] +) +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"]) +@pytest.mark.parametrize( + "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] +) +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"), + ], +) +@pytest.mark.parametrize( + "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] +) +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) + + +@pytest.mark.parametrize( + "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] +) +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"]) +@pytest.mark.parametrize( + "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] +) +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") diff --git a/tests/runtime/test_adaln_indexed.py b/tests/runtime/test_adaln_indexed.py index 643f3f964..20880e629 100644 --- a/tests/runtime/test_adaln_indexed.py +++ b/tests/runtime/test_adaln_indexed.py @@ -113,3 +113,17 @@ def test_cuda_indexed_resolution_can_execute_shared_without_stale_profile(): 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"] From b32cc62e73798bd41aa95dcba552ed476da163ee Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Fri, 23 Oct 2026 01:36:19 +0800 Subject: [PATCH 31/37] test(adaln): strengthen indexed cancellation checks and simplify fixtures Signed-off-by: hongleng-github --- .../validation/operators/operator_specs.py | 3 +- tests/helpers/adaln_indexed.py | 26 +++++--- tests/ops/norm/test_adaln_indexed.py | 66 ++++++++----------- tests/reference/test_adaln_indexed.py | 32 +++------ 4 files changed, 51 insertions(+), 76 deletions(-) diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 13c076d34..9a43d3923 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -362,10 +362,9 @@ class GtestAdaLNReference: 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: - norm = torch.nn.functional.layer_norm(x32, (x.shape[-1],), eps=eps) return _GtestSharedAdaLNAffine.apply(norm, mod32) - norm = torch.nn.functional.layer_norm(x32, (x.shape[-1],), eps=eps) return _GtestIndexedAdaLNAffine.apply(norm, mod32, modulate_index.cpu()) diff --git a/tests/helpers/adaln_indexed.py b/tests/helpers/adaln_indexed.py index dff728d4c..a369323ad 100644 --- a/tests/helpers/adaln_indexed.py +++ b/tests/helpers/adaln_indexed.py @@ -144,19 +144,16 @@ def check_cancellation(op, device, order): 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) - # Selected bank0: FP32 fold (2**26 + 1) - 2**26 = 0. - # Reordering to cancel the large terms first leaves 1 (scale: +/-3/4). + # 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) - u = torch.tensor([2**26, 4.0, middle, 5.0, last, 6.0, 0.0, 7.0], device=device) + 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 = ( - [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]] * 2 - if order == "lost_unit" - else [[1.0, 1.0, -0.75, 0.75, 1.0, 1.0], [2.0, 2.0, -1.5, 1.5, 3.0, 3.0]] - ) + 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( @@ -282,5 +279,14 @@ def check_shared_bank0(indexed_op, shared_op, device): 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 in ((y, ry), (dx, rdx), (dm, rdm)): - torch.testing.assert_close(actual.cpu(), expected, rtol=1e-5, atol=1e-6) + 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 + ) diff --git a/tests/ops/norm/test_adaln_indexed.py b/tests/ops/norm/test_adaln_indexed.py index 31b7becb0..7a8c836d1 100644 --- a/tests/ops/norm/test_adaln_indexed.py +++ b/tests/ops/norm/test_adaln_indexed.py @@ -7,62 +7,57 @@ from tests.helpers.adaln_indexed import check_small_case, require_cuda -@pytest.mark.cuda_only -def test_cuda_indexed_small_fp32_direct_backend(): +@pytest.fixture +def cuda_op(): require_cuda() from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp - check_small_case(CudaAdaLNModulationOp(), "cuda") + return CudaAdaLNModulationOp() + + +@pytest.mark.cuda_only +def test_cuda_indexed_small_fp32_direct_backend(cuda_op): + check_small_case(cuda_op, "cuda") @pytest.mark.cuda_only @pytest.mark.parametrize("route", ["all0", "all1", "alternating"]) -def test_cuda_indexed_route_identity_literal(route): - require_cuda() - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +def test_cuda_indexed_route_identity_literal(cuda_op, route): from tests.helpers.adaln_indexed import check_route_identity - check_route_identity(CudaAdaLNModulationOp(), "cuda", route) + check_route_identity(cuda_op, "cuda", route) @pytest.mark.cuda_only @pytest.mark.parametrize("explicit_zero", [False, True]) -def test_cuda_indexed_absent_vs_zero_dy_nonfinite_literal(explicit_zero): - require_cuda() - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +def test_cuda_indexed_absent_vs_zero_dy_nonfinite_literal(cuda_op, explicit_zero): from tests.helpers.adaln_indexed import check_absent_vs_zero_dy - check_absent_vs_zero_dy(CudaAdaLNModulationOp(), "cuda", explicit_zero) + check_absent_vs_zero_dy(cuda_op, "cuda", explicit_zero) @pytest.mark.cuda_only @pytest.mark.parametrize("order", ["lost_unit", "retained_unit"]) -def test_cuda_indexed_selected_token_cancellation_literal(order): - require_cuda() - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +def test_cuda_indexed_selected_token_cancellation_literal(cuda_op, order): from tests.helpers.adaln_indexed import check_cancellation - check_cancellation(CudaAdaLNModulationOp(), "cuda", order) + check_cancellation(cuda_op, "cuda", order) @pytest.mark.cuda_only -def test_cuda_indexed_dscale_product_rounds_before_add_literal(): - require_cuda() - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +def test_cuda_indexed_dscale_product_rounds_before_add_literal(cuda_op): from tests.helpers.adaln_indexed import check_product_rounding - check_product_rounding(CudaAdaLNModulationOp(), "cuda") + check_product_rounding(cuda_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_cuda_indexed_nonfinite_selected_bank_isolation_literal(bank, nonfinite): - require_cuda() - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +def test_cuda_indexed_nonfinite_selected_bank_isolation_literal(cuda_op, bank, nonfinite): from tests.helpers.adaln_indexed import check_nonfinite_isolation - check_nonfinite_isolation(CudaAdaLNModulationOp(), "cuda", bank, nonfinite) + check_nonfinite_isolation(cuda_op, "cuda", bank, nonfinite) @pytest.mark.cuda_only @@ -79,38 +74,29 @@ def test_cuda_indexed_nonfinite_selected_bank_isolation_literal(bank, nonfinite) ("modulation", "both"), ], ) -def test_cuda_indexed_whole_leaf_participation_literal(frozen, outputs): - require_cuda() - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +def test_cuda_indexed_whole_leaf_participation_literal(cuda_op, frozen, outputs): from tests.helpers.adaln_indexed import check_leaf_participation - check_leaf_participation(CudaAdaLNModulationOp(), "cuda", frozen, outputs) + check_leaf_participation(cuda_op, "cuda", frozen, outputs) @pytest.mark.cuda_only -def test_cuda_indexed_gate_copies_original_fp32_payload_literal(): - require_cuda() - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +def test_cuda_indexed_gate_copies_original_fp32_payload_literal(cuda_op): from tests.helpers.adaln_indexed import check_gate_payload - check_gate_payload(CudaAdaLNModulationOp(), "cuda") + check_gate_payload(cuda_op, "cuda") @pytest.mark.cuda_only @pytest.mark.parametrize("contribution", ["negative_zero", "subnormal"]) -def test_cuda_indexed_signed_zero_and_subnormal_fold_literal(contribution): - require_cuda() - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +def test_cuda_indexed_signed_zero_and_subnormal_fold_literal(cuda_op, contribution): from tests.helpers.adaln_indexed import check_tiny_contributions - check_tiny_contributions(CudaAdaLNModulationOp(), "cuda", contribution) + check_tiny_contributions(cuda_op, "cuda", contribution) @pytest.mark.cuda_only -def test_cuda_indexed_all0_matches_shared_arithmetic_and_independent_gold(): - require_cuda() - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp +def test_cuda_indexed_all0_matches_shared_arithmetic_and_independent_gold(cuda_op): from tests.helpers.adaln_indexed import check_shared_bank0 - op = CudaAdaLNModulationOp() - check_shared_bank0(op, op, "cuda") + check_shared_bank0(cuda_op, cuda_op, "cuda") diff --git a/tests/reference/test_adaln_indexed.py b/tests/reference/test_adaln_indexed.py index cdcf41b5d..5dc6427bc 100644 --- a/tests/reference/test_adaln_indexed.py +++ b/tests/reference/test_adaln_indexed.py @@ -11,6 +11,14 @@ 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): @@ -80,9 +88,6 @@ def test_indexed_byte_reference_preserves_selected_token_fold_and_sample_routing ) @pytest.mark.parametrize("route", ["all0", "all1", "alternating"]) -@pytest.mark.parametrize( - "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] -) def test_indexed_route_identity_literal(reference, route): from tests.helpers.adaln_indexed import check_route_identity @@ -90,18 +95,12 @@ def test_indexed_route_identity_literal(reference, route): @pytest.mark.parametrize("explicit_zero", [False, True]) -@pytest.mark.parametrize( - "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] -) 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( - "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] -) @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 @@ -109,9 +108,6 @@ def test_indexed_selected_token_cancellation_literal(reference, order): check_cancellation(reference, "cpu", order) -@pytest.mark.parametrize( - "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] -) def test_indexed_dscale_product_rounds_before_add_literal(reference): from tests.helpers.adaln_indexed import check_product_rounding @@ -120,9 +116,6 @@ def test_indexed_dscale_product_rounds_before_add_literal(reference): @pytest.mark.parametrize("bank", [0, 1]) @pytest.mark.parametrize("nonfinite", [float("nan"), float("inf")], ids=["nan", "inf"]) -@pytest.mark.parametrize( - "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] -) def test_indexed_nonfinite_selected_bank_isolation_literal(reference, bank, nonfinite): from tests.helpers.adaln_indexed import check_nonfinite_isolation @@ -142,18 +135,12 @@ def test_indexed_nonfinite_selected_bank_isolation_literal(reference, bank, nonf ("modulation", "both"), ], ) -@pytest.mark.parametrize( - "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] -) 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) -@pytest.mark.parametrize( - "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] -) def test_indexed_gate_copies_original_fp32_payload_literal(reference): from tests.helpers.adaln_indexed import check_gate_payload @@ -161,9 +148,6 @@ def test_indexed_gate_copies_original_fp32_payload_literal(reference): @pytest.mark.parametrize("contribution", ["negative_zero", "subnormal"]) -@pytest.mark.parametrize( - "reference", [IndexedAdaLNFixedReference(), GtestAdaLNReference().forward_fp32] -) def test_indexed_signed_zero_and_subnormal_fold_literal(reference, contribution): from tests.helpers.adaln_indexed import check_tiny_contributions From 872ea7c2c87b571c44846f451465eb0d243c3eb4 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Fri, 23 Oct 2026 14:28:41 +0800 Subject: [PATCH 32/37] feat(adaln): expand indexed CUDA input domain Signed-off-by: hongleng-github --- csrc/cuda/norm/adaln_modulation.cu | 151 +++++----- docs/operators/adaln-modulation.md | 50 +-- .../backends/cuda/norm/adaln_modulation.py | 55 ++-- rl_engine/reference/norm/adaln_indexed.py | 20 +- rl_engine/reference/norm/adaln_modulation.py | 10 +- rl_engine/runtime/adaln.py | 9 +- rl_engine/runtime/registry.py | 6 +- tests/helpers/adaln_indexed.py | 284 +++++++++++++++++- tests/ops/norm/test_adaln_indexed.py | 229 +++++++++++++- tests/ops/norm/test_adaln_modulation.py | 23 +- tests/reference/test_adaln_indexed.py | 117 +++++++- tests/runtime/test_adaln_indexed.py | 141 ++++++++- 12 files changed, 947 insertions(+), 148 deletions(-) diff --git a/csrc/cuda/norm/adaln_modulation.cu b/csrc/cuda/norm/adaln_modulation.cu index 09ad2f482..a53cc0715 100644 --- a/csrc/cuda/norm/adaln_modulation.cu +++ b/csrc/cuda/norm/adaln_modulation.cu @@ -77,10 +77,12 @@ 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) { + 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; + 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; @@ -159,90 +161,70 @@ void check(const torch::Tensor& x, const torch::Tensor& modulation, double eps, "eps must remain finite FP32 without positive underflow"); } -// First indexed slice: FP32 B2/S8/H2. Expand only with the domain tests. +template __global__ void indexed_forward_kernel( - const float* x, const float* m, const int* index, float* y, float* gate, float eps) { + 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 sample = row / 8; - const int bank_row = sample + index[row] * 2; - const auto stats = row_stats(x, scratch, row, 2, 2, eps); - if (threadIdx.x < 2) { - const int col = threadIdx.x; - const float n = __fmul_rn(__fsub_rn(x[row * 2 + col], stats.mean), stats.rstd); - y[row * 2 + col] = __fadd_rn( - __fmul_rn(n, __fadd_rn(1.0f, m[bank_row * 6 + 2 + col])), m[bank_row * 6 + col]); - gate[row * 2 + col] = m[bank_row * 6 + 4 + col]; - } -} - -__global__ void indexed_backward_rows_kernel( - const float* x, const float* m, const int* index, const float* dy, - float* dx, float* partial_shift, float* partial_scale, float eps) { - extern __shared__ float scratch[]; - const int row = blockIdx.x; - const int bank_row = row / 8 + index[row] * 2; - const auto stats = row_stats(x, scratch, row, 2, 2, eps); - if (threadIdx.x < 2) { - const int col = threadIdx.x; - const float n0 = __fmul_rn(__fsub_rn(x[row * 2], stats.mean), stats.rstd); - const float n1 = __fmul_rn(__fsub_rn(x[row * 2 + 1], stats.mean), stats.rstd); - const float dn0 = __fmul_rn(dy[row * 2], __fadd_rn(1.0f, m[bank_row * 6 + 2])); - const float dn1 = __fmul_rn(dy[row * 2 + 1], __fadd_rn(1.0f, m[bank_row * 6 + 3])); - const float mean_dn = __fdiv_rn(__fadd_rn(dn0, dn1), 2.0f); - const float mean_dnn = __fdiv_rn( - __fadd_rn(__fmul_rn(dn0, n0), __fmul_rn(dn1, n1)), 2.0f); - const float n = col == 0 ? n0 : n1; - const float dn = col == 0 ? dn0 : dn1; - dx[row * 2 + col] = __fmul_rn( - __fsub_rn(__fsub_rn(dn, mean_dn), __fmul_rn(n, mean_dnn)), stats.rstd); - partial_shift[row * 2 + col] = dy[row * 2 + col]; - partial_scale[row * 2 + col] = __fmul_rn(dy[row * 2 + col], n); + 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 float* dg, const int* index, float* dm) { - const int sample = blockIdx.x; - const int col = threadIdx.x; - if (col >= 2) return; + 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 < 8; ++token) { - const int row = sample * 8 + token; - const int bank = index[row]; - const int offset = row * 2 + col; - // The unselected bank performs no add and reads no contribution. + 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], dg[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 * 2) * 6 + col; - dm[offset] = ds[bank]; - dm[offset + 2] = dc[bank]; - dm[offset + 4] = dgate[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) == 2 && x.size(1) == 8 && x.size(2) == 2, - "indexed CUDA slice requires x [2,8,2]"); - TORCH_CHECK(m.dim() == 2 && m.size(0) == 4 && m.size(1) == 6, - "indexed CUDA slice requires modulation [4,6]"); - TORCH_CHECK(x.scalar_type() == at::kFloat && m.scalar_type() == at::kFloat, - "indexed CUDA slice requires FP32"); - // The shared preflight also checks CUDA, dtype/device, eps and launch settings. - check(x, m.narrow(0, 0, 2), eps, 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) == 2 && index.size(1) == 8, - "indexed CUDA slice requires contiguous int32 index [2,8] on x device"); + 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"); } @@ -306,10 +288,17 @@ 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); - indexed_forward_kernel<<<16, threads, 2 * sizeof(float), at::cuda::getCurrentCUDAStream()>>>( - x.data_ptr(), m.data_ptr(), index.data_ptr(), - y.data_ptr(), gate.data_ptr(), static_cast(eps)); + 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}; } @@ -320,27 +309,39 @@ std::vector adaln_indexed_backward_cuda( 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() == at::kFloat && + 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_like(x); - pc = torch::empty_like(x); - indexed_backward_rows_kernel<<<16, threads, 2 * sizeof(float), stream>>>( - x.data_ptr(), m.data_ptr(), index.data_ptr(), dy->data_ptr(), - dx.data_ptr(), ps.data_ptr(), pc.data_ptr(), static_cast(eps)); + ps = torch::empty(x.sizes(), x.options().dtype(torch::kFloat32)); + pc = torch::empty_like(ps); } - indexed_backward_reduce_kernel<<<2, threads, 0, stream>>>( - 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()); + 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/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index f4983dc05..2387c2844 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -71,18 +71,31 @@ 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. The initial dedicated -indexed CUDA slice accepts only FP32 x `[2,8,2]` and int32 index `[2,8]`, with -the usual `[4,6]` modulation. It requires fresh indexed extension symbols. -The small slice has local SM86 coverage using an AdaLN-only extension build; -full-extension and target H100 qualification remain separate. It does not enable strict. - -Other indexed inputs select PyTorch select01 with `fallback=true` and reason -`modulate_index_requires_select01`. If a resolved indexed CUDA slice receives -unsupported geometry, dtype or eps at execution, it explicitly executes the -framework path with reason `indexed_cuda_slice_unsupported`. -Expanded CUDA coverage and dedicated indexed Triton remain follow-up work. -Indexed framework backward has no shared replay-profile qualification. +CPU, ROCm, MUSA and NPU registry entries select PyTorch. Dedicated indexed CUDA +accepts 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. CUDA route storage +is converted to contiguous int32 after validation. + +Indexed CUDA checks 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`. Illegal inputs +raise and clear the previous successful execution observation. Direct CUDA calls +reject unsupported domains instead of silently falling back. + +Current source has CPU regression coverage and a fresh AdaLN-only SM86 local +run: indexed matrix 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 is scoped local evidence; historical +small-slice evidence does not qualify other sources. Dedicated indexed Triton, full-extension +and H100 qualification remain separate. Indexed strict stays closed. ### Strict mode and trace @@ -115,8 +128,8 @@ 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. The small indexed -CUDA slice uses the shared FP32 eps domain. Its per-sample/bank/channel dshift, +Indexed framework fallback checks only finite, nonnegative eps. Dedicated indexed +CUDA uses the shared FP32 eps domain. Its 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 @@ -141,19 +154,20 @@ 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 callable for indexed byte +`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 in the FP32 [2,8,2] slice. +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 fresh current-source GPU evidence. -Ticket 02 adds no wider input-domain or strict qualification claim; indexed strict -remains closed. +Ticket 03 implements the wider CUDA domain; GPU acceptance is pending and indexed +strict remains closed. The CLI records trace snapshots without automatically constructing C1 `BackendProvenance`; its accuracy report alone is not a WS1 system gate. diff --git a/rl_engine/backends/cuda/norm/adaln_modulation.py b/rl_engine/backends/cuda/norm/adaln_modulation.py index 001fec3d4..fea445f16 100644 --- a/rl_engine/backends/cuda/norm/adaln_modulation.py +++ b/rl_engine/backends/cuda/norm/adaln_modulation.py @@ -11,21 +11,35 @@ _eps_supported_fp32, _validate, _validate_select01, + _validate_select01_metadata, ) -def _indexed_slice_supported(x, index, eps): - # shortcut: only the first FP32 tracer bullet; expand with ticket 03 tests. - if not ( - x.is_cuda - and torch.version.hip is None - and x.shape == (2, 8, 2) - and x.dtype == torch.float32 - and index.shape == (2, 8) - and index.dtype == torch.int32 - ): - return False - return _eps_supported_fp32(eps) +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 class _AdaLNIndexedCuda(torch.autograd.Function): @@ -40,6 +54,9 @@ def forward(ctx, x, modulation, index, eps, threads): @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(), @@ -88,15 +105,19 @@ def __call__(self, x, modulation, *, eps=1e-6, modulate_index=None): 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 or not _indexed_slice_supported(x, modulate_index, eps): - raise RuntimeError( - "indexed CUDA slice requires FP32 x [2,8,2] and int32 index [2,8]" - ) + if not self.indexed_available: + raise RuntimeError("indexed CUDA extension symbols are unavailable") return _AdaLNIndexedCuda.apply( x.contiguous(), modulation.contiguous(), - modulate_index.contiguous(), + modulate_index.to(dtype=torch.int32).contiguous(), float(eps), self.threads, ) diff --git a/rl_engine/reference/norm/adaln_indexed.py b/rl_engine/reference/norm/adaln_indexed.py index 64e06c28a..719884f02 100644 --- a/rl_engine/reference/norm/adaln_indexed.py +++ b/rl_engine/reference/norm/adaln_indexed.py @@ -22,8 +22,9 @@ class _IndexedFixed(torch.autograd.Function): @staticmethod def forward(ctx, x, modulation, index, eps): batch, seq, hidden = x.shape - mean = _hidden_tree(x.float()) / hidden - centered = x.float() - mean[..., None] + 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 @@ -31,18 +32,23 @@ def forward(ctx, x, modulation, index, eps): 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 - return norm * (1.0 + scale) + shift, gate + 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 @@ -59,15 +65,15 @@ def backward(ctx, dy, dg): 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 dx, dm, None, None + return None if dx is None else dx.to(ctx.dtype), dm.to(ctx.dtype), None, None class IndexedAdaLNFixedReference: - """FP32 CPU callable for fixed-order byte tests, not a dispatch fallback.""" + """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 != torch.float32 or modulation.dtype != torch.float32: - raise TypeError("indexed byte reference requires FP32 inputs") + 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 index 120d57da0..d4f1934ba 100644 --- a/rl_engine/reference/norm/adaln_modulation.py +++ b/rl_engine/reference/norm/adaln_modulation.py @@ -95,7 +95,7 @@ def backward(ctx, grad_y: torch.Tensor, grad_gate: torch.Tensor): return dx, dm, None -def _validate_select01( +def _validate_select01_metadata( x: torch.Tensor, modulation: torch.Tensor, modulate_index: torch.Tensor, @@ -121,12 +121,16 @@ def _validate_select01( torch.int64, ): raise TypeError("modulate_index must have boolean or integer dtype") - if not torch.all((modulate_index == 0) | (modulate_index == 1)).item(): - raise ValueError("modulate_index values must be 0 or 1") 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] diff --git a/rl_engine/runtime/adaln.py b/rl_engine/runtime/adaln.py index e091a8143..d938b6ec8 100644 --- a/rl_engine/runtime/adaln.py +++ b/rl_engine/runtime/adaln.py @@ -80,12 +80,15 @@ def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): if indexed: implementation = self.implementation if backend in {"cuda", "triton"}: - from rl_engine.backends.cuda.norm.adaln_modulation import _indexed_slice_supported + from rl_engine.backends.cuda.norm.adaln_modulation import _indexed_domain_reason + from rl_engine.reference.norm.adaln_modulation import _validate_select01_metadata + + _validate_select01_metadata(x, modulation, modulate_index, eps) if ( backend == "cuda" and getattr(implementation, "indexed_available", False) - and _indexed_slice_supported(x, modulate_index, eps) + and _indexed_domain_reason(x, modulation, modulate_index, eps) is None ): profile = INDEXED_PROFILE else: @@ -108,7 +111,7 @@ def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): self.trace.update( fallback=True, fallback_reason=( - "indexed_cuda_slice_unsupported" + "indexed_cuda_domain_unsupported" if self._resolution_trace["arithmetic_profile"] == INDEXED_PROFILE else "modulate_index_requires_select01" ), diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index b5e8a0b73..9f563c500 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -1090,12 +1090,10 @@ def get_adaln_modulation_op( platform = self._platform_for_device(device) if modulate_index is not None: backend = OpBackend.PYTORCH_ADALN_MODULATION - # Only the first indexed CUDA slice is accelerated; execution rechecks x/eps. + # Resolution is provisional; execution checks actual geometry, device and eps. if ( platform == "cuda" - and hidden == 2 - and modulate_index.shape == (2, 8) - and modulate_index.dtype == torch.int32 + and hidden <= 4096 ): cuda = self._get_or_create_backend(OpBackend.CUDA_ADALN_MODULATION) if cuda is not None and getattr(cuda, "indexed_available", False): diff --git a/tests/helpers/adaln_indexed.py b/tests/helpers/adaln_indexed.py index a369323ad..9684265bf 100644 --- a/tests/helpers/adaln_indexed.py +++ b/tests/helpers/adaln_indexed.py @@ -287,6 +287,286 @@ def check_shared_bank0(indexed_op, shared_op, device): 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 + 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): + generator = torch.Generator().manual_seed(303) + x = torch.randn(3, 3, hidden, generator=generator).to(dtype) + m = (torch.randn(6, 3 * hidden, generator=generator) / 4).to(dtype) + index = torch.tensor([[0, 1, 0], [1, 0, 1], [0, 0, 1]], dtype=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, + ), + ) diff --git a/tests/ops/norm/test_adaln_indexed.py b/tests/ops/norm/test_adaln_indexed.py index 7a8c836d1..83e2c9ded 100644 --- a/tests/ops/norm/test_adaln_indexed.py +++ b/tests/ops/norm/test_adaln_indexed.py @@ -3,8 +3,9 @@ """L2: indexed AdaLN forward/backward accuracy and byte comparisons on CUDA.""" import pytest +import torch -from tests.helpers.adaln_indexed import check_small_case, require_cuda +from tests.helpers.adaln_indexed import assert_bytes, check_small_case, require_cuda @pytest.fixture @@ -100,3 +101,229 @@ def test_cuda_indexed_all0_matches_shared_arithmetic_and_independent_gold(cuda_o from tests.helpers.adaln_indexed import check_shared_bank0 check_shared_bank0(cuda_op, cuda_op, "cuda") + + +def test_cuda_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_cuda_indexed_bf16_final_rounding(cuda_op): + from tests.helpers.adaln_indexed import check_bf16_rounding + + check_bf16_rounding(cuda_op, "cuda") + + +@pytest.mark.cuda_only +def test_cuda_indexed_bf16_raw16_gate(cuda_op): + from tests.helpers.adaln_indexed import check_bf16_gate_payload + + check_bf16_gate_payload(cuda_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_cuda_indexed_full_route_domain(cuda_op, dtype, index_dtype, broadcast): + from tests.helpers.adaln_indexed import check_domain_routes + + check_domain_routes(cuda_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_cuda_indexed_source_leaf_gradients(cuda_op, dtype, layout): + from tests.helpers.adaln_indexed import check_view_edges + + check_view_edges(cuda_op, "cuda", dtype, layout) + + +@pytest.fixture +def metadata_op(): + 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_cuda_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_cuda_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_cuda_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_cuda_indexed_eps_arithmetic(cuda_op, eps, dtype): + from tests.helpers.adaln_indexed import check_eps_arithmetic + + check_eps_arithmetic(cuda_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 + + 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 = CudaAdaLNModulationOp(threads=threads)(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)) + + +@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_cuda_indexed_bf16_output_and_dx_ties(cuda_op): + from tests.helpers.adaln_indexed import check_bf16_ties + + check_bf16_ties(cuda_op, "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("invalid", ["value", "shape", "dtype", "eps"]) +def test_cuda_indexed_direct_rejects_invalid_inputs(cuda_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): + cuda_op(x, m, modulate_index=index, eps=eps) + + +def test_cuda_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_cuda_indexed_bf16_parameter_fold_final_cast(cuda_op): + from tests.helpers.adaln_indexed import check_bf16_parameter_fold_final_cast + + check_bf16_parameter_fold_final_cast(cuda_op, "cuda") + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("output", ["y", "gate"]) +def test_cuda_indexed_large_upstream_offset_rejected_before_copy(output): + 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" + ) + require_cuda() + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + 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 = CudaAdaLNModulationOp()(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) diff --git a/tests/ops/norm/test_adaln_modulation.py b/tests/ops/norm/test_adaln_modulation.py index d5a717b07..8f64e9196 100644 --- a/tests/ops/norm/test_adaln_modulation.py +++ b/tests/ops/norm/test_adaln_modulation.py @@ -392,9 +392,26 @@ def test_adaln_indexed_matches_independent_autograd(device, dtype, broadcast): _assert_accuracy(modulation.grad, ref_modulation.grad, judgment="gradient_accuracy") _assert_bytes_equal(gate, ref_gate.to(dtype)) assert x.grad.abs().max() > 0 - assert trace["fallback"] is True - assert trace["fallback_reason"] == "modulate_index_requires_select01" - assert "framework_select01_backward" in trace["reduction_order"] + accelerated = False + if device == "cuda" and torch.version.hip is None: + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp + + try: + accelerated = CudaAdaLNModulationOp().indexed_available + except RuntimeError: + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + raise + if accelerated: + assert trace["actual_backend"] == "cuda" + 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( diff --git a/tests/reference/test_adaln_indexed.py b/tests/reference/test_adaln_indexed.py index 5dc6427bc..53fbb169e 100644 --- a/tests/reference/test_adaln_indexed.py +++ b/tests/reference/test_adaln_indexed.py @@ -79,14 +79,17 @@ def test_indexed_byte_reference_preserves_selected_token_fold_and_sample_routing ) 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], - ]), + 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 @@ -158,3 +161,105 @@ 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") diff --git a/tests/runtime/test_adaln_indexed.py b/tests/runtime/test_adaln_indexed.py index 20880e629..9e4cd8087 100644 --- a/tests/runtime/test_adaln_indexed.py +++ b/tests/runtime/test_adaln_indexed.py @@ -30,9 +30,7 @@ def test_indexed_resolution_reports_backend_policy_without_execution(device): "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["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" ) @@ -42,7 +40,8 @@ def test_indexed_resolution_reports_backend_policy_without_execution(device): ) assert trace["reduction_order"] == ( "pairwise_lower_upper_H; ascending_selected_token_S_per_bank" - if accelerated else "pairwise_lower_upper_H; framework_select01_backward" + 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")) @@ -77,16 +76,19 @@ def test_cuda_indexed_execution_rechecks_support_and_clears_failed_observation(u op(x, m, modulate_index=index, eps=7.0) assert trace["actual_backend"] == "cuda" - fallback_x = x[:, :7] if unsupported == "geometry" else x - fallback_index = index[:, :7] if unsupported == "geometry" else index + 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, m, modulate_index=fallback_index, eps=eps) - expected = NativeAdaLNModulationOp()(fallback_x, m, modulate_index=fallback_index, eps=eps) + 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_slice_unsupported" + 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"): @@ -114,6 +116,7 @@ def test_cuda_indexed_resolution_can_execute_shared_without_stale_profile(): 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): @@ -127,3 +130,123 @@ def test_cuda_indexed_public_dispatch_preserves_absent_vs_zero_dy(explicit_zero) 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" From d6c894d916f47dcbfd2000a2b36ffcc009c75519 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Sat, 24 Oct 2026 00:29:57 +0800 Subject: [PATCH 33/37] feat(adaln): implement indexed Triton forward and backward Add the indexed FP32 tree and selected-token fold profile with BF16 bit-preserving gates, autograd participation, shared metadata checks, and truthful dispatch fallback traces. Validate public behavior with independent references and local SM86 GPU evidence: 297 indexed tests, 13 shared regressions, and 129 tests per sanitizer. Resource-limited oversized-upstream and real-workload acceptance remains pending; indexed strict stays closed. Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 67 ++++-- .../backends/cuda/norm/adaln_modulation.py | 30 +-- .../shared/triton/norm/adaln_indexed.py | 215 ++++++++++++++++++ .../shared/triton/norm/adaln_modulation.py | 14 +- rl_engine/reference/norm/adaln_modulation.py | 27 +++ rl_engine/runtime/adaln.py | 9 +- rl_engine/runtime/registry.py | 19 +- tests/helpers/adaln_indexed.py | 85 ++++++- tests/ops/norm/test_adaln_indexed.py | 148 ++++++------ tests/ops/norm/test_adaln_indexed_triton.py | 134 +++++++++++ tests/ops/norm/test_adaln_modulation.py | 16 +- tests/reference/test_adaln_indexed.py | 17 ++ tests/runtime/test_adaln_indexed_triton.py | 210 +++++++++++++++++ 13 files changed, 848 insertions(+), 143 deletions(-) create mode 100644 rl_engine/backends/shared/triton/norm/adaln_indexed.py create mode 100644 tests/ops/norm/test_adaln_indexed_triton.py create mode 100644 tests/runtime/test_adaln_indexed_triton.py diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index 2387c2844..c9d06e6b3 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -72,12 +72,12 @@ 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 -accepts FP32/BF16, H <= 4096, B <= 65535, both route shapes and +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. CUDA route storage +expanded; internal copies preserve gradients to source leaves. Indexed route storage is converted to contiguous int32 after validation. -Indexed CUDA checks element counts and accessed element offsets before copies or +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. @@ -85,17 +85,27 @@ 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`. Illegal inputs -raise and clear the previous successful execution observation. Direct CUDA calls +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. -Current source has CPU regression coverage and a fresh AdaLN-only SM86 local -run: indexed matrix 361 PASS (including CPU reference/metadata checks), small +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 is scoped local evidence; historical -small-slice evidence does not qualify other sources. Dedicated indexed Triton, full-extension -and H100 qualification remain separate. Indexed strict stays closed. +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 @@ -105,7 +115,8 @@ unknown backends and PyTorch execution on GPU are rejected. Strict mode rechecks the input device and hidden size at execution. It enforces the declared calculation policy, not hardware qualification. The indexed path -does not claim fixed-order CUDA/Triton backward equivalence. +declares its own profile but has no qualified CUDA/Triton backward equivalence +claim. The returned callable is an `AdaLNDispatch` wrapper with an `implementation` attribute. `selected_backend` records resolution; `actual_backend`, dtype and @@ -129,7 +140,7 @@ FP32. Positive values that round to zero are rejected; FP32 subnormals, includin 2^-149, are accepted. Indexed framework fallback checks only finite, nonnegative eps. Dedicated indexed -CUDA uses the shared FP32 eps domain. Its per-sample/bank/channel dshift, +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 @@ -165,9 +176,10 @@ 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 fresh current-source GPU evidence. -Ticket 03 implements the wider CUDA domain; GPU acceptance is pending and indexed -strict remains closed. +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. @@ -218,6 +230,31 @@ 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. Strict indexed execution stays rejected. + +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 diff --git a/rl_engine/backends/cuda/norm/adaln_modulation.py b/rl_engine/backends/cuda/norm/adaln_modulation.py index fea445f16..82fa6b312 100644 --- a/rl_engine/backends/cuda/norm/adaln_modulation.py +++ b/rl_engine/backends/cuda/norm/adaln_modulation.py @@ -8,40 +8,14 @@ from rl_engine.backends.extension import _C, _EXT_AVAILABLE from rl_engine.reference.norm.adaln_modulation import ( - _eps_supported_fp32, + _indexed_domain_reason, + _int32_addressable, _validate, _validate_select01, _validate_select01_metadata, ) -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 - - class _AdaLNIndexedCuda(torch.autograd.Function): @staticmethod def forward(ctx, x, modulation, index, eps, 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..66484d0ed --- /dev/null +++ b/rl_engine/backends/shared/triton/norm/adaln_indexed.py @@ -0,0 +1,215 @@ +# 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) + _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 index 2c8654cf5..8381b8661 100644 --- a/rl_engine/backends/shared/triton/norm/adaln_modulation.py +++ b/rl_engine/backends/shared/triton/norm/adaln_modulation.py @@ -182,10 +182,18 @@ def __init__(self, *, num_warps=4, reduction_tile=128): self.num_warps = num_warps self.reduction_tile = reduction_tile - def __call__(self, x, modulation, *, eps=1e-6): - return self.forward(x, modulation, eps=eps) + indexed_available = True - def forward(self, x, modulation, *, eps=1e-6): + 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") diff --git a/rl_engine/reference/norm/adaln_modulation.py b/rl_engine/reference/norm/adaln_modulation.py index d4f1934ba..e9da50567 100644 --- a/rl_engine/reference/norm/adaln_modulation.py +++ b/rl_engine/reference/norm/adaln_modulation.py @@ -44,6 +44,33 @@ def _eps_supported_fp32(eps: float) -> bool: 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]") diff --git a/rl_engine/runtime/adaln.py b/rl_engine/runtime/adaln.py index d938b6ec8..f9d27e71d 100644 --- a/rl_engine/runtime/adaln.py +++ b/rl_engine/runtime/adaln.py @@ -80,14 +80,13 @@ def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): if indexed: implementation = self.implementation if backend in {"cuda", "triton"}: - from rl_engine.backends.cuda.norm.adaln_modulation import _indexed_domain_reason + from rl_engine.reference.norm.adaln_modulation import _indexed_domain_reason from rl_engine.reference.norm.adaln_modulation import _validate_select01_metadata _validate_select01_metadata(x, modulation, modulate_index, eps) if ( - backend == "cuda" - and getattr(implementation, "indexed_available", False) + getattr(implementation, "indexed_available", False) and _indexed_domain_reason(x, modulation, modulate_index, eps) is None ): profile = INDEXED_PROFILE @@ -111,7 +110,7 @@ def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): self.trace.update( fallback=True, fallback_reason=( - "indexed_cuda_domain_unsupported" + f"indexed_{adaln_backend(self.implementation)}_domain_unsupported" if self._resolution_trace["arithmetic_profile"] == INDEXED_PROFILE else "modulate_index_requires_select01" ), @@ -123,7 +122,7 @@ def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): fallback=False, fallback_reason=None, reduction_order="pairwise_lower_upper_H; ascending_selected_token_S_per_bank", - kernel_fingerprint="adaln_indexed_cuda_v1", + kernel_fingerprint=f"adaln_indexed_{backend}_v1", ) elif self._resolution_trace["arithmetic_profile"] == INDEXED_PROFILE: self.trace.update( diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 9f563c500..e69f22e50 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -1095,13 +1095,20 @@ def get_adaln_modulation_op( platform == "cuda" and hidden <= 4096 ): - cuda = self._get_or_create_backend(OpBackend.CUDA_ADALN_MODULATION) - if cuda is not None and getattr(cuda, "indexed_available", False): - backend = OpBackend.CUDA_ADALN_MODULATION + 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 OpBackend.CUDA_ADALN_MODULATION + accelerated = backend is not OpBackend.PYTORCH_ADALN_MODULATION + actual = adaln_backend(op) trace = { "selected_backend": backend.name, "fallback": not accelerated, @@ -1117,11 +1124,11 @@ def get_adaln_modulation_op( "stream_k": False, "tf32": False, "kernel_fingerprint": ( - "adaln_indexed_cuda_v1" + f"adaln_indexed_{actual}_v1" if accelerated else "adaln_modulation_select01_native_v1" ), - "requested_backend": "cuda" if accelerated else "pytorch", + "requested_backend": actual, "arithmetic_profile": INDEXED_PROFILE if accelerated else profile, } return AdaLNDispatch(op, trace, hidden=hidden, strict=strict), trace diff --git a/tests/helpers/adaln_indexed.py b/tests/helpers/adaln_indexed.py index 9684265bf..802e429f1 100644 --- a/tests/helpers/adaln_indexed.py +++ b/tests/helpers/adaln_indexed.py @@ -454,11 +454,17 @@ def check_eps_arithmetic(op, device, dtype, eps): assert_bytes(dm[1], torch.zeros(6, dtype=dtype)) -def check_geometry_bytes(op, device, dtype, hidden): +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, 3, hidden, generator=generator).to(dtype) + x = torch.randn(3, seq, hidden, generator=generator).to(dtype) m = (torch.randn(6, 3 * hidden, generator=generator) / 4).to(dtype) - index = torch.tensor([[0, 1, 0], [1, 0, 1], [0, 0, 1]], dtype=torch.int32) + 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) @@ -570,3 +576,76 @@ def check_bf16_parameter_fold_final_cast(op, device): 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 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 index 83e2c9ded..40066d122 100644 --- a/tests/ops/norm/test_adaln_indexed.py +++ b/tests/ops/norm/test_adaln_indexed.py @@ -1,15 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""L2: indexed AdaLN forward/backward accuracy and byte comparisons on CUDA.""" +"""L2: indexed CUDA/Triton accuracy, bytes and public support boundaries.""" import pytest import torch -from tests.helpers.adaln_indexed import assert_bytes, check_small_case, require_cuda +from tests.helpers.adaln_indexed import check_small_case, require_cuda -@pytest.fixture -def cuda_op(): +@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 @@ -17,48 +24,48 @@ def cuda_op(): @pytest.mark.cuda_only -def test_cuda_indexed_small_fp32_direct_backend(cuda_op): - check_small_case(cuda_op, "cuda") +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_cuda_indexed_route_identity_literal(cuda_op, route): +def test_indexed_route_identity_literal(indexed_op, route): from tests.helpers.adaln_indexed import check_route_identity - check_route_identity(cuda_op, "cuda", route) + check_route_identity(indexed_op, "cuda", route) @pytest.mark.cuda_only @pytest.mark.parametrize("explicit_zero", [False, True]) -def test_cuda_indexed_absent_vs_zero_dy_nonfinite_literal(cuda_op, explicit_zero): +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(cuda_op, "cuda", explicit_zero) + check_absent_vs_zero_dy(indexed_op, "cuda", explicit_zero) @pytest.mark.cuda_only @pytest.mark.parametrize("order", ["lost_unit", "retained_unit"]) -def test_cuda_indexed_selected_token_cancellation_literal(cuda_op, order): +def test_indexed_selected_token_cancellation_literal(indexed_op, order): from tests.helpers.adaln_indexed import check_cancellation - check_cancellation(cuda_op, "cuda", order) + check_cancellation(indexed_op, "cuda", order) @pytest.mark.cuda_only -def test_cuda_indexed_dscale_product_rounds_before_add_literal(cuda_op): +def test_indexed_dscale_product_rounds_before_add_literal(indexed_op): from tests.helpers.adaln_indexed import check_product_rounding - check_product_rounding(cuda_op, "cuda") + 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_cuda_indexed_nonfinite_selected_bank_isolation_literal(cuda_op, bank, nonfinite): +def test_indexed_nonfinite_selected_bank_isolation_literal(indexed_op, bank, nonfinite): from tests.helpers.adaln_indexed import check_nonfinite_isolation - check_nonfinite_isolation(cuda_op, "cuda", bank, nonfinite) + check_nonfinite_isolation(indexed_op, "cuda", bank, nonfinite) @pytest.mark.cuda_only @@ -75,35 +82,35 @@ def test_cuda_indexed_nonfinite_selected_bank_isolation_literal(cuda_op, bank, n ("modulation", "both"), ], ) -def test_cuda_indexed_whole_leaf_participation_literal(cuda_op, frozen, outputs): +def test_indexed_whole_leaf_participation_literal(indexed_op, frozen, outputs): from tests.helpers.adaln_indexed import check_leaf_participation - check_leaf_participation(cuda_op, "cuda", frozen, outputs) + check_leaf_participation(indexed_op, "cuda", frozen, outputs) @pytest.mark.cuda_only -def test_cuda_indexed_gate_copies_original_fp32_payload_literal(cuda_op): +def test_indexed_gate_copies_original_fp32_payload_literal(indexed_op): from tests.helpers.adaln_indexed import check_gate_payload - check_gate_payload(cuda_op, "cuda") + check_gate_payload(indexed_op, "cuda") @pytest.mark.cuda_only @pytest.mark.parametrize("contribution", ["negative_zero", "subnormal"]) -def test_cuda_indexed_signed_zero_and_subnormal_fold_literal(cuda_op, contribution): +def test_indexed_signed_zero_and_subnormal_fold_literal(indexed_op, contribution): from tests.helpers.adaln_indexed import check_tiny_contributions - check_tiny_contributions(cuda_op, "cuda", contribution) + check_tiny_contributions(indexed_op, "cuda", contribution) @pytest.mark.cuda_only -def test_cuda_indexed_all0_matches_shared_arithmetic_and_independent_gold(cuda_op): +def test_indexed_all0_matches_shared_arithmetic_and_independent_gold(indexed_op): from tests.helpers.adaln_indexed import check_shared_bank0 - check_shared_bank0(cuda_op, cuda_op, "cuda") + check_shared_bank0(indexed_op, indexed_op, "cuda") -def test_cuda_indexed_two_bank_metadata_overflow_precedes_value_reads(metadata_op): +def test_indexed_two_bank_metadata_overflow_precedes_value_reads(metadata_op): import torch op = metadata_op @@ -116,17 +123,17 @@ def test_cuda_indexed_two_bank_metadata_overflow_precedes_value_reads(metadata_o @pytest.mark.cuda_only -def test_cuda_indexed_bf16_final_rounding(cuda_op): +def test_indexed_bf16_final_rounding(indexed_op): from tests.helpers.adaln_indexed import check_bf16_rounding - check_bf16_rounding(cuda_op, "cuda") + check_bf16_rounding(indexed_op, "cuda") @pytest.mark.cuda_only -def test_cuda_indexed_bf16_raw16_gate(cuda_op): +def test_indexed_bf16_raw16_gate(indexed_op): from tests.helpers.adaln_indexed import check_bf16_gate_payload - check_bf16_gate_payload(cuda_op, "cuda") + check_bf16_gate_payload(indexed_op, "cuda") @pytest.mark.cuda_only @@ -135,23 +142,28 @@ def test_cuda_indexed_bf16_raw16_gate(cuda_op): "index_dtype", [torch.bool, torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64] ) @pytest.mark.parametrize("broadcast", [False, True]) -def test_cuda_indexed_full_route_domain(cuda_op, dtype, index_dtype, broadcast): +def test_indexed_full_route_domain(indexed_op, dtype, index_dtype, broadcast): from tests.helpers.adaln_indexed import check_domain_routes - check_domain_routes(cuda_op, "cuda", dtype, index_dtype, broadcast) + 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_cuda_indexed_source_leaf_gradients(cuda_op, dtype, layout): +def test_indexed_source_leaf_gradients(indexed_op, dtype, layout): from tests.helpers.adaln_indexed import check_view_edges - check_view_edges(cuda_op, "cuda", dtype, layout) + 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 -@pytest.fixture -def metadata_op(): + return TritonAdaLNModulationOp() import os from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp @@ -172,7 +184,7 @@ def metadata_op(): ((1, 524289, 4096), "int32 indexing"), ], ) -def test_cuda_indexed_metadata_rejects_before_copy_or_allocation(metadata_op, shape, reason): +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") @@ -182,7 +194,7 @@ def test_cuda_indexed_metadata_rejects_before_copy_or_allocation(metadata_op, sh @pytest.mark.parametrize("which", ["x", "m", "index"]) -def test_cuda_indexed_metadata_rejects_strided_access_span(metadata_op, which): +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") @@ -197,7 +209,7 @@ def test_cuda_indexed_metadata_rejects_strided_access_span(metadata_op, which): @pytest.mark.parametrize("shape", [(1, 2**31 - 1, 1), (65535, 1, 4096)]) -def test_cuda_indexed_positive_metadata_bounds_reach_device_check(metadata_op, shape): +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( @@ -210,10 +222,10 @@ def test_cuda_indexed_positive_metadata_bounds_reach_device_check(metadata_op, s @pytest.mark.cuda_only @pytest.mark.parametrize("eps", [0.0, 2**-149, 2**-126]) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) -def test_cuda_indexed_eps_arithmetic(cuda_op, eps, dtype): +def test_indexed_eps_arithmetic(indexed_op, eps, dtype): from tests.helpers.adaln_indexed import check_eps_arithmetic - check_eps_arithmetic(cuda_op, "cuda", dtype, eps) + check_eps_arithmetic(indexed_op, "cuda", dtype, eps) @pytest.mark.cuda_only @@ -223,27 +235,9 @@ def test_cuda_indexed_maximum_batch_both_banks_nonzero_dx(dtype, threads): require_cuda() from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp - 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 = CudaAdaLNModulationOp(threads=threads)(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)) + from tests.helpers.adaln_indexed import check_maximum_batch + + check_maximum_batch(CudaAdaLNModulationOp(threads=threads), dtype) @pytest.mark.cuda_only @@ -259,15 +253,15 @@ def test_cuda_indexed_threads_geometry_batch_layout_bytes(threads, dtype, hidden @pytest.mark.cuda_only -def test_cuda_indexed_bf16_output_and_dx_ties(cuda_op): +def test_indexed_bf16_output_and_dx_ties(indexed_op): from tests.helpers.adaln_indexed import check_bf16_ties - check_bf16_ties(cuda_op, "cuda") + check_bf16_ties(indexed_op, "cuda") @pytest.mark.cuda_only @pytest.mark.parametrize("invalid", ["value", "shape", "dtype", "eps"]) -def test_cuda_indexed_direct_rejects_invalid_inputs(cuda_op, invalid): +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) @@ -284,10 +278,10 @@ def test_cuda_indexed_direct_rejects_invalid_inputs(cuda_op, invalid): else: eps, message = float("nan"), "finite and nonnegative" with pytest.raises(error, match=message): - cuda_op(x, m, modulate_index=index, eps=eps) + indexed_op(x, m, modulate_index=index, eps=eps) -def test_cuda_indexed_int32_maximum_access_offset_metadata(metadata_op): +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") @@ -297,33 +291,41 @@ def test_cuda_indexed_int32_maximum_access_offset_metadata(metadata_op): @pytest.mark.cuda_only -def test_cuda_indexed_bf16_parameter_fold_final_cast(cuda_op): +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(cuda_op, "cuda") + check_bf16_parameter_fold_final_cast(indexed_op, "cuda") @pytest.mark.cuda_only @pytest.mark.parametrize("output", ["y", "gate"]) -def test_cuda_indexed_large_upstream_offset_rejected_before_copy(output): +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" ) - require_cuda() - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp - + # 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 = CudaAdaLNModulationOp()(x, m, modulate_index=index, eps=7.0) + 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_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 index 8f64e9196..04f9486cc 100644 --- a/tests/ops/norm/test_adaln_modulation.py +++ b/tests/ops/norm/test_adaln_modulation.py @@ -392,17 +392,13 @@ def test_adaln_indexed_matches_independent_autograd(device, dtype, broadcast): _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 = False - if device == "cuda" and torch.version.hip is None: - from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp - - try: - accelerated = CudaAdaLNModulationOp().indexed_available - except RuntimeError: - if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": - raise + 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" + 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"] diff --git a/tests/reference/test_adaln_indexed.py b/tests/reference/test_adaln_indexed.py index 53fbb169e..fc03878a5 100644 --- a/tests/reference/test_adaln_indexed.py +++ b/tests/reference/test_adaln_indexed.py @@ -263,3 +263,20 @@ def test_indexed_reference_bf16_parameter_fold_cast_counterexample(): # 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_indexed_triton.py b/tests/runtime/test_adaln_indexed_triton.py new file mode 100644 index 000000000..17644e434 --- /dev/null +++ b/tests/runtime/test_adaln_indexed_triton.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""L3: real missing-extension resolution and indexed Triton execution traces.""" + +import json +import os +from pathlib import Path +import shutil +import subprocess +import sys + +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_available", [True, False]) +def test_indexed_missing_extension_real_resolution(tmp_path, triton_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__"), + ) + script = ( + f"TRITON_AVAILABLE = {triton_available!r}\n" + + """ +import pathlib, torch, rl_engine, importlib.util +assert (importlib.util.find_spec("triton") is not None) == TRITON_AVAILABLE +from rl_engine.backends.extension import _EXT_AVAILABLE +from rl_engine.runtime.registry import KernelRegistry +assert pathlib.Path(rl_engine.__file__).resolve().parent.parent == pathlib.Path.cwd() +assert not _EXT_AVAILABLE +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 +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 +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)] + 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 + + +@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 From 5490223a850fe9cb1efd5c3cceeb0fb964ee5c1c Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Sat, 24 Oct 2026 12:17:59 +0800 Subject: [PATCH 34/37] feat(adaln): add default-closed indexed strict activation Add explicit backend/device/runtime activation scopes and call-time strict policy guards. Preserve shared behavior, real Native fallback, input-error precedence and failed-call trace cleanup. Fix default ROCm device resolution and avoid an extra dispatch route scan. Cover real native-symbol/dependency isolation, mode/domain changes, route-scan observations, activated L3 and direct backend integration. Validation: - Focused CPU regression: 115 passed, 59 GPU/opt-in cases skipped. - Fresh local SM86 evidence: CUDA 278, Triton 319, integration 1 and shared 41/17 passed with zero skipped or xfailed cases. - Four sanitizer tools per backend: CUDA 199/Triton 153 each passed. - Ruff, isort, Black formatting, strict MkDocs build and diff checks. - Commit formatting preserves the GPU snapshot AST after normalizing import order; raw evidence retains its original hashes. Four large-offset/target suites remain resource-blocked. Complete instruction semantics and sanitizer kernel coverage review is pending. Indexed strict delivery qualification remains default-closed. Signed-off-by: hongleng-github --- docs/operators/adaln-modulation.md | 42 ++- rl_engine/runtime/adaln.py | 175 +++++++++-- rl_engine/runtime/registry.py | 61 +++- tests/helpers/adaln_indexed.py | 39 +++ .../norm/test_adaln_indexed_integration.py | 64 ++++ .../runtime/test_adaln_indexed_activation.py | 139 +++++++++ tests/runtime/test_adaln_indexed_triton.py | 117 +++++++- tests/runtime/test_adaln_strict_policy.py | 277 ++++++++++++++++++ 8 files changed, 856 insertions(+), 58 deletions(-) create mode 100644 tests/ops/norm/test_adaln_indexed_integration.py create mode 100644 tests/runtime/test_adaln_indexed_activation.py create mode 100644 tests/runtime/test_adaln_strict_policy.py diff --git a/docs/operators/adaln-modulation.md b/docs/operators/adaln-modulation.md index c9d06e6b3..e39af1f04 100644 --- a/docs/operators/adaln-modulation.md +++ b/docs/operators/adaln-modulation.md @@ -109,14 +109,31 @@ Full-extension and H100 qualification remain separate. Indexed strict stays clos ### Strict mode and trace -Pass `strict=True` when resolving the shared path to require its declared arithmetic -profile. CPU Native and NVIDIA CUDA/Triton are permitted; indexed modulation, -unknown backends and PyTorch execution on GPU are rejected. - -Strict mode rechecks the input device and hidden size at execution. It enforces -the declared calculation policy, not hardware qualification. The indexed path -declares its own profile but has no qualified CUDA/Triton backward equivalence -claim. +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 @@ -220,7 +237,11 @@ Shared CUDA and Triton require at most 2^31-1 elements in each input (`x` and contiguous copies or output/partial allocation. The CUDA extension additionally requires B <= 65535 and repeats its bounds checks in C++ before launch. -Exceeding these limits raises an error; resolution only uses H to select fallback. +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) @@ -248,7 +269,8 @@ equivalent final-result specialization, not multi-token low-precision accumulati 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. Strict indexed execution stays rejected. +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 diff --git a/rl_engine/runtime/adaln.py b/rl_engine/runtime/adaln.py index f9d27e71d..2fbd75830 100644 --- a/rl_engine/runtime/adaln.py +++ b/rl_engine/runtime/adaln.py @@ -5,6 +5,7 @@ from __future__ import annotations import sys +from dataclasses import dataclass from typing import Any import torch @@ -19,6 +20,63 @@ } +@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__) @@ -30,16 +88,26 @@ def adaln_backend(op: Any) -> str | None: class AdaLNDispatch: """A resolution-local trace, updated only after successful operator execution. - Strict enforces the declared arithmetic, not hardware qualification. CUDA and - Triton share the profile; Native on GPU is excluded pending its replay audit. + 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): + 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.arithmetic_profile = trace["arithmetic_profile"] + self.indexed_qualification = indexed_qualification + self.device = torch.device(device) if device is not None else None trace.update( actual_backend=None, observation=None, @@ -64,31 +132,69 @@ def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): backend = adaln_backend(self.implementation) indexed = modulate_index is not None profile = SELECT01_PROFILE if indexed else SHARED_PROFILE - if self.strict: - if indexed: - raise RuntimeError("strict AdaLN rejects select01 framework backward") - 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 self.arithmetic_profile != profile: - raise RuntimeError("strict AdaLN arithmetic profile mismatch") - if x.ndim != 3 or x.shape[-1] != self.hidden: - raise ValueError("strict AdaLN input hidden differs from resolved hidden") + implementation = self.implementation + shared_unsupported = False if indexed: - implementation = self.implementation - if backend in {"cuda", "triton"}: - from rl_engine.reference.norm.adaln_modulation import _indexed_domain_reason - from rl_engine.reference.norm.adaln_modulation import _validate_select01_metadata - - _validate_select01_metadata(x, modulation, modulate_index, eps) + from rl_engine.reference.norm.adaln_modulation import ( + _indexed_domain_reason, + _validate_select01, + _validate_select01_metadata, + ) - if ( - getattr(implementation, "indexed_available", False) - and _indexed_domain_reason(x, modulation, modulate_index, eps) is None + _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 @@ -97,7 +203,12 @@ def forward(self, x, modulation, *, eps=1e-6, modulate_index=None): backend = "pytorch" outputs = implementation(x, modulation, eps=eps, modulate_index=modulate_index) else: - outputs = self.implementation(x, modulation, eps=eps) + 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", @@ -124,11 +235,19 @@ def forward(self, x, modulation, *, eps=1e-6, modulate_index=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"] == INDEXED_PROFILE: + 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 diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index e69f22e50..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,9 +178,7 @@ 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" - ) + PYTORCH_ADALN_MODULATION = "rl_engine.reference.norm.adaln_modulation.NativeAdaLNModulationOp" TRITON_ADALN_MODULATION = ( "rl_engine.backends.shared.triton.norm.adaln_modulation.TritonAdaLNModulationOp" ) @@ -1071,8 +1070,9 @@ def get_adaln_modulation_op( hidden: int, modulate_index: torch.Tensor | None = None, strict: bool = False, + indexed_qualification: IndexedAdaLNQualification | None = None, ) -> tuple[Any, dict[str, Any]]: - """Resolve AdaLN; strict permits only the shared declared replay profile.""" + """Resolve AdaLN; indexed strict requires an explicit AdaLN-local scope.""" from rl_engine.runtime.adaln import ( INDEXED_PROFILE, @@ -1084,19 +1084,28 @@ def get_adaln_modulation_op( 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 - if strict and modulate_index is not None: - raise RuntimeError("strict AdaLN rejects select01 framework backward") 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 - ): + if platform == "cuda" and hidden <= 4096: for candidate in ( - OpBackend.CUDA_ADALN_MODULATION, OpBackend.TRITON_ADALN_MODULATION + OpBackend.CUDA_ADALN_MODULATION, + OpBackend.TRITON_ADALN_MODULATION, ): implementation = self._get_or_create_backend(candidate) if implementation is not None and getattr( @@ -1109,6 +1118,12 @@ def get_adaln_modulation_op( 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, @@ -1131,7 +1146,17 @@ def get_adaln_modulation_op( "requested_backend": actual, "arithmetic_profile": INDEXED_PROFILE if accelerated else profile, } - return AdaLNDispatch(op, trace, hidden=hidden, strict=strict), trace + 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: @@ -1160,7 +1185,17 @@ def get_adaln_modulation_op( "requested_backend": candidates[0].name.split("_")[0].lower(), "fallback_reason": "; ".join(rejected) if rejected else None, } - return AdaLNDispatch(op, trace, hidden=hidden, strict=strict), trace + 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: diff --git a/tests/helpers/adaln_indexed.py b/tests/helpers/adaln_indexed.py index 802e429f1..fc4a17131 100644 --- a/tests/helpers/adaln_indexed.py +++ b/tests/helpers/adaln_indexed.py @@ -584,6 +584,45 @@ def require_triton_gpu(): 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_() 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/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 index 17644e434..49d008bd4 100644 --- a/tests/runtime/test_adaln_indexed_triton.py +++ b/tests/runtime/test_adaln_indexed_triton.py @@ -2,12 +2,13 @@ # Copyright (c) 2026 RL-Kernel Contributors """L3: real missing-extension resolution and indexed Triton execution traces.""" +import hashlib import json import os -from pathlib import Path import shutil import subprocess import sys +from pathlib import Path import pytest import torch @@ -16,8 +17,19 @@ from tests.helpers.adaln_indexed import assert_bytes, check_small_case, small_case -@pytest.mark.parametrize("triton_available", [True, False]) -def test_indexed_missing_extension_real_resolution(tmp_path, triton_available): +@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: @@ -29,15 +41,53 @@ def test_indexed_missing_extension_real_resolution(tmp_path, triton_available): 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_AVAILABLE -from rl_engine.backends.extension import _EXT_AVAILABLE +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 not _EXT_AVAILABLE +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__ == ( @@ -49,6 +99,35 @@ def test_indexed_missing_extension_real_resolution(tmp_path, 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) @@ -68,6 +147,21 @@ def test_indexed_missing_extension_real_resolution(tmp_path, triton_available): 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: @@ -93,6 +187,13 @@ def test_indexed_missing_extension_real_resolution(tmp_path, triton_available): 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], @@ -103,6 +204,8 @@ def test_indexed_missing_extension_real_resolution(tmp_path, triton_available): 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 diff --git a/tests/runtime/test_adaln_strict_policy.py b/tests/runtime/test_adaln_strict_policy.py new file mode 100644 index 000000000..ec0d54717 --- /dev/null +++ b/tests/runtime/test_adaln_strict_policy.py @@ -0,0 +1,277 @@ +# 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 + + +def test_explicit_indexed_activation_resolves_without_enabling_cpu_select01(): + 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(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(mismatch): + from rl_engine.runtime.adaln import IndexedAdaLNQualification + + 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, + ) + 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) From f23b389fdad223b376688377ba9942dc42bf2a0e Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 28 Oct 2026 15:30:25 +0800 Subject: [PATCH 35/37] style(adaln): apply CI-pinned Black and isort formatting Signed-off-by: hongleng-github --- .../norm/benchmark_adaln_modulation.py | 4 +-- rl_engine/_C.pyi | 8 +++-- .../validation/operators/operator_specs.py | 4 +-- tests/ops/norm/test_adaln_indexed.py | 2 +- tests/ops/norm/test_adaln_modulation.py | 10 +++--- tests/reference/test_adaln_indexed.py | 3 ++ tests/runtime/test_adaln_dispatch.py | 17 +++++----- tests/runtime/test_adaln_indexed_triton.py | 3 +- .../validation/operators/test_adaln_gtest.py | 6 ++-- tests/validation/operators/test_op_checks.py | 32 +++++++++++++------ 10 files changed, 55 insertions(+), 34 deletions(-) diff --git a/benchmarks/operators/norm/benchmark_adaln_modulation.py b/benchmarks/operators/norm/benchmark_adaln_modulation.py index 605f02cae..fc7d7dc2e 100644 --- a/benchmarks/operators/norm/benchmark_adaln_modulation.py +++ b/benchmarks/operators/norm/benchmark_adaln_modulation.py @@ -10,11 +10,11 @@ import contextlib import hashlib import json -from pathlib import Path import platform import statistics import subprocess import sys +from pathlib import Path import torch @@ -103,8 +103,8 @@ def main(argv=None): 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.runtime.adaln import SHARED_PROFILE, AdaLNDispatch from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp + from rl_engine.runtime.adaln import SHARED_PROFILE, AdaLNDispatch backends = [("pytorch", NativeAdaLNModulationOp())] skipped = [] diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index bb194f8fb..8c1953fca 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -260,8 +260,12 @@ 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 + 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, diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 9a43d3923..5b2ac615d 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -40,9 +40,7 @@ def _load_object(path: str) -> Any: gold_path="rl_engine.validation.operators.operator_specs.GtestAdaLNReference", gold_method="forward_fp32", candidate_paths={ - "pytorch": ( - "rl_engine.reference.norm.adaln_modulation.NativeAdaLNModulationOp" - ), + "pytorch": ("rl_engine.reference.norm.adaln_modulation.NativeAdaLNModulationOp"), "triton": ( "rl_engine.backends.shared.triton.norm.adaln_modulation.TritonAdaLNModulationOp" ), diff --git a/tests/ops/norm/test_adaln_indexed.py b/tests/ops/norm/test_adaln_indexed.py index 40066d122..fbb07ea1c 100644 --- a/tests/ops/norm/test_adaln_indexed.py +++ b/tests/ops/norm/test_adaln_indexed.py @@ -165,6 +165,7 @@ def metadata_op(request): return TritonAdaLNModulationOp() import os + from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp try: @@ -234,7 +235,6 @@ def test_indexed_eps_arithmetic(indexed_op, eps, dtype): 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) diff --git a/tests/ops/norm/test_adaln_modulation.py b/tests/ops/norm/test_adaln_modulation.py index 04f9486cc..70c18945c 100644 --- a/tests/ops/norm/test_adaln_modulation.py +++ b/tests/ops/norm/test_adaln_modulation.py @@ -3,10 +3,10 @@ import pytest import torch -from rl_engine.validation.operators.operator_specs import GtestAdaLNReference 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): @@ -392,9 +392,11 @@ def test_adaln_indexed_matches_independent_autograd(device, dtype, broadcast): _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"} + 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" diff --git a/tests/reference/test_adaln_indexed.py b/tests/reference/test_adaln_indexed.py index fc03878a5..6010c440b 100644 --- a/tests/reference/test_adaln_indexed.py +++ b/tests/reference/test_adaln_indexed.py @@ -242,6 +242,7 @@ def test_indexed_reference_geometry_and_batch_layout(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 @@ -253,6 +254,7 @@ def test_indexed_reference_bf16_output_and_dx_ties(): 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) @@ -267,6 +269,7 @@ def test_indexed_reference_bf16_parameter_fold_cast_counterexample(): 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 diff --git a/tests/runtime/test_adaln_dispatch.py b/tests/runtime/test_adaln_dispatch.py index 2030a31aa..e69b7d964 100644 --- a/tests/runtime/test_adaln_dispatch.py +++ b/tests/runtime/test_adaln_dispatch.py @@ -3,15 +3,12 @@ """AdaLN dispatch policy and observations at the public resolver seam.""" import os +from argparse import Namespace import pytest import torch -from argparse import Namespace - -from rl_engine.validation.operators import CandidateSpec, run_operator_suite -from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case -from rl_engine.validation.operators.operator_specs import GtestAdaLNReference +from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp from rl_engine.contracts.numerical import ( BackendProvenance, ContractResolveError, @@ -19,10 +16,14 @@ resolve_tolerance, validate_backend_provenance, ) -from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp -from rl_engine.runtime.registry import OpBackend -from rl_engine.runtime.registry import KernelRegistry, kernel_registry +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 diff --git a/tests/runtime/test_adaln_indexed_triton.py b/tests/runtime/test_adaln_indexed_triton.py index 49d008bd4..73bdea9d3 100644 --- a/tests/runtime/test_adaln_indexed_triton.py +++ b/tests/runtime/test_adaln_indexed_triton.py @@ -64,7 +64,8 @@ def test_indexed_missing_extension_real_resolution(tmp_path, triton_state, missi 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" + """ + 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 diff --git a/tests/validation/operators/test_adaln_gtest.py b/tests/validation/operators/test_adaln_gtest.py index 0fcfa5af7..a5882e25d 100644 --- a/tests/validation/operators/test_adaln_gtest.py +++ b/tests/validation/operators/test_adaln_gtest.py @@ -9,6 +9,9 @@ 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 ( @@ -16,9 +19,6 @@ make_candidate, make_operator_case, ) -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 def _args(hidden=7): diff --git a/tests/validation/operators/test_op_checks.py b/tests/validation/operators/test_op_checks.py index b72e82959..311a7042d 100644 --- a/tests/validation/operators/test_op_checks.py +++ b/tests/validation/operators/test_op_checks.py @@ -125,22 +125,28 @@ def __call__(self, x): 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, + 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, + "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}, + {"sample": 1.0}, + {"sample": 3.0}, ] assert all("backend_provenance" not in case for case in result["cases"]) @@ -150,13 +156,19 @@ def __call__(self, x): 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, + 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, + "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) From 05f82211916ab828e350538289265d8de589e8ed Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 28 Oct 2026 16:40:23 +0800 Subject: [PATCH 36/37] fix(adaln): guard Triton devices and clarify validation backends Preserve the target device stream for shared and indexed Triton forward launches. Correct CUDA test prerequisites while retaining required-extension failures and CPU/HIP metadata coverage. Report observed backend and fallback per CLI case. Validation: 70 passed, 42 skipped with GPUs hidden; CI-pinned Black 24.4.2 and isort 5.13.2 checks pass. Real two-GPU and new H100 acceptance were not run. Signed-off-by: hongleng-github --- .../shared/triton/norm/adaln_indexed.py | 45 +++---- .../shared/triton/norm/adaln_modulation.py | 29 ++--- tests/ops/norm/test_adaln_modulation.py | 19 +-- tests/ops/norm/test_adaln_triton_device.py | 66 +++++++++++ tests/runtime/test_adaln_dispatch.py | 5 +- tests/runtime/test_adaln_strict_policy.py | 42 ++++++- tests/runtime/test_adaln_test_environments.py | 97 +++++++++++++++ .../operators/test_check_operator_cli.py | 112 ++++++++++++++++++ tools/validation/operators/check_operator.py | 12 +- 9 files changed, 370 insertions(+), 57 deletions(-) create mode 100644 tests/ops/norm/test_adaln_triton_device.py create mode 100644 tests/runtime/test_adaln_test_environments.py create mode 100644 tests/validation/operators/test_check_operator_cli.py diff --git a/rl_engine/backends/shared/triton/norm/adaln_indexed.py b/rl_engine/backends/shared/triton/norm/adaln_indexed.py index 66484d0ed..5d54c4a09 100644 --- a/rl_engine/backends/shared/triton/norm/adaln_indexed.py +++ b/rl_engine/backends/shared/triton/norm/adaln_indexed.py @@ -112,28 +112,29 @@ def forward(ctx, x, m, index, eps, num_warps, tile): b, s, h = x.shape y, gate = torch.empty_like(x), torch.empty_like(x) block = triton.next_power_of_2(h) - _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, - ) + 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 diff --git a/rl_engine/backends/shared/triton/norm/adaln_modulation.py b/rl_engine/backends/shared/triton/norm/adaln_modulation.py index 8381b8661..089a70fb1 100644 --- a/rl_engine/backends/shared/triton/norm/adaln_modulation.py +++ b/rl_engine/backends/shared/triton/norm/adaln_modulation.py @@ -115,20 +115,21 @@ 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) - _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, - ) + 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 diff --git a/tests/ops/norm/test_adaln_modulation.py b/tests/ops/norm/test_adaln_modulation.py index 70c18945c..b2de38c46 100644 --- a/tests/ops/norm/test_adaln_modulation.py +++ b/tests/ops/norm/test_adaln_modulation.py @@ -121,9 +121,7 @@ def test_adaln_modulation_accepts_qwen_six_chunk_view(): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_triton_adaln_modulation_forward_backward_and_batch_position(dtype): - try: - torch.cuda.init() - except RuntimeError: + if not torch.cuda.is_available(): pytest.skip("CUDA required") from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp @@ -161,9 +159,7 @@ def test_triton_adaln_modulation_forward_backward_and_batch_position(dtype): def test_triton_irregular_hidden_and_bf16_tie_rounding(): - try: - torch.cuda.init() - except RuntimeError: + if not torch.cuda.is_available(): pytest.skip("CUDA required") from rl_engine.backends.shared.triton.norm.adaln_modulation import TritonAdaLNModulationOp @@ -188,9 +184,7 @@ def test_triton_irregular_hidden_and_bf16_tie_rounding(): @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) def test_adaln_gpu_padding_and_batch_size_do_not_change_valid_bytes(dtype): - try: - torch.cuda.init() - except RuntimeError: + 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 @@ -237,12 +231,13 @@ def test_cuda_adaln_bit_equality_harness(hidden, dtype): from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp try: - torch.cuda.init() 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) @@ -268,9 +263,7 @@ def test_cuda_adaln_bit_equality_harness(hidden, dtype): reason="set RLK_ADALN_REAL_SHAPES=1 for Qwen-Image resolution coverage", ) def test_adaln_real_image_shapes_forward_backward(seq, dtype): - try: - torch.cuda.init() - except RuntimeError: + 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": 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/runtime/test_adaln_dispatch.py b/tests/runtime/test_adaln_dispatch.py index e69b7d964..5c7d6a400 100644 --- a/tests/runtime/test_adaln_dispatch.py +++ b/tests/runtime/test_adaln_dispatch.py @@ -79,10 +79,7 @@ def test_adaln_modulate_index_uses_explicit_native_fallback(): def test_adaln_cuda_registry_reports_real_fallback(): - try: - torch.cuda.init() - except RuntimeError: - pytest.skip("CUDA required") + 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": diff --git a/tests/runtime/test_adaln_strict_policy.py b/tests/runtime/test_adaln_strict_policy.py index ec0d54717..f9e430c10 100644 --- a/tests/runtime/test_adaln_strict_policy.py +++ b/tests/runtime/test_adaln_strict_policy.py @@ -109,7 +109,26 @@ def test_default_indexed_strict_rejection_clears_previous_shared_success(): assert trace["profiler_observed"] is False -def test_explicit_indexed_activation_resolves_without_enabling_cpu_select01(): +@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( @@ -162,7 +181,9 @@ def test_indexed_resolution_then_shared_execution_reports_actual_shared_contract "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(eps, reason): +def test_activated_strict_rechecks_actual_domain_before_execution( + indexed_cuda_symbols, eps, reason +): from rl_engine.runtime.adaln import IndexedAdaLNQualification scope = IndexedAdaLNQualification( @@ -221,9 +242,21 @@ def test_resolver_rejects_illegal_route_metadata_before_policy(index, error, str @pytest.mark.parametrize("mismatch", ["backend", "torch", "cuda", "device", "cpu", "hidden"]) -def test_activation_scope_does_not_qualify_another_backend_or_runtime(mismatch): +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", @@ -233,6 +266,9 @@ def test_activation_scope_does_not_qualify_another_backend_or_runtime(mismatch): 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", diff --git a/tests/runtime/test_adaln_test_environments.py b/tests/runtime/test_adaln_test_environments.py new file mode 100644 index 000000000..8e0d49391 --- /dev/null +++ b/tests/runtime/test_adaln_test_environments.py @@ -0,0 +1,97 @@ +# 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 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..2b6cda509 --- /dev/null +++ b/tests/validation/operators/test_check_operator_cli.py @@ -0,0 +1,112 @@ +# 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): + 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/tools/validation/operators/check_operator.py b/tools/validation/operators/check_operator.py index d6ae1e2fa..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( From 496324dbd1c85275f265f6ec2a8c39aeb5169974 Mon Sep 17 00:00:00 2001 From: hongleng-github Date: Wed, 28 Oct 2026 23:43:38 +0800 Subject: [PATCH 37/37] feat(tests): add CLI tests for optional CUDA extension handling Signed-off-by: hongleng-github --- tests/runtime/test_adaln_test_environments.py | 16 ++++++++++++++++ .../operators/test_check_operator_cli.py | 9 +++++++++ 2 files changed, 25 insertions(+) diff --git a/tests/runtime/test_adaln_test_environments.py b/tests/runtime/test_adaln_test_environments.py index 8e0d49391..86306de16 100644 --- a/tests/runtime/test_adaln_test_environments.py +++ b/tests/runtime/test_adaln_test_environments.py @@ -95,3 +95,19 @@ def test_indexed_qualification_tests_with_cpu_wheel_identity(hip): ) 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/validation/operators/test_check_operator_cli.py b/tests/validation/operators/test_check_operator_cli.py index 2b6cda509..427eead4d 100644 --- a/tests/validation/operators/test_check_operator_cli.py +++ b/tests/validation/operators/test_check_operator_cli.py @@ -24,6 +24,15 @@ "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",