| |
| """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() |
|
|
