ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

report(1H): add Agentic Workloads (§6) and HW-Spec Search (§7) sections

Browse Source 
- §6 Supporting Agentic Workloads: how the fused GQA design extends to
  agentic fan-out/fan-in; three-layer split (framework/runtime/kernel);
  Stationary-KV vs Distributed-Q execution policies.
- §7 Hardware Performance-Spec Search for GQA: WIP stub (sweep intent
  over GEMM TFLOPS, MATH-engine ALUs, CUBE↔CUBE and SIP↔SIP BW).
- Renumber Discussion/Conclusion/Future-Work to 08/09/10; update
  main.tex input order and toc.md.
- Add Agentic_Runtime_Architecture.md design note; rebuild main.pdf.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook Kang
2026-07-21 15:00:37 -07:00
parent d0db9ff40e
commit edb30326ce
9 changed files with 1268 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/report/1H-codesign-paper/Agentic_Runtime_Architecture.md
+848
View File
@@ -0,0 +1,848 @@
# AHBM Agentic Runtime Architecture
## Scope
This note organizes the current design discussion for executing agentic workloads on a memory-centric AHBM-style architecture.
The scope of this version is limited to:
1. Motivation
2. AHBM hardware assumptions
3. Attention execution
4. Attention execution policies
5. Fan-in and join
MoE execution and whole-model layer-aware scheduling are intentionally deferred to a later section.
---
# 1. Motivation
Agentic workload patterns: loop, fan-out / fan-in. Loop 은 일반적인 LLM execution 과 다르지 않다. Sub-agent 가 만들어질 때 발생하는 Fan out / Fan in 이 주로 고려되어야 할 대상이다.
## 1.1 Agentic fan-out
An agentic workload often starts from one shared conversation or task context and then forks into several specialized sub-agents.
```text
Shared prefix
├─ Agent A: Analyze performance
├─ Agent B: Check correctness
└─ Agent C: Find alternatives
```
Each sub-agent sees:
```text
Shared prefix
+
Private role or instruction
+
Private generated continuation
```
The shared prefix may be long, while each private branch may initially contain only a small number of new tokens.
This creates two important execution properties:
1. All branches reuse the same prefix KV cache.
2. New query tokens from multiple branches can potentially be batched.
## 1.2 Why the workload is different from ordinary batching
Ordinary batching groups unrelated requests that happen to arrive at similar times.
Agentic fan-out is different because the branches are structurally related:
```text
Agent A context = Shared prefix + A suffix
Agent B context = Shared prefix + B suffix
Agent C context = Shared prefix + C suffix
```
The requests therefore have:
- identical prefix KV,
- different private suffix KV,
- potentially synchronized execution points,
- a later fan-in stage that combines their results.
This structure creates opportunities that are not available in unrelated-request batching.
## 1.3 Main optimization opportunity
For attention, the key operation is:
```text
Q × Kᵀ
```
Multiple sub-agents can have different query rows while reading the same shared prefix KV.
Instead of executing:
```text
Q_A × K_shared
Q_B × K_shared
Q_C × K_shared
```
independently, the runtime can combine the query rows:
```text
Q_combined =
[
Q_A
Q_B
Q_C
]
```
and execute:
```text
Q_combined × K_shared
```
The arithmetic is still row-independent, but the larger GEMM can improve utilization and amortize scheduling overhead.
## 1.4 Main design questions
The architecture must answer:
- How should shared KV be distributed across CUBEs and PEs?
- How should query rows from multiple sub-agents be grouped?
- How should online softmax state be reduced across sequence shards?
- Should Q be replicated or partitioned?
- How should large fan-out results be joined back into the main agent?
- Which data should be reused, recomputed, summarized, or materialized on demand?
---
# 2. AHBM Hardware Assumptions
## 2.2 Sequence parallelism
The sequence dimension of one KV head is distributed across the 32 PEs.
```text
KV head sequence
├─ PE0 owns sequence shard 0
├─ PE1 owns sequence shard 1
├─ ...
└─ PE31 owns sequence shard 31
```
Each PE stores a different part of the sequence for the same KV head. A query row must attend to all 32 sequence shards, so every PE computes a partial attention result for its local KV shard.
## 2.3 Q placement
The baseline follows an AHBM-style replicated-Q execution model.
Logically:
```text
The same Q rows are visible to all 32 PEs.
```
This does not require 32 independent physical copies. A possible implementation is:
```text
Chip-level Q source
↓
CUBE multicast
↓
Shared Q buffer within each CUBE
↓
8 PEs consume the same Q tile
```
Thus Q is logically replicated across PEs while physical traffic is reduced through multicast and shared buffering.
## 2.4 KV placement
KV remains stationary near the owning PE.
```text
PE0 → KV shard 0
PE1 → KV shard 1
...
PE31 → KV shard 31
```
The baseline avoids:
- KV remapping,
- KV replication,
- remote KV reads,
- page-table reconstruction for every query group.
## 2.5 GEMM engine assumptions
Representative PE GEMM tile shapes include:
```text
16 × 16 × 16
```
and:
```text
8 × 64 × 8
```
For:
```text
8 sub-agents
20 query tokens per sub-agent
```
the combined number of Q rows is:
```text
M = 8 × 20 = 160
```
If each PE owns 256 KV columns, a representative local GEMM is:
```text
Q_local: 160 × d
K_local: d × 256
```
Both `M = 160` and `N = 256` align well with the assumed tile shapes.
---
# 3. Attention Execution
## 3.1 Fan-out context structure
Assume the parent agent has already prefetched a shared prefix.
```text
Shared prefix KV pages
```
The runtime forks the logical context:
```text
Branch A page table:
[Shared prefix pages] + [A private pages]
Branch B page table:
[Shared prefix pages] + [B private pages]
Branch C page table:
[Shared prefix pages] + [C private pages]
```
The shared pages are referenced by multiple branches without being copied. Only private suffix pages are branch-specific.
## 3.2 Sub-agent query batching
Assume:
```text
8 sub-agents
20 new tokens per sub-agent
```
The runtime forms:
```text
Q_A: 20 × d
Q_B: 20 × d
...
Q_H: 20 × d
```
and concatenates them along the row dimension:
```text
Q_combined: 160 × d
```
The combined operation is:
```text
(160 × d) × (d × S)
```
where `S` is the total sequence length represented by all KV shards.
Each query row remains logically independent. Batching changes the execution shape, not the attention semantics.
## 3.3 Per-PE local attention
Each PE owns only a local sequence shard.
For PE `p`:
```text
Scores_p = Q_combined × K_pᵀ
```
The PE then computes local online-softmax state:
```text
m_p[row]
l_p[row]
o_p[row, :]
```
For 160 rows, each PE conceptually produces:
```text
m_p[160]
l_p[160]
o_p[160, d_v]
```
These may be processed in smaller row tiles for pipelining.
## 3.4 Online softmax merge
Each row has an independent online-softmax state.
The reduction is always:
```text
same row index across different sequence shards
```
It is never:
```text
different query rows reduced together
```
Therefore 160 query rows do not imply 160 serialized communication rounds. The implementation exchanges vector or tiled payloads such as:
```text
m[160]
l[160]
o[160, d_v]
```
or:
```text
m[16]
l[16]
o[16, d_v]
```
The row states can be communicated and merged in parallel.
## 3.5 Hierarchical reduction
Because one KV head spans four CUBEs and 32 PEs, reduction is hierarchical.
```text
32 PE partial states
↓
8-PE reduction inside each CUBE
↓
4 CUBE-level states
↓
4-CUBE reduction
↓
Final attention outputs
```
The same online-softmax merge primitive is used at both levels.
## 3.6 Prefill versus decode
### Decode
Decode typically has a very small number of Q rows.
```text
small Q
↓ multicast
32 PEs read local KV shards
↓
hierarchical reduction
```
The dominant concerns are usually KV bandwidth and reduction overhead.
### Prefill
Fan-out can make the query dimension much larger.
```text
Per-agent Q rows: 20
Number of agents: 8
Combined Q rows: 160
```
The larger `M` dimension can produce a much better GEMM shape. Agentic batching is therefore especially attractive for prefill or multi-token private suffix processing.
## 3.7 Compute-cost clarification
Replicating Q across 32 PEs does not multiply total attention FLOPs by 32. Each PE computes a different KV-column region.
```text
32 PEs × 160 rows × 256 local columns
=
160 rows × 8192 total columns
```
If Q is temporally tiled into four groups of 40 rows:
```text
4 × 32 PEs × 40 rows × 256 columns
=
160 rows × 8192 columns
```
The total arithmetic is identical. The policy changes Q traffic, reduction traffic, scheduling granularity, utilization, and buffering—not the mathematical amount of attention work.
---
# 4. Attention Execution Policies
## 4.1 Policy A: Replicated Q, stationary KV
Policy A keeps the existing KV placement unchanged.
```text
Q_combined
↓ multicast
PE0 computes against KV shard 0
PE1 computes against KV shard 1
...
PE31 computes against KV shard 31
```
Each PE executes a local GEMM such as:
```text
(160 × d) × (d × 256)
```
### Advantages
- No KV movement
- No KV replication
- No remote KV access
- No page-table regrouping
- Natural compatibility with sequence-parallel attention
- Large local GEMM shapes
- Simple hierarchical softmax reduction
### Costs
- Q must be distributed to all CUBEs
- Every row requires a 32-way logical reduction
- Large Q batches increase multicast and softmax-state traffic
Policy A is the natural baseline for the assumed AHBM mapping.
## 4.2 Policy B: Partitioned Q groups
A possible alternative is to divide query rows among PE groups.
```text
Q group 0 → PE group 0
Q group 1 → PE group 1
...
```
However, every query row must still attend to the complete KV sequence. Because the 32 PEs already represent 32 sequence shards, spatially partitioning Q means each Q group must somehow access all KV shards.
This requires one of the following:
1. Regroup KV shards for each Q group.
2. Read remote KV through symmetric memory.
3. Replicate KV across Q-processing groups.
4. Time-multiplex the same PEs over Q groups.
The first three add memory-system complexity. The fourth is mainly temporal tiling and does not provide true spatial Q partitioning.
## 4.3 Why Policy B is not automatically better
The comparison is:
```text
Policy A cost
=
Q multicast
+
hierarchical reduction
```
versus:
```text
Policy B cost
=
KV regrouping, replication, or remote access
+
additional scheduling complexity
```
Policy B becomes attractive only if:
```text
Q distribution cost + reduction cost
>
remote or regrouped KV cost
```
Relevant factors include Q size, multicast bandwidth, reduction bandwidth, symmetric-memory bandwidth, remote-KV latency, KV replication capacity, page-table overhead, and GEMM utilization.
## 4.4 Recommended baseline
For:
```text
1 KV head = 4 CUBEs = 32 PEs
```
use:
```text
Policy A:
replicated Q + stationary sequence-sharded KV
```
Policy B should be treated as an adaptive or future policy for cases where Q batches become extremely large, multicast becomes a bottleneck, reduction traffic dominates, or remote KV access becomes inexpensive.
## 4.5 Runtime decision model
A future runtime can estimate:
```text
T_A =
T_Q_multicast
+
T_local_GEMM
+
T_hierarchical_reduce
```
and:
```text
T_B =
T_Q_partition
+
T_remote_or_replicated_KV
+
T_local_GEMM
+
T_group_reduce
+
T_remap
```
The runtime selects the lower-cost policy for the current sequence length, agent count, Q size, KV-head mapping, bandwidth state, and interconnect congestion.
---
# 5. Fan-In and Join
## 5.1 Fan-in problem
After fan-out, each sub-agent produces a private result.
```text
Agent A → result A
Agent B → result B
Agent C → result C
```
The main agent must synthesize these results into a final continuation.
## 5.2 Why branch KV cannot be concatenated
Each branch has a different causal token history.
```text
A history:
Shared prefix + A instruction + A reasoning
B history:
Shared prefix + B instruction + B reasoning
C history:
Shared prefix + C instruction + C reasoning
```
The K/V tensors of a token depend on token content, position, preceding causal context, and every transformer layer.
Therefore:
```text
A private KV
+
B private KV
+
C private KV
```
does not form the KV cache of any valid single token sequence.
## 5.3 Baseline text join
The standard framework behavior is:
```text
Sub-agent result text
↓
Framework gathers and formats results
↓
Main-agent join prompt
↓
Continuation prefill
↓
New main-agent KV suffix
```
The framework normally aggregates the inputs, while the main LLM performs final synthesis.
## 5.4 Main-agent KV after join
The main-agent page table becomes:
```text
[Shared prefix KV pages]
+
[Join-input KV pages]
+
[Main continuation KV pages]
```
The private A/B/C pages remain separate and are not attached directly.
## 5.5 Cost of full-text gather
Assume:
```text
8 sub-agents
1000 output tokens per agent
```
A full-text join creates:
```text
8000 join-input tokens
```
These tokens must pass through all transformer layers during continuation prefill. This increases prefill work, KV allocation, future decode KV reads, and context-window occupancy.
## 5.6 Schema-constrained join
A practical optimization is to constrain each sub-agent to a compact result schema.
```json
{
"claim": "memory_bandwidth_bottleneck",
"confidence": 0.91,
"evidence_ids": [17, 24]
}
```
The framework can serialize it compactly:
```text
Finding: memory bandwidth bottleneck
Confidence: 0.91
Evidence: E17, E24
```
Example:
```text
Before:
8 × 1000 = 8000 tokens
After:
8 × 50 = 400 tokens
```
The main LLM still performs final reasoning, but the continuation prefill is much shorter.
## 5.7 Deduplication and aggregation
Different sub-agents may produce overlapping findings.
```json
[
{"claim": "memory_bw", "confidence": 0.91},
{"claim": "memory_bw", "confidence": 0.87},
{"claim": "compute", "confidence": 0.42}
]
```
The framework can combine duplicates:
```text
Primary finding:
- Memory-bandwidth bottleneck
- Supported by 2 agents
- Maximum confidence: 0.91
Alternative:
- Compute bottleneck
- Confidence: 0.42
```
This is preprocessing, not final reasoning.
## 5.8 Pointer-based join
Detailed evidence can remain outside the initial main-agent context.
```text
Agent A:
- Finding: memory-bandwidth bottleneck
- Evidence handle: E17
Agent B:
- Finding: reduction error
- Evidence handle: E24
```
The main agent initially receives only summaries and handles. The framework retrieves and appends evidence only when requested.
```text
Effective input
=
summary tokens
+
tokens for evidence actually used
```
The transformer does not directly dereference the handle; the framework resolves it through retrieval or a tool call.
## 5.9 Hierarchical reduction
For large fan-out width, local reducer agents can summarize groups of branches.
```text
16 sub-agents
↓
4 local reducers
↓
4 summaries
↓
Main agent
```
Benefits:
- Reducers operate in parallel.
- The final join prompt is shorter.
- Main-agent KV growth is smaller.
- Fan-in traffic is organized hierarchically.
Costs:
- Additional reducer inference
- Possible loss of detail
- Need for fallback evidence retrieval
## 5.10 Latent-state join
A more aggressive approach is to replace long text outputs with learned latent representations.
```text
Sub-agent hidden states
↓
Compression or projection
↓
Small latent-token set
↓
Main model
```
Example:
```text
1000 text tokens
↓
16 latent tokens
```
This could reduce join prefill and KV growth, but it requires training the main model to consume latent tokens, aligning branch representations, defining causal and positional semantics, and validating accuracy. It is a model-system co-design direction rather than a drop-in runtime optimization.
## 5.11 KV and compute impact
The shared parent prefix KV is already present. The primary optimization target is the newly created join suffix.
```text
New main KV size
∝
number of join-input tokens
```
Reducing join input reduces continuation-prefill work, new KV allocation, Q rows processed during join, later decode-time KV reads, and context-window consumption.
A practical flow is:
```text
Sub-agent full outputs
↓
Schema-constrained results
↓
Deduplication and ranking
↓
Summaries + evidence handles
↓
Selective evidence materialization
↓
Main-agent continuation prefill
```
## 5.12 Recommended practical join design
For an unchanged LLaMA-style model on AHBM:
1. Reuse shared prefix KV through page-table references.
2. Keep branch-private KV isolated.
3. Require schema-constrained sub-agent outputs.
4. Deduplicate repeated claims in the framework.
5. Pass compact summaries, confidence values, and evidence handles.
6. Retrieve detailed evidence only when requested.
7. Introduce hierarchical reducer agents for wide fan-out.
8. Keep the main LLM responsible for final synthesis.
---
# Current Baseline Summary
```text
Shared prefix prefill
↓
Shared KV page reuse
↓
Agentic fan-out
↓
Combine private Q rows
↓
Replicated-Q attention over stationary sequence-sharded KV
↓
Hierarchical online-softmax reduction
↓
Independent private branch continuations
↓
Schema/pointer-based fan-in
↓
Main-agent continuation prefill
```
Main design principles:
- Reuse shared prefix KV without copying it.
- Batch sub-agent Q rows when they read the same KV.
- Keep KV stationary and multicast Q.
- Reduce softmax state hierarchically by matching row index.
- Do not directly merge branch-private KV.
- Reduce fan-in cost by shortening and selectively materializing join inputs.
docs/report/1H-codesign-paper/build/main.pdf
BIN
View File
Binary file not shown.
docs/report/1H-codesign-paper/main.tex
+6 -4
View File
@@ -28,7 +28,7 @@
\date{
\small
AGI Computing Lab, System Technology Group\\
2026 H1 Report
2026 Q1-Q3 Report
}
\begin{document}
@@ -40,8 +40,10 @@ AGI Computing Lab, System Technology Group\\
\input{sections/03-gemm}
\input{sections/04-allreduce}
\input{sections/05-gqa}
\input{sections/06-discussion}
\input{sections/07-conclusion}
\input{sections/08-future-work}
\input{sections/06-agentic}
\input{sections/07-hw-spec-search}
\input{sections/08-discussion}
\input{sections/09-conclusion}
\input{sections/10-future-work}
\end{document}
docs/report/1H-codesign-paper/sections/06-agentic.tex
+309
View File
@@ -0,0 +1,309 @@
\section{Supporting Agentic Workloads}
\label{sec:agentic}
The fused GQA kernel of \S\ref{sec:gqa} solved the efficient execution of a
\emph{single} logical attention stream: one query stream against one KV
cache, tiled and merged with an online softmax. Agentic inference
introduces a different execution model, in which one request dynamically
\emph{forks} into several cooperating reasoning branches and later
\emph{joins} them. The purpose of this section is to show that this
multi-stream model needs no new attention algorithm---the expensive
hardware machinery built for \S\ref{sec:gqa} is reused unchanged---and to
work out the resulting design across three implementation layers. The
central claim, stated once here and defended throughout, is:
\begin{quote}
\emph{The proposed agentic execution does not change the semantics of
transformer attention; it changes only how independent query rows are
scheduled over a shared KV cache.}
\end{quote}
\noindent Attention stays identical; only the scheduling differs. What
follows is a \emph{design} built on the measured \S\ref{sec:gqa} kernel; it
is not yet a KernBench measurement, and points that go beyond the
implemented path are flagged as such.
\subsection{Motivation: from one attention stream to many}
\label{sec:agentic-why}
Agentic execution has two recurring shapes: a \emph{loop} (an agent that
repeatedly generates, calls a tool, and continues) and a \emph{fan-out /
fan-in} (a parent agent forks several specialised sub-agents and later
synthesises their results). The loop is ordinary autoregressive decoding
and needs nothing new. The fan-out is the interesting case, because the
branches are \emph{structurally related} rather than independent:
\[
\begin{aligned}
\text{context}_i \;=\;& \underbrace{\text{shared prefix}}_{\text{common}} \\
&+\; \underbrace{\text{private role}_i + \text{private suffix}_i}_{\text{branch-specific}} .
\end{aligned}
\]
Ordinary batching groups unrelated requests that merely arrive together;
agentic fan-out groups requests that share an identical prefix KV cache and
differ only in a short private suffix. That structure creates two levers
that unrelated-request batching cannot pull:
\begin{enumerate}
\item \textbf{Shared-prefix KV reuse.} All branches attend to the same
prefix KV, so it is read (and stored) once, not once per branch.
\item \textbf{Query-row batching.} The new query rows of many sub-agents
read the \emph{same} shared KV, so they can be concatenated into one
taller GEMM.
\end{enumerate}
\noindent Both levers land directly on the \S\ref{sec:gqa} design, which
already multicasts a (small) query against a stationary KV cache and merges
the result hierarchically. Fan-out simply makes the query taller, and
fan-in adds a join step; neither touches the attention math.
\subsection{Architecture: three execution layers}
\label{sec:agentic-arch}
Before the mechanics, we fix the layering, because the recurring reader
question is ``whose job is this?''. Agentic execution divides cleanly along
the request-flow layers already used in this report: the \textbf{agentic
framework} decides \emph{what} to run and how to combine it, the
\textbf{runtime} decides \emph{how} to map it onto the hardware without
copying shared state, and the \textbf{kernel} does the actual math and
reduction on the PEs.
\[
\text{Framework} \;\longrightarrow\; \text{Runtime}
\;\longrightarrow\; \text{Kernel}
\]
\noindent Keeping the split this way preserves the topology-agnostic
runtime boundary: the framework never sees the SIP/CUBE/PE hierarchy, and
the kernel never sees the agent tree. The remainder of the section walks
the fan-out $\rightarrow$ runtime $\rightarrow$ kernel $\rightarrow$ fan-in
path through exactly these three layers.
Table~\ref{tab:agentic-levels} states each layer's responsibilities; the
kernel row is unchanged from \S\ref{sec:gqa}.
\begin{table*}[t]
\centering
\caption{Division of responsibilities for agentic attention. Only the
framework and runtime rows are new work; the kernel is the
\S\ref{sec:gqa} fused GQA kernel, unchanged in its math and reduction.}
\label{tab:agentic-levels}
\small
\begin{tabular}{@{}p{0.16\textwidth}p{0.40\textwidth}p{0.36\textwidth}@{}}
\toprule
\textbf{Level} & \textbf{Responsibilities} & \textbf{Explicitly not its job} \\
\midrule
\textbf{Agentic framework}
& Manage the agent tree: fork sub-agents and assign roles; identify the
shared prefix; mark which branches are co-schedulable (same shared prefix).
Fan-in: enforce schema-constrained results, deduplicate/rank findings,
hold evidence behind pointers, insert hierarchical reducers, and drive the
main agent's final synthesis.
& No topology, routing, KV placement, or scheduling. Sees agents and text,
not CUBEs or PEs. \\
\addlinespace
\textbf{Runtime}
& Fork the logical context by page table (shared-prefix pages $+$ private
suffix pages) with \emph{no copy} of shared KV. Concatenate co-scheduled
sub-agent query rows into $Q_{\text{cmb}}$. Select the execution policy
(Stationary-KV vs.\ Distributed-Q, \S\ref{sec:agentic-choice}) from the
cost model. Launch the
fused GQA kernel (composite command $+$ PE\_IPCQ) and hand the batched work
and policy to the simulation engine.
& No attention math and no per-hop routing (delegated to the engine and
policy); no agent-level semantics. \\
\addlinespace
\textbf{Kernel}
& Multicast $Q_{\text{cmb}}$ to all PEs; per-PE local attention over the
stationary KV shard via composite-command $Q\!\cdot\!K^{\top}$ and
$P\!\cdot\!V$; produce per-row online-softmax state; merge that state
hierarchically over PE\_IPCQ by matching row index (row-parallel, not
serialised).
& No knowledge of agents, page tables, or policy selection. Identical to
\S\ref{sec:gqa}; a taller $Q$ is the only difference. \\
\bottomrule
\end{tabular}
\end{table*}
\subsection{Stationary-KV Execution Policy}
\label{sec:agentic-policyA}
The core attention operation is $Q\!\cdot\!K^{\top}$, and under fan-out
multiple sub-agents contribute different query rows while reading a common
$K$. Rather than issuing $Q_A\!\cdot\!K_{\text{shared}}$,
$Q_B\!\cdot\!K_{\text{shared}}$, \dots\ as separate small GEMMs, the runtime
concatenates the rows,
\[
Q_{\text{cmb}} = \begin{bmatrix} Q_A \\ Q_B \\ \vdots \end{bmatrix},
\qquad
Q_{\text{cmb}}\!\cdot\!K_{\text{shared}},
\]
and issues one taller GEMM. The arithmetic is unchanged---each row is still
independent---but the $M$ dimension grows. For a representative fan-out of
$8$ sub-agents at $20$ new tokens each, $M = 8\times20 = 160$; with a per-PE
KV shard of $256$ sequence positions this is a $160\times d$ by $d\times256$
local product---an $M{=}160$, $N{=}256$ shape that sits well above the
per-command issue overhead the composite command (\S\ref{sec:gemm}) is built
to amortise. Fan-out is therefore especially valuable during \emph{prefill}
of the private suffixes, where the combined $M$ is large; during decode the
query stays short and the workload remains, as \S\ref{sec:gqa} found,
KV-bandwidth bound.
The Stationary-KV policy maps onto the \S\ref{sec:gqa} placement
essentially unchanged: logically replicated $Q$ over a stationary,
sequence-sharded KV cache. One
KV head spans four CUBEs and 32 PEs, with each PE owning a different KV
\emph{sequence} shard $K_p[S_p,d]$, $V_p[S_p,d_v]$. The combined $Q$ is
multicast to all of them; each PE forms a complete local score
$Q_{\text{cmb}}\!\cdot\!K_p^{\top}$ over its own shard, and the per-row
online-softmax state is reduced hierarchically---8-PE merge inside each
CUBE, then a 4-CUBE merge---using the same PE\_IPCQ merge primitive from
\S\ref{sec:allreduce}.
\paragraph{The reduction algorithm does not change.} This is the point to
emphasise. Agentic batching does \emph{not} require a new reduction
algorithm: the hierarchical online-softmax reduction of \S\ref{sec:gqa}
remains exactly as-is, and only the query batch becomes larger. The
reduction is always \emph{same row index across sequence shards}, never
\emph{different query rows against each other}, so $M$ batched rows do
\emph{not} become $M$ serialised communication rounds; the $(m,\ell,O)$ row
states are exchanged as vectors/tiles and merged in parallel. Because the
merge is row-indexed, a taller $Q$ widens each payload but adds no rounds.
This is what makes agentic support a reuse of \S\ref{sec:gqa} rather than a
redesign.
\paragraph{Logical vs.\ physical $Q$ replication.} ``Replicated $Q$'' is a
\emph{logical} statement. Because the shard axis is the KV \emph{sequence}
dimension $S$, every PE must form the full score $Q\!\cdot\!K_p^{\top}$
against its local $K_p$ and therefore needs $Q$'s \emph{entire} hidden
dimension $d$; what is partitioned across PEs is $K$/$V$ along $S$, never
$Q$ along its columns. Splitting $Q$ (and $K$) on the hidden dimension
would instead make each PE's product \emph{partial} and force a pre-softmax
hidden-dimension reduction ($QK^{\top}=\sum_i Q_iK_i^{\top}$)---that is
tensor-/head-parallel attention, a different structure from the
sequence-parallel one assumed here, and one that cannot coexist with using
the PE axis for sequence shards. Logical replication also does not mean 32
physical copies: $Q$ can be multicast once into a CUBE-local shared buffer
(shared SRAM) that all PEs in the CUBE read, and a large $Q$ can further be
\emph{row}-tiled in time ($Q[0{:}16,:],\,Q[16{:}32,:],\dots$)---row tiling
splits the $M$ dimension, not the hidden-dimension columns. In short:
the Stationary-KV policy uses logically replicated $Q$ across
sequence-parallel PEs while $K$ and $V$ are partitioned along the sequence
dimension; $Q$ may be
temporally row-tiled or physically shared through multicast buffers, but it
is not partitioned along the hidden-dimension columns.
\paragraph{Replication is not $32\times$ the compute work (attention
FLOPs).} Multicasting $Q$ to 32 PEs does not multiply attention FLOPs,
because each PE computes against a different KV sequence shard rather than
the same one. Let the KV sequence length be $S$; with sequence parallelism
over 32 PEs, each PE owns $S/32$ positions. The score GEMM
$Q[M,d]\!\cdot\!K^{\top}[d,S]$ costs $\propto M\,d\,S$, so each PE performs
$M\,d\,(S/32)$ and the 32 shards sum to
\[
32 \cdot M\,d\,\tfrac{S}{32} \;=\; M\,d\,S,
\]
identical to attention over one undivided sequence. Replication therefore
changes $Q$ distribution, reduction traffic, buffering, and scheduling---not
the total attention FLOPs.
\subsection{Distributed-Q Execution Policy (alternative)}
\label{sec:agentic-policyB}
The natural alternative is to partition (distribute) the query rows across
PE groups (\emph{the Distributed-Q policy}) rather than replicate them. It
is not automatically
better. Because the 32 PEs already shard the KV \emph{sequence}, every query
row must still attend to \emph{all} shards; partitioning $Q$ across PE
groups therefore forces each group to reach every KV shard, which requires
one of: regrouping KV shards per $Q$ group, replicating KV across groups, or
reading remote KV through symmetric memory. Each of these adds
memory-system complexity that the Stationary-KV policy avoids entirely.
Time-multiplexing the same PEs over $Q$ groups is the fourth option, but
that is temporal tiling---already available under the Stationary-KV policy
as row tiling---not true spatial $Q$ partitioning. The Distributed-Q policy
is thus a proposed adaptive extension, not the
baseline, and is only worth its complexity when $Q$ grows large enough that
multicast and reduction traffic dominate remote/regrouped-KV cost.
\subsection{Why the Stationary-KV policy is the baseline}
\label{sec:agentic-choice}
The choice reduces to a cost comparison the runtime can estimate. The
Stationary-KV policy pays for $Q$ multicast and a hierarchical reduction;
the Distributed-Q policy pays for moving or replicating KV plus extra
scheduling:
\[
T_{\text{SK}} = T_{Q\text{-mcast}} + T_{\text{local GEMM}} + T_{\text{hier.\ reduce}},
\]
\[
\begin{aligned}
T_{\text{DQ}} = {}& T_{Q\text{-part}} + T_{\text{remote/repl.\ KV}} + T_{\text{local GEMM}} \\
&+ T_{\text{grp.\ reduce}} + T_{\text{remap}}.
\end{aligned}
\]
For the assumed mapping (one KV head $=$ 4 CUBEs $=$ 32 PEs), the KV cache
is large and stationary while $Q$ is comparatively small, so $T_{\text{SK}}$
is the lower cost and the Stationary-KV policy is the recommended baseline.
A future runtime can compute both estimates per launch---from the current
agent count, $Q$ size, KV-head mapping, and interconnect state---and switch
to the Distributed-Q policy only in the regime where a very large $Q$ batch
makes multicast and reduction traffic outweigh the cost of remote or
regrouped KV. Until that regime is measured, the Distributed-Q policy
remains a designed, not-yet-implemented option.
\subsection{Fan-in: joining sub-agent branches}
\label{sec:agentic-fanin}
Fan-out is only half of the pattern; after it, each branch produces a
private continuation and the parent must synthesise them. We treat fan-in
as a runtime/framework design problem with a clear optimization ladder.
\paragraph{Problem.} The branch KV caches \emph{cannot} be concatenated.
Each token's K/V depends on its full causal history, so stacking several
branches' private KV does not form the cache of any single valid
sequence. Join must therefore happen at the token/text level, and its cost
is dominated by the number of join-input tokens, because the new main-agent
KV grows in proportion to them.
\paragraph{Baseline.} A naive full-text gather concatenates every branch's
raw output: $8$ agents $\times\ 1000$ tokens $=\ 8000$ tokens pushed through
every layer during continuation prefill---inflating prefill work, KV
allocation, and later decode-time KV reads. This is the cost the ladder
below drives down.
\paragraph{Optimization 1 --- schema-constrained results.} Constrain each
sub-agent to emit a compact structured result (claim, confidence, evidence
handles) instead of free text, cutting join input by an order of magnitude
($8\times1000 \to 8\times50$).
\paragraph{Optimization 2 --- deduplication and ranking.} Overlapping
findings across branches are merged and ranked in the framework before they
reach the main context. This is preprocessing, not reasoning, and shrinks
the input further without involving the main model.
\paragraph{Optimization 3 --- pointer-based evidence.} Detailed evidence
stays outside the main context behind handles, materialised only when the
main agent actually requests it, so the effective input is summaries plus
only the evidence truly used.
\paragraph{Optimization 4 --- hierarchical reducers.} For wide fan-out,
intermediate reducer agents summarise groups of branches, shrinking the
final join prompt and organising fan-in traffic hierarchically---mirroring
how the hierarchical online-softmax merge organises attention reduction.
\paragraph{Future --- latent-state join.} A more aggressive step replaces
text outputs with a few learned latent tokens, further cutting join prefill
and KV growth. This requires training the main model to consume latent
tokens and aligning branch representations; it is a model--system co-design
direction, not a drop-in runtime optimization, and is out of scope here.
\medskip
\noindent Throughout, the shared prefix KV is reused by page-table
reference and only the join suffix is new, so shortening join input is the
primary lever on continuation-prefill cost, KV growth, and later
decode-time KV reads. Taken together, \S\ref{sec:agentic-policyA}--\ref{sec:agentic-fanin}
show why this workload is a natural extension of the 1H design rather than a
new one: the decisive hardware levers of this report---cheap composite
issue and a fast on-device reduction path---are exactly what make agentic
fan-out efficient, and no part of the attention math is altered to get
there. Quantifying the fan-out speedup and the fan-in join savings on
KernBench is left as measured 2H work.
docs/report/1H-codesign-paper/sections/07-hw-spec-search.tex
+67
View File
@@ -0,0 +1,67 @@
\section{Hardware Performance-Spec Search for GQA}
\label{sec:hwspec}
\emph{Work in progress --- this section states the study's intent and
method; experimental data, figures, and conclusions are not yet
available and will be added in a later revision.}
The studies so far fixed the modeled hardware (\S\ref{sec:hw}) and
varied the \emph{software}: the composite command, the reduction path, and
the KV placement. This section inverts the question. Given the fused GQA
kernel of \S\ref{sec:gqa} as the target workload, \emph{which hardware
specification best serves it}, and where does spending more silicon stop
paying off? The discussion of \S\ref{sec:discussion} already gives a
strong prior---attention is data-movement bound, so raw MAC throughput is
not the binding constraint---but that conclusion was drawn at one operating
point. A systematic sweep is needed to turn it into a defensible
performance-spec recommendation across context lengths, agent counts, and
sharding regimes.
\subsection{Sweep axes}
\label{sec:hwspec-axes}
Four hardware knobs are varied, spanning the compute side and both levels
of the interconnect so that the compute/communication balance can be
located rather than assumed. The KernBench cost model already exposes each
as a parameter, so the sweep reuses the same deterministic engine as the
rest of the report; no production change is required.
\begin{table}[h]
\centering
\caption{Hardware knobs varied in the performance-spec search. Ranges and
step counts are to be finalised with the experiment.}
\label{tab:hwspec-axes}
\small
\begin{tabular}{@{}p{0.30\textwidth}p{0.14\textwidth}@{}}
\toprule
\textbf{Knob} & \textbf{Axis} \\
\midrule
GEMM throughput (TFLOPS) & compute \\
MATH-engine ALU count & compute \\
CUBE-to-CUBE (die-to-die) bandwidth & interconnect \\
SIP-to-SIP (card-to-card) bandwidth & interconnect \\
\bottomrule
\end{tabular}
\end{table}
The first two axes scale the on-PE compute engines; the latter two scale
the two inter-device links that carry the KV reduction and cross-device
softmax merge. Sweeping them jointly---rather than one at a time---is what
exposes the interactions: for example, whether added die-to-die bandwidth
only helps once card-to-card bandwidth is also raised, or whether GEMM
throughput is genuinely inert for attention once issue overhead is removed.
\subsection{Method and target output}
\label{sec:hwspec-method}
The intended procedure is a joint sweep over the four axes, running the
fused GQA kernel at representative decode and prefill configurations
(including the agentic fan-out shapes of \S\ref{sec:agentic}), and reading
latency and per-engine busy time from the engine's completion timestamps
and op log---the same measurement path used in \S\ref{sec:gqa}. The target
deliverables are: (i) the sensitivity of GQA latency to each knob in
isolation, (ii) the Pareto frontier of latency against a simple
area/cost proxy, and (iii) a recommended balanced specification---the knob
combination past which further investment does not move GQA latency. These
results are the natural quantitative counterpart to the qualitative
ranking in \S\ref{sec:discussion}, and completing them is a 2H objective.
docs/report/1H-codesign-paper/sections/06-discussion.tex → docs/report/1H-codesign-paper/sections/08-discussion.tex
View File
docs/report/1H-codesign-paper/sections/07-conclusion.tex → docs/report/1H-codesign-paper/sections/09-conclusion.tex
View File
docs/report/1H-codesign-paper/sections/08-future-work.tex → docs/report/1H-codesign-paper/sections/10-future-work.tex
+20 -9
View File
@@ -1,4 +1,4 @@
\section{Future Work --- 2H}
\section{Future Work}
\label{sec:future}
The 1H work covered attention end to end. The natural next step is to
@@ -30,11 +30,22 @@ vs.\ sequence parallelism---under a single, software-stack-independent
model, so the interconnect and memory implications of each choice are
measured rather than assumed.
\paragraph{Agentic workloads.} Agentic inference interleaves many
short, bursty decode requests with tool use and long shared contexts,
which stresses the system differently from a single long generation:
context reuse across requests, dynamic batching, and uneven expert load all
change how compute and data should be dispersed. Characterizing how total
compute and data movement distribute across the SIP/CUBE/PE hierarchy under
such workloads---and which of the 1H hardware levers still dominate when the
workload is this irregular---is the broader 2H agenda.
\paragraph{Agentic workloads: from design to implementation.}
Section~\ref{sec:agentic} established \emph{how} the fused GQA design
extends to agentic fan-out/fan-in and \emph{which} responsibilities fall to
the framework, the runtime, and the kernel. The 2H step is no longer to
analyse the workload but to \emph{build} that support: the detailed design
and implementation of all three levels. This is necessary work rather than
optional, because today's agentic frameworks are effectively all
closed-source---there is no open substrate that exposes shared-prefix KV
reuse, cross-agent query batching, and schema/pointer-based fan-in down to
the hardware. Concretely, 2H targets: a \textbf{framework} layer that
manages the agent tree, identifies co-schedulable branches, and enforces
compact fan-in; a \textbf{runtime} layer that forks logical contexts by
page table without copying shared KV, batches sub-agent query rows, and
selects the execution policy; and a \textbf{kernel} layer that runs the
batched replicated-Q attention and hierarchical online-softmax merge over
PE\_IPCQ. Bringing these up on KernBench turns the
Section~\ref{sec:agentic} design into a measured, end-to-end agentic path
and lets us confirm which of the 1H hardware levers still dominate once the
full framework/runtime/kernel stack is in place.
docs/report/1H-codesign-paper/toc.md
+18 -7
View File
@@ -45,17 +45,28 @@
6. **Fused Grouped Query Attention** (composite + PE_IPCQ) — `sections/05-gqa.tex`
necessity · design · results · analysis. (+ GQA seq/head/user configs)
7. **Discussion** — `sections/06-discussion.tex`
7. **Supporting Agentic Workloads** — `sections/06-agentic.tex`
How the §6 GQA design extends to agentic fan-out/fan-in: shared-prefix KV
reuse, batched replicated-Q attention, schema/pointer-based fan-in.
Implementation split across three levels (agentic framework / runtime /
kernel) with each component's responsibilities. Design, not yet measured.
8. **Hardware Performance-Spec Search for GQA** — `sections/07-hw-spec-search.tex`
*Work in progress.* Joint sweep of GEMM TFLOPS, MATH-engine ALU count,
CUBE↔CUBE (die-to-die) BW, SIP↔SIP (card-to-card) BW to find the balanced
spec for GQA. Intent + method only; data/figures/conclusions deferred.
9. **Discussion** — `sections/08-discussion.tex`
Which HW changes are meaningful, and under what regimes (cross-cutting).
8. **Conclusion** — `sections/07-conclusion.tex`
The codesign thesis, stated plainly, supported by the measured results.
10. **Conclusion** — `sections/09-conclusion.tex`
The codesign thesis, stated plainly, supported by the measured results.
9. **Future Work — 2H** — `sections/08-future-work.tex`
Add the FFN/MoE layer toward full LLM decoding; how compute & data should
be distributed for agentic / MoE workloads.
11. **Future Work — 2H** — `sections/10-future-work.tex`
Add the FFN/MoE layer toward full LLM decoding; how compute & data should
be distributed for agentic / MoE workloads.
10. **References** *(optional)* — external literature only (FlashAttention,
12. **References** *(optional)* — external literature only (FlashAttention,
Megatron-LM, GPT-3, Llama 3). No ADR/SPEC entries.
## Per-section structure (§4/§5/§6)