ywkang
6b6e29968a
§2 (KernBench Platform): - Fix Table 2 HBM aggregate BW (1024 → 2048 GB/s); drop stale hbm_total_bw_gbs from topology.yaml (never read by sim_engine) - Split PE_CPU / PE_SCHED fixed-cost row; disambiguate from the 40-cycle command-dispatch FIXED term - §2.2 two-pass: expand to describe Pass 1 timing and Pass 2 data data-correctness path - §2.3 dispatch model: add motivation sentence for the descriptor-size linear form - §2.4 Accuracy: reorder GEMM → All-reduce → Probe → Simplifications; drop FSIM aside (already covered by §4 Fig 5 caption); soften 'every ns' → 'every modeled latency contribution' - §2.5 HW config: add 64 TFLOP/s vs 2048 GB/s (~31 FLOP/byte) balance-point intuition and forward pointer to §5 GQA decode - Fig 1 caption: separate illustrative topology from experimental configuration §1: tighten GQA-as-primary-bandwidth-bottleneck framing. Figures: regenerate SIP / CUBE architecture (SVG sources + PDF + generator scripts). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
87 lines
4.8 KiB
TeX
87 lines
4.8 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. GQA
|
|
dominates inference-time memory traffic and KV-cache capacity in modern
|
|
LLM serving, making it the primary bandwidth bottleneck on memory-centric
|
|
architectures such as AHBM. 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.
|