84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
import torch
|
|
from torch import nn
|
|
|
|
from benchmark import (
|
|
benchmark_combined,
|
|
benchmark_forward,
|
|
benchmark_backward,
|
|
assert_close,
|
|
)
|
|
from fastvideo.ops.modulate.modulate import fused_modulate
|
|
|
|
|
|
def torch_modulate(x, scale, shift):
|
|
return x * (scale.unsqueeze(1) + 1) + shift.unsqueeze(1)
|
|
|
|
|
|
# from flash_attn import flash_attn_func
|
|
|
|
|
|
def time_fwd(func, *args, **kwargs):
|
|
time_fb = benchmark_forward(func, *args, **kwargs)
|
|
return time_fb[1].mean
|
|
|
|
|
|
def time_fwd_bwd(func, *args, **kwargs):
|
|
time_fb = benchmark_combined(func, *args, **kwargs)
|
|
return time_fb[1].mean
|
|
|
|
|
|
def time_bwd(func, *args, **kwargs):
|
|
time_fb = benchmark_backward(func, *args, **kwargs)
|
|
return time_fb[1].mean
|
|
|
|
|
|
device = "cuda"
|
|
dtype = torch.bfloat16
|
|
|
|
|
|
batch_sizes = [1, 8, 32]
|
|
seq_lengths = [128, 512, 1024]
|
|
hidden_dims = [768, 1024, 2048]
|
|
|
|
# methods = (["torch", "triton", "thunderkitten"])
|
|
methods = ["torch", "triton"]
|
|
time_f = {}
|
|
time_b = {}
|
|
time_f_b = {}
|
|
speed_f = {}
|
|
speed_b = {}
|
|
speed_f_b = {}
|
|
for B in batch_sizes:
|
|
for T in seq_lengths:
|
|
for D in hidden_dims:
|
|
config = (B, T, D)
|
|
|
|
norm_func = nn.LayerNorm(D, elementwise_affine=False, eps=1e-6)
|
|
x = torch.randn(B, T, D, device="cuda", requires_grad=True, dtype=dtype)
|
|
x = norm_func(x.to(torch.float32)).to(dtype)
|
|
shift = torch.randn(B, D, device="cuda", requires_grad=True, dtype=dtype)
|
|
scale = torch.randn(B, D, device="cuda", requires_grad=True, dtype=dtype)
|
|
# test torch
|
|
o_ref = torch_modulate(x, scale, shift)
|
|
o_ref.sum().backward(retain_graph=True)
|
|
f_b = time_fwd_bwd(torch_modulate, x, scale, shift, verbose=False)
|
|
time_f_b[config, "torch"] = f_b
|
|
# test triton
|
|
o2 = fused_modulate(x, scale, shift)
|
|
o2.sum().backward(retain_graph=True)
|
|
f_b = time_fwd_bwd(fused_modulate, x, scale, shift, verbose=False)
|
|
time_f_b[config, "triton"] = f_b
|
|
# test if the results are close
|
|
assert_close(" o", o_ref, o2, 0.005)
|
|
# time_f_b[config, "thunderkitten"] = f_b
|
|
|
|
print(f"### batch size={B}, seq length={T}, B={B}, hidden dim={D} ###")
|
|
for method in methods:
|
|
# time_f_b[config, method] = time_f[config, method] + time_b[config, method]
|
|
print(
|
|
f"{method:>50} fwd + bwd:\t {time_f_b[config, method]*1000:>6.4f} ms "
|
|
)
|
|
|
|
# with open('flash2_attn_time.plk', 'wb') as fp:
|
|
# pickle.dump((speed_f, speed_b, speed_f_b), fp, protocol=pickle.HIGHEST_PROTOCOL)
|