AMMA: A Multi-Chiplet Memory-Centric Architecture for Low-Latency 1M Context Attention Serving

Abstract (verbatim)

"All current LLM serving systems place the GPU at the center, from production-level attention-FFN disaggregation to NVIDIA's Rubin GPU-LPU heterogeneous platform. Even academic PIM/PNM proposals still treat the GPU as the central hub for cross-device communication. Yet the GPU's compute-rich architecture is fundamentally mismatched with the memory-bound nature of decode-phase attention, inflating serving latency while wasting power and die area on idle compute units. The problem is compounded as reasoning and agentic workloads push context lengths toward one million tokens, making attention latency the primary user-facing bottleneck. To address these inefficiencies, we present AMMA, a multi-chiplet, memory-centric architecture for low-latency long-context attention. AMMA replaces GPU compute dies with HBM-PNM cubes, roughly doubling the available memory bandwidth to better serve memory-bound attention workloads. To translate this bandwidth into proportional performance gains, we introduce (i) a logic-die microarchitecture that fully exploits per-cube internal bandwidth for decode attention under a minimal power and area budget, (ii) a two-level hybrid parallelism scheme, and (iii) a reordered collective flow that reduces intra-chip die-to-die communication overhead. We further conduct a design-space exploration over per-cube compute power and intra-chip D2D link bandwidth, providing actionable guidance for hardware designers. Evaluations show that AMMA achieves 15.5X lower attention latency and 6.9X lower energy consumption compared with the NVIDIA H100."

Key Contributions

• Architectural inversion: the GPU stops being the hub. Compute dies are replaced with HBM-PNM cubes, making memory bandwidth (not compute) the central resource — the sharpest statement yet of the memory-centric thesis for the specific workload of decode-phase attention.
• Names NVIDIA's Rubin GPU-LPU heterogeneous platform as the current GPU-centric baseline it is arguing against — a useful datapoint on how the frontier reads Rubin's architecture.
• Three mechanisms to convert doubled bandwidth into realized speedup: a logic-die microarch for per-cube decode attention, two-level hybrid parallelism, and a reordered collective flow to cut intra-package die-to-die traffic.
• A design-space exploration over per-cube compute power and D2D link bandwidth — explicit guidance for hardware designers (relevant to the chiplet/UCIe interconnect gap this KB flags).

Results (full-text magnitudes — deepened 2026-07-24 from the arXiv HTML)

(Preprint; simulated results vs H100 / Rubin / NeuPIMs baselines, not independently replicated.)

Headline vs NVIDIA H100 (1M-token decode attention): 15.5× lower attention latency, 6.9× lower energy, ~2× memory bandwidth.

Per-baseline detail:

• vs H100 — latency 12.0–16.3× faster at batch size 1; sustains 13–20× at BS=32. Energy efficiency 5.6–6.6× higher (Token/J).
• vs NVIDIA Rubin — 1.8–2.5× faster at short-to-medium sequences, narrowing to 1.1× at 1M tokens; energy 2.8× better at 1M (2.6–3.1× vs Rubin TP2).
• vs NeuPIMs (prior PIM) — 3.4× faster on Qwen3-235B; 1.4× on Llama4-Maverick.

Package / microarchitecture (the design the deepen was for):

• 16 HBM-PNM cubes in a 4×4 2D mesh, UCIe 3.0 D2D links (15 ns/hop, 1500 GB/s per link).
• Per cube: 96 small 16×16 systolic arrays → 96 TFLOPS @ 2 GHz (1536 TFLOPS aggregate); 2.75 TB/s HBM4 bandwidth; 384 KB SRAM (12 cores × 32 KB); 15 W.
• Package: 44 TB/s aggregate bandwidth, 240 W TDP, 67% smaller SRAM footprint than H100. Contrast: H100 = 3.35 TB/s / 700 W; Rubin = 22 TB/s (8 cubes) / 2200 W.

Design-space exploration (the "actionable guidance" the abstract promised):

• Per-cube compute saturates at ~96 TFLOPS for GQA decode; beyond that, bandwidth is the bottleneck and adding D2D bandwidth yields little (fixed per-hop latency dominates at small batch).
• Communication-reordering speedup grows with sequence length: 2.7× (8K) → 17.7× (256K) → 65.4× (1M) — the reordered collective flow (ReduceScatter replacing AllReduce+AllGather) is what makes 1M-context latency tractable.
• Batch scaling saturates ~BS=16 (throughput +2.14× from 1→32, but latency degrades ~30× over the same range — a latency-first design).
• Workloads: Qwen3-235B, Llama4-Maverick (GQA), DeepSeek-V3 (MLA); decode phase, 4K–1M sequences.

Why it matters for this KB

• Third distinct 2026 near-memory architecture (with TokenStack and StreamDQ) — and the most aggressive: TokenStack and StreamDQ add PNM logic alongside the GPU; AMMA removes the GPU compute die from the attention path entirely. Together they map a spectrum of how far memory-centric computing pushes into the decode phase.
• Sits at the intersection of two tracked gaps: multi-chiplet / interconnect (Tier 3 UCIe/chiplet gap) and memory-centric computing / PNM (frontier 8). It is memory-centric and chiplet.
• Directly extends the KV-cache-as-memory-object frontier: it is a hardware answer to memory-bound decode attention at 1M-token context.

Date caveat: submitted 2026-04-28 (v2 2026-04-30) — outside the strict 7-day discovery window. Ingested as a boundary case because it fills two standing gaps and is high quality, matching the precedent set by the TokenStack (2026-05) ingest. Not new-in-the-last-week; new-to-this-KB.

Source: AMMA: A Multi-Chiplet Memory-Centric Architecture for Low-Latency 1M Context Attention Serving — Yu, Ye, Zhou et al. (UC San Diego, Columbia, Yonsei, NVIDIA, Samsung Semiconductor). Full text ingested 2026-07-24 via the arXiv HTML (v2, 2604.26103v2); microarchitecture, per-baseline results, and DSE magnitudes captured. Provenance [preprint — authors' own simulation].