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

arXiv: https://arxiv.org/abs/2604.26103 · HTML: https://mlsanity.com/paper/2604.26103/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in KV cache optimization, one of 12 directions in Efficient LLM inference and serving.
- Builds on 25 of its 60 references that are in the map.
- First posted on arXiv 28 Apr 2026.

## Placed in

- Efficient LLM inference and serving / [KV cache optimization](https://mlsanity.com/inference/kv-cache-optimization/)

## Builds on (most cited first, 25 of 25)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Gemini 2.5: Pushing the Frontier with Advanced Reasoning, Multimodality, Long Context, and Next Generation Agentic Capabilities](https://mlsanity.com/paper/2507.06261/) (2025)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [WebShop: Towards Scalable Real-World Web Interaction with Grounded Language Agents](https://mlsanity.com/paper/2207.01206/) (2022)
- [DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving](https://mlsanity.com/paper/2401.09670/) (2024)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning](https://mlsanity.com/paper/2505.24298/) (2025)
- [SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills](https://mlsanity.com/paper/2308.16369/) (2023)
- [LoongServe: Efficiently Serving Long-Context Large Language Models with Elastic Sequence Parallelism](https://mlsanity.com/paper/2404.09526/) (2024)
- [KVFlow: Efficient Prefix Caching for Accelerating LLM-Based Multi-Agent Workflows](https://mlsanity.com/paper/2507.07400/) (2025)
- [Step-3 is Large yet Affordable: Model-system Co-design for Cost-effective Decoding](https://mlsanity.com/paper/2507.19427/) (2025)
- [NeuPIMs: NPU-PIM Heterogeneous Acceleration for Batched LLM Inferencing](https://mlsanity.com/paper/2403.00579/) (2024)
- [RecNMP: Accelerating Personalized Recommendation with Near-Memory Processing](https://mlsanity.com/paper/1912.12953/) (2019)
- [ASTRA-sim2.0: Modeling Hierarchical Networks and Disaggregated Systems for Large-model Training at Scale](https://mlsanity.com/paper/2303.14006/) (2023)
- [TensorDIMM: A Practical Near-Memory Processing Architecture for Embeddings and Tensor Operations in Deep Learning](https://mlsanity.com/paper/1908.03072/) (2019)
- [RollPacker: Mitigating Long-Tail Rollouts for Fast, Synchronous RL Post-Training](https://mlsanity.com/paper/2509.21009/) (2025)
- [Duplex: A Device for Large Language Models with Mixture of Experts, Grouped Query Attention, and Continuous Batching](https://mlsanity.com/paper/2409.01141/) (2024)
- [MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism](https://mlsanity.com/paper/2504.02263/) (2025)
- [Challenges and Research Directions for Large Language Model Inference Hardware](https://mlsanity.com/paper/2601.05047/) (2026)
- [LP-Spec: Leveraging LPDDR PIM for Efficient LLM Mobile Speculative Inference with Architecture-Dataflow Co-Optimization](https://mlsanity.com/paper/2508.07227/) (2025)
- [Kareus: Joint Reduction of Dynamic and Static Energy in Large Model Training](https://mlsanity.com/paper/2601.17654/) (2026)

## Abstract (arXiv)

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.