File size: 876 Bytes
88b8ef2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 | #!/usr/bin/env python3
"""Minimal Hub-style call for flashrt/fp4-gemm."""
from __future__ import annotations
import torch
from kernels import get_kernel
def main() -> None:
if not torch.cuda.is_available():
raise SystemExit("CUDA is required")
ops = get_kernel("flashrt/fp4-gemm", version=1, trust_remote_code=True)
x = torch.randn((32, 256), device="cuda", dtype=torch.float16)
w = torch.randn((512, 256), device="cuda", dtype=torch.float16)
a_packed, sfa = ops.quantize_fp4_sfa_fp16(x, is_sfb=False)
b_packed, sfb = ops.quantize_fp4_sfa_fp16(w, is_sfb=True)
y = ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, alpha=1.0)
print("a_packed", tuple(a_packed.shape), a_packed.dtype)
print("b_packed", tuple(b_packed.shape), b_packed.dtype)
print("output", tuple(y.shape), y.dtype)
if __name__ == "__main__":
main()
|