paper(exec-summary): edit-pass through Executive Summary done
- 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>
This commit is contained in:
Binary file not shown.
@@ -18,11 +18,15 @@
|
|||||||
|
|
||||||
\graphicspath{{figures/}}
|
\graphicspath{{figures/}}
|
||||||
|
|
||||||
\title{\textbf{Hardware--Software Codesign for Attention-Centric\\
|
\title{\textbf{Hardware--Software Co-Design for Grouped-Query Attention on AHBM}}
|
||||||
LLM Kernels: GEMM, Collective Communication, and\\
|
|
||||||
Fused Grouped-Query Attention}}
|
\author{Mukesh Garg \and Jiyoon Lee \and Yangwook Kang}
|
||||||
\author{KernBench Team}
|
|
||||||
\date{2026 First Half (1H)}
|
\date{
|
||||||
|
\small
|
||||||
|
AGI Computing Lab, System Technology Group\\
|
||||||
|
H1 2026
|
||||||
|
}
|
||||||
|
|
||||||
\begin{document}
|
\begin{document}
|
||||||
\maketitle
|
\maketitle
|
||||||
|
|||||||
@@ -1,39 +1,33 @@
|
|||||||
\section*{Executive Summary}
|
\section*{Executive Summary}
|
||||||
\addcontentsline{toc}{section}{Executive Summary}
|
\addcontentsline{toc}{section}{Executive Summary}
|
||||||
|
|
||||||
This report studies how to make attention-centric LLM kernels fast through
|
Grouped-Query Attention (GQA) has largely replaced GPT-3-style
|
||||||
hardware--software codesign, using KernBench, a discrete-event simulator
|
multi-head attention in modern decoder-only LLMs (e.g., Llama 3 and
|
||||||
that runs kernels at the source level---without a compiler or other
|
Mistral) due to its reduced memory and bandwidth requirements. Mapping
|
||||||
software-stack dependency---so that algorithm-level optimizations and the
|
GQA efficiently onto AHBM introduces three architectural requirements:
|
||||||
hardware features that support them can be isolated and measured.
|
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.
|
||||||
|
|
||||||
The 1H work centers on attention. Building on prior multi-head-attention
|
To address these requirements, this report introduces three
|
||||||
work, we target FlashAttention-style Grouped-Query Attention (GQA) and find
|
hardware--software co-design mechanisms: GQA-aware placement
|
||||||
that a fast fused attention kernel rests on two enabling optimizations,
|
of KV caches and weights across TCM, SRAM, and HBM; PE\_IPCQ,
|
||||||
each studied in its own right and then combined inside the fused kernel:
|
an efficient on-device collective communication primitive; and a
|
||||||
a \textbf{composite command} that issues a tiled GEMM as a single,
|
composite-command GEMM pipeline that tightly pipelines memory
|
||||||
self-routing pipeline, and \textbf{PE\_IPCQ}, a per-PE on-device collective
|
and compute operations within each PE under PE\_SCHEDULER control.
|
||||||
engine for all-reduce.
|
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.
|
||||||
|
|
||||||
The measured results form a single, consistent story. For GEMM, the
|
While GQA serves as the motivating workload in this study, the
|
||||||
composite command lets compute-rich shapes reach
|
resulting architectural mechanisms are broadly applicable across the
|
||||||
\textasciitilde\SI{78}{\percent} of peak MAC efficiency, with the simulator's
|
AHBM software stack. PE\_IPCQ provides a scalable communication
|
||||||
measured efficiency tracking the analytic prediction within
|
substrate for collective operations and communication-intensive workloads, while
|
||||||
\SIrange{10}{20}{\percent}---both a performance result and a validation of
|
composite-command execution enables efficient fusion of memory
|
||||||
the latency model. For all-reduce, PE\_IPCQ delivers on-device collectives
|
movement, GEMM kernels, normalization, and other element-wise
|
||||||
whose latency tracks the interconnect's physical limits; a 2D-torus
|
operations. Together with the hierarchical data-placement framework,
|
||||||
inter-device fabric is \SI{20}{}--\SI{25}{\percent} faster than a mesh or
|
these capabilities form a reusable foundation for future AI kernels
|
||||||
ring at scale, and staging the collective in on-PE TCM rather than HBM or
|
and communication libraries on AHBM.
|
||||||
SRAM saves a further \SI{14}{}--\SI{38}{\percent}. For fused GQA, the two
|
|
||||||
enablers combine: because the composite command keeps the MAC array nearly
|
|
||||||
idle, the kernel is exposed as overwhelmingly data-movement bound---the DMA
|
|
||||||
engine outweighs the compute engines by two to three orders of magnitude,
|
|
||||||
and on-device IPCQ reduction traffic scales with users and devices.
|
|
||||||
|
|
||||||
\textbf{Bottom line.} For an attention-dominated decoder, the hardware
|
|
||||||
changes worth keeping are the ones that move and reduce data faster: the
|
|
||||||
single-command self-routing GEMM pipeline (which makes the MAC array
|
|
||||||
usable and exposes the real bottleneck), the on-device PE\_IPCQ collective
|
|
||||||
with its compute/communication virtual-channel split, fast on-PE staging
|
|
||||||
memory, and wrap-around (torus) inter-device links. Additional raw MAC
|
|
||||||
throughput is not the limiting factor for this workload.
|
|
||||||
|
|||||||
Reference in New Issue
Block a user