#!/usr/bin/env python3 """Correctness tests for diffusion-step-ops.""" from __future__ import annotations import argparse import ctypes import ctypes.util import importlib import os import sys from pathlib import Path import torch ROOT = Path(__file__).resolve().parents[2] PACKAGE = ROOT / "diffusion-step-ops" REGISTRATION_INCLUDE = ( ROOT.parent / "kernels" / "kernel-builder" / "src" / "pyproject" / "templates" / "torch" ) class SourceOps: def __init__(self, namespace: str) -> None: self._ops = getattr(torch.ops, namespace) def add_bf16(self, a, b): out = torch.empty_like(a) self._ops.add_bf16_out(a, b, out) return out def euler_step_bf16(self, latent, velocity, dt): out = torch.empty_like(latent) self._ops.euler_step_bf16_out(latent, velocity, float(dt), out) return out def cfg_combine_into_residual_bf16(self, residual, v_cond, v_uncond, beta): self._ops.cfg_combine_into_residual_bf16(residual, v_cond, v_uncond, float(beta)) return residual def cfg_combine_into_residual_fp16(self, residual, v_cond, v_uncond, beta): self._ops.cfg_combine_into_residual_fp16(residual, v_cond, v_uncond, float(beta)) return residual def teacher_force_first_frame_bf16(self, video_latent, cond_latent): self._ops.teacher_force_first_frame_bf16(video_latent, cond_latent) return video_latent def motus_decode_postprocess_bf16_to_fp32(self, decoded): out = torch.empty( (decoded.shape[0], decoded.shape[1], decoded.shape[2] - 1, decoded.shape[3], decoded.shape[4]), device=decoded.device, dtype=torch.float32, ) self._ops.motus_decode_postprocess_bf16_to_fp32(decoded, out) return out def cast_bf16_to_fp32(self, src): dst = torch.empty_like(src, dtype=torch.float32) self._ops.cast_bf16_to_fp32(src, dst) return dst def pack_tail_bf16(self, tail, flat_dim, out=None): if out is None: out = torch.empty((flat_dim,), device=tail.device, dtype=tail.dtype) self._ops.pack_tail_bf16(tail, int(flat_dim), out) return out def add_bias_zero_tail_bf16(self, input, bias, valid_cols, out=None): if out is None: out = torch.empty_like(input) self._ops.add_bias_zero_tail_bf16(input, bias, int(valid_cols), out) return out def extract_tail_f32_to_bf16(self, flat, tail_numel, out=None): if out is None: out = torch.empty((tail_numel,), device=flat.device, dtype=torch.bfloat16) self._ops.extract_tail_f32_to_bf16(flat, int(tail_numel), out) return out def add_bias_pair_bf16(self, input, bias_a, bias_b): out = torch.empty_like(input) self._ops.add_bias_pair_bf16(input, bias_a, bias_b, out) return out def unipc_step_f32_bf16( self, sample, velocity, prev_m1, prev_m2, prev_last_sample, sigma, corrector_order, predictor_order, corrector_coefficients, predictor_coefficients, ): outputs = [torch.empty_like(sample) for _ in range(3)] self._ops.unipc_step_f32_bf16( sample, velocity, prev_m1, prev_m2, prev_last_sample, float(sigma), int(corrector_order), int(predictor_order), *map(float, corrector_coefficients), *map(float, predictor_coefficients), *outputs, ) return tuple(outputs) def _preload_cublaslt() -> None: for parent in Path(torch.__file__).resolve().parents: candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12" if candidate.exists(): ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL) return library = ctypes.util.find_library("cublasLt") if library: ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL) def _current_arch_list() -> str: major, minor = torch.cuda.get_device_capability(0) return f"{major}.{minor}" def load_source_ops() -> SourceOps: from torch.utils.cpp_extension import load if not REGISTRATION_INCLUDE.is_dir(): raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}") _preload_cublaslt() os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _current_arch_list()) namespace = "diffusion_step_ops_test" load( name=namespace, sources=[ str(PACKAGE / "torch-ext" / "torch_binding.cpp"), str(PACKAGE / "csrc" / "diffusion_step_ops.cu"), ], extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)], extra_cflags=["-O3", "-DCUDA_KERNEL"], extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"], verbose=False, ) return SourceOps(namespace) def load_installed_ops(artifact: str | None): if artifact: sys.path.insert(0, artifact) try: return importlib.import_module("diffusion_step_ops") finally: if artifact: sys.path.remove(artifact) def assert_close(name: str, got: torch.Tensor, ref: torch.Tensor, atol: float) -> None: diff = (got.float() - ref.float()).abs() max_err = diff.max().item() mean_err = diff.mean().item() cos = torch.nn.functional.cosine_similarity(got.float().flatten(), ref.float().flatten(), dim=0).item() if max_err > atol or cos < 0.9999: raise AssertionError(f"{name}: max_err={max_err:.8f}, mean_err={mean_err:.8f}, cos={cos:.8f}") def run_elementwise_tests(ops) -> int: count = 0 for shape in [(1024,), (1025,), (4, 4096), (2, 16, 32, 64)]: a = torch.randn(shape, device="cuda", dtype=torch.bfloat16) b = torch.randn(shape, device="cuda", dtype=torch.bfloat16) got = ops.add_bf16(a, b) ref = (a.float() + b.float()).to(torch.bfloat16) assert_close(f"add_bf16 shape={shape}", got, ref, 0.0) dt = -0.125 got = ops.euler_step_bf16(a, b, dt) ref = (a.float() + b.float() * dt).to(torch.bfloat16) assert_close(f"euler_step_bf16 shape={shape}", got, ref, 0.0) residual = torch.randn(shape, device="cuda", dtype=torch.bfloat16) residual_ref = residual.clone() beta = 4.5 got = ops.cfg_combine_into_residual_bf16(residual, a, b, beta) ref = (residual_ref.float() + b.float() + beta * (a.float() - b.float())).to(torch.bfloat16) assert_close(f"cfg_bf16 shape={shape}", got, ref, 0.0) ah = a.to(torch.float16) bh = b.to(torch.float16) residual_h = residual_ref.to(torch.float16) residual_h_ref = residual_h.clone() got = ops.cfg_combine_into_residual_fp16(residual_h, ah, bh, beta) ref = (residual_h_ref.float() + bh.float() + beta * (ah.float() - bh.float())).to(torch.float16) assert_close(f"cfg_fp16 shape={shape}", got, ref, 0.0) got = ops.cast_bf16_to_fp32(a) ref = a.float() assert_close(f"cast_bf16_to_fp32 shape={shape}", got, ref, 0.0) count += 5 return count def run_video_tests(ops) -> int: count = 0 for shape in [(1, 4, 5, 16, 16), (2, 8, 9, 8, 8), (1, 16, 17, 16, 24)]: video = torch.randn(shape, device="cuda", dtype=torch.bfloat16) cond = torch.randn((shape[0], shape[1], shape[3], shape[4]), device="cuda", dtype=torch.bfloat16) ref = video.clone() ref[:, :, 0] = cond got = ops.teacher_force_first_frame_bf16(video.clone(), cond) assert_close(f"teacher_force shape={shape}", got, ref, 0.0) decoded = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 3.0 got = ops.motus_decode_postprocess_bf16_to_fp32(decoded) ref = ((decoded[:, :, 1:].float() + 1.0) * 0.5).clamp(0.0, 1.0).contiguous() assert_close(f"motus_postprocess shape={shape}", got, ref, 0.0) count += 2 return count def run_tail_tests(ops) -> int: count = 0 for flat_dim, tail_numel in [(32, 7), (257, 51), (4096, 1024)]: tail = torch.randn((tail_numel,), device="cuda", dtype=torch.bfloat16) got = ops.pack_tail_bf16(tail, flat_dim) ref = torch.zeros((flat_dim,), device="cuda", dtype=torch.bfloat16) ref[-tail_numel:] = tail assert_close(f"pack_tail {flat_dim=} {tail_numel=}", got, ref, 0.0) flat = torch.randn((flat_dim,), device="cuda", dtype=torch.float32) got = ops.extract_tail_f32_to_bf16(flat, tail_numel) ref = flat[-tail_numel:].to(torch.bfloat16) assert_close(f"extract_tail {flat_dim=} {tail_numel=}", got, ref, 0.0) count += 2 for rows, cols, valid_cols in [(1, 16, 7), (51, 64, 32), (105, 257, 256)]: input = torch.randn((rows, cols), device="cuda", dtype=torch.bfloat16) bias = torch.randn((cols,), device="cuda", dtype=torch.bfloat16) got = ops.add_bias_zero_tail_bf16(input, bias, valid_cols) ref = (input.float() + bias.float()).to(torch.bfloat16) ref[:, valid_cols:] = 0 assert_close( f"add_bias_zero_tail {rows=} {cols=} {valid_cols=}", got, ref, 0.0, ) bias_b = torch.randn((cols,), device="cuda", dtype=torch.bfloat16) got = ops.add_bias_pair_bf16(input, bias, bias_b) ref = (input.float() + bias.float()).to(torch.bfloat16) ref = (ref.float() + bias_b.float()).to(torch.bfloat16) assert_close(f"add_bias_pair {rows=} {cols=}", got, ref, 0.0) count += 2 tail = torch.randn((51,), device="cuda", dtype=torch.bfloat16) input = torch.randn((51, 64), device="cuda", dtype=torch.bfloat16) bias_a = torch.randn((64,), device="cuda", dtype=torch.bfloat16) bias_b = torch.randn((64,), device="cuda", dtype=torch.bfloat16) def invoke(tail, input, bias_a, bias_b): return ( ops.pack_tail_bf16(tail, 257), ops.add_bias_pair_bf16(input, bias_a, bias_b), ) eager = invoke(tail, input, bias_a, bias_b) compiled = torch.compile(invoke, fullgraph=True)(tail, input, bias_a, bias_b) for got, expected in zip(compiled, eager): torch.testing.assert_close(got, expected, rtol=0.0, atol=0.0) print("PASS action-tail torch.compile fullgraph") return count + 1 def run_cosmos_edge_contract(ops) -> int: flat_dim = 1_201_920 tail_numel = 60 * 64 rows, cols, valid_cols = 60, 64, 9 tail = torch.randn((tail_numel,), device="cuda", dtype=torch.bfloat16) flat = torch.randn((flat_dim,), device="cuda", dtype=torch.float32) matrix = torch.randn((rows, cols), device="cuda", dtype=torch.bfloat16) bias = torch.randn((cols,), device="cuda", dtype=torch.bfloat16) packed = torch.empty((flat_dim,), device="cuda", dtype=torch.bfloat16) extracted = torch.empty((tail_numel,), device="cuda", dtype=torch.bfloat16) biased = torch.empty_like(matrix) ops.pack_tail_bf16(tail, flat_dim, out=packed) ops.extract_tail_f32_to_bf16(flat, tail_numel, out=extracted) ops.add_bias_zero_tail_bf16(matrix, bias, valid_cols, out=biased) expected_packed = torch.zeros_like(packed) expected_packed[-tail_numel:] = tail expected_extracted = flat[-tail_numel:].to(torch.bfloat16) expected_biased = (matrix.float() + bias.float()).to(torch.bfloat16) expected_biased[:, valid_cols:] = 0 torch.testing.assert_close(packed, expected_packed, rtol=0.0, atol=0.0) torch.testing.assert_close(extracted, expected_extracted, rtol=0.0, atol=0.0) torch.testing.assert_close(biased, expected_biased, rtol=0.0, atol=0.0) graph = torch.cuda.CUDAGraph() torch.cuda.synchronize() with torch.cuda.graph(graph): ops.pack_tail_bf16(tail, flat_dim, out=packed) ops.extract_tail_f32_to_bf16(flat, tail_numel, out=extracted) ops.add_bias_zero_tail_bf16(matrix, bias, valid_cols, out=biased) graph.replay() torch.cuda.synchronize() first = (packed.clone(), extracted.clone(), biased.clone()) graph.replay() torch.cuda.synchronize() second = (packed.clone(), extracted.clone(), biased.clone()) for got, expected in zip(second, first): torch.testing.assert_close(got, expected, rtol=0.0, atol=0.0) print("PASS Cosmos3-Edge action-tail contract and CUDA Graph replay") return 4 def run_unipc_tests(ops) -> int: count = 0 corrector = (0.75, 0.2, -0.1, 0.05, 0.4) predictor = (0.8, 0.3, -0.07) for shape in [(1,), (257,), (1, 16, 17, 8, 8)]: sample = torch.randn(shape, device="cuda", dtype=torch.float32) velocity = torch.randn( shape, device="cuda", dtype=torch.bfloat16 ) prev_m1 = torch.randn_like(sample) prev_m2 = torch.randn_like(sample) prev_last = torch.randn_like(sample) for corrector_order, predictor_order in [(0, 1), (1, 1), (1, 2), (2, 2)]: got_next, got_m, got_last = ops.unipc_step_f32_bf16( sample, velocity, prev_m1, prev_m2, prev_last, 0.37, corrector_order, predictor_order, corrector, predictor, ) sigma_velocity = (velocity.float() * 0.37).to( torch.bfloat16 ).float() expected_m = sample - sigma_velocity expected_last = corrector[0] * sample + corrector[4] * expected_m if corrector_order >= 1: expected_last = ( expected_last + corrector[1] * prev_last + corrector[2] * prev_m1 ) if corrector_order >= 2: expected_last = expected_last + corrector[3] * prev_m2 expected_next = ( predictor[0] * expected_last + predictor[1] * expected_m ) if predictor_order >= 2: expected_next = expected_next + predictor[2] * prev_m1 torch.testing.assert_close( got_m, expected_m, rtol=1e-6, atol=1e-6 ) torch.testing.assert_close( got_last, expected_last, rtol=2e-6, atol=2e-6 ) torch.testing.assert_close( got_next, expected_next, rtol=2e-6, atol=2e-6 ) count += 1 sample = torch.randn((257,), device="cuda", dtype=torch.float32) velocity = torch.randn( (257,), device="cuda", dtype=torch.bfloat16 ) history = [torch.randn_like(sample) for _ in range(3)] def invoke(sample, velocity, prev_m1, prev_m2, prev_last): return ops.unipc_step_f32_bf16( sample, velocity, prev_m1, prev_m2, prev_last, 0.37, 2, 2, corrector, predictor, ) eager = invoke(sample, velocity, *history) compiled = torch.compile(invoke, fullgraph=True)( sample, velocity, *history ) for got, expected in zip(compiled, eager): torch.testing.assert_close(got, expected, rtol=0.0, atol=0.0) print("PASS unipc_step torch.compile fullgraph") return count + 1 def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--backend", choices=["source", "installed"], default="source") parser.add_argument("--artifact", default=None) args = parser.parse_args() if not torch.cuda.is_available(): raise RuntimeError("CUDA is required") torch.manual_seed(0) ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact) total = ( run_elementwise_tests(ops) + run_video_tests(ops) + run_tail_tests(ops) + run_cosmos_edge_contract(ops) + run_unipc_tests(ops) ) torch.cuda.synchronize() print(f"diffusion-step-ops correctness passed: {total} checks") return 0 if __name__ == "__main__": raise SystemExit(main())