ywkang
d5267ec718
- §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>
84 lines
4.6 KiB
TeX
84 lines
4.6 KiB
TeX
\section{Introduction}
|
|
\label{sec:intro}
|
|
|
|
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.
|
|
|
|
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.
|
|
|
|
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.
|
|
|
|
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.
|