Near-memory scheduler improves session memory use for language model agents
UNISON: A Co-Designed Near-Memory Scheduler of Session KV Residency for LLM Agents
Hardware Architecture
Summary
When computer programs use large language models to complete tasks, they often have to remember lots of information between steps. This can slow things down because the memory used to store this information runs out or gets mixed up. The authors created a new memory manager called UNISON that smartly decides what information to keep or move near the processor, speeding up task completion. Their system works better than previous methods, reducing delays and improving how often needed information is readily available.
What this means in practice
- •For ai system architects: Improve memory efficiency and reduce latency in multi-step language model agents by integrating UNISON's near-memory scheduling.
- •For hardware designers: Use UNISON's compact CMOS scheduling core to manage high-speed key-value memory hierarchy with minimal area and power overhead.
Authors
Fan He, Yan Li, Xiaoyang Zeng
Abstract
Large language models are increasingly composed into agent loops that plan, call tools, and resume the same task after each action. These loops press a shared memory hierarchy harder than conventional multi-turn chat, because they hold a growing key-value (KV) prefix across tool waits and place many sessions on one SRAM/HBM pool, so that eviction and hierarchical placement become a session-level efficiency problem orthogonal to compute-mode optimization. Existing proxies based on recency, timeout, or identity miss the mechanism information of the loop and therefore treat a live wait as a cold, discardable unit. We present Unified Native Inter-turn Session Orchestration Nexus (UNISON), an event-driven near-memory scheduler in which Survival-Penalty Eviction for Agent Return-gap (SPEAR) and Tiering in Idle-window DMA Events (TIDE) share one live ranking. SPEAR selects who leaves from a gap average and a turn-indexed hazard, while TIDE spends the observed wait as a DMA budget for who sits in the fast tier. On coding and general-mission benchmarks with three model families, totaling 1,415 sessions and 33,596 turns, the joint policy is the best non-oracle entry on every trace, raising hit rate by 0.3% to 23.1%, reducing AMAT by 22% to 51%, and lowering TTFT by 58% to 89% on long-horizon traces. A structural necessity analysis shows that the unified near-memory design cannot be decomposed into independent IPs or realized in software without re-introducing documented failure modes. The 28-nm CMOS scheduling core occupies 0.169 mm^2 at 13.6 mW and 150 MHz, a negligible overhead relative to the KV hierarchy it manages, reproducing the floating-point ranking at Kendall tau exceeding 0.998.