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.

Shared prefix
   ├─ Agent A: Analyze performance
   ├─ Agent B: Check correctness
   └─ Agent C: Find alternatives

Each sub-agent sees:

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:

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:

Q × Kᵀ

Multiple sub-agents can have different query rows while reading the same shared prefix KV.

Instead of executing:

Q_A × K_shared
Q_B × K_shared
Q_C × K_shared

independently, the runtime can combine the query rows:

Q_combined =
[
  Q_A
  Q_B
  Q_C
]

and execute:

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.

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:

The same Q rows are visible to all 32 PEs.

This does not require 32 independent physical copies. A possible implementation is:

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.

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:

16 × 16 × 16

and:

8 × 64 × 8

For:

8 sub-agents
20 query tokens per sub-agent

the combined number of Q rows is:

M = 8 × 20 = 160

If each PE owns 256 KV columns, a representative local GEMM is:

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.

Shared prefix KV pages

The runtime forks the logical context:

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:

8 sub-agents
20 new tokens per sub-agent

The runtime forms:

Q_A: 20 × d
Q_B: 20 × d
...
Q_H: 20 × d

and concatenates them along the row dimension:

Q_combined: 160 × d

The combined operation is:

(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:

Scores_p = Q_combined × K_pᵀ

The PE then computes local online-softmax state:

m_p[row]
l_p[row]
o_p[row, :]

For 160 rows, each PE conceptually produces:

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:

same row index across different sequence shards

It is never:

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:

m[160]
l[160]
o[160, d_v]

or:

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.

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.

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.

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.

32 PEs × 160 rows × 256 local columns
=
160 rows × 8192 total columns

If Q is temporally tiled into four groups of 40 rows:

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.

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:

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

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:

Policy A cost
=
Q multicast
+
hierarchical reduction

versus:

Policy B cost
=
KV regrouping, replication, or remote access
+
additional scheduling complexity

Policy B becomes attractive only if:

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:

1 KV head = 4 CUBEs = 32 PEs

use:

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:

T_A =
T_Q_multicast
+
T_local_GEMM
+
T_hierarchical_reduce

and:

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.

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.

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:

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:

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:

[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:

8 sub-agents
1000 output tokens per agent

A full-text join creates:

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.

{
  "claim": "memory_bandwidth_bottleneck",
  "confidence": 0.91,
  "evidence_ids": [17, 24]
}

The framework can serialize it compactly:

Finding: memory bandwidth bottleneck
Confidence: 0.91
Evidence: E17, E24

Example:

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.

[
  {"claim": "memory_bw", "confidence": 0.91},
  {"claim": "memory_bw", "confidence": 0.87},
  {"claim": "compute", "confidence": 0.42}
]

The framework can combine duplicates:

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.

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.

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.

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.

Sub-agent hidden states
        ↓
Compression or projection
        ↓
Small latent-token set
        ↓
Main model

Example:

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.

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:

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

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.