46be74b4b4
Co-authored-by: 戚余航 <qiyuhang@bytedance.com> Co-authored-by: Qi Yuhang <45795032+HydraQYH@users.noreply.github.com>
115 lines
3.8 KiB
Python
115 lines
3.8 KiB
Python
import numpy as np
|
|
import pytest
|
|
import tabulate
|
|
import torch
|
|
from diffusers.models.embeddings import get_timestep_embedding
|
|
from sgl_kernel.elementwise import timestep_embedding as timestep_embedding_cuda
|
|
|
|
from sglang.multimodal_gen.runtime.layers.visual_embedding import timestep_embedding
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"batch_size", [1, 2, 8, 128, 256, 512, 1536, 2048, 4096, 11008, 16384]
|
|
)
|
|
@pytest.mark.parametrize("dim", [32, 128, 256, 512, 1536, 2048, 4096, 8192])
|
|
@pytest.mark.parametrize(
|
|
"dtype", [torch.int32, torch.int64, torch.bfloat16, torch.float16]
|
|
)
|
|
def test_timestep_embedding_correctness_with_sgld(batch_size, dim, dtype):
|
|
device = "cuda"
|
|
t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype)
|
|
torch_output = timestep_embedding(t, dim)
|
|
cuda_output = timestep_embedding_cuda(t, dim, flip_sin_to_cos=True)
|
|
torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
@pytest.mark.parametrize("batch_size", [1, 2, 8, 128, 256, 512, 1536, 2048, 16384])
|
|
@pytest.mark.parametrize("dim", [32, 256, 512, 1536, 8192])
|
|
@pytest.mark.parametrize("dtype", [torch.int32, torch.bfloat16])
|
|
@pytest.mark.parametrize("flip_sin_to_cos", [False, True])
|
|
@pytest.mark.parametrize("downscale_freq_shift", [0, 1])
|
|
@pytest.mark.parametrize("scale", [1, 0.01])
|
|
def test_timestep_embedding_correctness_with_diffusers(
|
|
batch_size, dim, flip_sin_to_cos, downscale_freq_shift, scale, dtype
|
|
):
|
|
device = "cuda"
|
|
t = torch.randint(low=0, high=1000, size=(batch_size,), device=device).to(dtype)
|
|
torch_output = get_timestep_embedding(
|
|
t,
|
|
dim,
|
|
flip_sin_to_cos=flip_sin_to_cos,
|
|
downscale_freq_shift=downscale_freq_shift,
|
|
scale=scale,
|
|
max_period=10000,
|
|
)
|
|
cuda_output = timestep_embedding_cuda(
|
|
t,
|
|
dim,
|
|
flip_sin_to_cos=flip_sin_to_cos,
|
|
downscale_freq_shift=downscale_freq_shift,
|
|
scale=scale,
|
|
max_period=10000,
|
|
)
|
|
torch.testing.assert_close(torch_output, cuda_output, atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_timestep_embedding_perf():
|
|
NUM_BATCH = [1, 2, 8, 63, 256, 512, 613, 1024, 1536]
|
|
NUM_DIM = [32, 64, 128, 256, 512, 1024, 2048, 4096]
|
|
|
|
def perf_kernel_fn(kernel_fn: callable, *args, **kwargs):
|
|
warmup_times = 4
|
|
repeat_times = 20
|
|
start = torch.cuda.Event(enable_timing=True)
|
|
end = torch.cuda.Event(enable_timing=True)
|
|
|
|
for _ in range(warmup_times):
|
|
output_fn = kernel_fn(*args, **kwargs)
|
|
torch.cuda.synchronize()
|
|
|
|
start.record()
|
|
for _ in range(repeat_times):
|
|
output_fn = kernel_fn(*args, **kwargs)
|
|
end.record()
|
|
end.synchronize()
|
|
return start.elapsed_time(end) / repeat_times
|
|
|
|
device = "cuda"
|
|
results = []
|
|
|
|
cuda_speedups = []
|
|
for B in NUM_BATCH:
|
|
for dim in NUM_DIM:
|
|
t = torch.linspace(0, max(100000, B), steps=B, device=device).to(
|
|
torch.int32
|
|
)
|
|
time_torch = perf_kernel_fn(timestep_embedding, t, dim)
|
|
time_cuda = perf_kernel_fn(timestep_embedding_cuda, t, dim)
|
|
speedup_cuda = time_torch / time_cuda
|
|
|
|
results.append(
|
|
{
|
|
"Batch Size": B,
|
|
"Dimension": dim,
|
|
"Torch Time (ms)": time_torch,
|
|
"CUDA Time (ms)": time_cuda,
|
|
"Speedup (CUDA)": speedup_cuda,
|
|
}
|
|
)
|
|
cuda_speedups.append(speedup_cuda)
|
|
|
|
print("=== Timestep Embedding Benchmark Results ===")
|
|
print(
|
|
tabulate.tabulate(
|
|
results,
|
|
headers="keys",
|
|
tablefmt="fancy_grid",
|
|
floatfmt=(".0f", ".0f", ".6f", ".6f", ".5f"),
|
|
)
|
|
)
|
|
print(f"Average Speedup(cuda): {np.mean(cuda_speedups):.4f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__])
|