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

Abstract: 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.

• The paper presents a memory-centric multi-chiplet architecture that replaces GPU compute with HBM-PNM cubes, achieving up to 16.3x lower latency in decode attention.
• It leverages an output-stationary dataflow and a hybrid tensor/context parallelism scheme to maximize HBM utilization while reducing communication overhead.
• Experimental evaluations on models like Qwen3-235B show impressive gains, including a 40 TB/s aggregate bandwidth and up to 6.6x improvement in energy efficiency.

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

Motivation and Architectural Shift

The AMMA architecture directly targets the inefficiency of GPU-centric long-context attention serving. Decode-phase attention in LLM inference exhibits highly memory-bound characteristics as context lengths scale to $\sim$1M tokens, rendering the expansive compute resources of current GPU platforms both underutilized and energetically wasteful. Profiling studies of NVIDIA H100 during decode attention (batch=1) reveal near-saturated HBM bandwidth usage ($>90\%$) with compute utilization dropping below $5\%$, resulting in significant area and power inefficiency and underlining the unsuitability of the GPU paradigm for ultra-long-context workloads.

Figure 1: Hardware utilization and power breakdown for Qwen3-235B attention demonstrates that GPU compute remains idle while memory bandwidth is saturated.

This mismatch is exacerbated in modern multi-agent and reasoning workloads, where user-facing latency is dominated by decode attention and datacenter power budgets are increasingly strained. AMMA proposes a radical architectural pivot: eliminate GPU compute dies entirely in favor of HBM-PNM cubes, each equipped with sophisticated logic dies on advanced process nodes (≤5 nm with HBM4 and beyond), roughly doubling aggregate package bandwidth and provisioning dedicated memory-side acceleration.

Figure 2: Conventional GPU-centric attention serving versus AMMA's PNM-based cube-centric approach.

PIM/PNM Integration and Microarchitecture

AMMA leverages processing-near-memory (PNM) on HBM logic dies, eschewing the typical constraints of processing-in-memory (PIM) approaches, which suffer from capacity penalties and limited DRAM-node compute density. Each HBM-PNM cube incorporates a microarchitecture optimized for memory-bound attention: 96 distributed $16\times16$ systolic arrays, hierarchical crossbars for inter-core communication, and LLC-free design. The trade-off between compute density and bandwidth utilization is addressed by deploying numerous small arrays (matching narrow $M$ dimensions inherent to decode workloads) and replacing energy/area-expensive LLCs with crossbar-based operand sharing.

Figure 3: Integration options and trade-offs for PIM (DRAM) and PNM (logic die) architectures.

Figure 4: Package-level cube arrangement and internal compute/memory hierarchy of each PNM logic die.

Dataflow, Tiling, and Core Scheduling

Addressing the low arithmetic intensity and sequential access patterns of attention, AMMA adopts output-stationary (OS) dataflow for its systolic arrays. Tiling strategies for GEMM dimensions prioritize maximizing utilization across the 96 SAs and minimize fill/drain overhead by continuous tiling. Split strategies for $K$ and $N$ dimensions activate idle SAs without inflating communication costs, and double-buffered DMA ensures latency hiding for memory accesses.

Figure 5: Dataflow options, tiling for small and large $N$, and continuous tiling in $N$ for high utilization.

Parallelism and Communication Optimization

AMMA introduces a two-level hybrid parallelism scheme combining tensor and context parallelism. KV heads are mapped to distinct cube groups via TP (Level 1), while sequence shards are distributed within groups via CP (Level 2), thereby localizing communication and avoiding sequence-length scaling in collective operations. This design fundamentally departs from both na\"ive TP across all cubes and prior GPU-centric PIM/PNM, which incur catastrophic communication latency when operating in the 1M-token regime.

Figure 6: Default TP16 strategy vs. AMMA’s hybrid TP/CP mapping, demonstrating localized collectives.

Further, AMMA deploys a reordered collective flow for the output projection, reducing intra-package synchronization overhead and eliminating redundant AllGather/AllReduce operations. Softmax correctness is preserved via online accumulation and commutativity of projection and reduction, following prior work in FlashAttention and Ring Attention.

Figure 7: Default versus reordered collective operation flow in output projection.

Experimental Evaluation and Numerical Impact

AMMA achieves significant improvements across core metrics. On GQA models (Qwen3-235B, Llama4-Maverick), single-chip decode latency speedup ranges from $12\times$–$>90\%$0 relative to H100 at batch size 1, sustained at $>90\%$1–$>90\%$2 up to batch size 32. Against NVIDIA Rubin (8 HBM4 cubes), AMMA maintains a $>90\%$3–$>90\%$4 latency advantage and, even in configurations matching aggregate bandwidth, retains $>90\%$5 gains in energy efficiency.

Figure 8: AMMA’s decode latency speedup relative to H100 across context lengths and batch sizes.

Figure 9: Energy and power analysis reveals $>90\%$6 higher energy efficiency vs. H100 and $>90\%$7 vs. Rubin.

Ablation studies confirm that most gains stem from reduced communication volume in hybrid parallelism and reordered collective scheduling, especially as sequence length grows toward $>90\%$8 tokens.

Figure 10: Ablation study on parallelism and collective operation strategies confirms communication speedup as the key benefit.

Per-layer timing breakdown further affirms that projection dominates at short sequences and attention at long sequences; AMMA delivers $>90\%$9 attention speedup against Rubin via $5\%$0 aggregate bandwidth.

Figure 11: Per-layer latency breakdown for Qwen3 shows AMMA’s speedup in projection and attention stages.

Batch size and hardware parameter exploration elucidate the trade-offs: AMMA’s compute power saturates at $5\%$1/cube, after which bandwidth remains the bottleneck, reinforcing its memory-centric efficiency.

Figure 12: Throughput and latency scaling with batch size; per-user generation speed reaches upper bounds at batch size 16.

Figure 13: Exploration of per-cube compute and D2D link bandwidth under extreme context length.

Related Work and Contrasts

AMMA stands in explicit contrast to GPU-centric optimizations (FlashDecoding++, MegaScale-Infer, Step-3, DistServe) and previous memory-side approaches (AttAcc, NeuPIMs, LoL-PIM, DReX, LongSight) by fully eliminating reliance on GPU as a communication hub and integrating advanced PNM logic dies for standalone memory-centric acceleration.

Implications and Future Directions

Practically, AMMA demonstrates that decoupling compute resources from memory-centric attention workloads yields substantial gains in latency, energy efficiency, and power consumption. Theoretically, this establishes memory-centric, multi-chiplet architectures as a distinct category in high-performance LLM inference hardware, suited for heterogeneous and disaggregated serving systems. The co-design of microarchitecture, parallelism, and collective scheduling is essential for achieving optimal utilization under strict area and power constraints as context lengths continue to grow.

Future research should explore deeper heterogeneity in multi-chiplet packages—e.g., integrating dedicated LPUs for FFN/MoE, further logic-die enhancements, and distributed memory-side collective optimizations. AMMA’s architectural principles are likely to inform new classes of accelerators for efficient reasoning, agentic computation, and broader multi-modal, multi-user deployments.

Conclusion

AMMA introduces a memory-centric, multi-chiplet PNM architecture that fundamentally rethinks the decode attention pipeline, eliminating GPU bottlenecks and establishing a scalable, efficient platform for extreme context-length LLM serving. Through tailored microarchitecture, hybrid parallelism, and communication optimization, AMMA achieves pronounced latency and energy improvements, demonstrating both practical and theoretical implications for future AI accelerator design.

(2604.26103)