Memory-Efficient Acceleration of Block Low-Rank Foundation Models on Resource Constrained GPUs

TL;DR

This work analyzes block low-rank (BLR) compressions, notably Monarch and BLAST, for large transformer models on memory-constrained GPUs. Through a systematic roofline analysis, it reveals that multi-token inference often becomes memory-bound due to intermediate data movement and layout/permutation overheads. To address this, the authors design memory-efficient Triton kernels with partial fusion and optimized memory layouts, achieving up to 3.76× end-to-end speedups and 3× model compression on edge and server GPUs, while supporting models like Llama-7B, GPT2-S, DiT-XL/2, and ViT-B. The results demonstrate that, when paired with hardware-aware kernel optimizations, BLR-based compression can enable practical deployment of foundation models on resource-constrained devices, narrowing the gap with dense baselines and outperforming existing BLR baselines in many settings.

Abstract

Recent advances in transformer-based foundation models have made them the default choice for many tasks, but their rapidly growing size makes fitting a full model on a single GPU increasingly difficult and their computational cost prohibitive. Block low-rank (BLR) compression techniques address this challenge by learning compact representations of weight matrices. While traditional low-rank (LR) methods often incur sharp accuracy drops, BLR approaches such as Monarch and BLAST can better capture the underlying structure, thus preserving accuracy while reducing computations and memory footprints. In this work, we use roofline analysis to show that, although BLR methods achieve theoretical savings and practical speedups for single-token inference, multi-token inference often becomes memory-bound in practice, increasing latency despite compiler-level optimizations in PyTorch. To address this, we introduce custom Triton kernels with partial fusion and memory layout optimizations for both Monarch and BLAST. On memory-constrained NVIDIA GPUs such as Jetson Orin Nano and A40, our kernels deliver up to $3.76\times$ speedups and $3\times$ model size compression over PyTorch dense baselines using CUDA backend and compiler-level optimizations, while supporting various models including Llama-7/1B, GPT2-S, DiT-XL/2, and ViT-B. Our code is available at https://github.com/pabillam/mem-efficient-blr .

Figures (11)

• Figure 1: Roofline model of BF16 operations for NVIDIA A40 GPU.
• Figure 2: Monarch (left) and BLAST (right) weight parametrization and linear layer execution for $b_1 = b_2 = 3$ blocks and rank $r = 4$.
• Figure 3: Performance of Llama-7B layers with low-rank methods versus dense, for single-token (left) and multi-token (right) inference on an NVIDIA A40 GPU.
• Figure 4: Roofline and runtime estimation of $\mathtt{Q/K/V/O_{proj}}$ during multi-token inference on A40.
• Figure 5: Pseudo-code for fused permutation and $\mathtt{bmm}$ Monarch kernel (❷).