12 lines
311 B
Python
12 lines
311 B
Python
import torch
|
|
from kernels import get_kernel
|
|
|
|
|
|
fastvideo_kernel = get_kernel("hao-ai-lab/fastvideo-kernel", version=1)
|
|
|
|
x = torch.randn((16, 1024), device="cuda", dtype=torch.float16)
|
|
weight = torch.ones((1024,), device="cuda", dtype=torch.float16)
|
|
y = fastvideo_kernel.rms_norm(x, 1e-6, weight)
|
|
|
|
print(y.shape)
|