Views
No views yet
mul_mat_vec — fused dequantize + gemv, for up to MAX_GEMV_ROWS rowsdequantize — blocks to valuesget_rows — gathers rows, unpacking as it goesmul_mat_id — one dispatch for a bank of routed experts, given the router's choicesGEMV_TYPES lists the quantization types this build has a gemv for.1import torch
2from kernels import get_kernel
3
4k = get_kernel("marcsun13/ggml-quantization", version=1)
5
6Q4_K = 12 # ggml type id; `k.GEMV_TYPES` lists what this build covers
7out_features = in_features = 4096
8# a GGUF weight as stored: one row per output feature, 144 bytes per 256-element Q4_K block
9blocks = torch.randint(0, 256, (out_features, in_features // 256 * 144), dtype=torch.uint8, device="mps")
10x = torch.randn(1, in_features, device="mps")
11
12y = k.mul_mat_vec(blocks, x, Q4_K, out_features) # (1, 4096) f32
13w = k.dequantize(blocks, Q4_K, out_features, in_features, torch.bfloat16) # (4096, 4096)
14rows = k.get_rows(blocks, torch.tensor([3, 7], device="mps"), Q4_K, in_features, torch.bfloat16)