| |
| """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()) |
|
|
