feat(adaln): add indexed CUDA/Triton modulation and explicit strict qualification

benchmarks/operators/norm/benchmark_adaln_modulation.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+"""Compare AdaLN backends on Qwen-Image image-token lengths.
+
+Run with --real for S=4096/6889/6032, H=3072. These S values are derived
+from 8x VAE downsampling and 2x2 latent packing, not specified by #386.
+"""
+
+import argparse
+import contextlib
+import hashlib
+import json
+import platform
+import statistics
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+
+def layer_norm_baseline(x, modulation, *, eps=1e-6):
+    """Ordinary PyTorch math for accuracy/performance, without a replay claim."""
+    shift, scale, _ = modulation.float().chunk(3, dim=-1)
+    norm = torch.nn.functional.layer_norm(x.float(), (x.shape[-1],), eps=eps)
+    return (
+        (norm * (1.0 + scale[:, None]) + shift[:, None]).to(x.dtype),
+        modulation[:, None, 2 * x.shape[-1] :],
+    )
+
+
+def _bench(op, x, modulation, *, backward, warmup, repeat):
+    def run():
+        x.grad = modulation.grad = None
+        if backward:
+            y, gate = op(x, modulation)
+            torch.autograd.backward((y, gate), (torch.ones_like(y), torch.ones_like(gate)))
+        else:
+            with torch.no_grad():
+                op(x, modulation)
+
+    for _ in range(warmup):
+        run()
+    torch.cuda.synchronize()
+    torch.cuda.reset_peak_memory_stats()
+    allocated_before = torch.cuda.memory_allocated()
+    samples = []
+    for _ in range(repeat):
+        start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+        start.record()
+        run()
+        end.record()
+        end.synchronize()
+        samples.append(start.elapsed_time(end))
+    return {
+        "median_ms": statistics.median(samples),
+        "samples_ms": samples,
+        "peak_allocated_bytes": torch.cuda.max_memory_allocated(),
+        "allocated_before_samples_bytes": allocated_before,
+    }
+
+
+def parse_args(argv=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--real", action="store_true", help="run all three Qwen-Image resolutions")
+    parser.add_argument("--warmup", type=int, default=3)
+    parser.add_argument("--repeat", type=int, default=10)
+    parser.add_argument("--dtype", choices=("bf16", "fp32"), default="bf16")
+    parser.add_argument("--seed", type=int, default=386)
+    parser.add_argument("--json", action="store_true", help="emit one JSON report on stdout")
+    args = parser.parse_args(argv)
+    if args.warmup < 0 or args.repeat < 1:
+        parser.error("warmup must be nonnegative and repeat must be positive")
+    return args
+
+
+def _command_output(command):
+    try:
+        return subprocess.check_output(
+            command, cwd=REPO_ROOT, text=True, stderr=subprocess.DEVNULL, timeout=5
+        ).strip()
+    except (OSError, subprocess.SubprocessError):
+        return None
+
+
+def _sha256(path):
+    if not path or not Path(path).is_file():
+        return None
+    with open(path, "rb") as source:
+        digest = hashlib.sha256()
+        for chunk in iter(lambda: source.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def main(argv=None):
+    args = parse_args(argv)
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA device required")
+    # Imports and extension diagnostics must not corrupt the JSON stdout stream.
+    with contextlib.redirect_stdout(sys.stderr):
+        from rl_engine.reference.norm.adaln_modulation import NativeAdaLNModulationOp
+        from rl_engine.runtime.adaln import SHARED_PROFILE, AdaLNDispatch
+
+        backends = [("pytorch", NativeAdaLNModulationOp())]
+        skipped = []
+        try:
+            from rl_engine.backends.shared.triton.norm.adaln_modulation import (
+                TritonAdaLNModulationOp,
+            )
+
+            backends.append(("triton", TritonAdaLNModulationOp()))
+        except ImportError as exc:
+            skipped.append({"backend": "triton", "reason": str(exc)})
+        try:
+            from rl_engine.backends.cuda.norm.adaln_modulation import CudaAdaLNModulationOp
+
+            backends.append(("cuda", CudaAdaLNModulationOp()))
+        except RuntimeError as exc:
+            skipped.append({"backend": "cuda", "reason": str(exc)})
+        try:
+            import triton
+
+            triton_version = triton.__version__
+        except ImportError:
+            triton_version = None
+        environment = {
+            "gpu": torch.cuda.get_device_name(),
+            "device": str(torch.cuda.current_device()),
+            "capability": list(torch.cuda.get_device_capability()),
+            "cuda_runtime": torch.version.cuda,
+            "torch": torch.__version__,
+            "python": platform.python_version(),
+            "triton": triton_version,
+            "driver_inventory": _command_output(
+                ["nvidia-smi", "--query-gpu=uuid,name,driver_version", "--format=csv,noheader"]
+            ),
+            "source_revision": _command_output(["git", "rev-parse", "HEAD"]),
+            "source_status": _command_output(["git", "status", "--porcelain=v1"]),
+        }
+        source_paths = [
+            "benchmarks/operators/norm/benchmark_adaln_modulation.py",
+            "rl_engine/runtime/adaln.py",
+            "rl_engine/reference/norm/adaln_modulation.py",
+            "rl_engine/backends/cuda/norm/adaln_modulation.py",
+            "rl_engine/backends/shared/triton/norm/adaln_modulation.py",
+            "rl_engine/ops/autograd/vjp_fp32.py",
+            "csrc/cuda/norm/adaln_modulation.cu",
+            "csrc/bindings/ops.cpp",
+            "setup.py",
+            "build_tools/extensions.py",
+        ]
+        environment["source_sha256"] = {path: _sha256(REPO_ROOT / path) for path in source_paths}
+        backends = [
+            (
+                name,
+                AdaLNDispatch(
+                    op,
+                    {
+                        "requested_backend": name,
+                        "fallback": False,
+                        "fallback_reason": None,
+                        "arithmetic_profile": SHARED_PROFILE,
+                    },
+                    hidden=3072,
+                ),
+            )
+            for name, op in backends
+        ]
+        backends.append(("layer_norm_baseline", layer_norm_baseline))
+        dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32
+        generator = torch.Generator(device="cuda").manual_seed(args.seed)
+        results = []
+        sizes = (4096, 6889, 6032) if args.real else (64,)
+        for seq in sizes:
+            x = torch.randn(
+                1, seq, 3072, device="cuda", dtype=dtype, generator=generator, requires_grad=True
+            )
+            modulation = torch.randn(
+                1, 9216, device="cuda", dtype=dtype, generator=generator, requires_grad=True
+            )
+            for name, op in backends:
+                # Observe outside timing; retain the original bare-operator event interval.
+                with torch.no_grad():
+                    op(x, modulation)
+                trace = (
+                    dict(op.trace)
+                    if isinstance(op, AdaLNDispatch)
+                    else {
+                        "actual_backend": "pytorch",
+                        "arithmetic_profile": "ordinary-layer-norm",
+                        "observation": "successful_forward",
+                        "strict": False,
+                        "profiler_observed": False,
+                        "fallback": False,
+                    }
+                )
+                trace["binary_sha256"] = _sha256(trace.get("binary_path"))
+                timed_op = op.implementation if isinstance(op, AdaLNDispatch) else op
+                for backward in (False, True):
+                    measured = _bench(
+                        timed_op,
+                        x,
+                        modulation,
+                        backward=backward,
+                        warmup=args.warmup,
+                        repeat=args.repeat,
+                    )
+                    results.append(
+                        {
+                            "backend": name,
+                            "shape": list(x.shape),
+                            "backward": backward,
+                            "trace": trace,
+                            **measured,
+                        }
+                    )
+    report = {
+        "schema_version": "adaln-benchmark-v1",
+        "environment": environment,
+        "dtype": args.dtype,
+        "seed": args.seed,
+        "warmup": args.warmup,
+        "repeat": args.repeat,
+        "timing": "CUDA events; forward or autograd forward+backward including ones_like upstreams",
+        "skipped_backends": skipped,
+        "results": results,
+    }
+    if args.json:
+        print(json.dumps(report, indent=2))
+    else:
+        print(json.dumps({key: value for key, value in report.items() if key != "results"}))
+        for result in results:
+            print(
+                f"S={result['shape'][1]} backend={result['backend']} "
+                f"backward={result['backward']} median_ms={result['median_ms']:.3f} "
+                f"peak_MB={result['peak_allocated_bytes'] / 2**20:.1f}"
+            )
+
+
+if __name__ == "__main__":
+    main()

build_tools/extensions.py

@@ -223,6 +223,7 @@ def get_extensions():
             "csrc/cuda/activation/activation.cu",
             "csrc/cuda/attention/deterministic_attention.cu",
         ]
+        adaln_cuda_enabled = not is_rocm and not envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH)
         if is_rocm:
             # ROCm-tuned WS2 vocab-parallel logprob kernels; the shared
             # deterministic_logp_kernel.cu keeps the SM90-tuned CUDA path.
@@ -235,6 +236,8 @@ def get_extensions():
         else:
             # CUDA IPC and the fixed-tree collective implementation are not
             # part of the ROCm extension.
+            if adaln_cuda_enabled:
+                cuda_sources.append("csrc/cuda/norm/adaln_modulation.cu")
             cuda_sources.append("csrc/cuda/collectives/deterministic_collective.cu")
             # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync).
             # The ROCm dispatcher falls back to PyTorch SDPA for this operator.
@@ -324,6 +327,8 @@ def get_extensions():
 
         platform_define = "-DKERNEL_ALIGN_WITH_ROCM" if is_rocm else "-DKERNEL_ALIGN_WITH_CUDA"
         cxx_flags = ["-O3", "-std=c++17", platform_define]
+        if adaln_cuda_enabled:
+            cxx_flags.append("-DRLK_ADALN_CUDA_ENABLED")
         extra_link_args = list(torch_rpath)
         if os.name != "nt" and not is_rocm:
             # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute).

ci/scripts/run_ws1_gtest.sh

@@ -23,8 +23,20 @@ echo "[ws1-gtest] interpreter=$PY out=$OUT"
   tests/validation/operators/test_four_judgment_matrix.py \
   tests/validation/operators/test_ws1_candidate_evidence.py \
   tests/validation/operators/test_op_checks.py \
+  tests/validation/operators/test_adaln_gtest.py \
+  tests/backends/test_adaln_shared_bounds.py \
+  tests/runtime/test_adaln_dispatch.py \
+  tests/ops/norm/test_adaln_modulation.py \
   tests/validation/operators/test_elementwise_inventory.py
 
+echo "[ws1-gtest] AdaLN shared outputs and gradients against independent CPU gold"
+for backend in cuda triton; do
+  for dtype in fp32 bf16; do
+    "$PY" tools/validation/operators/check_operator.py --op adaln_modulation --candidate "$backend" \
+      --device cuda --dtype "$dtype" --batch 2 --seq 5 --normalized-dim 3072 --check-grad
+  done
+done
+
 echo "[ws1-gtest] C3/C4 CUDA + Triton smoke (silu)"
 "$PY" tools/validation/operators/check_forward_invariance.py \
   --op silu --candidate cuda --backend-profile cuda_bf16

csrc/bindings/ops.cpp

@@ -256,6 +256,19 @@ std::vector<torch::Tensor> swiglu_packed_backward_cuda(
     torch::Tensor dy,
     torch::Tensor gate_up);
 
+#if defined(RLK_ADALN_CUDA_ENABLED)
+std::vector<torch::Tensor> adaln_modulation_forward_cuda(
+    torch::Tensor x, torch::Tensor modulation, double eps, int threads);
+std::vector<torch::Tensor> adaln_modulation_backward_cuda(
+    torch::Tensor dy, torch::Tensor dg, torch::Tensor x,
+    torch::Tensor modulation, double eps, int threads);
+std::vector<torch::Tensor> adaln_indexed_forward_cuda(
+    torch::Tensor x, torch::Tensor modulation, torch::Tensor index, double eps, int threads);
+std::vector<torch::Tensor> adaln_indexed_backward_cuda(
+    c10::optional<torch::Tensor> dy, c10::optional<torch::Tensor> dg,
+    torch::Tensor x, torch::Tensor modulation, torch::Tensor index, double eps, int threads);
+#endif
+
 // RMSNorm Declarations & Wrappers
 
 void rmsnorm_forward_cuda(
@@ -737,6 +750,23 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
           "Batch-invariant SwiGLU forward for [rows, 2 * intermediate]");
     m.def("swiglu_packed_backward", &swiglu_packed_backward,
           "Batch-invariant SwiGLU backward for [rows, 2 * intermediate]");
+#if defined(RLK_ADALN_CUDA_ENABLED)
+    m.def("adaln_modulation_forward", &adaln_modulation_forward_cuda,
+          "AdaLN modulation forward CUDA",
+          py::arg("x"), py::arg("modulation"), py::arg("eps"), py::arg("threads") = 256);
+    m.def("adaln_modulation_backward", &adaln_modulation_backward_cuda,
+          "AdaLN modulation backward CUDA",
+          py::arg("dy"), py::arg("dg"), py::arg("x"), py::arg("modulation"),
+          py::arg("eps"), py::arg("threads") = 256);
+    m.def("adaln_indexed_forward", &adaln_indexed_forward_cuda,
+          "Indexed AdaLN small FP32 CUDA slice",
+          py::arg("x"), py::arg("modulation"), py::arg("index"), py::arg("eps"),
+          py::arg("threads") = 256);
+    m.def("adaln_indexed_backward", &adaln_indexed_backward_cuda,
+          "Indexed AdaLN FP32 ascending-token VJP",
+          py::arg("dy"), py::arg("dg"), py::arg("x"), py::arg("modulation"),
+          py::arg("index"), py::arg("eps"), py::arg("threads") = 256);
+#endif
 
     // Deterministic standard-softmax attention (issue #147)
     m.def(