T-MAN: Enabling End-to-End Low-Bit LLM Inference on NPUs via Unified Table Lookup

TL;DR

T-MAN introduces a unified table lookup approach to enable efficient end-to-end low-bit LLM inference on NPUs, significantly improving speed and energy efficiency by overcoming hardware limitations.

Contribution

The paper proposes a novel unified table layout and tiling strategy that allows low-bit LLM inference to be performed entirely on NPUs, eliminating the need for CPU offloading.

Findings

1.4x speedup in prefill phase

3.1x speedup in decoding phase

84% energy savings compared to baseline methods

Abstract

Large language models (LLMs) are increasingly deployed on customer devices. To support them, current devices are adopting SoCs (System on Chip) with NPUs (Neural Processing Unit) installed. Although high performance is expected, LLM inference on NPUs is slower than its CPU counterpart. The reason is that NPUs have poor performance on computations other than GEMM, like dequantization. Current works either disaggregate prefill on the NPUs and decoding on the CPUs, or put both on the NPUs but with an accuracy loss. To solve this issue, based on the insight that low-bit can enable target computation encoded within an acceptably sized table, we propose table lookup to subsume hardware operations otherwise unsupported. To realize this, we overcome the conflicting hardware behavior of prefill and decoding to design a unified table layout and tiling through (1) fused two-level table-based dequantization and (2) concurrency-hierarchy-guided tiling. Based on that, we implement the prefill phase by three-stage pipeline and map the table-lookup-based decoding to NPU's vector units. Results show 1.4x and 3.1x speedup for prefill and decoding respectively, and 84% energy savings compared to the baseline NPU methods. The code is available at https://github.com/microsoft/T-MAC/tree/main/t-man.