from kernels import get_kernel
import torch

ops = get_kernel("flashrt/transformer-fused-ops", version=1, trust_remote_code=True)
x = torch.randn((8, 128), device="cuda", dtype=torch.bfloat16)
gate = torch.randn_like(x)
weight = torch.ones((128,), device="cuda", dtype=torch.bfloat16)
y = ops.rms_norm_gated_silu_bf16(x, gate, weight)
print(y.shape)