Accelerating Block Low-Rank Foundation Model Inference on Memory-Constrained GPUs | Proceedings of the 35th International Symposium on High-Performance Parallel and Distributed Computing

AbstractRecent 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 × speedups and 3 × 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.1 IntroductionLarge-scale transformer-based foundation models have achieved remarkable success across language understanding, image classification, and generative tasks [5, 13, 18, 35, 44]. However, their rapid growth in size is increasingly outpacing the capacity of available hardware. For example, Llama-70B [44] requires over 140 GB of memory simply to load its weights in half-precision format, yet some of today’s most powerful commercial GPUs provide only 80 GB. Beyond sheer memory constraints, the reliance of transformer models with dense matrix multiplications introduces significant computational and memory bandwidth bottlenecks during inference. These challenges limit the deployment of models at scale and also their accessibility on resource-constrained devices.A widely adopted strategy to address these bottlenecks is to approximate weight matrices using low-rank factorizations [19, 24, 48]. By representing a dense weight matrix as the product of two smaller matrices, the computational and memory complexity of linear layers can be substantially reduced.

However, traditional low-rank decompositions often exhibit sharp accuracy degradation at high compression ratios [25], which limits their practicality. To overcome this limitation, structured decompositions such as Monarch [11] and BLAST [26] have been proposed. They leverage block low-rank (BLR) structures to capture the underlying representation of weight matrices more effectively, thereby better preserving accuracy while offering memory savings and computational reduction.Despite these algorithmic advances, the expected end-to-end speedups often fail to materialize in practice on GPUs. Performance on modern GPUs is governed by a roofline model [46, 47] that balances memory bandwidth, peak computational throughput, and arithmetic intensity (i.e., the ratio of operations to memory traffic). Figure 1 illustrates the roofline model of an NVIDIA A40 GPU for 16-bit brain floating-point (BF16) operations. While structured low-rank (LR) decompositions reduce the nominal compute requirements, we first show that they also introduce additional intermediate data movement, particularly in long-sequence scenarios such as the pre-fill stage of large language model (LLM) inference [22, 23]. This shift can move linear layers from the compute-bound regime into the memory-bound regime, creating a gap between algorithmic promise and system-level reality. The issue arises specifically for devices with limited memory subsystems, such as edge GPUs with small L2 caches (4–6 MB) and DRAM based on DDR technology (Jetson Orin Nano) as well as datacenter GPUs (A40). In such cases, (B)LR decompositions paradoxically degrade performance, despite reducing floating-point operations (FLOP) and model size.Figure 1:Roofline model of 16-bit brain floating-point (BF16) operations for a NVIDIA A40 GPU.In this work, we analyze the performance of LR, Monarch, and BLAST matrix multiplications for efficient transformer-based foundation model inference, with a particular emphasis on long sequences. We identify and characterize key bottlenecks arising from data movement, suboptimal memory layouts, and compiler limitations, and we introduce optimized implementations using Triton [43], an open-source intermediate language designed for writing efficient GPU kernels.

Our proposed kernels exploit partial fusion, operation reordering, and tailored memory layouts to mitigate the overheads of (B)LR structured matrix multiplications in multi-token inference. Through extensive evaluation, we demonstrate that these optimizations deliver substantial performance gains across diverse transformer models, including GPT2-S, Llama-1/7B, and DiT-XL/2 on both server-grade and edge-class GPUs. Our results establish that BLR-based model compression, when paired with hardware-aware optimizations, offers a viable path toward practical deployment of foundation models in resource-constrained environments. Overall, our contributions can be summarized as follows:•We provide the first systematic roofline analysis of BLR (Monarch, BLAST) matrix multiplications, showing that whi-le they reduce FLOP, block structures introduce intermediate data movement and uncover PyTorch compiler limitations that push multi-token inference into the memory-bound regime compared to traditional low-rank and dense methods.•We design Triton kernels with partial fusion, operation reordering, and tensor-core-friendly layouts that eliminate redundant data movement and restore efficiency to BLR inference.•We release1 our optimized kernels and benchmark on server and edge-grade GPUs, demonstrating up to 3.76 × speedups and 3 × model compression over PyTorch CUDA dense baselines with compiler optimizations, fostering reproducibility and rendering structured compression practical.2 Background2.1 Weight StructuresWe introduce the weight matrix structures considered in this work, together with their computational properties and modeling capabilities. Let i, o, n, and r denote the number of input features, output features, sequence length, and rank, respectively. The input to all linear layers is \({\boldsymbol {X}}\in \mathbb {R}^{n\times i}\), and we assume r ≪ i, o.2.1.1 Dense.A dense weight matrix \({\boldsymbol {W}}\in \mathbb {R}^{i\times o}\) has i × o parameters, and the corresponding linear layer Y = XW requires n × i × o FLOP. Dense matrices can represent arbitrary linear maps and therefore provide the highest expressiveness for foundation models.2.1.2 Low-Rank (LR).Dense weights can be factorized as W = VU, where \({\boldsymbol {V}}\in \mathbb {R}^{i\times r}\) and \({\boldsymbol {U}}\in \mathbb {R}^{r\times o}\).

Instead of materializing W, the factorization is stored and used directly in computation. This reduces parameter count to r(i + o) and computation to nr(i + o) FLOP. The low-rank assumption can cause accuracy degradation if the chosen rank r does not capture the true structure of W. In practice, r ≪ i, o is selected such that r(i + o) < io, yielding both memory and computational savings [20, 24, 45].Figure 2:Monarch (left) and BLAST (right) parametrization and linear layer execution for b1 = b2 = 3 blocks and rank r = 4.Table 1:MethodSmallMediumLargeViT-BGPT2-SDiT-XL/2Llama-3.2-1BLlama-7B(CF = 3 ×)(CF = 1.85 ×)(CF = 2 ×)(CF = 2 ×)(CF = 2 ×)ImageNetWikiText-103FID (↓)sFID (↓)IS (↑)WikiText-2Avg. 0-shotWikiText-2Avg. 0-shotAccuracy (%)Perplexity (↓)Perplexity (↓)Accuracy (%)Perplexity (↓)Accuracy (%)Dense78.720.29.626.85121.5011.5756.549.3766.07BLAST79.320.710.456.72111.0520.1046.3714.2156.23Monarch79.221.1---22.1744.3519.5449.78Low-Rank78.921.748.0711.4426.0921.9244.7126.3348.40Accuracy of foundation models using different (B)LR model compression factors (CF).2.1.3 Monarch.

Introduced by Dao et al. [11], Monarch divides a dense weight into b2 × b1 blocks, yielding a BLR representation2 with uniform per-block rank r′. Each block Wl, k is factorized as \begin{equation*} {\boldsymbol {W}}_{l,k} = {\boldsymbol {V}}_{l,k} {\boldsymbol {U}}_{l,k} \in \mathbb {R}^{p\times q}, \quad \begin{array}{l}l \in \lbrace 1,\dots ,b_1\rbrace , \\ k \in \lbrace 1,\dots ,b_2\rbrace , \\ p = i/b_1, \quad q = o/b_2 \end{array} \end{equation*} A Monarch layer has b1b2r′(p + q) parameters. The k -th output block Yk is computed as \begin{equation*} {\boldsymbol {Y}}_k = \sum \nolimits _l {\boldsymbol {X}}_l {\boldsymbol {W}}_{l,k}, \qquad {\boldsymbol {X}}_l \in \mathbb {R}^{n \times p}, \qquad {\boldsymbol {Y}}_k \in \mathbb {R}^{n \times q}, \end{equation*} which requires nb1b2r′(p + q) FLOP in total.Figure 2 (left) illustrates the execution of a Monarch layer for b1 = b2 = 3, where a permutation (b1⇔b2) separates two batched matrix multiplications (\(\mathtt {bmm}\)). In practice, the common setting b1 = b2 = b with r = r′b recovers the same complexity as low-rank layers, r(i + o) parameters and nr(i + o) FLOP. Monarch weights are stored as tensors : \(\mathcal {V} \in \mathbb {R}^{b_1\times (r^\prime b_2) \times p}\) and \(\mathcal {U} \in \mathbb {R}^{b_2\times q \times (b_1 r^\prime)}\).2.1.4 BLAST.Introduced by Lee et al. [26], BLAST represents a weight matrix using a generalized BLR structure.

Unlike Monarch, each block Wl, k shares a pair of matrices Vl and Uk while retaining a unique diagonal matrix Sl, k such that the block factorization becomes \begin{equation*} {\boldsymbol {W}}_{l,k} = {\boldsymbol {V}}_l {\boldsymbol {S}}_{l,k} {\boldsymbol {U}}_k, \qquad {\boldsymbol {V}}_l \in \mathbb {R}^{p\times r}, \; {\boldsymbol {S}}_{l,k} \in \mathbb {R}^{r\times r}, \; {\boldsymbol {U}}_k \in \mathbb {R}^{r \times q}. \end{equation*} This structure generalizes multiple families of structured low-rank matrices. In particular, low-rank and Monarch layers can be recovered by setting the values of Sl, k appropriately for all l and k.A BLAST layer has r(p + q + b1b2) parameters. The k -th output block Yk is computed as \begin{equation*} {\boldsymbol {Y}}_k = \Bigl (\sum \nolimits _l ({\boldsymbol {X}}_l {\boldsymbol {V}}_l) {\boldsymbol {S}}_{l,k}\Bigr) {\boldsymbol {U}}_k, \qquad {\boldsymbol {X}}_l \in \mathbb {R}^{n\times p}, \qquad {\boldsymbol {Y}}_k \in \mathbb {R}^{n\times q}, \end{equation*} which requires nr(p + q + b1b2) FLOP in total.Typically, b1 = b2 = b ≤ 16, which yields r(i + o + b2) parameters and nr(i + o + b2) FLOP. Since b ≪ i, o, BLAST achieves the same asymptotic savings as low-rank and Monarch layers using a slightly higher compression ratio. Figure 2 (right) illustrates the execution of a BLAST layer for b1 = b2 = 3 and r = 4. In practice, BLAST parameters are stored as \(\mathcal {V} \in \mathbb {R}^{b_1\times p\times r}\), \(\mathcal {S} \in \mathbb {R}^{b_1\times b_2 \times r}\), and \(\mathcal {U} \in \mathbb {R}^{b_2\times r \times q}\).2.1.5 Motivation.