From c88f7c0fe35d6544e092a16097a094d323b35581 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Sun, 11 Oct 2026 00:26:15 +0000 Subject: [PATCH 1/6] feat(qwenimage): add CUDA AdaLN modulation projection Add the Qwen-Image Phase A `adaln_mod_projection` operator (RFC #386): per-stream SiLU + Linear 3072->18432 with a fixed-order FP32 forward, compensated input gradients and ascending FP32 Neumaier parameter gradients, plus an independent CPU FP32 reference. Sources live under csrc/cuda/model_specific/qwenimage, rl_engine/backends/cuda/model_specific/qwenimage and rl_engine/reference/model_specific/qwenimage. Register the CUDA and CPU slots (other platforms fail closed), the native bindings, the WS1 operator spec/inputs and the build source, and add tests, evidence runners under benchmarks/models/qwenimage, and the operator page. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- benchmarks/models/qwenimage/__init__.py | 2 + benchmarks/models/qwenimage/adaln_evidence.py | 286 ++++++++ .../qwenimage/adaln_parameter_screen.py | 43 ++ .../models/qwenimage/adaln_prior_art.py | 111 +++ .../models/qwenimage/adaln_reuse_bi_audit.py | 439 ++++++++++++ .../models/qwenimage/adaln_reuse_candidate.py | 40 ++ .../models/qwenimage/plot_adaln_audit.py | 248 +++++++ .../models/qwenimage/plot_adaln_evidence.py | 49 ++ benchmarks/models/qwenimage/streams.py | 8 + .../models/qwenimage/validate_adaln_repair.py | 86 +++ build_tools/extensions.py | 5 + csrc/bindings/ops.cpp | 17 + .../model_specific/qwenimage/det_linear.cu | 663 ++++++++++++++++++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + .../qwenimage-adaln-mod-projection.md | 74 ++ rl_engine/_C.pyi | 22 + .../backends/cuda/model_specific/__init__.py | 2 + .../cuda/model_specific/qwenimage/__init__.py | 2 + .../qwenimage/adaln_mod_projection.py | 116 +++ .../reference/model_specific/__init__.py | 2 + .../model_specific/qwenimage/__init__.py | 2 + .../qwenimage/adaln_mod_projection.py | 90 +++ rl_engine/runtime/registry.py | 18 + .../validation/operators/operator_inputs.py | 16 + .../validation/operators/operator_specs.py | 16 + tests/models/qwenimage/__init__.py | 1 + tests/models/qwenimage/test_adaln.py | 209 ++++++ .../qwenimage/test_parameter_reduction.py | 44 ++ tests/models/qwenimage/test_reference.py | 35 + 30 files changed, 2648 insertions(+) create mode 100644 benchmarks/models/qwenimage/__init__.py create mode 100644 benchmarks/models/qwenimage/adaln_evidence.py create mode 100644 benchmarks/models/qwenimage/adaln_parameter_screen.py create mode 100644 benchmarks/models/qwenimage/adaln_prior_art.py create mode 100644 benchmarks/models/qwenimage/adaln_reuse_bi_audit.py create mode 100644 benchmarks/models/qwenimage/adaln_reuse_candidate.py create mode 100644 benchmarks/models/qwenimage/plot_adaln_audit.py create mode 100644 benchmarks/models/qwenimage/plot_adaln_evidence.py create mode 100644 benchmarks/models/qwenimage/streams.py create mode 100644 benchmarks/models/qwenimage/validate_adaln_repair.py create mode 100644 csrc/cuda/model_specific/qwenimage/det_linear.cu create mode 100644 docs/operators/qwenimage-adaln-mod-projection.md create mode 100644 rl_engine/backends/cuda/model_specific/__init__.py create mode 100644 rl_engine/backends/cuda/model_specific/qwenimage/__init__.py create mode 100644 rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py create mode 100644 rl_engine/reference/model_specific/__init__.py create mode 100644 rl_engine/reference/model_specific/qwenimage/__init__.py create mode 100644 rl_engine/reference/model_specific/qwenimage/adaln_mod_projection.py create mode 100644 tests/models/qwenimage/__init__.py create mode 100644 tests/models/qwenimage/test_adaln.py create mode 100644 tests/models/qwenimage/test_parameter_reduction.py create mode 100644 tests/models/qwenimage/test_reference.py diff --git a/benchmarks/models/qwenimage/__init__.py b/benchmarks/models/qwenimage/__init__.py new file mode 100644 index 000000000..cdd927da2 --- /dev/null +++ b/benchmarks/models/qwenimage/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Model benchmark extension point.""" diff --git a/benchmarks/models/qwenimage/adaln_evidence.py b/benchmarks/models/qwenimage/adaln_evidence.py new file mode 100644 index 000000000..b1143363b --- /dev/null +++ b/benchmarks/models/qwenimage/adaln_evidence.py @@ -0,0 +1,286 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Real-weight accuracy, deterministic replay, and alternating CUDA measurements. + +Run from the repository root with PYTHONPATH=. and a compiled local extension. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import subprocess +from pathlib import Path + +import torch +from safetensors import safe_open + +from benchmarks.models.qwenimage.streams import STREAMS +from rl_engine.backends.cuda.model_specific.qwenimage.adaln_mod_projection import ( + QwenImageAdaLNModProjectionCudaOp, +) +from rl_engine.backends.extension import _C +from rl_engine.contracts.numerical import load_contract +from rl_engine.reference.model_specific.qwenimage.adaln_mod_projection import ( + QwenImageAdaLNModProjectionReference, + reference_projection, +) + +CHECKPOINT_REVISION = "75e0b4be04f60ec59a75f475837eced720f823b6" +DIFFUSERS_REVISION = "d961a388fd02e4db38d17350c8dd9b8abe642e05" + + +def equal_bytes(left, right): + return ( + left.dtype == right.dtype + and left.shape == right.shape + and torch.equal(left.contiguous().view(torch.uint8), right.contiguous().view(torch.uint8)) + ) + + +def tensor_hash(tensor): + return hashlib.sha256(tensor.contiguous().view(torch.uint8).numpy().tobytes()).hexdigest() + + +QWENIMAGE_SOURCE_ROOTS = ( + "csrc/cuda/model_specific/qwenimage", + "rl_engine/backends/cuda/model_specific/qwenimage", + "rl_engine/reference/model_specific/qwenimage", + "benchmarks/models/qwenimage", +) + + +def source_state(): + base = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() + diff = subprocess.check_output(["git", "diff", "--binary", "HEAD"]) + untracked = subprocess.check_output( + ["git", "ls-files", "--others", "--exclude-standard", "-z"] + ).split(b"\0") + hashes = {} + for name in untracked: + if not name: + continue + path = Path(name.decode()) + if path.suffix in (".py", ".cu", ".cpp", ".h", ".pyi"): + hashes[str(path)] = hashlib.sha256(path.read_bytes()).hexdigest() + fingerprint = hashlib.sha256(diff + json.dumps(hashes, sort_keys=True).encode()).hexdigest() + status = subprocess.check_output(["git", "status", "--porcelain", "--untracked-files=all"]) + sources = { + str(path): hashlib.sha256(path.read_bytes()).hexdigest() + for root in QWENIMAGE_SOURCE_ROOTS + for path in sorted(Path(root).rglob("*")) + if path.suffix in (".py", ".cu") and "__pycache__" not in path.parts + } + return { + "base_sha": base, + "tracked_diff_sha256": hashlib.sha256(diff).hexdigest(), + "untracked_source_sha256": hashes, + "workspace_source_sha256": fingerprint, + "qwenimage_source_sha256": sources, + "extension_sha256": hashlib.sha256(Path(_C.__file__).read_bytes()).hexdigest(), + "clean_commit": not status.strip(), + } + + +def load_weights(model, stem): + index = json.loads( + (model / "transformer/diffusion_pytorch_model.safetensors.index.json").read_text() + ) + tensors, manifest = {}, {} + for suffix in ("weight", "bias"): + name = f"{stem}.{suffix}" + shard = index["weight_map"][name] + metadata = model / ".cache/huggingface/download/transformer" / (shard + ".metadata") + revision, etag, *_ = metadata.read_text().splitlines() + if revision != CHECKPOINT_REVISION: + raise ValueError(f"unverified local checkpoint revision: {revision}") + with safe_open(model / "transformer" / shard, framework="pt", device="cpu") as handle: + tensor = handle.get_tensor(name) + tensors[suffix] = tensor + manifest[name] = { + "shape": list(tensor.shape), + "dtype": str(tensor.dtype), + "sha256": tensor_hash(tensor), + "shard": shard, + "download_etag": etag, + } + return tensors, manifest + + +def compare(actual, expected, dtype, judgment): + spec = load_contract()["judgments"][judgment]["by_op_class"]["reduction"][ + str(dtype).split(".")[-1] + ] + a, b = actual.detach().float().cpu(), expected.detach().float().cpu() + error = (a - b).abs() + limit = spec["atol"] + spec["rtol"] * b.abs() + return { + "passed": bool( + torch.isfinite(a).all() and torch.isfinite(b).all() and (error <= limit).all() + ), + "max_abs": error.max().item(), + "max_tolerance_ratio": (error / limit).max().item(), + "violations": (error > limit).sum().item(), + "atol": spec["atol"], + "rtol": spec["rtol"], + } + + +def correctness(candidate, args, kind): + dtype = args["temb"].dtype + reference = reference_projection + cpu = {k: v.detach().cpu().float().requires_grad_() for k, v in args.items()} + y = candidate(**args) + gold = reference(**cpu, execution_dtype=dtype) + gen = torch.Generator(device=y.device).manual_seed(127) + grad = torch.randn(y.shape, device=y.device, dtype=y.dtype, generator=gen) + grads = torch.autograd.grad(y, tuple(args.values()), grad) + gold_grads = torch.autograd.grad(gold, tuple(cpu.values()), grad.float().cpu()) + metrics = { + "forward": compare(y, gold, dtype, "forward_accuracy"), + "gradients": { + k: compare(a, b, dtype, "gradient_accuracy") for k, a, b in zip(args, grads, gold_grads) + }, + } + again = candidate(**args) + repeat = torch.autograd.grad(again, tuple(args.values()), grad) + metrics["repeat_bitwise"] = bool( + equal_bytes(y, again) and all(equal_bytes(a, b) for a, b in zip(grads, repeat)) + ) + del gold, gold_grads, cpu + return metrics + + +def measure(function, args, mode): + values = tuple(args.values()) + if mode == "backward": + output = function(**args) + upstream = torch.ones_like(output) + torch.cuda.synchronize() + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + if mode == "forward": + with torch.no_grad(): + output = function(**args) + elif mode == "backward": + torch.autograd.grad(output, values, upstream) + else: + output = function(**args) + torch.autograd.grad(output, values, torch.ones_like(output)) + end.record() + end.synchronize() + return start.elapsed_time(end) + + +def benchmark(functions, args, repeats): + result = {} + for mode in ("forward", "backward", "forward_backward"): + for fn in functions.values(): + for _ in range(3): + measure(fn, args, mode) + samples = {name: [] for name in functions} + orders = [] + for iteration in range(repeats): + order = list(functions) + if iteration % 2: + order.reverse() + orders.append(order) + for name in order: + samples[name].append(measure(functions[name], args, mode)) + result[mode] = { + "samples_ms": samples, + "order": orders, + "median_ms": {k: statistics.median(v) for k, v in samples.items()}, + } + result["peak_increment_bytes"] = {} + for name, fn in functions.items(): + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + before = torch.cuda.memory_allocated() + measure(fn, args, "forward_backward") + result["peak_increment_bytes"][name] = torch.cuda.max_memory_allocated() - before + return result + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--repeats", type=int, default=10) + opts = parser.parse_args() + torch.set_num_threads(8) + torch.manual_seed(71) + report = { + "checkpoint": "Qwen/Qwen-Image", + "checkpoint_revision": CHECKPOINT_REVISION, + "diffusers_revision": DIFFUSERS_REVISION, + "diffusers_source_sha256": { + "transformer_qwenimage.py": ( + "0bad11901deae69f05aec01fd464e51e38582d455068f2851ee665036179321f" + ), + "normalization.py": "a3c40cee9308a0e339cca21c14bf51f394c61ae68e5b9a578e5636441f5eaa84", + }, + "source": source_state(), + "environment": { + "torch": torch.__version__, + "cuda": torch.version.cuda, + "driver": subprocess.check_output( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], text=True + ).splitlines()[0], + "gpu": torch.cuda.get_device_name(), + "tf32": torch.backends.cuda.matmul.allow_tf32, + }, + "weights": {}, + "cases": [], + } + for label, stem in STREAMS: + weights, manifest = load_weights(opts.model, stem) + report["weights"].update(manifest) + kind = "projection" + candidate = QwenImageAdaLNModProjectionCudaOp() + native = QwenImageAdaLNModProjectionReference() + for dtype in (torch.bfloat16, torch.float32): + for batch, seq in [(1, 1), (2, 1), (4, 1), (9, 1)]: + args = { + "temb": torch.randn(batch, 3072, device="cuda", dtype=dtype).requires_grad_(), + **{ + k: v.to(device="cuda", dtype=dtype).requires_grad_() + for k, v in weights.items() + }, + } + entry = { + "stream": label, + "dtype": str(dtype), + "batch": batch, + "seq": seq, + "readback": candidate.implementation_info(**args), + "correctness": correctness(candidate, args, kind), + "timing": benchmark( + {"candidate": candidate, "native": native}, args, opts.repeats + ), + } + report["cases"].append(entry) + opts.output.parent.mkdir(parents=True, exist_ok=True) + opts.output.write_text(json.dumps(report, indent=2) + "\n") + print(label, str(dtype), batch, seq, entry["correctness"], flush=True) + del args + report["passed"] = all( + c["correctness"]["forward"]["passed"] + and c["correctness"]["repeat_bitwise"] + and all(v["passed"] for v in c["correctness"]["gradients"].values()) + and all( + v["passed"] + for group in c["correctness"].get("fp64_diagnostic", {}).values() + for v in group.values() + ) + for c in report["cases"] + ) + opts.output.write_text(json.dumps(report, indent=2) + "\n") + if not report["passed"]: + raise SystemExit("accuracy gate failed; see report") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/models/qwenimage/adaln_parameter_screen.py b/benchmarks/models/qwenimage/adaln_parameter_screen.py new file mode 100644 index 000000000..88d80add6 --- /dev/null +++ b/benchmarks/models/qwenimage/adaln_parameter_screen.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Check the fixed parameter reducer on cancellation cases that defeat naive sums.""" +import argparse +import json +from pathlib import Path + +import torch + +from benchmarks.models.qwenimage.adaln_evidence import compare, source_state +from rl_engine.backends.extension import _C + +p = argparse.ArgumentParser() +p.add_argument("--output", type=Path, required=True) +a = p.parse_args() +report = dict( + source=source_state(), + torch=torch.__version__, + gpu=torch.cuda.get_device_name(), + construction="grad[B,2]=1; grad[0]=2**26; grad[-1]=-2**26; act[B,7]=1; act[:,1::2]=1.5", + cases=[], +) +for dtype in [torch.float32, torch.bfloat16]: + for batch in [255, 256, 257, 2048]: + g = torch.ones(batch, 2, device="cuda", dtype=torch.float32) + g[0] = 2**26 + g[-1] = -(2**26) + x = torch.ones(batch, 7, device="cuda", dtype=dtype) + x[:, 1::2] = 1.5 + golden = g.double().cpu().T @ x.double().cpu() + baseline = _C.qwenimage_det_linear_backward_weight(g, x, dtype, True)[0] + report["cases"].append( + dict( + dtype=str(dtype), + batch=batch, + fp64=golden.tolist(), + fixed_actual=baseline.float().cpu().tolist(), + fixed_contract=compare(baseline, golden, dtype, "gradient_accuracy"), + ) + ) +a.output.parent.mkdir(parents=True, exist_ok=True) +a.output.write_text(json.dumps(report, indent=2) + "\n") +print("fixed failures", sum(not c["fixed_contract"]["passed"] for c in report["cases"])) diff --git a/benchmarks/models/qwenimage/adaln_prior_art.py b/benchmarks/models/qwenimage/adaln_prior_art.py new file mode 100644 index 000000000..2f95934f4 --- /dev/null +++ b/benchmarks/models/qwenimage/adaln_prior_art.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Compare native, reused fixed GEMV, and an optional existing QKV cuBLASLt build.""" + +import argparse +import hashlib +import importlib.util +import json +from pathlib import Path + +import torch +import torch.nn.functional as F + +from benchmarks.models.qwenimage.adaln_evidence import equal_bytes, load_weights, source_state +from benchmarks.models.qwenimage.streams import STREAMS +from rl_engine.backends.cuda.model_specific.qwenimage.adaln_mod_projection import projection_forward + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", type=Path, required=True) + parser.add_argument("--fast-extension", type=Path) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + torch.set_num_threads(8) + torch.manual_seed(29) + fast = None + if args.fast_extension: + spec = importlib.util.spec_from_file_location("qwenimage_prior_art._C", args.fast_extension) + fast = importlib.util.module_from_spec(spec) + spec.loader.exec_module(fast) + report = {"source": source_state(), "cases": [], "tf32": torch.backends.cuda.matmul.allow_tf32} + if args.fast_extension: + report["existing_qkv_extension_sha256"] = hashlib.sha256( + args.fast_extension.read_bytes() + ).hexdigest() + report["environment"] = { + "torch": torch.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(), + } + for label, stem in STREAMS: + weights, _ = load_weights(args.model, stem) + for dtype in (torch.bfloat16, torch.float32): + w, b = [weights[k].to(device="cuda", dtype=dtype) for k in ("weight", "bias")] + t = torch.randn(64, 3072, device="cuda", dtype=dtype) + funcs = { + "native": lambda z: F.linear(F.silu(z), w, b), + "fixed_gemv": lambda z: projection_forward(z, w, b)[0], + } + if fast: + funcs["existing_qkv_cublaslt"] = lambda z: ( + fast.qwenimage_fast_gemm_fp32_cuda(F.silu(z).contiguous(), w) + b.float() + ).to(dtype) + timings = {name: [] for name in funcs} + orders = [] + for fn in funcs.values(): + for _ in range(4): + fn(t[:1]) + for iteration in range(15): + order = list(funcs) + if iteration % 2: + order.reverse() + orders.append(order) + for name in order: + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event( + enable_timing=True + ) + start.record() + funcs[name](t[:1]) + end.record() + end.synchronize() + timings[name].append(start.elapsed_time(end)) + for name, fn in funcs.items(): + base = fn(t[:1]) + violations = [] + for batch in [1, 2, 3, 4, 7, 8, 9, 16, 33, 64]: + for pos in sorted(set([0, batch // 2, batch - 1])): + z = t[:batch].clone() + z[pos] = t[0] + y = fn(z)[pos : pos + 1] + if not equal_bytes(y, base): + violations.append( + { + "batch": batch, + "position": pos, + "differing_bytes": ( + y.contiguous().view(torch.uint8) + != base.contiguous().view(torch.uint8) + ) + .sum() + .item(), + } + ) + report["cases"].append( + { + "stream": label, + "dtype": str(dtype), + "implementation": name, + "batch_invariance_violations": violations, + "forward_b1_ms": timings[name], + "timing_order": orders, + } + ) + print(label, dtype, name, "violations", len(violations), flush=True) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/models/qwenimage/adaln_reuse_bi_audit.py b/benchmarks/models/qwenimage/adaln_reuse_bi_audit.py new file mode 100644 index 000000000..712e3772e --- /dev/null +++ b/benchmarks/models/qwenimage/adaln_reuse_bi_audit.py @@ -0,0 +1,439 @@ +#!/usr/bin/env python3 +# isort: skip_file +# SPDX-License-Identifier: Apache-2.0 +"""Full-op reuse screening, followed by bytewise sample-batch invariance. + +Run each mode in a fresh process. cublas_policy reproduces vLLM's two cuBLAS +workspace settings, not the complete vLLM integration. The vLLM, SGLang and +Megatron modes use each project's batch-invariant persistent matmul with the +same FP32 VJP adapter. No backend math flags are changed. FP64 comparisons use +the declared storage boundaries/STE. +""" +import argparse +import importlib.metadata +import json +import os +import subprocess +import sys +from pathlib import Path + +p = argparse.ArgumentParser(description=__doc__) +p.add_argument( + "--mode", + choices=[ + "native", + "cublas_policy", + "fixed", + "vllm_local", + "sglang_local", + "megatron_local", + ], + required=True, +) +p.add_argument( + "--phase", + choices=["reuse", "bi", "workload", "accuracy", "accuracy_large", "parameters"], + required=True, +) +p.add_argument("--model", type=Path, required=True) +p.add_argument("--output", type=Path, required=True) +p.add_argument("--seed", type=int, default=0) +p.add_argument("--capture-failures", action="store_true") +p.add_argument("--megatron-src", type=Path, help="Megatron-LM checkout for megatron_local") +a = p.parse_args() +if a.mode == "cublas_policy": + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" + os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1" + +if a.mode in ("vllm_local", "sglang_local", "megatron_local"): + os.environ["TRITON_F32_DEFAULT"] = "ieee" +if a.mode == "megatron_local": + assert a.megatron_src, "--megatron-src is required for megatron_local" + sys.path.insert(0, str(a.megatron_src)) +if a.mode == "sglang_local": + os.environ["SGLANG_ENABLE_JIT_DEEPGEMM"] = "0" + os.environ["SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_DEEPGEMM"] = "0" + +# Imports follow explicit subprocess environment setup. +import torch # noqa: E402 +from benchmarks.models.qwenimage.streams import STREAMS # noqa: E402 +from benchmarks.models.qwenimage.adaln_evidence import ( # noqa: E402 + load_weights, + source_state, + benchmark, + equal_bytes, + compare, +) +from rl_engine.backends.cuda.model_specific.qwenimage.adaln_mod_projection import ( # noqa: E402 + QwenImageAdaLNModProjectionCudaOp, +) +from rl_engine.reference.model_specific.qwenimage.adaln_mod_projection import ( # noqa: E402 + QwenImageAdaLNModProjectionReference, + reference_projection, +) + +torch.set_num_threads(8) +assert not torch.backends.cuda.matmul.allow_tf32, "TF32 must already be disabled" +report = dict( + mode=a.mode, + phase=a.phase, + accuracy_seed=a.seed, + source=source_state(), + cases=[], + environment=dict( + torch=torch.__version__, + cuda=torch.version.cuda, + gpu=torch.cuda.get_device_name(), + workspace={ + k: os.environ.get(k) for k in ["CUBLAS_WORKSPACE_CONFIG", "CUBLASLT_WORKSPACE_SIZE"] + }, + tf32=torch.backends.cuda.matmul.allow_tf32, + ), + scope=("sample batch BI; " "shared parameter gradients: same ordered inputs, repeat only"), +) +report["packages"] = {} +for lib in [ + "torch", + "triton", + "diffusers", + "vllm", + "sglang", + "transformer_engine", + "megatron-core", +]: + try: + report["packages"][lib] = importlib.metadata.version(lib) + except importlib.metadata.PackageNotFoundError: + report["packages"][lib] = None +if a.mode == "megatron_local": + report["packages"]["megatron-lm-source"] = subprocess.check_output( + ["git", "-C", str(a.megatron_src), "rev-parse", "HEAD"], text=True + ).strip() + + +def save(): + a.output.parent.mkdir(parents=True, exist_ok=True) + a.output.write_text(json.dumps(report, indent=2) + "\n") + + +def run(fn, inp, upstream, parameters=False): + names = list(inp) if parameters else [k for k in inp if k not in ("weight", "bias")] + vals = {k: v.detach().requires_grad_(k in names) for k, v in inp.items()} + y = fn(**vals) + gs = torch.autograd.grad(y, tuple(vals[k] for k in names), upstream) + return {"output": y.detach(), **{k: g.detach() for k, g in zip(names, gs)}} + + +def sliced(inp, idx): + return {k: (v if k in ("weight", "bias") else v[idx]) for k, v in inp.items()} + + +def stats(left, right): + return { + k: int( + ( + v.contiguous().view(torch.uint8).reshape(v.shape[0], -1) + != right[k].contiguous().view(torch.uint8).reshape(v.shape[0], -1) + ) + .any(1) + .sum() + ) + for k, v in left.items() + } + + +def metric(x, y): + x = x.double().cpu() + y = y.double().cpu() + return float((x - y).abs().max() / y.abs().max().clamp_min(1e-30)) + + +for stream, stem in STREAMS: + weights, manifest = load_weights(a.model, stem) + fn = ( + QwenImageAdaLNModProjectionCudaOp() + if a.mode == "fixed" + else QwenImageAdaLNModProjectionReference() + ) + if a.mode in ("vllm_local", "sglang_local", "megatron_local"): + from benchmarks.models.qwenimage.adaln_reuse_candidate import Projection + + if a.mode == "sglang_local": + import benchmarks.models.qwenimage.adaln_reuse_candidate as adapter + from sglang.srt.batch_invariant_ops.batch_invariant_ops import ( + matmul_persistent, + ) + + adapter.matmul_persistent = matmul_persistent + if a.mode == "megatron_local": + import benchmarks.models.qwenimage.adaln_reuse_candidate as adapter + from megatron.core.transformer.custom_layers.batch_invariant_kernels import ( + matmul_persistent, + ) + + adapter.matmul_persistent = matmul_persistent + fn = Projection() + for dtype in [torch.bfloat16, torch.float32]: + w, b = [weights[k].cuda().to(dtype) for k in ("weight", "bias")] + case = dict(stream=stream, dtype=str(dtype), weights=manifest, checks=[]) + report["cases"].append(case) + + def inputs(batch, seq=3, seed=0): + gen = torch.Generator(device="cuda").manual_seed(seed) + inp = dict( + temb=torch.randn(batch, 3072, device="cuda", dtype=dtype, generator=gen), + weight=w, + bias=b, + ) + shape = (batch, w.shape[0]) + upstream = torch.randn(shape, device="cuda", dtype=dtype, generator=gen) + return inp, upstream + + if a.phase == "parameters": + for seed in [0, 127, 2026]: + for batch in [1, 255, 256, 257, 2048]: + inp, g = inputs(batch, seed=seed) + left = run(fn, inp, g, True) + right = run(fn, inp, g, True) + case["checks"].append( + dict( + kind="parameter_repeat", + batch=batch, + seed=seed, + all_equal={ + k: equal_bytes(left[k], right[k]) for k in ("weight", "bias") + }, + ) + ) + inp, g = inputs(257, seed=seed) + g[255:] = 0 + base = run(fn, sliced(inp, slice(0, 255)), g[:255], True) + for batch in [256, 257]: + padded = run(fn, sliced(inp, slice(0, batch)), g[:batch], True) + case["checks"].append( + dict( + kind="ordered_contributions_zero_padding", + batch=batch, + seed=seed, + all_equal={ + k: equal_bytes(base[k], padded[k]) for k in ("weight", "bias") + }, + ) + ) + save() + print(a.mode, stream, dtype, "parameters complete", flush=True) + elif a.phase in ("accuracy", "accuracy_large"): + shapes = [(256, 1), (2048, 1)] if a.phase == "accuracy_large" else [(1, 1), (3, 1)] + for batch, seq in shapes: + inp, g = inputs(batch, seq, a.seed) + out = run(fn, inp, g, True) + hi = {k: v.double().cpu().requires_grad_() for k, v in inp.items()} + gold = reference_projection(**hi, execution_dtype=dtype) + gg = torch.autograd.grad(gold, tuple(hi.values()), g.double().cpu()) + expected = {"output": gold.detach(), **dict(zip(hi, gg))} + case["checks"].append( + dict( + kind="fp64_contract", + batch=batch, + sequence=seq, + errors={k: metric(v, expected[k]) for k, v in out.items()}, + contract={ + k: compare( + v, + expected[k], + dtype, + ("forward_accuracy" if k == "output" else "gradient_accuracy"), + ) + for k, v in out.items() + }, + ) + ) + low = {k: v.float().cpu().requires_grad_() for k, v in inp.items()} + low_y = reference_projection(**low, execution_dtype=dtype) + low_g = torch.autograd.grad(low_y, tuple(low.values()), g.float().cpu()) + low_expected = {"output": low_y.detach(), **dict(zip(low, low_g))} + case["checks"][-1]["cpu_fp32_contract"] = { + k: compare( + v, + low_expected[k], + dtype, + "forward_accuracy" if k == "output" else "gradient_accuracy", + ) + for k, v in out.items() + } + case["checks"][-1]["cpu_fp32_vs_fp64"] = { + k: compare( + v, + expected[k], + dtype, + "forward_accuracy" if k == "output" else "gradient_accuracy", + ) + for k, v in low_expected.items() + } + del low, low_y, low_g, low_expected + if a.capture_failures: + details = {} + for name, value in out.items(): + check = case["checks"][-1]["contract"][name] + actual = value.detach().float().cpu() + expected_value = expected[name].detach().float().cpu() + limit = check["atol"] + check["rtol"] * expected_value.abs() + coords = ((actual - expected_value).abs() > limit).nonzero() + details[name] = [ + dict( + index=row.tolist(), + actual=actual[tuple(row)].item(), + fp64_reference=expected[name].detach().cpu()[tuple(row)].item(), + comparison_reference=expected_value[tuple(row)].item(), + tolerance=limit[tuple(row)].item(), + ) + for row in coords[:1000] + ] + case["checks"][-1]["failure_elements"] = details + save() + print(a.mode, stream, dtype, "accuracy complete", flush=True) + elif a.phase == "reuse": + inp, g = inputs(3) + out = run(fn, inp, g, True) + hi = {k: v.double().cpu().requires_grad_() for k, v in inp.items()} + gold = reference_projection(**hi, execution_dtype=dtype) + gg = torch.autograd.grad(gold, tuple(hi.values()), g.double().cpu()) + expected = {"output": gold.detach(), **dict(zip(hi, gg))} + case["fp64_relative_max_error"] = {k: metric(v, expected[k]) for k, v in out.items()} + case["repeat_all_gradients"] = { + k: equal_bytes(v, run(fn, inp, g, True)[k]) for k, v in out.items() + } + case["performance"] = [] + for batch, seq in [(1, 1), (3, 1), (256, 1), (2048, 1)]: + inp, g = inputs(batch, seq) + inp = {k: v.detach().requires_grad_() for k, v in inp.items()} + case["performance"].append( + dict( + batch=batch, + sequence=seq, + measurements=benchmark({a.mode: fn}, inp, 10), + ) + ) + save() + print(a.mode, stream, dtype, "reuse complete", flush=True) + elif a.phase == "workload": + for seed in [0, 127, 2026]: + inp, g = inputs(4096, seed=seed) + whole = run(fn, inp, g) + for chunk in [7, 64, 257, 1024, 2048]: + counts = {k: 0 for k in whole} + for start in range(0, 4096, chunk): + stop = min(start + chunk, 4096) + got = run(fn, sliced(inp, slice(start, stop)), g[start:stop]) + for k, n in stats( + got, {k: v[start:stop] for k, v in whole.items()} + ).items(): + counts[k] += n + case["checks"].append( + dict( + kind="workload_4096", + seed=seed, + batch=4096, + chunk=chunk, + rows=4096, + mismatched_rows=counts, + ) + ) + save() + print(a.mode, stream, dtype, "workload seed", seed, flush=True) + else: + sizes = sorted( + set(range(1, 10)) + | {2**k + d for k in range(4, 12) for d in [-1, 0, 1] if 2**k + d <= 2048} + ) + for seed in [0, 127, 2026]: + inp, g = inputs(2048, seed=seed) + # One singleton reference per row, reused across full-batch checks. + singles = {} + for i in range(2048): + one = run(fn, sliced(inp, slice(i, i + 1)), g[i : i + 1]) + for k, v in one.items(): + if k not in singles: + singles[k] = torch.empty( + (2048, *v.shape[1:]), device=v.device, dtype=v.dtype + ) + singles[k][i : i + 1] = v + for batch in [64, 257, 2048]: + got = run(fn, sliced(inp, slice(0, batch)), g[:batch]) + case["checks"].append( + dict( + kind="every_row_singleton", + seed=seed, + batch=batch, + rows=batch, + mismatched_rows=stats(got, {k: v[:batch] for k, v in singles.items()}), + ) + ) + # Every partition covers every row; compare per-sample gradients only. + whole = run(fn, inp, g) + for chunk in [1, 7, 64, 257, 1024]: + counts = {k: 0 for k in whole} + for start in range(0, 2048, chunk): + stop = min(start + chunk, 2048) + got = ( + {k: v[start:stop] for k, v in singles.items()} + if chunk == 1 + else run(fn, sliced(inp, slice(start, stop)), g[start:stop]) + ) + for k, n in stats( + got, {k: v[start:stop] for k, v in whole.items()} + ).items(): + counts[k] += n + case["checks"].append( + dict( + kind="full_partition", + seed=seed, + batch=2048, + chunk=chunk, + rows=2048, + mismatched_rows=counts, + ) + ) + # Front/middle/back target, then reverse entire batch; raw bytes. + for batch in sizes: + for pos in sorted({0, batch // 2, batch - 1}): + idx = torch.arange(batch, device="cuda") + idx[0], idx[pos] = idx[pos].clone(), idx[0].clone() + got = run(fn, sliced(inp, idx), g[idx]) + case["checks"].append( + dict( + kind="dense_position", + seed=seed, + batch=batch, + position=pos, + rows=1, + mismatched_rows=stats( + {k: v[pos : pos + 1] for k, v in got.items()}, + {k: v[:1] for k, v in singles.items()}, + ), + ) + ) + idx = torch.arange(batch - 1, -1, -1, device="cuda") + got = run(fn, sliced(inp, idx), g[idx]) + case["checks"].append( + dict( + kind="reverse", + seed=seed, + batch=batch, + rows=batch, + mismatched_rows=stats(got, {k: v[idx] for k, v in singles.items()}), + ) + ) + small, sg = inputs(9, seed=seed) + first = run(fn, small, sg, True) + again = run(fn, small, sg, True) + case["checks"].append( + dict( + kind="ordered_repeat", + seed=seed, + all_equal={k: equal_bytes(v, again[k]) for k, v in first.items()}, + ) + ) + save() + print(a.mode, stream, dtype, "seed", seed, "BI complete", flush=True) + save() diff --git a/benchmarks/models/qwenimage/adaln_reuse_candidate.py b/benchmarks/models/qwenimage/adaln_reuse_candidate.py new file mode 100644 index 000000000..6f249dd4f --- /dev/null +++ b/benchmarks/models/qwenimage/adaln_reuse_candidate.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Audit-only composition of the repository's vendored vLLM persistent GEMM. + +This adapter supplies autograd; the upstream GEMM itself is forward-only. +No operator registration or global dispatch override is performed. FP32 +Triton products require TRITON_F32_DEFAULT=ieee in the audit subprocess. +""" +import os + +import torch +import torch.nn.functional as F + +from benchmarks.layers.vllm_batch_invariant_matmul import matmul_persistent + + +class _Linear(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias): + act = F.silu(temb).contiguous() + ctx.save_for_backward(temb, act, weight) + return matmul_persistent(act, weight.T, bias) + + @staticmethod + def backward(ctx, grad): + temb, act, weight = ctx.saved_tensors + # Match the fixed CUDA operator's declared FP32 VJP intermediates. + grad = grad.float().contiguous() + da = matmul_persistent(grad, weight.float()) + dw = matmul_persistent(grad.T, act.float()).to(weight.dtype) + db = grad.sum(0).to(weight.dtype) + z = temb.float() + sigmoid = torch.sigmoid(z) + dt = da * (sigmoid * (1 + z * (1 - sigmoid))) + return dt.to(temb.dtype), dw, db + + +class Projection: + def __call__(self, temb, weight, bias): + assert os.environ.get("TRITON_F32_DEFAULT") == "ieee" + return _Linear.apply(temb, weight, bias) diff --git a/benchmarks/models/qwenimage/plot_adaln_audit.py b/benchmarks/models/qwenimage/plot_adaln_audit.py new file mode 100644 index 000000000..58e7e40da --- /dev/null +++ b/benchmarks/models/qwenimage/plot_adaln_audit.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +# isort: skip_file +# SPDX-License-Identifier: Apache-2.0 +"""Export review figures and CSV from raw audit reports (no recomputation).""" +import argparse +import csv +import json +from pathlib import Path +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + +from benchmarks.models.qwenimage.streams import STREAMS # noqa: E402 + +p = argparse.ArgumentParser() +p.add_argument("directory", type=Path) +a = p.parse_args() +modes = ["native", "cublas_policy", "fixed", "vllm_local", "sglang_local", "megatron_local"] +labels = [ + "Native", + "cuBLAS policy", + "Fixed CUDA", + "vLLM GEMM adapter", + "SGLang GEMM adapter", + "Megatron GEMM adapter", +] +keys = [(s, "torch." + d) for d in ["bfloat16", "float32"] for s, _ in STREAMS] +reports = {} +for mode in modes: + for phase in ["reuse", "bi", "accuracy", "workload"]: + path = a.directory / f"{mode}-{phase}.json" + if path.exists(): + reports[mode, phase] = { + (c["stream"], c["dtype"]): c for c in json.loads(path.read_text())["cases"] + } +rows = [] +fig, axes = plt.subplots(1, 2, figsize=(16, 5.5), layout="constrained") +for ax, component, title in zip( + axes, + ["output", "temb"], + ["Forward: singleton vs full batch", "d_temb: singleton vs full batch"], +): + data = np.full((6, len(modes)), np.nan) + annotations = {} + for j, mode in enumerate(modes): + for i, key in enumerate(keys): + case = reports.get((mode, "bi"), {}).get(key) + if case is None: + continue + checks = [c for c in case["checks"] if c["kind"] == "every_row_singleton"] + n = sum(c["rows"] for c in checks) + bad = sum(c["mismatched_rows"][component] for c in checks) + data[i, j] = bad / n if n else np.nan + annotations[i, j] = f"{bad:,}/{n:,}" + rows.append( + dict( + mode=mode, + stream=key[0], + dtype=key[1], + metric=component + "_BI_rows", + bad=bad, + total=n, + ) + ) + im = ax.imshow(data, vmin=0, vmax=1, cmap="RdYlGn_r", aspect="auto") + ax.set_xticks(range(len(modes)), labels, rotation=20, ha="right", fontsize=9) + ax.set_yticks(range(len(keys)), [f'{s} / {d.removeprefix("torch.")}' for s, d in keys]) + ax.set_title(title) + for i in range(len(keys)): + for j in range(len(modes)): + ax.text( + j, + i, + annotations.get((i, j), "not tested"), + ha="center", + va="center", + fontsize=8, + color=( + "white" + if np.isfinite(data[i, j]) and (data[i, j] < 0.1 or data[i, j] > 0.9) + else "black" + ), + ) +fig.suptitle( + "Qwen-Image / B200 — raw-byte BI mismatches (lower is better)\n" + "3 seeds; every row at B=64,257,2048. This panel alone is not the full BI verdict.", + fontsize=13, +) +fig.savefig(a.directory / "bi.png", dpi=180) +plt.close(fig) + +fig, axes = plt.subplots(2, 2, figsize=(16, 10), layout="constrained") +colors = [ + "#64748b", + "#eab308", + "#2563eb", + "#9333ea", + "#059669", + "#ea580c", + "#db2777", + "#0891b2", +] +width = 0.8 / len(modes) +for ax, batch, title in zip( + axes[0], + [1, 2048], + ["Small batch: full forward + backward", "Large batch: full forward + backward"], +): + for j, mode in enumerate(modes): + vals = [] + lows = [] + highs = [] + for key in keys: + case = reports.get((mode, "reuse"), {}).get(key) + match = ( + next((x for x in case["performance"] if x["batch"] == batch), None) + if case + else None + ) + samples = match["measurements"]["forward_backward"]["samples_ms"][mode] if match else [] + vals.append(np.median(samples) if samples else np.nan) + lows.append(np.percentile(samples, 10) if samples else np.nan) + highs.append(np.percentile(samples, 90) if samples else np.nan) + vals = np.array(vals) + x = np.arange(len(keys)) + (j - (len(modes) - 1) / 2) * width + ax.bar(x, vals, width=0.95 * width, label=labels[j], color=colors[j]) + ax.errorbar( + x, + vals, + yerr=np.array([vals - np.array(lows), np.array(highs) - vals]), + fmt="none", + ecolor="black", + linewidth=0.6, + ) + ax.set_yscale("log") + ax.set_ylabel("ms (median; whiskers p10–p90)") + ax.set_title(title) + ax.set_xticks( + range(len(keys)), [f'{s}\n{d.removeprefix("torch.")}' for s, d in keys], fontsize=9 + ) + ax.grid(axis="y", alpha=0.2) + ax.set_axisbelow(True) +axes[0, 0].legend(fontsize=8) +ax = axes[1, 0] +for j, mode in enumerate(modes): + vals = [] + for key in keys: + c = reports.get((mode, "reuse"), {}).get(key) + vals.append( + max(v for k, v in c["fp64_relative_max_error"].items() if k != "output") + if c + else np.nan + ) + ax.bar( + np.arange(len(keys)) + (j - (len(modes) - 1) / 2) * width, + vals, + width=0.95 * width, + color=colors[j], + ) +ax.set_yscale("log") +ax.set_ylabel("max |error| / max |FP64 reference|") +ax.set_title("Worst gradient normalized error at B=3, S=3") +ax.set_xticks(range(len(keys)), [f'{s}\n{d.removeprefix("torch.")}' for s, d in keys], fontsize=9) +ax.grid(axis="y", alpha=0.2) +ax.set_axisbelow(True) +ax = axes[1, 1] +ax.axis("off") +lines = ["Full-BI checks beyond the heatmap:"] +for mode, label in zip(modes, labels): + cs = list(reports.get((mode, "bi"), {}).values()) + n = sum(len(c["checks"]) for c in cs) + bad = sum( + any(x.get("mismatched_rows", {}).values()) or not all(x.get("all_equal", {}).values()) + for c in cs + for x in c["checks"] + ) + lines.append(f"{label}: {bad}/{n} check records fail") +lines += [ + "", + "Latency includes Python launch gaps; separate processes.", + "Raw samples include fwd, bwd, total and peak memory.", + "FP64 uses explicit storage casts and straight-through VJP.", + "Strict elementwise tolerance results: *-accuracy.json.", + "Source commit and file hashes are recorded in each report.", +] +ax.text(0, 1, "\n".join(lines), va="top", fontsize=11, linespacing=1.6) +fig.suptitle("Reuse audit: speed does not replace accuracy and BI acceptance", fontsize=15) +fig.savefig(a.directory / "reuse.png", dpi=180) +plt.close(fig) +with (a.directory / "bi-summary.csv").open("w") as f: + writer = csv.DictWriter(f, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + +# Strict acceptance is deliberately separate from normalized error at small B. +fig, axes = plt.subplots(1, 2, figsize=(15, 6), layout="constrained") +for ax, dtype in zip(axes, ["torch.bfloat16", "torch.float32"]): + axis_keys = [(s, b) for s, _ in STREAMS for b in [256, 2048]] + for j, mode in enumerate(["fixed", "vllm_local"]): + path = a.directory / f"{mode}-accuracy_large.json" + if not path.exists(): + continue + cases = json.loads(path.read_text())["cases"] + vals, bads = [], [] + for stream, batch in axis_keys: + checks = [ + x + for c in cases + if c["stream"] == stream and c["dtype"] == dtype + for x in c["checks"] + if x["batch"] == batch + ] + vals.append( + max( + (v["max_tolerance_ratio"] for x in checks for v in x["contract"].values()), + default=np.nan, + ) + ) + bads.append(sum(v["violations"] for x in checks for v in x["contract"].values())) + xs = np.arange(len(axis_keys)) + (j - 0.5) * 0.35 + ax.bar( + xs, + vals, + 0.33, + label=["Fixed CUDA", "vLLM GEMM adapter"][j], + color=["#2563eb", "#9333ea"][j], + ) + for x, val, bad in zip(xs, vals, bads): + if np.isfinite(val): + ax.text( + x, + val + 0.025, + f"{val:.2f}\n{bad} fail", + ha="center", + fontsize=8, + color="#b91c1c" if bad else "#1f2937", + ) + ax.axhline(1, color="#b91c1c", linestyle="--", label="Acceptance limit") + ax.set_xticks(range(len(axis_keys)), [f"{s}\nB={b}" for s, b in axis_keys]) + ax.set_title(dtype.removeprefix("torch.")) + ax.set_ylabel("max error / (atol + rtol × |reference|)") + ax.set_ylim(top=max(1.25, ax.get_ylim()[1] * 1.2)) + ax.legend(fontsize=9) +fig.suptitle("Large-batch FP64 check: BI passes do not imply numerical acceptance", fontsize=14) +fig.savefig(a.directory / "accuracy.png", dpi=180) +plt.close(fig) diff --git a/benchmarks/models/qwenimage/plot_adaln_evidence.py b/benchmarks/models/qwenimage/plot_adaln_evidence.py new file mode 100644 index 000000000..d8f1d86be --- /dev/null +++ b/benchmarks/models/qwenimage/plot_adaln_evidence.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Plot public-operator medians from an adaln_evidence JSON report.""" +import argparse +import json + +import matplotlib.pyplot as plt +import numpy as np + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("report") + parser.add_argument("output") + args = parser.parse_args() + plt.switch_backend("Agg") + with open(args.report) as handle: + report = json.load(handle) + cases = [ + c + for c in report["cases"] + if c["dtype"] == "torch.bfloat16" and c["batch"] == 1 and c["seq"] in (1, 6889) + ] + labels = [f'{c["stream"]}\nB={c["batch"]}, S={c["seq"]}' for c in cases] + fig, axes = plt.subplots(1, 3, figsize=(11, 3.7), layout="constrained") + positions = np.arange(len(cases)) + for axis, mode, title in zip( + axes, + ("forward", "backward", "forward_backward"), + ("Forward", "Backward only", "Forward + backward"), + ): + for offset, name, label in [ + (-0.18, "native", "Native PyTorch"), + (0.18, "candidate", "Strict CUDA"), + ]: + values = [c["timing"][mode]["median_ms"][name] for c in cases] + bars = axis.bar(positions + offset, values, width=0.36, label=label) + axis.bar_label(bars, fmt="%.3f", fontsize=8, padding=2) + axis.set_xticks(positions, labels) + axis.set_title(title) + axis.set_ylabel("Median latency (ms)") + axis.margins(y=0.25) + axes[0].legend(fontsize=8) + fig.suptitle("Qwen-Image AdaLN • B200 • BF16 • real checkpoint weights") + fig.savefig(args.output, dpi=180) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/models/qwenimage/streams.py b/benchmarks/models/qwenimage/streams.py new file mode 100644 index 000000000..4167ce8cb --- /dev/null +++ b/benchmarks/models/qwenimage/streams.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Checkpoint streams shared by the evidence runners, validator and plots.""" + +# Checkpoint tensors exercised per stream: block-0 image and text projections. +STREAMS = ( + ("image", "transformer_blocks.0.img_mod.1"), + ("text", "transformer_blocks.0.txt_mod.1"), +) diff --git a/benchmarks/models/qwenimage/validate_adaln_repair.py b/benchmarks/models/qwenimage/validate_adaln_repair.py new file mode 100644 index 000000000..c3309f5d3 --- /dev/null +++ b/benchmarks/models/qwenimage/validate_adaln_repair.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Check completeness and acceptance without reinterpreting tolerance failures.""" +import argparse +import json +from pathlib import Path + +from benchmarks.models.qwenimage.streams import STREAMS + +# Checks per (stream, dtype) case. +CHECKS_PER_CASE = { + "accuracy": {"image": 2, "text": 2}, + "accuracy_large": {"image": 2, "text": 2}, + "bi": {"image": 402, "text": 402}, + "workload": {"image": 15, "text": 15}, + "parameters": {"image": 21, "text": 21}, +} + +p = argparse.ArgumentParser() +p.add_argument("directory", type=Path) +p.add_argument("--allow-fail", action="store_true") +p.add_argument("--allow-dirty", action="store_true", help="Accept reports from an uncommitted tree") +a = p.parse_args() +rows, failures, extensions, commits = [], [], set(), set() +streams = [stream for stream, _ in STREAMS] +for phase, seeds in [ + ("accuracy", [0, 127, 2026]), + ("accuracy_large", [0, 127, 2026]), + ("bi", [0]), + ("workload", [0]), + ("parameters", [0]), +]: + expected = 2 * sum(CHECKS_PER_CASE[phase][stream] for stream in streams) + for seed in seeds: + path = a.directory / f"{phase}-{seed}.json" + report = json.loads(path.read_text()) + assert a.allow_dirty or report["source"]["clean_commit"], (path, "dirty tree") + commits.add(report["source"]["base_sha"]) + extensions.add(report["source"]["extension_sha256"]) + assert [c["stream"] for c in report["cases"][::2]] == streams, path + checks = sum(len(c["checks"]) for c in report["cases"]) + assert checks == expected, (path, checks, expected) + bad = [] + for c in report["cases"]: + for check in c["checks"]: + for field in ["contract", "cpu_fp32_contract"]: + for tensor, result in check.get(field, {}).items(): + if not result["passed"]: + bad.append( + dict( + stream=c["stream"], + dtype=c["dtype"], + batch=check.get("batch"), + sequence=check.get("sequence"), + field=field, + tensor=tensor, + result=result, + ) + ) + if any(check.get("mismatched_rows", {}).values()) or not all( + check.get("all_equal", {}).values() + ): + bad.append(dict(stream=c["stream"], dtype=c["dtype"], check=check)) + rows.append(dict(file=path.name, checks=checks, failures=len(bad))) + failures.extend(dict(file=path.name, **b) for b in bad) +screen = json.loads((a.directory / "cancellation-screen.json").read_text()) +assert a.allow_dirty or screen["source"]["clean_commit"], "dirty tree" +commits.add(screen["source"]["base_sha"]) +extensions.add(screen["source"]["extension_sha256"]) +assert len(screen["cases"]) == 8 +bad = [c for c in screen["cases"] if not c["fixed_contract"]["passed"]] +rows.append(dict(file="cancellation-screen.json", checks=8, failures=len(bad))) +failures.extend(dict(file="cancellation-screen.json", check=c) for c in bad) +assert len(extensions) == 1, "Do not combine results from different extension builds" +assert len(commits) == 1, "Do not combine results from different commits" +result = dict( + qualified=not failures, + commit=next(iter(commits)), + extension_sha256=next(iter(extensions)), + reports=rows, + failures=failures, +) +(a.directory / "acceptance.json").write_text(json.dumps(result, indent=2) + "\n") +print(json.dumps(result, indent=2)) +if failures and not a.allow_fail: + raise SystemExit(1) diff --git a/build_tools/extensions.py b/build_tools/extensions.py index db079e828..bce062253 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -236,6 +236,11 @@ def get_extensions(): # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. cuda_sources.append("csrc/cuda/collectives/deterministic_collective.cu") + cuda_sources.extend( + [ + "csrc/cuda/model_specific/qwenimage/det_linear.cu", + ] + ) # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index 0c7c3e7f4..f29fb5abb 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -4,6 +4,15 @@ #include #include + +#ifdef KERNEL_ALIGN_WITH_CUDA +bool qwenimage_strict_math_compiled(); +torch::Tensor qwenimage_det_linear_backward_bias(torch::Tensor, c10::ScalarType); +std::vector qwenimage_det_linear_forward(torch::Tensor, torch::Tensor, c10::optional, int64_t, bool); +torch::Tensor qwenimage_det_linear_backward_input(torch::Tensor, torch::Tensor, c10::ScalarType); +std::vector qwenimage_det_linear_backward_weight(torch::Tensor, torch::Tensor, c10::ScalarType, bool); +#endif + // Fused LogP Declarations torch::Tensor fused_logp_forward(torch::Tensor logits, torch::Tensor token_ids); #if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) @@ -502,6 +511,14 @@ at::Tensor prefix_shared_attention( // PyBind11 Module Registration PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { +#ifdef KERNEL_ALIGN_WITH_CUDA + m.def("qwenimage_det_linear_forward", &qwenimage_det_linear_forward); + m.def("qwenimage_det_linear_backward_input", &qwenimage_det_linear_backward_input); + m.def("qwenimage_det_linear_backward_weight", &qwenimage_det_linear_backward_weight); + m.def("qwenimage_strict_math_compiled", &qwenimage_strict_math_compiled); + m.def("qwenimage_det_linear_backward_bias", &qwenimage_det_linear_backward_bias); +#endif + m.doc() = "RL-Kernel High-Performance Operator Extension Library"; m.def("fused_logp", &fused_logp_forward, "Fused logp forward fallback"); diff --git a/csrc/cuda/model_specific/qwenimage/det_linear.cu b/csrc/cuda/model_specific/qwenimage/det_linear.cu new file mode 100644 index 000000000..4117e9ad6 --- /dev/null +++ b/csrc/cuda/model_specific/qwenimage/det_linear.cu @@ -0,0 +1,663 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Deterministic, batch-invariant row-wise linear layers for the Qwen-Image +// conditioning path (RFC #386: modulation and final projections). +// +// These layers run on a handful of rows (one per distinct timestep) against +// large weights, so they are GEMVs bound by weight bandwidth. The reduction +// order is fixed and depends only on K (and N for the input gradient), never +// on the number of rows or their position. +// +// Reduction contract qwenimage-det-linear-v1 +// ----------------------------------- +// forward y[t, n] = act(bias[n] + dot(x[t, :], W[n, :])) +// * one warp per output column n; lane l owns the 16-byte K chunks +// c = l, l + 32, l + 64, ... (VEC = 16 / sizeof(w) elements each); +// * each lane accumulates its chunks in ascending order and the elements of +// a chunk in ascending order with fmaf into an FP32 accumulator from 0; +// * the 32 lane sums are combined with an xor butterfly (16, 8, 4, 2, 1), +// whose result is identical on every lane (FP add is commutative); +// * the bias is added once after the butterfly; the activation runs in FP32 +// (SiLU as v / (1 + expf(-v)), PyTorch's formula); one cast at the store. +// * this FP32 forward is used for FP32 weights; BF16 weights take the +// tensor-core forward of contract qwenimage-det-linear-bf16-mma-v1 below. +// d_input dx[t, k] = sum_n g[t, n] * W[n, k] +// * N is split into fixed chunks of 64 rows; each chunk uses ascending +// FP32 Neumaier summation with FMA product residuals; the chunk partials are then Kahan-folded in +// ascending chunk order and cast once. +// d_weight dW[n, k] = sum_t g[t, n] * x[t, k]; d_bias db[n] = sum_t g[t, n] +// * ascending-t Neumaier FP32 summation; dW includes FMA product residuals. +// One final cast. The logical row order determines the cross-row reduction. +// +// Reduction contract qwenimage-det-linear-bf16-mma-v1 (BF16 x and W, FP32 accumulate) +// ----------------------------------------------------------------------------- +// forward y[t, n] = act(bias[n] + dot(x[t, :], W[n, :])) +// * a warp owns 16 output columns (MMA rows) and 8 input rows (MMA columns; +// rows past T are zero), so every launch runs the same instruction +// sequence for every T <= 8, and an MMA column never sees another's data; +// * K is visited in groups of 16 in ascending order. Within a 32-wide step, +// quad q supplies k = 32s + 8q + {0..3} to the first mma.sync m16n8k16 and +// k = 32s + 8q + {4..7} to the second (the same k map for W and x); +// * every mma.sync starts from a zero accumulator, so the tensor core only +// sums 16 products; its FP32 result is added to the running FP32 sum with +// an IEEE add, group by group in ascending k; +// * then bias, FP32 activation and one cast, as above. + +#include +#include +#include +#include +#include + +#include +#include + +bool qwenimage_strict_math_compiled() { +#ifdef __FAST_MATH__ + return false; +#else + return true; +#endif +} + +namespace { + +#ifndef QWENIMAGE_FWD_F32_COLS +#define QWENIMAGE_FWD_F32_COLS 2 +#define QWENIMAGE_FWD_F32_AHEAD 6 +#endif +#ifndef QWENIMAGE_FWD_CONFIGS +#define QWENIMAGE_FWD_CONFIGS QWENIMAGE_FWD_CASE(QWENIMAGE_FWD_F32_COLS, QWENIMAGE_FWD_F32_AHEAD) +#endif + +constexpr int kWarp = 32; +constexpr int kWarpsPerBlock = 8; +constexpr int kRowTile = 4; +constexpr int kDInputChunk = 64; + +enum Activation : int64_t { kActNone = 0, kActSilu = 1 }; + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ float from_float(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { + return __float2bfloat16(v); // round to nearest even, like at::BFloat16 +} + +// 16-byte vector load converted to FP32. +template +struct Vec16 { + static constexpr int kN = 16 / sizeof(T); + __device__ __forceinline__ static void load(const T* ptr, float (&out)[kN]) { + unpack(load_raw(ptr), out); + } + __device__ __forceinline__ static uint4 load_raw(const T* ptr) { + return *reinterpret_cast(ptr); + } + // Streaming weights are read once: bypass L1 allocation. + __device__ __forceinline__ static uint4 load_stream(const T* ptr) { + return __ldcs(reinterpret_cast(ptr)); + } + __device__ __forceinline__ static void unpack(const uint4& raw, float (&out)[kN]) { + const T* vals = reinterpret_cast(&raw); +#pragma unroll + for (int i = 0; i < kN; ++i) out[i] = to_float(vals[i]); + } +}; + +__device__ __forceinline__ float warp_butterfly_sum(float v) { +#pragma unroll + for (int offset = kWarp / 2; offset > 0; offset >>= 1) { + v += __shfl_xor_sync(0xffffffffu, v, offset); + } + return v; +} + +// Each warp owns kCols adjacent output columns so that kLoadAhead x kCols +// 16-byte weight loads are in flight per lane. Grouping columns and chunks +// changes only when loads are issued: every column still accumulates its own +// chunks in ascending order, exactly as a one-column-per-warp kernel would. +template +__global__ void __launch_bounds__(kWarpsPerBlock * kWarp) + det_linear_forward_kernel(const x_t* __restrict__ x, const w_t* __restrict__ w, + const w_t* __restrict__ bias, out_t* __restrict__ out, + float* __restrict__ pre_act, int64_t rows, int64_t n_out, + int64_t k_in, int64_t activation) { + using WV = Vec16; + using XV = Vec16; + constexpr int kVec = WV::kN; + static_assert(kCols * kRows <= kWarp, "one warp cannot store more than kWarp outputs per tile"); + static_assert(XV::kN == kVec, "x and weight share a dtype"); + constexpr int64_t kStride = static_cast(kWarp) * kVec; + + const int lane = threadIdx.x % kWarp; + const int64_t warp_global = + static_cast(blockIdx.x) * kWarpsPerBlock + threadIdx.x / kWarp; + const int64_t warp_stride = static_cast(gridDim.x) * kWarpsPerBlock; + + for (int64_t n0 = warp_global * kCols; n0 < n_out; n0 += warp_stride * kCols) { + for (int64_t t0 = 0; t0 < rows; t0 += kRows) { + float acc[kCols][kRows]; +#pragma unroll + for (int c = 0; c < kCols; ++c) { +#pragma unroll + for (int r = 0; r < kRows; ++r) acc[c][r] = 0.0f; + } + for (int64_t g0 = static_cast(lane) * kVec; g0 < k_in; + g0 += kStride * kLoadAhead) { + uint4 raw[kLoadAhead][kCols]; +#pragma unroll + for (int u = 0; u < kLoadAhead; ++u) { + const int64_t k0 = g0 + u * kStride; +#pragma unroll + for (int c = 0; c < kCols; ++c) { + if (k0 < k_in && n0 + c < n_out) raw[u][c] = WV::load_stream(w + (n0 + c) * k_in + k0); + } + } +#pragma unroll + for (int u = 0; u < kLoadAhead; ++u) { + const int64_t k0 = g0 + u * kStride; + if (k0 >= k_in) break; +#pragma unroll + for (int r = 0; r < kRows; ++r) { + if (t0 + r < rows) { + float xv[kVec]; + XV::load(x + (t0 + r) * k_in + k0, xv); +#pragma unroll + for (int c = 0; c < kCols; ++c) { + float wv[kVec]; + WV::unpack(raw[u][c], wv); +#pragma unroll + for (int j = 0; j < kVec; ++j) acc[c][r] = fmaf(xv[j], wv[j], acc[c][r]); + } + } + } + } + } +#pragma unroll + for (int c = 0; c < kCols; ++c) { + const int64_t n = n0 + c; +#pragma unroll + for (int r = 0; r < kRows; ++r) { + const int64_t t = t0 + r; + const float sum = warp_butterfly_sum(acc[c][r]); + if (lane == c * kRows + r && t < rows && n < n_out) { + float v = sum + (bias != nullptr ? to_float(bias[n]) : 0.0f); + if (pre_act != nullptr) pre_act[t * n_out + n] = v; + if (activation == kActSilu) v = v / (1.0f + expf(-v)); + out[t * n_out + n] = from_float(v); + } + } + } + } + } +} + +__device__ __forceinline__ void mma_m16n8k16_bf16(float (&c)[4], const uint32_t (&a)[4], + uint32_t b0, uint32_t b1) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + // BF16 mma.sync needs SM80+; the host refuses to launch below that. + __trap(); +#else + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +#endif +} + +constexpr int kMmaCols = 16; // output columns per warp (MMA M) +constexpr int kMmaRows = 8; // input rows per pass (MMA N) +constexpr int kMmaAhead = 4; // 32-wide K steps loaded before use + +__device__ __forceinline__ void mma_group_accumulate(float (&acc)[4], const uint32_t (&a)[4], + uint32_t b0, uint32_t b1) { + float part[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + mma_m16n8k16_bf16(part, a, b0, b1); +#pragma unroll + for (int i = 0; i < 4; ++i) acc[i] += part[i]; +} + +// Contract qwenimage-det-linear-bf16-mma-v1. Lane = 4 * g + q: g picks the MMA row +// (output columns n0 + g and n0 + g + 8) and the MMA column (input row t0 + g). +__global__ void __launch_bounds__(kWarpsPerBlock * kWarp) + det_linear_forward_bf16_mma_kernel(const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ out, float* __restrict__ pre_act, + int64_t rows, int64_t n_out, int64_t k_in, + int64_t activation) { + const int lane = threadIdx.x % kWarp; + const int g = lane >> 2; + const int q = lane & 3; + const int64_t warp_global = + static_cast(blockIdx.x) * kWarpsPerBlock + threadIdx.x / kWarp; + const int64_t warp_stride = static_cast(gridDim.x) * kWarpsPerBlock; + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + + for (int64_t n0 = warp_global * kMmaCols; n0 < n_out; n0 += warp_stride * kMmaCols) { + const bool lo_ok = n0 + g < n_out; + const bool hi_ok = n0 + g + 8 < n_out; + const __nv_bfloat16* w_lo = w + (n0 + g) * k_in; + const __nv_bfloat16* w_hi = w + (n0 + g + 8) * k_in; + for (int64_t t0 = 0; t0 < rows; t0 += kMmaRows) { + const bool x_ok = t0 + g < rows; + const __nv_bfloat16* x_row = x + (t0 + g) * k_in; + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int64_t kb = 0; kb < k_in; kb += 32 * kMmaAhead) { + uint4 a_lo[kMmaAhead], a_hi[kMmaAhead], bx[kMmaAhead]; +#pragma unroll + for (int u = 0; u < kMmaAhead; ++u) { + const int64_t k = kb + 32 * u + 8 * q; + const bool k_ok = k < k_in; + a_lo[u] = (k_ok && lo_ok) ? __ldcs(reinterpret_cast(w_lo + k)) : zero; + a_hi[u] = (k_ok && hi_ok) ? __ldcs(reinterpret_cast(w_hi + k)) : zero; + bx[u] = (k_ok && x_ok) ? *reinterpret_cast(x_row + k) : zero; + } +#pragma unroll + for (int u = 0; u < kMmaAhead; ++u) { + if (kb + 32 * u >= k_in) break; + const uint32_t first[4] = {a_lo[u].x, a_hi[u].x, a_lo[u].y, a_hi[u].y}; + const uint32_t second[4] = {a_lo[u].z, a_hi[u].z, a_lo[u].w, a_hi[u].w}; + mma_group_accumulate(acc, first, bx[u].x, bx[u].y); + mma_group_accumulate(acc, second, bx[u].z, bx[u].w); + } + } + // acc = {(n0+g, t0+2q), (n0+g, t0+2q+1), (n0+g+8, t0+2q), (n0+g+8, t0+2q+1)} +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int64_t n = n0 + g + (i >= 2 ? 8 : 0); + const int64_t t = t0 + 2 * q + (i & 1); + if (n < n_out && t < rows) { + float v = acc[i] + (bias != nullptr ? __bfloat162float(bias[n]) : 0.0f); + if (pre_act != nullptr) pre_act[t * n_out + n] = v; + if (activation == kActSilu) v = v / (1.0f + expf(-v)); + out[t * n_out + n] = __float2bfloat16(v); + } + } + } + } +} + +// partial[c, t, k] = sum_{n in chunk c} g[t, n] * w[n, k] +template +__global__ void det_linear_dinput_partial_kernel(const g_t* __restrict__ grad, + const w_t* __restrict__ w, + float* __restrict__ partial, int64_t rows, + int64_t n_out, int64_t k_in) { + const int64_t k = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t chunk = blockIdx.y; + if (k >= k_in) return; + const int64_t n_begin = chunk * kDInputChunk; + const int64_t n_end = min(n_begin + kDInputChunk, n_out); + for (int64_t t0 = 0; t0 < rows; t0 += kRowTile) { + float acc[kRowTile]; + float correction[kRowTile]; +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + acc[r] = 0.0f; + correction[r] = 0.0f; + } + for (int64_t n = n_begin; n < n_end; ++n) { + const float wv = to_float(w[n * k_in + k]); +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + const int64_t t = t0 + r; + if (t < rows) { + const float gv = to_float(grad[t * n_out + n]); + const float product = __fmul_rn(gv, wv); + const float product_error = __fmaf_rn(gv, wv, -product); + const float next = __fadd_rn(acc[r], product); + const float sum_error = fabsf(acc[r]) >= fabsf(product) + ? __fadd_rn(__fsub_rn(acc[r], next), product) + : __fadd_rn(__fsub_rn(product, next), acc[r]); + correction[r] = __fadd_rn(correction[r], __fadd_rn(sum_error, product_error)); + acc[r] = next; + } + } + } +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + const int64_t t = t0 + r; + if (t < rows) partial[(chunk * rows + t) * k_in + k] = __fadd_rn(acc[r], correction[r]); + } + } +} + +template +__global__ void det_fold_chunks_kernel(const float* __restrict__ partial, out_t* __restrict__ out, + int64_t chunks, int64_t elems) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= elems) return; + float acc = 0.0f; + float correction = 0.0f; + for (int64_t c = 0; c < chunks; ++c) { + const float value = __fsub_rn(partial[c * elems + i], correction); + const float next = __fadd_rn(acc, value); + correction = __fsub_rn(__fsub_rn(next, acc), value); + acc = next; + } + out[i] = from_float(acc); +} + +// dW[n, k]: ascending-t compensated FP32 dot, including product residuals. +template +__global__ void det_linear_dweight_kernel(const g_t* __restrict__ grad, const x_t* __restrict__ x, + w_t* __restrict__ dw, int64_t rows, int64_t n_out, + int64_t k_in) { + const int64_t total = n_out * k_in; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < total; + i += static_cast(gridDim.x) * blockDim.x) { + const int64_t n = i / k_in; + const int64_t k = i - n * k_in; + float acc = 0.0f; + float correction = 0.0f; + for (int64_t t = 0; t < rows; ++t) { + const float gv = to_float(grad[t * n_out + n]); + const float xv = to_float(x[t * k_in + k]); + const float product = __fmul_rn(gv, xv); + const float product_error = __fmaf_rn(gv, xv, -product); + const float next = __fadd_rn(acc, product); + const float sum_error = fabsf(acc) >= fabsf(product) + ? __fadd_rn(__fsub_rn(acc, next), product) + : __fadd_rn(__fsub_rn(product, next), acc); + correction = __fadd_rn(correction, __fadd_rn(sum_error, product_error)); + acc = next; + } + dw[i] = from_float(__fadd_rn(acc, correction)); + } +} + +template +__global__ void det_linear_dbias_kernel(const g_t* __restrict__ grad, w_t* __restrict__ db, + int64_t rows, int64_t n_out) { + const int64_t n = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (n >= n_out) return; + float acc = 0.0f; + float correction = 0.0f; + for (int64_t t = 0; t < rows; ++t) { + const float value = to_float(grad[t * n_out + n]); + const float next = __fadd_rn(acc, value); + const float error = fabsf(acc) >= fabsf(value) + ? __fadd_rn(__fsub_rn(acc, next), value) + : __fadd_rn(__fsub_rn(value, next), acc); + correction = __fadd_rn(correction, error); + acc = next; + } + db[n] = from_float(__fadd_rn(acc, correction)); +} + +template +struct CudaType { + using type = T; +}; +template <> +struct CudaType { + using type = __nv_bfloat16; +}; + +template +const typename CudaType::type* cptr(const torch::Tensor& t) { + return reinterpret_cast::type*>(t.data_ptr()); +} +template +typename CudaType::type* mptr(torch::Tensor& t) { + return reinterpret_cast::type*>(t.data_ptr()); +} + +void check_matrix(const torch::Tensor& t, const char* name) { + TORCH_CHECK(qwenimage_strict_math_compiled(), "Qwen-Image requires a build without fast math"); + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.dim() == 2, name, " must be 2-D"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(t.scalar_type() == at::kFloat || t.scalar_type() == at::kBFloat16, name, + " must be float32 or bfloat16, got ", t.scalar_type()); + TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0, name, + " must be 16-byte aligned"); +} + +int64_t grid_for(int64_t work, int64_t per_block) { + const int64_t blocks = (work + per_block - 1) / per_block; + return std::max(1, std::min(blocks, 1 << 20)); +} + +// Launch shape only: the row tile, the columns per warp and the load depth +// change which loads are in flight, never a column's accumulation order. +struct ForwardConfig { + int rows; + int cols; + int ahead; +}; + +ForwardConfig pick_forward_config(int64_t rows) { + const int row_tile = rows <= 1 ? 1 : (rows <= 2 ? 2 : 4); + return {row_tile, QWENIMAGE_FWD_F32_COLS, QWENIMAGE_FWD_F32_AHEAD}; +} + +template +void launch_forward_impl(const T* x, const T* w, const T* bias, T* out, float* pre, + int64_t rows, int64_t n_out, int64_t k_in, int64_t activation, + cudaStream_t stream) { + const int64_t blocks = grid_for((n_out + kCols - 1) / kCols, kWarpsPerBlock); + det_linear_forward_kernel + <<(blocks), kWarpsPerBlock * kWarp, 0, stream>>>( + x, w, bias, out, pre, rows, n_out, k_in, activation); +} + +template +void launch_forward_rows(const ForwardConfig& cfg, const T* x, const T* w, const T* bias, T* out, + float* pre, int64_t rows, int64_t n_out, int64_t k_in, + int64_t activation, cudaStream_t stream) { +#define QWENIMAGE_FWD_CASE(C, A) \ + if (cfg.cols == C && cfg.ahead == A) { \ + launch_forward_impl(x, w, bias, out, pre, rows, n_out, k_in, activation, \ + stream); \ + return; \ + } + QWENIMAGE_FWD_CONFIGS +#undef QWENIMAGE_FWD_CASE + TORCH_CHECK(false, "unsupported det_linear forward config cols=", cfg.cols, + " ahead=", cfg.ahead); +} + +template +void launch_forward(const ForwardConfig& cfg, const T* x, const T* w, const T* bias, T* out, + float* pre, int64_t rows, int64_t n_out, int64_t k_in, int64_t activation, + cudaStream_t stream) { + if (cfg.rows == 1) { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } else if (cfg.rows == 2) { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } else { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } +} + +} // namespace + +std::vector qwenimage_det_linear_forward(torch::Tensor x, torch::Tensor weight, + c10::optional bias, + int64_t activation, bool save_pre_activation) { + check_matrix(x, "x"); + check_matrix(weight, "weight"); + TORCH_CHECK(x.device() == weight.device(), "x and weight must be on the same device"); + TORCH_CHECK(x.scalar_type() == weight.scalar_type(), + "x and weight must share a dtype (cast at the declared boundary first), got ", + x.scalar_type(), " and ", weight.scalar_type()); + TORCH_CHECK(activation == kActNone || activation == kActSilu, "unknown activation ", + activation); + const int64_t rows = x.size(0); + const int64_t k_in = x.size(1); + const int64_t n_out = weight.size(0); + TORCH_CHECK(rows > 0, "x must have at least one row"); + TORCH_CHECK(weight.size(1) == k_in, "weight is [", n_out, ", ", weight.size(1), + "] but x has K=", k_in); + const int64_t vec = 16 / x.element_size(); + TORCH_CHECK(k_in % vec == 0, "K=", k_in, " must be a multiple of ", vec, + " for 16-byte vector loads"); + if (bias.has_value()) { + TORCH_CHECK(bias->is_cuda() && bias->dim() == 1 && bias->size(0) == n_out && + bias->is_contiguous() && bias->scalar_type() == weight.scalar_type(), + "bias must be a contiguous [N] CUDA tensor with the weight dtype"); + TORCH_CHECK(bias->device() == weight.device(), + "bias and weight must be on the same device"); + } + + const c10::cuda::CUDAGuard device_guard(x.device()); + auto out = torch::empty({rows, n_out}, x.options()); + torch::Tensor pre; + if (save_pre_activation) pre = torch::empty({rows, n_out}, x.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + if (x.scalar_type() == at::kFloat) { + const ForwardConfig cfg = pick_forward_config(rows); + launch_forward(cfg, cptr(x), cptr(weight), + bias.has_value() ? cptr(*bias) : nullptr, mptr(out), + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, + k_in, activation, stream); + } else { + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 8, + "the BF16 Qwen-Image det_linear forward uses mma.sync and needs SM80 or newer"); + const int64_t blocks = grid_for((n_out + kMmaCols - 1) / kMmaCols, kWarpsPerBlock); + det_linear_forward_bf16_mma_kernel<<(blocks), kWarpsPerBlock * kWarp, + 0, stream>>>( + cptr(x), cptr(weight), + bias.has_value() ? cptr(*bias) : nullptr, mptr(out), + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (save_pre_activation) return {out, pre}; + return {out}; +} + +// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype. +torch::Tensor qwenimage_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && + grad.scalar_type() == at::kFloat, + "grad must be a contiguous 2-D float32 CUDA tensor"); + check_matrix(weight, "weight"); + TORCH_CHECK(grad.device() == weight.device(), "grad and weight must be on the same device"); + TORCH_CHECK(grad.size(1) == weight.size(0), "grad N != weight N"); + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + const int64_t rows = grad.size(0); + TORCH_CHECK(rows > 0, "grad must have at least one row"); + const int64_t n_out = weight.size(0); + const int64_t k_in = weight.size(1); + const c10::cuda::CUDAGuard device_guard(grad.device()); + auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); + if (k_in == 0) return out; + if (n_out == 0) return out.zero_(); + const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; + auto partial = torch::empty({chunks, rows, k_in}, grad.options()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + dim3 grid(static_cast((k_in + threads - 1) / threads), static_cast(chunks)); + if (weight.scalar_type() == at::kFloat) { + det_linear_dinput_partial_kernel<<>>( + grad.data_ptr(), cptr(weight), partial.data_ptr(), rows, n_out, + k_in); + } else { + det_linear_dinput_partial_kernel<<>>( + grad.data_ptr(), cptr(weight), partial.data_ptr(), rows, + n_out, k_in); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + const int64_t elems = rows * k_in; + const unsigned fold_blocks = static_cast((elems + threads - 1) / threads); + if (out_dtype == at::kFloat) { + det_fold_chunks_kernel<<>>( + partial.data_ptr(), mptr(out), chunks, elems); + } else { + det_fold_chunks_kernel<__nv_bfloat16><<>>( + partial.data_ptr(), mptr(out), chunks, elems); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// grad [T, N] float32, x [T, K] (float32 or bfloat16) -> (dW [N, K], dbias [N]) in w_dtype. +std::vector qwenimage_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, + c10::ScalarType w_dtype, + bool with_bias) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && + grad.scalar_type() == at::kFloat, + "grad must be a contiguous 2-D float32 CUDA tensor"); + check_matrix(x, "x"); + TORCH_CHECK(grad.device() == x.device(), "grad and x must be on the same device"); + TORCH_CHECK(grad.size(0) == x.size(0), "grad rows != x rows"); + TORCH_CHECK(w_dtype == at::kFloat || w_dtype == at::kBFloat16, + "w_dtype must be float32 or bfloat16"); + const int64_t rows = grad.size(0); + TORCH_CHECK(rows > 0, "grad must have at least one row"); + const int64_t n_out = grad.size(1); + const int64_t k_in = x.size(1); + const c10::cuda::CUDAGuard device_guard(grad.device()); + auto dw = torch::empty({n_out, k_in}, grad.options().dtype(w_dtype)); + if (n_out == 0) { + if (!with_bias) return {dw}; + return {dw, torch::empty({n_out}, grad.options().dtype(w_dtype))}; + } + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const int64_t total = n_out * k_in; + const unsigned blocks = static_cast(grid_for(total, threads)); + const float* g = grad.data_ptr(); +#define QWENIMAGE_DW_LAUNCH(XT, XC, WT, WC) \ + det_linear_dweight_kernel<<>>( \ + g, cptr(x), mptr(dw), rows, n_out, k_in) + if (x.scalar_type() == at::kFloat && w_dtype == at::kFloat) { + QWENIMAGE_DW_LAUNCH(float, float, float, float); + } else if (x.scalar_type() == at::kFloat) { + QWENIMAGE_DW_LAUNCH(float, float, at::BFloat16, __nv_bfloat16); + } else if (w_dtype == at::kFloat) { + QWENIMAGE_DW_LAUNCH(at::BFloat16, __nv_bfloat16, float, float); + } else { + QWENIMAGE_DW_LAUNCH(at::BFloat16, __nv_bfloat16, at::BFloat16, __nv_bfloat16); + } +#undef QWENIMAGE_DW_LAUNCH + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (!with_bias) return {dw}; + auto db = torch::empty({n_out}, grad.options().dtype(w_dtype)); + const unsigned bias_blocks = static_cast((n_out + threads - 1) / threads); + if (w_dtype == at::kFloat) { + det_linear_dbias_kernel<<>>( + g, mptr(db), rows, n_out); + } else { + det_linear_dbias_kernel<<>>( + g, mptr(db), rows, n_out); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dw, db}; +} + +// Bias-only entrypoint for evaluating a different dW kernel without recomputing dW. +torch::Tensor qwenimage_det_linear_backward_bias(torch::Tensor grad, c10::ScalarType dtype) { + check_matrix(grad, "grad"); + TORCH_CHECK(grad.scalar_type() == at::kFloat && grad.size(0) > 0, + "grad must be nonempty float32 [B,N]"); + TORCH_CHECK(dtype == at::kFloat || dtype == at::kBFloat16, + "bias gradient dtype must be float32 or bfloat16"); + const c10::cuda::CUDAGuard guard(grad.device()); + auto db = torch::empty({grad.size(1)}, grad.options().dtype(dtype)); + if (grad.size(1) == 0) return db; + const int threads = 256; + const unsigned blocks = static_cast((grad.size(1) + threads - 1) / threads); + auto stream = at::cuda::getCurrentCUDAStream(); + if (dtype == at::kFloat) { + det_linear_dbias_kernel<<>>( + grad.data_ptr(), db.data_ptr(), grad.size(0), grad.size(1)); + } else { + det_linear_dbias_kernel<<>>( + grad.data_ptr(), mptr(db), grad.size(0), grad.size(1)); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return db; +} diff --git a/docs/.nav.yml b/docs/.nav.yml index 29b5ead37..2e2aa6e71 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -16,6 +16,7 @@ nav: - Hardware Profiling Guide: getting_started/hardware-profiling.md - Operators: - operators/README.md + - operators/qwenimage-adaln-mod-projection.md - operators/activation.md - operators/attention.md - operators/fused-logp.md diff --git a/docs/operators/README.md b/docs/operators/README.md index 00f4cbb45..c3de2fdb5 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -32,4 +32,5 @@ Every operator page should include: - [Matmul](matmul.md) - [Sampling](sampling.md) - [Token Embedding](embedding.md) +- [Qwen-Image AdaLN Modulation Projection](qwenimage-adaln-mod-projection.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/qwenimage-adaln-mod-projection.md b/docs/operators/qwenimage-adaln-mod-projection.md new file mode 100644 index 000000000..76cea7462 --- /dev/null +++ b/docs/operators/qwenimage-adaln-mod-projection.md @@ -0,0 +1,74 @@ +# Qwen-Image AdaLN modulation projection (CUDA) + +> Status: validated on B200 (SM100). Accuracy, batch-invariance, reuse-survey and timing evidence, with reproduction commands, is in [B200 evidence](../usage/evidence/qwenimage-adaln-mod-projection-b200/README.md). + +RFC #386 `adaln_mod_projection`, registered as `qwenimage_adaln_mod_projection` +to keep the H3 and Qwen-Image contracts separate. This is a Phase A single-GPU +operator, not the block's LayerNorm/modulation or gated residual. + +## Interface + +```python +from rl_engine.runtime.registry import KernelRegistry + +op = KernelRegistry().get_op("qwenimage_adaln_mod_projection", device=temb.device) +params = op(temb=temb, weight=weight, bias=bias) +shift_attn, scale_attn, gate_attn, shift_mlp, scale_mlp, gate_mlp = params.chunk(6, -1) +trace = op.implementation_info(temb=temb, weight=weight, bias=bias) +``` + +`temb` is `[B,3072]`, `weight` is `[18432,3072]`, `bias` is `[18432]` and +the output is `[B,18432]`. Run this interface separately for `img_mod.1` and +`txt_mod.1`; there is no shared weight assumption. All inputs share a device +and dtype (BF16 or FP32), and B must be positive. Noncontiguous views are copied +explicitly. Backward returns gradients for all three inputs. First-order +autograd only; gradient accumulation into leaves follows PyTorch's usual rules. + +The CUDA registry slot has no reference fallback. NVIDIA SM80+ is required for +the implementation; only B200/SM100 is validated here. A missing extension or +fast-math build fails closed. CPU reference classes are separate; no optimized +ROCm, Ascend, or MUSA slot is registered. + +## Arithmetic and reuse + +Contract `qwenimage-adaln-v2`, backward `qwenimage-parameter-grad-neumaier-v2`: + +1. Native CUDA SiLU, with a declared output store in the execution dtype. +2. Fixed-order projection with FP32 accumulation, one bias addition, and an + execution-dtype output store. FP32 uses one warp's fixed vector/FMA order and + XOR reduction. BF16 uses zero-initialized `mma.sync` K=16 partials summed in + a fixed order. Neither depends on B or logical batch position. +3. Backward keeps intermediates FP32. Input gradients use 64-column compensated dot + partials and an ascending compensated (Kahan) fold. Parameter gradients use + ascending FP32 Neumaier accumulation, including FMA product residuals for + `d_weight`; `d_bias` uses ascending Neumaier summation. Only the final + result is cast. SiLU's analytic VJP is evaluated in + FP32, and leaf gradients are cast to each input's dtype. + +The projection source is adapted from the H3 deterministic GEMV at +`af3311ab3cf22e0d77ddd523cd4283990a9f8377`, retaining the forward arithmetic and +using compensated FP32 summation for the input-gradient dot and final fold. There +is no dependency on H3 operators or their stacked branches. + +`adaln_prior_art.py` compares native `SiLU + F.linear`, this reused GEMV and, +optionally, the existing Qwen-Image QKV cuBLASLt extension. The native FP32 +projection changes bytes across batch sizes in the measured sweep. The QKV +forward is invariant in that sweep but is slower for these small conditioning +batches. Native BF16 forward passes the measured sweep; this does not establish +its backward or an all-dtype strict contract. The prior-art report makes only +a forward claim for the QKV extension. No process-wide cuBLAS/TF32 settings are +changed by this implementation. + +The independent reference uses CPU FP32 SiLU/linear algebra with explicit +storage rounding and straight-through derivatives at cast boundaries. Its +input-gradient dot is blocked into 128-column CPU BLAS operations and pairwise +FP32 addition: a single long CPU BLAS dot can itself miss the repository's +FP32 gradient tolerance. FP64 diagnostics check this reference independently. +This reference does not call the candidate CUDA kernel. + +## Validation and reproduction + +See the [B200 evidence and review notes](../usage/evidence/qwenimage-adaln-mod-projection-b200/README.md). +Reports use the existing `reduction` forward/gradient tolerance rows, without +private overrides. Strict replay and batch invariance are checked with +`torch.equal` on contiguous raw-byte views, separately from agreement with the CPU reference. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index fb3c4379c..4fb9ab443 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -325,3 +325,25 @@ def deterministic_collective_rocm_ipc_all_gather_input( input: torch.Tensor, output: torch.Tensor, ) -> None: ... + +# Qwen-Image conditioning; NVIDIA CUDA only, strict math build. +def qwenimage_det_linear_forward( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + activation: int, + save_pre_activation: bool, +) -> list[torch.Tensor]: ... +def qwenimage_det_linear_backward_input( + grad: torch.Tensor, + weight: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: ... +def qwenimage_det_linear_backward_weight( + grad: torch.Tensor, + x: torch.Tensor, + out_dtype: torch.dtype, + with_bias: bool, +) -> list[torch.Tensor]: ... +def qwenimage_strict_math_compiled() -> bool: ... +def qwenimage_det_linear_backward_bias(grad: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: ... diff --git a/rl_engine/backends/cuda/model_specific/__init__.py b/rl_engine/backends/cuda/model_specific/__init__.py new file mode 100644 index 000000000..854372bd1 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Model-specific CUDA implementations.""" diff --git a/rl_engine/backends/cuda/model_specific/qwenimage/__init__.py b/rl_engine/backends/cuda/model_specific/qwenimage/__init__.py new file mode 100644 index 000000000..1d9e4c538 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/qwenimage/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen-Image Phase A conditioning operators.""" diff --git a/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py b/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py new file mode 100644 index 000000000..273fdc971 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Fixed-order CUDA Qwen-Image modulation projection, RFC #386.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch.autograd.function import once_differentiable + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.reference.model_specific.qwenimage.adaln_mod_projection import validate_projection + +CONTRACT = "qwenimage-adaln-v2" +_SYMBOLS = ( + "qwenimage_det_linear_forward", + "qwenimage_det_linear_backward_input", + "qwenimage_det_linear_backward_weight", + # Introduced with v2; reject stale v1 extensions with uncompensated dW/db. + "qwenimage_det_linear_backward_bias", + "qwenimage_strict_math_compiled", +) + + +def require_extension(): + if not _EXT_AVAILABLE or not all(hasattr(_C, name) for name in _SYMBOLS): + raise RuntimeError("rebuild rl_engine._C with Qwen-Image strict CUDA kernels") + if not _C.qwenimage_strict_math_compiled(): + raise RuntimeError("Qwen-Image requires a build without fast math") + + +def require_device(tensor): + if tensor.device.type != "cuda" or torch.version.hip is not None: + raise ValueError("Qwen-Image CUDA requires NVIDIA CUDA tensors") + if torch.cuda.get_device_capability(tensor.device) < (8, 0): + raise RuntimeError("Qwen-Image CUDA requires SM80 or newer") + + +def projection_forward(temb, weight, bias): + act = F.silu(temb).contiguous() + table = _C.qwenimage_det_linear_forward(act, weight.contiguous(), bias.contiguous(), 0, False)[ + 0 + ] + return table, act + + +def parameter_gradients(grad, act, dtype): + return _C.qwenimage_det_linear_backward_weight(grad, act, dtype, True) + + +def projection_backward(grad, temb, act, weight): + grad = grad.float().contiguous() + dw, db = parameter_gradients(grad, act, weight.dtype) + da = _C.qwenimage_det_linear_backward_input(grad, weight.contiguous(), torch.float32) + z = temb.float() + sigmoid = torch.sigmoid(z) + dt = da * (sigmoid * (1 + z * (1 - sigmoid))) + return dt.to(temb.dtype), dw, db + + +def implementation_info(op, tensor): + require_extension() + require_device(tensor) + return { + "operator": op, + "requested_backend": "cuda", + "actual_backend": "cuda", + "contract": CONTRACT, + "fallback": False, + "accumulation_dtype": "float32", + "projection_kernel": "rl_engine._C.qwenimage_det_linear_forward", + "projection_tree": ("mma_k16_ascending" if tensor.dtype == torch.bfloat16 else "warp_xor"), + "backward": ( + "compensated_64_column_partials_kahan_fold; " + "ascending_batch_neumaier_dW_product_residual_db" + ), + "parameter_gradient_kernel": "rl_engine._C.qwenimage_det_linear_backward_weight", + "parameter_gradient_backend": "cuda_compensated", + "backward_contract": "qwenimage-parameter-grad-neumaier-v2", + "activation": "aten.silu; fp32 analytic VJP", + "execution_dtype": str(tensor.dtype), + "device": str(tensor.device), + "capability": list(torch.cuda.get_device_capability(tensor.device)), + } + + +class _Projection(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias): + table, act = projection_forward(temb, weight, bias) + ctx.save_for_backward(temb, act, weight) + return table + + @staticmethod + @once_differentiable + def backward(ctx, grad): + return projection_backward(grad, *ctx.saved_tensors) + + +class QwenImageAdaLNModProjectionCudaOp: + op_class = "reduction" + kernel_id = "aten.silu+rl_engine._C.qwenimage_det_linear_forward" + + def __init__(self): + require_extension() + + def forward(self, temb, weight, bias): + validate_projection(temb, weight, bias, 6) + require_device(temb) + return _Projection.apply(temb, weight, bias) + + __call__ = forward + + def implementation_info(self, temb, weight, bias): + validate_projection(temb, weight, bias, 6) + return implementation_info("qwenimage_adaln_mod_projection", temb) diff --git a/rl_engine/reference/model_specific/__init__.py b/rl_engine/reference/model_specific/__init__.py new file mode 100644 index 000000000..20aaa7470 --- /dev/null +++ b/rl_engine/reference/model_specific/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Model-specific mathematical references.""" diff --git a/rl_engine/reference/model_specific/qwenimage/__init__.py b/rl_engine/reference/model_specific/qwenimage/__init__.py new file mode 100644 index 000000000..1d9e4c538 --- /dev/null +++ b/rl_engine/reference/model_specific/qwenimage/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen-Image Phase A conditioning operators.""" diff --git a/rl_engine/reference/model_specific/qwenimage/adaln_mod_projection.py b/rl_engine/reference/model_specific/qwenimage/adaln_mod_projection.py new file mode 100644 index 000000000..33fde0588 --- /dev/null +++ b/rl_engine/reference/model_specific/qwenimage/adaln_mod_projection.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Qwen-Image modulation projection and independent precision reference.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +HIDDEN = 3072 + + +def validate_projection(temb, weight, bias, multiplier): + if temb.ndim != 2 or temb.shape[0] < 1 or temb.shape[1] != HIDDEN: + raise ValueError("temb must be nonempty [B,3072]") + if tuple(weight.shape) != (multiplier * HIDDEN, HIDDEN): + raise ValueError(f"weight must be [{multiplier * HIDDEN},3072]") + if tuple(bias.shape) != (multiplier * HIDDEN,): + raise ValueError(f"bias must be [{multiplier * HIDDEN}]") + if temb.dtype not in (torch.bfloat16, torch.float32): + raise TypeError("only bfloat16 and float32 are supported") + if any(t.dtype != temb.dtype for t in (weight, bias)): + raise TypeError("temb, weight and bias must share a dtype") + if any(t.device != temb.device for t in (weight, bias)): + raise ValueError("all inputs must be on the same device") + + +class _RoundStraightThrough(torch.autograd.Function): + @staticmethod + def forward(ctx, x, dtype): + return x.to(dtype).to(x.dtype) + + @staticmethod + def backward(ctx, grad): + return grad, None + + +def round_boundary(x, dtype): + """Model a declared storage rounding while retaining an FP32/FP64 VJP.""" + return _RoundStraightThrough.apply(x, dtype) + + +class _ReferenceLinear(torch.autograd.Function): + """Independent FP32 reference with a balanced, blocked input-gradient sum. + + A single CPU BLAS dot over 6144/18432 terms can itself miss the gradient + tolerance under cancellation. The 128-output blocks are independent BLAS + calls, followed by pairwise FP32 addition; no CUDA kernel is used here. + """ + + @staticmethod + def forward(ctx, x, weight, bias): + ctx.save_for_backward(x, weight) + return F.linear(x, weight, bias) + + @staticmethod + def backward(ctx, grad): + x, weight = ctx.saved_tensors + parts = [ + grad[:, start : start + 128] @ weight[start : start + 128] + for start in range(0, weight.shape[0], 128) + ] + while len(parts) > 1: + combined = [parts[i] + parts[i + 1] for i in range(0, len(parts) - 1, 2)] + if len(parts) % 2: + combined.append(parts[-1]) + parts = combined + return parts[0], grad.T @ x, grad.sum(0) + + +def reference_projection(temb, weight, bias, execution_dtype): + act = round_boundary(F.silu(temb), execution_dtype) + return round_boundary(_ReferenceLinear.apply(act, weight, bias), execution_dtype) + + +class QwenImageAdaLNModProjectionReference: + """One stream: SiLU + Linear, returning two (shift, scale, gate) triples.""" + + op_class = "reduction" + kernel_id = "pytorch.qwenimage.adaln_mod_projection.reference" + + def forward(self, temb, weight, bias): + validate_projection(temb, weight, bias, 6) + return F.linear(F.silu(temb), weight, bias) + + __call__ = forward + + def forward_fp32(self, temb, weight, bias): + validate_projection(temb, weight, bias, 6) + return reference_projection(temb.float(), weight.float(), bias.float(), temb.dtype) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 7751d5901..d24902853 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -107,6 +107,15 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): # Variable-length packing (pack-and-pad), [B,S,...] -> [Total_Active,...] PYTORCH_PACK = "rl_engine.kernels.ops.pytorch.packing.pack.NativePackOp" + # Qwen-Image conditioning (RFC #386). + CUDA_QWENIMAGE_ADALN_MOD_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.qwenimage." + "adaln_mod_projection.QwenImageAdaLNModProjectionCudaOp" + ) + PYTORCH_QWENIMAGE_ADALN_MOD_PROJECTION = ( + "rl_engine.reference.model_specific.qwenimage." + "adaln_mod_projection.QwenImageAdaLNModProjectionReference" + ) # 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" @@ -785,6 +794,15 @@ def __init__(self): self._priority_map["npu"]["det_gemm"] = [ OpBackend.ASCEND_DET_GEMM, ] + # CUDA strict operators deliberately have no reference fallback. + for platform in ("rocm", "npu", "musa"): + self._priority_map[platform]["qwenimage_adaln_mod_projection"] = [] + self._priority_map["cuda"]["qwenimage_adaln_mod_projection"] = [ + OpBackend.CUDA_QWENIMAGE_ADALN_MOD_PROJECTION + ] + self._priority_map["cpu"]["qwenimage_adaln_mod_projection"] = [ + OpBackend.PYTORCH_QWENIMAGE_ADALN_MOD_PROJECTION + ] logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") self._adjust_priority_for_hardware() self._adjust_priority_from_env() diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index ca3b7c120..00a33b1bb 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 = { + "qwenimage_adaln_mod_projection": _make_qwenimage_projection_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 = { + "qwenimage_adaln_mod_projection": f"{batch}x3072x18432", "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)}", @@ -395,3 +397,17 @@ def _arg_int(args: argparse.Namespace, name: str, default: int) -> int: def _arg_str(args: argparse.Namespace, name: str, default: str) -> str: return str(getattr(args, name, default)) + + +def _qwenimage_inputs(args, dtype, device, multiplier): + batch, _ = _batch_seq(args) + return { + "temb": _floating_tensor((batch, 3072), args, dtype, device, offset=0), + "weight": _floating_tensor((multiplier * 3072, 3072), args, dtype, device, offset=1) + / 3072**0.5, + "bias": _floating_tensor((multiplier * 3072,), args, dtype, device, offset=2) * 0.01, + } + + +def _make_qwenimage_projection_inputs(args, dtype, device): + return _qwenimage_inputs(args, dtype, device, 6) diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 2323f4787..67a901784 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -32,6 +32,22 @@ def _load_object(path: str) -> Any: OP_SPECS = { + "qwenimage_adaln_mod_projection": OperatorSpec( + name="qwenimage_adaln_mod_projection", + op_class="reduction", + gold_path=( + "rl_engine.reference.model_specific.qwenimage." + "adaln_mod_projection.QwenImageAdaLNModProjectionReference" + ), + gold_method="forward_fp32", + candidate_paths={ + "cuda": ( + "rl_engine.backends.cuda.model_specific.qwenimage." + "adaln_mod_projection.QwenImageAdaLNModProjectionCudaOp" + ) + }, + grad_input_names=("temb", "weight", "bias"), + ), "rms_norm": OperatorSpec( name="rms_norm", op_class="reduction", diff --git a/tests/models/qwenimage/__init__.py b/tests/models/qwenimage/__init__.py new file mode 100644 index 000000000..e2891cf63 --- /dev/null +++ b/tests/models/qwenimage/__init__.py @@ -0,0 +1 @@ +"""tests.models.qwenimage package.""" diff --git a/tests/models/qwenimage/test_adaln.py b/tests/models/qwenimage/test_adaln.py new file mode 100644 index 000000000..fd832fcad --- /dev/null +++ b/tests/models/qwenimage/test_adaln.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.qwenimage.adaln_mod_projection import ( + QwenImageAdaLNModProjectionCudaOp, +) +from rl_engine.contracts.numerical import load_contract +from rl_engine.reference.model_specific.qwenimage.adaln_mod_projection import reference_projection + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +KINDS = ["projection"] + + +def equal_bytes(left, right): + return ( + left.dtype == right.dtype + and left.shape == right.shape + and torch.equal(left.contiguous().view(torch.uint8), right.contiguous().view(torch.uint8)) + ) + + +def inputs(kind, dtype, batch=2, seq=33, device="cuda:0", seed=19): + gen = torch.Generator(device=device).manual_seed(seed) + + def rand(shape, scale=1): + return ( + torch.randn(shape, generator=gen, device=device, dtype=dtype) * scale + ).requires_grad_() + + args = { + "temb": rand((batch, 3072)), + "weight": rand((6 * 3072, 3072), 0.02), + "bias": rand((6 * 3072,), 0.01), + } + return args + + +def operator(kind): + return QwenImageAdaLNModProjectionCudaOp() + + +def random_upstream(y, seed=127): + generator = torch.Generator(device=y.device).manual_seed(seed) + return torch.randn(y.shape, device=y.device, dtype=y.dtype, generator=generator) + + +def evaluate(fn, args, grad=None): + y = fn(**args) + if grad is None: + grad = random_upstream(y) + grads = torch.autograd.grad(y, tuple(args.values()), grad) + return y.detach(), dict(zip(args, grads)) + + +def assert_accurate(actual, expected, dtype, judgment): + spec = load_contract()["judgments"][judgment]["by_op_class"]["reduction"][ + str(dtype).split(".")[-1] + ] + torch.testing.assert_close( + actual.float().cpu(), expected.float().cpu(), atol=spec["atol"], rtol=spec["rtol"] + ) + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("seq", [33, 257]) +def test_cpu_reference_all_gradients(kind, dtype, seq): + args = inputs(kind, dtype, seq=seq) + gold = {k: v.detach().cpu().float().requires_grad_() for k, v in args.items()} + ref = reference_projection + y = operator(kind)(**args) + gy = ref(**gold, execution_dtype=dtype) + grad = random_upstream(y) + grads = torch.autograd.grad(y, tuple(args.values()), grad) + gold_grads = torch.autograd.grad(gy, tuple(gold.values()), grad.cpu().float()) + assert_accurate(y, gy, dtype, "forward_accuracy") + for a, b in zip(grads, gold_grads): + assert_accurate(a, b, dtype, "gradient_accuracy") + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_batch_position_and_repeat_bytes(kind, dtype): + args = inputs(kind, dtype, batch=9) + op = operator(kind) + grad = random_upstream(op(**args)) + y, grads = evaluate(op, args, grad) + y2, grads2 = evaluate(op, args, grad) + assert equal_bytes(y, y2) + for k in grads: + assert equal_bytes(grads[k], grads2[k]), k + for order in ([0], [4], [8], [8, 0], [8, 4, 0], list(reversed(range(9)))): + sub = { + k: v[order].detach().requires_grad_() if k in ("x", "temb") else v + for k, v in args.items() + } + sy, sg = evaluate(op, sub, grad[order]) + assert equal_bytes(sy, y[order]) + for k in ("x", "temb"): + if k in args: + assert equal_bytes(sg[k], grads[k][order]), (k, order) + + +@pytest.mark.parametrize("kind", KINDS) +def test_noncontiguous(kind): + args = inputs(kind, torch.bfloat16) + strided = { + k: torch.stack((v, v), -1)[..., 0].detach().requires_grad_() for k, v in args.items() + } + op = operator(kind) + grad = random_upstream(op(**args)) + a, ag = evaluate(op, args, grad) + b, bg = evaluate(op, strided, grad) + assert equal_bytes(a, b) + for k in ag: + assert equal_bytes(ag[k], bg[k]) + + +@pytest.mark.parametrize("kind", KINDS) +def test_validation_and_readback(kind): + args = inputs(kind, torch.float32) + op = operator(kind) + assert op.implementation_info(**args)["fallback"] is False + for name in args: + bad = dict(args) + bad[name] = args[name].cpu() + with pytest.raises(ValueError): + op(**bad) + bad = dict(args) + bad["temb"] = args["temb"][:0] + with pytest.raises(ValueError): + op(**bad) + bad = dict(args) + bad["weight"] = args["weight"][:-1] + with pytest.raises(ValueError): + op(**bad) + bad = {k: v.half() for k, v in args.items()} + with pytest.raises(TypeError): + op(**bad) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="two CUDA devices required") +@pytest.mark.parametrize("kind", KINDS) +def test_device_one_under_device_zero(kind): + args = inputs(kind, torch.bfloat16, device="cuda:1") + op = operator(kind) + with torch.cuda.device(0): + y, grads = evaluate(op, args) + assert y.device.index == 1 + assert all(g.device.index == 1 for g in grads.values()) + args["bias"] = args["bias"].to("cuda:0") + with pytest.raises(ValueError): + op(**args) + + +@pytest.mark.parametrize("kind", KINDS) +def test_registry_and_gtest(kind): + from argparse import Namespace + + from rl_engine.runtime.registry import KernelRegistry + from rl_engine.validation.operators.op_checks import run_operator_suite + from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case + + name = "qwenimage_adaln_mod_projection" + registry = KernelRegistry() + op = registry.get_op(name, device="cuda:0") + assert isinstance(op, type(operator(kind))) + args = Namespace(op=name, candidate="cuda", arch_key="sm100", batch=1, seq=33, seed=7) + case = make_operator_case(args, torch.bfloat16, torch.device("cuda")) + report = run_operator_suite( + name, candidates=[make_candidate(args)], cases=[case], check_grad=True + ) + assert report.passed, report.to_dict() + + +@pytest.mark.parametrize("kind", KINDS) +def test_missing_extension_fails_closed(kind, monkeypatch): + import rl_engine.backends.cuda.model_specific.qwenimage.adaln_mod_projection as module + + monkeypatch.setattr(module, "_C", None) + with pytest.raises(RuntimeError, match="rebuild"): + operator(kind) + + +@pytest.mark.parametrize("kind", KINDS) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_same_order_parameter_contributions(kind, dtype): + args = inputs(kind, dtype, batch=3) + op = operator(kind) + grad = random_upstream(op(**args)) + _, small_grads = evaluate(op, args, grad) + expanded = { + k: torch.cat([v, v[:2]], 0).detach().requires_grad_() if k in ("x", "temb") else v + for k, v in args.items() + } + padded_grad = torch.cat([grad, torch.zeros_like(grad[:2])], 0) + _, large_grads = evaluate(op, expanded, padded_grad) + for k in ("weight", "bias"): + assert equal_bytes(small_grads[k], large_grads[k]), k + + +def test_projection_chunk_order(): + args = inputs("projection", torch.float32, batch=1) + args["weight"] = torch.zeros_like(args["weight"]) + args["bias"] = torch.arange(6, device="cuda", dtype=torch.float32).repeat_interleave(3072) + out = operator("projection")(**args) + for i, chunk in enumerate(out.chunk(6, -1)): + assert equal_bytes(chunk, torch.full_like(chunk, i)) diff --git a/tests/models/qwenimage/test_parameter_reduction.py b/tests/models/qwenimage/test_parameter_reduction.py new file mode 100644 index 000000000..33e3f73d7 --- /dev/null +++ b/tests/models/qwenimage/test_parameter_reduction.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cancellation regressions for the shared parameter-gradient reduction.""" +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.qwenimage.adaln_mod_projection import require_extension +from rl_engine.backends.extension import _C + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("rows", [255, 256, 257, 2048]) +def test_parameter_sum_retains_small_contributions(dtype, rows): + require_extension() + # Ascending ordinary FP32 summation loses all the unit contributions. + grad = torch.ones(rows, 2, device="cuda", dtype=torch.float32) + grad[0] = 2**26 + grad[-1] = -(2**26) + act = torch.ones(rows, 7, device="cuda", dtype=dtype) + act[:, 1::2] = 1.5 + dw, db = _C.qwenimage_det_linear_backward_weight(grad, act, dtype, True) + golden = grad.double().cpu().T @ act.double().cpu() + golden_bias = grad.double().cpu().sum(0) + assert torch.equal(dw.cpu(), golden.to(dtype)) + assert torch.equal(db.cpu(), golden_bias.to(dtype)) + again = _C.qwenimage_det_linear_backward_weight(grad, act, dtype, True) + for left, right in zip((dw, db), again): + assert torch.equal(left.view(torch.uint8), right.view(torch.uint8)) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_bias_only_entrypoint_matches_joint_reduction(dtype): + grad = torch.randn(257, 17, device="cuda", dtype=torch.float32) + act = torch.randn(257, 3, device="cuda", dtype=dtype) + _, expected = _C.qwenimage_det_linear_backward_weight(grad, act, dtype, True) + actual = _C.qwenimage_det_linear_backward_bias(grad, dtype) + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + + +def test_stale_v1_extension_is_rejected(monkeypatch): + monkeypatch.delattr(_C, "qwenimage_det_linear_backward_bias") + with pytest.raises(RuntimeError, match="rebuild"): + require_extension() diff --git a/tests/models/qwenimage/test_reference.py b/tests/models/qwenimage/test_reference.py new file mode 100644 index 000000000..35c5963c4 --- /dev/null +++ b/tests/models/qwenimage/test_reference.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only checks for the independent reference and strict dispatch gates.""" +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.reference.model_specific.qwenimage.adaln_mod_projection import _ReferenceLinear + + +def test_reference_linear_backward_against_fp64_autograd(): + gen = torch.Generator().manual_seed(2026) + x = torch.randn(2, 7, generator=gen, dtype=torch.float64, requires_grad=True) + weight = torch.randn(385, 7, generator=gen, dtype=torch.float64, requires_grad=True) + bias = torch.randn(385, generator=gen, dtype=torch.float64, requires_grad=True) + assert torch.autograd.gradcheck(_ReferenceLinear.apply, (x, weight, bias), fast_mode=True) + + +def test_fast_math_build_refused(monkeypatch): + from rl_engine.backends.cuda.model_specific.qwenimage import adaln_mod_projection as module + + extension = SimpleNamespace(**{name: lambda: None for name in module._SYMBOLS}) + extension.qwenimage_strict_math_compiled = lambda: False + monkeypatch.setattr(module, "_C", extension) + monkeypatch.setattr(module, "_EXT_AVAILABLE", True) + with pytest.raises(RuntimeError, match="without fast math"): + module.QwenImageAdaLNModProjectionCudaOp() + + +def test_other_platform_slots_fail_closed(): + from rl_engine.runtime.registry import KernelRegistry + + registry = KernelRegistry() + for platform in ("rocm", "npu", "musa"): + assert registry._priority_map[platform]["qwenimage_adaln_mod_projection"] == [] From 442e003166c14a12ff6c55032b2f0e1e391ef689 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Sun, 11 Oct 2026 00:28:55 +0000 Subject: [PATCH 2/6] feat(qwenimage): add CUDA final AdaLN output Add the Qwen-Image Phase A `final_adaln_out` operator (RFC #386): image-stream final modulation, LayerNorm(eps 1e-6) with SiLU + Linear 3072->6144 `(scale, shift)`, reusing the fixed-order projection from `adaln_mod_projection` and adding a fixed-tree FP32 token reduction for the modulation parameter gradients. Register the CUDA and CPU slots (other platforms fail closed), the token-sum binding, the WS1 operator spec/inputs and the build source, and extend the tests, evidence runners and documentation to the final stream. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- benchmarks/models/qwenimage/adaln_evidence.py | 52 +++++++-- .../models/qwenimage/adaln_reuse_bi_audit.py | 100 +++++++++++++++--- .../models/qwenimage/adaln_reuse_candidate.py | 48 +++++++++ .../models/qwenimage/plot_adaln_audit.py | 14 ++- benchmarks/models/qwenimage/streams.py | 3 +- .../models/qwenimage/validate_adaln_repair.py | 12 +-- build_tools/extensions.py | 1 + csrc/bindings/ops.cpp | 2 + .../model_specific/qwenimage/token_sum.cu | 56 ++++++++++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/qwenimage-final-adaln-out.md | 64 +++++++++++ rl_engine/_C.pyi | 1 + .../qwenimage/adaln_mod_projection.py | 1 + .../qwenimage/final_adaln_out.py | 72 +++++++++++++ .../qwenimage/final_adaln_out.py | 79 ++++++++++++++ rl_engine/runtime/registry.py | 15 +++ .../validation/operators/operator_inputs.py | 9 ++ .../validation/operators/operator_specs.py | 16 +++ tests/models/qwenimage/test_adaln.py | 52 +++++++-- tests/models/qwenimage/test_reference.py | 27 ++++- 21 files changed, 589 insertions(+), 37 deletions(-) create mode 100644 csrc/cuda/model_specific/qwenimage/token_sum.cu create mode 100644 docs/operators/qwenimage-final-adaln-out.md create mode 100644 rl_engine/backends/cuda/model_specific/qwenimage/final_adaln_out.py create mode 100644 rl_engine/reference/model_specific/qwenimage/final_adaln_out.py diff --git a/benchmarks/models/qwenimage/adaln_evidence.py b/benchmarks/models/qwenimage/adaln_evidence.py index b1143363b..cf045809a 100644 --- a/benchmarks/models/qwenimage/adaln_evidence.py +++ b/benchmarks/models/qwenimage/adaln_evidence.py @@ -21,12 +21,19 @@ from rl_engine.backends.cuda.model_specific.qwenimage.adaln_mod_projection import ( QwenImageAdaLNModProjectionCudaOp, ) +from rl_engine.backends.cuda.model_specific.qwenimage.final_adaln_out import ( + QwenImageFinalAdaLNOutCudaOp, +) from rl_engine.backends.extension import _C from rl_engine.contracts.numerical import load_contract from rl_engine.reference.model_specific.qwenimage.adaln_mod_projection import ( QwenImageAdaLNModProjectionReference, reference_projection, ) +from rl_engine.reference.model_specific.qwenimage.final_adaln_out import ( + QwenImageFinalAdaLNOutReference, + reference_final, +) CHECKPOINT_REVISION = "75e0b4be04f60ec59a75f475837eced720f823b6" DIFFUSERS_REVISION = "d961a388fd02e4db38d17350c8dd9b8abe642e05" @@ -128,9 +135,9 @@ def compare(actual, expected, dtype, judgment): } -def correctness(candidate, args, kind): +def correctness(candidate, args, kind, fp64_check=False): dtype = args["temb"].dtype - reference = reference_projection + reference = reference_projection if kind == "projection" else reference_final cpu = {k: v.detach().cpu().float().requires_grad_() for k, v in args.items()} y = candidate(**args) gold = reference(**cpu, execution_dtype=dtype) @@ -149,6 +156,20 @@ def correctness(candidate, args, kind): metrics["repeat_bitwise"] = bool( equal_bytes(y, again) and all(equal_bytes(a, b) for a, b in zip(grads, repeat)) ) + if fp64_check and kind == "final" and dtype == torch.float32: + high = {k: v.detach().cpu().double().requires_grad_() for k, v in args.items()} + high_output = reference(**high, execution_dtype=dtype) + high_grads = torch.autograd.grad(high_output, tuple(high.values()), grad.cpu().double()) + metrics["fp64_diagnostic"] = { + "candidate_gradients": { + k: compare(a, b, dtype, "gradient_accuracy") + for k, a, b in zip(args, grads, high_grads) + }, + "cpu_fp32_gradients": { + k: compare(a, b, dtype, "gradient_accuracy") + for k, a, b in zip(args, gold_grads, high_grads) + }, + } del gold, gold_grads, cpu return metrics @@ -209,6 +230,7 @@ def main(): parser.add_argument("--model", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--repeats", type=int, default=10) + parser.add_argument("--fp64-check", action="store_true") opts = parser.parse_args() torch.set_num_threads(8) torch.manual_seed(71) @@ -238,11 +260,23 @@ def main(): for label, stem in STREAMS: weights, manifest = load_weights(opts.model, stem) report["weights"].update(manifest) - kind = "projection" - candidate = QwenImageAdaLNModProjectionCudaOp() - native = QwenImageAdaLNModProjectionReference() + kind = "final" if label == "final" else "projection" + candidate = ( + QwenImageFinalAdaLNOutCudaOp() + if kind == "final" + else QwenImageAdaLNModProjectionCudaOp() + ) + native = ( + QwenImageFinalAdaLNOutReference() + if kind == "final" + else QwenImageAdaLNModProjectionReference() + ) for dtype in (torch.bfloat16, torch.float32): - for batch, seq in [(1, 1), (2, 1), (4, 1), (9, 1)]: + for batch, seq in ( + [(1, 33), (2, 4096), (1, 6889), (1, 6032)] + if kind == "final" + else [(1, 1), (2, 1), (4, 1), (9, 1)] + ): args = { "temb": torch.randn(batch, 3072, device="cuda", dtype=dtype).requires_grad_(), **{ @@ -250,13 +284,17 @@ def main(): for k, v in weights.items() }, } + if kind == "final": + args["x"] = torch.randn( + batch, seq, 3072, device="cuda", dtype=dtype + ).requires_grad_() entry = { "stream": label, "dtype": str(dtype), "batch": batch, "seq": seq, "readback": candidate.implementation_info(**args), - "correctness": correctness(candidate, args, kind), + "correctness": correctness(candidate, args, kind, opts.fp64_check), "timing": benchmark( {"candidate": candidate, "native": native}, args, opts.repeats ), diff --git a/benchmarks/models/qwenimage/adaln_reuse_bi_audit.py b/benchmarks/models/qwenimage/adaln_reuse_bi_audit.py index 712e3772e..db72c2823 100644 --- a/benchmarks/models/qwenimage/adaln_reuse_bi_audit.py +++ b/benchmarks/models/qwenimage/adaln_reuse_bi_audit.py @@ -6,8 +6,9 @@ Run each mode in a fresh process. cublas_policy reproduces vLLM's two cuBLAS workspace settings, not the complete vLLM integration. The vLLM, SGLang and Megatron modes use each project's batch-invariant persistent matmul with the -same FP32 VJP adapter. No backend math flags are changed. FP64 comparisons use -the declared storage boundaries/STE. +same FP32 VJP adapter; vllm_liger and vllm_flashattn also replace the final +LayerNorm with Liger's or flash-attention's Triton LayerNorm. No backend math +flags are changed. FP64 comparisons use the declared storage boundaries/STE. """ import argparse import importlib.metadata @@ -25,6 +26,8 @@ "cublas_policy", "fixed", "vllm_local", + "vllm_liger", + "vllm_flashattn", "sglang_local", "megatron_local", ], @@ -40,12 +43,13 @@ p.add_argument("--seed", type=int, default=0) p.add_argument("--capture-failures", action="store_true") p.add_argument("--megatron-src", type=Path, help="Megatron-LM checkout for megatron_local") +p.add_argument("--flash-attn-src", type=Path, help="flash-attention checkout for vllm_flashattn") a = p.parse_args() if a.mode == "cublas_policy": os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1" -if a.mode in ("vllm_local", "sglang_local", "megatron_local"): +if a.mode in ("vllm_local", "vllm_liger", "vllm_flashattn", "sglang_local", "megatron_local"): os.environ["TRITON_F32_DEFAULT"] = "ieee" if a.mode == "megatron_local": assert a.megatron_src, "--megatron-src is required for megatron_local" @@ -67,10 +71,17 @@ from rl_engine.backends.cuda.model_specific.qwenimage.adaln_mod_projection import ( # noqa: E402 QwenImageAdaLNModProjectionCudaOp, ) +from rl_engine.backends.cuda.model_specific.qwenimage.final_adaln_out import ( # noqa: E402 + QwenImageFinalAdaLNOutCudaOp, +) from rl_engine.reference.model_specific.qwenimage.adaln_mod_projection import ( # noqa: E402 QwenImageAdaLNModProjectionReference, reference_projection, ) +from rl_engine.reference.model_specific.qwenimage.final_adaln_out import ( # noqa: E402 + QwenImageFinalAdaLNOutReference, + reference_final, +) torch.set_num_threads(8) assert not torch.backends.cuda.matmul.allow_tf32, "TF32 must already be disabled" @@ -89,7 +100,10 @@ }, tf32=torch.backends.cuda.matmul.allow_tf32, ), - scope=("sample batch BI; " "shared parameter gradients: same ordered inputs, repeat only"), + scope=( + "sample batch BI; final uses S=3 for large B, long S tested separately; " + "shared parameter gradients: same ordered inputs, repeat only" + ), ) report["packages"] = {} for lib in [ @@ -98,6 +112,7 @@ "diffusers", "vllm", "sglang", + "liger-kernel", "transformer_engine", "megatron-core", ]: @@ -105,6 +120,11 @@ report["packages"][lib] = importlib.metadata.version(lib) except importlib.metadata.PackageNotFoundError: report["packages"][lib] = None +if a.mode == "vllm_flashattn": + assert a.flash_attn_src, "--flash-attn-src is required for vllm_flashattn" + report["packages"]["flash-attention-source"] = subprocess.check_output( + ["git", "-C", str(a.flash_attn_src), "rev-parse", "HEAD"], text=True + ).strip() if a.mode == "megatron_local": report["packages"]["megatron-lm-source"] = subprocess.check_output( ["git", "-C", str(a.megatron_src), "rev-parse", "HEAD"], text=True @@ -149,14 +169,19 @@ def metric(x, y): for stream, stem in STREAMS: + if a.mode in ("vllm_liger", "vllm_flashattn") and stream != "final": + continue weights, manifest = load_weights(a.model, stem) + final = stream == "final" fn = ( - QwenImageAdaLNModProjectionCudaOp() + (QwenImageFinalAdaLNOutCudaOp() if final else QwenImageAdaLNModProjectionCudaOp()) if a.mode == "fixed" - else QwenImageAdaLNModProjectionReference() + else ( + QwenImageFinalAdaLNOutReference() if final else QwenImageAdaLNModProjectionReference() + ) ) - if a.mode in ("vllm_local", "sglang_local", "megatron_local"): - from benchmarks.models.qwenimage.adaln_reuse_candidate import Projection + if a.mode in ("vllm_local", "vllm_liger", "vllm_flashattn", "sglang_local", "megatron_local"): + from benchmarks.models.qwenimage.adaln_reuse_candidate import Final, Projection if a.mode == "sglang_local": import benchmarks.models.qwenimage.adaln_reuse_candidate as adapter @@ -172,7 +197,15 @@ def metric(x, y): ) adapter.matmul_persistent = matmul_persistent - fn = Projection() + fn = Final() if final else Projection() + if a.mode == "vllm_liger": + from benchmarks.models.qwenimage.adaln_reuse_candidate import FinalLiger + + fn = FinalLiger() + if a.mode == "vllm_flashattn": + from benchmarks.models.qwenimage.adaln_reuse_candidate import FinalFlashAttn + + fn = FinalFlashAttn(a.flash_attn_src) for dtype in [torch.bfloat16, torch.float32]: w, b = [weights[k].cuda().to(dtype) for k in ("weight", "bias")] case = dict(stream=stream, dtype=str(dtype), weights=manifest, checks=[]) @@ -185,7 +218,12 @@ def inputs(batch, seq=3, seed=0): weight=w, bias=b, ) - shape = (batch, w.shape[0]) + if final: + inp = { + "x": torch.randn(batch, seq, 3072, device="cuda", dtype=dtype, generator=gen), + **inp, + } + shape = (batch, seq, 3072) if final else (batch, w.shape[0]) upstream = torch.randn(shape, device="cuda", dtype=dtype, generator=gen) return inp, upstream @@ -223,12 +261,20 @@ def inputs(batch, seq=3, seed=0): save() print(a.mode, stream, dtype, "parameters complete", flush=True) elif a.phase in ("accuracy", "accuracy_large"): - shapes = [(256, 1), (2048, 1)] if a.phase == "accuracy_large" else [(1, 1), (3, 1)] + shapes = ( + [(256, 3 if final else 1), (2048, 3 if final else 1)] + if a.phase == "accuracy_large" + else ( + [(1, 3), (3, 3), (1, 4096), (1, 6889), (1, 6032)] if final else [(1, 1), (3, 1)] + ) + ) for batch, seq in shapes: inp, g = inputs(batch, seq, a.seed) out = run(fn, inp, g, True) hi = {k: v.double().cpu().requires_grad_() for k, v in inp.items()} - gold = reference_projection(**hi, execution_dtype=dtype) + gold = (reference_final if final else reference_projection)( + **hi, execution_dtype=dtype + ) gg = torch.autograd.grad(gold, tuple(hi.values()), g.double().cpu()) expected = {"output": gold.detach(), **dict(zip(hi, gg))} case["checks"].append( @@ -249,7 +295,9 @@ def inputs(batch, seq=3, seed=0): ) ) low = {k: v.float().cpu().requires_grad_() for k, v in inp.items()} - low_y = reference_projection(**low, execution_dtype=dtype) + low_y = (reference_final if final else reference_projection)( + **low, execution_dtype=dtype + ) low_g = torch.autograd.grad(low_y, tuple(low.values()), g.float().cpu()) low_expected = {"output": low_y.detach(), **dict(zip(low, low_g))} case["checks"][-1]["cpu_fp32_contract"] = { @@ -296,7 +344,7 @@ def inputs(batch, seq=3, seed=0): inp, g = inputs(3) out = run(fn, inp, g, True) hi = {k: v.double().cpu().requires_grad_() for k, v in inp.items()} - gold = reference_projection(**hi, execution_dtype=dtype) + gold = (reference_final if final else reference_projection)(**hi, execution_dtype=dtype) gg = torch.autograd.grad(gold, tuple(hi.values()), g.double().cpu()) expected = {"output": gold.detach(), **dict(zip(hi, gg))} case["fp64_relative_max_error"] = {k: metric(v, expected[k]) for k, v in out.items()} @@ -304,7 +352,11 @@ def inputs(batch, seq=3, seed=0): k: equal_bytes(v, run(fn, inp, g, True)[k]) for k, v in out.items() } case["performance"] = [] - for batch, seq in [(1, 1), (3, 1), (256, 1), (2048, 1)]: + for batch, seq in ( + [(1, 33), (2, 4096), (1, 6889), (1, 6032), (256, 3), (2048, 3)] + if final + else [(1, 1), (3, 1), (256, 1), (2048, 1)] + ): inp, g = inputs(batch, seq) inp = {k: v.detach().requires_grad_() for k, v in inp.items()} case["performance"].append( @@ -436,4 +488,22 @@ def inputs(batch, seq=3, seed=0): ) save() print(a.mode, stream, dtype, "seed", seed, "BI complete", flush=True) + if final: + for seq in [31, 32, 33, 4096, 6889, 6032]: + inp, g = inputs(2, seq, 127) + whole = run(fn, inp, g) + singles = [ + run(fn, sliced(inp, slice(i, i + 1)), g[i : i + 1]) for i in range(2) + ] + case["checks"].append( + dict( + kind="long_sequence", + sequence=seq, + rows=2, + mismatched_rows=stats( + whole, + {k: torch.cat([s[k] for s in singles]) for k in whole}, + ), + ) + ) save() diff --git a/benchmarks/models/qwenimage/adaln_reuse_candidate.py b/benchmarks/models/qwenimage/adaln_reuse_candidate.py index 6f249dd4f..785740fbb 100644 --- a/benchmarks/models/qwenimage/adaln_reuse_candidate.py +++ b/benchmarks/models/qwenimage/adaln_reuse_candidate.py @@ -5,7 +5,10 @@ No operator registration or global dispatch override is performed. FP32 Triton products require TRITON_F32_DEFAULT=ieee in the audit subprocess. """ +import importlib import os +import sys +import types import torch import torch.nn.functional as F @@ -38,3 +41,48 @@ class Projection: def __call__(self, temb, weight, bias): assert os.environ.get("TRITON_F32_DEFAULT") == "ieee" return _Linear.apply(temb, weight, bias) + + +class Final: + def __call__(self, x, temb, weight, bias): + scale, shift = Projection()(temb, weight, bias).chunk(2, -1) + # Native modulation autograd: explicitly a candidate, not a certified + # replacement for the existing fixed FP32 token-reduction backward. + norm = F.layer_norm(x.float(), (3072,), eps=1e-6).to(x.dtype) + return norm * (1 + scale)[:, None, :] + shift[:, None, :] + + +class FinalLiger: + def __call__(self, x, temb, weight, bias): + from liger_kernel.ops.layer_norm import LigerLayerNormFunction + + scale, shift = Projection()(temb, weight, bias).chunk(2, -1) + gain = torch.ones(3072, device=x.device, dtype=torch.float32) + offset = torch.zeros_like(gain) + norm = LigerLayerNormFunction.apply(x.float(), gain, offset, 1e-6).to(x.dtype) + return norm * (1 + scale)[:, None, :] + shift[:, None, :] + + +class FinalFlashAttn: + """Final modulation with flash-attention's Triton LayerNorm (forward and backward). + + Only ``flash_attn.ops.triton.layer_norm`` is loaded from the checkout; the + package ``__init__`` (which needs the compiled attention extension) is skipped. + """ + + def __init__(self, source): + if "flash_attn" not in sys.modules: + package = types.ModuleType("flash_attn") + package.__path__ = [str(source / "flash_attn")] + sys.modules["flash_attn"] = package + self.layer_norm_fn = importlib.import_module( + "flash_attn.ops.triton.layer_norm" + ).layer_norm_fn + + def __call__(self, x, temb, weight, bias): + scale, shift = Projection()(temb, weight, bias).chunk(2, -1) + # Parameter-free LayerNorm: constant unit gain and zero offset, as for Liger. + gain = torch.ones(3072, device=x.device, dtype=torch.float32) + offset = torch.zeros_like(gain) + norm = self.layer_norm_fn(x.float(), gain, offset, eps=1e-6).to(x.dtype) + return norm * (1 + scale)[:, None, :] + shift[:, None, :] diff --git a/benchmarks/models/qwenimage/plot_adaln_audit.py b/benchmarks/models/qwenimage/plot_adaln_audit.py index 58e7e40da..ed7157e29 100644 --- a/benchmarks/models/qwenimage/plot_adaln_audit.py +++ b/benchmarks/models/qwenimage/plot_adaln_audit.py @@ -17,7 +17,16 @@ p = argparse.ArgumentParser() p.add_argument("directory", type=Path) a = p.parse_args() -modes = ["native", "cublas_policy", "fixed", "vllm_local", "sglang_local", "megatron_local"] +modes = [ + "native", + "cublas_policy", + "fixed", + "vllm_local", + "sglang_local", + "megatron_local", + "vllm_liger", + "vllm_flashattn", +] labels = [ "Native", "cuBLAS policy", @@ -25,6 +34,8 @@ "vLLM GEMM adapter", "SGLang GEMM adapter", "Megatron GEMM adapter", + "vLLM + Liger adapter", + "vLLM + flash-attn LN adapter", ] keys = [(s, "torch." + d) for d in ["bfloat16", "float32"] for s, _ in STREAMS] reports = {} @@ -180,6 +191,7 @@ lines += [ "", "Latency includes Python launch gaps; separate processes.", + "Small-batch final: S=33; large-batch final: S=3.", "Raw samples include fwd, bwd, total and peak memory.", "FP64 uses explicit storage casts and straight-through VJP.", "Strict elementwise tolerance results: *-accuracy.json.", diff --git a/benchmarks/models/qwenimage/streams.py b/benchmarks/models/qwenimage/streams.py index 4167ce8cb..7a1b147d8 100644 --- a/benchmarks/models/qwenimage/streams.py +++ b/benchmarks/models/qwenimage/streams.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 """Checkpoint streams shared by the evidence runners, validator and plots.""" -# Checkpoint tensors exercised per stream: block-0 image and text projections. +# Checkpoint tensors exercised per stream: block-0 projections and final modulation. STREAMS = ( ("image", "transformer_blocks.0.img_mod.1"), ("text", "transformer_blocks.0.txt_mod.1"), + ("final", "norm_out.linear"), ) diff --git a/benchmarks/models/qwenimage/validate_adaln_repair.py b/benchmarks/models/qwenimage/validate_adaln_repair.py index c3309f5d3..2d23b1b91 100644 --- a/benchmarks/models/qwenimage/validate_adaln_repair.py +++ b/benchmarks/models/qwenimage/validate_adaln_repair.py @@ -7,13 +7,13 @@ from benchmarks.models.qwenimage.streams import STREAMS -# Checks per (stream, dtype) case. +# Checks per (stream, dtype) case; final modulation adds long-sequence cases. CHECKS_PER_CASE = { - "accuracy": {"image": 2, "text": 2}, - "accuracy_large": {"image": 2, "text": 2}, - "bi": {"image": 402, "text": 402}, - "workload": {"image": 15, "text": 15}, - "parameters": {"image": 21, "text": 21}, + "accuracy": {"image": 2, "text": 2, "final": 5}, + "accuracy_large": {"image": 2, "text": 2, "final": 2}, + "bi": {"image": 402, "text": 402, "final": 408}, + "workload": {"image": 15, "text": 15, "final": 15}, + "parameters": {"image": 21, "text": 21, "final": 21}, } p = argparse.ArgumentParser() diff --git a/build_tools/extensions.py b/build_tools/extensions.py index bce062253..56b56b163 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -239,6 +239,7 @@ def get_extensions(): cuda_sources.extend( [ "csrc/cuda/model_specific/qwenimage/det_linear.cu", + "csrc/cuda/model_specific/qwenimage/token_sum.cu", ] ) # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index f29fb5abb..a8b30db38 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -11,6 +11,7 @@ torch::Tensor qwenimage_det_linear_backward_bias(torch::Tensor, c10::ScalarType) std::vector qwenimage_det_linear_forward(torch::Tensor, torch::Tensor, c10::optional, int64_t, bool); torch::Tensor qwenimage_det_linear_backward_input(torch::Tensor, torch::Tensor, c10::ScalarType); std::vector qwenimage_det_linear_backward_weight(torch::Tensor, torch::Tensor, c10::ScalarType, bool); +torch::Tensor qwenimage_token_sum(torch::Tensor); #endif // Fused LogP Declarations @@ -515,6 +516,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("qwenimage_det_linear_forward", &qwenimage_det_linear_forward); m.def("qwenimage_det_linear_backward_input", &qwenimage_det_linear_backward_input); m.def("qwenimage_det_linear_backward_weight", &qwenimage_det_linear_backward_weight); + m.def("qwenimage_token_sum", &qwenimage_token_sum); m.def("qwenimage_strict_math_compiled", &qwenimage_strict_math_compiled); m.def("qwenimage_det_linear_backward_bias", &qwenimage_det_linear_backward_bias); #endif diff --git a/csrc/cuda/model_specific/qwenimage/token_sum.cu b/csrc/cuda/model_specific/qwenimage/token_sum.cu new file mode 100644 index 000000000..8311aedfd --- /dev/null +++ b/csrc/cuda/model_specific/qwenimage/token_sum.cu @@ -0,0 +1,56 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +#include +#include +#include +#include + +namespace { +// Recursive fixed 32-token balanced trees. Each +// sample/channel has its own sequence; batching never changes the tree. +__global__ void token_partials(const float* x, float* p, int64_t seq, int64_t hidden, + int64_t tiles) { + const int64_t h = blockIdx.x * blockDim.x + threadIdx.x; + const int64_t tile = blockIdx.y; + const int64_t batch = blockIdx.z; + if (h >= hidden) return; + float v[32]; +#pragma unroll + for (int i = 0; i < 32; ++i) { + int64_t row = tile * 32 + i; + v[i] = row < seq ? x[(batch * seq + row) * hidden + h] : 0.0f; + } +#pragma unroll + for (int stride = 1; stride < 32; stride *= 2) { +#pragma unroll + for (int i = 0; i < 32; i += 2 * stride) v[i] = __fadd_rn(v[i], v[i + stride]); + } + p[(batch * tiles + tile) * hidden + h] = v[0]; +} + +} + +torch::Tensor qwenimage_token_sum(torch::Tensor x) { +#ifdef __FAST_MATH__ + TORCH_CHECK(false, "Qwen-Image requires a build without fast math"); +#endif + TORCH_CHECK(x.is_cuda() && x.scalar_type() == at::kFloat && x.dim() == 3 && x.is_contiguous(), + "token_sum requires contiguous CUDA float32 [B,S,H]"); + const auto b = x.size(0), s = x.size(1), h = x.size(2); + TORCH_CHECK(b > 0 && b <= 65535 && s > 0 && s <= 32 * 65535 && h > 0, + "token_sum geometry out of bounds"); + const c10::cuda::CUDAGuard guard(x.device()); + auto current = x; + auto stream = at::cuda::getCurrentCUDAStream(); + int64_t rows = s; + while (rows > 1) { + const auto tiles = (rows + 31) / 32; + auto partials = torch::empty({b, tiles, h}, x.options()); + token_partials<<>>( + current.data_ptr(), partials.data_ptr(), rows, h, tiles); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + current = partials; + rows = tiles; + } + return current.view({b, h}); +} diff --git a/docs/.nav.yml b/docs/.nav.yml index 2e2aa6e71..aa184cdb5 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -17,6 +17,7 @@ nav: - Operators: - operators/README.md - operators/qwenimage-adaln-mod-projection.md + - operators/qwenimage-final-adaln-out.md - operators/activation.md - operators/attention.md - operators/fused-logp.md diff --git a/docs/operators/README.md b/docs/operators/README.md index c3de2fdb5..0dadfb4b2 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -33,4 +33,5 @@ Every operator page should include: - [Sampling](sampling.md) - [Token Embedding](embedding.md) - [Qwen-Image AdaLN Modulation Projection](qwenimage-adaln-mod-projection.md) +- [Qwen-Image Final AdaLN Output](qwenimage-final-adaln-out.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/qwenimage-final-adaln-out.md b/docs/operators/qwenimage-final-adaln-out.md new file mode 100644 index 000000000..47f80c0dc --- /dev/null +++ b/docs/operators/qwenimage-final-adaln-out.md @@ -0,0 +1,64 @@ +# Qwen-Image final AdaLN output (CUDA) + +> Status: validated on B200 (SM100). Accuracy, batch-invariance, reuse-survey and timing evidence, with reproduction commands, is in [B200 evidence](../usage/evidence/qwenimage-final-adaln-out-b200/README.md). + +RFC #386 `final_adaln_out`, registered as `qwenimage_final_adaln_out`. This covers +`norm_out` only; the subsequent 3072→64 `proj_out` belongs to `latent_proj_gemm`. + +## Interface and formula + +```python +from rl_engine.runtime.registry import KernelRegistry + +op = KernelRegistry().get_op("qwenimage_final_adaln_out", device=x.device) +y = op(x=x, temb=temb, weight=weight, bias=bias) +trace = op.implementation_info(x=x, temb=temb, weight=weight, bias=bias) +``` + +- `x`: `[B,S,3072]`; `temb`: `[B,3072]`. +- `weight`: `[6144,3072]`; `bias`: `[6144]`. +- Output: `[B,S,3072]`. B and S must be positive. +- All inputs share a CUDA device and BF16/FP32 dtype. Gradients are provided for + all four inputs. First-order autograd only. + +The formula is `LayerNorm(x) * (1 + scale) + shift`, with `scale, shift` from +`Linear(SiLU(temb)).chunk(2)`. LayerNorm uses eps=1e-6 and **no learned affine**; +H3's RMSNorm and `(shift, scale)` order are not interchangeable with this op. + +## Rounding and deterministic gradients + +Projection uses the [shared conditioning implementation](qwenimage-adaln-mod-projection.md). +The final operator is deliberately unfused: native CUDA LayerNorm computes in +FP32, then stores the normalized values in the execution dtype. Separate +execution-dtype stores occur after `1 + scale`, multiplication, and addition. +Those BF16 stores are part of the forward contract, and the reference models +each one. This is not a claim of bitwise equality to arbitrary Diffusers/cuBLAS +versions; strict training and rollout must both bind this same operator. + +Backward uses FP32 intermediates and straight-through derivatives at storage +rounding boundaries. Native LayerNorm backward computes `dx`; scale/shift +contributions reduce over S with recursively balanced 32-token FP32 trees. +The independent CPU reference uses sequential Neumaier summation over tokens +in FP32 (or FP64 for diagnostics). This also retains small contributions under +cancellation, independently of the unchanged CUDA token tree. There are no atomics +or collective reductions. The projection receives FP32 +table gradients before the final casts to leaf dtypes. Batch size/position does +not alter a sample's token reduction tree. + +Readback identifies the native LayerNorm and elementwise stages, CUDA projection +kernel, recursive token-sum tree, dtype, device, and lack of fallback. No fused +kernel or speedup over the native composition is claimed. + +## Supported scope and evidence + +Standard Phase A image stream only; no edit `modulate_index`, additional +conditioning, distributed execution, or output projection. NVIDIA SM80+ code, +with measured validation limited to B200/SM100. Missing/fast-math builds and +unsupported inputs fail explicitly. + +Tests cover token counts 4096, 6889 and 6032, short cases, tile boundaries, +noncontiguous inputs, forward/backward accuracy, sample batching, repeatability, +parameter contribution ordering, and device isolation. Reports retain existing +tolerances and include FP64 diagnostics of the CPU FP32 reference. + +See [B200 evidence and reproduction](../usage/evidence/qwenimage-final-adaln-out-b200/README.md). diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 4fb9ab443..f3e3589c4 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -345,5 +345,6 @@ def qwenimage_det_linear_backward_weight( out_dtype: torch.dtype, with_bias: bool, ) -> list[torch.Tensor]: ... +def qwenimage_token_sum(x: torch.Tensor) -> torch.Tensor: ... def qwenimage_strict_math_compiled() -> bool: ... def qwenimage_det_linear_backward_bias(grad: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: ... diff --git a/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py b/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py index 273fdc971..57a902cbe 100644 --- a/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py +++ b/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py @@ -18,6 +18,7 @@ "qwenimage_det_linear_backward_weight", # Introduced with v2; reject stale v1 extensions with uncompensated dW/db. "qwenimage_det_linear_backward_bias", + "qwenimage_token_sum", "qwenimage_strict_math_compiled", ) diff --git a/rl_engine/backends/cuda/model_specific/qwenimage/final_adaln_out.py b/rl_engine/backends/cuda/model_specific/qwenimage/final_adaln_out.py new file mode 100644 index 000000000..8507cb323 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/qwenimage/final_adaln_out.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Unfused strict final modulation using native CUDA LayerNorm and fixed sums.""" + +import torch +from torch.autograd.function import once_differentiable + +from rl_engine.backends.extension import _C +from rl_engine.reference.model_specific.qwenimage.final_adaln_out import EPS, HIDDEN, validate_final + +from .adaln_mod_projection import ( + implementation_info, + projection_backward, + projection_forward, + require_device, + require_extension, +) + + +class _Final(torch.autograd.Function): + @staticmethod + def forward(ctx, x, temb, weight, bias): + table, act = projection_forward(temb, weight, bias) + scale, shift = table.chunk(2, -1) + xf = x.float().contiguous() + norm, mean, rstd = torch.native_layer_norm(xf, (HIDDEN,), None, None, EPS) + norm = norm.to(x.dtype) + gain = 1 + scale + out = norm * gain[:, None, :] + shift[:, None, :] + ctx.save_for_backward(xf, mean, rstd, norm, gain, temb, act, weight) + ctx.dtype = x.dtype + return out + + @staticmethod + @once_differentiable + def backward(ctx, grad): + xf, mean, rstd, norm, gain, temb, act, weight = ctx.saved_tensors + gf = grad.float().contiguous() + dn = (gf * gain.float()[:, None, :]).contiguous() + dx = torch.ops.aten.native_layer_norm_backward( + dn, xf, (HIDDEN,), mean, rstd, None, None, (True, False, False) + )[0] + ds = _C.qwenimage_token_sum((gf * norm.float()).contiguous()) + db = _C.qwenimage_token_sum(gf) + dt, dw, dbias = projection_backward(torch.cat((ds, db), -1), temb, act, weight) + return dx.to(ctx.dtype), dt, dw, dbias + + +class QwenImageFinalAdaLNOutCudaOp: + op_class = "reduction" + kernel_id = "qwenimage_det_linear+aten.native_layer_norm+aten.mul+aten.add" + + def __init__(self): + require_extension() + + def forward(self, x, temb, weight, bias): + validate_final(x, temb, weight, bias) + require_device(x) + return _Final.apply(x, temb, weight, bias) + + __call__ = forward + + def implementation_info(self, x, temb, weight, bias): + validate_final(x, temb, weight, bias) + info = implementation_info("qwenimage_final_adaln_out", x) + info.update( + norm="aten.native_layer_norm[fp32,eps=1e-6]", + token_gradient_tree="recursive_32_token_balanced", + cast_boundaries="silu,projection,norm,1+scale,product,output", + fused=False, + ) + return info diff --git a/rl_engine/reference/model_specific/qwenimage/final_adaln_out.py b/rl_engine/reference/model_specific/qwenimage/final_adaln_out.py new file mode 100644 index 000000000..45a399877 --- /dev/null +++ b/rl_engine/reference/model_specific/qwenimage/final_adaln_out.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Qwen-Image final LayerNorm modulation; excludes the output projection.""" + +import torch +import torch.nn.functional as F + +from .adaln_mod_projection import HIDDEN, reference_projection, round_boundary, validate_projection + +EPS = 1e-6 + + +def validate_final(x, temb, weight, bias): + validate_projection(temb, weight, bias, 2) + if x.ndim != 3 or x.shape[0] != temb.shape[0] or x.shape[1] < 1 or x.shape[2] != HIDDEN: + raise ValueError("x must be [B,S,3072], S > 0, with the temb batch size") + if x.dtype != temb.dtype: + raise TypeError("x and temb must share a dtype") + if x.device != temb.device: + raise ValueError("all inputs must be on the same device") + + +def _reference_token_sum(values): + # Independent sequential Neumaier sum, not the CUDA 32-token tree. + # Keep the requested FP32/FP64 arithmetic; long cancellation-heavy sums + # otherwise make the reference itself an unreliable tolerance oracle. + total = torch.zeros_like(values[:, 0]) + correction = torch.zeros_like(total) + for value in values.unbind(1): + next_total = total + value + error = torch.where( + total.abs() >= value.abs(), + (total - next_total) + value, + (value - next_total) + total, + ) + correction = correction + error + total = next_total + return total + correction + + +class _ReferenceModulate(torch.autograd.Function): + @staticmethod + def forward(ctx, norm, scale, shift, dtype): + gain = round_boundary(1 + scale, dtype) + ctx.save_for_backward(norm, gain) + product = round_boundary(norm * gain[:, None, :], dtype) + return round_boundary(product + shift[:, None, :], dtype) + + @staticmethod + def backward(ctx, grad): + norm, gain = ctx.saved_tensors + return ( + grad * gain[:, None, :], + _reference_token_sum(grad * norm), + _reference_token_sum(grad), + None, + ) + + +def reference_final(x, temb, weight, bias, execution_dtype): + scale, shift = reference_projection(temb, weight, bias, execution_dtype).chunk(2, -1) + norm = round_boundary(F.layer_norm(x, (HIDDEN,), eps=EPS), execution_dtype) + return _ReferenceModulate.apply(norm, scale, shift, execution_dtype) + + +class QwenImageFinalAdaLNOutReference: + op_class = "reduction" + kernel_id = "pytorch.qwenimage.final_adaln_out.reference" + + def forward(self, x, temb, weight, bias): + validate_final(x, temb, weight, bias) + scale, shift = F.linear(F.silu(temb), weight, bias).chunk(2, -1) + return F.layer_norm(x, (HIDDEN,), eps=EPS) * (1 + scale)[:, None, :] + shift[:, None, :] + + __call__ = forward + + def forward_fp32(self, x, temb, weight, bias): + validate_final(x, temb, weight, bias) + return reference_final(x.float(), temb.float(), weight.float(), bias.float(), x.dtype) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index d24902853..db1815db6 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -112,10 +112,18 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): "rl_engine.backends.cuda.model_specific.qwenimage." "adaln_mod_projection.QwenImageAdaLNModProjectionCudaOp" ) + CUDA_QWENIMAGE_FINAL_ADALN_OUT = ( + "rl_engine.backends.cuda.model_specific.qwenimage." + "final_adaln_out.QwenImageFinalAdaLNOutCudaOp" + ) PYTORCH_QWENIMAGE_ADALN_MOD_PROJECTION = ( "rl_engine.reference.model_specific.qwenimage." "adaln_mod_projection.QwenImageAdaLNModProjectionReference" ) + PYTORCH_QWENIMAGE_FINAL_ADALN_OUT = ( + "rl_engine.reference.model_specific.qwenimage." + "final_adaln_out.QwenImageFinalAdaLNOutReference" + ) # 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" @@ -797,12 +805,19 @@ def __init__(self): # CUDA strict operators deliberately have no reference fallback. for platform in ("rocm", "npu", "musa"): self._priority_map[platform]["qwenimage_adaln_mod_projection"] = [] + self._priority_map[platform]["qwenimage_final_adaln_out"] = [] self._priority_map["cuda"]["qwenimage_adaln_mod_projection"] = [ OpBackend.CUDA_QWENIMAGE_ADALN_MOD_PROJECTION ] + self._priority_map["cuda"]["qwenimage_final_adaln_out"] = [ + OpBackend.CUDA_QWENIMAGE_FINAL_ADALN_OUT + ] self._priority_map["cpu"]["qwenimage_adaln_mod_projection"] = [ OpBackend.PYTORCH_QWENIMAGE_ADALN_MOD_PROJECTION ] + self._priority_map["cpu"]["qwenimage_final_adaln_out"] = [ + OpBackend.PYTORCH_QWENIMAGE_FINAL_ADALN_OUT + ] logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") self._adjust_priority_for_hardware() self._adjust_priority_from_env() diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index 00a33b1bb..fe93717e2 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -26,6 +26,7 @@ def make_operator_inputs( ) -> dict[str, Any]: builders = { "qwenimage_adaln_mod_projection": _make_qwenimage_projection_inputs, + "qwenimage_final_adaln_out": _make_qwenimage_final_inputs, "rms_norm": _make_rms_norm_inputs, "qk_norm": _make_qk_norm_inputs, "pack": _make_pack_inputs, @@ -55,6 +56,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: vocab = _arg_int(args, "vocab", DEFAULT_VOCAB) names = { "qwenimage_adaln_mod_projection": f"{batch}x3072x18432", + "qwenimage_final_adaln_out": f"{batch}x{seq}x3072", "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)}", @@ -411,3 +413,10 @@ def _qwenimage_inputs(args, dtype, device, multiplier): def _make_qwenimage_projection_inputs(args, dtype, device): return _qwenimage_inputs(args, dtype, device, 6) + + +def _make_qwenimage_final_inputs(args, dtype, device): + inputs = _qwenimage_inputs(args, dtype, device, 2) + batch, seq = _batch_seq(args) + inputs["x"] = _floating_tensor((batch, seq, 3072), args, dtype, device, offset=3) + return inputs diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 67a901784..f4b93428d 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -48,6 +48,22 @@ def _load_object(path: str) -> Any: }, grad_input_names=("temb", "weight", "bias"), ), + "qwenimage_final_adaln_out": OperatorSpec( + name="qwenimage_final_adaln_out", + op_class="reduction", + gold_path=( + "rl_engine.reference.model_specific.qwenimage." + "final_adaln_out.QwenImageFinalAdaLNOutReference" + ), + gold_method="forward_fp32", + candidate_paths={ + "cuda": ( + "rl_engine.backends.cuda.model_specific.qwenimage." + "final_adaln_out.QwenImageFinalAdaLNOutCudaOp" + ) + }, + grad_input_names=("x", "temb", "weight", "bias"), + ), "rms_norm": OperatorSpec( name="rms_norm", op_class="reduction", diff --git a/tests/models/qwenimage/test_adaln.py b/tests/models/qwenimage/test_adaln.py index fd832fcad..8066a901d 100644 --- a/tests/models/qwenimage/test_adaln.py +++ b/tests/models/qwenimage/test_adaln.py @@ -5,11 +5,15 @@ from rl_engine.backends.cuda.model_specific.qwenimage.adaln_mod_projection import ( QwenImageAdaLNModProjectionCudaOp, ) +from rl_engine.backends.cuda.model_specific.qwenimage.final_adaln_out import ( + QwenImageFinalAdaLNOutCudaOp, +) from rl_engine.contracts.numerical import load_contract from rl_engine.reference.model_specific.qwenimage.adaln_mod_projection import reference_projection +from rl_engine.reference.model_specific.qwenimage.final_adaln_out import reference_final pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -KINDS = ["projection"] +KINDS = ["projection", "final"] def equal_bytes(left, right): @@ -30,14 +34,20 @@ def rand(shape, scale=1): args = { "temb": rand((batch, 3072)), - "weight": rand((6 * 3072, 3072), 0.02), - "bias": rand((6 * 3072,), 0.01), + "weight": rand(((6 if kind == "projection" else 2) * 3072, 3072), 0.02), + "bias": rand(((6 if kind == "projection" else 2) * 3072,), 0.01), } + if kind == "final": + args["x"] = rand((batch, seq, 3072)) return args def operator(kind): - return QwenImageAdaLNModProjectionCudaOp() + return ( + QwenImageAdaLNModProjectionCudaOp() + if kind == "projection" + else QwenImageFinalAdaLNOutCudaOp() + ) def random_upstream(y, seed=127): @@ -68,7 +78,7 @@ def assert_accurate(actual, expected, dtype, judgment): def test_cpu_reference_all_gradients(kind, dtype, seq): args = inputs(kind, dtype, seq=seq) gold = {k: v.detach().cpu().float().requires_grad_() for k, v in args.items()} - ref = reference_projection + ref = reference_projection if kind == "projection" else reference_final y = operator(kind)(**args) gy = ref(**gold, execution_dtype=dtype) grad = random_upstream(y) @@ -102,6 +112,25 @@ def test_batch_position_and_repeat_bytes(kind, dtype): assert equal_bytes(sg[k], grads[k][order]), (k, order) +@pytest.mark.parametrize("seq", [4096, 6889, 6032]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("seed", [0, 127, 2026]) +def test_long_sequence_gradients(seq, dtype, seed): + args = inputs("final", dtype, batch=1, seq=seq) + gold = {k: v.detach().cpu().float().requires_grad_() for k, v in args.items()} + y = operator("final")(**args) + gy = reference_final(**gold, execution_dtype=dtype) + grad = random_upstream(y, seed) + a = torch.autograd.grad(y, tuple(args.values()), grad) + b = torch.autograd.grad(gy, tuple(gold.values()), grad.float().cpu()) + assert_accurate(y, gy, dtype, "forward_accuracy") + for k, av, bv in zip(args, a, b): + try: + assert_accurate(av, bv, dtype, "gradient_accuracy") + except AssertionError as exc: + raise AssertionError(k) from exc + + @pytest.mark.parametrize("kind", KINDS) def test_noncontiguous(kind): args = inputs(kind, torch.bfloat16) @@ -162,7 +191,7 @@ def test_registry_and_gtest(kind): from rl_engine.validation.operators.op_checks import run_operator_suite from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case - name = "qwenimage_adaln_mod_projection" + name = "qwenimage_adaln_mod_projection" if kind == "projection" else "qwenimage_final_adaln_out" registry = KernelRegistry() op = registry.get_op(name, device="cuda:0") assert isinstance(op, type(operator(kind))) @@ -207,3 +236,14 @@ def test_projection_chunk_order(): out = operator("projection")(**args) for i, chunk in enumerate(out.chunk(6, -1)): assert equal_bytes(chunk, torch.full_like(chunk, i)) + + +def test_final_scale_then_shift_and_layernorm(): + args = inputs("final", torch.float32, batch=1) + args["weight"] = torch.zeros_like(args["weight"]) + args["bias"] = torch.cat( + [torch.full((3072,), 2.0, device="cuda"), torch.full((3072,), 7.0, device="cuda")] + ) + args["x"] = torch.full_like(args["x"], 13.0) + out = operator("final")(**args) + assert equal_bytes(out, torch.full_like(out, 7.0)) diff --git a/tests/models/qwenimage/test_reference.py b/tests/models/qwenimage/test_reference.py index 35c5963c4..151b33f24 100644 --- a/tests/models/qwenimage/test_reference.py +++ b/tests/models/qwenimage/test_reference.py @@ -32,4 +32,29 @@ def test_other_platform_slots_fail_closed(): registry = KernelRegistry() for platform in ("rocm", "npu", "musa"): - assert registry._priority_map[platform]["qwenimage_adaln_mod_projection"] == [] + for op in ("qwenimage_adaln_mod_projection", "qwenimage_final_adaln_out"): + assert registry._priority_map[platform][op] == [] + + +def test_reference_modulation_backward_against_fp64_autograd(): + from rl_engine.reference.model_specific.qwenimage.final_adaln_out import _ReferenceModulate + + gen = torch.Generator().manual_seed(2026) + norm = torch.randn(2, 257, 7, generator=gen, dtype=torch.float64, requires_grad=True) + scale = torch.randn(2, 7, generator=gen, dtype=torch.float64, requires_grad=True) + shift = torch.randn(2, 7, generator=gen, dtype=torch.float64, requires_grad=True) + assert torch.autograd.gradcheck( + lambda n, s, b: _ReferenceModulate.apply(n, s, b, torch.float64), + (norm, scale, shift), + fast_mode=True, + ) + + +def test_fp32_reference_retains_small_token_contributions(): + from rl_engine.reference.model_specific.qwenimage.final_adaln_out import _reference_token_sum + + values = torch.ones(1, 6032, 3, dtype=torch.float32) + values[:, 0] = 2**26 + values[:, -1] = -(2**26) + golden = values.double().sum(1).float() + assert torch.equal(_reference_token_sum(values), golden) From a73c66c5e2db7ec844e30ac763ce85d410179259 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:32 +0000 Subject: [PATCH 3/6] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/qwenimage-adaln-mod-projection.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/operators/qwenimage-adaln-mod-projection.md b/docs/operators/qwenimage-adaln-mod-projection.md index 76cea7462..fd200e50c 100644 --- a/docs/operators/qwenimage-adaln-mod-projection.md +++ b/docs/operators/qwenimage-adaln-mod-projection.md @@ -1,6 +1,6 @@ # Qwen-Image AdaLN modulation projection (CUDA) -> Status: validated on B200 (SM100). Accuracy, batch-invariance, reuse-survey and timing evidence, with reproduction commands, is in [B200 evidence](../usage/evidence/qwenimage-adaln-mod-projection-b200/README.md). +> Status: validated on B200 (SM100). Accuracy, batch-invariance, reuse-survey and timing evidence, with reproduction commands, is in [B200 evidence](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr537/docs/usage/evidence/qwenimage-adaln-mod-projection-b200/README.md). RFC #386 `adaln_mod_projection`, registered as `qwenimage_adaln_mod_projection` to keep the H3 and Qwen-Image contracts separate. This is a Phase A single-GPU @@ -68,7 +68,7 @@ This reference does not call the candidate CUDA kernel. ## Validation and reproduction -See the [B200 evidence and review notes](../usage/evidence/qwenimage-adaln-mod-projection-b200/README.md). +See the [B200 evidence and review notes](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr537/docs/usage/evidence/qwenimage-adaln-mod-projection-b200/README.md). Reports use the existing `reduction` forward/gradient tolerance rows, without private overrides. Strict replay and batch invariance are checked with `torch.equal` on contiguous raw-byte views, separately from agreement with the CPU reference. From 2d9680d7b8b9c0f7d150bd7cb296cf496652963a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:58:42 +0000 Subject: [PATCH 4/6] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/qwenimage-adaln-mod-projection.md | 4 ++-- docs/operators/qwenimage-final-adaln-out.md | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/operators/qwenimage-adaln-mod-projection.md b/docs/operators/qwenimage-adaln-mod-projection.md index 76cea7462..fd200e50c 100644 --- a/docs/operators/qwenimage-adaln-mod-projection.md +++ b/docs/operators/qwenimage-adaln-mod-projection.md @@ -1,6 +1,6 @@ # Qwen-Image AdaLN modulation projection (CUDA) -> Status: validated on B200 (SM100). Accuracy, batch-invariance, reuse-survey and timing evidence, with reproduction commands, is in [B200 evidence](../usage/evidence/qwenimage-adaln-mod-projection-b200/README.md). +> Status: validated on B200 (SM100). Accuracy, batch-invariance, reuse-survey and timing evidence, with reproduction commands, is in [B200 evidence](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr537/docs/usage/evidence/qwenimage-adaln-mod-projection-b200/README.md). RFC #386 `adaln_mod_projection`, registered as `qwenimage_adaln_mod_projection` to keep the H3 and Qwen-Image contracts separate. This is a Phase A single-GPU @@ -68,7 +68,7 @@ This reference does not call the candidate CUDA kernel. ## Validation and reproduction -See the [B200 evidence and review notes](../usage/evidence/qwenimage-adaln-mod-projection-b200/README.md). +See the [B200 evidence and review notes](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr537/docs/usage/evidence/qwenimage-adaln-mod-projection-b200/README.md). Reports use the existing `reduction` forward/gradient tolerance rows, without private overrides. Strict replay and batch invariance are checked with `torch.equal` on contiguous raw-byte views, separately from agreement with the CPU reference. diff --git a/docs/operators/qwenimage-final-adaln-out.md b/docs/operators/qwenimage-final-adaln-out.md index 47f80c0dc..56d6203a4 100644 --- a/docs/operators/qwenimage-final-adaln-out.md +++ b/docs/operators/qwenimage-final-adaln-out.md @@ -1,6 +1,6 @@ # Qwen-Image final AdaLN output (CUDA) -> Status: validated on B200 (SM100). Accuracy, batch-invariance, reuse-survey and timing evidence, with reproduction commands, is in [B200 evidence](../usage/evidence/qwenimage-final-adaln-out-b200/README.md). +> Status: validated on B200 (SM100). Accuracy, batch-invariance, reuse-survey and timing evidence, with reproduction commands, is in [B200 evidence](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr538/docs/usage/evidence/qwenimage-final-adaln-out-b200/README.md). RFC #386 `final_adaln_out`, registered as `qwenimage_final_adaln_out`. This covers `norm_out` only; the subsequent 3072→64 `proj_out` belongs to `latent_proj_gemm`. @@ -61,4 +61,4 @@ noncontiguous inputs, forward/backward accuracy, sample batching, repeatability, parameter contribution ordering, and device isolation. Reports retain existing tolerances and include FP64 diagnostics of the CPU FP32 reference. -See [B200 evidence and reproduction](../usage/evidence/qwenimage-final-adaln-out-b200/README.md). +See [B200 evidence and reproduction](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr538/docs/usage/evidence/qwenimage-final-adaln-out-b200/README.md). From 274b6c5228e1689d3da46496862b986f7dd13907 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:15:17 +0000 Subject: [PATCH 5/6] fix(qwenimage): measure the fast-math build and align weight views explicitly Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../model_specific/qwenimage/det_linear.cu | 38 +++++++++++++--- .../qwenimage/adaln_mod_projection.py | 20 ++++++--- tests/models/qwenimage/test_adaln.py | 26 +++++++++++ .../qwenimage/test_strict_math_build.py | 44 +++++++++++++++++++ 4 files changed, 117 insertions(+), 11 deletions(-) create mode 100644 tests/models/qwenimage/test_strict_math_build.py diff --git a/csrc/cuda/model_specific/qwenimage/det_linear.cu b/csrc/cuda/model_specific/qwenimage/det_linear.cu index 4117e9ad6..c674538ad 100644 --- a/csrc/cuda/model_specific/qwenimage/det_linear.cu +++ b/csrc/cuda/model_specific/qwenimage/det_linear.cu @@ -50,15 +50,43 @@ #include #include +#include #include #include +namespace { +// Runs in this translation unit, so it is built with the same nvcc flags as the kernels. A +// subnormal survives a multiply only without flush-to-zero, and IEEE division differs from +// div.full (what --use_fast_math emits) for the chosen operands. Inputs come from device memory +// so the compiler cannot fold either result. +__global__ void strict_math_probe(const float* in, float* out) { + out[0] = in[0] * in[3]; + out[1] = in[1] / in[2]; +} +} // namespace + +// Report whether this extension was built without --use_fast_math by executing the probe above. +// nvcc defines no macro for that flag (neither __FAST_MATH__ nor __USE_FAST_MATH__ is visible to +// device code), so the build is measured, not detected. Fails closed when the probe cannot run. bool qwenimage_strict_math_compiled() { -#ifdef __FAST_MATH__ - return false; -#else - return true; -#endif + static std::atomic verdict{0}; // 0 unknown, 1 strict, 2 fast math + if (int known = verdict.load()) return known == 1; + try { + // 0x1.ed8256p+0 / 0x1.5e809ep+0 rounds differently under div.full. + volatile float subnormal_source = 1e-40f, a_source = 0x1.ed8256p+0f, b_source = 0x1.5e809ep+0f; + const float subnormal = subnormal_source, a = a_source, b = b_source, quotient = a / b; + auto in = torch::tensor({subnormal, a, b, 1.0f}, torch::kFloat).to(torch::kCUDA); + auto out = torch::empty({2}, in.options()); + strict_math_probe<<<1, 1, 0, at::cuda::getCurrentCUDAStream()>>>(in.data_ptr(), + out.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + auto result = out.cpu(); + const bool strict = result[0].item() == subnormal && result[1].item() == quotient; + verdict.store(strict ? 1 : 2); + return strict; + } catch (const c10::Error&) { + return false; + } } namespace { diff --git a/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py b/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py index 273fdc971..c217a31d9 100644 --- a/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py +++ b/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py @@ -36,11 +36,19 @@ def require_device(tensor): raise RuntimeError("Qwen-Image CUDA requires SM80 or newer") +def aligned(tensor): + """Return ``tensor`` contiguous and 16-byte aligned for the vector loads of the kernels. + + ``contiguous()`` keeps a view whose storage offset is not a multiple of 16 bytes, so such a + view is copied explicitly. + """ + tensor = tensor.contiguous() + return tensor.clone() if tensor.data_ptr() % 16 else tensor + + def projection_forward(temb, weight, bias): - act = F.silu(temb).contiguous() - table = _C.qwenimage_det_linear_forward(act, weight.contiguous(), bias.contiguous(), 0, False)[ - 0 - ] + act = aligned(F.silu(temb)) + table = _C.qwenimage_det_linear_forward(act, aligned(weight), aligned(bias), 0, False)[0] return table, act @@ -49,9 +57,9 @@ def parameter_gradients(grad, act, dtype): def projection_backward(grad, temb, act, weight): - grad = grad.float().contiguous() + grad = aligned(grad.float()) dw, db = parameter_gradients(grad, act, weight.dtype) - da = _C.qwenimage_det_linear_backward_input(grad, weight.contiguous(), torch.float32) + da = _C.qwenimage_det_linear_backward_input(grad, aligned(weight), torch.float32) z = temb.float() sigmoid = torch.sigmoid(z) dt = da * (sigmoid * (1 + z * (1 - sigmoid))) diff --git a/tests/models/qwenimage/test_adaln.py b/tests/models/qwenimage/test_adaln.py index fd832fcad..885ff40a5 100644 --- a/tests/models/qwenimage/test_adaln.py +++ b/tests/models/qwenimage/test_adaln.py @@ -207,3 +207,29 @@ def test_projection_chunk_order(): out = operator("projection")(**args) for i, chunk in enumerate(out.chunk(6, -1)): assert equal_bytes(chunk, torch.full_like(chunk, i)) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_misaligned_storage_views_match_aligned_copies(dtype): + rows, cols = 6 * 3072, 3072 + generator = torch.Generator(device="cuda").manual_seed(5) + flat = torch.randn(rows * cols + 8, generator=generator, device="cuda", dtype=dtype) * 0.02 + bias_flat = torch.randn(rows + 8, generator=generator, device="cuda", dtype=dtype) * 0.01 + # Contiguous views whose storage offset is one element: contiguous() keeps them misaligned. + weight = flat[1 : 1 + rows * cols].view(rows, cols).requires_grad_() + bias = bias_flat[1 : 1 + rows].requires_grad_() + assert weight.is_contiguous() and weight.data_ptr() % 16 + temb = torch.randn(2, cols, generator=generator, device="cuda", dtype=dtype).requires_grad_() + copies = { + "temb": temb, + "weight": weight.detach().clone().requires_grad_(), + "bias": bias.detach().clone().requires_grad_(), + } + assert copies["weight"].data_ptr() % 16 == 0 + op = operator("projection") + grad = random_upstream(op(temb=temb, weight=weight, bias=bias)) + y, grads = evaluate(op, {"temb": temb, "weight": weight, "bias": bias}, grad) + y_copy, grads_copy = evaluate(op, copies, grad) + assert equal_bytes(y, y_copy) + for name in grads: + assert equal_bytes(grads[name], grads_copy[name]), name diff --git a/tests/models/qwenimage/test_strict_math_build.py b/tests/models/qwenimage/test_strict_math_build.py new file mode 100644 index 000000000..b552b5ce5 --- /dev/null +++ b/tests/models/qwenimage/test_strict_math_build.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Build the real kernel source with and without --use_fast_math and ask it which one it is.""" + +import shutil +from pathlib import Path + +import pytest +import torch + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or shutil.which("nvcc") is None, + reason="CUDA and nvcc required", +) + +ROOT = Path(__file__).resolve().parents[3] +BINDING = """ +#include +bool qwenimage_strict_math_compiled(); +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("strict", &qwenimage_strict_math_compiled); } +""" + + +def build(name, flags, tmp_path): + from torch.utils.cpp_extension import load + + major, minor = torch.cuda.get_device_capability() + (tmp_path / name).mkdir() + binding = tmp_path / f"{name}.cpp" + binding.write_text(BINDING) + return load( + name=name, + sources=[str(binding), str(ROOT / "csrc/cuda/model_specific/qwenimage/det_linear.cu")], + extra_cuda_cflags=["-O3", f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}"] + + flags, + build_directory=str(tmp_path / name), + ) + + +@pytest.mark.parametrize( + ("flags", "strict"), [([], True), (["--use_fast_math"], False)], ids=["default", "fast_math"] +) +def test_probe_reports_the_actual_build_flags(flags, strict, tmp_path): + name = "qwenimage_math_" + ("default" if strict else "fast_math") + assert build(name, flags, tmp_path).strict() is strict From 366bb9a2ce362d9198934bc47ed1bcb4561ed53e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 28 Oct 2026 16:15:18 +0000 Subject: [PATCH 6/6] fix(qwenimage): compensate the token sum so small contributions survive cancellation Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../model_specific/qwenimage/token_sum.cu | 48 +++++++++++++------ docs/operators/qwenimage-final-adaln-out.md | 11 +++-- .../qwenimage/final_adaln_out.py | 2 +- tests/models/qwenimage/test_adaln.py | 13 +++++ 4 files changed, 53 insertions(+), 21 deletions(-) diff --git a/csrc/cuda/model_specific/qwenimage/token_sum.cu b/csrc/cuda/model_specific/qwenimage/token_sum.cu index 8311aedfd..cd8fc8f16 100644 --- a/csrc/cuda/model_specific/qwenimage/token_sum.cu +++ b/csrc/cuda/model_specific/qwenimage/token_sum.cu @@ -6,26 +6,38 @@ #include namespace { -// Recursive fixed 32-token balanced trees. Each -// sample/channel has its own sequence; batching never changes the tree. -__global__ void token_partials(const float* x, float* p, int64_t seq, int64_t hidden, - int64_t tiles) { +// Neumaier step: add v to the running sum s and accumulate the rounding error in c. +__device__ __forceinline__ void neumaier_add(float& s, float& c, float v) { + const float t = __fadd_rn(s, v); + c = __fadd_rn(c, fabsf(s) >= fabsf(v) ? __fadd_rn(__fsub_rn(s, t), v) + : __fadd_rn(__fsub_rn(v, t), s)); + s = t; +} + +// Fixed 32-token tiles, folded level by level. Inside a tile the tokens are added in ascending +// order with Neumaier compensation, and a tile's carried correction joins the next level's +// correction, so small contributions survive next to large cancelling ones. Each +// sample/channel has its own sequence; batching never changes the order. +// The first level reads x only; later levels read the previous level's sums and corrections. +__global__ void token_partials(const float* x, const float* x_corr, float* p, float* p_corr, + int64_t seq, int64_t hidden, int64_t tiles) { const int64_t h = blockIdx.x * blockDim.x + threadIdx.x; const int64_t tile = blockIdx.y; const int64_t batch = blockIdx.z; if (h >= hidden) return; - float v[32]; -#pragma unroll + float s = 0.0f, c = 0.0f; +#pragma unroll 4 for (int i = 0; i < 32; ++i) { - int64_t row = tile * 32 + i; - v[i] = row < seq ? x[(batch * seq + row) * hidden + h] : 0.0f; - } -#pragma unroll - for (int stride = 1; stride < 32; stride *= 2) { -#pragma unroll - for (int i = 0; i < 32; i += 2 * stride) v[i] = __fadd_rn(v[i], v[i + stride]); + const int64_t row = tile * 32 + i; + if (row < seq) { + const int64_t at = (batch * seq + row) * hidden + h; + neumaier_add(s, c, x[at]); + if (x_corr != nullptr) c = __fadd_rn(c, x_corr[at]); + } } - p[(batch * tiles + tile) * hidden + h] = v[0]; + const int64_t out = (batch * tiles + tile) * hidden + h; + p[out] = s; + p_corr[out] = c; } } @@ -41,16 +53,22 @@ torch::Tensor qwenimage_token_sum(torch::Tensor x) { "token_sum geometry out of bounds"); const c10::cuda::CUDAGuard guard(x.device()); auto current = x; + torch::Tensor correction; auto stream = at::cuda::getCurrentCUDAStream(); int64_t rows = s; while (rows > 1) { const auto tiles = (rows + 31) / 32; auto partials = torch::empty({b, tiles, h}, x.options()); + auto corrections = torch::empty({b, tiles, h}, x.options()); token_partials<<>>( - current.data_ptr(), partials.data_ptr(), rows, h, tiles); + current.data_ptr(), correction.defined() ? correction.data_ptr() : nullptr, + partials.data_ptr(), corrections.data_ptr(), rows, h, tiles); C10_CUDA_KERNEL_LAUNCH_CHECK(); current = partials; + correction = corrections; rows = tiles; } + // One IEEE add of the final sum and its carried correction. + if (correction.defined()) current = current + correction; return current.view({b, h}); } diff --git a/docs/operators/qwenimage-final-adaln-out.md b/docs/operators/qwenimage-final-adaln-out.md index 56d6203a4..0af753641 100644 --- a/docs/operators/qwenimage-final-adaln-out.md +++ b/docs/operators/qwenimage-final-adaln-out.md @@ -37,16 +37,17 @@ versions; strict training and rollout must both bind this same operator. Backward uses FP32 intermediates and straight-through derivatives at storage rounding boundaries. Native LayerNorm backward computes `dx`; scale/shift -contributions reduce over S with recursively balanced 32-token FP32 trees. -The independent CPU reference uses sequential Neumaier summation over tokens -in FP32 (or FP64 for diagnostics). This also retains small contributions under -cancellation, independently of the unchanged CUDA token tree. There are no atomics +contributions reduce over S in fixed 32-token tiles. Inside a tile the tokens +are added in ascending order with FP32 Neumaier compensation, and the tile sums +and corrections are folded level by level, so small contributions survive next +to large cancelling ones. The independent CPU reference uses sequential Neumaier +summation over tokens in FP32 (or FP64 for diagnostics). There are no atomics or collective reductions. The projection receives FP32 table gradients before the final casts to leaf dtypes. Batch size/position does not alter a sample's token reduction tree. Readback identifies the native LayerNorm and elementwise stages, CUDA projection -kernel, recursive token-sum tree, dtype, device, and lack of fallback. No fused +kernel, compensated token-sum tiles, dtype, device, and lack of fallback. No fused kernel or speedup over the native composition is claimed. ## Supported scope and evidence diff --git a/rl_engine/backends/cuda/model_specific/qwenimage/final_adaln_out.py b/rl_engine/backends/cuda/model_specific/qwenimage/final_adaln_out.py index 8507cb323..268b882ed 100644 --- a/rl_engine/backends/cuda/model_specific/qwenimage/final_adaln_out.py +++ b/rl_engine/backends/cuda/model_specific/qwenimage/final_adaln_out.py @@ -65,7 +65,7 @@ def implementation_info(self, x, temb, weight, bias): info = implementation_info("qwenimage_final_adaln_out", x) info.update( norm="aten.native_layer_norm[fp32,eps=1e-6]", - token_gradient_tree="recursive_32_token_balanced", + token_gradient_tree="neumaier_32_token_tiles_folded", cast_boundaries="silu,projection,norm,1+scale,product,output", fused=False, ) diff --git a/tests/models/qwenimage/test_adaln.py b/tests/models/qwenimage/test_adaln.py index 8066a901d..89b55c8d0 100644 --- a/tests/models/qwenimage/test_adaln.py +++ b/tests/models/qwenimage/test_adaln.py @@ -247,3 +247,16 @@ def test_final_scale_then_shift_and_layernorm(): args["x"] = torch.full_like(args["x"], 13.0) out = operator("final")(**args) assert equal_bytes(out, torch.full_like(out, 7.0)) + + +def test_token_sum_retains_small_contributions_under_cancellation(): + from rl_engine.backends.extension import _C + + # 2**26 absorbs a unit token in FP32, so a plain tree returns 6020 instead of 6030. + values = torch.ones(2, 6032, 3, dtype=torch.float32) + values[:, 0] = 2**26 + values[:, -1] = -(2**26) + golden = values.double().sum(1).float() + assert equal_bytes(_C.qwenimage_token_sum(values.cuda()).cpu(), golden) + # The same sample reduced alone gives the same bytes. + assert equal_bytes(_C.qwenimage_token_sum(values[1:].cuda()).cpu(), golden[1:])