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