Show HN: Fast NF4 dequantization Triton kernel (1.41x faster than bitsandbytes)

A developer released an open-source Triton GPU kernel that dequantizes NF4 4-bit weights 1.27x–1.72x faster than the existing bitsandbytes C++ library, using a single fused kernel with inline PTX assembly and L2 cache eviction optimizations. The kernel eliminates CPU dispatch overhead and supports arbitrary tensor sizes, BF16/FP16 output, and torch.compile compatibility.

Optimized NF4 NormalFloat 4-bit dequantization via Triton GPU kernels Unsloth AI Founding Engineer Challenge 1— Convert NF4/BnB 4-bit dequantization from C++ to Triton. Achieves1.27x–1.41x speedupover bitsandbytes C++ across all tensor sizes. pip install git+https://github.com/Griffith-7/nf4-triton-kernel.git python from nf4 kernel import dequant nf4 output = dequant nf4 packed weights, absmax scales, dtype=torch.bfloat16 LLMs use 4-bit NF4 quantization to fit in VRAM. At inference time, weights must be dequantized back to 16-bit. The existing bitsandbytes library does this in C++, but suffers from CPU dispatch overhead. This project implements the same operation as a single fused Triton kernel — eliminating the overhead and achieving consistent speedups. Tested against bitsandbytes v0.49.2 on RTX 3050 Laptop GPU PyTorch 2.13.0, Triton 3.6.0 : | Tensor Size | Triton Kernel | bitsandbytes C++ | Speedup | |---|---|---|---| | 4,096 | 0.021 ms | 0.030 ms | 1.43x | | 16,384 | 0.018 ms | 0.030 ms | 1.68x | | 65,536 | 0.017 ms | 0.029 ms | 1.72x | | 262,144 | 0.025 ms | 0.035 ms | 1.39x | All sizes pass the 1.15x threshold. On Tesla T4, the original author reported 2.09x . A single fused Triton kernel performs three operations in one GPU pass: uint8 packed bytes → bit unpack → NF4 table lookup PTX ASM → absmax scale → BF16/FP16 output Bit unpacking — splits each uint8 into two 4-bit nibbles via bit-shifting NF4 lookup — maps each 4-bit index to a float using inline PTX assembly 16 values hardcoded in GPU registers Absmax scaling — multiplies by block-wise scaling factor and stores as BF16 or FP16 | Optimization | Detail | |---|---| Inline PTX Assembly | NF4 lookup table lives in GPU registers via tl.inline asm elementwise — zero memory reads, no branch mispredictions | L2 Cache Eviction | evict first for streaming packed weights, evict last for shared absmax — prevents cache thrashing | Single Fused Kernel | Bit-unpack + lookup + scale in one pass — no intermediate memory round-trips | Large Block Size | 1024 elements per thread block for optimal GPU occupancy | Dynamic Shapes | Handles arbitrary tensor sizes without hardcoded grid/block dimensions | | Feature | Status | Points | |---|---|---| | Single Fused Kernel | Passed | +3 | | Inline PTX Assembly | Passed | +3 | | L2 Cache Eviction | Passed | +1 | | BF16/FP16 Output | Passed | +1 | | torch.compile Compatible | Passed inductor | +1 | | 1.15x Speedup vs BnB | Passed 1.39x – 1.72x | +5 | Total | 14/14 | From GitHub pip install git+https://github.com/Griffith-7/nf4-triton-kernel.git Editable dev git clone https://github.com/Griffith-7/nf4-triton-kernel.git cd nf4-triton-kernel pip install -e ". dev-all " - Python 3.10+ - PyTorch 2.13+ - Triton 3.6+ - CUDA 12.4+ GPU nf4-triton-kernel/ ├── src/nf4 kernel/ │ ├── init .py Package exports, version │ └── kernel.py Core Triton kernel + quantize/dequantize utils ├── tests/ │ └── test nf4 kernel.py 21 unit tests correctness, edge cases, validation ├── benchmarks/ │ └── benchmark.py Performance benchmark vs bitsandbytes ├── .github/workflows/ │ └── ci.yml GitHub Actions CI lint + test ├── pyproject.toml PEP 621 project config ├── .pre-commit-config.yaml Pre-commit hooks ruff, codespell, etc. ├── Dockerfile Multi-stage build CPU + CUDA ├── colab notebook.py Google Colab notebook cell-by-cell └── README.md python import torch from nf4 kernel import dequant nf4 packed weights: uint8 tensor two 4-bit values per byte absmax: block-wise scaling factors FP8 or float32 output = dequant nf4 packed weights, absmax, group size=64, dtype=torch.bfloat16 output fp16 = dequant nf4 packed weights, absmax, dtype=torch.float16 python @torch.compile def compiled dequant pw, am : return dequant nf4 pw, am output = compiled dequant packed weights, absmax python -m pytest tests/ -v 21 tests covering: - Functional correctness BF16, FP16 - Edge cases small/large/odd/non-aligned tensors, zeros, uniform, alternating - Quantize utility validation - Input validation errors - Multiple group sizes 32, 64, 128 git clone https://github.com/Griffith-7/nf4-triton-kernel.git %cd nf4-triton-kernel pip install -e . exec open "colab notebook.py" .read MIT License — see LICENSE /Griffith-7/nf4-triton-kernel/blob/main/LICENSE .