Skip to content

MLX: Multi-Layer Execution for Structured LLM Workload Acceleration on Spatial Architectures

Jun 2026 · International Symposium on Computer Architecture · pp. 1-16 · 0 citations · 78 references

Abstract

Structured sparsity is a promising approach to scaling large-language-model (LLM) inference, but existing forms such as butterfly-structured sparse projections and transformations often map inefficiently to GPUs due to deep stage dependencies and limited bulk parallelism. This paper presents MLX, an algorithm–architecture co-design for structured LLM inference. MLX couples semantic-aware FFT compression and hierarchical sparse projections with spatial dataflow execution, enabling staged structured operators to run efficiently on compact arrays. MLX defines Closed Dependency Components (CDCs) to capture deterministic forward-only dataflow regions that can be folded across layers and pipelined on compact arrays. It then realizes CDCs through a multi-layer execution architecture with bounded-hop skip-hop routing, tag-based scheduling, and decoupled compute/transfer pipelines to overlap communication and computation across deep operators. We prototype MLX in 12 nm and show that it achieves $3.2 \times$ hardware speedup and $3.1 \times$ energy savings over Jetson Xavier. A transformer-specialized reduced design further delivers up to 5.7× speedup over prior sparse accelerators. MLX also scales nearly linearly to $8 \times 8$ meshes and remains effective for long sequences from $\mathbf{1 K}$ to 4 K, demonstrating that structured operator semantics can be translated into efficient spatial execution for sparse LLMs.

View source

Similar papers

Preprint Jul 2026

HEMERA: A Heterogeneous Memory-Centric Accelerator with Recursive Dataflow for Edge-Constrained State-Space-Duality Models Inference

Structured State Space Models (SSMs), such as Mamba, enable efficient long-sequence modeling with linear time complexity. Recent implementations realize this capability through Structured State Space Duality (SSD), which transforms recursive state evolution into matrix-form computations. However, SSD introduces substantial system-level overheads, including quadratic intermediate materialization, irregular data movement, and prefix-dependent execution, leading to excessive memory traffic and bandwidth demand on conventional architectures. Although prior accelerators mitigate these overheads through optimized dataflows or compute-in-memory techniques, they largely retain matrix-oriented SSD execution and cannot simultaneously avoid quadratic intermediate storage and efficiently map dependency-bound state propagation. This paper presents HEMERA, a heterogeneous memory-centric accelerator for efficient Mamba-2 inference. Rather than directly executing the matrix-form SSD computation, HEMERA reformulates it into an algebraically equivalent streaming-recursive dataflow that avoids quadratic intermediate storage while preserving the original computation. The resulting heterogeneous execution paradigm maps dense linear operations onto in-memory computing units and recursive state updates onto a dedicated streaming engine. Across Mamba-2 models ranging from 130M to 2.8B, HEMERA achieves average latency speedups of 1.4x-3.6x and energy-efficiency improvements of 12.2x-27.0x over the official optimized fused Mamba-2 kernel on NVIDIA A100. It further reduces the average SSD-related execution-time ratio across model scales to 14.12% during long-sequence inference, demonstrating its potential for efficient deployment under edge constraints.

Hao Ding, Ling Liang, Ruitong Qiao et al. · 0 citations
Preprint Aug 2026

Celty: SpMspV GPU Kernel and SIMT Co-Design for Efficient Dual-Sparse LLM Inference

Large Language Models (LLMs) increasingly rely on sparsity to reduce inference cost, but most prior work targets a single sparsity source-either weight or activation-and optimizes for batched multi-user inference. Dual-sparsity, which combines unstructured weight pruning with runtime activation sparsity, offers a compelling tradeoff among model size, accuracy, and latency for single-user decoding, but formulates as a Sparse Matrix-Sparse Vector (spMspV) workload that existing GPU kernels handle poorly. We propose Celty, a co-designed sparse format, GPU kernel, and SIMT microarchitecture for efficient spMspV in LLM inference. At the kernel level, Celty introduces a Run-Length Compressed CSC (RLC-CSC) format that enables vectorized loading of compressed weight columns and exploits both sparsity sources to skip unnecessary memory accesses, with shared memory used for scattered partial-product accumulation. At the microarchitecture level, the Celty Sparse SIMT Core integrates a pipelined RLC decoder to eliminate software-level index reconstruction and repurposes local register files for conflict-free accumulation-operating directly on the same RLC-CSC format without data layout changes. The Celty GPU kernel achieves up to 2.8x speedup over cuBLAS and 2.4x over Flash-LLM. With the Sparse SIMT Core, speedups reach up to 5.3x over cuBLAS at 70% dual-sparsity.

Ruokai Yin, Priyadarshini Panda · 0 citations
Book Open access Jul 2026

Toward Low-Latency and Memory-Efficient Deployment of Irregular Sparse Deep Learning Workloads

Our work introduces a tile-aware scheduling framework for efficient sparse Vision Transformer execution on GPUs. Sparse attention reduces the cost of high-resolution Vision Transformers, but its irregular masks produce blocks with diverse sizes, densities, and locations. Existing FlashAttention-style kernels rely on fixed tile configurations and cannot fully exploit these sparse patterns, leading to wasted computation and underutilized GPU resources. Our framework bridges this gap through four key techniques: sparse attention is represented as an adjacency matrix, structure-aware reordering algorithms improve locality, locally dense blocks are extracted as scheduling units, and offline profiling with integer linear programming (ILP) selects hardware-feasible tile assignments. Results show that our inference scheduler achieves up to 2.13 × end-to-end speedup over fixed-tile FlashAttention and up to 4.6 × speedup in high-resolution images. We further introduce a training-aware extension that reuses the inference tile schedule and augments it with backward computation and activation-memory strategies.

Changxin Li · 0 citations
Open access Aug 2026

Structure-Derived Bottleneck-Aware Scheduling for Multitasking MCM-GPUs

Multi-chip-module (MCM) GPUs improve scalability by integrating multiple GPU chiplets, but spatial multitasking introduces severe performance interference due to non-uniform memory access and asymmetric inter-chip bandwidth. Existing GPU multitasking schedulers rely on runtime profiling metrics to guide placement decisions. These reactive approaches observe the consequence of interference rather than its structural cause, vary across inputs and co-running workloads, and incur non-trivial overhead. This paper presents SA-Scheduler, a structure-derived bottleneck-aware scheduling framework for multitasking MCM-GPUs. Instead of profiling execution behavior, SA-Scheduler derives hardware bottleneck regimes directly from kernel structure prior to execution. Through offline architecture-aware modeling, we extract iteration-normalized features such as warp-level working-set footprint and global memory intensity (GMI). A fine-tuned lightweight code model (Qwen-2.5-Coder-1.5B) assists semantic structural extraction for CUDA kernels. These features are translated via an analytical bottleneck model into distinct interference regimes (C, M1, M2, and M3), characterizing how kernels interact with hierarchical and asymmetric memory resources under co-execution. Based on the inferred regimes, SA-Scheduler determines chip placement without hardware modification or runtime bottleneck profiling. For streaming multiprocessor (SM) allocation and optional power-gating decisions, SA-Scheduler further employs lightweight counter-based feedback to tune resource usage under bounded overhead. Across 105 two-program workloads, SA-Scheduler improves system throughput (STP) by 17.63% and 10.48% over intra-chip and inter-chip scheduling, respectively, while increasing fairness by up to 33.27%. An energy-aware Power Mode further reduces average power by up to 25.39% under bounded throughput degradation. These results demonstrate that structure-derived bottleneck inference provides a principled and scalable foundation for multitasking on future MCM-GPUs.

Tiejian Zhang, Guangda Zhang, Lu Wang et al. · 0 citations
Book Open access Jul 2026

SPPO: Making Million-Token LLM Training Practical on Modest GPU Clusters

In recent years, Large Language Models (LLMs) have exhibited remarkable capabilities, driving advancements in real-world applications. However, training LLMs on increasingly long input sequences imposes significant challenges due to high GPU memory and computational demands. Existing solutions face two key limitations: (1) memory reduction techniques, such as activation recomputation and CPU offloading, compromise training efficiency; (2) distributed parallelism strategies require excessive GPU resources, limiting the scalability of input sequence length. To address these gaps, we propose Adaptive Sequence Pipeline Parallel Offloading (SPPO), a novel framework that optimizes memory and computational resource efficiency for long-sequence LLM training. SPPO introduces adaptive offloading, leveraging sequence-aware offloading, and two-level activation management to reduce GPU memory consumption without degrading the training efficiency. Additionally, it develops an adaptive pipeline scheduling approach with a heuristic solver and multiplexed sequence partitioning to improve computational resource efficiency. We perform comprehensive experiments to demonstrate the superiority of SPPO: it achieves up to 3.38 × throughput improvement over state-of-the-art methods like Megatron-LM, DeepSpeed and TeraPipe, realizing efficient training of a 7B LLM with sequence lengths of up to 4M tokens on only 128 A100 GPUs.

Qiaoling Chen, Shenggui Li, Wei Gao et al. · 0 citations
Book Open access Jul 2026

Aurora: A Disaggregated GPU-PNM-PIM System for High-Throughput Mixed-Length LLM Inference

Modern large language model (LLM) inference is increasingly dominated by memory-bound operations, making processing-near-memory with processing-in-memory (PNM-PIM) an attractive approach for accelerating the decode phase. However, recent long-context LLMs adopt interleaved local and global attention (ILGA), which introduces heterogeneous execution characteristics across transformer blocks. In realistic serving environments, this heterogeneity is further amplified by mixed-length requests, where inputs with widely different context lengths are processed concurrently. These trends break a key assumption underlying existing PNM-PIM systems that transformer blocks exhibit similar latency and resource demands and can be efficiently mapped to a uniform pipeline. Under ILGA and mixed-length workloads, this assumption no longer holds, leading to severe pipeline imbalance and low utilization in prior PNM-PIM designs. Moreover, PNM-PIM–only systems struggle to efficiently support long-context prefill, which remains compute-intensive and is better suited for GPU execution. In this paper, we propose Aurora, a GPU–PNM–PIM disaggregated system designed to efficiently serve mixed-length LLM inference under ILGA. Aurora introduces an ILGA-aware multi-path PNM-PIM pipeline that explicitly accounts for block-level heterogeneity and request-length diversity, improving pipeline utilization without overprovisioning tensor parallelism. Also, Aurora further adopts a stream-oriented Softmax design to reduce stage-level latency imbalance during decoding. To enable end-to-end inference under disaggregation, Aurora treats KV cache transfer as a pipeline operation and coordinates request issuance and scheduling to avoid transfer-induced stalls. Our experimental results show that Aurora achieves up to 8.5 × and 2.2 × higher throughput than the GPU-only baseline and prior PNM-PIM systems, respectively, in end-to-end Llama4-Scout inference, while requiring a comparable or smaller number of devices.

Hyeonu Kim, Seunghyuk Yu, Minjeoung Kim et al. · 0 citations