paper(intro): edit-pass through §1 Introduction done
- §1 Introduction rewritten with AHBM-first opening, in-paper KernBench platform overview, GQA motivation and three architectural requirements, three matching HW-SW co-design mechanisms (GQA-aware placement, PE_IPCQ, composite-command GEMM under PE_SCHEDULER), broader- applicability outlook (PE_IPCQ as collective substrate, composite command for fused operator pipelines), and outline. Consistent with the Executive Summary structure. - Executive Summary: extended broader-applicability closing with a forward-looking line previewing FFN/MoE integration. - §5 section heading renamed from "All-Reduce Acceleration via PE_IPCQ" to "PE_IPCQ and Collective Communication". - Title page date line refined to "2026 H1 Report". Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Binary file not shown.
@@ -25,7 +25,7 @@
|
||||
\date{
|
||||
\small
|
||||
AGI Computing Lab, System Technology Group\\
|
||||
H1 2026
|
||||
2026 H1 Report
|
||||
}
|
||||
|
||||
\begin{document}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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}
|
||||
|
||||
Reference in New Issue
Block a user