ywkang
a965db7056
- Title: "Hardware-Software Codesign for Grouped-Query Attention on AHBM"
(drop "Attention-Centric LLM Kernels" framing; subtitle pillars removed)
- Authors: Mukesh Garg, Jiyoon Lee, Yangwook Kang
- Affiliation: System Technology Group, AGI Computing Lab — 2026 1H
- Executive Summary rewritten into a tight three-paragraph structure:
(1) GQA motivation + three AHBM architectural requirements (placement,
inter-PE comm, pipelined memory-compute);
(2) three matching HW-SW co-design mechanisms (GQA-aware placement,
PE_IPCQ, composite-command GEMM under PE_SCHEDULER) and headline
results (38% / 25-35% / 78%), closing on efficient HBM-bandwidth use;
(3) broader applicability across the AHBM software stack.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
34 lines
1.8 KiB
TeX
34 lines
1.8 KiB
TeX
\section*{Executive Summary}
|
|
\addcontentsline{toc}{section}{Executive Summary}
|
|
|
|
Grouped-Query Attention (GQA) has largely replaced GPT-3-style
|
|
multi-head attention in modern decoder-only LLMs (e.g., Llama 3 and
|
|
Mistral) due to its reduced memory and bandwidth requirements. Mapping
|
|
GQA efficiently onto AHBM introduces three architectural requirements:
|
|
optimized placement of KV caches and weights to minimize inter-PE
|
|
communication, low-overhead support for unavoidable cross-PE traffic,
|
|
and efficient pipelining of memory accesses and computation within each
|
|
PE.
|
|
|
|
To address these requirements, this report introduces three
|
|
hardware--software co-design mechanisms: GQA-aware placement
|
|
of KV caches and weights across TCM, SRAM, and HBM; PE\_IPCQ,
|
|
an efficient on-device collective communication primitive; and a
|
|
composite-command GEMM pipeline that tightly pipelines memory
|
|
and compute operations within each PE under PE\_SCHEDULER control.
|
|
Kernbench-based evaluation show up to \SI{38}{\percent} lower collective
|
|
communication overhead, \SIrange{25}{35}{\percent} lower all-reduce
|
|
latency than a mesh-based interconnect, and up to \SI{78}{\percent} of
|
|
peak MAC efficiency for compute-intensive GEMM kernels. These results
|
|
demonstrate that the fused GQA kernel can efficiently exploit AHBM's HBM bandwidth.
|
|
|
|
While GQA serves as the motivating workload in this study, the
|
|
resulting architectural mechanisms are broadly applicable across the
|
|
AHBM software stack. PE\_IPCQ provides a scalable communication
|
|
substrate for collective operations and communication-intensive workloads, while
|
|
composite-command execution enables efficient fusion of memory
|
|
movement, GEMM kernels, normalization, and other element-wise
|
|
operations. Together with the hierarchical data-placement framework,
|
|
these capabilities form a reusable foundation for future AI kernels
|
|
and communication libraries on AHBM.
|