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:
2026-06-11 15:13:01 -07:00
parent a965db7056
commit d5267ec718
5 changed files with 83 additions and 46 deletions
@@ -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.