PAPER2026-05-07·Peking University (per author affiliations)·arXiv 2605.05639

TokenStack: A Heterogeneous HBM-PIM Architecture and Runtime for Efficient LLM Inference

Zhuoran Li, Zhuohang Bian, Zihao Huang, Guangyu Sun, Yun Liang, Youwei Zhuo
COMPILED NOTES

Concrete HBM4-logic-die PIM architecture: separates memory stacks into capacity-focused vs. compute-enabled layers, using the HBM4 base die as a local controller; claims 1.62x geomean token throughput, 1.70x latency-sensitive serving capacity, 30-47% per-token energy reduction on production-derived workloads. Directly answers two open Knowledge Gaps already flagged in this topic's frontier.md ('Commercial PIM product specs' and 'C-HBM4E NMC workload adoption'). Published 2026-05-07 — just outside the strict 60-day discovery window but flagged given the direct gap match.

TokenStack: Heterogeneous HBM-PIM for LLM Inference

The problem, stated exactly

"Large language model (LLM) serving is now limited by the key-value (KV) cache. During decode, each new token rereads prior KV state, so attention becomes a bandwidth- and capacity-heavy memory task."

This is the memory-wall thesis (hardware T1, conviction 7) arriving at the serving layer with a specific mechanism: decode re-reads prior KV state per token, so attention stops being compute and becomes memory traffic.

The design

Vertical heterogeneity inside the HBM-PIM stack — dense capacity layers separated from compute-enabled layers, with the logic base die acting as a control point for cross-layer data movement without host intervention. The interesting move is architectural rather than numerical: it stops treating a memory stack as uniform.

Measured against AttAcc (a dedicated-PIM baseline)

MetricResult
Token throughput1.62× geometric mean
SLO-compliant serving capacity1.70×
Per-token energy30–47% reduction

Evaluated on production-derived traces across four models with multi-QPS runs — a stronger evaluation posture than a single synthetic benchmark.

What it bears on

  • hardware T1 — "the AI compute shortage is substantially a memory-movement shortage." This is a serving-layer design whose entire premise is that claim, and its energy result (30–47% per token) is a direct measurement of movement cost.
  • The two cost-decline calls"Inference Costs Will Fall 90% by 2028 — The Memory Stack Is the Mechanism" names exactly this mechanism.

Limits

Compared against one baseline (AttAcc), itself a dedicated-PIM design rather than a deployed GPU stack — so the 1.62× is PIM-against-PIM, not PIM-against-what-you-run-today. A preprint, and PIM's commercial adoption remains the open question hardware T2 addresses.

Provenance

preprint, arXiv 2605.05639, submitted 2026-05-07. Ingested 2026-08-06 — listed without a close read since before the rung split.

RELATED · IN THE BASE
TokenStack: A Heterogeneous HBM-PIM Architecture and Runtime for Efficient LLM Inference | Knowledge Base | MenFem