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/}}
|
||||
|
||||
\title{\textbf{Hardware--Software Codesign for Attention-Centric\\
|
||||
LLM Kernels: GEMM, Collective Communication, and\\
|
||||
Fused Grouped-Query Attention}}
|
||||
\author{KernBench Team}
|
||||
\date{2026 First Half (1H)}
|
||||
\title{\textbf{Hardware--Software Co-Design for Grouped-Query Attention on AHBM}}
|
||||
|
||||
\author{Mukesh Garg \and Jiyoon Lee \and Yangwook Kang}
|
||||
|
||||
\date{
|
||||
\small
|
||||
AGI Computing Lab, System Technology Group\\
|
||||
H1 2026
|
||||
}
|
||||
|
||||
\begin{document}
|
||||
\maketitle
|
||||
|
||||
@@ -1,39 +1,33 @@
|
||||
\section*{Executive Summary}
|
||||
\addcontentsline{toc}{section}{Executive Summary}
|
||||
|
||||
This report studies how to make attention-centric LLM kernels fast through
|
||||
hardware--software codesign, using KernBench, a discrete-event simulator
|
||||
that runs kernels at the source level---without a compiler or other
|
||||
software-stack dependency---so that algorithm-level optimizations and the
|
||||
hardware features that support them can be isolated and measured.
|
||||
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.
|
||||
|
||||
The 1H work centers on attention. Building on prior multi-head-attention
|
||||
work, we target FlashAttention-style Grouped-Query Attention (GQA) and find
|
||||
that a fast fused attention kernel rests on two enabling optimizations,
|
||||
each studied in its own right and then combined inside the fused kernel:
|
||||
a \textbf{composite command} that issues a tiled GEMM as a single,
|
||||
self-routing pipeline, and \textbf{PE\_IPCQ}, a per-PE on-device collective
|
||||
engine for all-reduce.
|
||||
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.
|
||||
|
||||
The measured results form a single, consistent story. For GEMM, the
|
||||
composite command lets compute-rich shapes reach
|
||||
\textasciitilde\SI{78}{\percent} of peak MAC efficiency, with the simulator's
|
||||
measured efficiency tracking the analytic prediction within
|
||||
\SIrange{10}{20}{\percent}---both a performance result and a validation of
|
||||
the latency model. For all-reduce, PE\_IPCQ delivers on-device collectives
|
||||
whose latency tracks the interconnect's physical limits; a 2D-torus
|
||||
inter-device fabric is \SI{20}{}--\SI{25}{\percent} faster than a mesh or
|
||||
ring at scale, and staging the collective in on-PE TCM rather than HBM or
|
||||
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.
|
||||
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.
|
||||
|
||||
Reference in New Issue
Block a user