296 lines
11 KiB
Python
296 lines
11 KiB
Python
import torch
|
|
import pytest
|
|
import sys
|
|
import os
|
|
|
|
# Ensure local package is imported
|
|
# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../python")))
|
|
|
|
try:
|
|
from fastvideo_kernel._C import fastvideo_kernel_ops
|
|
except ImportError:
|
|
fastvideo_kernel_ops = None
|
|
|
|
from fastvideo_kernel import turbodiffusion_ops
|
|
|
|
|
|
# Helper for RMS Norm reference
|
|
def rms_norm_ref(x, w, eps=1e-6):
|
|
dtype = x.dtype
|
|
x = x.float()
|
|
variance = x.pow(2).mean(-1, keepdim=True)
|
|
x = x * torch.rsqrt(variance + eps)
|
|
return (x * w.float()).to(dtype)
|
|
|
|
|
|
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
|
class TestTurboDiffusion:
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("shape", [(16, 128), (32, 256), (1, 1024)])
|
|
def test_quant_correctness(self, dtype, shape):
|
|
if turbodiffusion_ops.quant_cuda is None:
|
|
pytest.skip("quant_cuda not available")
|
|
|
|
x = torch.randn(shape, dtype=dtype, device="cuda")
|
|
x_q, x_scale = turbodiffusion_ops.int8_quant(x)
|
|
|
|
assert x_q.dtype == torch.int8
|
|
assert x_scale.dtype == torch.float32
|
|
|
|
# Simple check: dequantize and compute error
|
|
# Note: The quantization scheme details matter here (per block? per tensor?).
|
|
# Looking at quant.cu, it seems to be block-based but the output scale shape isn't immediately obvious from python signature
|
|
# without looking at C++ code deeper.
|
|
# But let's check shapes at least.
|
|
|
|
# If we can't easily dequantize without knowing block size logic in python,
|
|
# checking that it runs and produces valid shapes is a good start.
|
|
assert x_q.shape == shape
|
|
# x_scale shape depends on block size, usually smaller than x
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
def test_gemm_correctness(self, dtype):
|
|
if turbodiffusion_ops.gemm_cuda is None:
|
|
pytest.skip("gemm_cuda not available")
|
|
|
|
M, N, K = 32, 64, 128
|
|
x = torch.randn(M, K, dtype=dtype, device="cuda")
|
|
|
|
# Create weights
|
|
# For simplicity in testing, let's create random int8 weights and scales
|
|
w_q = torch.randint(-127, 127, (N, K), dtype=torch.int8, device="cuda")
|
|
|
|
# Scale shape: The Int8Linear class uses:
|
|
# row_blocks = cdiv(out_features, b=128)
|
|
# col_blocks = cdiv(in_features, b=128)
|
|
# scale shape: (row_blocks, col_blocks)
|
|
|
|
row_blocks = (N + 127) // 128
|
|
col_blocks = (K + 127) // 128
|
|
w_s = torch.randn(row_blocks, col_blocks, dtype=torch.float32, device="cuda").abs()
|
|
|
|
# Run int8_linear
|
|
output = turbodiffusion_ops.int8_linear(x, w_q, w_s)
|
|
|
|
assert output.shape == (M, N)
|
|
assert output.dtype == dtype
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
def test_gemm_backward(self, dtype):
|
|
if turbodiffusion_ops.gemm_cuda is None:
|
|
pytest.skip("gemm_cuda not available")
|
|
|
|
M, N, K = 32, 64, 128
|
|
x = torch.randn(M, K, dtype=dtype, device="cuda", requires_grad=True)
|
|
|
|
# Weights (frozen)
|
|
w_q = torch.randint(-127, 127, (N, K), dtype=torch.int8, device="cuda")
|
|
row_blocks = (N + 127) // 128
|
|
col_blocks = (K + 127) // 128
|
|
w_s = torch.randn(row_blocks, col_blocks, dtype=torch.float32, device="cuda").abs()
|
|
|
|
bias = torch.randn(N, dtype=dtype, device="cuda", requires_grad=True)
|
|
|
|
# Use Int8LinearFunction
|
|
output = turbodiffusion_ops.Int8LinearFunction.apply(x, w_q, w_s, bias)
|
|
loss = output.sum()
|
|
loss.backward()
|
|
|
|
assert x.grad is not None
|
|
assert bias.grad is not None
|
|
assert x.grad.shape == (M, K)
|
|
|
|
# Check correctness against dequantized weight
|
|
w_float = turbodiffusion_ops.dequantize_weight(w_q, w_s, dtype)
|
|
|
|
# Reference
|
|
x_ref = x.detach().clone().requires_grad_()
|
|
bias_ref = bias.detach().clone().requires_grad_()
|
|
|
|
# Note: Int8Linear forward uses quantized input, so output won't match exactly reference with float input.
|
|
# But backward gradients should be consistent with the logic we implemented:
|
|
# grad_input = grad_output @ w_float
|
|
|
|
# Let's verify our backward logic matches standard matmul backward with dequantized weights
|
|
# We manually compute expected gradient given grad_output = ones
|
|
grad_output = torch.ones_like(output)
|
|
expected_x_grad = grad_output @ w_float
|
|
expected_bias_grad = grad_output.sum(0)
|
|
|
|
torch.testing.assert_close(x.grad, expected_x_grad, atol=1e-2, rtol=1e-2)
|
|
torch.testing.assert_close(bias.grad, expected_bias_grad, atol=1e-2, rtol=1e-2)
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("shape", [(2, 16, 128), (4, 32, 256)])
|
|
def test_rms_norm_triton(self, dtype, shape):
|
|
x = torch.randn(shape, dtype=dtype, device="cuda")
|
|
dim = shape[-1]
|
|
w = torch.randn(dim, dtype=dtype, device="cuda")
|
|
eps = 1e-5
|
|
|
|
# Triton implementation
|
|
# Note: rmsnorm returns tuple now
|
|
res = turbodiffusion_ops.rmsnorm(x, w, eps)
|
|
if isinstance(res, tuple):
|
|
out_triton = res[0]
|
|
else:
|
|
out_triton = res
|
|
|
|
# Reference
|
|
out_ref = rms_norm_ref(x, w, eps)
|
|
|
|
torch.testing.assert_close(out_triton, out_ref, atol=1e-2, rtol=1e-2)
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("shape", [(2, 16, 128)])
|
|
def test_rms_norm_backward(self, dtype, shape):
|
|
x = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
|
|
dim = shape[-1]
|
|
w = torch.randn(dim, dtype=dtype, device="cuda", requires_grad=True)
|
|
eps = 1e-5
|
|
|
|
# Forward via Function
|
|
y = turbodiffusion_ops.RMSNormFunction.apply(x, w, eps)
|
|
loss = y.sum()
|
|
loss.backward()
|
|
|
|
x_grad = x.grad
|
|
w_grad = w.grad
|
|
|
|
# Reference
|
|
x_ref = x.detach().clone().requires_grad_()
|
|
w_ref = w.detach().clone().requires_grad_()
|
|
|
|
# Custom RMSNorm ref in pytorch
|
|
def rms_norm_ref_grad(x, w, eps):
|
|
x_float = x.float()
|
|
var = x_float.pow(2).mean(-1, keepdim=True)
|
|
rstd = torch.rsqrt(var + eps)
|
|
return (x_float * rstd * w.float()).to(x.dtype)
|
|
|
|
y_ref = rms_norm_ref_grad(x_ref, w_ref, eps)
|
|
loss_ref = y_ref.sum()
|
|
loss_ref.backward()
|
|
|
|
torch.testing.assert_close(x_grad, x_ref.grad, atol=1e-2, rtol=1e-2)
|
|
torch.testing.assert_close(w_grad, w_ref.grad, atol=2e-2, rtol=2e-2)
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
def test_rms_norm_cuda(self, dtype):
|
|
if fastvideo_kernel_ops is None or not hasattr(fastvideo_kernel_ops, "rms_norm_cuda"):
|
|
pytest.skip("rms_norm_cuda not available")
|
|
|
|
shape = (16, 128)
|
|
x = torch.randn(shape, dtype=dtype, device="cuda")
|
|
dim = shape[-1]
|
|
w = torch.randn(dim, dtype=dtype, device="cuda")
|
|
eps = 1e-5
|
|
|
|
# C++ implementation
|
|
# Signature: rms_norm_cuda(Input, eps, Weight, Output) -> Output
|
|
out_cuda = torch.empty_like(x)
|
|
fastvideo_kernel_ops.rms_norm_cuda(x, eps, w, out_cuda)
|
|
|
|
# Reference
|
|
out_ref = rms_norm_ref(x, w, eps)
|
|
|
|
torch.testing.assert_close(out_cuda, out_ref, atol=1e-2, rtol=1e-2)
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("shape", [(2, 16, 128), (4, 32, 256)])
|
|
def test_layer_norm_triton(self, dtype, shape):
|
|
x = torch.randn(shape, dtype=dtype, device="cuda")
|
|
dim = shape[-1]
|
|
eps = 1e-5
|
|
|
|
# With affine
|
|
w = torch.randn(dim, dtype=dtype, device="cuda")
|
|
b = torch.randn(dim, dtype=dtype, device="cuda")
|
|
|
|
# Triton implementation
|
|
# Note: layernorm returns tuple now
|
|
res = turbodiffusion_ops.layernorm(x, w, b, eps, elementwise_affine=True)
|
|
if isinstance(res, tuple):
|
|
out_triton = res[0]
|
|
else:
|
|
out_triton = res
|
|
out_triton = out_triton.to(dtype)
|
|
|
|
# Reference
|
|
ln = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=dtype).cuda()
|
|
ln.weight.data.copy_(w)
|
|
ln.bias.data.copy_(b)
|
|
out_ref = ln(x)
|
|
|
|
torch.testing.assert_close(out_triton, out_ref, atol=1e-2, rtol=1e-2)
|
|
|
|
# Without affine
|
|
res_no_affine = turbodiffusion_ops.layernorm(x, None, None, eps, elementwise_affine=False)
|
|
if isinstance(res_no_affine, tuple):
|
|
out_triton_no_affine = res_no_affine[0]
|
|
else:
|
|
out_triton_no_affine = res_no_affine
|
|
out_triton_no_affine = out_triton_no_affine.to(dtype)
|
|
ln_no_affine = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=False, dtype=dtype).cuda()
|
|
out_ref_no_affine = ln_no_affine(x)
|
|
|
|
torch.testing.assert_close(out_triton_no_affine, out_ref_no_affine, atol=1e-2, rtol=1e-2)
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
@pytest.mark.parametrize("shape", [(2, 16, 128)])
|
|
def test_layer_norm_backward(self, dtype, shape):
|
|
x = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
|
|
dim = shape[-1]
|
|
eps = 1e-5
|
|
w = torch.randn(dim, dtype=dtype, device="cuda", requires_grad=True)
|
|
b = torch.randn(dim, dtype=dtype, device="cuda", requires_grad=True)
|
|
|
|
# Triton implementation via Function
|
|
y = turbodiffusion_ops.LayerNormFunction.apply(x, w, b, eps, True)
|
|
loss = y.sum()
|
|
loss.backward()
|
|
|
|
x_grad = x.grad
|
|
w_grad = w.grad
|
|
b_grad = b.grad
|
|
|
|
# Reference
|
|
x_ref = x.detach().clone().requires_grad_()
|
|
ln = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=dtype).cuda()
|
|
ln.weight.data.copy_(w.detach())
|
|
ln.bias.data.copy_(b.detach())
|
|
|
|
y_ref = ln(x_ref)
|
|
loss_ref = y_ref.sum()
|
|
loss_ref.backward()
|
|
|
|
torch.testing.assert_close(x_grad, x_ref.grad, atol=1e-2, rtol=1e-2)
|
|
torch.testing.assert_close(w_grad, ln.weight.grad, atol=1e-2, rtol=1e-2)
|
|
torch.testing.assert_close(b_grad, ln.bias.grad, atol=1e-2, rtol=1e-2)
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
|
def test_layer_norm_cuda(self, dtype):
|
|
if fastvideo_kernel_ops is None or not hasattr(fastvideo_kernel_ops, "layer_norm_cuda"):
|
|
pytest.skip("layer_norm_cuda not available")
|
|
|
|
shape = (16, 128)
|
|
x = torch.randn(shape, dtype=dtype, device="cuda")
|
|
dim = shape[-1]
|
|
eps = 1e-5
|
|
w = torch.randn(dim, dtype=dtype, device="cuda")
|
|
b = torch.randn(dim, dtype=dtype, device="cuda")
|
|
|
|
# C++ implementation
|
|
# Signature: layer_norm_cuda(Input, eps, W, B, Output) -> Output
|
|
out_cuda = torch.empty_like(x)
|
|
fastvideo_kernel_ops.layer_norm_cuda(x, eps, w, b, out_cuda)
|
|
|
|
# Reference
|
|
ln = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=dtype).cuda()
|
|
ln.weight.data.copy_(w)
|
|
ln.bias.data.copy_(b)
|
|
out_ref = ln(x)
|
|
|
|
torch.testing.assert_close(out_cuda, out_ref, atol=1e-2, rtol=1e-2)
|