Accelerating a Triton Fused Kernel for W4A16 Quantized Inference with SplitK work decomposition
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866913241045663744 |
|---|---|
| author | Hoque, Adnan Wright, Less Yang, Chih-Chieh Srivatsa, Mudhakar Ganti, Raghu |
| author_facet | Hoque, Adnan Wright, Less Yang, Chih-Chieh Srivatsa, Mudhakar Ganti, Raghu |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2402_00025 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | Accelerating a Triton Fused Kernel for W4A16 Quantized Inference with SplitK work decomposition Hoque, Adnan Wright, Less Yang, Chih-Chieh Srivatsa, Mudhakar Ganti, Raghu Distributed, Parallel, and Cluster Computing Artificial Intelligence 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. |
| title | Accelerating a Triton Fused Kernel for W4A16 Quantized Inference with SplitK work decomposition |
| topic | Distributed, Parallel, and Cluster Computing Artificial Intelligence |
| url | https://arxiv.org/abs/2402.00025 |