diff --git a/docs/report/1H-codesign-paper/build/main.pdf b/docs/report/1H-codesign-paper/build/main.pdf index 023f499..07267be 100644 Binary files a/docs/report/1H-codesign-paper/build/main.pdf and b/docs/report/1H-codesign-paper/build/main.pdf differ diff --git a/docs/report/1H-codesign-paper/main.tex b/docs/report/1H-codesign-paper/main.tex index 37a15ce..c83fd4a 100644 --- a/docs/report/1H-codesign-paper/main.tex +++ b/docs/report/1H-codesign-paper/main.tex @@ -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 diff --git a/docs/report/1H-codesign-paper/sections/00-exec-summary.tex b/docs/report/1H-codesign-paper/sections/00-exec-summary.tex index 61c0b19..995c595 100644 --- a/docs/report/1H-codesign-paper/sections/00-exec-summary.tex +++ b/docs/report/1H-codesign-paper/sections/00-exec-summary.tex @@ -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.