AMEND: Audited Margins Enable Nonblocking Drops in GPU-PIM LLM Decoding

2026-09-09

AI Summary

为解决长上下文注意力受GPU内存带宽限制问题，AMEND通过审计边距预测并筛选关键值块，结合GPU-PIM设计提高解码速度和降低能耗。

Abstract

Autoregressive large language model (LLM) decoding re-reads a growing key-value (KV) cache at every step, so long-context attention is bound by graphics processing unit (GPU) memory bandwidth. Block-sparse attention skips low-contribution KV blocks, but a selector that decides after the current query-key (QK) product, such as max-relative block thresholding (BLASST), still reads every K block, and a processing-in-memory (PIM) filter that decides from the current query places a serial PIM stage on the critical path. We present AMEND, a GPU-PIM attention design that removes both dependencies. AMEND predicts each block's BLASST verdict from margins audited at earlier steps, so the GPU fetches only predicted survivors while near-bank PIM units in high-bandwidth memory (HBM) concurrently score the omitted complement. A stack-level controller merges both observations, updates the predictor, and eagerly generates the next step's mask, so every predicted drop is re-observed without blocking the current token. Operating points are selected offline by constrained Bayesian optimization under a false-drop budget. Across LongBench and RULER runs, AMEND preserves near-baseline task quality; in simulation at batch size 8, it achieves $1.40$-$3.63\times$ end-to-end decode speedup and 28-66% lower dynamic decode energy than dense attention across 8K-64K contexts.

Problem

• GPU-PIM
• large language model decoding
• memory bandwidth
• block-sparse attention

Innovation

• AMEND
• GPU-PIM
• Nonblocking Drops
• Audited Margins
• High-bandwidth Memory (HBM)