| from kernels import get_kernel | |
| import torch | |
| try: | |
| ops = get_kernel( | |
| "flashrt/grouped-moe-gemv", version=2, trust_remote_code=True | |
| ) | |
| except TypeError: # kernels==0.12.x does not expose trust_remote_code. | |
| ops = get_kernel("flashrt/grouped-moe-gemv", version=2) | |
| K, N = 256, 128 | |
| x = torch.ones((K,), device="cuda", dtype=torch.bfloat16) | |
| w = torch.full((N, K // 2), 0x11, device="cuda", dtype=torch.uint8) | |
| sfb = torch.full((512,), 0x38, device="cuda", dtype=torch.uint8) | |
| y = ops.w4a16_decode_gemv_bf16(x, w, sfb) | |
| print(y.shape) | |