HydraPIM: A Heterogeneous PIM Architecture for Efficient Attention in Long-Context LLMs

Abstract

The growing demand for long-context LLM inference has exposed a critical bandwidth–capacity trade-off in memory systems, rendering single-tier PIM architectures ineffective. HBM-PIMs offer high bandwidth but limited capacity, while DIMM-PIMs provide scalability at the cost of lower bandwidth; neither satisfies the throughput-latency requirements of long-context serving workloads for LLMs. To address this, we propose <bold>HydraPIM</bold>, a heterogeneous multi-tier PIM architecture that decomposes attention computation across HBM-PIM and DIMM-PIM tiers by exploiting the <bold>inherent sparsity</bold> of attention. HydraPIM introduces <bold>HydraAttention</bold>, a tiled attention mechanism with hierarchical reduction, enabling efficient cross-tier execution through lightweight on-chip reduction units. To maintain load balance under dynamic access patterns of sparse attention, HydraPIM features an <bold>importance-aware KV migration</bold> mechanism that monitors token relevance and relocates hot tokens to high-bandwidth tiers at runtime. This software-hardware co-design helps improve the utilization of both bandwidth and capacity. Evaluations show that HydraPIM achieves 1.66<inline-formula><tex-math notation="LaTeX">$\boldsymbol{\times}$</tex-math><alternatives><mml:math><mml:mo mathvariant="bold">×</mml:mo></mml:math><inline-graphic xlink:href="wang-ieq1-3710733.gif"/></alternatives></inline-formula> higher throughput than HBM-based NPU-PIM and 1.96<inline-formula><tex-math notation="LaTeX">$\boldsymbol{\times}$</tex-math><alternatives><mml:math><mml:mo mathvariant="bold">×</mml:mo></mml:math><inline-graphic xlink:href="wang-ieq2-3710733.gif"/></alternatives></inline-formula> lower latency than DIMM-based NPU-PIM across diverse long-context workloads.