PaperlayerView exact revision on arXiv

Source-linked AI summary

LLM Inference on IMC-NoC Architecture with Balanced Dataflow and Fine-Grained Parallelism

Yimin Wang, Yue Jiet Chong, Xuanyao Fong

arXiv:2609.00857v1cs.AR

TL;DR

LLM inference stresses memory bandwidth, computation, and communication because static weights coexist with dynamic data and large models require scale-up interconnects. The paper proposes LEAP, a hardware–software co-design integrating specialized compute and communication resources with partitioning, mapping, scheduling, and prefill-decode disaggregation. Evaluations report 24× to 72× energy-efficiency gains over commercial GPU platforms.

  • Problem

    LLM inference combines static weights with dynamic runtime data and requires large-scale architectures whose on-chip communication can bottleneck execution.

  • Method

    The paper develops LEAP, integrating IMC, NMC, and INC with partitioning, mapping, scheduling, and deterministic dataflow routing.

  • Results

    24× to 72× energy-efficiency improvement over commercial GPU platforms is reported for LEAP.

  • Takeaways & Limitations

    LEAP provides a unified compute–memory–communication fabric with phase-specific organization for balancing LLM inference throughput and latency.

Abstract

from arXiv · show

LLM inference has become an essential service, yet it imposes unprecedented demands on memory bandwidth, computational density, and communication efficiency. While IMC is a promising solution to the memory wall issue, the heterogeneous data dynamicity of LLM requires complementary resources to handle intermediate data generated during run-time. Furthermore, the massive number of parameters in LLM necessitates scale-up architectures where on-chip data movement is often the primary performance bottleneck. This article presents a hardware-software co-design framework that unifies distributed compute, memory, and communication into a seamless processing-communication fabric. On the hardware side, we propose a scalable architecture, named LEAP, that integrates IMC PE, NMC PE, and INC. This allows each hardware layer to execute specialized tasks: IMC for static weights, NMC for dynamic data, and INC for partial result reduction. On the software side, we introduce a partitioning, mapping, and scheduling framework optimized for key metrics in LLM serving, including throughput and latency. To address the distinct computational intensities of the prefill and decode phases, we present a prefill-decode disaggregation approach that dynamically reconfigures PE organizations to maximize resource utilization. Compared to commercial GPU platforms, the proposed architecture provides a throughput and an energy efficiency improvement of $\geq{}1.52\times$ and $24.91\times$, respectively.

I. INTRODUCTION

The paper presents LEAP, a hardware–software co-design framework addressing heterogeneous LLM data, scale-up communication, and software orchestration challenges. It combines specialized hardware resources with partitioning, mapping, scheduling, and prefill-decode disaggregation strategies.

  • Motivation: LLM inference combines static pretrained weights with dynamic runtime intermediate data, creating demands beyond IMC-only architectures.IMC processes static data effectively, while complementary resources handle dynamic activations.
  • Motivation: Scaling IMC processing across thousands of spatial PEs burdens the NoC with broadcasting inputs and aggregating partial results.INC is identified as a way to integrate arithmetic reduction into NoC routers.
  • Architecture: LEAP unifies IMC PEs, NMC PEs, and INC in a scalable processing-communication fabric for specialized matrix operations.IMC handles static weights, NMC dynamic data, and INC partial-result reduction.
  • Software framework: The framework coordinates partitioning, spatial mapping, deterministic dataflow routing, and scheduling across operation, phase, and request levels.Its optimization targets key LLM serving metrics, including throughput and latency.
  • Prefill-decode orchestration: Prefill-decode disaggregation reconfigures PE organizations for the distinct phases and exposes tradeoffs between throughput and hardware resource provisioning.The paper extends its design space with dedicated phase handling and heterogeneous IMC/NMC configurations.
  • Evaluation: 24× to 72× energy-efficiency improvement over commercial GPU platforms is demonstrated for LEAP.The evaluations characterize performance and energy-efficiency frontiers.

II. PRELIMINARIES

The preliminaries characterize LLM operations, hardware scaling constraints, and the distinct serving requirements of prefill and decode. They motivate LEAP’s unified architecture and its aggregated and disaggregated execution modes.

  • LLM hardware accelerators: LLM decoder workloads contain dynamic-dynamic matrix multiplications in attention and dynamic-static matrix multiplications in attention and FFNs.These operation types are distinguished by whether their operands are runtime-generated or pretrained static matrices.
  • LLM hardware accelerators: 128 × 128-cell IMC crossbars require large LLM weight matrices to be partitioned across spatial arrays, creating broadcast and partial-result aggregation traffic.Imbalanced partial-result aggregators can bottleneck overall performance.
  • LLM hardware accelerators: LEAP integrates IMC, NMC, and INC to address dynamic-static operations, dynamic-dynamic operations, and partial-result reduction, respectively.The architecture is presented as a flexible, scalable alternative to accelerators custom-designed for specific models.
  • LLM inference serving: Interactive services prioritize small TTFT, whereas longform generation prioritizes small TPOT.The relevant metric depends on the application.
  • LLM inference serving: Prefill is primarily compute-intensive while decode is constrained by memory bandwidth, motivating temporal balancing and spatial disaggregation strategies.Hybrid strategies can introduce tunable variables at operation, phase, or request levels.
  • LLM inference serving: LEAP-A couples prefill and decode on shared hardware, while LEAP-D assigns them to dedicated hardware partitions.The modes implement aggregated and disaggregated approaches using LEAP’s distributed memory organization.

III. FRAMEWORK OVERVIEW

LEAP combines distributed processing elements with a mesh NoC and supports distinct hardware organizations for aggregated and disaggregated execution. Its framework covers weight and dynamic-matrix partitioning and mapping across the architecture.

  • III. FRAMEWORK OVERVIEW: LEAP combines PE arrays with a mesh NoC, using each PE-router pair as a macro.The router integrates NMC scratchpad and matrix-multiplication logic with INC interconnection and reduction resources.
  • III. FRAMEWORK OVERVIEW: IMC PEs store static data and perform DSMMs, while router scratchpads store dynamic data and IRCUs perform DDMMs and partial-result reduction.
  • III. FRAMEWORK OVERVIEW: LEAP-A uses homogeneous macros, whereas LEAP-D separates prefill and decode regions with different macro organizations.The supplied figures cover weight-matrix partitioning and mapping for LEAP-A and the prefill region of LEAP-D, plus dynamic-matrix mapping and whole-model placement.

B. Partitioning

The framework partitions static and dynamic matrices across row and column dimensions, maps them to distributed on-chip memories, and uses heuristic placement to regularize communication. Dynamic mapping additionally manages KV-cache growth through cyclic storage.

  • B. Partitioning: Static weights WQ, WK, WV, and WO and dynamic matrices Q, K, and V are partitioned along both row and column dimensions.The scheme resembles tensor and context parallelism but uses finer-grained partitions and more frequent on-chip communication.
  • B. Partitioning: Static matrix partitions map to IMC PEs, while dynamic partitions, including KV cache and temporary Q/O storage, map to scratchpads.
  • B. Partitioning: Static mapping groups submatrices contiguously, preferably in rectangular regions, with row-major or column-major traversal to reduce the mapping space and irregular dataflow.
  • B. Partitioning: Q, K, and V use column-major inter-tile mapping, O uses row-major inter-tile mapping, and V is transposed intra-tile but not inter-tile.
  • B. Partitioning: KV-cache tokens are appended cyclically within designated scratchpad regions to balance local-memory utilization as decode context length grows.

2) Inter-layer mapping:

Inter-layer placement uses a serpentine topology to relay autoregressive outputs across decoders while supporting regular layer-wise resource planning. The architecture distinguishes coupled and disaggregated prefill-decode operation.

  • 2) Inter-layer mapping:: 18: Each decoder places an attention layer followed by FFNs, with Llama models using three FFN layers per decoder.The FFN hidden dimension is approximately four times the embedding dimension, making its IMC capacity comparable to an attention layer.
  • 2) Inter-layer mapping:: The serpentine topology relays the latest output token from the last decoder back to the first decoder during autoregressive decode.
  • 2) Inter-layer mapping:: After partitioning, on-chip communication follows deterministic Broadcast, Reduce, or ReduceScatter patterns that are fixed at compilation time.This reduces run-time control and arbitration burden across the massive PE array.
  • 2) Inter-layer mapping:: Prefill is more computationally intensive, whereas decode places greater demands on memory bandwidth.
  • 2) Inter-layer mapping:: LEAP-A couples prefill and decode on shared hardware, while LEAP-D disaggregates them for dedicated throughput improvement.

1) Prefill Dataflow:

Prefill broadcasts tokens to weight-hosting IMC PEs, reduces partial results through RPUs, and performs subsequent attention operations in NMC resources. Communication and computation can overlap through pipelining, while LEAP-D transfers KV data concurrently.

  • 1) Prefill Dataflow:: Prefill broadcasts each input token to IMC PEs hosting Q, K, and V weights, then performs MAC-1 on those regions.
  • 1) Prefill Dataflow:: Q/K outputs undergo row-major ReduceScatter-1 and V outputs undergo element-wise ReduceScatter-1 among RPUs.The resulting K data is then distributed to Q-region macros before the next attention computation.
  • 1) Prefill Dataflow:: NMC PEs storing Q compute across locally stored Q vectors, followed by AllReduce when attention heads are partitioned along the head-width dimension.
  • 1) Prefill Dataflow:: Pipelining overlaps collective communication with concurrent DDMM-MAC or reduction, including AllGather-2 with MAC-3.
  • 1) Prefill Dataflow:: In decode, one token is processed per iteration and prefill’s ReduceScatter-1 becomes Reduce-1; in LEAP-D, KV transfer overlaps AllGather-1 and AllGather-2 without time overhead.

2) Decode Dataflow:

LEAP-D separates prefill and decode execution across different on-chip regions, using decode-region parallelism and fine-grained communication scheduling to reduce latency. The design transfers KV caches without adding critical-path overhead and differs from prior disaggregation approaches in its operation grouping and resource allocation.

  • Decode-region organization: LEAP-D processes prefill and decode phases in different on-chip core regions, allocating more cores to decode to increase concurrent scratchpad access.Prefill cores retain IMC PE and router resources, whereas decode cores host routers only.
  • Decode dataflow: Broadcasting, MAC, and reduction operations are scheduled across decode macros, with loop unrolling reducing latency despite one-time Q/K/V broadcast overhead.The current token’s Q vector is broadcast across the K channel, whose macros perform MAC-2 concurrently.
  • Communication scheduling: KV-cache transfer between prefill and decode regions overlaps existing prefill steps, producing no critical-path overhead.The design uses fine-grained computation–communication scheduling rather than overprovisioned links to maintain low latency.
  • Relation to prior work: Unlike one prior approach that disaggregates projections and FFNs while aggregating attention, this work keeps attention disaggregated and other operations aggregated.The stated goal is to fit the IMC resource limit.

C. Sequence Batching

The batching evaluation studies two simultaneous requests under local-capacity constraints, requiring coordinated scheduling to manage interference between prefill and decode. Hardware implementation and system-level measurements cover representative Llama models, area and power, throughput, energy efficiency, and scalability-related figures.

  • Batching assumptions: The study uses a batch size of two because local scratchpad capacity and the disaggregated architecture constrain larger batches.The authors note that this scale is not comparable to cloud-level applications.
  • Request scheduling: Concurrent requests require coordinated scheduling because decode and prefill phases can interfere when accessing shared IMC PE resources.A decode phase may be stalled during attention, while prefill may be halted between vector-loop iterations to prioritize decode.
  • Evaluation setup: The evaluation uses an instruction-level accurate NoC simulator and representative Llama 3.2-1B, Llama 3-8B, and Llama 2-13B models.Batched request timing is modeled from layer-wise and operation-level estimates.
  • Power and area: 56.32% of total energy is consumed by the router despite occupying 17.51% of macro area, reflecting its role in data movement and on-the-fly processing.The 2D mesh preserves the original area distribution as tiles are replicated during expansion.
  • GPU comparison: 2.55× higher throughput and 71.94× better energy efficiency are reported for LEAP-A relative to A100, while LEAP-D achieves 1.52× throughput and 24.91× energy-efficiency improvements over H100.The comparison uses a 2048-token context window with 1024 input and 1024 output tokens.
  • System implications: The reported efficiency advantages are attributed to minimized data movement enabled by distributed compute–memory organization and optimized dataflow.The authors state that mapping and scheduling preserve high throughput despite tensor partitioning caused by capacity constraints.

C. Single Request Performance

This subsection analyzes the performance of processing a single request.

  • Evaluation scope: The subsection evaluates performance when processing a single request.It establishes the evaluation scope for the associated single-request results.
  • Evaluation scope: The analysis focuses on request-level execution rather than concurrent batched processing.This distinguishes the subsection from evaluations involving multiple simultaneous requests.
  • Evaluation scope: Single-request performance is treated as a separate evaluation setting in the system analysis.The subsection’s stated subject is processing one request.

1) Prefill-decode in LEAP-A:

Single-request execution exposes a large decode-throughput disadvantage relative to prefill, while LEAP-D reduces that disparity through broader decode parallelism and effective memory-bandwidth use. Scalability analysis attributes favorable growth to partitioning that limits communication critical paths, while IRCU data movement and DDMM operations remain key latency constraints.

  • Single-request performance: 4–6× lower decode throughput than prefill is observed for single requests as context length grows and only one token vector flows through the mesh.The passage identifies underutilized pipeline stages as the more important factor.
  • LEAP-D comparison: LEAP-D significantly reduces prefill–decode throughput disparity by broadcasting the current token vector across decode-region PEs.This increases computational parallelism and effective memory bandwidth by addressing decode-region pipeline bottlenecks.
  • Model scalability: Model size can increase by approximately 8× from Llama 3.2-1B to Llama 3-8B, while the proposed partitioning keeps critical-path scaling tied to selected dimension–layer products.The communication path scales with (se · sl) or (sh · sl), rather than the full product of all three scaling factors.
  • Critical-path analysis: IMC-side operations seldom dominate critical-path latency because computation overlaps communication and IMC PEs provide high parallelism.Latency is primarily constrained by data movement and DDMM operations within the IRCUs.
  • Execution-mode comparison: Throughput comparisons between LEAP-A and LEAP-D are presented to assess the effect of the two execution modes.The cited figure is the comparative analysis referenced for the execution-mode evaluation.
  • Design-space projection: A 64-bit packet width with 16-way IRCU parallelism lies near the performance frontier and achieves near-optimal throughput without excessive resource cost.The result is presented as a trend-projection finding.

D. Batched Request Performance

Batched execution improves TTFT by overlapping prefill and decode, but increases TPOT because requests compete for IMC access. Memory utilization also differs across phases, making batching and non-batching complementary modes.

  • Batching substantially improves TTFT by overlapping prefill and decode, but increases TPOT because sequences compete for IMC units.The first request’s latency remains similar, while coordination adds latency for the batched execution.
  • Batching is suited to TTFT-critical applications, whereas non-batching is preferred when optimizing TPOT.LEAP therefore treats batching and non-batching as complementary execution modes rather than universally ranking one above the other.
  • Scratchpad utilization is dynamic during prefill but more consistent during decode, while IMC memory usage remains largely stable.Prefill scratchpads temporarily buffer K/V data and briefly store Q/O intermediates; decode scratchpads have more continuous utilization.
  • The framework balances throughput and latency through coordinated computation, communication scheduling, and locality-aware mapping.These design choices support efficient LLM serving on a unified compute–memory–communication fabric.
Loading 2609.00857v1…