Fused Triton Kernels for 4-Bit and Ternary Weight Only Inference: A TinyLlama-1.1B Case Study
Autoregressive decoding repeatedly reads model weights from device memory, making weight compression valuable on accelerators with limited memory capacity and bandwidth. This work studies two custom Triton implementations for TinyLlama-1.1B: a zero-centered 4-bit format and a ternary format stored with two bits per weight. Both kernels load packed INT32 weights, decode them on chip, reconstruct FP16 tiles, and perform matrix multiplication with FP32 accumulation. On an NVIDIA T4, the 4-bit implementation reduced allocated inference memory from 2.21 GB to 1.05 GB while obtaining a sampled WikiText-2 perplexity of 18.90 compared with 17.77 for FP16. The ternary representation reduced memory to 0.80 GB but suffered substantial quality degradation under naive groupwise absmean ternarization. A lightweight post-quantization affine adaptation improved perplexity substantially but did not restore useful generation. The study provides an implementation-level reference and negative baseline for future work on activation-aware and optimization-based low-precision inference.
Authors
- Anshuman Agrawal
Institutions
- University of Petroleum and Energy Studies (IN)
Publication Details
- Journal
- Zenodo (CERN European Organization for Nuclear Research)
- Published
- 2026-08-14
- DOI
- https://doi.org/10.5281/zenodo.21936390
- Primary Topic
- Parallel Computing and Optimization Techniques
- Type
- preprint