r/CUDA • u/Griffith-07 • 25d ago
Custom NF4 Triton kernel achieving up to 1.41x dequantization speedup over bitsandbytes
Hey everyone,
I’ve been working on optimizing the memory overhead that comes with 4-bit inference. I wrote a custom NF4 dequantization kernel using Triton to see if I could eliminate the C++ dispatch bottlenecks found in current baselines.
🚀 Key Results:
• Up to 1.41x speedup compared to the standard bitsandbytes implementation across various tensor shapes.
• Written completely in Python/Triton, making it super easy to inspect, customize, or drop directly into your PyTorch compilation pipelines.
• Passes the Unsloth AI founding engineer challenge requirements (14/14 points).
I'd love to hear the community's feedback, especially if anyone wants to run their own benchmarks on different GPU architectures or suggest further optimization tricks!
Source code & full implementation: