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

HEMERA is presented, a heterogeneous memory-centric accelerator for efficient Mamba-2 inference that reformulates the matrix-form SSD computation into an algebraically equivalent streaming-recursive dataflow that avoids quadratic intermediate storage while preserving the original computation.

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

Celty, a co-designed sparse format, GPU kernel, and SIMT microarchitecture for efficient spMspV in LLM inference is proposed, which introduces a Run-Length Compressed CSC format that enables vectorized loading of compressed weight columns and exploits both sparsity sources to skip unnecessary memory accesses.

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

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

This work introduces a tile-aware scheduling framework for efficient sparse Vision Transformer execution on GPUs and introduces 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

SA-Scheduler is presented, a structure-derived bottleneck-aware scheduling framework for multitasking MCM-GPUs that determines chip placement without hardware modification or runtime bottleneck profiling, and 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