kernbench2/docs/report/1H-codesign-paper/Agentic_Runtime_Architecture.md

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