Accelerating a Triton Fused Kernel for W4A16 Quantized Inference with SplitK work decomposition

2024/01/05 by Hoque, Adnan, Wright, Less, Yang, Chih-Chieh +2 · 2 citations
#Artificial Intelligence (cs.AI) #Distributed #FOS: Computer and information sciences #Parallel #and Cluster Computing (cs.DC)

paper · doi:10.48550/arxiv.2402.00025

Abstract

We propose an implementation of an efficient fused matrix multiplication kernel for W4A16 quantized inference, where we perform dequantization and GEMM in a fused kernel using a SplitK work decomposition. Our implementation shows improvement for the type of skinny matrix-matrix multiplications found in foundation model inference workloads. In particular, this paper surveys the type of matrix multiplication between a skinny activation matrix and a square weight matrix. Our results show an average of 65% speed improvement on A100, and an average of 124% speed improvement on H100 (with a peak of 295%) for a range of matrix dimensions including those found in a llama-style model, where m < n = k.

Cited by

• Celty: SpMspV GPU Kernel and SIMT Co-Design for Efficient Dual-Sparse LLM Inference
• Tilus: A Tile-Level GPGPU Programming Language for Low-Precision Computation

Related

• LUT-GEMM: Quantized Matrix Multiplication based on LUTs for Efficient Inference in Large-Scale Generative Language Models
• Fast Matrix Multiplications for Lookup Table-Quantized LLMs
• A Speed Odyssey for Deployable Quantization of LLMs
• QUICK: Quantization-aware Interleaving and Conflict-free Kernel for efficient LLM inference
• Tender: Accelerating Large Language Models via Tensor Decomposition and Runtime Requantization