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..c674538ad --- /dev/null +++ b/csrc/cuda/model_specific/qwenimage/det_linear.cu @@ -0,0 +1,691 @@ +// 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 +#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() { + 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 { + +#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..fd200e50c --- /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](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 +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](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/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..c217a31d9 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/qwenimage/adaln_mod_projection.py @@ -0,0 +1,124 @@ +# 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 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 = aligned(F.silu(temb)) + table = _C.qwenimage_det_linear_forward(act, aligned(weight), aligned(bias), 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 = aligned(grad.float()) + dw, db = parameter_gradients(grad, act, weight.dtype) + 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))) + 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..885ff40a5 --- /dev/null +++ b/tests/models/qwenimage/test_adaln.py @@ -0,0 +1,235 @@ +# 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)) + + +@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_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"] == [] 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