Dong1017andamy-why-3459
ee9ab9ea88
Signed-off-by: GUOGUO <xwdong1998@163.com> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
235 lines
10 KiB
Python
235 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Candidate-added tests for the fused Triton AdaLayerNorm CUDA path.
|
|
|
|
Complements (never replaces) the frozen suite: exercises the fused fast path
|
|
and the forward_native fallback branches of forward_cuda. Uses the same
|
|
tolerance scheme as the frozen suite; golden = forward_native real behavior.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.layers.adalayernorm import (
|
|
AdaLayerNorm,
|
|
_adaln_fused_forward,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
|
|
|
|
TOL_STRICT = {"bf16": (2e-2, 2e-2), "fp32": (1e-3, 1e-3)}
|
|
# Matches the frozen suite: the double-rounded golden chain deviates from a
|
|
# pure-fp32 reference more than the strict band, so the fp32-reference check
|
|
# uses the same loose band as the frozen tests.
|
|
TOL_LOOSE = {"bf16": (5e-2, 2e-2), "fp32": (5e-3, 5e-3)}
|
|
DTYPES = [torch.bfloat16, torch.float32]
|
|
|
|
|
|
def fp32_reference(x, scale, shift, eps):
|
|
xf = x.float()
|
|
xn = torch.nn.functional.layer_norm(xf, (x.shape[-1],), None, None, eps)
|
|
return xn * (1 + scale.float().reshape(1, 1, -1)) + shift.float().reshape(1, 1, -1)
|
|
|
|
|
|
def make_module(hidden, affine, eps, device, dtype):
|
|
m = AdaLayerNorm(hidden, elementwise_affine=affine, eps=eps)
|
|
if affine:
|
|
m = m.to(dtype)
|
|
return m.to(device)
|
|
|
|
|
|
def assert_close(a, b, dtype, loose=False):
|
|
atol, rtol = (TOL_LOOSE if loose else TOL_STRICT)["bf16" if dtype == torch.bfloat16 else "fp32"]
|
|
torch.testing.assert_close(a.float(), b.float(), atol=atol, rtol=rtol)
|
|
|
|
|
|
def make_inputs(bs, seq, hidden, dtype, device, seed=0):
|
|
g = torch.Generator(device=device).manual_seed(seed)
|
|
x = torch.randn(bs, seq, hidden, generator=g, device=device, dtype=dtype)
|
|
scale = torch.randn(1, hidden, generator=g, device=device, dtype=dtype)
|
|
shift = torch.randn(1, hidden, generator=g, device=device, dtype=dtype)
|
|
return x, scale, shift
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", DTYPES)
|
|
@pytest.mark.parametrize("affine", [False, True])
|
|
@pytest.mark.parametrize("bs,seq,hidden", [(1, 512, 3072), (2, 128, 1536), (1, 3, 1000)])
|
|
@torch.no_grad()
|
|
def test_fused_fast_path_matches_native(dtype, affine, bs, seq, hidden):
|
|
# hidden=1000 exercises BLOCK_C masking (next_pow2(1000)=1024 > 1000)
|
|
device = "cuda"
|
|
m = make_module(hidden, affine, 1e-6, device, dtype)
|
|
x, scale, shift = make_inputs(bs, seq, hidden, dtype, device)
|
|
fused = _adaln_fused_forward(m, x, scale, shift)
|
|
assert fused is not None, "supported inputs must take the fused path so BLOCK_C masking is exercised"
|
|
out = fused
|
|
assert out.shape == x.shape and out.dtype == dtype and out.device == x.device
|
|
assert_close(out, m.forward_native(x, scale, shift), dtype)
|
|
assert_close(out, fp32_reference(x, scale, shift, 1e-6), dtype, loose=True)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", DTYPES)
|
|
def test_fused_noncontiguous_x(dtype):
|
|
# Non-contiguous x routes to the forward_native fallback (the fused kernel
|
|
# requires contiguous x); result must still match native exactly.
|
|
device = "cuda"
|
|
hidden = 1536
|
|
m = make_module(hidden, False, 1e-6, device, dtype)
|
|
g = torch.Generator(device=device).manual_seed(3)
|
|
x_big = torch.randn(1, 1024, hidden, generator=g, device=device, dtype=dtype)
|
|
x = x_big[:, ::2, :]
|
|
scale = torch.randn(1, hidden, generator=g, device=device, dtype=dtype)
|
|
shift = torch.randn(1, hidden, generator=g, device=device, dtype=dtype)
|
|
out = m.forward_cuda(x, scale, shift)
|
|
assert_close(out, m.forward_native(x, scale, shift), dtype)
|
|
assert_close(out, fp32_reference(x.contiguous(), scale, shift, 1e-6), dtype, loose=True)
|
|
|
|
|
|
def test_normalized_shape_mismatch_preserves_native_error():
|
|
m = AdaLayerNorm(16, elementwise_affine=False).cuda()
|
|
x = torch.randn(1, 2, 32, device="cuda")
|
|
scale = torch.zeros(32, device="cuda")
|
|
shift = torch.zeros_like(scale)
|
|
assert _adaln_fused_forward(m, x, scale, shift) is None
|
|
with pytest.raises(RuntimeError) as native_error:
|
|
m.forward_native(x, scale, shift)
|
|
with pytest.raises(RuntimeError) as cuda_error:
|
|
m.forward_cuda(x, scale, shift)
|
|
assert str(cuda_error.value) == str(native_error.value)
|
|
|
|
|
|
@pytest.mark.parametrize("grad_target", ["x", "scale", "shift", "weight", "bias"])
|
|
def test_grad_fallback_matches_native_backward(grad_target):
|
|
m = make_module(32, grad_target in ("weight", "bias"), 1e-6, "cuda", torch.float32)
|
|
m.requires_grad_(False)
|
|
x, scale, shift = make_inputs(1, 2, 32, torch.float32, "cuda", seed=12)
|
|
targets = {"x": x, "scale": scale, "shift": shift, "weight": m.layernorm.weight, "bias": m.layernorm.bias}
|
|
target = targets[grad_target]
|
|
target.requires_grad_(True)
|
|
with torch.enable_grad():
|
|
assert _adaln_fused_forward(m, x, scale, shift) is None
|
|
native = m.forward_native(x, scale, shift)
|
|
out = m.forward_cuda(x, scale, shift)
|
|
assert out.requires_grad
|
|
torch.testing.assert_close(out, native)
|
|
upstream = torch.randn_like(out)
|
|
native_grad = torch.autograd.grad(native, target, upstream)[0]
|
|
cuda_grad = torch.autograd.grad(out, target, upstream)[0]
|
|
torch.testing.assert_close(cuda_grad, native_grad)
|
|
|
|
|
|
def test_no_grad_keeps_fused_path_for_grad_requiring_tensors():
|
|
m = make_module(32, True, 1e-6, "cuda", torch.float32)
|
|
x, scale, shift = make_inputs(1, 2, 32, torch.float32, "cuda", seed=14)
|
|
for tensor in (x, scale, shift):
|
|
tensor.requires_grad_(True)
|
|
with torch.no_grad():
|
|
fused = _adaln_fused_forward(m, x, scale, shift)
|
|
assert fused is not None
|
|
assert not fused.requires_grad
|
|
torch.testing.assert_close(fused, m.forward_native(x, scale, shift), atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_1d_scale_shift_fast_path():
|
|
# (C,) scale/shift broadcast per-channel exactly like (1, C)
|
|
device = "cuda"
|
|
hidden = 3072
|
|
m = make_module(hidden, False, 1e-6, device, torch.bfloat16)
|
|
g = torch.Generator(device=device).manual_seed(10)
|
|
x = torch.randn(1, 256, hidden, generator=g, device=device, dtype=torch.bfloat16)
|
|
scale = torch.randn(hidden, generator=g, device=device, dtype=torch.bfloat16)
|
|
shift = torch.randn(hidden, generator=g, device=device, dtype=torch.bfloat16)
|
|
out = _adaln_fused_forward(m, x, scale, shift)
|
|
assert out is not None, "(C,) scale/shift must take the fused path"
|
|
assert_close(out, m.forward_native(x, scale, shift), torch.bfloat16)
|
|
assert_close(out, fp32_reference(x, scale, shift, 1e-6), torch.bfloat16, loose=True)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", DTYPES)
|
|
def test_per_row_scale_falls_back_to_native(dtype):
|
|
# (L, C) scale with B=1 is a valid per-row torch broadcast; it is not
|
|
# per-channel, so forward_cuda must route it to forward_native.
|
|
device = "cuda"
|
|
hidden = 512
|
|
m = make_module(hidden, False, 1e-6, device, dtype)
|
|
g = torch.Generator(device=device).manual_seed(4)
|
|
x = torch.randn(1, 64, hidden, generator=g, device=device, dtype=dtype)
|
|
scale = torch.randn(64, hidden, generator=g, device=device, dtype=dtype)
|
|
shift = torch.randn(64, hidden, generator=g, device=device, dtype=dtype)
|
|
out = m.forward_cuda(x, scale, shift)
|
|
ref = m.forward_native(x, scale, shift)
|
|
assert out.dtype == ref.dtype
|
|
assert_close(out, ref, dtype)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", DTYPES)
|
|
def test_mixed_dtype_falls_back_to_native(dtype):
|
|
# scale/shift dtype != x.dtype: native type-promotes; fused path requires
|
|
# matching dtypes, so this must go through forward_native.
|
|
other = torch.float32 if dtype == torch.bfloat16 else torch.bfloat16
|
|
device = "cuda"
|
|
hidden = 256
|
|
m = make_module(hidden, False, 1e-6, device, dtype)
|
|
g = torch.Generator(device=device).manual_seed(5)
|
|
x = torch.randn(1, 32, hidden, generator=g, device=device, dtype=dtype)
|
|
scale = torch.randn(1, hidden, generator=g, device=device, dtype=other)
|
|
shift = torch.randn(1, hidden, generator=g, device=device, dtype=other)
|
|
out = m.forward_cuda(x, scale, shift)
|
|
ref = m.forward_native(x, scale, shift)
|
|
assert out.dtype == ref.dtype
|
|
torch.testing.assert_close(out, ref)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", DTYPES)
|
|
def test_zero_dim_scale_falls_back_to_native(dtype):
|
|
# 0-d scalar modulation is a valid torch broadcast; not per-channel.
|
|
device = "cuda"
|
|
hidden = 256
|
|
m = make_module(hidden, False, 1e-6, device, dtype)
|
|
g = torch.Generator(device=device).manual_seed(6)
|
|
x = torch.randn(1, 32, hidden, generator=g, device=device, dtype=dtype)
|
|
scale = torch.tensor(0.5, device=device, dtype=dtype)
|
|
shift = torch.tensor(-0.25, device=device, dtype=dtype)
|
|
out = m.forward_cuda(x, scale, shift)
|
|
assert_close(out, m.forward_native(x, scale, shift), dtype)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", DTYPES)
|
|
def test_per_sample_scale_raises_like_native(dtype):
|
|
# Frozen behavior: B>1 with (B, C) scale/shift -> RuntimeError.
|
|
device = "cuda"
|
|
hidden = 3072
|
|
m = make_module(hidden, False, 1e-6, device, dtype)
|
|
g = torch.Generator(device=device).manual_seed(7)
|
|
x = torch.randn(4, 64, hidden, generator=g, device=device, dtype=dtype)
|
|
scale = torch.randn(4, hidden, generator=g, device=device, dtype=dtype)
|
|
shift = torch.randn(4, hidden, generator=g, device=device, dtype=dtype)
|
|
with pytest.raises(RuntimeError):
|
|
m.forward_native(x, scale, shift)
|
|
with pytest.raises(RuntimeError):
|
|
m.forward_cuda(x, scale, shift)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", DTYPES)
|
|
def test_fused_determinism_repeat_calls(dtype):
|
|
device = "cuda"
|
|
hidden = 3072
|
|
m = make_module(hidden, False, 1e-6, device, dtype)
|
|
x, scale, shift = make_inputs(1, 1024, hidden, dtype, device, seed=8)
|
|
out1 = _adaln_fused_forward(m, x, scale, shift)
|
|
out2 = _adaln_fused_forward(m, x, scale, shift)
|
|
assert out1 is not None and out2 is not None, "determinism must exercise the fused path"
|
|
assert torch.equal(out1, out2)
|
|
|
|
|
|
@pytest.mark.parametrize("eps", [1e-5, 1e-3])
|
|
def test_fused_eps_variants(eps):
|
|
device = "cuda"
|
|
hidden = 3072
|
|
m = make_module(hidden, False, eps, device, torch.bfloat16)
|
|
x, scale, shift = make_inputs(1, 256, hidden, torch.bfloat16, device, seed=9)
|
|
out = _adaln_fused_forward(m, x, scale, shift)
|
|
assert out is not None, "supported eps variants must take the fused path"
|
|
assert_close(out, m.forward_native(x, scale, shift), torch.bfloat16)
|
|
assert_close(out, fp32_reference(x, scale, shift, eps), torch.bfloat16, loose=True)
|