diff --git a/docs/report/1H-codesign-paper/build/main.pdf b/docs/report/1H-codesign-paper/build/main.pdf index 07267be..b977c8d 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 c83fd4a..5d3fa3f 100644 --- a/docs/report/1H-codesign-paper/main.tex +++ b/docs/report/1H-codesign-paper/main.tex @@ -25,7 +25,7 @@ \date{ \small AGI Computing Lab, System Technology Group\\ -H1 2026 +2026 H1 Report } \begin{document} 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 995c595..e5cd414 100644 --- a/docs/report/1H-codesign-paper/sections/00-exec-summary.tex +++ b/docs/report/1H-codesign-paper/sections/00-exec-summary.tex @@ -30,4 +30,7 @@ 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. +and communication libraries on AHBM. Looking ahead, these same mechanisms provide +the basis for optimizing FFN-intensive workloads such as Mixture-of-Experts (MoE) +layers and for developing integrated optimization strategies spanning both attention +and MoE components within next-generation AI models. diff --git a/docs/report/1H-codesign-paper/sections/01-introduction.tex b/docs/report/1H-codesign-paper/sections/01-introduction.tex index 2a4ab0a..b1e140b 100644 --- a/docs/report/1H-codesign-paper/sections/01-introduction.tex +++ b/docs/report/1H-codesign-paper/sections/01-introduction.tex @@ -1,49 +1,83 @@ \section{Introduction} \label{sec:intro} -The performance of a large language model on an accelerator is decided less -by peak FLOPS than by how well the kernel exploits the memory system and -the interconnect. This is especially true of attention, which in the -decode phase reads a large KV cache to do a small amount of arithmetic. -Optimizing such kernels is fundamentally a \emph{codesign} problem: the -algorithm (how to tile, fuse, and reduce) and the hardware (what issue -mechanism, what collective engine, what memory hierarchy) have to be -designed against each other. Optimizing only the software leaves the -hardware idle; adding hardware that the software cannot reach is wasted -area. +AHBM integrates compute units directly into the HBM stack. Each +processing element (PE) is paired with a dedicated slice of HBM and uses +local TCM and SRAM to stage data between memory and the MAC array. On +this memory-centric architecture, kernel performance depends not only on +compute throughput but also on how effectively data is placed, moved, +and shared across the memory hierarchy and between PEs. Optimizing AI +kernels on AHBM therefore requires hardware--software co-design, in +which kernel algorithms and architectural mechanisms are developed +together. -\paragraph{The 1H focus: attention.} This half's work targets attention. -Multi-head attention (MHA) was studied previously and is taken here as the -established baseline. The new work is FlashAttention-style tiling and -Grouped-Query Attention (GQA)---the form used by modern long-context -decoders, in which several query heads share a single KV head to shrink the -KV cache. Realizing a fast \emph{fused} GQA kernel, however, is not a -single optimization. It decomposes into two enabling optimizations that we -studied separately and then brought together inside the fused kernel: +To enable detailed performance analysis and rapid design exploration, we +developed \textbf{KernBench}, a source-level discrete-event simulation +platform for AHBM. KernBench implements the AHBM execution model, +including memory-system latencies, the PE execution model, inter-PE +communication, and host-side orchestration, while executing both kernel +and host software directly from source code. This provides fine-grained +visibility into execution behavior. It enables systematic evaluation of hardware--software co-design choices +independently of higher-level software stacks such as compilers and +runtimes. All +results presented in this report are obtained using KernBench, which +serves as the common evaluation platform for all mechanisms and kernels +discussed in this study. -\begin{enumerate} - \item \textbf{GEMM optimization} via a composite command - (\S\ref{sec:gemm}): a way to issue a tiled matrix multiply as one - self-routing pipeline so the MAC array stays fed without per-tile command - overhead. Attention's $Q\!\cdot\!K^{\top}$ and $P\!\cdot\!V$ products are - exactly such GEMMs. - \item \textbf{Communication optimization} via PE\_IPCQ - (\S\ref{sec:allreduce}): a per-PE on-device collective engine that - performs all-reduce overlapped with compute. Attention's multi-user and - sequence-parallel KV reductions are exactly such collectives. -\end{enumerate} +This report focuses on Grouped-Query Attention (GQA), one of the most +performance- and bandwidth-critical components of LLM inference. Modern +decoder-only models such as Llama~3 and Mistral have largely transitioned +from GPT-3-style multi-head attention (MHA) to GQA, in which multiple +query heads share a single KV head to reduce KV cache capacity and +memory-bandwidth requirements. While GQA improves system efficiency at +the model level, mapping it 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 report therefore reads as \emph{two enablers building to one capstone}: -GEMM (\S\ref{sec:gemm}) and all-reduce (\S\ref{sec:allreduce}) are -developed and measured on their own, and then fused GQA -(\S\ref{sec:gqa}) shows them operating together. Before the results, -\S\ref{sec:platform} describes the KernBench platform---why a -software-stack-independent, source-level simulator is the right tool for -this question, how it executes a kernel, how it computes latency and how -accurate that is, and the exact hardware configuration used throughout. -We close with a cross-cutting discussion of which hardware changes are -worth their cost (\S\ref{sec:discussion}), a conclusion -(\S\ref{sec:conclusion}), and the 2H agenda -(\S\ref{sec:future}): extending from attention to the feed-forward and -mixture-of-experts layers, and to the compute/data distribution questions -that full LLM decoding and agentic workloads raise. +To address these requirements, this report introduces three +hardware--software co-design mechanisms. First, GQA-aware data placement +distributes KV caches and weights across the TCM/SRAM/HBM hierarchy to +reduce communication overhead and improve data locality. Second, +PE\_IPCQ provides an efficient on-device collective communication +primitive for reductions and other communication-intensive operations. +Third, a composite-command GEMM pipeline tightly pipelines memory +movement and computation within each PE under PE\_SCHEDULER control, +reducing command overhead while keeping the MAC array efficiently +utilized. + +The compute enabler (the composite-command GEMM pipeline, +\S\ref{sec:gemm}) and the communication enabler (PE\_IPCQ, +\S\ref{sec:allreduce}) are first developed and evaluated independently. +The fused GQA kernel (\S\ref{sec:gqa}) then combines them with +GQA-aware data placement to demonstrate an end-to-end attention +implementation on AHBM. The correspondence is direct: attention's +$QK^{\top}$ and $PV$ products are precisely the GEMMs that benefit from +the composite-command pipeline, while its KV reductions are precisely +the collective operations that benefit from PE\_IPCQ. Together, these +mechanisms enable the fused GQA kernel to efficiently exploit AHBM's HBM +bandwidth. + +Although GQA serves as the motivating workload for this study, the +resulting mechanisms are intended as reusable building blocks for a much +broader class of AI kernels. PE\_IPCQ can support collective +communication across distributed and communication-intensive workloads, +while composite-command execution can be applied to GEMM-based kernels, +feed-forward networks (FFNs), normalization, and other fused operator +pipelines. Together with the hierarchical data-placement framework, +these mechanisms form a foundation for future AI kernels and +communication libraries on AHBM. In the second half of 2026, this +foundation will be extended to FFN- and MoE-dominated workloads and to +end-to-end optimization of complete LLM execution. + +The remainder of this report is organized as follows. +Section~\ref{sec:platform} describes the KernBench platform and the AHBM +configuration used throughout this study. Sections~\ref{sec:gemm}, +\ref{sec:allreduce}, and \ref{sec:gqa} present the composite-command +GEMM pipeline, PE\_IPCQ collective communication, and the fused GQA +kernel, respectively. Section~\ref{sec:discussion} discusses the broader +architectural implications of these results. Finally, +Sections~\ref{sec:conclusion} and \ref{sec:future} summarize the key +findings and outline future work on FFN, MoE, and full-model +optimization. diff --git a/docs/report/1H-codesign-paper/sections/04-allreduce.tex b/docs/report/1H-codesign-paper/sections/04-allreduce.tex index 03e5d46..cc65c28 100644 --- a/docs/report/1H-codesign-paper/sections/04-allreduce.tex +++ b/docs/report/1H-codesign-paper/sections/04-allreduce.tex @@ -1,4 +1,4 @@ -\section{All-Reduce Acceleration via PE\_IPCQ} +\section{PE\_IPCQ and Collective Communication} \label{sec:allreduce} \subsection{Why it is needed}