Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Beyond GEMM-Centric NPUs: Enabling Efficient Diffusion LLM Sampling

Binglei Lou, Haoran Wu, Yao Lai, Jiayi Nie, Can Xiao, Xuan Guo, Rika Antonova, Robert Mullins, Aaron Zhao

TL;DR

This work identifies diffusion LLM sampling as a distinct, memory-bound workload that becomes a bottleneck even as model-phase latency improves. It introduces d-PLENA, a vector-scalar NPU extension with a decoupled mixed-precision memory hierarchy and hardware primitives tailored to non-GEMM sampling tasks, plus ISA extensions for fast ArgMax, Top-$k$, and masked updates. Through cycle-accurate simulation and RTL verification, it demonstrates up to $2.53\times$ speedup over an RTX A6000 on diffusion-sampling workloads and provides a co-design framework linking PyTorch configurations to hardware execution. The results highlight architectural patterns that can be generalized to diffusion-based inference and pave the way for niche NPUs optimized for reduction-heavy, memory-irregular workloads in diffusion LLMs.

Abstract

Diffusion Large Language Models (dLLMs) introduce iterative denoising to enable parallel token generation, but their sampling phase displays fundamentally different characteristics compared to GEMM-centric transformer layers. Profiling on modern GPUs reveals that sampling can account for up to 70% of total model inference latency-primarily due to substantial memory loads and writes from vocabulary-wide logits, reduction-based token selection, and iterative masked updates. These processes demand large on-chip SRAM and involve irregular memory accesses that conventional NPUs struggle to handle efficiently. To address this, we identify a set of critical instructions that an NPU architecture must specifically optimize for dLLM sampling. Our design employs lightweight non-GEMM vector primitives, in-place memory reuse strategies, and a decoupled mixed-precision memory hierarchy. Together, these optimizations deliver up to a 2.53x speedup over the NVIDIA RTX A6000 GPU under an equivalent nm technology node. We also open-source our cycle-accurate simulation and post-synthesis RTL verification code, confirming functional equivalence with current dLLM PyTorch implementations.

Beyond GEMM-Centric NPUs: Enabling Efficient Diffusion LLM Sampling

TL;DR

This work identifies diffusion LLM sampling as a distinct, memory-bound workload that becomes a bottleneck even as model-phase latency improves. It introduces d-PLENA, a vector-scalar NPU extension with a decoupled mixed-precision memory hierarchy and hardware primitives tailored to non-GEMM sampling tasks, plus ISA extensions for fast ArgMax, Top-, and masked updates. Through cycle-accurate simulation and RTL verification, it demonstrates up to speedup over an RTX A6000 on diffusion-sampling workloads and provides a co-design framework linking PyTorch configurations to hardware execution. The results highlight architectural patterns that can be generalized to diffusion-based inference and pave the way for niche NPUs optimized for reduction-heavy, memory-irregular workloads in diffusion LLMs.

Abstract

Diffusion Large Language Models (dLLMs) introduce iterative denoising to enable parallel token generation, but their sampling phase displays fundamentally different characteristics compared to GEMM-centric transformer layers. Profiling on modern GPUs reveals that sampling can account for up to 70% of total model inference latency-primarily due to substantial memory loads and writes from vocabulary-wide logits, reduction-based token selection, and iterative masked updates. These processes demand large on-chip SRAM and involve irregular memory accesses that conventional NPUs struggle to handle efficiently. To address this, we identify a set of critical instructions that an NPU architecture must specifically optimize for dLLM sampling. Our design employs lightweight non-GEMM vector primitives, in-place memory reuse strategies, and a decoupled mixed-precision memory hierarchy. Together, these optimizations deliver up to a 2.53x speedup over the NVIDIA RTX A6000 GPU under an equivalent nm technology node. We also open-source our cycle-accurate simulation and post-synthesis RTL verification code, confirming functional equivalence with current dLLM PyTorch implementations.
Paper Structure (15 sections, 2 equations, 5 figures, 4 tables, 2 algorithms)

This paper contains 15 sections, 2 equations, 5 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background and Motivation
  3. PLENA Accelerator
  4. Diffusion-based LLMs and Blocked Sampling
  5. Methodology
  6. Proposed NPU Architecture
  7. Hardware-Friendly Sampling Flow
  8. ISA Extensions for dLLM and Phased Sampling Execution
  9. Co-design and Verification Framework
  10. Experimental Results
  11. Experimental Setup
  12. Latency and Memory Utilization
  13. Execution Efficiency and Bottleneck Analysis
  14. Post-Synthesis Area and Power Evaluation
  15. Conclusion

Figures (5)

  • Figure 1: Latency breakdown of the LLaDA model on an A6000 GPU, evaluated across a broad parameter space: batch sizes from 1 to 32, denoising steps from 1 to 32, generation lengths from 64 to 1024 tokens, and block sizes from 8 to 64.
  • Figure 2: Diffusion Large Language Model. For each timestep, the transformer predicts token confidences and updates the top-$k$ masked positions to progressively refine the sequence.
  • Figure 3: Proposed NPU architecture for diffusion sampling.
  • Figure 4: Hardware--software co-design and verification workflow with cross-layer accuracy and latency checks.
  • Figure 5: Latency and Memory utilization of diffusion sampling.