\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.