| from __future__ import annotations |
|
|
| import torch |
| from kernels import get_kernel |
|
|
|
|
| def main() -> None: |
| ops = get_kernel("flashrt/adaptive-layernorm-producers", version=1, trust_remote_code=True) |
|
|
| rows, dim = 2520, 3072 |
| x = torch.randn((rows, dim), device="cuda", dtype=torch.bfloat16) |
| scale = torch.zeros((dim,), device="cuda", dtype=torch.bfloat16) |
| shift = torch.zeros((dim,), device="cuda", dtype=torch.bfloat16) |
| act_scale = torch.tensor([0.025], device="cuda", dtype=torch.float32) |
|
|
| out = torch.empty_like(x, dtype=torch.float8_e4m3fn) |
| ops.ada_layer_norm_quant_fp8_bf16(x, scale, shift, act_scale, out=out) |
|
|
| packed, sf = ops.ada_layer_norm_quant_nvfp4_swizzled_bf16(x, scale, shift) |
| print(out.shape, out.dtype) |
| print(packed.shape, packed.dtype, sf.numel()) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
