WaveCut/orbitquant-packed-matmul

1# OrbitQuant Packed Matmul2 3Packed low-bit matrix multiplication kernel for OrbitQuant inference.4 5This kernel consumes OrbitQuant packed weight indices, per-row norms, and Lloyd-Max6centroids directly, avoiding a full BF16/FP16 dequantized weight cache before the7linear projection.8 9## API10 11```python12import torch13from orbitquant_packed_matmul import matmul_packed_weight14 15out = matmul_packed_weight(16    x,17    packed_weight_indices,18    row_norms,19    centroids,20    bits=4,21    out_features=3072,22    in_features=3072,23    bias=bias,24)25```26 27Inputs:28 29- `x`: contiguous or reshapeable tensor with shape `[..., in_features]`.30- `packed_weight_indices`: `uint8` low-bit packed row-major codebook indices.31- `row_norms`: row norms with shape `[out_features]`; CUDA consumes the32  artifact's `bfloat16` values directly, while Metal uses `float32` internally.33- `centroids`: `float32` Lloyd-Max centroids with shape `[2**bits]`.34- `bias`: optional projection bias.35 36`x` may be `float32`, `float16`, or `bfloat16`. The output has shape37`[..., out_features]` and the same dtype as `x`.38 39The CUDA package also exports the operations used by OrbitQuant's W4A4 runtime:40 41- `quantize_activations_int8`: token norm, RPBH/FWHT, nearest-codebook42  assignment, and INT8-surrogate output in one native launch.43- `quantize_activations_packed_w4`: the same activation path with packed 4-bit44  output for the direct packed matmul fallback.45- `matmul_packed_w4a4_int8`: direct packed A4/W4 CUDA MMA with fused token norm,46  row norm, surrogate scales, and bias epilogue.47 48On CUDA compute capability 8.0 or newer, OrbitQuant normally combines49`quantize_activations_int8` with chunked packed-weight decode and Torch's50CUTLASS-backed INT8 matmul. The direct packed MMA operation remains available51when that path is unsupported. Neither path materializes a complete BF16/FP1652weight matrix.53 54## Build And Test55 56```bash57nix --option sandbox relaxed --option max-jobs 1 --option cores 8 \58  run .#build-and-copy -L59nix --option sandbox relaxed --option max-jobs 1 --option cores 8 \60  run .#ci-test -L61```62 63The build produces ABI3 Hugging Face Kernels artifacts under `build/` for the64supported backend variants on the current platform. On macOS, `sandbox relaxed`65or enabled Nix sandboxing is required by `kernel-builder`. The commands build66local files only; they do not upload to Kernel Hub.67 68For a faster CUDA-only development build on a machine with a matching Torch and69CUDA toolchain:70 71```bash72cargo install --git https://github.com/huggingface/kernels hf-kernel-builder73kernel-builder check-config .74kernel-builder create-pyproject -f .75TORCH_CUDA_ARCH_LIST="8.9" CUDACXX=/usr/local/cuda/bin/nvcc \76  python setup.py build_kernel77```78 79For a local Metal build compatible with macOS 15 and newer:80 81```bash82cargo install --git https://github.com/huggingface/kernels hf-kernel-builder83kernel-builder check-config .84kernel-builder create-pyproject -f .85MACOSX_DEPLOYMENT_TARGET=15.0 \86  CMAKE_ARGS="-DCMAKE_OSX_DEPLOYMENT_TARGET=15.0" \87  python setup.py build_kernel88kernel-builder check-abi --macos 15.0 --python-abi 3.9 .89```90 91This generated project is for local testing and must not be committed or92distributed without a successful `kernel-builder check-abi`. Use the Nix build93for redistributable variants.94 95For direct local imports, add the matching `build/torch*-<backend>-<platform>`96directory to `PYTHONPATH`; the `torch*` variant must match the runtime PyTorch97version:98 99```bash100export PYTHONPATH="/path/to/build/torch212-metal-aarch64-darwin:$PYTHONPATH"101python -c "import orbitquant_packed_matmul; print(orbitquant_packed_matmul)"102```103 104For PyTorch 2.9 CUDA inference, set the allocator before starting Python when105minimum reserved memory is important:106 107```bash108PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python generate.py109```110 111OrbitQuant detects that importable package before trying any Hub loader. For112Hugging Face `kernels` local loading instead, set `LOCAL_KERNELS` to the same113built variant directory containing `metadata.json`, not to the source package114root:115 116```bash117export LOCAL_KERNELS="WaveCut/orbitquant-packed-matmul=/path/to/build/torch212-metal-aarch64-darwin"118```119 120## Benchmark121 122The benchmark reports two PyTorch references:123 124- `predequantized_f_linear_seconds_per_iter`: `torch.nn.functional.linear`125  over a full dequantized weight matrix that was materialized before timing.126- `dequantize_then_f_linear_seconds_per_iter`: materialize the full127  dequantized weight matrix inside each timed iteration, then call128  `torch.nn.functional.linear`.129 130```bash131PYTHONPATH=/path/to/build/torch212-metal-aarch64-darwin \132python benchmarks/benchmark.py \133  --device mps \134  --bits 4 \135  --rows 512 \136  --in-features 3072 \137  --out-features 3072 \138  --iters 20139```140 141`--rows` accepts a comma-separated sweep (the default covers decode-bound142small batches and GEMM-bound large batches), and `--dtype` selects the143activation dtype explicitly. Headline timings are hot-loop medians.144 145The script prints JSON with `packed_seconds_per_iter`,146`predequantized_f_linear_seconds_per_iter`,147`dequantize_then_f_linear_seconds_per_iter` (all hot-loop medians), the148per-path `*_hot_mean_seconds`, `*_hot_median_seconds`, and149`*_hot_p95_seconds` distributions,150`packed_vs_predequantized_f_linear_speedup`,151`packed_vs_dequantize_then_f_linear_speedup`, compatibility aliases152`reference_seconds_per_iter` and `packed_vs_reference_speedup`, and153`max_abs_error`.154 155It also reports storage accounting for the packed weight path:156`packed_weight_indices_bytes`, `row_norms_bytes`, `centroid_bytes`,157`packed_weight_path_bytes`, `materialized_weight_bytes`, and158`packed_weight_path_vs_materialized_weight_ratio`. These values describe only159the weight-side storage used by this operator; they are not end-to-end model160VRAM measurements.161 162### Metal reference results163 164Measured on an Apple M2 Max with Torch 2.12.1, FP16 activations, W4 packed165weights, transformer-scale shapes (hot-loop medians over 30 iterations; each166iteration synchronizes, so sub-millisecond rows include the MPS submit167floor):168 169| Shape (rows x in x out) | Packed Metal | Resident FP16 `F.linear` | Materialize + `F.linear` | Packed vs resident | Packed vs materialize |170| --- | ---: | ---: | ---: | ---: | ---: |171| 1 x 3072 x 3072 | 0.330 ms | 0.200 ms | 1.791 ms | 0.60x | 5.42x |172| 4 x 3072 x 3072 | 0.377 ms | 0.215 ms | 1.768 ms | 0.57x | 4.69x |173| 32 x 3072 x 3072 | 0.367 ms | 0.312 ms | 1.700 ms | 0.85x | 4.63x |174| 512 x 3072 x 3072 | 1.355 ms | 1.077 ms | 2.331 ms | 0.79x | 1.72x |175| 512 x 3072 x 12288 | 5.138 ms | 3.907 ms | 9.312 ms | 0.76x | 1.81x |176| 1 x 3072 x 12288 | 0.417 ms | 0.383 ms | 6.202 ms | 0.92x | 14.88x |177 178Batches of at most four rows dispatch a skinny-batch GEMV (one simdgroup per179output column, decoded weight segments reused across the batch); larger180batches dispatch the simdgroup-matrix tiles. The packed weight payload, row181norms, and centroids occupy about 25% of the materialized FP16 weight size at182W4. The resident reference excludes weight materialization time and retains183the complete FP16 matrix in memory — it is the throughput ceiling for a184kernel that decodes weights on the fly, not a like-for-like memory185configuration.186 187End-to-end FLUX.2 Klein 9B measurements and the SDNQ comparison are recorded in188[`docs/flux2-klein-9b-sdnq-vs-orbitquant.md`](../../docs/flux2-klein-9b-sdnq-vs-orbitquant.md).189