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

Compare commits

base: ywkang/kernbench2:39fc2e953f8a4f1c7c8c1391fd1e6684bdf91260
Branches Tags
ywkang/kernbench2:master ywkang/kernbench2:feat/gqa-composite-commands ywkang/kernbench2:sccl-distributed-allreduce ywkang/kernbench2:kernbench_Apr16
...
compare: ywkang/kernbench2:master
Branches Tags
ywkang/kernbench2:master ywkang/kernbench2:feat/gqa-composite-commands ywkang/kernbench2:sccl-distributed-allreduce ywkang/kernbench2:kernbench_Apr16

164 Commits
39fc2e953f ... master

Author SHA1 Message Date
mukesh fc4747668e paper(§5): add roofline, capacity planning, parallelism selection subsections 
Three new subsections at the top of §5, motivated by the deployment-sizing
questions the analytical visualization tool exposes:

- 5.1 Roofline Analysis — B*, L*, short vs long context batch effect
  (measured: 19× per-token cost reduction at short context, 28% at long)
- 5.2 Capacity Planning — 3-axis sizing (weights / KV / SLO), SLO targets,
  regime rules, deployment templates, per-axis playbook
- 5.3 Parallelism Selection — TP × CP × PP × DP × EP comparison,
  symptom-driven axis selection, add/stop criteria, misconceptions

Restructure:
- 05-gqa.tex trimmed to section header + intro
- Existing 5.4-5.7 content (placement, short/long ctx, composite) moved
  to 05x-fused-kernel.tex
- Summary extracted to 05z-summary.tex, cross-refs updated
- main.tex \input order sets the requested subsection sequence
- lmodern loaded to satisfy microtype font-expansion

Two new figures generated from tests/analytical_visualization/chip_roofline:
- roofline_short_context.png (S_kv=8K, batch wins)
- roofline_long_context.png (S_kv=1M, batch stalls at KV floor)
- Generator: tests/analytical_visualization/_gen_roofline_paper_figs.py

4 tables (regime rules, deployment templates, playbook, parallelism
criteria) use table* so they span both columns without overflow.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-30 15:56:41 -07:00
mukesh bfe3e0e8d1 analytical-viz: parallelism rules — Section 11 TP vs CP shared-budget 
New Section 11 'TP vs CP: the shared-budget tradeoff' on the
Parallelism rules tab. Captures the core insight that fixing TP·CP
doesn't equalize the two configurations — larger TP additionally
divides model-weight traffic, larger CP only divides the sequence
dimension.

Five sub-tables:

  11a. Same TP·CP = 16 → different weight sharding
       (TP=8/CP=2 reads 4× fewer weight bytes per PE than
        TP=2/CP=8 at the same KV division)
  11b. What each dimension actually shards
       (TP·CP·PP·DP × Weights·KV·Attention-math)
  11c. When to prefer TP vs CP (6 situations mapped to the winning
       axis with reasoning)
  11d. CP communication profile (prefill vs decode) — bytes/hop,
       overlap potential, real cost, notes on kv vs qoml variants
  11e. Decode cost decomposition mental model:
       T_decode = T_weight_read(TP) + T_KV_read(TP,CP)
                + T_TP_comm(TP) + T_CP_merge(CP)

Reframes the 'is CP heavier than TP' question as 'which cost
dominates T_decode, and which axis attacks it?'.

Smoke test asserts the section title appears exactly once.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-30 09:26:22 -07:00
mukesh 55653cbb3e analytical-viz: new 'Parallelism rules' tab 
New last tab distilling the parallelism-selection discussion into
tables. Framed up-front as heuristics + hard constraints, not a
production procedure — with public counterexamples (DeepSeek-V3 PP16
no TP; vLLM PP-on-non-NVLink) that break the naive 'start TP at 8'
rule.

Ten sections, all table-first:

  1. The core principle (chattiest comm on fastest links)
  2. Parallelism techniques × comm profile (DP/TP/PP/CP/EP with
     Shards, Pattern, Frequency, Per-link volume as degree ↑,
     Hard ceiling)
  3. Dominant problem → first axis to try (11 rows: weight fit,
     activation memory, single-seq KV, many-seq KV, MoE experts,
     TPOT, aggregate throughput, multi-node, expert layer, global
     batch cap)
  4. When to add / when to stop, per axis (DP/TP/PP/CP/EP/FSDP/SP)
  5. Common misconceptions vs reality (9 rows: KV-head hard
     ceiling, always start TP at 8, EP comm falling, attn-DP
     shards KV for free, DP > CP always, PP inside domain
     worthless, memory is only feasibility, FSDP as substrate,
     pipeline bubble = (P-1)/m)
  6. Inference decision procedure (7 ordered steps)
  7. Inference cases → recommended shape (A–E: fits on 1 GPU /
     within NVLink / multi-node / extreme context / many moderate)
  8. Training baseline recipe (6 ordered steps)
  9. Sanity checks for any candidate config (7 checks)
 10. One-paragraph rule + honest caveat block

All content is st.markdown + st.dataframe — no new pure functions.
Smoke test asserts the tab title + `with tab_prules:` block are both
present in app.py.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-30 09:11:08 -07:00
mukesh 5ce7a5b30b analytical-viz: capacity planning — symbol legend + plain-English formulas 
Section 1 rewritten so a layman can read it without a glossary:

1. **Symbol legend table** listing every letter that appears in the
   three-axis formulas (N, b, N·b, HBM_per_PE, S_kv, kv_bpt,
   users_per_replica, TPOT SLO, step_latency, B_at_SLO, n_users,
   N_replicas, ⌈ ⌉). Each row: Symbol / Plain English / Value now
   for the current sidebar model + chip. For example:
     - N = 'Total number of model parameters (attention weights +
       FFN weights, across ALL layers)' → 6.98 billion params for
       Llama 3 8B
     - kv_bpt = 'KV cache bytes per token per user, summed across all
       layers (= 2·H_kv·d_head·b·layers)' → 128 KB/token

2. **Three-axis table** widened from (Formula / Meaning / Grows with)
   to (In symbols / In plain English / What it says / Value now).
   Axis A now shows the actual substituted arithmetic:
   ⌈ 13.96 GB / 6.0 GB ⌉ = ⌈ 2.33 ⌉ = 3 PEs.
   B and C flag 'see calculator' because they depend on user inputs.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-30 08:35:59 -07:00
mukesh bc5e704572 analytical-viz: Physical Layout — TTFT & TPOT full-model latency section 
New section at the bottom of the Physical Layout tab that sums the
per-stage latencies into end-to-end metrics for the current sharding:

- **TPOT** = one decode-step time × N_layers  (per output token)
- **TTFT** = one prefill pass time × N_layers  (for a prompt of the
  input length, computed in prefill mode)
- **E2E** = TTFT + (N_out − 1) × TPOT
- **Throughput/replica** = B / TPOT  (tokens/sec)

Two number inputs: prompt length (default = sidebar S_kv) and target
N_out tokens (default 100). Four KPI cards + a breakdown table with
per-layer and all-layer µs/ms for attention/FFN in both modes.

Both TTFT and TPOT are computed by cloning the sidebar TopologyConfig
with mode swapped ('decode' for TPOT, 'prefill' for TTFT with T_q =
prompt length), then reusing all_stages / all_ffn_stages — no new
math. Caveat noted in caption: PP=1 assumption (no microbatch
overlap / bubble modeling).

Sanity: Llama 3 8B on CP=2/TP=8/default SIP → TPOT ≈ 4.2 ms/token,
TTFT ≈ 1.1 s for an 8k prompt. Cross-checkable on the Per-stage
latency tab.

Smoke test in test_stage_shapes asserts the section title appears
exactly once and sits inside tab_layout (before tab_memory).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 15:51:06 -07:00
mukesh c39bbb16aa analytical-viz: capacity planning tab — add public-example B/replica table 
Section 6 (new): 'Actually-provisioned batch sizes (public examples)'.
Ballpark B/replica figures for Gemini 1.5 Pro/Flash 1M, GPT-4.1 1M,
Claude 3.5 Sonnet 200k, and enterprise dedicated tiers, with the
inference basis for each ('pricing ~2x standard suggests low B',
etc). Caption flags that none are officially published; numbers are
backed out from public pricing + SLAs.

Binding-axis playbook renumbered #6 → #7.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 15:44:35 -07:00
mukesh 77abc95d78 analytical-viz: capacity planning tab — add rules of thumb + deployment templates 
Two more table sections on the Capacity planning tab:

**Section 4 — Practical rules of thumb (by context regime)**
- Short (S_kv < L*): pack HBM, B ~ 2·B*, high util, cheap tier
- Long (S_kv > L*): fewer users/replica, low B, util drops as
  1/(1 + S_kv/L*), pricier
- Extreme (S_kv >> L*): dedicated pool, heavy CP, disaggregated
  prefill, very low util without sparse attention

Columns: Regime, Batch strategy, Utilization, Cost/token, Deployment.

**Section 5 — Sample deployment templates**
Same base model, three different sharding recipes routed to by the
API gateway based on request context length:
- Config_small : CP=1,  TP=8, PP=1 → 8 GPUs, up to 32k, B=64
- Config_medium: CP=4,  TP=8, PP=1 → 32 GPUs, up to 128k, B=32
- Config_large : CP=32, TP=8, PP=1 → 256 GPUs, up to 1M, B=4

Columns: Tier, CP, TP, PP, Total GPUs/replica, Max context, Typical
B, Best for.

Playbook section renumbered to #6.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 15:13:08 -07:00
mukesh f54822d247 analytical-viz: move capacity planning to its own tab, table-first layout 
New 'Capacity planning' tab (last one, after 'Chip roofline & B*').
Everything below is a dataframe table where a bullet list could have
gone, matching the user's 'mostly try to make tables to describe the
information' preference:

1. **Three-axis GPU sizing formula** (table: A/B/C rows with Formula,
   Meaning, 'Grows with' columns)
2. **Live sizing calculator** (KPI cards + 3 number inputs — this is
   the only interactive block on the tab)
3. **SLO standards by workload** (table: TTFT / TPOT targets for
   voice/code/chat/agentic/batch, with 'what binds' column)
4. **Hybrid deployment layers** (table: Layer 1 elastic pool /
   Layer 2 length-tier routing / Layer 3 disaggregated prefill/decode
   with Serves / Key mechanism / Admission rule / Typical config)
5. **Binding-axis playbook** (table: for each of A/B/C — Symptom,
   First lever, Second lever, Cost impact)

Removed the sizing calculator + 3 expanders from the bottom of the
Chip Roofline tab. Roofline tab is back to pure chip-vs-model /
formulas territory; planning tab has the workload / SLO / deployment
material.

Calculator uses the base sidebar chip (_default_machine) rather than
the roofline tab's knob-scaled chip, so results are stable across
navigation.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 15:09:58 -07:00
mukesh b7c4c680aa analytical-viz: roofline tab — 'How hyperscalers pick GPU count' section 
New section at the bottom of the Chip Roofline tab implementing the
three-axis GPU count decision:

  N_GPUs = max(capacity_floor, KV_headroom, throughput_SLO) × N_replicas

Interactive inputs: concurrent users, avg context/user, TPOT SLO.
Four KPI cards report each axis + the recommended total, with a
'binds:' delta chip showing which axis is the bottleneck.

  - Axis A (capacity):  PEs to hold weights alone
  - Axis B (KV):        PEs to hold weights + all KV of one replica's users
  - Axis C (throughput): replicas needed at B_at_SLO users each
  - Total: max(A,B) × N_replicas

Contextual caption explains what to do about the binding axis. If SLO
is infeasible (even B=1 exceeds SLO), an error block explains the
options (loosen SLO, shorten context, or scale up FLOPs/BW knobs).

Plus three collapsed expanders explaining the hybrid deployment
pattern hyperscalers use:
  - Layer 1: one elastic pool + PagedAttention + continuous batching
  - Layer 2: length-tier routing (standard vs long-context)
  - Layer 3: disaggregated prefill/decode (DistServe, Splitwise)

Pure additions in chip_roofline.py:
- max_batch_within_slo(machine, model, s_kv, slo_s) — analytical
  inverse of step_latency to find the largest per-replica B under SLO
- size_deployment(machine, model, n_users, avg_ctx, slo_s) — returns
  GpuSizingResult with all three axes + binding info

6 new tests cover: capacity scales with model size, KV grows with
users/context, binding axis flips with workload shape, tight SLO
shrinks max batch, weight-time-exceeds-SLO returns 0, total ==
per-replica × replicas.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 15:07:20 -07:00
mukesh 6410c0843c analytical-viz: roofline tab — fix NameError from knob-rename refactor 
Previous commit renamed slider vars (_rf_flops_mult -> _flops_mult_now,
etc.) and split _ai/_b_star into base vs scaled, but the AI-sensitivity
plot's vertical markers + caption still referenced the old names, so
loading the tab crashed with:

    NameError: name '_rf_flops_mult' is not defined

Fixed all four call sites:
- _axAI1: 'Base AI' reference line now uses _ai_base (was _ai, which
  is now the scaled value).
- _axAI2: 'Base B*' reference line uses _b_star_base; vertical
  markers use _flops_mult_now / _bw_mult_now.
- Caption 'Your knob position' uses the new names.

Grep confirms no _rf_flops_mult / _rf_bw_mult / _ai_scaled /
_b_star_scaled references remain.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 14:51:01 -07:00
mukesh 97e765b58f analytical-viz: roofline knobs — absolute TFLOPS / GB/s, drive whole tab 
Two changes:

1. **Knob units and ranges**
   - FLOPs slider: absolute TFLOPS, 1..32 (step 0.5). Was 0.25..8× multiplier.
   - HBM BW slider: absolute GB/s, 128..1024 (step 16). Was 0.25..8× multiplier.
   Defaults still initialize to the sidebar Hardware panel's base chip.
   Caption also shows the implied × multipliers of the base chip.

2. **Knobs now drive every roofline number on the tab.**
   Previously FLOPs/BW only affected the AI-sensitivity plot. Now the
   scaled machine feeds:
     - AI / B* / L* KPI cards
     - Regime KPI
     - Good-B / Good-L cards + utilization
     - Headline plots (step latency + cost/token)
     - Regime formulas table (t_com, t_mem_short, t_mem_long values)
     - Formula table (C, W, knee substitution)
   Left unchanged (intentionally):
     - PE memory budget stacks — depend on HBM capacity, not BW/FLOPs.
     - AI-sensitivity plot itself — sweeps around the SIDEBAR base
       chip so the curves are comparable across knob positions.

Header caption now shows both effective chip (from knobs) and base
chip (from sidebar) side by side.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 14:47:30 -07:00
mukesh bda76cbd66 analytical-viz: roofline tab — drop formula from step-latency header 
Revert the multi-line formula annotation on the 'Latency per decode
step' header; make it a single-line markdown title matching the
'Cost per token' header on the right. Both headline plots now render
at the same height so the plot area lines up.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 14:14:24 -07:00
mukesh a1c3c28f62 analytical-viz: roofline tab — PE memory budget, AI sensitivity, FLOPs/BW knobs 
Three additions to the Chip Roofline tab:

1. **PE memory budget** section (two side-by-side stacked-area plots)
   - Left: per-PE memory vs S_kv (128..1M, log). Stacks: weights,
     KV, transient. HBM budget line + current-S_kv marker.
   - Right: per-PE memory vs B (1..256). Same stacks, current-B
     marker. KV grows linearly with either axis; weights are
     invariant of B and S_kv (fixed by CP·TP·PP sharding).
   Uses actual sharded per-PE quantities from memory_layout so the
   numbers match what the deployment would really allocate.

2. **AI sensitivity** section (two side-by-side line plots)
   - Left: AI vs chip-parameter multiplier. Two curves: scale FLOPs
     (AI grows linearly), scale HBM BW (AI shrinks inversely).
   - Right: B* vs multiplier. Same shape as AI (for BF16 where B*=AI).
   Vertical markers show current FLOPs/BW knob positions.

3. **Two new knobs** in the top row: FLOPs × multiplier and HBM BW ×
   multiplier (0.25..8, step 0.25). The AI-sensitivity plot markers
   move with them; the header caption shows scaled AI + B* for the
   selected point.

Also: formula annotation on the headline 'Latency per decode step'
plot header showing t_step = N·b/W + 2·N·B/C + B·S_kv·kv_bpt/W.

Pure additions in chip_roofline.py (all testable):
- MemoryBudgetPoint dataclass
- memory_budget_curve_vs_skv, memory_budget_curve_vs_batch
- AISensitivityPoint dataclass
- ai_sensitivity_curve(machine, model, mults, axis='flops'|'bw')

9 new tests cover linear scaling of KV with S_kv/B, weight
invariance across sweeps, over-budget flag flip, free_gb clamp,
FLOPs multiplier → AI linear, BW multiplier → AI inverse, B*
tracks AI for BF16, bad axis raises.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 14:10:02 -07:00
mukesh 9b4aee83da analytical-viz: roofline tab — trim to 4 plots, cap B at 256 / S_kv at 1M 
- Slider caps: B_max = 256, S_kv_max = 1M (matches practical decode
  operating ranges).
- Batch axis on every plot capped at 256 to match the slider.
- Removed the two redundant plots (old Plot 1 and old Plot 4). Both
  were re-drawing the same weight/compute/KV decomposition as the new
  headline 'Cost per token' plot at the top — the only distinguishing
  bits (extra 'Memory total' curve, B* marker) have been folded into
  the headline plot already, so keeping them was pure clutter.
- Every remaining plot now reacts to the S_kv / B knobs:
    * Plot A (step latency): S_kv drives curves, B marker drawn.
    * Plot B (cost/token): same.
    * Plot 2 (no-knee overlay): S_kv drives the L* ratios shown,
      B marker drawn.
    * Plot 3 (knee vs context): S_kv AND B markers drawn on both axes.

Net: 4 plots (headline pair + no-knee + knee-vs-context), all live-
interactive against both sliders.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 13:52:40 -07:00
mukesh 221097ef08 analytical-viz: roofline tab — headline step-latency + cost-per-token plots, B knob 
Two new plots at the top of the Chip Roofline tab (right after the
knob row), side by side:

1. **Latency per decode step** — one forward pass's total time as B
   grows. weight_fetch (flat in B), compute (linear), KV (linear),
   total. The SLO / TTFT view — bigger B = longer step.

2. **Cost per token** = step latency ÷ B. weight (shrinks 1/B),
   compute (flat), KV (flat), total. The efficiency / cost view —
   bigger B (up to ~2·B*) = cheaper per token.

Same underlying decomposition, two divisors. Puts the classic
throughput ↔ latency tradeoff in one glance.

Also added a **B (batch size) slider** next to the existing S_kv
slider. Both plots draw a vertical marker at the current B; the
regime KPI ("memory-bound / compute-bound / kv-bound") now uses the
slider B, so you can sweep B live and watch the label flip when you
cross B*.

step_latency_curve + StepLatencyPoint added to chip_roofline; 5 new
tests cover: weight_s batch-invariant, compute_s/kv_s linear in B,
step_total == per_token_total × B (definitional), and step ==
per_token at B=1.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 13:40:37 -07:00
mukesh c99a238826 analytical-viz: roofline tab — S_kv knob, regime formulas, good B/L, decomp plot, split table 
Bundle of roofline-tab enhancements requested in-thread:

1. **S_kv slider at top of tab** — override the sidebar's S_kv locally
   so all four plots update as you drag it. Handy for exploring how
   the KV wall arrives without disturbing the rest of the app config.

2. **Good B / Good L KPI cards** — 3 metrics under the main KPI row:
     - Good B = 2·B* (Pope's rule of thumb)
     - Good L ceiling = L* (compute-friendly context ceiling)
     - Utilization @ current S_kv = 1 / (1 + S_kv/L*)

3. **Plot 4 — latency-decomposition per decode step** — five curves
   on one axis:
     - Compute (dashed, flat)
     - Weight fetch (triangles, shrinks 1/B)
     - KV fetch (squares, flat — batching doesn't help)
     - Memory total (weights + KV, purple)
     - Total (compute + memory, black bold)
   Makes "which term dominates at this B?" visible at a glance.

4. **Regime formulas table** — one row per cost term (t_com, weight
   fetch, KV fetch, bottleneck) × two columns (short-context vs
   long-context regime), plus a 'value now' column using the
   current S_kv slider.

5. **How to pick B and S_kv — one-paragraph guidance** with the
   numeric recommendations plugged in.

6. **Split formula table** — was cramming symbolic + substituted
   form into one cell with newlines; now has four proper columns:
   Symbol / Formula / With numbers / Meaning / Value.

New pure functions in chip_roofline.py: t_mem_short, t_mem_long,
t_com, good_batch, good_context, utilization_at. All roofline math
stays in the pure module; app.py just calls them and formats.

10 new tests bring the roofline test count to 27 (all 67 tests still
green): t_mem_long(L*) == t_com, doubling B halves t_mem_short,
good_batch scales with sparsity, utilization at 0/L*/2L* returns
1.0/0.5/0.333, monotonic decrease with context.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 13:27:17 -07:00
mukesh 04c7d247f0 analytical-viz: roofline formula table shows substituted numeric form 
Each row in the 'Formulas & interpretation' table now shows both the
symbolic form AND the actual numbers plugged in, on a second line in
the Formula column. Row height bumped to fit two lines. Rendering:

  AI:      C / W
           = 8.00e+12 / 2.56e+11
           = 31.25 FLOPs/byte

  L*:      2 · N_active · W / (C · kv_bpt)
           = 2 · 6.98e+09 · 2.56e+11 / (8.00e+12 · 131,072)
           = 3,407 tokens

  B_knee:  B* / (1 - S_kv/L*)
           = 31 / (1 - 8,192/3,407)
           = 31 / (1 - 2.404)   [<= 0 -> no knee]

Makes the derivation traceable at a glance without needing to plug the
numbers back into the abstract formula. Purely a display change; no
runtime behavior shifts, all 57 tests still pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 12:04:47 -07:00
mukesh 1ddd1baa7a analytical-viz: 'Chip roofline & B*' tab — AI, B*, L*, cost curves 
New tab surfaces the arithmetic-intensity story from LLM-serving
practice for the current sidebar chip + model. All derived from
MachineParams + ModelConfig; no new configuration.

chip_roofline module (pure functions):
- arithmetic_intensity      = C / W   (FLOPs per byte HBM BW)
- critical_batch (B*)       = C * b / (2 * W) * sparsity
- balance_context (L*)      = 2 * N_active * W / (C * kv_bpt)
- knee_batch (B_knee)       = B* / (1 - S_kv/L*)  (None past L*)
- per_token_latency_curve   = weight/B + compute + KV, per B
- bound_regime              = which term dominates now

Peak roofline convention (no compute_util factor) so weight_s ==
compute_s exactly at B*. Comm + TP/CP sharding intentionally
excluded — stage_latencies is the full latency model; this is the
back-of-envelope chip-vs-model view.

Tab shows:
- 4 KPI cards: AI, B*, L*, current regime at (B, S_kv)
- MoE hint when preset flags MoE ('300 x sparsity' rule)
- Plot 1: cost vs B at current S_kv (weight, compute floor, KV
  floor, total) with B* marker line
- Plot 2: cost curves at S_kv/L* = 0.25, 0.5, 1, 2, 5 — shows the
  no-knee regime past L*
- Plot 3: B_knee vs S_kv — knee slides right, diverges at L*
- Formulas + interpretation table

test_chip_roofline covers B* against H100 reference (295), sparsity
scaling, L* scaling with HBM BW, knee divergence at L*, per-token
curve monotonicity + asymptote to compute+KV floor, regime
classification. Smoke test guards the tab is wired in app.py.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 11:59:19 -07:00
mukesh 5e92b89821 analytical-viz: PE view keeps CP rank locator visible at TP>=8 
The Per-PE layout diagram used to color-code + label CP ranks only
when cp_placement=='pe' (CP packed intra-cube). At TP>=8 the TP dim
fills the cube, cp_placement gets pushed to 'cube', the intra-cube CP
branch is skipped, and the CP rank tag silently disappeared from
cubes, per-PE headers, and the figure title. Users at TP>=8 had no
way to tell which of the CP groups the diagram was drawing.

Fix: when cp_placement=='cube' and CP>1, add
  - cube header suffix:   '[CP rank 0 of {CP}]'
  - per-PE header suffix: '  |  CP=0'
  - figure title:         'CP={CP} (showing 1 of CP groups; others identical)'

Behavior at cp_placement=='pe' is unchanged — the existing per-CP-rank
palette + 'CP={rank}' label still fire. CP=1 adds no locators at all.

test_pe_weight_layout covers four cases: TP=8/CP=4 (cube placement,
locator appears), TP=2/CP=4 forced to pe placement (regression:
per-rank labels still there, cube tag not added), CP=1 (no locators),
TP=16/CP=2 (both spilled cubes carry the CP-rank-0 tag).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 10:19:39 -07:00
mukesh 3360be2488 analytical-viz: consolidate shape tables under Physical Layout tab 
Move the per-stage shape tables (attention + FFN) off the Per-stage
latency tab and combine them with the weight + KV shape tables into a
single 'All tensor shapes (per PE)' expander at the bottom of the
Physical Layout tab. Five subsections in one place:
  - Attention weights (global + per-PE shape, per layer)
  - FFN weights (same schema)
  - KV cache (same schema)
  - Per-stage attention shapes (input / weight / output)
  - Per-stage FFN shapes (same)

Rationale: the shape catalog belongs next to the physical-layout view
that shows how the tensors are placed on hardware. The Per-stage
latency tab now focuses purely on time-per-stage. Weight/KV tables
stay on the Memory Breakdown tab too because the summary metrics
underneath still consume the same rows.

Guard test verifies the section string appears exactly once, is
positioned inside tab_layout (before tab_memory / tab_stages), and
that neither attn_stage_shape_rows nor ffn_stage_shape_rows is called
inside the tab_stages block anymore.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 10:14:08 -07:00
mukesh adc14c84af analytical-viz: per-stage tensor shape tables (attention + FFN) 
New stage_shapes module reports per-PE input / weight / output tensor
shapes for every attention (S1..S10, C1..C3) and FFN (F1..F5, CF1)
stage. Streamlit's Stages tab picks these up below the existing
latency table + bar chart as two separate dataframes ('Attention
stages' and 'FFN stages'), so anyone who wants to see the exact
sharded tensor dims per stage can, without cluttering the latency
view.

Rows mirror the conditional structure of all_stages / all_ffn_stages:
S8 merge only when CP>1; C1 CP ring only in prefill with CP>1; C3
Score AR only when kv_shard_mode='split' and TP>H_kv; C2 TP AR only
when TP>1; CF1 FFN AR only when the FFN divisor > 1. Batch B flows
into every activation dim; weights stay B-invariant.

test_stage_shapes covers column presence, conditional row insertion
(decode omits C1; CP=1 omits S8; prefill+CP>1 inserts C1), batch
scaling of the leading dim, and that S5 output references both T_q
and S_local for the current cfg.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-29 10:05:02 -07:00
mukesh 9c5062bdfb analytical-viz: memory-min buttons reset stale placement/mode session keys 
auto_suggest only explores (cp, tp, pp, cp_placement); every other
TopologyConfig field is left at its dataclass default when the score
is computed. The sidebar's memory-min apply used to update only
cp/tp/pp/dp/cp_placement, so any stale session_state value from a
previous user interaction (most damaging: tp_placement="cube") would
carry over into the rerun and produce a topology the memory-min search
never scored — visibly, cubes_used could balloon (e.g. Llama 3 70B
Attn scope: auto_suggest picks 1 cube; stale tp_placement="cube" made
the applied config span 8 cubes).

Fix: after picking a Suggestion, reset tp_placement, kv_mode,
cp_ring_variant, ep back to their TopologyConfig defaults, and drop
the ffn_scope_label key (dynamic string; falls back to the TP+CP
default index on re-render). Now the sliders + placement + mode state
after apply exactly reproduce the topology auto_suggest scored.

test_auto_suggest_cubes_match_default_topology asserts, across four
model/skv combinations and all three scopes, that a TopologyConfig
built from (cp, tp, pp, cp_placement) plus dataclass defaults has the
same cubes_used the Suggestion reports.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 23:45:38 -07:00
mukesh 4ded67a443 analytical-viz: sidebar memory-min buttons are scope-aware (Attn / FFN / Attn+FFN) 
Replace the three memory-optimal Pareto-search buttons in the Parallelism
sidebar with three scope-aware memory-min autosuggest buttons. Each
button sizes for a different per-PE memory footprint:

  - Attn      -> attention weights + KV cache (no FFN weights)
  - FFN/MoE   -> FFN weights only (no attention, no KV)
  - Attn+FFN  -> everything (traditional autosuggest)

memory_layout: per_pe_weight_bytes and compute_memory take
include_attention / include_ffn flags; KV cache is zeroed when attention
scope is excluded.

autosuggest: _score_candidate and auto_suggest forward the flags into
compute_memory, so the smallest-fit search now respects the chosen
scope.

app.py: single "Apply memory-min" caption + 3 column buttons. Each
button runs auto_suggest with its scope filter, snaps the sliders to
the picked (CP, TP, PP), and sets the Physical Layout scope filter to
match. Removes the standalone Apply memory-min button (was duplicating
the Attn+FFN case) and the Pareto-search import.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 23:34:05 -07:00
mukesh 47c4f40f4d analytical-viz: batch B in every formula string + rename 'latency-optimal' → 'memory-optimal' 
Two related fixes.

1) Formula strings show B alongside every other symbol. Previously the
   numeric FLOPs/mem_bytes/comm_bytes were scaled by B, but the printed
   formula strings still read like B=1 amounts. That made the per-stage
   table confusing at high B: numbers moved but formulas didn't. Updated
   every attention + FFN stage's `formula` / `flops_formula` /
   `mem_formula` / `comm_formula` to include B explicitly. Weight-bytes
   lines now say "(shared across batch)" / "(weight, B-invariant)" so it's
   obvious what does and doesn't scale.

   Stages touched: S1..S10 (attention), C1/C2/C3 (attn comm), F1..F5
   (FFN), CF1 (FFN AR).

2) Renamed the button/spinner/help copy from "latency-optimal" to
   "memory-optimal" to match the smallest-fit pick semantics I put in
   place earlier. The buttons already picked smallest-fit; the labels
   just still said "latency-optimal" from the previous iteration.

   Left the Auto Hardware sensitivity chart's "latency-optimal baseline"
   label alone — that panel is a HW co-design view where "how fast can
   each HW go?" is the intended question.

Verified: 24 pytest tests still pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 23:20:32 -07:00
mukesh 50aab8d591 analytical-viz: complete batch (B) scaling in remaining stages + comm 
Finishes the batch scaffold. All stages that produce per-request work
now scale flops / activation-memory / comm-bytes by B; weight bytes stay
fixed (shared across the batch).

Attention block, previously partial:
  - S6 stage_softmax: elems + bytes × B
  - S8 stage_merge: flops × B; O/m/l AR bytes M × B
  - S9 stage_normalize: flops × B
  - S10 stage_wo: flops × B
  - C1 comm_cp_ring: decode O/m/l AR M × B; prefill K/V ring + Q+O/m/l
    variants both × B
  - C2 comm_tp_allreduce (W_O output): bytes × B
  - C3 comm_kv_split_allreduce (head-split scores): bytes_per_hop × B

FFN block, previously untouched:
  - F1 stage_ffn_rmsnorm: activation × B (weight fixed)
  - F2 stage_ffn_gate (via _ffn_gemm): flops × B
  - F3 stage_ffn_up: same
  - F4 stage_ffn_swiglu: elems + flops × B
  - F5 stage_ffn_down: flops × B
  - CF1 comm_ffn_allreduce (batched FFN output): bytes × B

Verified with a smoke check on Qwen 3 8B / 128K decode:
  B=1  per-layer visible latency:  363 us
  B=8                                716 us  (sub-linear — many stages
                                              stay weight-bound)
  B=64                              4023 us  (approaching linear scaling
                                              as batch dominates)

All 24 pytest tests still pass at default B=1 (backward compat).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 23:10:39 -07:00
mukesh fc1dcdde24 analytical-viz: tensor sharding — show batch B on KV cache 
Three additions to the tensor sharding view when batch B > 1:
  1. Shape label appends "B={B}": e.g. (128, S_kv=131072, B=8).
  2. Note under the KV cache: "batch B={B} => {B}x KV bytes per PE".
  3. Visual "stack" — up to 5 dashed offset rectangles drawn behind
     the KV cache to hint at the batch stacking dimension. Capped at 5
     so a big B doesn't overwhelm the diagram.
  4. Title also gets B={B} between CP and FFN scope.

Attention/FFN weight tensors are NOT stacked — weights are shared
across the batch (correct: only activations + KV scale with B).

At B=1, all four additions are no-ops so the diagram looks unchanged
from before.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 22:56:48 -07:00
mukesh fd45bd2408 analytical-viz: add sidebar Batch B selectbox + wire through all suggesters 
New sidebar widget in the Workload expander:
  Batch B (concurrent requests) — selectbox 1/2/4/…/256, default 1.

Wired end-to-end:
  - app.py sidebar → forwards b_batch when constructing the main topo
  - auto_suggest(..., b=1) — passes b to _score_candidate, which builds
    TopologyConfig with b=b
  - run_auto_explore(..., b=1) — sets topo.b = b for every enumerated
    candidate before scoring
  - joint_explore(..., b=1) — forwards b to _best_parallelism_for_hw
    and _best_parallelism_two_stage; both set topo.b = b before scoring

All button handlers (sidebar Apply-scope, Physical Layout scope,
Auto Suggest tab sweep, Auto Hardware tab sweep) now pass b_batch.

Combined with the earlier partial-scaffold commits (memory_layout
scales KV by B; stage_rmsnorm / stage_wq / stage_wkv / stage_kv_append /
_per_hop_qkT_pv scale flops + activation memory by B), changing B in
the sidebar now affects reported latency and per-PE memory footprint
in the visible parts of the pipeline. The remaining FFN + comm-AR
stages still ignore B (they'll be next); their contribution is small
for the memory-bound decode case that matters most, but latency for
FFN-heavy configs at high B will be slightly under-reported until
those are scaled too.

Verified: 24 pytest tests pass at default B=1 (backward compat).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 22:51:56 -07:00
mukesh 674e194dc9 analytical-viz: hot-reload guard — include model_config + model_presets, ordered 
The reload guard didn't include model_config.py. When the batch (B)
field was added to TopologyConfig, Streamlit kept the stale
pre-B dataclass in sys.modules and every downstream module hit
"AttributeError: 'TopologyConfig' object has no attribute 'b'".

Fix: add model_config and model_presets to the reload list, and order
model_config FIRST so downstream modules that reload after it pick up
the new dataclass definition. Also reordered the rest so upstream
dependencies (autosuggest, memory_layout, stage_latencies) reload
before their consumers (auto_explore, auto_hardware).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 22:43:47 -07:00
mukesh 770aed6291 analytical-viz: Auto Suggest tab — add Top-10 by memory footprint table 
Adds a second table beneath the Pareto configurations table, sorted by
(hbm_utilization ↑, pes_used ↑, latency ↑) — top 10 smallest-memory
Pareto configs first.

Columns highlight the memory story:
  HBM % | weights (GB) | KV (GB) | PEs | SIPs | lat (ms) | CP TP PP DP | kv | ffn | tp_place | cp_place

Useful when memory pressure is the real binding constraint of a
deployment — e.g., picking a config for a per-PE HBM-limited SIP,
or spotting configs that trade small latency headroom for large
memory savings.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 22:24:38 -07:00
mukesh 1bce080d79 analytical-viz: buttons pick smallest-fit Pareto point, not fastest 
Sidebar Attention/FFN/Attn+FFN buttons and Physical Layout tab buttons
now select the SMALLEST-FIT config from the scope's Pareto set — key
changed from (latency ↑, pes_used ↑, hbm_utilization ↑) to (pes_used ↑,
hbm_utilization ↑, latency ↑). Answers "smallest deployment that's
still Pareto-optimal for this scope" instead of "fastest deployment
achievable".

Left unchanged:
  - "Best latency" metric displays in Auto Suggest / Auto Hardware tabs
    still show the fastest number (informational; the metric labels say
    "Best latency").
  - auto_hardware._best_parallelism_for_hw stays latency-min: it's
    inside HW co-design, where "how fast can each HW candidate go" is
    the primary question.

Verified: 24 pytest tests pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 22:23:34 -07:00
mukesh 5d7ef48b14 analytical-viz: partial batch (B) scaffold — TopologyConfig field + KV / early stages 
WIP toward batch-size support. This first commit is behavior-preserving
at the default B=1 (every new multiplication is by max(1, cfg.topo.b) =
1 today) so all existing tests pass. Follow-up commits will:
  - scale the remaining stages (S6/S7/S8/S9/S10 + C1/C2/C3 + FFN + FFN AR)
  - add a batch selectbox to the sidebar
  - forward b through auto_suggest / auto_explore / auto_hardware

Changes so far:
  - TopologyConfig: new b: int = 1 field (batch size).
  - memory_layout.per_pe_kv_cache_bytes: * B (each concurrent request
    keeps its own KV cache slice).
  - stage_rmsnorm / stage_wq / stage_wkv / stage_kv_append /
    _per_hop_qkT_pv: FLOPs and activation memory scaled by B; weight
    bytes stay fixed (weights shared across the batch).

Verified: 24 pytest tests still pass at default B=1.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 22:21:37 -07:00
mukesh 2a0bc26db0 analytical-viz: compound tiebreaker — fewer PEs / lower HBM% when latency ties 
Every place that picks a single "best" config from the Pareto set now
uses a compound sort key: (latency_ns ↑, pes_used ↑, hbm_utilization ↑).
Primary = latency; when configs tie on latency (common — cp_ring
variants, some placement variants often produce identical numbers),
prefer smaller-footprint picks.

Places updated:
  - app.py:  sidebar Apply Attn/FFN/Attn+FFN button
  - app.py:  Physical Layout tab Attn/FFN/Attn+FFN button
  - app.py:  Auto Suggest tab "Best latency" metric
  - app.py:  Auto Hardware tab "Best latency" metric (uses parallelism.pes_used
             + parallelism.hbm_utilization since JointScore wraps ConfigScore)
  - auto_hardware.py:  _best_parallelism_for_hw iteration key

No behavior change when there's a strict latency winner. When there are
ties, the picked config uses fewer PEs and lower HBM utilization.

Verified: all 24 pytest tests pass (default include_attention=True and
include_ffn=True paths unchanged).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 22:16:02 -07:00
mukesh 6f61657ba4 analytical-viz: topology map — color PEs by CP rank when cp_placement=pe 
The topology map previously colored a whole cube (border + all PE
fills) by the single cp_rank assigned to that cube in the cube→(pp,cp)
mapping. Under cp_placement=cube this is right (each cube = one CP
rank). Under cp_placement=pe, however, multiple CP ranks are PACKED
into the same cube's PEs, so the whole-cube coloring makes every PE
look identical (defaulting to cp_rank=0's palette entry — bright red).

Fix: in the per-PE loop, if cp_packed = (cp_placement=="pe" and cp>1),
compute each PE's own cp_rank = pe_id // tp and look up its color via
_cp_color(pp_stage, pe_cp_rank, cp_size). Border still uses the
cube-level color (cp_rank=0), so the outer bounding box is unchanged,
but the interior PEs now show all four (or however many) CP-rank
colors visibly.

For cp_placement=cube: unchanged (single pe_fill per cube).

Verified with a headless render at Qwen 3 8B, CP=4, TP=2,
cp_placement=pe: 145 patches → 8 distinct facecolors (4 for the CP
ranks in the packed cube + inactive/border tints), where before it
was fewer distinct colors.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 22:07:50 -07:00
mukesh da1909fdf2 analytical-viz: extend hot-reload guard to all viz sub-modules 
Previously only reloaded auto_explore / auto_hardware / autosuggest /
stage_latencies / memory_layout. When we edit pe_weight_layout,
tensor_sharding, topology_map, pipeline_diagram, or optimization_report,
Streamlit was still holding the pre-edit versions until a full restart.

Add all five to the reload list so any visualization-module edit lands
on the next Streamlit rerun without needing Ctrl+C.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 22:03:42 -07:00
mukesh ed28eea156 analytical-viz: color-code CP groups in PE layout when cp_placement=pe 
When cp_placement=pe packs multiple CP ranks intra-cube, the PE-level
layout now:

  1. Colors each PE's background by its CP rank (8-color palette wraps
     for cp > 8). Cp_placement=cube keeps the historical light-blue
     styling (only one CP rank per cube tile).

  2. Adds "TP=N  |  CP=M" to the per-PE header so users can read the
     (cp_rank, tp_rank) pair without inferring from position.

  3. Fixes a head-assignment bug: _q_heads_for_pe / _kv_heads_for_pe /
     _per_pe_bytes used to be called with the raw PE-in-group index,
     which treated every PE as a distinct TP rank. Under cp_placement=pe
     this gave wrong Q/KV head lists for PEs beyond the first TP group
     (PE 2..7 with tp=2, cp=4 would ask for head slots 32..127 in a
     32-head model). Now called with tp_rank = pe_id % tp, so all CP
     ranks in the cube share the correct head split.

Verified via a headless matplotlib smoke test with CP=4/TP=2 on Qwen 3
8B: 8 PE patches + 1 cube patch → 5 distinct facecolors (cube + 4 CP
ranks), no errors.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 21:59:18 -07:00
mukesh f1cf0257d3 analytical-viz: restore Apply memory-min button 
An "Apply memory-min" button was removed when the 3 latency-optimal
buttons were added. Consequence: after clicking any latency-optimal
button, the sidebar sliders drift from the memory-min caption and
there was no way to sync them back without restarting the app or
manually adjusting each slider.

Restore the button, positioned below the caption and above the 3
latency-optimal buttons. Clicking it snaps cp/tp/pp/dp/cp_placement
to what the caption shows and clears _pl_active_scope to "full" so
the Physical Layout tab stops filtering.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 21:55:29 -07:00
mukesh afe35ee0b5 Revert "analytical-viz: default intra-cube (PE↔PE) BW to 128 GB/s" 
This reverts commit 84abf847c4.
 2026-07-28 21:50:17 -07:00
mukesh 84abf847c4 analytical-viz: default intra-cube (PE↔PE) BW to 128 GB/s 
Was 512 GB/s in the analytical viz MachineParams default. Bringing the
default down to 128 GB/s to match what the modelled physical link now
represents (parity with inter-cube D2D and topology.yaml).

Files touched (all three needed to keep defaults coherent):
  - model_config.py: MachineParams.bw_intra_gbs = 128.0 (was 512.0)
  - app.py: sidebar selectbox default index = 0 (128 GB/s) instead of 2 (512)
  - auto_hardware.py: _HW_KNOB_DEFAULTS + BALANCED + COARSE bw_intra_gbs
    baselines shifted so cost_score at defaults remains 6.0 and the
    sensitivity sweep starts from the new baseline.

Verified: 24 pytest tests pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 21:44:21 -07:00
mukesh a7f39ade7e analytical-viz: autosuggest prefers fewer cubes, tries cp_placement=pe 
Two changes to the memory-only auto_suggest so a "fits" that packs into
one cube is preferred over a "fits" that spreads across multiple cubes.

Suggestion dataclass gains cubes_used and cp_placement fields.

_score_candidate: for each (CP,TP,PP) triple, try cp_placement="cube"
(historical default: CP spans cubes) and, when CP·TP fits in one cube's
PE count, also cp_placement="pe" (pack CP into intra-cube PEs). Keep
the placement with fewer cubes; break ties toward "cube".

auto_suggest: switch sort key from (pes_used ↑, pp ↑, tp ↑, cp ↑) to
(cubes_used ↑, pes_used ↑, pp ↑, tp ↑, cp ↑). Fewer cubes wins first
because a cube is the physical die-level unit; PE count is the
tiebreaker.

Sidebar caption now also displays cubes_used + cp_placement so the
user can see the packed layout at a glance. Preset-change auto-reset
also applies the picked cp_placement.

Observed effect (verified):
  Qwen 3 8B / 128K decode:  CP=4 TP=2, "pe" → 1 cube, 8 PEs
                             (was 4 cubes with default "cube")
  Llama 3.1 70B / 128K:      CP=4 TP=16, "cube" → 8 cubes (unchanged;
                             TP=16 > 8 PEs/cube, can't pack)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 16:04:17 -07:00
mukesh c2b42999d8 analytical-viz: reorder tabs — Physical layout back to first 
New tab order:
  Physical layout / Auto Suggest Parallelism / Memory breakdown /
  Per-stage latency / Save & compare / Auto Hardware

Physical layout is what a user typically wants to see first when
loading the app, so it takes the leftmost slot again. Auto Suggest
Parallelism moves to second — still prominent, still usable as a
first step, but doesn't push the layout view behind it.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 14:26:07 -07:00
mukesh 9afeb6bc16 analytical-viz: 3 sidebar apply-buttons (Attn / FFN / Attn+FFN) 
Replaces the single memory-only "Apply auto-suggest" button in the
sidebar Parallelism expander with three latency-optimal buttons:
"Attention", "FFN/MoE", "Attn+FFN".

Each clicked button runs run_auto_explore for its scope, picks the
latency-minimum Pareto config, snaps the parallelism knobs to the
sidebar's selectbox option sets (CP/TP/PP/DP), and loads the other
knobs directly (tp_placement, cp_placement, cp_ring_variant, kv_mode,
ffn_scope_label). Also sets _pl_active_scope so the Physical Layout
tab's stage-table filter follows the scope automatically.

The caption above the buttons still shows the memory-only autosuggest
values as a reference — separately labeled "(memory-min)" to avoid
confusion with the latency-optimal buttons below.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 14:24:06 -07:00
mukesh bccc90db82 analytical-viz: Physical Layout scope filter for stage tables 
Clicking one of the three Physical Layout tab buttons now persists the
chosen scope in session_state["_pl_active_scope"] and filters the
per-stage latency tables accordingly:

  - attn scope → show only "Attention" stage table
  - ffn scope  → show only "FFN" stage table
  - full scope → show both (default; also matches a fresh sidebar-driven
                  config with no button click yet)

Added a "Layout scope: {label}" header so the user can tell at a glance
which scope's config is loaded.

The pipeline diagram + topology map + weight/tensor sharding + PE
layout continue to show the full model layout — those diagrams give
context that stays useful even when the user is focusing on attn or
FFN. Deeper filtering (e.g., attention-only pipeline stripe) can come
later if needed.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 14:21:27 -07:00
mukesh 1a6d52c58a analytical-viz: force-reload our modules on Streamlit rerun 
Streamlit's hot-reload re-executes app.py on every interaction, but
sub-modules imported from app.py stay cached in sys.modules across
reruns. When we edit auto_explore.py or auto_hardware.py while
Streamlit is alive, the app keeps using the pre-edit version until a
full Ctrl+C + restart — leading to confusing "unexpected keyword
argument" errors after a signature change.

Force importlib.reload() on our own modules at the top of app.py so
future signature changes land without a full restart. Only reloads if
the module is already in sys.modules (first run just imports normally).

Applied to:
  - tests.analytical_visualization.auto_explore
  - tests.analytical_visualization.auto_hardware
  - tests.analytical_visualization.autosuggest
  - tests.analytical_visualization.stage_latencies
  - tests.analytical_visualization.memory_layout

Third-party modules (streamlit, matplotlib, pandas, numpy) NOT reloaded
— unnecessary and slow.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 14:17:13 -07:00
mukesh 09721040ca analytical-viz: 3 sweep buttons at top of Physical Layout tab 
Shortcut: click any of {Attention, FFN/MoE, Attn+FFN/MoE} → runs
run_auto_explore for that scope, picks the latency-minimum Pareto config,
loads it into the sidebar's session_state (cp, tp, pp, dp, tp_placement,
cp_placement, cp_ring_variant, kv_mode, ffn_scope_label), then st.rerun().

Because the Physical Layout tab reads model + machine + cfg from the
sidebar, the sidebar update automatically redraws the pipeline diagram,
per-stage table, and everything below. No preview / parallel-display
state — WYSIWYG.

Ffn_scope_label reconstruction handles the dynamic "(div=…)" suffix the
sidebar shows (same pattern as auto_explore + auto_hardware tabs' "Load
into sidebar" widgets).

Verified: app.py parses; 24 pytest tests still pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 14:11:30 -07:00
mukesh bf5b659a3d analytical-viz: add FFN-only scope + 3rd sweep button 
Both auto tabs now offer three sweep scopes via three buttons instead
of two:
  - Run sweep — Attention          → include_attention=T, include_ffn=F
  - Run sweep — FFN/MoE            → include_attention=F, include_ffn=T
  - Run sweep — Attn + FFN/MoE     → include_attention=T, include_ffn=T

Each button caches its result under its own session_state key; the most
recently clicked button drives the display. All three caches persist so
users can flip between scopes without re-running.

Core changes:
  auto_explore.py + auto_hardware.py:
    - New include_attention: bool = True param alongside include_ffn
    - _sum_visible_latency, _efficiency, score_config, run_auto_explore,
      compute_parallelism_sensitivity, joint_explore, compute_sensitivity,
      _best_parallelism_for_hw, _best_parallelism_two_stage all wired.
    - Attention and FFN are additive with no overlap: measured 7.35 ms
      (attn) + 5.50 ms (ffn) = 12.85 ms (full) for Llama 70B decode 128K.

Bug fix (drive-by): the display block in both tab render functions was
incorrectly nested inside the "cached ctx is stale" warning branch, so
metrics/scatter/table/sensitivity/load only rendered when the cache
was stale. Un-indented and removed the dead trailing else-info block.

Verified:
  - 24 pytest tests pass (added 1 new for FFN-only scope invariants)
  - Smoke: Llama 70B decode 128K all three scopes produce sensible Pareto:
      * Attention only: 7.35 ms (14 pareto configs)
      * FFN / MoE only: 5.50 ms (34 pareto configs)
      * Attn + FFN/MoE: 12.85 ms (6 pareto configs)
    Sums add up exactly, confirming no overlap in stage summation.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 14:09:24 -07:00
mukesh b8e1a3322f analytical-viz: attention-only vs attn+FFN via two sweep buttons 
Both auto tabs now accept a scope choice at run time:
  - "Run sweep — Attention"        → include_ffn=False
  - "Run sweep — Attn + FFN/MoE"   → include_ffn=True

Each button runs an independent sweep and caches its result under its
own session_state key. The most recently clicked button determines the
displayed view; both caches persist so users can flip between the two
scopes without re-running.

Tabs:
  - Renamed "Auto Explore" → "Auto Suggest Parallelism"
    (accurately reflects that it only varies parallelism knobs; HW is
    held at the sidebar values).
  - "Auto Hardware" tab unchanged.
  - Still 6 top-level tabs; no additional tabs added.

Core changes:
  auto_explore.py:
    - New include_ffn: bool = True parameter on _sum_visible_latency,
      _efficiency, score_config, run_auto_explore, compute_parallelism_
      sensitivity. False drops all FFN stages from the summed latency.

  auto_hardware.py:
    - New include_ffn: bool = True parameter on joint_explore,
      compute_sensitivity, _best_parallelism_for_hw, _best_parallelism_
      two_stage. Forwards to score_config.

Both defaults keep existing tests byte-identical.

Verified:
  - 23 pytest tests pass (added 3 new: attn-only latency lower, attn-only
    Pareto non-empty, joint HW attn-only faster than full).
  - Smoke: Llama 70B decode 128K:
      * Attn+FFN best latency: 12.85 ms (unchanged)
      * Attention-only best:    7.35 ms (~57% of full)
      * Both sensitivities top-rank bw_hbm_gbs (physics preserved).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 13:58:29 -07:00
mukesh 91b63eb9f6 analytical-viz: parallelism sensitivity chart in Auto Explore tab 
Adds compute_parallelism_sensitivity() + ParallelismSensitivityRow to
auto_explore.py. For a baseline ConfigScore (picked from the Pareto set),
sweeps each parallelism knob (CP, TP, PP, DP, EP) individually while
holding others fixed. Reports latency + memory-fit per value.

Sweep values start at 1 and step by 2 (multiples of 2 rather than only
powers of 2), giving finer granularity for the visual than the enumerator's
sparser set. CP goes up to 256 (per real-deployment scale), TP to 64,
PP to 32.

New UI panel in Auto Explore tab: 5 subplots (log/log), one per knob:
  - Solid line = latency where the config fits memory
  - Red X markers = infeasible (out of budget)
  - Dotted vertical line = baseline value
User picks which Pareto row is the "baseline" via a number input; the
sensitivity chart re-computes around it.

Behaviour caveat noted during verification: for large PP the FFN AR can
cross a SIP boundary (uses stage_latencies.py:730 sips_used tier
selection), producing a step-up in latency. This is the existing model's
choice, honestly reflected in the chart. Whether the FFN AR should span
only one PP stage's ranks (thus not cross SIPs) is a separate discussion
about stage_latencies.py.

Verified:
- 11 pytest tests pass (added 2 new for parallelism sensitivity)
- Smoke: at Llama 70B decode 128K with baseline CP=8/TP=16/PP=1/DP=1:
  * CP sweep: 4 fits, 8 = baseline optimum, 16+ overshoots memory
  * TP sweep: 8/16 fit, 16 = baseline optimum, 32 slower (spans SIPs)
  * PP sweep: 1 = baseline, 2+ slower due to sips_used tier drop for FFN AR
  * DP sweep: same shape as PP
  * EP sweep: monotone decreasing (bigger EP = smaller per-PE FFN)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 13:39:51 -07:00
mukesh 55c20bd1a6 analytical-viz: add Auto Hardware Streamlit tab 
Consumes auto_hardware.joint_explore(). UI:

  - Sweep-depth radio: two_stage / balanced (default) / coarse
  - Run joint sweep button + spinner
  - 4 metric cards: HW candidates, feasible joint, Pareto count,
    best latency
  - Panel 1 (Pareto scatter): latency vs hardware cost proxy, feasible
    in grey, Pareto colored by PE count, dashed line connects Pareto in
    cost order — this is the "optimal HW config for the model" plot
  - Panel 2 (sensitivity bar chart): per-knob relative speedup when
    doubled from the baseline. Green bars, sorted biggest first;
    annotated with the baseline → doubled values. Answers "where to
    invest next?"
  - Panel 3 (table): sortable Pareto joint configs with parallelism +
    HW spec columns (PE HBM, HBM BW, TFLOPs, PE↔PE, D2D, C2C)
  - Load into sidebar: picks a Pareto row and syncs both the HW
    selectboxes (Per-PE + Interconnect sub-tabs) AND the parallelism
    sliders. Values only get loaded when the target is one of the
    selectbox options; otherwise the sidebar keeps its current value.

Session-state cache under _hw_result keyed by (model, s_kv, mode, depth);
if any of these drift, a warning suggests refreshing.

Verified: app.py parses; auto_hardware smoke run on Llama 70B decode
128K in ~20s produces sensible HW co-design signal (HBM BW 33% speedup,
everything else <1.5%).

Next: verification across Qwen 3 8B, Mixtral 8x7B (Commit 6).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 13:19:13 -07:00
mukesh 6ef09dd6f5 analytical-viz: auto_hardware.py — joint HW×parallelism explore 
New module extends auto_explore into the hardware co-design space. For a
fixed model + workload, sweeps hardware knobs (pe_hbm_gb, bw_hbm_gbs,
peak_tflops_f16, bw_intra_gbs, bw_inter_gbs, bw_intersip_gbs) and, for
each hardware candidate, searches parallelism for the latency-minimum
that fits memory. Returns:

  - all_scores: every (hw, parallelism) pair that fits, sorted by latency
  - pareto_scores: 2D Pareto frontier on (latency ↓, cost_score ↓)
  - sensitivity: per-knob rel_speedup when doubled from the best-fast HW
    baseline. Ranks which HW knob gives the biggest speedup — a co-design
    signal.

Three sweep depths trade coverage for time:
  - two_stage: 1 HW candidate (defaults) × autosuggest's memory-min
    parallelism. Fast (~1s), useful for the sensitivity ranking alone.
  - balanced: 64 HW × ~2k reduced-parallelism configs = ~130k joint evals,
    ~10-20s. Default UI setting.
  - coarse:   729 HW × ~2k configs = ~1.4M joint evals, ~2-5 min.

Reduced parallelism sweep for the inner loop: CP × TP × PP × DP ×
kv_shard_mode (1,920 configs), other 4 knobs held at latency-friendly
defaults (ffn_shard_scope='TP+CP', tp_placement='cube', cp_placement='pe',
cp_ring_variant='qoml' for decode, 'kv' for prefill). Full 28,800-config
auto_explore per HW would take 6+ minutes — too slow.

Cost proxy: sum of (knob / knob_default). 6.0 at defaults. Not dollars —
a rough capability score where higher = "more spec'd hardware".

Verified:
- 9 pytest tests pass:
    * enumeration counts match expected (1, 64, 729)
    * default cost_score = 6.0
    * Pareto non-dominated + subset of all_scores
    * every knob is monotone-non-worsening when doubled
    * for Llama 70B decode, bw_hbm_gbs tops the sensitivity ranking
      (physically correct: memory-bound workload)
- Smoke: Llama 70B decode 128K balanced sweep in ~20s produces
  5 Pareto configs; best 7.57 ms with 1024 GB/s HBM BW. Doubling
  HBM BW gives 33% additional speedup; every other knob < 1.5%.

Next: Streamlit UI tab consuming this in Commit 5.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 13:16:40 -07:00
mukesh 2834383700 analytical-viz: add Auto Explore Streamlit tab 
New tab consumes auto_explore.run_auto_explore(). UI:
  - Header showing current model + workload + per-PE HBM budget
  - Run sweep button (spinner while ~28k configs run in ~5-10s)
  - 4 metric cards: enumerated, feasible, pareto count, best latency
  - 2-panel scatter:
    * latency vs PEs (feasible in grey, Pareto coloured by efficiency,
      dashed line connects Pareto in PE order to show the trade-off curve)
    * HBM utilization vs latency for Pareto, coloured by PE count
  - Sortable Pareto table
  - "Load into sidebar" widget: pick a row, sets the sidebar
    session_state keys (cp, tp, pp, dp, tp_placement, cp_placement,
    cp_ring_variant, kv_mode, ffn_scope_label) and st.rerun()s so the
    user can flip to another tab and see the full breakdown.

Session-state caches the last sweep result under _auto_explore_result;
if the model/workload/HBM change without re-running, a warning suggests
refreshing.

Ffn_scope_label mapping is dynamic (contains substituted divisors), so
the Load button reconstructs the exact label using the target
cp/tp/dp values before assigning to session_state["ffn_scope_label"].

Verified:
- app.py parses cleanly
- All 9 pytest tests in test_auto_explore.py still pass
- Smoke: matplotlib + pandas + auto_explore imports round-trip

Next: verification pass across Qwen 3 8B and Mixtral 8x7B; polish.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 13:02:08 -07:00
mukesh 85f6716fc1 analytical-viz: auto_explore.py — 9-knob Pareto search 
Extends the memory-only autosuggest to search the full 9-dimensional
parallelism space (CP, TP, PP, DP, kv_shard_mode, ffn_shard_scope,
tp_placement, cp_placement, cp_ring_variant) and rank feasible configs
on a 3D Pareto frontier: (latency ↓, pes_used ↓, efficiency ↑).

Throughput is stored on ConfigScore for display but is deliberately NOT
a Pareto axis because for a single-request analysis it collapses to
1 / latency, which would collapse the frontier.

Reuses existing physics (stage_latencies.all_stages + all_ffn_stages,
memory_layout.compute_memory) — no new formulas.

Single-request latency formula fix: PP does NOT reduce single-request
decode/prefill latency because the request has to traverse every layer
sequentially regardless of pipeline depth. The initial version had
latency ~ per_layer × layers_per_stage, which incorrectly rewarded high
PP. Corrected to latency ~ per_layer × model.layers.

Enumerator prunes:
  - PP > model.layers
  - TP > 4 × h_q
  - ffn_shard_scope contains 'DP' when dp=1 (redundant)
  - cp_ring_variant='qoml' when cp=1 (no-op)

Full sweep on Llama 3.1 70B: ~28,800 configs enumerated in ~7s, ~7k-10k
feasible (varies with S_kv), 2-7 unique Pareto configs. Faster context
lengths produce richer frontiers; at 1M, memory forces a single
dominant config (128 PEs, HBM 85%).

Verified:
- 9 pytest tests pass (enumerate, score, Pareto, subset invariants)
- Manual: Llama 70B decode at 8K/64K/128K/1M produces physically
  sensible Pareto (CP=8/TP=16 wins latency; smaller-PE options
  appear at longer context up to memory limits)

Next: Streamlit tab UI in app.py + verification against more presets.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 12:58:10 -07:00
mukesh 09e9331026 topology: set flit_bytes=4096 (was implicit 256 default) 
Wire per-flit loop in engine._wire (ADR-0033 Phase 2c) iterates
nbytes/flit_bytes times per hop, firing ~2 SimPy events per flit.
At long-context decode, K/V tile loads are hundreds of KB per shard,
producing thousands of flit-iterations per hop. The default 256-byte
flit size made the wire loop the dominant wall-clock cost.

Raising flit_bytes to 4096 (still a multiple of hbm_ctrl.burst_bytes=256
per ADR-0033 D1) cuts flit iterations 16x with no meaningful change to
modeled latency. Wormhole pipelining and the available_at BW-share
chain in _wire behave identically at any flit size.

Measured impact on Case 4 (Cube-SP x PE-SP) decode, C=8, P=8, T_q=1,
LLaMA-3.1-70B single-KV-group, enable_data=False:

  S_kv     flit=256 wall    flit=4096 wall    speedup    latency drift
   128K       233.71 s          18.16 s        12.9x     +0.05% (0.25us)
   256K       485.75 s          25.31 s        19.2x     +0.01% (0.09us)
   512K      1045.64 s          48.45 s        21.6x     +0.01% (0.17us)
    1M      ~2250 s (proj)     104.75 s        ~21x      within extrap

Total sweep at flit=4096: ~3.3 min. Baseline would have been ~65 min.
1M decode iteration is now practical.

op_log records are byte-identical (same counts, same components); the
tiny latency drift (order of a few ns) comes from the last flit's
env.timeout landing on a slightly different granularity when flits are
larger. Well within noise for any test that asserts on structural
invariants (dma_write cubes, ipcq_copy counts).

Verified:
- tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py: 18/18 pass
  (also 7x faster: 174s -> 24.83s per full run)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-28 11:35:04 -07:00
mukesh f9d0077472 router: memoize Dijkstra path lookups 
Per-instance cache keyed by (id(adj), start, goal). All 4 adj dicts
(_adj, _adj_all, _adj_local, _adj_mcpu_dma) are built in __init__ and
never mutated (topology is static per ADR-0006 / SPEC §0.1), so id(adj)
is stable for the router's lifetime. Cache is populated on the
successful-return paths; RoutingError paths intentionally re-run each
call (rare, keeps error semantics unchanged).

Motivation: cProfile of Case 4 decode at S_kv=8K showed 1,142
_run_dijkstra_with_dist calls consuming ~1.95s tottime. Paths depend
only on (adj, src, dst) so ~99% of those calls are recomputing the same
result.

Verified:
- tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py: 18/18 pass
- 128K decode wall: 237.28s -> 233.71s (-1.5%)
- Modeled kernel latency: 461.13us (byte-identical before/after)
- op_log_len: 3057 (unchanged)

Small win but no risk: memoization returns byte-identical results and
paths cannot change during a sim run. Larger event-count reductions
require touching the per-hop hot path (zero-latency chain collapse
etc.).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-27 23:39:39 -07:00
mukesh 65f358fdfa sim: report correct kernel latency under enable_data=False 
Three coupled fixes so enable_data=False now reports byte-identical
kernel-body sim latency (max(t_end) - min(t_start) over op_log) as
enable_data=True, and skips setup-write sim events that were pure
wall-clock overhead.

Before this change, at Case 4 (Cube-SP x PE-SP) decode, S_kv=8K:
  enable_data=False  -> latency_ns = 0  (op_log empty)
  enable_data=True   -> latency_ns = 30.646 us  (correct)

After:
  enable_data=False  -> latency_ns = 30.646 us  (byte-equal)
  enable_data=True   -> latency_ns = 30.646 us  (unchanged)

engine.py: always create OpLogger. Previously OpLogger was gated on
enable_data=True together with MemoryStore, so op_log stayed empty when
data mode was off. Decouple: MemoryStore is gated (Phase 2 DataExecutor
still needs it) but OpLogger runs unconditionally, receiving
memory_store=None when data mode is off. OpLogger already guards its
arr.copy() snapshot paths on `if self._memory_store is not None`.

pe_cpu.py: always use the ADR-0020 greenlet execution path. The
if store is not None: _execute_greenlet(); else: _execute_legacy branch
gated the *execution model* on data-mode presence. The legacy
command-list path predates ADR-0020 and doesn't route IpcqSendCmd
through PE_IPCQ (it goes to PE_SCHEDULER instead), so all 189 Case-4
ipcq_copy fabric transfers were silently no-op'd in that mode. Forcing
greenlet gives identical sim behavior in both modes. KernelRunner
already guards its store reads on `self._store is not None`.

context.py: gate setup MemoryWriteMsg on memory_store presence. Under
enable_data=False there is no MemoryStore to populate and no Phase 2
replay, so the per-shard sim events for Q/K/V deploy are pure wall-clock
overhead with no effect on reported kernel latency (Yangwook's max-min
formula excludes pre-kernel ops). Handle count for Case 4 at S_kv=8K
drops 229 -> 37 under enable_data=False.

Verified:
- tests/attention/test_milestone_gqa_decode_long_ctx_4cases.py: 18/18 pass
- Parity probe: enable_data=False/True both report 30.646 us kernel sim
  time, 1649 op_log records, 189 ipcq_copy events

Known follow-up: pe_cpu._execute_legacy is now dead code but left in
place; separate cleanup once we confirm no downstream caller.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-27 22:05:39 -07:00
mukesh 9fdde44922 analytical-viz: interactive Streamlit tool for SIP transformer analysis 
New tests/analytical_visualization/ module - an interactive dashboard
for exploring memory / latency tradeoffs of transformer inference on
the SIP architecture.

Highlights:
- 30+ model presets (Qwen, Llama 2/3/3.1, Mistral, Gemma 2, Phi 3,
  DeepSeek MLA, Mixtral/Qwen 3 MoE, Grok-1, ...)
- Placement toggles: TP and CP each on PE-level vs cube-level
- CP ring variant: K/V ring vs Q+O/m/l ring (prefill); in decode the
  O/m/l all-reduce is folded into S8 (no separate C1 row)
- SIP interconnect: ring / mesh2d / torus2d with matching link drawing
- Per-stage latency table with compute + memory + comm formulas,
  auto-scaled ns/us/ms, colored by dominant bound
- Ring attention loop indicator on the pipeline diagram (purple arc
  over S5-S8 with 'xN hops' badge)
- Tensor sharding view with optional physical PE/cube annotations
- Replication-waste + optimization-hints panel
- Save & compare configurations (config1, config2, ...): summary table
  plus side-by-side per-stage attention and FFN latency, best-in-row
  highlighting
- Symbol glossary with current values for every symbol used in formulas

Not tied to production sim_engine or runtime API; purely analytical
tooling for design-space exploration.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-07-26 22:12:08 -07:00
ywkang ae131aa2af report(S5.2): explain why dense mapping wins throughput (per-PE efficiency) 
Both 1-kv (2 users x 64 PE) and 8-kv (16 users x 8 PE) fill all 128 PEs of
the SIP; the throughput gap is per-PE HBM efficiency. Splitting a head over
8 PEs leaves each a single tile, so pipeline-fill and the 8-PE softmax
reduce dominate -> 46% HBM util; one head per PE streams contiguously with
no reduce -> 76%. Throughput ratio (1.6x) tracks the utilization ratio
(1.7x): decode is bandwidth-bound, so throughput follows total HBM
efficiency, not PE count. 1-kv spends hardware on latency (8 PE/head for a
4.8x speedup = 60% strong-scaling efficiency).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-07-22 17:46:54 -07:00
ywkang baf22f01ab plot: extend batch-scaling curves with saturated-throughput ceiling 
Beyond a mapping's SIP capacity (16/C) the throughput plateaus: extra
users run in waves at the saturated rate, they are not un-runnable. Draw
a dashed horizontal extension at the ceiling so 1-kv-per-cube (cap 2)
reads as saturating, not stopping, at B=2. Caption updated.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-07-22 17:39:13 -07:00
ywkang 13f5cbc317 gqa decode: fix cube_base output-store; measure batch scaling on one SIP 
Concurrency fix (completes the cube_base change): the decode kernel's
output store o_base still used the global cube_id while every load used
the user-local cube_local. For a single user (cube_base=0) they coincide,
so it was masked; a second batched user (cube_base>0) overshot its o shard
into an unmapped VA, mis-decoded to a phantom sip0.cube0.pe8 and raised a
RoutingError. Switch o_base to cube_local (one line); single-user results
are byte-identical (decode smoke/correctness pass).

Batch harness (un-skipped): create all users' tensors first, then submit
all launches deferred and drain together, so deploys don't drive an
already-submitted launch and serialize the batch. Concurrent users on
disjoint CUBE groups now overlap to within 3-5% of the single-user
latency. Swept B in {1,2,capacity} at 8K (high-B dense runs are the
expensive ones; throughput is near-linear in B).

Measured result: aggregate throughput at a full SIP rises from 0.86
(1-kv-per-cube, 2 users) to 1.41 requests/us (8-kv-per-cube, 16 users) —
the dense mapping wins throughput, the spread mapping wins latency.
S5.2 batch paragraph updated projected -> measured; add batch_scaling
figure and plot_batch_scaling.py.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-07-22 17:37:10 -07:00
ywkang 8c108154a7 plot: fix Fig 16 suptitle to match KV-footprint panel 
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-07-22 16:25:50 -07:00
ywkang 711a9a257f gqa short-context: exclude KV deploy from wall; correct HBM peak; batch scaffold 
Metric fixes (test harnesses, deploy artifact + wrong constant):
- wall = max(t_end) - min(t_start): exclude the one-time KV-cache deploy
  from the measured decode/prefill step (it was 92-99% of wall, mapping-
  invariant, and masked the real per-mapping separation).
- PEAK_PE_HBM_BPS 128 -> 256 B/ns (8 channels/PE x 32 GB/s); the old value
  was half the modeled per-PE HBM BW, so hbm_bw_util read >1.0 once wall
  was corrected. All six short-context sweep CSVs regenerated/repatched.

Result: the four KV mappings now separate along a {64,32,16,8}-active-PE
ladder (decode 8-kv/1-kv = 4.8x at 8K, 7.4x at 64K), not "modest/tied" as
before; decode is bandwidth-bound at 46-76% of the 256 GB/s per-PE ceiling.

Report (S5.2 rewrite):
- Replace the tied-wall / 8x-per-PE-util claims (both deploy artifacts)
  with the corrected separation and a density trade-off (per-CUBE KV
  footprint, Fig 16 top panel switched to a wall-invariant metric).
- Add a projected latency-vs-batched-throughput analysis (marked
  projected, not measured): dense mappings win throughput, 1-kv wins
  latency; converges at long context.
- Regenerate Fig 15/16/17/18; fix plot script hardcoded ROOT path.

Batch experiment (Part 2, cube_base):
- Add backward-compatible cube_base=0 scalar to the decode kernel so a
  batched user placed at DPPolicy.cube_start addresses 0-based shards.
  Default preserves single-user behavior (32 decode tests pass unchanged).
- New batch harness (skipped): concurrent B>=2 launches hit a sim routing
  issue (sip0.cube0.pe8); single-user path verified. Concurrency fix next.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-07-22 16:25:04 -07:00
ywkang cb6d21f145 report(1H): tighten §5 GQA (mapping tie, taxonomy bridge, dedup, Summary) 
- Soften short-context claim: 1-kv-/8-kv-per-cube are tied at 8K and
  separate as context grows (1-kv faster at the large end).
- Add a bridge noting the six-case placement taxonomy (§5.1) is the
  long-context lens; short context turns on head-to-CUBE distribution.
- Remove caption-duplicating numbers from the long-context conclusion.
- Reduce the closing subsection to a GQA-scoped Summary; defer the
  cross-cutting hardware-investment thesis to the Discussion.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-07-22 10:33:56 -07:00
ywkang d88e6cc72e report(1H): rebuild main.pdf 
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-07-22 10:22:06 -07:00
ywkang edb30326ce report(1H): add Agentic Workloads (§6) and HW-Spec Search (§7) sections 
- §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>
 2026-07-22 10:20:27 -07:00
mukesh d0db9ff40e gqa-decode: topology-agnostic reduce + multi-model attention bench 
Make reduce_mlo derive its submesh dimensions (sub_w, sub_h, root_col,
root_row, root_cube) from C at call time, with peer-existence guards on
every inter-cube send/recv so any C ≥ 1 completes cleanly (previously
hardcoded to C=8's 4×2 mesh; non-rectangular C like 7 or 12 hit
IpcqDeadlock at the root cube's east receive). Byte-equal at C=8 —
existing composite digest tests still pass 8/8.

Callsites in the four Case-6 kernels (primitive / primitive-tiled /
composite / composite_extended) updated to pass C into reduce_mlo and
use root_cube_for(C) for the final tl.store gate.

Add the multi-model attention bench: sweeps the composite kernel at
S_kv=128K across six GQA models (Gemma-2 27B, LLaMA-3 8B, Qwen 2.5 7B,
LLaMA-3 70B, Qwen 2.5 72B, Command R+), with per-model topology
C = h_q (cubes per KV group = query heads per KV group). Captures
per-op-kind occupancy (matmul GEMM+MATH, comm DMA) alongside latency
and PE_CPU dispatch. Three-panel plot writes to bench-output and
paper-figures dir.

Headline result: same-h_q-different-family models land within 0.2 µs
(model-agnostic given fixed h_kv=1 per KV group + d_head=128); latency
climbs 263 → 492 µs as C climbs 2 → 12, driven by reduce depth
(comm/matmul ratio 0.4 → 0.6 across the sweep).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-07-22 09:54:31 -07:00
mukesh 23a2bc3fb3 gqa-decode-composite: measure primitive hand-tiled (16×16×16) latency 
Switch primitive hand-tiled Q·Kᵀ / P·V from a deferred-K-sum tl.dot
chain to per-block GemmCmd writes into a shared (M, N) out handle
(implicit MAC-side accumulation). Full-shape coarse Q / K_T / V loads
carry data-mode correctness; per-block DMAs remain for the streaming
architecture dispatch story. The kernel now runs in engine mode
end-to-end, so its 128K latency is measured instead of derived as
primitive + Δdispatch. Drop the coarse-primitive baseline from the
sweep and plots; keep the kernel file for reference.

At S_kv=128K: primitive hand-tiled = 959.5 µs vs composite = 460.7 µs
(2.08× from 5 235 vs 94 PE_CPU commands).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-07-13 21:49:57 -07:00
mukesh eec61838b5 sweep+plot: add primitive hand-tiled (16×16×16) variant to Case-6 decode composite study 
Wires a fourth variant into the Cube-SP × PE-SP long-context decode
composite command-form study to surface the worst-case PE_CPU dispatch
inflation that the coarse composite forms delegate to PE_SCHEDULER
(ADR-0065): per-block DMA of Q/K/V slices, tl.dot per 16³ block,
deferred K-inner sum outside the K loop.

Renamed _tiled.py → _hand_tiled_16x16x16.py; coarse-primitive kernel
retained for the 4-cases bench, paper scripts, and the golden byte-equal
regression guard.

New breakdown bar chart at S_kv=128K shows engine (~460 μs) dominates
all three variants; hand-tiled adds ~68 μs PE_CPU dispatch on top —
composite forms sit essentially on the memory-bound floor.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-07-13 09:28:57 -07:00
eusonice f8caf2e16c report(gqa-short): draft results subsection (mode + composite-tier comparison, 4 figures) 2026-06-26 16:01:32 -07:00
eusonice 09e70bd516 test: fix CSV output paths in 3 sweep files (each variant now writes a unique CSV) 2026-06-26 15:59:21 -07:00
eusonice e20e1f7362 plot: merge plot scripts into single plot_short_context.py 2026-06-26 15:59:21 -07:00
eusonice 3437d7e079 sweep: 6 short-context CSVs (3 variants x prefill+decode) + 6 wall/util figures 2026-06-26 15:59:21 -07:00
eusonice 3da116d40d adr-0070: gqa short context attention prefill/decode + 3 composite tiers 2026-06-26 15:59:21 -07:00
eusonice f7aa8134b9 test: 3-variant sweep (baseline + composite + composite-fused) for prefill/decode 2026-06-26 15:59:21 -07:00
eusonice a455ae61b4 sim: pin compute outputs and on-chip recv slots for composite operands 2026-06-26 15:59:21 -07:00
eusonice 1971ac731c add three gqa attention kernel variants for short context length 2026-06-26 15:59:21 -07:00
eusonice a98e93110b short context length composite command 2026-06-26 15:59:21 -07:00
eusonice f08dda3bd4 test short context length attention kernel 2026-06-26 15:59:21 -07:00
eusonice 7d90d53d4d short context length attention kernel 2026-06-26 15:59:21 -07:00
ywkang faf011dae5 paper(ipcq): regen alternatives figures + comparison plot + allreduce section 
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-23 19:19:55 -07:00
ywkang 1dade267ff experiment(gqa-decode): 16x16x16 MAC-block dispatch model 
Research artifact (NOT wired into the production sweep) for the decode composite-vs-primitive investigation. Models a primitive decode kernel that hand-blocks each tl.dot into 16x16x16 GemmCmds, charging per-MAC-block PE_CPU dispatch -- testing whether faithfully charging the dispatch that the composite form offloads to PE_SCHEDULER flips the 'composite gives no decode latency benefit' result.

Finding: it does NOT. Even with up to 512 serialized blocking GemmCmds per matmul (vs 2 coarse tl.dot), end-to-end decode latency is unchanged (8192: 28.9us, 32768: 114.9us) -- the inter-cube (m,l,O) reduce DMA tail dominates and the local-attention GEMM/dispatch sits in critical-path slack. Confirms the two-regime conclusion (24c7054): composite helps compute-bound prefill, not memory/reduce-bound decode.

Caveats: measured at S_kv 8K/32K (reduce-dominated); large-S_kv streaming regime extrapolated only. The -5.5% tiled<untiled is reduce-tail ordering noise, not fully nailed down.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-23 18:24:26 -07:00
ywkang 24c705419e paper(gqa): two-regime "use of composite commands" — decode + compute-bound prefill 
Complete the composite-command study across both attention regimes and
write up the result, resolving when the composite command helps latency.

Decode (memory-bound, T_q=1, M=8): the command form is latency-neutral —
the kernel is bound by KV streaming and the MAC array is near-idle, so the
composite's only benefit here is host-issue offload (O(n_tiles) -> O(1)
PE_CPU commands). Figure: gqa_decode_long_ctx_composite (latency flat,
command count saturates).

Prefill (compute-bound, large M=G*T_q tile-filling): add three command-form
variants of a single-rank FlashAttention prefill kernel
(_gqa_prefill_compute_bound: primitive / composite / composite_extended).
Here the composite's scheduler-internal per-tile DMA<->compute pipeline
keeps the MAC array fed while the primitive's blocking tl.dot starves it,
so composite wins on MAC utilization (68% flat -> 83%) and wall-clock
(19% faster at ctx=1024), with the margin growing with context (deeper
P.V reduction = more tiles to pipeline) — the compute-bound mirror of the
GEMM result in section 3. Figure: gqa_prefill_compute_bound (latency +
MAC util). Sweep wired as GQA_1H_SWEEPS=prefill_cb; tests cover structure,
e2e, and composite<primitive in the compute-bound regime.

The synthesis: composite has two benefits set by roofline position —
host-issue offload (always) and MAC-array feeding (compute-bound only).
Decode exercises the first, prefill both.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-23 17:04:26 -07:00
ywkang cd6f0ed91d feat(gqa): Case-6 composite-command decode variants + on-chip-operand DMA fix 
Add two command-form variants of the Case-6 (Cube-SP x PE-SP) long-context
decode kernel alongside the primitive baseline:
  - composite: per-tile GEMMs issued as one coarse tl.composite over the
    full S_local (PE_SCHEDULER tiles/streams K,V) -- O(1) PE_CPU commands
    vs the primitive kernel's O(n_tiles).
  - composite_extended: Q.K^T composite + softmax_merge recipe composite
    (ADR-0065), folding the per-tile online merge + P.V.

Extract the shared two-level (m,l,O) reduce into _gqa_mlo_reduce and
refactor the baseline to use it (byte-equal, guarded by a 64-rank
command-stream digest test).

Engine fix: _make_compute_out omitted pinned=True, so an on-chip TCM
compute result consumed as a composite GEMM operand was scheduled as an
HBM DMA_READ of its bit-61 scratch address -> PhysAddrError in data mode.
Mark compute outputs pinned (resident), matching the composite
auto-output and recipe scratch handles. .pinned is read only by the
tiling DMA gate, so this is byte-equal for existing benches.

Add the composite sweep (emit-level PE_CPU dispatch to 1M + data-mode
latency to 128K), its plot, umbrella wiring (GQA_1H_SWEEPS=composite),
and tests (refactor guard, dispatch saturation, engine-fix, e2e).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-23 16:07:41 -07:00
ywkang e9a5c438e3 paper(1H): 6-case GQA long-ctx + IPCQ design alternatives + GEMM terminology 
- §6 GQA: rewrite long-context decode from 4-case to 6-case. Data
  Placement Policy now presents the six placement options and their
  intrinsic per-PE memory / per-token comm costs (no winner predicted);
  the long-context subsection selects the best placement for that regime
  (Case 6 ★, both-axes S_kv shard) from the measured 6-case sweep. Add
  KV-sharding diagram + analytical budget/summary figures.
- Regenerate the 6-row decode sweep (milestone-1h-gqa) and the
  sweep-dependent decode panels (latency/traffic/parallelism/memory) so
  figures, prose, and sweep_decode.json are mutually consistent.
- §5 All-Reduce: add IPCQ design alternatives (architecture + decision
  matrix) as design-rationale schematics (illustrative step-counts, not
  measured) and the bench-generated topology diagram.
- §4 GEMM: rename "user-orchestrated/user-level" -> "kernel-orchestrated/
  kernel-level" (orchestration runs in the kernel program vs the
  scheduler-orchestrated composite); minor accuracy fixes (7.18 TFLOP/s
  ~10% below peak; ~781 ns DMA).
- Recompile build/main.pdf.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-22 00:55:54 -07:00
mukesh bb668a3ac3 paper(gqa): derive (m,l,O) AR per-PE cost from kernel topology 
Replaces the unsourced 32 KB/layer placeholder in the analytical chart
with T*(N-1)/N where T = h_q * S_q * (d_head*2 + 8) ~ 2.1 KB
(per-PE (m, l, O) payload) and N is the participant count of the
reduce-only chain per stage. Cases 2/3/6 analytical bars now match
the simulator-measured numbers within ~10% (was 6-30x over).

Also:
- Bumps the measurement S_kv from 8 K to 64 K to verify Cases 1, 2, 3, 6
  are genuinely S_kv-independent and Cases 4, 5 scale linearly.
- In-bar descriptor on the paired chart is now wrapped to fit the bar
  width, black, no background box, and centered between the analytical
  and measured bars so the label applies to both.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-18 15:38:20 -07:00
mukesh 0662c76fa7 paper(ipcq): de-emphasize flow-diagram outer frame so inner flow stands out 
Drops the per-case panel outline from linewidth 2.2 to 0.9 with 85% alpha.
The case-color frame still anchors each panel but no longer competes
with the swimlanes, flit packets, and annotations inside.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-18 15:01:09 -07:00
mukesh 37dc48fd27 paper(ipcq): add latency-model-style flow + stacked architecture diagrams 
Adds two new IPCQ architecture views in the latency_model.png aesthetic:
a 2x2 flow figure (one per case) and a 4x1 stacked figure for taller
horizontal panels. Refactors per-case panel drawing into shared helpers
and disambiguates the flit labels (HMQ desc -> Hd so it does not
collide with data D). Two-row legend keys.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-18 14:55:59 -07:00
mukesh 7c45047e93 paper(ipcq): swap control-data overlay for latency PNG in figures/ 
Drops the control-data overlay from the paper figures folder and
adds the per-send latency cycle-stack PNG instead.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-18 14:06:32 -07:00
mukesh 73e0b315fe paper(gqa): table cleanup — drop measured-comm column, narrower fig, wrapped title 
Removes the (measured) comm column, widens the Notes column with
shorter wrapped text, scales the table figure down, and wraps the
heading to two lines so it sits flush with the table.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-18 14:02:42 -07:00
mukesh 1552028c25 paper(ipcq): add IPCQ vs HW alternatives comparative study 
6 generated figures comparing IPCQ (chosen) against the 3 HW alternatives
(Doorbell+Polling, HMQ, RDMA-CQ): architecture blocks, latency stack,
op counts, sequence flow, decision matrix, and control/data plane overlay.

Three figures also copied into the paper's figures/ folder for inclusion.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-18 13:27:23 -07:00
mukesh 56c51b184a paper(gqa): refresh 4 decode 6-case PNGs in figures/ with actual 6-case content 
The previous commit (8102ddb) renamed the figures-dir
gqa_decode_long_ctx_4cases_*.png → 6cases_*.png but only the rename
was staged; the cp that brought the regenerated 6-case content
landed after `git add`, so the committed bytes were still the old
4-case ones. This commit syncs the bytes.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-18 11:51:30 -07:00
mukesh 8102ddbe30 paper(gqa): rename long-ctx artifacts to gqa_long_ctx_6cases_* and add 4 decode 6-case charts 
Filename cleanup so every long-ctx GQA artifact has a consistent
"gqa_long_ctx_6cases_*" prefix (or "gqa_decode_long_ctx_6cases_*"
for decode-only charts). Old "4cases" / mixed names retired.

Renames (long_ctx + figures, content unchanged):
  gqa_hbm_budget.png                  -> gqa_long_ctx_6cases_hbm_budget.png
  gqa_4cases_summary.png              -> gqa_long_ctx_6cases_summary.png
  gqa_4cases_memory_comm_analytical   -> gqa_long_ctx_6cases_memory_comm_analytical.png
  gqa_4cases_memory_comm_paired       -> gqa_long_ctx_6cases_memory_comm_paired.png
  gqa_kv_sharding_6cases_diagram      -> gqa_long_ctx_6cases_kv_sharding_diagram.png
  gqa_kv_sharding_6cases_table        -> gqa_long_ctx_6cases_kv_sharding_table.png
  gqa_3cases_measured_comm.json       -> gqa_long_ctx_6cases_measured_comm.json
  gqa_decode_long_ctx_4cases_*.png    -> gqa_decode_long_ctx_6cases_*.png
                                         (figures dir; long_ctx never had old)

New 4-chart 6-case set in long_ctx output dir (regenerated by
paper_plot_gqa_decode_long_ctx_4cases.py, which now reads all 6
sweep_decode.json panels — Cases 1-6 with the same colour scheme
used elsewhere: red = overflow per-PE HBM, grey = neutral, blue
= Pareto-best ★):

  gqa_decode_long_ctx_6cases_latency.png
  gqa_decode_long_ctx_6cases_memory.png
  gqa_decode_long_ctx_6cases_parallelism.png
  gqa_decode_long_ctx_6cases_traffic.png

Generator scripts updated to write the new filenames + handle the
two new d_head-TP variants (Cases 4, 5) in their per-PE memory and
active-PE-count helpers. Figure widths bumped 10 -> 12 in to fit 6
multi-line case labels.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-18 11:49:03 -07:00
mukesh 84bb418e1e gqa(decode): add d_head-TP kernels + measurement runner + figure generators 
Two new long-ctx decode attention kernels for the d_head-TP sharding
variants the 6-case chart predicts:

  · _gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead.py
        Cube-SP × PE-TP-d_head (Case 4 in chart). Per cube holds
        S_kv/C tokens of full d_head; per PE holds same tokens but
        only d_head/P dims. Partial Q·Kᵀ scores reduced intra-cube
        before softmax; outer (m,ℓ,O) merge two-phase (intra+inter).

  · _gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp.py
        Cube-TP-d_head × PE-SP (Case 5). Per cube holds full S_kv
        with only d_head/C dims; per PE holds S_kv/P of those dims.
        Partial scores reduced inter-cube (UCIe) before softmax.

Sweep dispatch (gqa_decode_long_ctx_4cases.py) extended with two
new panels so the milestone-1h-gqa sweep covers all 6 cases.

Smoke test scripts/verify_case4_dhead_tp.py runs Cases 4/5/6 at
S_kv=2K to validate the kernels load and execute end-to-end.

Plus the figure-generation toolchain that produced the committed
PNGs in the prior commit (dd3337f):

  · paper_plot_gqa_4cases_summary.py    - 3-panel summary +
        2-panel (analytical / paired-measured) chart generator.
        _plot_comm now takes mode="analytical" | "paired".
  · paper_plot_gqa_kv_sharding_diagram.py  - 6-case 2-D KV-tensor
        diagram + companion comparison-table PNG.
  · measure_gqa_decode_placement_comm.py   - runs all 6 kernels at
        S_kv=8K, sums actual IPCQ-copy bytes from engine.op_log,
        scales partial-score AR ×128 to S_kv=1M, writes
        gqa_3cases_measured_comm.json (committed).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-18 11:26:05 -07:00
mukesh dd3337f2e4 paper(gqa): add 6-case KV-sharding figures (analytical + measured) 
Six PNGs covering the 6-case decode KV-placement story, copied
into both the kernbench 1H_milestone_output staging dir and the
1H-codesign-paper figures/ dir:

  · gqa_hbm_budget.png                 - per-PE HBM budget bar
  · gqa_4cases_summary.png             - 3-panel: HBM + KV mem + comm
  · gqa_4cases_memory_comm_analytical  - 2-panel: KV mem + comm (analytical)
  · gqa_4cases_memory_comm_paired      - 2-panel: + measured overlay
  · gqa_kv_sharding_6cases_diagram     - 2-D KV-tensor view (Y=S_kv, X=d_head)
  · gqa_kv_sharding_6cases_table       - companion comparison table

The original 4-case gqa_4cases_summary.png is updated to the new
6-case 3-panel layout. Generator scripts and Phase-1 d_head-TP
kernels remain untracked - separate commit scope.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-18 11:22:57 -07:00
ywkang b6315c3c90 paper(1H): reflect D8 single-op cost model in §2/§3.4 + figure/diagram regen 
Two strands bundled as the 1H-codesign-paper refresh unit:

(A) This session — single-op cost-model reflection (depends on 2d8271c):
  - §2 Table 2 (tab:hw): split "FIXED per command" into "FIXED per
    single-op command" (8 cycles) and "FIXED per composite command"
    (40 cycles); §2 dispatch-overhead prose updated to the two-class split.
  - §3.4 (sec:gemm-vs-async): rename paragraph headers + prose to
    async-full / async-tiled; "atomic" -> "single-op" throughout; reframe
    mechanism #3 from the old DMA-only fast-path to the single-op
    fast-path. Headline narrative now: even with EVERY single-op cmd
    (96 DMA + 48 dot + 47 add) charged the light 8-cycle FIXED, composite
    still wins ~2.8x at K=3072 purely on command-count structure (1 vs
    192 commands) -- down from the pre-D8 ~6.3x, and explicitly NOT a
    modelling artifact. Numbers refreshed from the regenerated sweep:
    async-full 3.83->3.91, async-tiled 1.14->~2.53, under-tile corner
    1.06->1.21, depth-2 vs depth-inf spread <1%. New figure wired in.
  - build/main.pdf rebuilt (tectonic); pdftotext-verified (no broken
    refs; Table 2 split, single-op terms, 2.8x/2.53/192-host-commands
    all present).

(B) Prior-session paper work riding along uncommitted: §4 all-reduce
    deep-edit, §5 GQA, §6 discussion trims; milestone_1h_ccl.py plot
    label "FSIM" -> "H2 2025 SW queue baseline"; regenerated diagrams
    under docs/diagrams/** and gemm output PNGs under
    1H_milestone_output/gemm/. (Composite-window gemm plots are
    unaffected by D8 — D8 only changes single-op dispatch FIXED, which
    the composite window excludes.)

All TODO items for the D8 single-op extension are now complete and
pushed across 3 commits (2d8271c cost-model+ADR+tests, 821bbf2 bench
harness, this paper refresh). Full regression green (826 passed, 1
skipped). No remaining work.

NOTE for review (carried from 2d8271c): ADR-0065's "2x CPU-offload win"
headline for GQA decode opt2 may want a refresh to the post-D8 ~1.87x.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-17 15:12:31 -07:00
ywkang 821bbf26a2 gemm(bench): composite-vs-async dispatch comparison harness + single-op-model sweep 
Adds the user-orchestrated async GEMM variants used in paper §3.4 to
contrast with the composite (load_ref) path:
  - matmul-async              : naive (tl.load full A+B, one tl.dot)
  - matmul-async-chunked      : depth-inf prefetch (all B-tiles queued)
  - matmul-async-chunked-db   : depth-2 double-buffer (TCM-bounded)
plus scripts/paper/paper_plot_gemm_async_vs_composite.py (4-way per-PE
TFLOPS sweep + figure) and the regenerated outputs under
1H_milestone_output/gemm/.

Sweep regenerated under the ADR-0064 D8 single-op cost model (commit
2d8271c). At K=3072 the composite-vs-async-tiled gap narrows from the
pre-D8 ~6.3x to ~2.8x (composite 7.18 / async-tiled 2.54 TFLOPS): D8
makes the async-tiled kernel's ~191 single-op dispatches 5x cheaper, so
its dispatch overhead drops to ~1.5us. Qualitative order persists:
composite > async-full (3.91) > async-tiled (2.54).

test_bench_registry.py: register matmul-async / -chunked / -chunked-db
(also picks up the earlier milestone-gqa-headline -> milestone-1h-gqa
rename already present in the tree).

--- Remaining work (resume here if interrupted) ---
- Paper §3.4 (03-gemm.tex sec:gemm-vs-async): finish naive->async-full /
  chunked->async-tiled rename AND reframe "FIXED_DMA=8 for DMA descriptors"
  to single-op(8) vs composite(40). Figure already copied to
  docs/report/1H-codesign-paper/figures/gemm_composite_vs_async_tflops.png.
- Paper §3.4 K=3072 para + mechanism #3: update dispatch breakdown to new
  model (191 single-op * 8c ~= 1.5us, was 4.6us) and headline gap
  ~6.3x -> ~2.8x.
- Rebuild build/main.pdf (tectonic) + verify via pdftotext.
- Then commit Group 3 (paper + diagrams + PDF).
- Table 2 / §2 prose / ADR-0064 D8 already done in 2d8271c.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-17 15:06:36 -07:00
ywkang 2d8271c981 perf(cost-model): D8 single-op-cmd fast-path (FIXED=8 single-op / 40 composite) 
ADR-0064 D8 broadened from "DMA fast-path" to "single-op-cmd fast-path":
every single-op command (DmaRead/DmaWrite/Gemm/Math/Copy) now pays the
lighter FIXED=8; only CompositeCmd keeps the 40-cycle control-path FIXED
(it alone needs scheduler plan generation + per-tile RW-hazard tracking +
completion wiring). Renamed knob fixed_per_dma_cmd_cycles ->
fixed_per_single_op_cmd_cycles. dispatch_cycles now branches on
"is CompositeCmd" rather than enumerating DMA types.

Term choice: "single-op" (not "atomic", which read as sync/async) — the
axis is composition (one engine op vs fused multi-op plan), orthogonal to
timing. single-op <-> composite.

Tests: test_pe_cost_model.py updated to the single-op surface (defaults,
fast-path over all 5 single-op cmd types, composite general path, yaml
override). All green.

Recalibrated tests/attention/test_gqa_decode_opt2.py
::test_opt3_dispatch_exceeds_opt2 — NOT a regression: D8 makes single-op
cmds 5x cheaper, so opt2's two-composite fusion win over opt3's many
single-ops narrowed from pre-D8 ~3.7x to ~1.87x (opt3=224 > opt2=120).
The CPU-offload invariant (opt2 cheaper) still holds; only the model-
dependent ">2x" constant was over-fit to the old uniform-40 model. Gate
now: direction + >1.5x margin (matches sibling R-sweep test's stated
"absolute ratio informative-only" philosophy).
NOTE for review: ADR-0065's "2x CPU-offload win" headline may want a
refresh to reflect the post-D8 ~1.87x — left to user (architectural doc).

Full regression: 826 passed, 1 skipped (tests/ excl. tests/gemm).

--- Remaining work (resume here if interrupted) ---
5. Re-run scripts/paper/paper_plot_gemm_async_vs_composite.py with new
   cost model; verify async-tiled dispatch overhead drops (~4576ns ->
   ~1536ns expected) and the composite-vs-async-tiled gap narrows from
   the prior ~6.3x at K=3072.
6. Copy regenerated gemm_composite_vs_async_tflops.png to
   docs/report/1H-codesign-paper/figures/.
7. Paper §3.4 (03-gemm.tex sec:gemm-vs-async): finish naive->async-full /
   chunked->async-tiled rename AND reframe FIXED_DMA wording to single-op
   vs composite (currently still says "lighter FIXED for DMA descriptors,
   FIXED_DMA=8"). Table 2 (02-platform) + §2 dispatch prose already done.
8. Paper §3.4 K=3072 corner para + mechanism #3: update dispatch breakdown
   to new model (96 DMA*8 + 95 single-op*8 ≈ 1.5us vs old 4.6us); update
   headline ratio if it changed.
9. Rebuild docs/report/1H-codesign-paper/build/main.pdf (tectonic) +
   verify via pdftotext.
10. Then this is the bench-harness + paper commits (Groups 2 & 3).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-17 15:04:53 -07:00
ywkang 0c6ca0aaed paper(gemm): inline §3.1 pipeline-stage chain; drop display equation 
The five-stage chain was set as a centered display equation, which
overflowed the column and clipped the last 'DMA_WRITE' label. The
inline prose right below it also repeated the same chain in a
shorter form. Consolidate into one inline parenthetical listing the
five stage names, and reuse 'flows through the chain' for the
self-routing description — no clipping, no repetition.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-17 10:40:23 -07:00
ywkang b64c43b947 paper(gemm): drop in-chart titles + subtitles, let LaTeX caption carry that 
The matplotlib title and gray subtitle were duplicating what the
figure caption already says in the LaTeX source, and crowding the
top of every chart. Drop both — chart now leaves the title/subtitle
slots blank and the caption is the single source of truth for
chart titling.

Re-render all 5 GEMM figures (mac util measured, mac util theoretical
vs measured, stage breakdown, HBM BW util, per-PE TFLOPS) and resync
into docs/report/1H-codesign-paper/figures/. No data changes.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-17 10:29:13 -07:00
ywkang e8b7f4f064 paper(gemm): per-PE HBM BW + TFLOPS evaluation plots; unify peak to 8 TFLOP/s 
Two new evaluation plots for §3 GEMM, both using the composite window
as the denominator (the up-front tl.load of A in load_ref is therefore
excluded — only HBM traffic and compute inside the composite count):

- gemm_hbm_bw_util.png: per-PE HBM bandwidth utilization against the
  256 GB/s per-PE ceiling. Answers "is this workload memory-bound on
  this configuration?" — small/single-tile shapes sit at 22–46%
  (pipeline-fill-limited), large or output-write-dominated shapes
  saturate (≥85%).

- gemm_per_pe_tflops.png: per-PE GEMM throughput against the 8 TFLOP/s
  engine peak. The deep-K K=3072 load_ref case lands at
  7.18 TFLOP/s (~90% of peak) — the configuration's clean win.
  ref_ref at the same shape drops to 3.83 TFLOP/s because the second
  operand doubles HBM pressure and the kernel hits the BW ceiling.
  The variant gap on this chart is the operational cost of NOT
  pre-staging the activation.

Both charts compare two operand-staging variants on every shape:
load_ref ("activation pre-staged", weight-only HBM streaming) and
ref_ref (both A and W streamed from HBM). load_load is omitted —
with both operands pre-staged the composite carries no HBM traffic
and the BW metric collapses to 0.

Analytic / measured peak unified to the single hardware spec
(8.0 TFLOP/s = 8000 flops/ns). T_STAGE is now derived as
tile_flops / peak (= 16.384 ns for the 32×64×32 tile) rather than
the previous hardcoded 16 ns, so MAC efficiency and per-PE TFLOPS%
share the same denominator and never disagree. Side effect: max
analytic-vs-measured gap tightens from 2.2 ppt to 1.4 ppt across
the sweep.

§3 Results restructured: lead with the two new evaluation plots
(BW saturation, achieved throughput), then the MAC-utilization
analytic-vs-measured chart serves as the validation step, then the
per-stage engine wall-clock as the supporting diagnostic. §3
Analysis updated to use the load_ref / ref_ref contrast as the
hardware-software boundary line that motivates the GQA kernel of §5.

§2.4 Accuracy claim updated to match: GEMM analytic-vs-measured
agreement is now "within 1.4 ppt across every swept shape", from
~7.7% at single-tile up to ~90% at K=3072. Subtitle/title vertical
stacking glitch in the matplotlib charts fixed in passing.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-17 10:24:42 -07:00
ywkang 23992548f7 gemm(perf): fix back-to-back DMA pipelining + tighten analytic model 
Three coupled fixes that recover small-tile GEMM pipeline efficiency
from 53% to 88% (32x3072x32 load_ref, composite_window basis).

1. PE_DMA channel-hold (ADR-0014 D4 clarified): both the
   _handle_with_hooks (PeInternalTxn) and _pipeline_process
   (TileToken) paths used to hold the cap=1 DMA channel through the
   full HBM round-trip, which double-serialized with the HBM_CTRL's
   own per-PC `available_at` model and prevented back-to-back tile
   DMAs from amortizing their per-request head latency. Channel is
   now released after the request is enqueued onto the next hop;
   HBM serialization is HBM_CTRL's responsibility alone.

   Tests: new test_pe_dma_back_to_back_pipelining as the oracle
   (asserts wall < 75% of strict-serialized N x single_op). Existing
   test_pe_dma_record_start_after_channel_acquire rewritten to assert
   t_start clustering (channel released fast) instead of the old
   round-trip-hold invariant. test_pe_dma_same_channel_serializes
   still passes — HBM_CTRL preserves ordering.

   Probe regression: PE→local-HBM 32 KiB stays at 141 ns
   (single-request, unaffected).

2. milestone_1h_gemm bench: matmul_composite was reading MATMUL_M/K/N
   env vars at module load, so every sweep row replayed the cached
   256³ result; values now read inside run(). Drops the stale
   sys.modules deletion hack.

3. Analytic ideal-pipeline model: dropped the (n_mn-1)·dma_w_per_pair
   penalty (over-pessimistic for under-tile shapes — it pushed
   measured > theoretical) and replaced the D_STAGES-derived head
   with empirical T_PIPELINE_FILL=60 ns / T_PIPELINE_TAIL=30 ns.
   Max analytic-vs-measured gap across all 7 swept shapes now 2.2 ppt
   (was 9-44 ppt under the old constants).

Paper updates:
- §3 (GEMM): 78%→88% measured at 48 tiles, 23%→15% at 1 tile,
  stage breakdown numbers refreshed (DMA in / Fetch / GEMM all
  ~785 ns at K=3072), analytic-vs-measured agreement tightened
  to "within 2.2 ppt".
- §2.4 (Accuracy): GEMM tracking claim refreshed accordingly.
- §5 (GQA): restore long-ctx 4-cases figures into figures/
  (they were dropped from bench output dir as derived artifacts in
  92b9221 / e45626c but §5 still cites them by name).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-17 09:56:05 -07:00
mukesh 92b9221533 paper(gqa): drop committed long-ctx 4-cases derived chart PNGs 
These 8 prefill/decode 4-case PNGs are derived artifacts regenerated
on demand from scripts/paper/paper_plot_gqa_{prefill,decode}_long_ctx_4cases.py
and no longer belong in the committed bench output.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-17 09:39:26 -07:00
mukesh 694e0cc9b9 paper(gqa): add 6-case per-PE memory + comm summary chart 
Analytical plot generator + PNG covering the 4 sharding cases from
slide 17 plus the 2 d_head-TP variants (Cases 4, 5). Three panels:
per-PE HBM budget breakdown (Wq + Wk + Wv + Wo + 4.24 GB KV
headroom), per-PE KV memory at S_kv=1M, and per-PE communication
per output token (decode). Numbers match slide 17 (max KV context
~7.1M for the 64-way sharded cases).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-16 16:12:26 -07:00
ywkang 6b6e29968a paper: §2 platform deep-edit + §1 GQA framing + figure regen 
§2 (KernBench Platform):
- Fix Table 2 HBM aggregate BW (1024 → 2048 GB/s); drop stale
  hbm_total_bw_gbs from topology.yaml (never read by sim_engine)
- Split PE_CPU / PE_SCHED fixed-cost row; disambiguate from the
  40-cycle command-dispatch FIXED term
- §2.2 two-pass: expand to describe Pass 1 timing and Pass 2 data
  data-correctness path
- §2.3 dispatch model: add motivation sentence for the
  descriptor-size linear form
- §2.4 Accuracy: reorder GEMM → All-reduce → Probe →
  Simplifications; drop FSIM aside (already covered by §4 Fig 5
  caption); soften 'every ns' → 'every modeled latency
  contribution'
- §2.5 HW config: add 64 TFLOP/s vs 2048 GB/s (~31 FLOP/byte)
  balance-point intuition and forward pointer to §5 GQA decode
- Fig 1 caption: separate illustrative topology from experimental
  configuration

§1: tighten GQA-as-primary-bandwidth-bottleneck framing.

Figures: regenerate SIP / CUBE architecture (SVG sources + PDF +
generator scripts).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-16 16:00:23 -07:00
mukesh 77eece81c4 gqa(milestone-1h-gqa): add umbrella bench + first sweep outputs 
Single @bench entry that drives the prefill + decode long-context
4-cases sweeps in one invocation, mirroring milestone-1h-gemm. The
individual gqa_helpers/long_ctx/gqa_{prefill,decode}_long_ctx_4cases
modules are now pure helpers (no @bench), reached only via this
umbrella.

Env var contract:
  GQA_1H_RUN=1                    (required gate)
  GQA_1H_SWEEPS=prefill,decode    (default: both; pick subset to skip)
  GQA_1H_TOPOLOGY=topology.yaml   (override)

Output layout:
  benches/1H_milestone_output/gqa/long_ctx/
    sweep_prefill.json
    sweep_decode.json
    gqa_prefill_long_ctx_4cases_{latency,traffic,memory,parallelism}.png
    gqa_decode_long_ctx_4cases_{latency,traffic,memory,parallelism}.png

The decode bench config drops S_kv from 131_072 to 8_192 so the umbrella
finishes in minutes — the comparative-story bars are the same shape at
8K. The 131K production headline number is recoverable by overriding
the dispatch; the helper docstring notes this.

Outputs (2 sweep JSONs + 8 PNGs) ship in this commit alongside the
umbrella that produced them, following the milestone-1h-gemm pattern
where derived artifacts live with the bench that emits them.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-16 13:06:33 -07:00
mukesh 443ede99c7 gqa(prefill-4cases): add long-context prefill 4-cases comparative study 
Mirrors the existing decode 4-cases study for prefill on the LLaMA-3.1-70B
single-KV-head group (h_q=8, h_kv=1, d_head=128) across 8 cubes × 8 PEs:

  Case 1  Cube-SP × PE-TP  → KV S_kv-Ring across cubes; Q/O T_q-split 64-way
  Case 2  Cube-Repl × PE-TP → full KV per cube; only CUBE 0's P PEs split T_q
  Case 3  Cube-Repl × PE-SP → full KV per cube; PEs SP on S_kv (intra-cube AR)
  Case 4  Cube-SP × PE-SP  → KV split 64-way + Q T_q-split across cubes
                              (Ring KV + per-cube intra-CUBE AR)  ★ optimal

Cases 2 and 3 use Q-axis tiling (TILE_Q) so the per-PE scratch is bounded
by TILE_Q × TILE_S_KV regardless of T_q (the full Q tensor would be
G·T_q·d_head·2 bytes which scales linearly with T_q and blew the 1 MB
budget at T_q≥128 without tiling).

Per-panel try/except in run_sweep tolerates per-panel failures so
sweep.json always lands with whichever cases succeed plus a failures
list. Case 4 uses configure_sfr_intercube_ring with the snake submesh
(2,4) which installs both the Ring E/W lanes and the intra_* lanes
needed for the intra-CUBE reduce — single SFR config for all 4 cases.

The plot script generates 4 PNGs (latency, traffic, memory, parallelism)
into 1H_milestone_output/gqa/long_ctx/. The parallelism chart shows
total compute work (active_PE × T_q_per_PE × S_kv_processed) — exposes
Case 3's 8× redundancy vs. Cases 1/2/4 which all do unique work.

gqa_prefill_long_ctx_4cases.py exposes run_sweep() rather than a @bench
decorator — the umbrella bench in the next commit invokes it.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-16 13:06:08 -07:00
mukesh e45626c036 gqa: reorganize benches into gqa_helpers/ subpackage; drop legacy headline 
Splits the GQA helpers into a dedicated subpackage to make room for the
prefill 4-cases study (next commit) and a single umbrella bench
(milestone-1h-gqa, after that).

Layout:
  benches/gqa_helpers/
    long_ctx/    — decode 4-cases kernels + sweep runner
    short_ctx/   — prefill/decode short-context kernels
    shared/      — _gqa_panel_helpers + decode_opt2 (context-agnostic)

The registry audit now skips subpackages so gqa_helpers/ (without a
leading underscore) doesn't get audited for @bench decorators.

Also drops the legacy milestone-gqa-headline bench, its
_gqa_attention_prefill_long kernel, 6 dependent prefill tests, the
paper_gqa_latency.py report harness, and the 3 stale headline-derived
PNGs the §6 wire-up referenced (paper will re-pull from the new
1H_milestone_output/gqa/long_ctx/ once §6 is updated).

The _ccl_cfg and _summarize_op_log helpers used to live in the
headline bench; extracted them to gqa_helpers/shared/_gqa_panel_helpers.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-16 13:05:41 -07:00
ywkang 359a0eaa44 paper(platform): Table 1 -> table*, shrink Fig 1 row 1, move Fig 1/2 source up 
- Probe table (Table 1, tab:probe-pe-dma) promoted to table* so it
  spans both columns and no longer gets clipped at the column edge.
- Fig 1 (hw-arch) row 1 (SIP, CUBE subfigures) trimmed: widths
  0.495 -> 0.42 textwidth and a height=0.6\linewidth cap added so
  the figure occupies less vertical space and floats can land
  earlier in the document.
- Fig 1 and Fig 2 source blocks relocated to the top of §2 so they
  enter LaTeX's float queue before any §2 body text. As a result
  Fig 1 now lands on page 3 and Fig 2 on page 4 (was 4 and 5).
- Table 2 (tab:hw, modeled hardware config) reverted to plain
  \begin{table}; the previous \begin{table*} reshape was not
  requested.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-16 11:58:48 -07:00
ywkang 1cf8dd7868 paper(platform): widen SIP/CUBE subfigures, drop height cap 
- subfig widths 0.46 -> 0.495 textwidth so the \hfill gap between
  (a) SIP and (b) CUBE shrinks to ~2% textwidth instead of ~8%.
- Removed the height=0.85\linewidth cap on both subfigures; with the
  cap, the (roughly 1:1) SIP and CUBE diagrams were being scaled
  down to 85% even when the column had room. Let them render at
  their natural aspect ratio.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-16 11:22:01 -07:00
ywkang b715002b5a paper: §2 figure layout + Accuracy subsection + Table 1 → 2-col; §6 trim historical 4-panel content 
§2 platform:
- Accuracy promoted from \paragraph to \subsection (sec:accuracy);
  it sits at the same heading level as Why / Device / Latency / HW
  config and so reads as a first-class component of the platform
  description rather than a tail-end footnote.
- Three architecture diagrams (SIP, CUBE, PE) collapsed into one
  figure* using subcaption: Fig.~\ref{fig:hw-arch}(a) SIP and (b)
  CUBE side-by-side on row 1, (c) PE wide on row 2, total height
  capped at ~40% of the page. Subfigure cross-references rewritten
  to Fig.~ref{fig:hw-arch}\subref{...} in the body text. main.tex
  now pulls in the subcaption package.
- Table 1 (modeled hardware configuration) promoted to table*
  (two-column / full page width) and the row layout rewritten as a
  4-column tabular so Hierarchy + PE + Command-issue sit on the
  left and Memory + Interconnect sit on the right. The previous
  single-column rendering was getting cut at the right edge of the
  printed column.

§6 GQA:
- Removed the historical four-panel headline table (tab:gqa), the
  two figures (fig:gqa-lat, fig:gqa-break), and the prose paragraph
  that cited their per-panel numbers. Reason: the underlying
  milestone-gqa-headline bench has been simplified in collaborator
  commit 65c365f ("drop misleading single_user_/multi_user_
  panels") and no longer reproduces that data, so the section was
  left referencing a dataset the current bench cannot regenerate.
- The previously-added 4-cases long-context decode comparison is
  now §6's only Results subsection, retitled to
  "Results: long-context decode and parallelism strategies".

The headline GQA result that survives into the paper is therefore
the parallelism trade-off study (Case 4 / Cube-SP × PE-SP at 34 us
buying an 8x KV-memory reduction over the latency leaders).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-16 11:16:16 -07:00
ywkang 22f3968b51 paper(gqa): wire the 4-cases long-context decode comparison into §6 
The collaborator commit 7c346de added the comparative figures
(gqa_decode_long_ctx_4cases_{latency,memory,traffic}.png) but the
fused-GQA section still only referenced the four headline panels.
This commit closes that loop:

- New §6 subsection "Long-context decode: parallelism strategies"
  added between Results and Analysis. It lays out the four
  parallelism strategies (Cube-SP/Repl x PE-TP/SP) and pulls in the
  three new figures.
- The discussion highlights the central trade: the fastest strategy
  (Case 3, Cube-Repl x PE-SP at 20.2 us) requires replicating the
  full KV cache to every CUBE, while Case 4 (Cube-SP x PE-SP, the
  chosen design marked *) gives back ~14 us in exchange for an 8x
  KV-memory reduction.  Case 4's ~190 IPCQ copies + ~190 DMA reads
  are precisely the on-device collective traffic PE_IPCQ and the
  torus links of §5 are provisioned to absorb -- a direct payoff
  of the communication-side codesign work.
- Connects back to §5 (PE_IPCQ / all-reduce) so the reader sees the
  capstone arc: the GEMM enabler exposes the data-movement bound,
  the communication enabler attacks it, and the long-context
  parallelism study shows how the choice between the two extremes is
  framed by KV memory vs. on-device collective traffic.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-16 10:44:32 -07:00
ywkang e3f08972da paper: §2 polish + §5 fresh all-reduce data + kernbench probe table 
§2 platform:
- Add a "memory-centric AHBM" sentence to the device-and-execution
  intro: each PE is paired with dedicated HBM bandwidth and on-PE
  TCM, so the performance question is about feeding compute from
  locally-attached memory + moving the unavoidable inter-PE traffic
  efficiently. Makes the AHBM character of the platform visible
  before §2.3 starts unpacking the latency model.
- Promote "Congestion and contention modeling" from \paragraph to
  \subsubsection: this is the platform's core differentiator over a
  peak-BW roofline, so it deserves its own heading.
- Rename "Control-plane (issue) cost model" to "Command dispatch
  overhead model" -- describes what it actually charges (the PE_CPU
  paying a fixed + per-byte cost to push a command to one of the
  accelerator engines) without the more abstract "control-plane"
  framing.
- Add a third Accuracy cross-check from kernbench probe: a
  PE→HBM DMA-distance sweep that confirms monotonic hop progression,
  bandwidth saturation matching the per-edge model, and the built-in
  invariants (D2H >= H2D, cross-CUBE best < worst). New
  Table~\ref{tab:probe-pe-dma} reports per-traversal latency and
  utilisation at 32KiB / 1MiB across five hop classes.
- Captured probe output as figures/probe_pe_dma_summary.txt for
  reproducibility.

§5 PE_IPCQ / all-reduce:
- Re-ran milestone-1h-ccl on current sim_engine (post-ADR-0064 Rev2
  and the IPCQ slot-wrap Phase-2 race fix). Updated the topology-
  comparison table and the buffer-kind caption to the fresh raw
  latencies. Headline ratios are preserved:
    torus vs mesh saving at 96 KB/PE: 24.8% (was 24.8%) -> ~25%
    torus vs ring saving at 96 KB/PE: 19.3% (was 19.3%) -> ~19%
    TCM vs HBM saving at 64 KB/PE: 13.4% (was 13.9%) -> ~13%
    TCM vs SRAM saving at 64 KB/PE: 37.2% (was 38.3%) -> ~37%
  The Executive Summary's "up to ~14%" / "up to ~38%" framing stays
  consistent with these post values.

Artifacts refreshed in src/kernbench/benches/1H_milestone_output/:
- ccl/summary.csv + per-topology PNGs + buffer-kind CSV/PNG
- ccl/comparison_mesh_vs_ring_vs_2DTorus_vs_theoretical_vs_fsim.png
- gemm/* (re-run yields identical pe_window structure; PNGs
  refreshed)

Paper figures synced to the fresh artifacts:
- figures/allreduce_comparison.png
- figures/allreduce_buffer_kind.png

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-16 10:38:47 -07:00
mukesh 7c346dec1b paper(figures): GQA long-context decode 4-cases comparative figures 
Generated by scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py
from the bench's sweep.json (S_kv=8192, the headline 128K run was
sidestepped — Case 2's single-PE 128-tile sweep is the long pole
and shape/ordering is preserved across scales).

  gqa_decode_long_ctx_4cases_latency.png   end-to-end latency / case
  gqa_decode_long_ctx_4cases_traffic.png   ipcq / dma op-count breakdown
  gqa_decode_long_ctx_4cases_memory.png    KV bytes per cube (slide-11)

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 17:48:34 -07:00
mukesh 1fbe833992 gqa(decode-4cases): wire latency + engine occupancy into sweep.json; add comparative plot script (5C.F) 
Bench changes:
  - new _end_to_end_ns(op_log) and _engine_occupancy_ns(op_log) helpers
    in milestone_gqa_decode_long_ctx_4cases.py (mirror paper_gqa_latency.py)
  - _run_panel return dict now carries latency_ns + engine_occupancy_ns
    alongside op_log_summary, so sweep.json is the single source of truth
    for the comparative figures

Plot script:
  - new scripts/paper/paper_plot_gqa_decode_long_ctx_4cases.py reads
    sweep.json and emits 3 PNGs to docs/report/1H-codesign-paper/figures/:
      gqa_decode_long_ctx_4cases_latency.png   (end-to-end latency / case)
      gqa_decode_long_ctx_4cases_traffic.png   (ipcq/dma op counts / case)
      gqa_decode_long_ctx_4cases_memory.png    (KV bytes per cube / case)

Test changes:
  - 2 new tests verifying the helpers + _run_panel dict shape
  - lower smoke S_kv from 8192 -> 2048 (4x faster Case 2; assertions
    are S_kv-independent; one fold-loop iteration preserved)

18/18 tests pass in ~4 min.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 17:45:44 -07:00
ywkang 7f437a20bd paper(platform): edit-pass through §2 KernBench Platform done 
- Reordered §2 subsections to follow the SIP → CUBE → PE → graph
  flow: Why KernBench → Device and execution model → Latency model
  → Modeled hardware configuration. Readers now meet the device
  hierarchy before the graph abstraction that re-uses it.
- §2.2 Device and execution model: starts with the SIP/CUBE/PE
  hierarchy paragraphs (each anchoring fig:sip-arch, fig:cube-arch,
  fig:pe-arch); then the runtime-API/sim-engine/components bullet
  list; then the atomic-vs-composite command distinction (corrects
  the prior over-narrow framing that read every PE command as
  composite -- atomic single-engine commands exist too, and PE_CPU
  itself runs control-plane work directly).
- §2.3 Latency model: opens with the four-contribution decomposition
  (per-node overhead, per-edge transmission, drain, queuing delay)
  and the latency_model schematic; retains existing The hardware as
  a graph / From graph to DES / Latency contributions / Congestion
  / Control-plane cost model / Accuracy paragraphs. Accuracy
  paragraph now closes on KernBench's sufficiency for *relative*
  HW/SW design trade-offs given analytic + external-simulator
  agreement.
- New figures and assets:
  - figures/sip_architecture.pdf  (SIP-level graph view)
  - figures/cube_architecture.pdf (CUBE-level zoom-in)
  - figures/latency_model.png      (conceptual latency-model
                                   schematic with per-node /
                                   per-edge / drain / queuing-delay
                                   colour coding)
  - figures/pe_architecture.png    (carried over)
- Source-of-truth generator for the latency schematic:
  scripts/paper/paper_latency_model_diagram.py (a report-only
  harness under scripts/paper/ per the /paper isolation rule).
- main.tex preamble: \usepackage{tikz} added (kept from prior
  sequence-diagram draft -- harmless now that the latency model is
  a PNG; left in to keep paragraph numbering stable).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-15 17:17:19 -07:00
mukesh 756680f4e6 gqa: fold decode_long into the Case 4 kernel; drop 1D-chain dead code 
Option B fold: the lrab math previously in
``_gqa_attention_decode_long.py`` (used only as a thin re-export by
the Case 4 wrapper and as the import source for ``_merge_running``
in Cases 1 / 3) now lives inline in the Case 4 kernel file. ``sub_w``
is hardcoded to 4 (the 4×2 cube sub-mesh geometry; root cube 6);
the legacy ``sub_w=0`` 1D-chain backward-compat path is removed
along with its dedicated tests — the dropped ``single_user_*`` /
``multi_user_*`` panels that exercised it are already gone.

Production changes:
  - inline lrab math into _gqa_attention_decode_long_ctx_cube_sp_pe_sp.py
    (drops sub_w param; _ROOT_CUBE=6 baked in)
  - inline _merge_running into Cases 1 (cube_sp_pe_tp) and 3
    (cube_repl_pe_sp) so they no longer depend on the deleted file
  - delete src/kernbench/benches/_gqa_attention_decode_long.py
  - remove dead _run_decode_panel / _DECODE_* constants /
    decode-side _PANEL_DISPATCH / _make_bench_fn decode branch
    from milestone_gqa_headline.py (only the prefill panel remains)
  - update _gqa_attention_decode_opt2.py docstring reference

Test changes:
  - delete 7 legacy test_gqa_*.py files that pre-dated the 4-cases
    architectural split (coverage now subsumed by the 16 4-cases tests)
  - remove test_opt2_matches_opt3_data_mode + _run_decode_data helper
    from test_gqa_decode_opt2.py (the parity check required sub_w=0
    which no longer exists; opt2's other 4 tests preserved)

20/20 tests pass (16 4-cases + 4 opt2 smoke/dispatch).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 16:32:37 -07:00
mukesh ddee28a499 gqa(decode-4cases): rename bench/kernels/panels with long_ctx token 
The 4-cases comparative study is specifically about long-context
decode (LLaMA-3.1-70B, S_kv=128K, T_q=1); the long_ctx token in the
names makes the scope explicit and aligns with the existing
_gqa_attention_decode_long convention.

Renames (mechanical; no semantic change):
  - bench file:   milestone_gqa_decode_4cases.py
                  → milestone_gqa_decode_long_ctx_4cases.py
  - bench name:   milestone-gqa-decode-4cases
                  → milestone-gqa-decode-long-ctx-4cases
  - output dir:   1H_milestone_output/gqa_decode_4cases/
                  → 1H_milestone_output/gqa_decode_long_ctx_4cases/
  - env vars:     GQA_DECODE_4CASES_RUN / _TOPOLOGY
                  → GQA_DECODE_LONG_CTX_4CASES_RUN / _TOPOLOGY
  - 4 kernel files _gqa_attention_decode_<case>.py
                  → _gqa_attention_decode_long_ctx_<case>.py
  - 4 kernel functions gqa_attention_decode_<case>_kernel
                  → gqa_attention_decode_long_ctx_<case>_kernel
  - 4 dispatch kinds  decode_<case> → decode_long_ctx_<case>
  - 4 panel names     single_kv_group_decode_gqa_<case>
                      → single_kv_group_decode_long_ctx_gqa_<case>
  - 4 helper functions _run_decode_panel_<case>
                       → _run_decode_panel_long_ctx_<case>
  - test file renamed in lockstep

Also: Case 4 smoke test now uses its case-specific helper
_run_decode_panel_long_ctx_cube_sp_pe_sp (consistent with Cases 1-3)
instead of the legacy _run_decode_panel from milestone_gqa_headline.
16 tests pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 16:09:21 -07:00
mukesh 3c155be8e6 gqa(decode-4cases): give Case 4 its own thin kernel file for naming symmetry 
Cases 1-3 each have a dedicated _gqa_attention_decode_<case>.py
kernel file; Case 4 previously reached into _gqa_attention_decode_long.py
via the headline bench's _run_decode_panel helper, breaking the
one-file-per-case convention. Adds _gqa_attention_decode_cube_sp_pe_sp.py
as a 20-line wrapper that bakes in sub_w=4 (the C=8 lrab geometry)
and gives Case 4 its own kind ("decode_cube_sp_pe_sp") and helper
(_run_decode_panel_cube_sp_pe_sp). decode_long.py is unchanged
(still serves the legacy decode_long tests). 16 tests pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 15:58:11 -07:00
mukesh 0ef4fde5d8 gqa(decode-4cases): Case 1 — Cube-SP × PE-TP (inter-CUBE lrab only; PE-TP B=1 waste) (5C.A) 
Adds the fourth and final 4-cases panel: KV split S_kv-wise across
the 8 cubes (cube=row_wise, S_local = S_kv/C), replicated within
each cube (pe=replicate). PE-TP at B=1 means only PE 0 of each cube
has work; PEs 1-7 early-return (slide-11 PE-TP-at-B=1 waste). No
intra-CUBE comm; inter-CUBE 8-way reduce reuses the lrab-adapted
center-root pattern (root cube 6) — same structural cost as Case 4's
inter-CUBE phase (21 ipcq_copy). 16 tests pass (4 per case).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 15:30:00 -07:00
mukesh ae942f6959 gqa(decode-4cases): Case 3 — Cube-Repl × PE-SP (intra-CUBE AR only; 8× memory) (5C.C) 
Adds the third 4-cases panel: KV replicated per cube (8× memory
waste), PEs SP on S_kv within each cube, intra-CUBE 8-way reduce on
(m, ℓ, O), and no inter-CUBE comm (every cube ends with full answer;
designated writer = cube 0). Reuses the row-chain + col-bridge
intra-CUBE pattern that anchors Case 4 (21 ipcq_copy per cube × 8
cubes = 168 total). 12 tests pass (4 Case 4 + 4 Case 2 + 4 Case 3).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 15:18:41 -07:00
mukesh 5672c8f3ef gqa(decode-4cases): Case 2 — Cube-Repl × PE-TP (no comm; 8× memory) (5C.B) 
Second case in the GQA decode 4-cases comparative study per
GQA_full_deck.pptx slide 11. Case 2 replicates K, V across all 8
cubes × 8 PEs (the slide-11 8 KB/tok/PE memory waste) and has zero
inter-rank comm. For B=1 (single-user decode default), only PE 0
of CUBE 0 has work — the inherent PE-TP waste slide 11 calls out.

Changes:
- New kernel: src/kernbench/benches/_gqa_attention_decode_cube_repl_pe_tp.py
  Simplest of the 4 cases. Active rank loads full Q/K/V from HBM,
  computes attention via S_kv tile sweep with online-softmax merge,
  writes O. All non-(0,0) ranks early-return. No tl.send/recv.
- src/kernbench/benches/milestone_gqa_decode_4cases.py:
    - Add panel single_kv_group_decode_gqa_cube_repl_pe_tp (Case 2)
      to _PANELS and _PANEL_DISPATCH.
    - Add _run_decode_panel_cube_repl_pe_tp helper: DPPolicy K/V/Q/O
      = cube=replicate, pe=replicate (models 8× memory waste).
    - Extend _make_bench_fn to dispatch kind="decode_cube_repl_pe_tp"
      to the new runner.
- tests/attention/test_milestone_gqa_decode_4cases.py:
    4 new tests assert Case 2 contract: panel registered, smoke
    completion, zero ipcq_copy (no comm), single dma_write from cube 0.

Verification: 8/8 tests pass (4 Case 4 anchor + 4 new Case 2).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 14:59:58 -07:00
mukesh 65c365f858 bench(milestone-gqa-headline): drop misleading single_user_/multi_user_ panels 
The legacy panel names suggested batched serving semantics they never
had — all four modeled a single user with KV sharded differently
(C=1 single-cube; C=4 multi-cube Cube-SP), at toy dims (T_q=4, S_kv≤128).
The single-KV-group C=8 panel + the new milestone-gqa-decode-4cases
bench cover the meaningful comparisons; pytest regression already
covers C=1/C=4 configurations end-to-end at richer scale.

Changes:
- milestone_gqa_headline.py: drop the 4 legacy panels; _PANELS now
  contains only single_kv_group_prefill_gqa_c8_p8. Update docstring.
- tests/attention/test_milestone_gqa_headline.py: drop the 3 legacy-
  panel architectural tests (Ring-KV traffic, root-only decode write,
  per-CUBE distributed output) and test_decode_panels_use_real_gqa
  (no decode panels in this bench anymore). Equivalent properties
  are asserted in test_milestone_gqa_single_kv_group_prefill_panel.py
  (64 dma_writes, 896 ipcq_copy) and test_milestone_gqa_decode_4cases.py
  (1 dma_write at cube 6, 189 ipcq_copy).
- tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py:
  drop test_existing_prefill_panel_runner_backward_compat (it exercised
  multi_user_prefill_gqa which no longer exists).
- scripts/paper/paper_plot_gqa.py: replace the 4 legacy _LABELS entries
  with the single single_kv_group_prefill_gqa_c8_p8 label.
- Regenerate 1H_milestone_output/gqa_headline/sweep.json from the new
  panel set.

Verification: 9/9 tests pass across the 3 affected test files.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 14:50:25 -07:00
mukesh c164645aee gqa(decode-4cases): Case 4 anchor in dedicated bench (5C.D) 
First milestone of the decode 4-cases comparative study per
GQA_full_deck.pptx slides 11-17. Case 4 (Cube-SP × PE-SP, the optimal
case per slide 11) is structurally the existing _gqa_attention_decode_long
at sub_w=4 (Increment 2's lrab-adapted center-root reduce). This commit
wires it into a dedicated bench so the remaining cases land alongside.

Changes:
- New bench: src/kernbench/benches/milestone_gqa_decode_4cases.py
  Houses the 4 case panels under one milestone-gqa-decode-4cases entry
  (gated by GQA_DECODE_4CASES_RUN=1; output to
   1H_milestone_output/gqa_decode_4cases/sweep.json). Cases 1-3 are
  TBD in subsequent sub-increments (5C.A/B/C).
- New panel: single_kv_group_decode_gqa_cube_sp_pe_sp
  C=8, P=8, sub_w=4, T_q=1, S_kv=131_072, d_head=128, h_q=8, h_kv=1.
- src/kernbench/benches/milestone_gqa_headline.py: _run_decode_panel
  extended with keyword-only sub_w/T_q/d_head/h_q/h_kv overrides
  (defaults preserve existing-panel behaviour).
- tests/attention/test_milestone_gqa_decode_4cases.py: 4 new tests
  asserting registration, smoke completion, reduce-to-root at the lrab
  center cube (cube 6), and the predicted 189-ipcq Case-4 traffic
  pattern (168 intra-CUBE + 21 inter-CUBE lrab Phase 1+2).
- tests/attention/test_milestone_gqa_headline.py: rename
  test_sweep_json_has_four_panels -> test_sweep_json_has_expected_panels
  and switch hardcoded 4 to len(PANELS) (the panel set grew to 5
  with Increment 5's single_kv_group_prefill_gqa_c8_p8).

Deviation noted: slide 13 prescribes AllReduce on (m,ℓ,O); our kernel
does reduce-to-root (only the lrab center cube has the answer) per
ADR-0060 §4. Treated as the kernbench Case-4 baseline.

Verification: all 4 new tests pass; 90 regression tests pass; the
previously-failing test_sweep_json_has_four_panels now passes under
its renamed form.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 14:31:44 -07:00
mukesh 5b4d9cb597 bench(milestone-gqa-headline): scale single_kv_group prefill panel T_q=S_kv=1K (scratch budget) 
The headline T_q=S_kv=32K target overflows the 1 MB per-PE scratch
pool: at T_q_local=4K and tile_s=1024 the scores matrix alone is 8 MB.
The prefill kernel's bootstrap section also leaves K_t/V_t/scores/
exp_scores persistent (outside tl.scratch_scope), inflating baseline.

Scale-down to T_q=S_kv=1K (T_q_local=128, fits comfortably) preserves
the C=8 + P=8 architecture demonstration; the true LLaMA 32K headline
awaits a future increment to add Q-axis tiling and tighten bootstrap
scratch discipline.

Verified end-to-end: kernbench run --bench milestone-gqa-headline now
produces sweep.json with all 5 panels. The new panel shows
ipcq_copy=896 (matches (C-1)·n_tiles·2·C·P = 7·1·2·8·8) and
dma_write=64 (one per PE, head-parallel + intra-CUBE PE-SP).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 14:00:30 -07:00
mukesh 9e1242039b gqa: single-KV-group LLaMA-3.1-70B prefill milestone (Increments 1-5) 
End-to-end wires C=8 P=8 d_head=128 prefill with snake-ring inter-CUBE
SFR, intra-CUBE PE-SP (all 64 ranks active), and the milestone bench
panel. Decode kernel gains lrab-adapted center-root reduce for the
2×4 sub-mesh per ADR-0060 §4.2.

Increment 1 — SFR multi-row snake
  src/kernbench/ccl/sfr_config.py: configure_sfr_intercube_ring gains
  submesh_shape / submesh_origin kwargs; installs a Hamiltonian snake
  ring through a rectangular sub-mesh (every hop is 1-hop physical
  neighbour). Backward-compat: 1D-row behaviour preserved when
  submesh_shape is None.
  tests/test_intercube_snake_ring.py (12 tests)

Increment 2 — Decode lrab-adapted center-root reduce
  src/kernbench/benches/_gqa_attention_decode_long.py: new sub_w param
  (default 0 = existing 1D-chain). sub_w >= 2 selects the ADR-0060
  §4.2 prescribed lrab-adapted Phase 1+2 reduce (bidirectional row +
  bidirectional col converge to the center cube), with log-sum-exp
  _merge_running replacing the plain + of lrab.
  tests/attention/test_gqa_decode_long_2d_reduce.py (4 tests)

Increment 3 — Prefill kernel at C=8 (no production change)
  Verified by inspection that the existing prefill_long kernel +
  Increment 1's snake SFR already work at C=8 without any kernel
  edit. The kernel speaks logical W/E; the snake routes it.
  tests/attention/test_gqa_prefill_long_c8_snake.py (3 tests)

Increment 4 — Intra-CUBE PE-SP in prefill (all 64 ranks)
  src/kernbench/benches/_gqa_attention_prefill_long.py: new P param
  (default 1 = existing PE-0-only). P > 1 splits T_q query-axis-wise
  across the P PEs of each CUBE; output rows are disjoint per PE so
  no intra-CUBE reduce is needed; each PE drives its own same-lane
  ring (P parallel rings).
  tests/attention/test_gqa_prefill_long_pe_sp.py (5 tests)

Increment 5 — LLaMA-scale milestone bench panel
  src/kernbench/benches/milestone_gqa_headline.py: new panel
  single_kv_group_prefill_gqa_c8_p8 (C=8, P=8, T_q=S_kv=32K,
  d_head=128). _run_prefill_panel extended with P/T_q/d_head
  defaults; routes snake SFR when C > mesh_w.
  tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py (3 tests)

Total: 4 production files modified, 5 new test files, 27 new tests.
Followed the Phase 1/2 protocol per CLAUDE.md throughout.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-15 13:42:31 -07:00
ywkang 9270d3435a paper(platform): WIP — §2 KernBench Platform mid-edit 
Mid-pass through §2. Captured progress so far:

- §2 intro previews the three new threads (graph view, DES engine,
  congestion) the rewritten latency section weaves together.
- §2.3 retitled to "Latency model: a graph traversed by events" and
  restructured into bold-led paragraphs:
  - The hardware as a graph (nodes = components, edges = links).
  - From graph to discrete-event simulation (node/edge events,
    deterministic ordering, correlation-ID trace).
  - Latency contributions (per-node fixed, per-edge size-aware,
    per-service occupancy).
  - Congestion: per-edge FIFO BW occupancy, HBM per-PC parallelism,
    component serial workers — the mechanisms that surface real
    bottlenecks instead of peak-BW roofline.
  - Control-plane cost model (FIXED + b·R) — unchanged.
  - Accuracy: extended with a second cross-check from the all-reduce
    study (torus vs. analytic startup-plus-per-packet model + FSIM
    external single-device reference). Existing GEMM analytic-vs-
    measured 10-20% check retained.
- §2.4 Hardware-configuration table unchanged.

Still TODO: cross-check accuracy numbers against a fresh
milestone-1h-gemm/ccl re-run; the current artifacts predate the
ADR-0064 Rev2 cost-model and IPCQ Phase-2 race fix and may need a
refresh before §2.3 final lock-in.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-11 16:31:42 -07:00
ywkang d5267ec718 paper(intro): edit-pass through §1 Introduction done 
- §1 Introduction rewritten with AHBM-first opening, in-paper KernBench
  platform overview, GQA motivation and three architectural requirements,
  three matching HW-SW co-design mechanisms (GQA-aware placement,
  PE_IPCQ, composite-command GEMM under PE_SCHEDULER), broader-
  applicability outlook (PE_IPCQ as collective substrate, composite
  command for fused operator pipelines), and outline. Consistent with
  the Executive Summary structure.
- Executive Summary: extended broader-applicability closing with a
  forward-looking line previewing FFN/MoE integration.
- §5 section heading renamed from
  "All-Reduce Acceleration via PE_IPCQ" to
  "PE_IPCQ and Collective Communication".
- Title page date line refined to "2026 H1 Report".

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-11 15:13:01 -07:00
ywkang a965db7056 paper(exec-summary): edit-pass through Executive Summary done 
- Title: "Hardware-Software Codesign for Grouped-Query Attention on AHBM"
  (drop "Attention-Centric LLM Kernels" framing; subtitle pillars removed)
- Authors: Mukesh Garg, Jiyoon Lee, Yangwook Kang
- Affiliation: System Technology Group, AGI Computing Lab — 2026 1H
- Executive Summary rewritten into a tight three-paragraph structure:
  (1) GQA motivation + three AHBM architectural requirements (placement,
      inter-PE comm, pipelined memory-compute);
  (2) three matching HW-SW co-design mechanisms (GQA-aware placement,
      PE_IPCQ, composite-command GEMM under PE_SCHEDULER) and headline
      results (38% / 25-35% / 78%), closing on efficient HBM-bandwidth use;
  (3) broader applicability across the AHBM software stack.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-11 14:44:45 -07:00
ywkang 724a40453d test: fix 3 stale tests from the prior pull (registry list + memory_store msg) 
test_bench_registry EXPECTED_NAMES: add the 3 milestone benches the pull registered (milestone-1h-ccl, milestone-1h-gemm, milestone-gqa-headline) — 11 benches, alphabetical. test_memory_store::test_shape_mismatch_raises: the pull made same-size reshapes byte-conserving (allowed) and an over-large read raise 'Out-of-bounds read' (not 'Shape mismatch') — match the new message. Full suite now 820 passed / 1 skipped / 0 failed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 23:33:34 -07:00
ywkang 1a4ecb7aac gqa(ddd): DDD-0065 as-built content (prior commit only recorded the rename) 
The reconciliation table, corrected file/test lists, measured 5.22x ratio, and the new section 15 (data-mode numerics D8 + N1-N4) — the content the git-mv in 30a0451 left uncommitted.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 23:27:43 -07:00
ywkang 30a0451481 gqa(ddd): DDD-0065 reconciled to as-built + promoted to docs/ddd/ 
Updated the detailed design to match the shipped implementation: an as-built reconciliation table (P1-first, pinned-based DMA, D7 hard-cap-error, D8 output handle, prologue 1-stage-tile feed, D6.7 prologue-scoped), corrected file/test lists, the measured 5.22x dispatch ratio, and a new section 15 documenting the data-mode numeric subsystem (D8 + N1-N4) the original DDD omitted. Status Draft->Accepted; moved docs/adr-proposed -> docs/ddd/.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 23:27:13 -07:00
ywkang f0c5b9c33f gqa(adr): ADR-0065 Proposed -> Accepted 
Flat-ops CompositeCmd + softmax_merge recipe is fully implemented and tested: structural restructure (P1-P4), dispatch-cost measurement (P5b/P6, 5.22x per-tile win), and data-mode numerics (D8 output handle/space + N1-N4 composite GEMM / recipe MATH / accumulate / opt2 parity). git mv EN -> docs/adr/, KO -> docs/adr-ko/ (verify_adr_lang_pairs OK).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 23:16:52 -07:00
ywkang e9d62b908c gqa(adr-0065): N4 — opt2 numeric parity vs opt3 in data mode 
composite() now awaits its prologue recipe operands (e.g. the score tile from a prior #1 composite) so the cross-composite RAW (#1 writes scores -> #2 reads them) is ordered: #2 waits for #1 before reading. With N1-N3, opt2 (the two-composite recipe kernel) now computes a correct attention output end-to-end in data mode and matches opt3 within fp tolerance.

Scope note: the e2e opt2<->opt3 parity holds in the near-uniform-attention regime (small inputs). kernbench's K reshape-as-transpose (opt3 kernel docstring: 'correct for zero/symmetric inputs') makes opt2's per-sub-tile K interpretation differ from opt3's single-tile one for high-contrast inputs — a pre-existing kernbench limitation shared by both kernels, not the recipe. The EXACT recipe math (online-softmax merge m/l + O = O*corr + P@V) is verified non-trivially and exactly by N2/N3 (test_recipe_data_mode.py). Suite 817 pass / 3 pre-existing fail.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 23:13:16 -07:00
ywkang 5bbe293bea gqa(adr-0065): N3 — GEMM-epilogue accumulate dataflow (O = O.corr + P.V) 
Two fixes make the recipe's P.V GEMM fold into the rescaled accumulator: (1) op_log marks the composite gemm record accumulate=True when an 'add' epilogue targets the GEMM's own output; data_executor._execute_gemm then does O = O_existing + a@b instead of overwriting. (2) PE_SCHEDULER records the composite's numeric op at COMPLETION (in _retire_on_done) instead of at dispatch, so the GEMM's t_start sorts AFTER its prologue MATH ops -> it reads the computed P and the rescaled O.

Verified end-to-end (DataExecutor): the full recipe produces O = O_old*corr + P@V matching a numpy flash-attention reference. Suite 816 pass / 3 pre-existing fail.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 23:06:01 -07:00
ywkang 3d4a3d43c4 gqa(adr-0065): N2 — softmax_merge recipe MATH ops compute in data mode 
data_executor._compute_math gains the 5 recipe ops (rmax/rsum as keepdims reductions, max_elem/exp_diff/mul_bcast as binary; numpy broadcasting covers bcast_axis). tiling._math_stage now carries operand+output addrs/shapes/spaces + axis on the recipe prologue MATH stages. op_log.record_end promotes a MATH stage to op_kind='math' ONLY when it carries input_addrs -> the DataExecutor runs it; legacy epilogue MATH stages (bias/relu, no addrs) stay op_kind-opaque -> byte-equal.

Fixed a latent P2 lowering bug exposed by data mode: _resolve_recipe_dst gave reductions shape (G,) (axis removed) but max_elem broadcasts them against the running m=(G,1); reductions now keep the reduced axis as size 1 (keepdims) -> (G,1). Verified end-to-end: running the recipe's 8 MATH ops through the DataExecutor produces m, l (fully updated online-softmax) and O (rescaled by corr) matching a numpy reference. Suite 815 pass / 3 pre-existing fail.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 22:56:25 -07:00
ywkang 4089e18770 gqa(adr-0065): N1 — composite GEMM computes numerically in data mode 
Two changes make a composite's matmul replay in Phase 2: (1) op_log _extract_op_info(CompositeCmd) now SCANS ops for the gemm op (it is ops[0] for a legacy composite but sits after the prologue MATH ops in a recipe), emitting one composite_gemm record with the gemm's a/b/out addrs+spaces; (2) PE_SCHEDULER._dispatch_composite records the composite as a zero-duration numeric op (no-op without an op_logger, so latency-only mode is unchanged). Verified: a composite GEMM with seeded inputs writes a@b to its HBM output.

DataExecutor._execute_gemm is now best-effort: if an operand's data was never produced (torch.empty bench input) or a shard read is out-of-bounds (column-wise sharded operand read at full shape), skip the matmul instead of crashing. This keeps composite-using benches (qkv-gemm, composite-epilogue) green now that their previously-uncomputed GEMM actually fires. Suite 813 pass / 3 pre-existing fail.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 22:46:39 -07:00
ywkang e8d6c283d8 gqa(adr-0065): D8 — composite returns output handle + output-space DMA 
tl.composite now returns the output TensorHandle (not a CompletionHandle) so its result chains like tl.dot's; the handle carries the completion in a CompositeFuture pending so downstream ops and tl.wait auto-await it. out is a handle: out=tl.ref(addr,shape) (HBM, DMA_WRITE inside the composite) or an in-place TCM handle (STORE only); omitted -> TLContext auto-allocates a TCM scratch. out_ptr kept as HBM shorthand (= out=tl.ref(out_ptr, shape)) to avoid churning ~30 existing call sites.

tl.ref now returns space=hbm (it references HBM data; operand-input DMA stays pinned-based per D4 so input streaming is unchanged). tiling: the tile loop's DMA_WRITE is gated on out.space==hbm (out analog of the operand pinned rule) — a TCM output stays on-chip (chainable) and its high-bit scratch address no longer hits the DMA PA decoder.

Fixes the opt2 data-mode crash: the recipe accumulator O is TCM -> no DMA_WRITE -> opt2 now RUNS end-to-end in data mode (enable_data=True). Numeric parity of the recipe MATH ops is the next step. Suite 812 pass / 3 pre-existing fail; existing composite benches use out_ptr->hbm->DMA_WRITE unchanged (byte-equal).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 22:29:10 -07:00
ywkang dd525bfcb7 paper: add /paper skill + 1H HW-SW codesign report (GEMM, All-Reduce, fused GQA) 
New `/paper` slash-command skill that synthesizes ADR/SPEC content and live
KernBench benchmark results into a sectioned LaTeX technical paper compiled
to PDF with Tectonic (auto-installed). The skill negotiates a TOC, grounds
every number in committed artifacts or fresh bench runs, and keeps
report-only benches isolated.

This commit also includes the first generated report:
- docs/report/1H-codesign-paper/ — main.tex + per-section .tex, figures,
  toc.md contract, and the built 8-page main.pdf. Covers the platform
  (source-level kernels, latency model + accuracy, HW config from
  topology.yaml), GEMM via composite command, All-Reduce via PE_IPCQ, and
  fused GQA combining both, plus discussion/conclusion/2H future work.
- scripts/paper/ — isolated report harnesses (not registered benches):
  paper_gqa_latency.py harvests per-panel GQA end-to-end latency + engine
  occupancy (the milestone only emitted op-counts); paper_plot_gqa.py
  renders the GQA figures.

GEMM/All-Reduce reuse committed milestone figures/CSVs; GQA results are
generated fresh. Honest flags retained: PE_CPU dispatch cost is 0 in this
config, and the proposed two-composite softmax_merge decode is marked
designed-not-measured.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 22:15:14 -07:00
ywkang 4a55ae5c0b gqa(adr-0065): P6 — close dispatch-ratio loop (ADR-0064 Test #9 direction) 
Extend the opt2 R-sensitivity test to assert the opt3/opt2 dispatch ratio increases as R decreases (4.20x -> 5.22x -> 5.45x at R=0.25/0.0625/0.03125) — the cost is FIXED-dominated (command-count-driven), per ADR-0064 Test #9. Test-only; no production change (the P0 cost model + P2 recipe already make the ratio measurable).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 21:04:45 -07:00
ywkang 35453cc4fe gqa(adr-0065): P5(B) — decode opt2 two-composite kernel + dispatch measurement 
New _gqa_attention_decode_opt2.py: per-tile attention as #1 Q.Kt GEMM composite + #2 softmax_merge recipe composite (online merge + P.V + add). Runnable in op_log mode; full data-mode numeric parity is a deferred follow-up (P5-numerics).

Measured per-tile PE_CPU dispatch (ADR-0064 Rev2 default): opt3=960ns vs opt2=184ns = 5.22x (gate >2x), FIXED-dominated (command-count reduction) — the ADR-0060/0064 CPU-offload win. opt2<opt3 across R in {0.25,0.0625,0.03125}. K-before-V: softmax_merge prologue carries no DMA; V (ref) streams only in the GEMM tile loop.

Fix latent P3 bug surfaced by the first e2e recipe run: prologue MATH stages were FOLDED into the first GEMM tile, but a folded MATH->DMA_READ boundary needs a PE_MATH->PE_DMA token route the pipeline never wires (KeyError pe_dma). Now prologue/post-loop stages are fed as standalone 1-stage tiles (each completes on its component); completion count includes them. Existing benches have no prologue -> tiles + feed order unchanged (byte-equal); full suite 806 pass / 3 pre-existing fail.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 21:02:27 -07:00
ywkang 17fb94086e gqa(adr-0065): P4 — strict-FIFO RW hazard tracker 
_RwHazardTracker (pe_scheduler): a composite registers its write set (rw_handles) as in-flight; a new composite whose read (op operands) or write handles intersect an in-flight write set waits on those composites' done events before admission (ADR-0065 D6.3 / DDD 8). _dispatch_composite calls admit() before feeding and retires on the done event.

Legacy composites (rw_handles=()) never register and never block -> existing benches untouched (byte-equal). Strict FIFO is conservative (partial overlap waits for full drain); RW-aware reorder is deferred (A4).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 20:36:31 -07:00
ywkang 55f025c4b1 gqa(adr-0065): P3 — flat-ops PE_SCHEDULER plan (position-scan + prologue stages) 
tiling.generate_plan_from_ops: scan flat ops for the GEMM (<=1); pre-GEMM KERNEL ops become single-shot prologue MATH stages, post-GEMM ops split by scope (K_TILE/OUTPUT_TILE epilogue + KERNEL post-loop). MATH-only composite reproduces the legacy math-head plan. Prologue/post-loop stages fold into the first/last tile so the feeder + completion counting are untouched (existing benches have neither -> byte-equal op_log).

PipelinePlan gains prologue_stages/epilogue_stages. pe_scheduler._generate_plan delegates to generate_plan_from_ops. DMA keeps the existing pinned signal (NOT a space flip); recipe scratch/primary-out handles are pinned=True so the head GEMM's auto-bound a (=P) is consumed in place.

ADR-0065 D4/D6.7/Test#5 amended (EN+KO): DMA decision from pinned, not space; prologue recipe ops TCM-only (head op exempt -- it may stream from HBM).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 20:24:33 -07:00
ywkang 184f654295 gqa(adr-0065): P2 — softmax_merge recipe + TLContext prologue lowering 
New triton_emu/tl_recipes.py: RecipeDescriptor/EngineOp/PrimaryOutSpec + RECIPE_DESCRIPTORS['softmax_merge'] (8-step engine_seq). PE_SCHEDULER does not import it (ADR-0065 D5 boundary).

TLContext.composite(): add prologue=[...] + out=TensorHandle kwargs, a optional. _expand_prologue lowers a recipe into flat MATH OpSpecs (scope=KERNEL), allocates TCM scratch, derives the primary-out slot 'P' and auto-binds it into the head GEMM a (D6.6 conflict check); rw_handles=(m,l,O). decode-opt2 #2 lowers to 10 ops [rmax,max_elem,exp_diff,exp_diff,rsum,fma,mul_bcast,copy,gemm,add].

D6.7 (MATH operand TCM-only) scoped to prologue recipe ops only — the head op (gemm or math) keeps existing DMA-staged-from-HBM behavior. D6.1 (GEMM count <=1) on the whole composite. Host-side lowering only; PE_SCHEDULER position-scan is P3.

Also commit the ADR-0064 Rev2 promotion content that the prior commit's git mv dropped: Status Proposed->Accepted + D7 amended to hard-cap ValueError (no segmentation), EN+KO.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 19:48:29 -07:00
ywkang 47e2c78c66 gqa(adr-0064/0065): flat-ops CompositeCmd (P1) + structural dispatch cost (ADR-0064 Rev2); promote ADR-0064 
ADR-0065 P1: CompositeCmd -> flat ordered ops list (drop legacy op/a/b/out_addr fields); OpSpec.operands dict + out handle. Meaning-preserving (op_log byte-equal); pe_scheduler + op_log read the head op.

ADR-0064 Rev2: replace Rev1 per-op cost table with structural FIXED + logical_bytes*R formula. logical_bytes on every PeCommand; new common/pe_cost_model.py; cost centralized in TLContext._emit (load/recv_async charge explicitly); pe_cpu/kernel_runner wire the per-PE model + clock. D7: cap exceeded -> ValueError (no auto-segmentation). Remove Rev1 cpu_issue_cost.py + its tests. No goldens churn.

Promote ADR-0064 Rev2 Proposed->Accepted (docs/adr/ + docs/adr-ko/); amend D7 (error not segmentation) + record P1-before-P0 ordering in ADR-0064/0065 Migration notes (EN+KO).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-10 19:18:04 -07:00
mukesh 79ddb12b42 gqa: explain why Tile-0 / bootstrap can't be folded into the tile loop 
Comment-only change. Adds a kernbench-only limitation note at the
Tile-0 / bootstrap block in all four GQA attention kernels:

  - persistent (m, l, O) must live outside tl.scratch_scope so
    subsequent tiles' merges can read them;
  - kernbench has no scratch-backed initializer (tl.zeros / tl.full
    return addr=0 handles), so we can't seed (-inf, 0, 0) and rely
    on tl.copy_to;
  - therefore Tile 0 must compute the initial running state directly.

prefill_long adds a third bullet noting the ring-step ordering reason
(k=0 must send W before any k>0 can recv E, so the (t=0, k=0) step has
to run outside the scoped loop where the send is conditional on C > 1).

Each block also notes that a real Triton port collapses Tile 0 into a
unified loop (SSA tensors stay live across iterations).

No behavior change; tests unchanged.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-10 16:35:43 -07:00
ywkang 3647e2d8f5 gqa(adr): ADR-0064/0065 second-pass review — 1024B rationale, FIFO as ordering source, robust formula-based tests, D5/D6 auto-bind alignment 
ADR-0064 Revision 2 second-pass review fixes:
- D7 added 1024-byte cap rationale: explicitly framed as a *safe
  engineering limit* (intentionally above all known composites at ~322
  bytes, finite descriptor capacity placeholder), not a measured HW
  number. Default is the discipline; topology override for real HW.
- D7 segment ordering: strict FIFO is the *ordering source*; rw_handles
  is dependency metadata, not the ordering primitive. Explicit note for
  the future RW-aware reorder migration path.
- Tests rewritten around the formula (D1) instead of specific numbers:
  - Test #1 → "formula preservation" (parametrized across composites,
    asserts dispatch == FIXED + bytes × R, not anchor 43 ns).
  - Test #2 → robust "opt3 > 2 × opt2" qualitative gate; ≈ 4.0× is
    informative only, not the gate, so calibration changes don't break.
  - Test #9 → qualitative (opt2 < opt3 at all R), not a numeric ratio.
- Anchor description in D3 stays as *informative* (shows the model
  produces ≈ 43 ns at defaults), but is no longer a test gate.

ADR-0065 second-pass review:
- D5 step 5 aligned with D6.6 invariant: explicit conflict check before
  auto-binding (previously written as if auto-bind always happens).
  Mirrors D6.6 wording: error if both kernel explicit `a` and prologue
  primary_out are present.
- Test #7 made robust: `opt3 > 2 × opt2` gate (was ≈ 4.0×).
  Default-calibration model expectation (≈ 4.0×) recorded as informative
  reference, not the gate.

DDD-0065 follow-on robust thresholds:
- §1.4 success criteria: ratio condition rewritten as `opt3 > 2 × opt2`
  with model-expected ≈ 4.0× as informative.
- §9 P6 phase gate: `opt3 > 2 × opt2` (was 4.0× ± 15%).
- §10 dispatch ratio test: assert `opt3 > 2 × opt2`; record observed
  ratio for performance tracking but do not gate on it.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-10 16:34:34 -07:00
ywkang a8d04750e6 gqa(adr): ADR-0065 review fixes — KERNEL phase/repetition split, auto-bind conflict + MATH-TCM invariants, FIFO conservatism note, logical_bytes per-op summation 
ADR-0065 review fixes:
- D3 semantics clarified: position determines *phase* (pre-loop / head /
  in-accumulation / per-output-tile / post-loop), scope determines
  *repetition*. KERNEL = "once per CompositeCmd invocation", not "once
  per kernel launch" — removes ambiguity around softmax_merge's KERNEL
  scope at different positions before/after GEMM.
- D4 MATH-TCM invariant: Phase 1 limits DMA auto-insertion to GEMM
  operands and head outputs; MATH operand with space="hbm" raises
  validation error at TLContext emit. Catches kernel misuse early; a
  MATH-with-HBM path is future work (explicit tl.load + MATH composite).
- D6 #6 NEW: auto-bind conflict — explicit GEMM `a=` + prologue
  primary_out simultaneously → validation error (prevents "which value
  wins" ambiguity).
- D6 #7 NEW: MATH operand TCM-only restatement.
- Consequences/Negative: strict-FIFO conservatism note — safe but may
  under-expose composite-level overlap when handle-sharing composites
  could in principle run in parallel.
- Test Requirements #10, #11: auto-bind conflict + MATH-HBM validation
  error tests.

ADR-0064 Revision 2 clarification:
- D2 counting rule: per-op summation, no deduplication. Same handle
  appearing in multiple OpSpecs or in rw_handles is counted
  independently — matches HW reality (each descriptor field is a
  distinct address slot, rw_handles is separate metadata block).
  Example walked through (mul_bcast with O in src_a + out + rw_handles).

DDD-0065 follow-on:
- §3.2: new test file test_tl_composite_validation.py covers auto-bind
  conflict, MATH-TCM, GEMM-count invariants.
- §4.3: phase vs repetition split aligned with ADR-0065 D3.
- §6 lowering pipeline: step 5 adds auto-bind conflict check; step 8
  NEW — operand space validation; step 9 (size cap); step 10 (emit).
- §8 strict-FIFO conservatism note: explains why next-tile #1 cross-tile
  pipelining still works under FIFO (no rw overlap), and where the
  conservatism would hurt (future handle-sharing recipes).
- §10 verification: invariant-guard tests (auto-bind conflict, MATH
  TCM, per-op summation accuracy).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-10 16:29:39 -07:00
ywkang 95cecccd8e gqa(adr): ADR-0064/0065 review fixes — composite size cap, type-aware logical_bytes, R recalibration 
ADR-0064 Revision 2 review fixes:
- D7 NEW: composite size cap (MAX_COMPOSITE_LOGICAL_BYTES default 1024
  bytes). Oversized recipes deterministically segmented into N
  CompositeCmds; each segment incurs its own dispatch cost. Models real
  HW limits (descriptor queue entry, scheduler parser buffer, command
  SRAM) and prevents the model from rewarding pathologically-large
  fused composites.
- D2: type-aware extra-field byte counting (int/float=4, bool=1,
  tuple/list=1+4N, str=1) — replaces uniform 4 bytes per extra.
- D3: recalibrated defaults to FIXED=40 cycles, R=0.0625 cycles/byte
  (16 B/cycle — typical on-die descriptor queue width); anchor stays at
  ~43 ns for typical 1-OpSpec composite. Clarified anchor description:
  DMA stages do not appear in logical_bytes (auto-inserted by
  PE_SCHEDULER from operand.space per ADR-0065 D4).
- D4: removed clock_freq_ghz from pe_cost_model: override block;
  conversion uses the PE node's existing clock_freq_ghz attr. Added
  max_composite_logical_bytes knob.
- Context: emphasized command-count reduction (FIXED) as the primary
  signal; byte term as secondary refinement.
- Open review: added large-composite scheduler-cost stress test.
- Test req: added composite-size-cap (#8) and R-sensitivity sweep (#9).

ADR-0065 + DDD-0065 follow-on updates:
- opt2 vs opt3 dispatch ratio updated 2.4× → ≈4.0× under new defaults
  (FIXED-dominated, reflecting the corrected framing).
- Test req #9: decode opt2 composite fits within 1024-byte cap; no
  segmentation needed for the GQA workload.
- DDD §6: TLContext lowering checks logical_bytes against cap (step 8).
- DDD §11: performance model recomputed with new defaults + sensitivity
  table across R ∈ {0.25, 0.0625, 0.03125} confirming opt2 < opt3 holds.
- DDD §9 P6 gate: ratio band 2.4×±10% → 4.0×±15%; sensitivity sweep added.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-10 16:21:58 -07:00
mukesh 7fad0371c5 gqa: tile-granular Ring KV (P3c) + rename to gqa_attention_* + ADR-0060/62/63/64 → Accepted 
Three logically distinct changes, bundled for atomic test green:

1. **P3c — prefill_long tile-granular Ring KV** (ADR-0060 §5.5.1 amendment).
   Convert the ring from slice-granular (one full ``(d_head, S_local)``
   KV slice per step) to tile-granular (``n_tiles`` tiles of
   ``TILE_S_KV`` per step). Nested loop with outer tile, inner ring step:
   each tile propagates through all C ring positions before the next
   tile starts, so IPCQ in-flight depth stays at 1 per direction.
   Bootstrap at ``(t=0, k=0)`` outside the scratch_scope establishes the
   persistent ``(m, ℓ, O)``; every other iteration scope-wraps + persists
   via ``copy_to``. Per-rank persistent scratch shrinks to ~1 KB; per-tile
   scope bounded by TILE_S_KV regardless of S_local. Headline:
   prefill_long now completes at S_kv=128K (previously overflowed).
   New: ``tests/attention/test_gqa_prefill_long_tile_ring.py``
   (3 tests — ceiling-lift + tile-granular ipcq_copy count +
   per-CUBE distributed output regression guard).

2. **Rename ``gqa_*`` → ``gqa_attention_*``** across kernel files,
   function names, and importers. The "attention" name makes the role
   explicit (GQA is grouped-query attention) and matches upstream Triton
   FlashAttention naming conventions. Renames:
     _gqa_decode_long.py        -> _gqa_attention_decode_long.py
     _gqa_decode_short.py       -> _gqa_attention_decode_short.py
     _gqa_prefill_long.py       -> _gqa_attention_prefill_long.py
     _gqa_prefill_short.py      -> _gqa_attention_prefill_short.py
   And function names ``gqa_<phase>_<context>_kernel`` →
   ``gqa_attention_<phase>_<context>_kernel``. Updated 1 bench file
   (milestone_gqa_headline.py) and 10 test files.

3. **ADR-0060 / 0062 / 0063 / 0064: Proposed → Accepted**.
   All four are reflected in production code and covered by tests:
   - ADR-0060 (GQA fused attention): 4 kernels deployed; §5.5.1
     amendment added for the tile-granular Ring KV introduced by P3c
     (EN + KO mirror).
   - ADR-0062 (lazy tl.load): LoadFuture + _await_pending live in
     tl_context.py.
   - ADR-0063 (tl.scratch_scope + tl.copy_to): used in every chain
     reduce + tile sweep + ring step. EN-only previously; KO
     translation authored as part of this commit (CLAUDE.md
     bidirectional rule).
   - ADR-0064 (per-op-type CPU issue cost): cpu_issue_cost.py +
     issue_cost_table wiring in tl_context.py (Phase E).
   Files git mv'd from docs/adr-proposed/ to docs/adr/ (EN) and
   docs/adr-ko/ (KO). ADR-0061 (tl.broadcast) stays Proposed — no
   implementation; documented as optional convenience primitive in
   the ADR itself.

Tests: 88/88 focused regression green
(tests/attention/ + Phase E + TL discipline).
ADR pair verification: ``python tools/verify_adr_lang_pairs.py`` OK.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-10 16:20:00 -07:00
ywkang a8c50238c6 gqa(adr): ADR-0064 Rev2 (structural dispatch cost) + ADR-0065/DDD-0065 (flat-ops composite + softmax_merge recipe) 
ADR-0064 Revision 2: replace per-op-type calibration table with a
structural formula `FIXED_PER_CMD + cmd.logical_bytes × R`, defaults
anchored at typical composite ≈ 45 ns. Topology yaml override under
`pe_cost_model:` block; logical_bytes property per PE command.

ADR-0065: implement ADR-0060 §5.6 / §8 item 4 carve-out as a flat-ops
CompositeCmd (no head/epilogue structural fields — position + scope
drives placement) + first stateful recipe `softmax_merge` (MATH-only
8-step). RECIPE_DESCRIPTORS lives in TLContext-adjacent module only;
PE_SCHEDULER stays recipe-free and auto-inserts DMAs from operand
`space`. Strict-FIFO RW hazard tracker; ≤1 GEMM per composite
invariant. User-facing `tl.composite(prologue=[...], op=, epilogue=[...])`
API preserved; existing benches unchanged (meaning-preserving refactor).

DDD-0065: implementation-ready phased plan (P0=ADR-0064 Rev2 first,
P1-P6 for ADR-0065), file plan, recipe engine sequence, scheduler
plan-gen algorithm, RW tracker design, test matrix.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-10 16:11:20 -07:00
mukesh 91cdfebb67 gqa: S_kv tile sweep for decode_long/decode_short/prefill_short (ADR-0063 §A.2) 
Replace each kernel's local one-shot partial with a Tile-0 bootstrap + a
``scratch_scope``-wrapped merge loop. Per-rank scratch is now bounded by
``TILE_S_KV = 1024`` regardless of ``S_local``, lifting the long-context S
ceiling. Backward compatible at ``S_local <= TILE_S_KV`` (loop body is
empty; op_log structurally identical to today).

Framework: extend ``memory_store.read`` to support partial reads at
offsets within a stored region. Models real Triton ``tl.load(ptr+offset,
shape=...)`` — needed because each tile loads ``k_ptr + tile_start*row_bytes``
for its sub-slice of the per-rank KV region.

prefill_long is intentionally left untiled. Its ring loop carries
full-slice ``Kc``/``Vc`` for ``tl.send`` between CUBEs, so tile-sweeping
step 0 wouldn't shrink the kernel's actual scratch footprint. Lifting
prefill_long's ceiling requires a tile-granular ring rewrite — separate
phase. Docstring + inline comments cleaned to reflect the current shape.

Docstrings across all 4 GQA kernels: drop P1a/P2a/P2b/P6a/P6b lineage
paragraphs and historical deviation lists; describe each kernel by what
it does, not how it got here. Add explicit ``# Local attention`` /
``# Communication`` section headers.

Tests: new ``tests/attention/test_gqa_tile_sweep.py`` with 5 tests
covering single-tile op_log stability, multi-tile ``copy_to`` emission
across all 3 refactored kernels, and the headline ceiling-lift
(decode_long at S_kv=256K). Full focused regression green: 85/85 across
``tests/attention/`` + Phase E + TL discipline tests.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-10 15:33:29 -07:00
mukesh 5a76ed4f6a gqa: rename long-context kernels to *_long for symmetry with *_short 
Make the file + function naming symmetric:
  _gqa_decode.py        -> _gqa_decode_long.py
  _gqa_prefill.py       -> _gqa_prefill_long.py
  gqa_decode_kernel     -> gqa_decode_long_kernel
  gqa_prefill_kernel    -> gqa_prefill_long_kernel

Mirrors the existing _gqa_{decode,prefill}_short.py naming. Updates the
two imports + two call sites in milestone_gqa_headline.py and the 9
attention tests that import the kernels.

Tests: 72/72 focused regression green (tests/attention/ + Phase E + TL
discipline). milestone-gqa-headline bench passes its 7 panel/schema
assertions.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-10 10:05:54 -07:00
mukesh d282144339 gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model 
Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.

ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
  intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
  online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
  (m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.

ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
  owns whole KV heads; PE-parallel heads with intra-group chain
  reduce. Prefill has no Ring KV (each head fully resident).

ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).

ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.

ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
  primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
  consults table with dispatch_cycles fallback (ADR-0046 §D6
  back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
  DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
  measurable end-to-end.

P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).

Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
  _attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
  (EN + KO kept in sync).

Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
 2026-06-09 18:15:59 -07:00
mukesh b3730a33eb policy(adr): ADR-0026 Revision 6 — document DPPolicy.cube_start 
Amends ADR-0026 to reflect the cube_start field added in e2fe331.
The production code shipped without an ADR update; this fills that
gap. Documentation-only change (no production code, no test code).

Changes (mirrored in both EN and KO):
- Status: Revision 5 → Revision 6
- D1: add ``cube_start: int = 0`` to the canonical DPPolicy dataclass
- D3: ``cube = policy.cube_start + cube_id`` in resolve_dp_policy
- D8: new section explaining purpose (disjoint cube sub-meshes for
      GQA Llama-70B 8-KV-group headline), semantics, default-0 backward
      compatibility, intra-device constraint, design rationale
      (scalar vs 2D origin vs cube_ids list), and the kernel-side
      cube_start subtraction needed to compensate for ADR-0022's
      physical-cube-id ``program_id(axis=1)`` semantics.

tools/verify_adr_lang_pairs.py passes (EN/KO Status keyword and
title in sync).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-08 14:18:27 -07:00
ywkang ab324c155c gqa(adr): review fixes — unify KV placement to contiguous, ring/reduce cost model, opt3 wording, rank def 
Per review:
- Placement unified to *contiguous* C×P blocks throughout (S0/S2.1/S2.2/S0.5);
  round-robin demoted to the rejected alternative. Adds driver SP-enable
  threshold fallback (smaller C / C=P=1 for short/early decode).
- Ring-vs-reduce cost model in S5.5: reduce ~ G*T_q*log(C*P) (O dominant;
  m,l scalars), ring ~ 2*S (total K+V bytes a CUBE injects over C-1
  rotations; recv_async pipelines so latency ~ max-step) -> ring wins when
  T_q > 2S/(G*log(C*P)).
- opt3 'removes the bubble' -> 'hides (subject to scheduler+engine balance)'
  everywhere; table 'hidden*' with footnote.
- 'rank' defined (SP participant = a PE in a CUBE, KV shard in its HBM->TCM).
- out-proj handoff contract (S0.5.4); S11 gate-type note (absolute latency
  deferred to ADR-0064; structural + relative-to-baseline gates now).
- greenlet-as-contrast tightened to 'primitive-op (tl.dot) path' (S1, SB).
KO mirror synced; DDD gets the SP-enable threshold fallback. Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-05 00:52:38 -07:00
ywkang 9103661098 gqa(adr): mark DDD-0060 as synced in ADR-0060 SB (EN+KO) 
DDD-0060 was rewritten to the two-kernel composite-hybrid design, so SB
item 1 ('DDD not yet synced') is stale. Update EN + KO mirror to record the
DDD as synced (ADR authoritative, DDD = impl how-to). Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-04 23:09:14 -07:00
ywkang c699e2ddfc gqa(ddd): rewrite DDD-0060 to match current ADR (two-kernel composite hybrid) 
Supersede the old greenlet-primitive/load_async/single-kernel DDD. Align to
ADR-0060's composite-hybrid + hierarchical CUBE-Group SP with two kernels:
decode+SP (head-replicated, contiguous C×P static shard, 2-level
reduce-to-root) and prefill+SP (1 Q head/CUBE, Ring KV, no reduce). Updated
file plan (lazy tl.load, scratch_scope, broadcast-opt; evolve
_attention_mesh_mlo_2d/_kv; DPPolicy.cube_start; llama70b_4sip.yaml),
placement (contiguous shared KV), phase plan (P1..P8), verification, perf
model, risks, glossary. Defers design rationale to ADR-0060 (now
authoritative); open items point to ADR-0060 §B. ADR-0064 cost model noted.
Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-04 23:08:02 -07:00
ywkang 973a452b77 gqa(adr): sync ADR-0060 KO mirror to current EN (two-kernel design) 
Regenerate the KO mirror to match the current EN: two SP kernels
(decode=reduce / prefill=ring), TL;DR full code for the 3 decode variants
+ prefill, 'Q replicated / M-fold' and '1 Q head per CUBE' terminology,
contiguous shared KV layout, opt2 tl.wait, and all SB items. Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-04 23:02:49 -07:00
ywkang 8de13efb9d gqa(adr): opt2 ex_composite decode needs tl.wait() before the reduction 
The ex_composite #2 updates acc=(m,l,O) asynchronously on the scheduler and
the kernel never reads its output, so no auto-wait fires; acc is only final
after the last #2 in the serial chain. Add tl.wait() (drain all composites)
before hierarchical_reduce_and_store reads acc. opt1/opt3 don't need it
(their composite outputs are consumed in-iteration -> auto-wait). Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-04 22:49:59 -07:00
ywkang 9c554e3f64 gqa(adr): fix prefill TL;DR pseudocode — K/V are TCM handles for the ring 
The ring rotates K/V over IPCQ, so they must be kernel-held TCM handles
(tl.load), not HBM tl.ref streamed inside the composite. Correct the
prefill kernel: load K pre-transposed [d, S/C] and pass the TCM-resident
Kc/Vc directly to the composite (drop the bogus tl.ref(Kc,(d,TILE_S)) and
TILE_S); recv shapes match ([d,S/C] for K, [S/C,d] for V). Comments note
why K/V live in TCM. Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-04 22:40:12 -07:00
ywkang c6bc286740 gqa(adr): TL;DR — full code for 3 decode variants + prefill (per request) 
Move the 3 decode CPU-pipelining variants (opt1 current / opt3 sw-pipe /
opt2 ex_composite) up into the TL;DR as full standalone kernels alongside
the full prefill ring kernel, with the comparison table. S5.6 is reduced to
a brief anchor (still referenced by S8/SB) pointing to the TL;DR code +
keeping the recommend/cost-model linkage. Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-04 22:05:15 -07:00
ywkang 444810c6e3 gqa(adr): TL;DR — Q-replicated/M-fold wording + pointer to S5.6 decode variants 
The prior commit updated S0/S10/S5.5 terminology and added S5.6 (3 decode
variants) but left the TL;DR with the old 'all G heads replicated' wording
and no pointer to S5.6. Sync the TL;DR: 'Q replicated (G heads stacked into
the GEMM M-dim, M-fold)', '1 Q head per CUBE', and a one-line pointer to the
3 CPU-pipelining variants in S5.6. Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-04 22:01:25 -07:00
ywkang 943626d758 gqa(adr): ADR-0060 — 1 Q head/CUBE term, decode 3 variants (S5.6), shared contiguous KV layout 
- Terminology: 'Q replicated' (all G query heads stacked into the GEMM M-dim;
  M-fold explained) for decode; 'one Q head per CUBE' precise for prefill.
- New S5.6: three decode CPU-pipelining variants — opt1 current CompositeCmd
  (has GEMM-engine bubble), opt3 software pipelining (issue next Q.Kt before
  this tile's softmax; Sj in persistent double buffer; ships now, no new cmd),
  opt2 ex_composite split into two (#1 = existing GEMM+scale reads K first;
  #2 = softmax+P.V+accumulator merge, the only new flash-epilogue machinery,
  gives DMA K-before-V priority). MATH engine already has max/sum/exp — the
  new part is the stateful flash accumulator, not the ops.
- S2.1/SB: shared prefill/decode KV layout = contiguous CxP blocks (prefill
  causal-skip needs contiguous; avoids prefill->decode reshard; short-context
  under-use caveat). S8 item 4 sizing note for the two-composite split.

Prefill note: opt2/opt3 give little for prefill (causal if can't enter a
composite; recv_async already overlaps). Docs only; KO mirror deferred.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-04 21:59:03 -07:00
ywkang 8176cdf287 gqa(adr): split ADR-0060 SP into two kernels — decode=reduce, prefill=ring 
Decode-SP and prefill-SP are structurally different and cannot share one
kernel (principle: move the smaller thing):
- Decode: O=[G,d] tiny, KV cache big -> keep KV statically sharded/resident,
  G heads replicated (M-fold), move only (m,l,O) via the 2-level reduce (S4).
- Prefill: O=[S,d] big -> shard heads (1 query head per CUBE, C=G),
  rotate KV (Ring KV, S5.5), no (m,l,O) reduce; each CUBE writes its own head.

Rewrites TL;DR (two kernels), S0 (head map differs by case), S0.5.4 (output
head distribution differs -> downstream out-proj impact), S4 (scoped to
decode; S4.1 = intra-CUBE KV-split + PE reduce, the only way decode uses P
PEs), S5.1 (decode skeleton), S5.5 (head-parallel Ring KV), S9/S10/S11, and
adds SB items (two head mappings, output asymmetry, prefill within-CUBE PE,
C=G coupling, reconcile with _attention_mesh_mlo_2d). KO mirror deferred
until the design stabilizes (adr-proposed is mirror-exempt). Docs only.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-04 21:32:08 -07:00
ywkang a61fed3ce0 gqa(adr): ADR-0060 hierarchical CUBE-Group SP — 2-level reduce-to-root 
Placement pivot (user-approved): a CUBE Group = C CUBEs within one SIP that
jointly own one KV head + its G=8 query heads. KV sequence sharded 2 levels
(Level-1 inter-CUBE over C, Level-2 intra-CUBE PE over P=8 = C*P ranks); G
folded into matmul M. C is a knob (8/4; C=1 = single-CUBE). Reverses the old
'a query head never spans CUBEs' non-goal (held only because the baseline
was H_kv=1). device=SIP; the for-kv loop is gone (head picked by CUBE coord).

Reduction is a 2-level reduce-to-root (not all-reduce): Level-2 PE tree ->
Level-1 center-root CUBE-mesh reduce, adapting lrab_hierarchical_allreduce's
inter-CUBE pattern as reduce-only + log-sum-exp. Data-driven (send on local
P.V completion, no global barrier) + level-pipelined; per-level topology
configurable (tree for decode, ring for long prefill).

Rewrites SS0/2/4/5/0.5/8-11, pseudocode, and adds SSB items (4-SIP config,
mesh partition, C knob, invariant-reversal check, index-math test). KO mirror
updated. Topology grounded in topology.yaml (4x4 CUBE mesh/SIP, 8 PE/cube).
Docs only; no production code changed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
 2026-06-04 17:36:18 -07:00
mukesh f009ba7b1f topology: add llama70b_4sip.yaml for the 8-KV-group diag harness 
Adds the 4-SIP topology referenced by tests/attention/test_attention_8kv_groups_diag.py
(landed in 39fc2e9). Identical to the repo-root topology.yaml except
``system.sips.count: 2 -> 4``, matching the GQA Llama-70B sharding
study's 1 Q-head-per-cube baseline (64 cubes = 8 KV-groups x 8
cubes/group across 4 SIPs).

Opt-in only — only the diag harness and milestone-gqa-llama70b use
this file via env var. The repo-root topology.yaml is unchanged and
continues to drive milestone-1h-ccl / -gemm at their 2-SIP scale.

Without this commit the diag tests pytest.skip on a fresh checkout
("4-SIP topology missing"). With it, step 1 / step 2 / step 3 all run
and pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
 2026-06-04 12:42:40 -07:00
314 changed files with 68029 additions and 5727 deletions
Expand all files Collapse all files
.claude/commands/paper.md
+417
View File
@@ -0,0 +1,417 @@
---
description: Generate the 1H HW-SW codesign technical paper (PDF) — propose/confirm a TOC, then synthesize ADRs + live KernBench benchmark results into a sectioned LaTeX paper compiled with Tectonic.
---
# `/paper` — HW-SW Codesign Technical Report Generator
Produces a **technical-paper-style PDF** at
`docs/report/1H-codesign-paper/build/main.pdf`, synthesizing the
ADR corpus, SPEC.md, and **live KernBench benchmark results** into a
publication-shaped report on the three 1H optimizations:
1. **GEMM Acceleration** using the composite command
2. **All-Reduce Acceleration** using PE_IPCQ
3. **Fused Grouped Query Attention** using both composite command and PE_IPCQ
plus an explanation of the simulation platform itself (why KernBench
exists, its execution model, and how it computes latency / how accurate
that is).
Audience is the **internal team** (engineers + architects), but the
report reads like a real technical paper: **do not cite `ADR-NNNN` or
`SPEC §X.Y` anywhere** — not in prose, not in a reference list. ADRs and
SPEC are *grounding sources only* (use their content; never name them).
The **References** section, if used, lists **external literature only**
(e.g. FlashAttention, Megatron-LM, GPT-3, Llama 3) — actual papers, not
internal docs.
---
## Operating principles (read first)
- **Confirm the TOC before writing.** The first thing `/paper` does is
present the outline (saved `toc.md` if present, else the default
below) and **stop for the user's confirmation / adjustments**. Do not
generate sections until the user approves the TOC.
- **Ground every number.** Every quantitative result in prose, tables,
or captions MUST trace to (a) a committed bench artifact
(`*.csv` / `*.json` / `*.png` under `1H_milestone_output/` or
`docs/diagrams/`), (b) a freshly-run bench output, or (c) a specific
ADR/SPEC quotation. **Never invent latencies, speedups, or
utilizations.** If a needed number does not exist, generate it via a
bench run (below) or state the gap — do not guess.
- **Reuse existing results by default.** The three milestone benches
already have committed outputs. Do **not** re-run a bench unless the
user asks to regenerate, or the artifact is missing/stale.
- **Design / proposal status honesty.** Some designs are `Proposed`,
not yet implemented (e.g. ADR-0061, ADR-0065). Present these as
proposed and clearly separate "modeled & measured" from
"designed, not yet measured." Do not report measurements for
unimplemented paths.
- **English prose.** Architectural terms (composite command, PE_IPCQ,
GQA, TCM, SIP, CUBE, PE) used as-is.
---
## Invocation
- `/paper` — **default**. Resolve the period, present the TOC
(`toc.md` if it exists, else the default TOC), and **stop** for
confirmation. Once the user approves (possibly after editing the TOC
in conversation), persist `toc.md` and proceed to full build.
- `/paper toc` — only (re)negotiate and persist the TOC. No sections,
no build.
- `/paper build` — assume the saved `toc.md` is approved; run the full
pipeline (sections → figures → compile PDF) without re-confirming.
- `/paper section <id>` — regenerate a single section (`<id>` = a
filename stem under `sections/`, e.g. `04-allreduce`) and recompile.
- `/paper figures` — (re)generate only the benchmark figures/data
(runs/refreshes benches), no prose changes, then recompile.
- `/paper pdf` — only recompile the existing `.tex` to PDF (Tectonic).
If `toc.md` does not yet exist, any mode except `/paper`/`/paper toc`
first falls back to presenting the default TOC for confirmation.
---
## Output layout
```
docs/report/1H-codesign-paper/
toc.md # the agreed outline (persisted; human-editable)
main.tex # preamble + \input of each section
refs.bib # ADR/SPEC bibliography (optional; or inline refs)
sections/
00-exec-summary.tex
01-introduction.tex
02-platform.tex # includes 3.4 Modeled Hardware Configuration
03-gemm.tex
04-allreduce.tex
05-gqa.tex
06-discussion.tex
07-conclusion.tex
08-future-work.tex
figures/ # PNGs copied in from bench outputs (portable)
build/ # Tectonic output; main.pdf lives here
```
Report-specific benchmark/plot harnesses (when a figure must be
generated, see *Generating missing results*) live **isolated** under:
```
scripts/paper/ # report-only harnesses — NOT registered benches
```
This directory is intentionally separate from `src/kernbench/benches/`
so the report's harnesses never collide with other people's benches
(the bench registry only audits `src/kernbench/benches/`). Treat
`scripts/paper/*` and everything under `docs/report/1H-codesign-paper/`
as **derived artifacts** (like `docs/diagrams/`): creating/updating them
does not require Phase 2 approval, but they MUST stay consistent with
SPEC.md and ADRs.
---
## Default TOC
If `toc.md` is absent, propose this and ask the user to confirm or edit:
1. **Executive Summary** — the attention-optimization goal, the two
enabling optimizations (GEMM composite command + PE_IPCQ communication),
the fused GQA capstone, the headline results, and the bottom-line
recommendation on which HW changes are worth keeping. May run a
paragraph or two longer than a terse abstract; front-loads conclusions.
2. **Introduction** — the 1H focus is **attention-kernel optimization**:
FlashAttention-style tiling and Grouped Query Attention (GQA), building
on prior work on Multi-Head Attention (MHA, studied earlier — reference
it as the established baseline, not re-derived here). Optimizing a fused
attention kernel requires two enabling optimizations, each studied in
its own right and then **combined inside the fused kernel**:
(a) **GEMM optimization** (the composite command) and
(b) **communication optimization** (PE_IPCQ). Frame the three result
sections as *two enablers → the fused GQA capstone that uses both*.
Motivate why HW-SW codesign (not software alone) is required to realize
the gains.
3. **The KernBench Platform**
1. *Why KernBench* — source-level kernel execution with **no compiler /
software-stack dependency**; isolate algorithm-level optimization from
other layers so a kernel can fully exploit the hardware without E2E
cross-layer entanglement.
2. *Execution model* — discrete-event graph; `runtime API → sim_engine →
components`; PE pipeline + composite commands; 2-pass data/timing model.
3. *Latency model & accuracy* — traversal-based golden invariants
(every hop pays > 0, deterministic, explicit connectivity); the
structural CPU-dispatch cost model; per-edge BW occupancy, HBM
pseudo-channel parallelism, flit streaming; known simplifications &
what calibration anchors the constants.
4. *Modeled Hardware Configuration* — the shared, platform-wide
quantitative config used across **all** experiments, read from
`topology.yaml`: SIP/CUBE/PE hierarchy (counts & nesting), PE clock &
compute throughput, TCM/SRAM/HBM capacities & bandwidths, NOC /
inter-CUBE / inter-SIP link bandwidths & latencies, and the cost-model
constants. Per-experiment workload parameters (GEMM shapes, All-Reduce
topologies, GQA seq/head/user counts) stay in each evaluation section,
not here.
4. **GEMM Acceleration via Composite Command** — necessity · design · results · analysis.
5. **All-Reduce Acceleration via PE_IPCQ** — necessity · design · results · analysis.
6. **Fused Grouped Query Attention** (composite + PE_IPCQ) — necessity · design · results · analysis.
7. **Discussion** — which HW changes are meaningful, and why (cross-cutting).
8. **Conclusion**.
9. **Future Work — 2H** — add the FFN/MoE layer toward full LLM decoding;
analyze how compute & data should be distributed for agentic / MoE
workloads.
10. **References** *(optional)* — **external literature only** (e.g.
FlashAttention, Megatron-LM, GPT-3, Llama 3). **No ADR/SPEC entries.**
Each of §4/§5/§6 follows the same four-beat structure:
**(a) why the optimization is needed**, **(b) the design (how it works)**,
**(c) experimental results** (run/plot/measure), **(d) analysis & meaning**.
---
## Source map (ADR / SPEC → section)
Use these as **grounding sources only** — read them for content, but
**never name an ADR or SPEC in the paper**. Verify each quantitative
claim against the actual file before asserting it (CLAUDE.md: no
asserting architectural numbers from memory).
- **Modeled Hardware Configuration (§3.4)**: `topology.yaml`
(authoritative for counts, clocks, BW, link params, cost-model
constants) is the **primary** source — read it directly and build the
HW table from it. Cross-check against ADR-0003 (hierarchy/scope),
ADR-0004 (memory BW guarantee), ADR-0033 (latency model), ADR-0064
(cost-model defaults). If `topology.yaml` and an ADR disagree, the
loaded `topology.yaml` is what the sim actually runs — report that
value and flag the mismatch.
- **Platform — why/execution/latency**: `SPEC §0`, `SPEC §0.1`
(golden invariants), `SPEC §1 R1/R2/R3`, `SPEC §2.1` (graph execution);
ADR-0003 (hierarchy/scope), ADR-0007 (runtime↔sim_engine boundary),
ADR-0014 (PE pipeline / composite), ADR-0020 (2-pass data/timing),
ADR-0033 (latency model assumptions & simplifications),
ADR-0064 (structural CPU-dispatch cost model: `FIXED + logical_bytes×R`).
- **GEMM / composite**: ADR-0014 (composite command D3.2/D3.3),
ADR-0042 (tile-plan generators), ADR-0044 (GEMM eval harness),
ADR-0064 (cost model — composite's per-command BW reduction),
ADR-0065 (flat-ops `CompositeCmd`, *Proposed*).
- **All-Reduce / PE_IPCQ**: ADR-0023 (PE_IPCQ — control plane, ring
buffers, VC split), ADR-0025 (direction addressing), ADR-0032
(intercube all-reduce), ADR-0043 (allreduce eval harness),
ADR-0047 (AHBM CCL backend), ADR-0050 (CCL algorithm contract).
- **GQA / fused attention**: ADR-0060 (fused GQA kernel),
ADR-0061 (`tl.broadcast`/`tl.repeat`, *Proposed*),
ADR-0062 (lazy `tl.load`), ADR-0063 (`tl.scratch_scope`),
ADR-0064 (cost model), ADR-0065 (flat-ops composite + `softmax_merge`,
*Proposed*). KO mirrors in `docs/adr-ko/`; DDD-0060/0065 design docs if present.
- **Eval/milestone pattern**: ADR-0054 (self-contained milestone benches);
ADR-0045 (bench contract).
---
## Result map (bench → artifact → section) + run commands
All runs use the repo's `topology.yaml` unless a section needs a
specific topology. Reuse committed artifacts unless told to regenerate.
### §4 GEMM — `milestone-1h-gemm`
- Artifacts (committed): `src/kernbench/benches/1H_milestone_output/gemm/`
→ `gemm_sweep.json`, `gemm_stage_breakdown.png`,
`gemm_mac_utilization_measured.png`,
`gemm_mac_utilization_theoretical_vs_measured.png`.
- Regenerate: `kernbench run --topology topology.yaml --bench milestone-1h-gemm`
(full 24-shape sweep ≈ minutes; `MILESTONE_FAST=1` reuses committed JSON).
### §5 All-Reduce — `milestone-1h-ccl`
- Artifacts (committed): `src/kernbench/benches/1H_milestone_output/ccl/`
→ `summary.csv`, `topology.png`, per-topology
`AllReduce_LRAB_*` PNGs, buffer-kind sweep
`AllReduce_LRAB_2Dtorus_..._with_TCM_SRAM_HBM.{png,csv}`,
`comparison_mesh_vs_ring_vs_2DTorus_vs_theoretical_vs_fsim.png`.
- Regenerate: `kernbench run --topology topology.yaml --bench milestone-1h-ccl`.
### §6 GQA — `milestone-gqa-headline`
- Artifact (committed): `src/kernbench/benches/1H_milestone_output/gqa_headline/sweep.json`.
**Caveat:** this currently holds only `op_log_summary` **op-counts**
(`gemm_count`, `ipcq_copy_count`, `dma_read_count`, `dma_write_count`)
for 4 panels (single/multi-user × prefill/decode) — **no latency,
no figures**.
- Regenerate counts:
`GQA_HEADLINE_RUN=1 kernbench run --topology topology.yaml --bench milestone-gqa-headline`.
- For a real GQA **results** section, generate latency + figures via an
isolated report harness — see next.
---
## Generating missing results (isolation rule)
When a section needs an artifact that does not yet exist (today: GQA
latency and any GQA figure), create a **report-dedicated harness** under
`scripts/paper/` — never under `src/kernbench/benches/`. Rules:
- One file per concern, prefixed `paper_`, e.g.
`scripts/paper/paper_gqa_latency.py` (drives the engine, emits
latency + a baseline-vs-optimized comparison),
`scripts/paper/paper_plot_gqa.py` (renders figures from that data).
- Follow the **ADR-0054 self-contained pattern**: the harness builds its
own `GraphEngine`(s), imports the existing GQA kernel modules
(`_gqa_attention_{prefill,decode}_{short,long}.py`,
`milestone_gqa_headline.py`), runs them, and reads latency from the
engine's completion timestamps / op log. It does **not** register a
bench and does **not** import or mutate other people's benches.
- Output figures land in `docs/report/1H-codesign-paper/figures/` and,
if useful as a shared diagram, mirror into `docs/diagrams/gqa_plots/`.
- **Only measure what the current code actually runs.** ADR-0065's GQA
opt2 (`softmax_merge`, two-composite decode) is *Proposed*; if it is
not wired in production, present it as designed-not-measured and base
the measured results on the implemented path. Verify implementation
status by reading the code before claiming a measurement.
- These harnesses are derived-artifact tooling (no Phase 2 needed). **But**
if the report ever requires changing a *registered production bench*
under `src/kernbench/benches/` or any production module, that is a
non-trivial production change — STOP and follow CLAUDE.md Phase 1 →
approval → Phase 2.
---
## Procedure
### Step 0 — Toolchain (Tectonic)
Ensure a working Tectonic before any compile step. In order:
1. If `tectonic` is on PATH (`tectonic --version` succeeds), use it.
2. Else try, in order, whichever succeeds (capture output, don't hang):
- `winget install --id TectonicProject.Tectonic --silent --accept-package-agreements --accept-source-agreements`
- `scoop install tectonic`
- Direct download (no admin): fetch the latest Windows `x86_64-pc-windows-msvc`
release zip from `https://github.com/tectonic-typesetting/tectonic/releases`
via PowerShell `Invoke-WebRequest`, expand it, and place `tectonic.exe`
under `docs/report/1H-codesign-paper/build/.tools/`. Use that absolute
path for compilation.
3. Re-verify with `tectonic --version`. If all methods fail, write the
`.tex` files anyway, **report the failure**, and tell the user the
exact command to compile manually (or to use Overleaf). Do not block
the rest of the pipeline.
Tectonic fetches LaTeX packages on first compile (needs network once).
### Step 1 — Period + TOC
- Period from system date: month 1–6 → `{YYYY}-1H`, else `{YYYY}-2H`.
(Report dir is fixed at `1H-codesign-paper` for this half; adjust the
title page text with the period.)
- If `toc.md` exists, present it; else present the **Default TOC**.
- **Stop and ask the user to confirm or edit.** Incorporate edits.
Persist the agreed outline to `toc.md`. (In `/paper build` mode, skip
the stop and trust the saved `toc.md`.)
### Step 2 — Platform section (§3, incl. §3.4 HW config)
Write `sections/02-platform.tex` from the platform source map, covering
all four sub-parts:
- §3.1–3.3 (why / execution model / latency model + accuracy). Include
the cost-model formula `dispatch_cycles = FIXED + logical_bytes×R` with
the committed defaults (verify against ADR-0064 before quoting), and
explicitly address **how accurate** the latency is: what is modeled
precisely (per-edge BW occupancy, HBM pseudo-channels, flit streaming,
per-component overhead) and the known simplifications.
- §3.4 *Modeled Hardware Configuration*: read `topology.yaml` directly
and build a compact, table-driven summary — SIP/CUBE/PE hierarchy
(counts & nesting), PE clock & compute throughput, TCM/SRAM/HBM
capacity & bandwidth, NOC / inter-CUBE / inter-SIP link bandwidth &
latency, and the cost-model constants (`FIXED_PER_CMD`, `R`, composite
cap). Every number from `topology.yaml`. This is the shared config;
later sections say "the modeled hardware" and only add their own
workload params.
Do not name the source ADRs/SPEC in the prose.
### Step 3 — The three result sections (§4 GEMM, §5 All-Reduce, §6 GQA)
For each, in order, produce the four beats:
- **(a) Necessity** — from the ADR Context: the bottleneck the
optimization targets (e.g. per-command dispatch overhead for GEMM;
collective comm cost / HoL blocking for All-Reduce; KV-load-bound
decode + softmax merge for GQA).
- **(b) Design** — from the ADR Decisions: composite command tile
pipeline (ADR-0014); PE_IPCQ control/data split + ring buffers + VC
(ADR-0023); fused attention with online-softmax, lazy load, scratch
recycling (ADR-0060/0062/0063). Note Proposed-vs-Accepted status.
- **(c) Results** — reuse committed artifacts (GEMM, All-Reduce). For
GQA, generate latency + figures via the isolated `scripts/paper/`
harness (above). Copy chosen PNGs into `figures/`, `\includegraphics`
them, and build result tables from the CSV/JSON — quoting actual
numbers, with a one-line note of the run command and source artifact.
- **(d) Analysis** — interpret: what the curve/table shows, why the HW
feature produces it, where it saturates or fails to help, and the
cost (area/complexity) implied. Tie back to the cost model.
### Step 4 — Discussion, Conclusion, Future Work (§6/§7/§8)
- **Discussion**: synthesize across the three — which HW changes
(composite command, PE_IPCQ, the cost-model-exposed BW reduction)
carry their weight, and under what regimes.
- **Conclusion**: the codesign thesis, supported by the measured
results; state plainly which HW changes are meaningful.
- **Future Work (2H)**: FFN / MoE layer toward full LLM decoding; how
compute & data should be distributed for agentic / MoE workloads
(expert routing, token dispersion, all-to-all vs all-reduce, capacity).
Frame as the next codesign questions, not claims.
### Step 5 — Assemble & compile
- `main.tex`: article/IEEEtran-style preamble, title with the period,
author, executive-summary input, then `\input{sections/...}` in TOC
order, then references. Keep the preamble minimal (graphicx, booktabs,
hyperref, amsmath) so Tectonic compiles cleanly.
- Compile: `tectonic main.tex --outdir build` (or the resolved Tectonic
path) from `docs/report/1H-codesign-paper/`. Re-run if references need
a second pass. Confirm `build/main.pdf` exists and report its size.
### Step 6 — Chat report
Emit to chat (not to any file):
```
## /paper — Build Summary
**PDF:** docs/report/1H-codesign-paper/build/main.pdf (<N> pages, <size>)
**Toolchain:** tectonic <version | how installed>
**Sections written:** <list>
**Benches: reused vs regenerated**
- GEMM: <reused committed | regenerated via milestone-1h-gemm>
- All-Reduce: <reused | regenerated>
- GQA: <generated via scripts/paper/paper_gqa_latency.py | counts-only>
**Figures included:** <list of figures/*.png>
**Grounding check:** every result number traced to <artifact list>
**Proposed-not-measured items flagged:** <e.g. ADR-0065 opt2, ADR-0061>
**Open gaps / suggested next runs:** <…>
```
---
## Constraints (do not violate)
1. **TOC first.** Never write sections before the TOC is confirmed
(default mode). `toc.md` is the contract.
2. **No fabricated numbers.** Every quantitative claim traces to a
committed artifact, a fresh bench run, or an ADR/SPEC quote. Verify
architectural constants against the source file, not memory.
3. **Reuse by default.** Do not re-run a bench unless asked or the
artifact is missing/stale.
4. **Isolation.** Report harnesses live only in `scripts/paper/`; never
register a bench, never touch other people's benches.
5. **Production-change gate.** Any change to a registered bench or
production module triggers CLAUDE.md Phase 1 → approval → Phase 2.
Report tooling and `docs/report/...` are derived artifacts (exempt).
6. **Status honesty.** Mark Proposed designs as proposed; never present
an unimplemented path as measured.
7. **No internal references.** Never name `ADR-NNNN` or `SPEC §X.Y` in
prose, captions, or the reference list — they are grounding sources
only. A References section, if present, holds **external literature
only**. Hardware numbers come from `topology.yaml`.
8. **Determinism.** Sections in TOC order; same inputs → same paper.
9. **English prose**, internal-audience tone.
```
docs/adr-ko/ADR-0026-par-dppolicy-intra-device.md
+45 -2
View File
@@ -2,7 +2,7 @@
## Status ## Status
Accepted (Revision 5 — Phase 2 landed 2026-04-14, 523 passed + 1 strict xfail) Accepted (Revision 6 — cube_start 추가 2026-06-04; Revision 5 landed 2026-04-14)
## Context ## Context
@@ -30,9 +30,11 @@ class DPPolicy:
pe: Literal["replicate", "column_wise", "row_wise"] = "replicate" pe: Literal["replicate", "column_wise", "row_wise"] = "replicate"
num_pes: int | None = None num_pes: int | None = None
num_cubes: int | None = None num_cubes: int | None = None
cube_start: int = 0 # Revision 6 — SIP 내 첫 cube 인덱스; D8 참조
``` ```
제거되는 필드: `sip`, `num_sips`. 제거되는 필드: `sip`, `num_sips`.
추가되는 필드 (Revision 6): `cube_start` — D8 참조.
### D2. `ShardSpec` — structural (sip, cube, pe) 좌표, `pe_index` 완전 제거 ### D2. `ShardSpec` — structural (sip, cube, pe) 좌표, `pe_index` 완전 제거
@@ -114,7 +116,7 @@ def resolve_dp_policy(
for ls in local_shards: for ls in local_shards:
all_shards.append(ShardSpec( all_shards.append(ShardSpec(
sip=target_sip, # from caller (current_device) sip=target_sip, # from caller (current_device)
cube=cube_id, # local within SIP cube=policy.cube_start + cube_id, # Rev 6: D8 cube_start만큼 shift; 기본 0 = local within SIP
pe=ls.local_pe, # local within cube (explicit name) pe=ls.local_pe, # local within cube (explicit name)
offset_bytes=cube_offset + ls.offset_bytes, offset_bytes=cube_offset + ls.offset_bytes,
nbytes=ls.nbytes, nbytes=ls.nbytes,
@@ -224,6 +226,47 @@ KernBench는 사내 프로젝트로 call site가 한정되어 있어 한 번에
**Silent drift 차단**이 property 완전 제거의 주된 이점: global flat을 기대한 **Silent drift 차단**이 property 완전 제거의 주된 이점: global flat을 기대한
코드가 SIP-local 결과를 받아 조용히 잘못된 인덱싱을 할 가능성 제거. 코드가 SIP-local 결과를 받아 조용히 잘못된 인덱싱을 할 가능성 제거.
### D8. SIP 내 분리된 cube sub-mesh를 위한 `cube_start: int = 0` 추가
**Revision 6 추가 사항 (2026-06-04)**.
**목적**: 한 SIP 내에서 분리된 cube sub-mesh 지정 (예: 0..7과 8..15을 함께
사용). GQA Llama-70B 8-KV-group 헤드라인 타겟 — SIP당 2×4 KV-group 2개 ×
4 SIP = 64 cubes — 의 두 번째 KV-group을 기본값 0..7이 아닌 cubes 8..15에
배치하기 위해 필요.
**Semantics**: `resolve_dp_policy`가
`cube = policy.cube_start + cube_id`로 ShardSpec 생성 (cube_id는 launch 내에서
`0..num_cubes-1` 순회). 선택된 cube는 target SIP 내의
`[cube_start, cube_start + num_cubes)` 범위.
**기본값**: `cube_start = 0`은 기존 모든 호출 사이트 동작을 bit-for-bit
보존 (이전처럼 `ShardSpec.cube ∈ [0, num_cubes)`). CCL milestone bench의
full-SIP `DPPolicy(num_cubes=16)`은 그대로 cubes 0..15 생성.
**제약**: intra-device 불변성 보존.
`cube_start ∈ [0, cubes_per_sip)`이며
`cube_start + num_cubes ≤ cubes_per_sip`. SIP 경계 교차는 여전히
`ahbm.set_device(rank)` (ADR-0024)의 책임.
**왜 scalar인가** (2D `cube_mesh_origin`이나 임의 `cube_ids` list가 아닌):
consumer 커널 (예: `_attention_mesh_mlo_2d`)이 row-major contiguous cube를
가정; scalar `cube_start`는 `num_cubes`와 자연스럽게 쌍을 이룸 (범위 =
`[start, start + count)`); 기본값 0에서 마이그레이션 부담 없음. 비연속
케이스가 필요해지면 향후 더 일반적인 설계를 위에 얹을 수 있음.
**하위 호환성**: 기본값이 있는 가산적 변경. 기존 호출 사이트의 변경을
강제하지 않음. D7의 "breaking change" 입장은 원래 sip/num_sips 제거에만
해당하며, `cube_start`는 기존 호출자를 깨지 않음.
**커널 측 참고**: kernbench의 `tl.program_id(axis=1)`은 launch-local rank가
아닌 물리 cube id를 반환 (ADR-0022). `program_id(axis=1)`에서 ring 위치를
도출하는 커널은 `cube_start > 0`일 때 launch-local rank를 복원하기 위해
`cube_start`를 빼야 함 — 그렇지 않으면 out-of-bounds sub-mesh 위치를
계산하여 deadlock 발생. 참고 패턴은 `_attention_mesh_mlo_2d.py` 참조.
ADR-0022의 향후 개정에서 `program_id(axis=1)`을 launch-local rank semantics로
이동하면 커널 측 명시적 차감이 제거될 수 있음.
## Dependencies ## Dependencies
- **ADR-0024** (launcher): `set_device(rank)` 및 current-device scoping이 - **ADR-0024** (launcher): `set_device(rank)` 및 current-device scoping이
docs/adr-ko/ADR-0060-algo-gqa-fused-attention-ahbm.md
+1114
View File
File diff suppressed because it is too large Load Diff
docs/adr-proposed/ADR-0062-prog-tl-async-load-ko.md → docs/adr-ko/ADR-0062-prog-tl-async-load.md
+1 -1
View File
@@ -2,7 +2,7 @@
## Status ## Status
Proposed Accepted
> **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. Decode와 long-context > **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. Decode와 long-context
> 어텐션은 **KV-load-bound**이며, load/compute 오버랩이 지배적 레버다. 본 > 어텐션은 **KV-load-bound**이며, load/compute 오버랩이 지배적 레버다. 본
docs/adr-ko/ADR-0063-prog-tl-scratch-scope.md
+222
View File
@@ -0,0 +1,222 @@
# ADR-0063: `tl.scratch_scope` — long-context를 위한 tile 단위 scratch 재활용
## Status
Accepted
> **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. Long-context attention
> 은 수많은 K/V tile을 sweep하며, 매 tile의 중간 결과
> (`scores`, `P`, `exp`, partial `O`, …) 는 현재 kernel 호출 동안 회수되지
> 않는 fresh scratch를 새로 할당한다. PE당 1 MiB scratch 예산은 실제로
> 의미 있는 context length에 도달하기 훨씬 전에 소진된다.
## Context
### Bump allocator
`TLContext` 는 모든 math/compute output handle을 PE-local scratch pool
에서 **linear bump cursor** 로 할당한다
(`src/kernbench/triton_emu/tl_context.py`; `_scratch_alloc`):
```python
def _scratch_alloc(self, nbytes):
aligned = (nbytes + 15) & ~15
addr = self._scratch_base + self._scratch_cursor
self._scratch_cursor += aligned
if self._scratch_cursor > self._scratch_size: # default 1 << 20 = 1 MiB
raise RuntimeError("TLContext scratch overflow: ...")
return addr
```
Docstring은 cursor가 **"매 kernel 호출마다 reset된다"** 고 명시한다 —
즉 kernel entry에서만 reset, kernel body 내에서는 절대 reset되지 않는다.
모든 `tl.dot`, `tl.softmax`, `tl.exp`, `tl.sum`, `a - b`, `a * b`, … 는
fresh slice를 가져가고 kernel이 return될 때까지 아무것도 회수되지 않는다.
### 왜 GQA kernel에서 문제가 되는가
현재 milestone bench는 이 한계 **때문에** `S_q = S_kv_per_rank = 16` 으로
명시적으로 고정해 두었다
(`tests/attention/test_milestone_gqa_llama70b.py:123-148`):
> "S_q_prefill and S_kv_per_rank are deliberately small (16 each) so the
> simulator's 1 MB per-PE TCM kernel scratch is not exhausted by the
> bump-allocated handle outputs of softmax/exp/dot/sum chains over
> n_ranks ring steps."
`n_tiles`개의 tile sweep에 대해 FlashAttention은 O(`n_tiles` ×
tile당 중간값) 만큼 할당한다. 의미 있는 context에서는 overflow된다.
Flash attention의 *수학* 자체는 **O(1)** 의 live scratch만 필요로 한다 —
running `(m, l, O)` + 현재 tile의 working set — 매 tile의 임시값은
running state에 fold된 직후 dead가 되기 때문이다. 다만 allocator가 그
사실을 모를 뿐이다.
## Decision
Kernel이 할당 영역을 **재활용 가능(reclaimable)** 으로 표시할 수 있는
**scratch scope** 을 추가한다. 이를 통해 tile당 임시값은 재활용되고,
live running state(및 in-flight prefetch buffer)는 보존된다.
### D1. `tl` 표면 — context manager
```python
with tl.scratch_scope():
s = tl.dot(q, k_t) # `with` 내부의 모든 handle은
p = tl.softmax(s) # 종료 시 rewind되는 영역을 공유
o_j = tl.dot(p, v)
# ... o_j를 외부 영역에 사는 running (m,l,O) 에 fold ...
# __exit__: cursor가 __enter__ 시점 값으로 rewind
```
Semantics:
- `__enter__` 는 현재 `_scratch_cursor` 를 save-point로 기록.
- `__exit__` 는 cursor를 save-point로 복원, 내부에서 할당된 모든 것을
free.
- Scope **외부**에서 할당된 handle (running `m,l,O`, ADR-0062의
`LoadFuture` 가 보유한 prefetch buffer) 은 주소를 유지 — save-point
이전에 또는 enclosing scope에서 할당되었기 때문.
### D2. 안전 계약
Scope 내부에서 할당된 handle은 scope exit 이후 **읽으면 안 된다** —
그 바이트는 다음 scope의 할당으로 덮어쓰여질 수 있다. Flash loop은
이를 자연스럽게 지킨다: tile iteration 사이에 살아남는 값은 running
accumulator 뿐이고, 이들은 per-tile scope **외부**에서 할당되며
*내부*의 op들이 외부 주소로 쓰면서 갱신된다(merge가 새 running state
를 쓰는 방식 — D3 참조).
### D3. Running accumulator와의 상호작용
Online-softmax merge는 이전 running `(m, l, O)` 와 현재 tile의
`(m_j, l_j, O_j)` 를 읽어 새 running 값을 만든다. 새 running 값을
재활용 영역 바깥에 두기 위해, merge는 결과를 loop 시작 전 1회 할당된
**stable scratch** (per-tile scope과는 별개의 작은 고정 "running-state"
arena) 로 쓴다. 구체적으로 kernel은 두 개의 arena를 둔다:
- **persistent arena** (1회 할당): `m, l, O` (merge용 double-buffer가
필요한 경우 그것까지).
- **scoped arena** (매 tile rewind): `scores, P, exp, O_j, scale_*`.
이는 실제 flash-attention SRAM budgeting의 모습과 같다: 작은 persistent
accumulator 영역 + 재활용되는 tile working set.
#### D3.1 Persistent-arena 쓰기 mechanism: `tl.copy_to(dst, src)`
Merge op (`tl.maximum`, `tl.exp`, binary `*` / `+`) 모두 D1의 bump
cursor로부터 할당하는 `_make_compute_out(...)` 을 호출한다. 따라서
`scratch_scope` 내부에서 그 결과 handle들은 scope **내부**에 살고
`__exit__` 시점에 사라진다. D3의 two-arena 분리를 실현하려면 kernel은
**scoped 결과의 바이트를 persistent 주소로 쓰는** 방법이 필요하다.
이 gap을 메우는 primitive:
```python
def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None:
"""``src`` 의 바이트를 ``dst`` 의 주소로 복사한다 (양쪽 모두 TCM).
Shape, dtype 일치 필요. ``dst`` 는 통상 어떤 활성 ``scratch_scope`` 도
바깥에서 할당된 handle (persistent arena); ``src`` 는 scope ``__exit__``
이후에도 바이트가 살아 있어야 하는 scoped handle.
"""
```
`tl.store` (HBM 측 바이트 복사) 와 대칭이지만, running-state writeback이
op_log를 불필요한 DMA로 오염시키지 않도록 TCM 전용으로 유지한다.
**Mechanics:**
- 신규 `CopyCmd(src, dst, nbytes, data_op=True)` command.
- op_log: `op_kind="math"`, `op_name="copy"` — vector engine에서 실행.
- Latency: `pe_math._compute_ns(prod(shape))` — HBM 전송이 아니라
on-chip register writeback을 모델링.
- Emit-time 검증: `dst.shape == src.shape`, `dst.dtype == src.dtype`,
`dst.space == "tcm"`, `src.space == "tcm"`. 작성자 오류는 Phase 2
data execution 깊숙한 곳이 아니라 Phase 1에서 노출.
**호출 패턴:**
```python
m, l, O = init_running(...) # persistent (scope 외부)
for j in range(n_tiles):
with tl.scratch_scope():
... # tile 작업 (재활용)
m_new = tl.maximum(m, mj) # scoped scratch
l_new = l * scale_old + l_step * scale_step
O_new = O * scale_old + O_step * scale_step
tl.copy_to(m, m_new) # ← 새 running state를 persist
tl.copy_to(l, l_new)
tl.copy_to(O, O_new)
# exit: scoped m_new/l_new/O_new 는 사라지고, 그 바이트는
# persistent m/l/O 에 살아있다
```
`copy_to` 는 `__exit__` **이전** 에 실행되므로 `src` (scoped) 의 read는
유효하며, exit 이후에는 `dst` (persistent) 만 read되어 D2의
"exit 이후에는 읽지 말 것" 안전 계약을 만족한다.
**왜 모든 math op에 `dst=` kwarg를 추가하지 않고 전용 primitive 인가**
(검토 후 기각): `tl.maximum`, `tl.exp`, `_binary_math`, `_unary_math`,
`_reduction` 에 `dst=` 를 추가하는 것은 5개 op-family 에 걸쳐 ~25 LOC
이며, "호출은 fresh handle을 return한다" 라는 일관된 패턴을 깬다.
`tl.copy_to` 는 단일 primitive, 단일 command, 단일 executor handler —
같은 효과에 비해 최소한의 surface area.
### D4. Nesting
Scope는 nest 가능 (save-point의 stack). Inner scope exit은 inner
save-point로 rewind, outer exit은 더 멀리 rewind한다. Prefill의
"per-query-block" 외부 scope을 둘러싼 "per-KV-tile" 내부 scope을
지원한다.
## Alternatives
### A1. Temporary를 HBM에 round-trip
중간값을 HBM에 store하고 reload하여 TCM scratch를 "free" 한다. 기각:
TCM-resident streaming kernel을 HBM-bandwidth-bound로 만들어 목적에
정반대이며, op_log를 불필요한 DMA로 오염시킨다.
### A2. Tiled `tl.composite` (scheduler-관리 scratch)
`tl.composite` 는 PE_SCHEDULER 내부에서 tile당 scratch를 자동 재활용
한다. ADR-0062 A1과 마찬가지로 매력적이지만 flash-capable composite
kind (두 GEMM + carry되는 `(m,l,O)`) 가 필요하므로 훨씬 큰 변경에
종속된다. `scratch_scope` 는 작고 일반적인 primitive로 greenlet kernel
에 동일한 memory 거동을 제공한다. 두 방식은 공존 가능.
### A3. Scratch 예산 증가
`pe_tcm.kernel_scratch_mb` 를 증가. 해결책으로는 기각: O(`n_tiles`)
scratch leak이 여전히 존재하면서 context-length ceiling만 선형적으로
밀어내며, hardware 실제 모습 (실 TCM은 작고, flash attention의 핵심은
O(1) working set) 을 잘못 표현한다. 거친 조절 knob으로만 유용, 재활용의
대체재는 아니다.
## Consequences
### Positive
- 인위적인 `S = 16` validation-scale ceiling 제거; 시간/데이터 mode
양쪽에서 의미 있는 context length 가능.
- Flash attention의 O(1) working-set 특성을 충실히 모델링.
- 작고 일반적인 primitive (모든 tiled kernel이 이점).
### Negative
- Use-after-scope 버그는 stale 바이트를 조용히 읽는다. D2 계약, attention
helper 안에서 scope discipline 공유, (선택적으로) rewound 영역을
poison시키는 debug build로 완화.
- ADR-0062와 조정 필요: prefetch buffer는 tile iteration 사이에 살아
있어야 하므로 scoped arena가 아니라 persistent arena에 속한다.
## Test Requirements
1. **Recycling**: `tl.scratch_scope()` 내부의 `N` tile loop에서 peak
`_scratch_cursor` 가 `N`과 무관하게 한 tile의 footprint로 bound됨
(오늘날은 선형 증가 후 overflow).
2. **Correctness**: scope이 있는 flash sweep이 scope이 없는 동일 sweep
과 (재활용 없이도 수렴하는 작은 `N` 에서) 동일한 `O` 를 산출
(Phase 2).
3. **Long context**: scope 없이 1 MiB를 초과하는 `N` 의 sweep이 완료
(오늘날 `S=16` cap이 회피하는 그 실패).
4. **Persistent-vs-scoped isolation**: scope 외부에서 할당된 running
`(m,l,O)` 가 `__exit__` 이후에도 올바른 값을 유지.
5. **Nesting**: 중첩된 scope이 올바른 save-point로 rewind.
docs/adr-ko/ADR-0064-perf-cpu-issue-cost-model.md
+434
View File
@@ -0,0 +1,434 @@
# ADR-0064: 구조적 CPU dispatch cost 모델 (`logical_bytes` + FIXED + R)
## Status
Accepted (Revision 2)
> **ADR-0060** (AHBM GQA Fused Attention) 와 **ADR-0065** (flat-ops
> composite + 첫 stateful recipe) 의 보조 ADR. 그 hybrid 의 핵심 — "GEMM 은
> `tl.composite` 로, softmax merge 는 커널에서" — 은 **PE_SCHEDULER 로
> tiling 을 offload 해서 CPU 가 굵은 descriptor 만 issue 하고 앞서가며
> 엔진을 포화시킴** 으로 이깁니다. 그 win 이 현재 시뮬레이터에선 보이지
> 않습니다 (per-op CPU issue cost = 0).
>
> Revision 2 는 원안의 **op-type calibration 표** 를 **logical_bytes 기반
> 구조 공식** 으로 대체합니다 — op-type 별 calibration 불필요, 새 op_kind
> 도 자동 적용.
## Context
### 현재 코드
- 모든 `tl.*` op 가 cmd emit 전에 `_emit_dispatch_overhead()` 호출
(`tl_context.py:196-212`) → `dispatch_cycles > 0` 일 때만
`PeCpuOverheadCmd(cycles=dispatch_cycles)` 발행. knob 은 op kind 무관
**uniform** 이며 두 실행 경로 모두 **0** 하드코딩
(`pe_cpu.py:101` greenlet, `:195` replay).
- ⇒ 명령 issue (descriptor 구성 + scheduler 큐 push) 가 PE_CPU 에서
**0 ns** 비용.
- `PeCpuOverheadCmd` 는 PE_CPU 에서 `yield env.timeout(cmd.cycles)` 로
소비 (`kernel_runner.py:131-132`).
### uniform-zero 가 hybrid 에 부적합한 이유
ADR-0060 §1 의 핵심: composite 1 개가 `N_tiles` 분량의 GEMM tiling 을
offload → CPU 는 `O(1)` 개의 굵은 cmd 만 issue, primitive 경로는
`O(N_tiles × ops/tile)`. issue cost = 0 이면 모델은:
- primitive 경로의 *CPU 가 push 를 못 따라가 엔진 idle* 을 못 보임
- composite 의 *구성 비용 ≫ primitive 1 개이지만 ≪ 대체된 primitive 들의 합* 을 못 보임
### Revision 1 의 op-type calibration 이 과한 이유
원안은 `cost_table[kind]` (kind = `composite`, `load`, `dot`, `math`, …)
형태였습니다. 비용:
- kind 마다 값 필요 (calibration ≥ |kinds|)
- 새 kind 추가 시마다 entry 추가
- 그런데 잡으려던 *ratio* — "composite ≫ primitive, ≪ 대체된 primitive 들의 합" —
는 *op kind* 가 아니라 *cmd 가 들고 있는 필드 수* 의 함수.
cmd 의 *byte 발자국* 이 자연 proxy: N OpSpec composite 는 1 OpSpec primitive 의 ~N 배 bytes.
*고정* 부분 (큐 tail 업데이트, completion 등록, MMIO-급 latency) 은 cmd 마다.
둘 합치면: `FIXED + bytes × R`.
### 이 모델이 실제로 노출하는 것
**1차** 신호는 per-cmd FIXED 비용에 의한 **command-count 감소**.
**byte** 항은 *비정상적으로 큰* composite 가 공짜로 보이지 않도록 막는
2차 보정. 현실적 on-die 큐 대역폭 (16 B/cycle, D3) 에서 FIXED 가
decode opt2 의 opt3 (≈10 cmds/tile) vs opt2 (≈2 cmds/tile) 의 dispatch
cost 차이 중 ≥85% 차지.
이 framing 이 모델의 한계도 정의: 단일 composite 가 무한히 커지면 byte
항만으로는 HW 가 받지 못할 정도의 fused command 를 모델이 보상하는 것을
막지 못함 — 그래서 **D7** 의 descriptor 크기 cap 이 필요.
## Decision
### D1. 구조적 dispatch cost 공식
PE_SCHEDULER 로 가는 모든 PE command 가 PE_CPU 에서 dispatch cycles 소모:
```
dispatch_cycles(cmd) = FIXED_PER_CMD + cmd.logical_bytes × R
```
- `FIXED_PER_CMD` (cycles/cmd): 큐 tail 업데이트, MMIO-급 RTT,
completion-event 등록 — cmd 크기와 무관 고정 비용.
- `R` (cycles/byte): scheduler 큐에 cmd 를 직렬화하는 큐-write 대역폭.
- `cmd.logical_bytes` (int): cmd 의 *HW-논리* byte 수 — D2 룰로 계산,
Python `sys.getsizeof` 가 아님.
PE_CPU 는 기존 hook 그대로 `PeCpuOverheadCmd(cycles=dispatch_cycles(cmd))`
를 dispatch 전에 발행 — 사이클 값만 변경.
### D2. `logical_bytes` 룰
각 PE command dataclass 가 `logical_bytes: int` (property) 노출. 룰
(HW 친화적, Python overhead 무시):
| 필드 종류 | bytes |
|---|---|
| cmd framing (cmd 종류 discriminator + completion id 참조) | 4 |
| Opcode (op kind enum) | 1 |
| Enum (scope 등) | 1 |
| `TensorHandle` 참조 (address only — shape/dtype 은 descriptor table 가정) | 8 |
| Scalar (int/float) | 4 |
| Tuple length marker | 1 |
`CompositeCmd` 는 `ops` 와 `rw_handles` 재귀 합산:
```python
@property
def logical_bytes(self) -> int:
return (
4 # framing
+ 1 + sum(op.logical_bytes for op in self.ops)
+ 1 + 8 * len(self.rw_handles)
)
```
`OpSpec`:
```python
@property
def logical_bytes(self) -> int:
return (
1 + 1 # opcode + scope
+ 1 + 8 * len(self.operands) # named operand handles
+ (8 if self.out is not None else 0) # out handle
+ 1 + sum(_extra_bytes(v) for v in self.extra.values())
)
def _extra_bytes(v) -> int:
"""OpSpec.extra 값의 타입-aware byte 카운트."""
if isinstance(v, bool): return 1
if isinstance(v, (int, float)): return 4
if isinstance(v, (tuple, list)): return 1 + 4 * len(v) # shape, axes, …
if isinstance(v, str): return 1 # opcode-like tag
return 4 # default scalar
# extra 의 예시 타입:
# m, k, n int → 4 each
# reduce_axis int → 4
# shape=(64, 64) tuple → 1 + 8 = 9
# factor=1.0 float → 4
```
(`DmaReadCmd`, `MathCmd` 등도 같은 룰 — dataclass 당 property 1 개,
약 3 줄.)
**카운팅 룰 — per-op 합산, dedup 없음.** `OpSpec` 안의 각 `TensorHandle`
참조는 독립적으로 카운트 — 같은 핸들이 여러 OpSpec 에 등장하거나
`rw_handles` 에도 있어도 마찬가지. `rw_handles` block 은 cross-composite
해저드 추적 (ADR-0065 D6.3) 의 메타데이터로 `ops` 의 operand 참조와
**별도** 카운트. 중복 제거 없음. HW 와 일치: descriptor 가 각 operand
slot 을 독립 address field 로 인코딩하고 dispatcher 가 `rw_handles` 를
별도 metadata block 으로 추적. 예: `OpSpec(kind="mul_bcast", operands=
{"src_a": O, "src_b": corr}, out=O)` 는 핸들 `O` 를 *두 번* 카운트
(`src_a` 한 번, `out` 한 번); 만약 `O` 가 둘러싼 CompositeCmd 의
`rw_handles` 에도 있으면 세 번째 카운트.
### D3. Default — 일반 composite ≈ 43 ns 기준
Anchor: **DMA-staged GEMM 경로용 단일-OpSpec composite** —
`OpSpec(kind="gemm", ...)` 1 개; DMA stage 는 PE_SCHEDULER 가 operand
`space` 에서 자동 삽입 (ADR-0065 D4), **`logical_bytes` 에 등장하지 않음**
(커널이 별도 cmd 로 emit 안 함). 분해:
```
framing 4
ops tuple length 1
GEMM OpSpec 40 (opcode 1 + scope 1 + 1 + 2 handles 16
+ out 8 + 1 + extra m/k/n 12)
rw_handles tuple length 1
rw_handles content 8 (출력 RW 핸들 1 개)
─────────────────────────────
total ~54 bytes
```
Dispatch 목표 ≈43 ns. on-die producer→consumer 큐 16 bytes/cycle
(일반 on-die descriptor queue 폭) 가정.
```
FIXED_PER_CMD = 40 cycles
R = 0.0625 cycles/byte (= 16 bytes/cycle)
```
검증: `40 + 54 × 0.0625 = 43.375 cycles ≈ 43 ns` ✓
Cycle→ns 변환은 PE 노드의 기존 `clock_freq_ghz` attr (PE_MATH `_compute_ns`
가 쓰는 같은 것) 사용. cost-model knob 은 **cycle-domain 만** — clock
설정 중복 안 함.
### D4. topology config override
default 는 `pe_cpu.py` 에 내장. topology yaml 의 PE 노드 attrs 아래
`pe_cost_model:` 절로 override (cycle-domain knob 만; clock 은 PE 의
기존 `clock_freq_ghz` 사용):
```yaml
pe:
attrs:
clock_freq_ghz: 1.0 # 기존, cycle→ns 변환용
pe_cost_model:
fixed_per_cmd_cycles: 40
byte_cycles_recip: 0.0625 # = 16 bytes/cycle
max_composite_logical_bytes: 1024 # D7 — descriptor 크기 cap
```
누락된 키는 default. dispatch 공식은 PE_CPU init 시
`node.attrs["pe_cost_model"]` 에서 읽음.
### D5. Scope — 무엇이 cost 를 내고 무엇이 안 내나
| 경로 | dispatch cost? |
|---|---|
| PE_CPU → PE_SCHEDULER 의 모든 `PeCommand` | **예** |
| `PeCpuOverheadCmd` 자체 (이미 cycles 명시) | **아니오** (공식 우회) |
| PE_SCHEDULER 가 자동 생성한 Stage (DMA_READ/WRITE/FETCH/STORE) | **아니오** (PE_SCHEDULER 내부) |
| 엔진 compute latency (PE_DMA `drain_ns`, GEMM/MATH `_compute_ns`) | **변화 없음** — 엔진에 그대로 (SPEC §0.1) |
"latency 는 모델링된 컴포넌트에서" invariant 유지 — dispatch cost 는
*추가* CPU-side 시간, 엔진 시간에 fold 안 함.
### D6. 설정 가능 값; goldens 재생성
issue cost 를 0 → non-zero 로 바꾸면 **모든** bench 의 latency 변화.
이 ADR land 시 골든 latency **한 번** 재생성 — ADR-0062 D3 lazy-load 와
같은 패턴. 재생성 후 동일 calibration 이 ADR-0065 opt2 측정에 적용.
### D7. Composite 크기 cap (하드 제한 — validation error)
각 `CompositeCmd` 의 `logical_bytes` 가 **`max_composite_logical_bytes`**
(default **1024 bytes**, D4 로 override) 로 제한. `logical_bytes` 가 cap
을 초과하는 composite 는 host-side TLContext 가 emit 시점에 `ValueError`
로 **거부** — **자동 segmentation 없음**. 커널 작성자가 각 `CompositeCmd`
가 descriptor capacity 에 맞도록 recipe 를 재구성(예: 여러 개의 더 작은
composite 로 명시적으로 분할)해야 함.
**cap 이 필요한 이유.** 실제 하드웨어는 descriptor queue entry 크기,
scheduler parser buffer, command SRAM, firmware 입력의 하드 제한 보유.
cap 없으면 `FIXED + bytes × R` 모델이 하드웨어가 받지 못할 정도로 큰
fused composite 를 보상 (예: primitive 100 개를 composite 1 개로 fuse,
FIXED 1 회만 지불).
**1024 bytes 라는 *특정* 숫자의 근거.** 이는 측정된 HW 숫자가 아니라
**safe engineering limit** — kernbench 코드베이스의 현재 모든 알려진
composite (가장 큰 것이 decode opt2 의 `#2` ~322 bytes) 보다 훨씬 위이면서
미래 recipe 도 지켜야 할 *유한* descriptor capacity 를 의도적으로 표현.
값은 topology 별 override (D4); 실제 HW reference 등장 시 재캘리브레이션.
이 default 의 역할은 cap 을 *원칙으로 존재시키는 것*, 특정 HW 에 맞추는
것이 아님.
**자동 segmentation 대신 하드 에러를 택한 근거.** 초과 composite 를 자동
분할(원래 Revision 2 제안)하는 것은 inter-segment 순서, `rw_handles`
공유, completion 체이닝 등 emitter 복잡도를 추가하면서, 실질적으로는
하드웨어가 허용하는 것보다 큰 descriptor 를 요청한 커널을 덮어주는
것일 뿐. 명시적 에러로 표면화하면 emitter 가 단순해지고, HW 제약이
작업을 어떻게 분할할지 가장 잘 아는 커널 작성자에게 보임.
Decode opt2 의 `#2` composite (10 ops, ~322 bytes) 는 1024 cap 안에
편안히 — GQA workload 에 에러 없음.
### D8. Single-op-cmd fast-path (cmd 타입별 FIXED)
모든 **single-op** cmd — DMA descriptor (`DmaReadCmd`, `DmaWriteCmd`)
*그리고* single-op compute dispatch (`GemmCmd`, `MathCmd`, `CopyCmd`)
— 는 일반 control-path 비용과 분리된, 가벼운 FIXED 를 지불한다:
**`fixed_per_single_op_cmd_cycles`** (기본 **8 cycles**) 가 이들 전부에서
`fixed_per_cmd_cycles` 를 대체한다. **`CompositeCmd` 만이 일반
control-path 비용** `fixed_per_cmd_cycles` (= 40 cycles) 를 유지하는
유일한 cmd 다.
구체적으로 D1 의 공식이 다음과 같이 확장된다:
```
dispatch_cycles(cmd) = FIXED(type(cmd)) + cmd.logical_bytes × R
FIXED(CompositeCmd) = fixed_per_cmd_cycles (= 40)
FIXED(나머지 전부, 즉 모든 single-op cmd:
DmaReadCmd / DmaWriteCmd / GemmCmd / MathCmd / CopyCmd)
= fixed_per_single_op_cmd_cycles (= 8)
```
**왜 single-op cmd 에 가벼운 FIXED 를 따로 두나.** D1 은 모든 발행 cmd 에
동일한 FIXED 를 매겼다 — 가장 단순한 모델이라서. 그러나 실제로 single-op
cmd — DMA descriptor 든, 엔진에 단일 GEMM / elementwise / copy 를
dispatch 하는 것이든 — 는 40-cycle FIXED 가 모델링하던 일반 control-path
비용보다 훨씬 가볍다. scheduler 측 plan 을 만들 필요 없이, descriptor 나
instruction 하나를 엔진 queue 에 push 하는 것뿐이다:
- NVIDIA Hopper TMA: 단일 PTX 명령 (`cp.async.bulk`) 으로 bulk async DMA
발사 — SM 측에선 ~1 ISA cycle, TMA 엔진이 fan-out / chunking 자체 처리.
- NVIDIA Ampere `cp.async`: warp-level async copy ~1–2 cycles.
- 단일 MMA / tensor-core 발행 (`wgmma`, `mma.sync`): math pipe 에 명령
하나, control-path round trip 아님.
- AMD AQL packet 을 HSA queue 에 write: ~5–15 cycles.
- 일반 descriptor-ring-push 디자인 (Synopsys/Xilinx-style DMA IP): MMIO
write 들 합쳐 ~5–20 cycles.
- Tenstorrent tile descriptor 발행, Habana Gaudi TPC descriptor RAM
+ start register: descriptor 당 한 자릿수 cycles.
40-cycle FIXED 는 `CompositeCmd` 에 한해 정당하다. composite 는 단일
dispatch 가 아니기 때문이다: scheduler 가 tile-feeder plan 을 만들고,
내부 `DMA_READ → FETCH → GEMM → STORE → DMA_WRITE` stage 들에 걸쳐
per-tile read/write hazard 를 추적하고, completion handle 을 배선해야
한다. single-op `GemmCmd` 는 그런 기구 없이 엔진 발행 하나일 뿐이다 —
이를 composite 와 같은 40 cycles 로 매기면 단일 명령을 스케줄된 plan
전체와 혼동하는 것이다.
**왜 이게 커널 평가에 중요한가.** 일정한 40-cycle FIXED 는 user 가 직접
orchestrate 하는 single-op-primitive 커널의 dispatch overhead 를 부풀린다.
`K = N_K · TILE_K` 의 chunked-prefetching async GEMM 은 K-tile 마다
`DmaReadCmd` (B-chunk load) **와** `GemmCmd` (per-chunk `tl.dot`) **와**
`MathCmd` (누적 `tl.add`) 를 발행하므로, per-cmd FIXED 가 dispatch
예산을 지배한다. single-op fast-path 는 D1 이 잡으려던 *구조적 cmd 수
신호* (`N` 개의 single-op cmd 를 발행하는 recipe 가 `N` tile 을 내부로
묶는 composite 하나보다 더 비싸야 한다는 점) 를 유지하면서, 40-cycle
일반 control-path 비용을 덧씌우는 모델링 과잉 청구는 분리한다.
**왜 굳이 8 cycles 인가.** 위 조사 범위 (TMA ~1, descriptor-ring ~15) 의
중간값. 의도적인 *모델링 디폴트* — 측정된 HW 수치가 아님. 이 디폴트의 역할
은 정성적 동작에서 single-op dispatch path 를 composite control path 와
구별하는 것. Topology override 가능 (D4).
**Composite 수치에 미치는 영향.** Composite tile 내부 stage (HW tile 당
`DMA_READ`, `FETCH`, `GEMM`, `STORE`, `DMA_WRITE`) 는 scheduler 의
tile-feeder loop 가 발행하므로 host-side TLContext 의 `_charge_dispatch`
를 안 거친다. 단일 `CompositeCmd` 자체는 여전히 40-cycle control-path
FIXED 를 지불한다. D8 은 user-side single-op primitive (`tl.load` /
`tl.store` / `tl.dot` / `tl.add` / …) 에 매기는 FIXED 만 바꾼다.
Composite 측정치는 안정 유지.
## Alternatives
### A1. Revision 1 의 op-type calibration 표 유지
기각: calibration 비용이 |kinds| 에 비례, 그리고 표가 잡으려던 *ratio* 는
구조적으로 cmd 크기의 함수. 구조 공식은 같은 정성적 동작을 N 개가 아닌
2 개의 calibratable 숫자로 달성.
### A2. byte-only 공식 (FIXED 항 없음)
기각. FIXED = 0 이면 opt2 (ADR-0065 Option Y) 가 opt3 를 **이기지 못함**
— per-tile dispatch *총 bytes* 가 비슷 (opt3 ≈ 232, opt2 ≈ 380); win 은
전적으로 *per-cmd fixed cost 횟수 감소* 에서 옴. byte-only 는 모델이
보여야 할 신호를 지움.
### A3. PE_SCHEDULER 에 dispatch 비용 부과
기각: saturation 질문은 *"CPU 가 descriptor 를 push 하는 속도가 엔진을
계속 바쁘게 할 수 있는가?"* — **PE_CPU** 의 issue-bandwidth 특성.
scheduler 에 부과하면 CPU back-pressure 모델링 안 됨.
### A4. DMA program/setup 시간을 별도 fixed per-descriptor 비용으로 모델
연기: 처음엔 descriptor-program 비용을 **issuing op 의** dispatch cost 에
fold. calibration 이 분리 필요성을 보이면 PE_DMA fixed setup 으로 분리.
## Consequences
### Positive
- hybrid 의 CPU-offload / saturation win (ADR-0060 §1) 이 **측정 가능**,
구조적으로 정직한 모델 (calibration 표 없음).
- 새 op_kind 추가 (예: ADR-0065 의 `softmax_merge` 8-step recipe) 가
*zero cost* — 같은 공식 자동 적용.
- 하드웨어에 더 충실 (queue-head MMIO RTT + queue-write 대역폭).
### Negative
- **모든** bench goldens 변화 → 일회성 재생성 (D6); CI 골든 fixtures 업데이트.
- 두 calibration knob (FIXED, R) 필요; default 는 문서화된 가정에 기반 —
HW reference 없는 동안엔 절대 latency 를 잠정 취급, **ratio** 만 방어.
- 각 PE command dataclass 에 작은 `logical_bytes` property 추가.
## Open review items
1. **FIXED 와 R 의 calibration 출처.** default 는 "일반 composite ≈ 43 ns
+ on-die 큐 16 bytes/cycle"; on-die descriptor 큐로 합리.
HW reference 등장 시 재방문.
2. **대형 composite 에서 Scheduler plan-gen 비용.** 0 유지 — D5 가
PE_SCHEDULER 의 plan 생성을 dispatch 공식 밖에 둠. D7 cap (1024 bytes
≈ 30–35 OpSpec) 이 최악을 묶지만, cap 근처 composite 도 현재
zero-cost 모델에서 1-op composite 와 같은 scheduler 비용. stress test
(대형 composite microbench) 가 scheduler-bound 행동을 보이면
per-op-count `overhead_ns` 노출.
3. **Override 위치.** topology yaml 의 PE 노드 attrs 아래 `pe_cost_model:`
block — 모든 knob 을 한 곳에, 리뷰 가능. clock 은 PE 의 기존
`clock_freq_ghz` 에서 — 여기 중복 안 함.
4. **경로 parity.** greenlet (`_execute_legacy`, `kernel_runner`) 와
replay 둘 다 같은 cost model 읽어야 함. 검증.
5. **R 에 대한 결론의 민감도.** opt2 < opt3 가 합리적 `R` 범위 (큐 대역폭
가정) 에서 유지되어야 함. 민감도 sweep 이 Test Requirements (#9) 의
일부.
## Test Requirements
테스트는 specific 숫자 anchor 가 아닌 **공식 (D1)** 에 대해 작성 —
calibration 이 바뀌거나 OpSpec/CompositeCmd 필드가 추가되어도 유효.
1. **공식 보존.** 임의 `CompositeCmd` `c` 에 대해 PE_CPU 의 기록된
dispatch overhead 가 `FIXED_PER_CMD + c.logical_bytes × R`
(floor/round-off 의 ±1 cycle 내). 여러 composite parametrize:
1-OpSpec GEMM composite, 5-op MATH chain, 10-op recipe composite.
Default-calibration 숫자 (1-OpSpec composite 의 anchor ≈43 ns) 는
*informative* 참고, 테스트 gate 가 아님 — gate 는 공식 등가.
2. **정성적 ratio (robust).** opt3 per-tile PE_CPU dispatch 가 opt2
per-tile dispatch 를 최소 2× 차이로 엄격히 초과 — `opt3 > 2 × opt2`.
Default-calibration 모델은 ≈4× 예측; gate 는 느슨한 2× 한계라
calibration 이동 (예: HW reference 가 default 대체) 시 깨지지 않음.
Informative 숫자 — ADR-0065 §verification 와 DDD-0065 §11 의 모델
기대치 참조.
3. **Override 경로.** topology yaml 의 `pe_cost_model:` block 이 per-PE
dispatch cost 변경; block 누락 시 default 복귀. #1 의 공식 등가가
override 값에서도 성립.
4. **`PeCpuOverheadCmd` 우회.** 수동 `tl.cycles(n)` 는 정확히 `n` cycles,
`n + dispatch_cycles(...)` 아님.
5. **double-count 없음.** PE_DMA `drain_ns`, PE_GEMM/MATH `_compute_ns` 가
pre-ADR 값과 동일.
6. **결정성.** 동일 입력 → 동일 op_log + latency (SPEC §0.1).
7. **경로 parity.** greenlet 과 replay 가 동일 커널에 대해 동일
dispatch-cycle 회계.
8. **Composite 크기 cap (D7).** `logical_bytes >
max_composite_logical_bytes` 가 될 composite 는 emit 시점에 `ValueError`
(segmentation 없음). cap 이내 composite 는 정상 emit.
9. **민감도 (정성적).** `R ∈ {0.25, 0.0625, 0.03125}` cycles/byte 에서
3 지점 모두 `opt3 > opt2`. 방향 (R 감소 시 ratio 단조 증가) 도 단언,
단 절대 ratio 값은 *불필요*.
## Migration
ADR-0064 Revision 2 는 단일 PR 로 land:
- 각 `PeCommand` dataclass 의 `logical_bytes` property (D2 의 타입-aware
extra 카운트)
- `pe_cpu.py` dispatch 경로의 공식 적용
- PE_CPU init 의 `pe_cost_model:` override read (cycle-domain knob +
`max_composite_logical_bytes`)
- **composite 크기 cap (D7)** — TLContext-side 하드 cap,
`max_composite_logical_bytes` default 1024: cap 초과 composite 는 emit
시점에 `ValueError` (자동 segmentation 없음). 기존 bench 는 트리거
안 됨 (현재 최대 composite 가 ~322 bytes), 하지만 recipe 가 커질 때
descriptor-capacity 제한이 강제되도록 체크 land.
- 일회성 골든 재생성
land 후 ADR-0065 가 위에 쌓임 — 기존 bench 추가 골든 churn 없음
(ADR-0065 는 기존 경로의 `CompositeCmd` 의미 보존 refactor; opt2 만 새 bench).
docs/adr-ko/ADR-0065-prog-flat-ops-composite-softmax-merge.md
+440
View File
@@ -0,0 +1,440 @@
# ADR-0065: 평평한 ops `CompositeCmd` + 첫 stateful recipe (`softmax_merge`)
## Status
Accepted
> **ADR-0060** (AHBM GQA Fused Attention) 의 보조 ADR. §5.6 / §8 item 4 의
> carve-out 을 정확히 구현: decode opt2 = (#1 기존 Q·Kᵀ GEMM composite) +
> (#2 softmax + P·V + online-softmax merge 를 담은 단일 composite). ADR-0060
> §5.6 의 "ex_composite" / "flash_pv_merge" 명칭은 *retire* — 본 ADR 이
> 정식 모양을 확정: 기존 `tl.composite` 입구를 그대로 사용하면서 두 가지
> 구조적 추가 — (a) 평평한 `ops` 튜플, (b) MATH micro-op 시퀀스로
> 펼쳐지는 첫 stateful recipe `softmax_merge`.
## Context
### ADR-0060 가 이 ADR 에 남긴 것
ADR-0060 §5.6 는 opt3 (software pipelining) 를 *지금 ship* 으로 권고,
opt2 는 — per-tile dispatch 적음, K-before-V DMA priority,
**ADR-0064 의 cost model 이후에만 측정 가능** — 후속으로 carve out.
§8 item 4 가 carve-out 의 sizing: "재방문 시 **두 개** composite 로 분할 —
`#1` = Q·Kᵀ (기존 composite + `scale`), `#2` = softmax + P·V + 온라인-softmax
누산기 merge". 본 ADR 이 `#2` 를 구현.
### composite 가 구조 변경이 필요한 이유
현재 `CompositeCmd` 는 `(op, a, b, out_addr, ops=(head, *epi))` 모양 + 암묵 컨벤션:
- `op` 가 head 엔진 선택 (gemm | math)
- `ops[0]` 는 head OpSpec, `ops[1:]` 는 epilogue OpSpec 으로 head *이후*
per-output-tile 또는 per-K-tile 에 실행 (scope 결정)
decode opt2 는 MATH op 이 head GEMM **이전** 에 실행 (`m, l, O` 의
online-softmax 갱신 + GEMM 입력 `P` emit) 이 필요. 현 모양엔 자리 없음.
이전에 검토 후 기각된 경로:
- softmax_merge 를 #1 (Q·Kᵀ) 의 *epilogue* 로 두기 — `Sj` 를 읽고 *다음*
GEMM (#2 P·V) 가 소비할 `P` 를 쓰는 구조라 순환.
- 단일 mixed-engine recipe `flash_pv_merge` 가 MATH + GEMM + MATH 흡수 —
engine boundary 위반, PE_SCHEDULER 복잡화.
깔끔한 모양: **command 포맷에서 head/epilogue 구분 제거**. `ops` 를 평평한
순서 튜플로, 각 op 의 `scope` + position 이 tile-loop 배치 결정. "prologue"
개념은 *사용자 API* 의 `tl.composite(...)` 에 ergonomic 으로 남되 —
**HW command 는 보지 않음**.
### recipe 가 (generic op list 가 아닌) 이유
decode opt2 의 `#2` MATH chain 은 8 단계 고정 구조 (online softmax merge).
커널 작성자에게 8 단계 직접 와이어링은 verbose + correctness hazard
(단계 순서 의존). *Recipe* — TLContext 가 8 단계로 (주소 미리 채워)
펼치는 단일 명명 op (`softmax_merge`) — 가 커널을 간결히 + 구조를 단일
source 에 보존. recipe 표는 TLContext (compiler analog) 에 위치;
PE_SCHEDULER 는 안 읽음. D5 boundary 참조.
## Decision
### D1. `CompositeCmd` 는 평평한 ordered op 리스트
```python
@dataclass(frozen=True)
class CompositeCmd:
completion: CompletionHandle
ops: tuple[OpSpec, ...] # ordered MATH/GEMM ops
rw_handles: tuple[TensorHandle, ...] = () # cross-composite hazard
data_op: bool = True
@property
def logical_bytes(self) -> int: ... # ADR-0064 D2
```
이전의 `(op, a, b, out_addr, out_nbytes, math_op)` 필드는 `ops` 에 흡수.
`ops` 의 첫 OpSpec 이 `kind == "gemm"` 이면 head GEMM; 앞의 OpSpec 들이
pre-loop, 뒤의 OpSpec 들이 `scope` 에 따라 post-loop / per-tile epilogue.
### D2. `OpSpec` 이 named operand 보유
```python
@dataclass(frozen=True)
class OpSpec:
kind: str # "gemm" | "rmax" | ...
scope: Scope # KERNEL | K_TILE | OUTPUT_TILE
operands: dict[str, TensorHandle] # named — D5 참조
extra: dict[str, Any] # scalars, axes, m/k/n 등
out: TensorHandle | None = None # 명시적 write-back handle
@property
def logical_bytes(self) -> int: ...
```
named operand 로 PE_SCHEDULER 가 (positional 컨벤션 없이) 엔진 port 에
입력 라우팅 (예: GEMM 은 `operands["a"]`, `operands["b"]`).
### D3. Position + scope 가 tile-loop 배치 결정
**Semantics 분리.** *Position* 이 **phase** (tile-loop 수명주기 어디서
실행되는지) 결정; *scope* 가 그 phase 안에서의 **repetition** 결정.
`Scope.KERNEL` 의미는 "**`CompositeCmd` invocation 당 1 회**" — kernel
launch 당 1 회 아님. 그 1 회 실행이 op 의 시퀀스 내 *position* 에서
발생. 같은 KERNEL 값이 OpSpec 이 GEMM op 의 앞/위치/뒤에 있느냐에 따라
*phase* 가 달라짐.
PE_SCHEDULER 가 `cmd.ops` 에서 GEMM op 검색 (composite 당 ≤ 1, D6 참조).
인덱스 `g`:
| OpSpec position | scope | Phase | plan 내 배치 |
|---|---|---|---|
| `0 .. g-1` | KERNEL | pre-tile-loop | tile loop 진입 전 1 회 |
| `g` | KERNEL | head GEMM | `extra["m"], extra["k"], extra["n"]` 로 tile loop drive |
| `g+1 .. ` | K_TILE | in-accumulation | per K-tile epilogue (K 누산 loop 안) |
| `g+1 .. ` | OUTPUT_TILE | per-output-tile | per (m, n) tile epilogue (K 누산 후) |
| `g+1 .. ` | KERNEL | post-tile-loop | tile loop 종료 후 1 회 |
GEMM op 가 없는 composite (예: MATH-only) 는 모든 op 를 KERNEL-scope
직렬로 처리 — phase 가 "순서대로 single-shot" 으로 collapse.
### D4. PE_SCHEDULER 가 operand `pinned` 플래그 보고 DMA 자동 삽입
PE_SCHEDULER 가 GEMM 의 `operands` 검사. 각각:
- **not `pinned`** (데이터가 HBM 에 있어 스트리밍 필요) → FETCH/GEMM Stage
앞에 DMA_READ Stage 삽입
- **`pinned`** (이전 `tl.load` 로 이미 TCM 에 staged, 또는 recipe 의 TCM
scratch / primary-out `P`) → DMA Stage 없음, in-place 소비
커널은 composite operand 의 명시적 DMA cmd 를 절대 emit 안 함;
PE_SCHEDULER 가 핸들에서 추론. `tl.ref` operand 는 not pinned (DMA_READ);
`tl.load` 결과와 recipe scratch / primary-out 은 pinned (in-place).
> **as-built 노트 (P3).** D4 의 초안은 DMA 결정을 `space` 태그
> (`hbm`→DMA, `tcm`→in-place) 에서 도출하며 현재 `space` 값
> (오늘 `tl.load`=`hbm`, `tl.ref`=`tcm`) 을 뒤집어야 했음. 구현은 기존
> **`pinned` 플래그**를 DMA 신호로 그대로 유지 — `pinned` 이 이미
> "이미 TCM 에 있음, DMA skip" 을 인코딩하고, `space` 뒤집기는 기존 bench
> Stage 시퀀스를 바꾸면서 기능적 이득이 없음. `space` 통일 + `pinned`
> 폐기는 후순위 cleanup 으로 연기.
**Prologue recipe op 은 TCM-only (Phase 1 한계, TLContext emit 시 강제).**
*prologue recipe* (비 GEMM) OpSpec 의 모든 operand 는 **반드시** TCM 상주;
HBM 핸들을 recipe operand 로 전달하면 emit 시 validation error. Phase 1
에서 DMA staging 을 head op 의 operand 로 제한. **head op 자체** (gemm
또는 math) 는 기존 DMA-staged-from-HBM 동작 유지 — D4 의 TCM-only 룰은
head 에 적용 **안 됨**. 이 invariant 는 D6 #7 에 재기술.
### D5. RECIPE_DESCRIPTORS — TLContext 내부, `pe_commands.py` 가 아님
Recipe 는 새 TLContext-adjacent 모듈
(`src/kernbench/triton_emu/tl_recipes.py`) 에 위치. PE_SCHEDULER 는
import **안 함**. 첫 recipe:
```python
@dataclass(frozen=True)
class RecipeDescriptor:
operands: dict[str, str] # name → "R" | "RW"
primary_out: PrimaryOutSpec | None # 암묵 primary output 의 type/shape 룰
tile_alignment: Literal["single_shot", "tile_aligned"]
internal_scratch_bytes_fn: Callable[..., int]
engine_seq: tuple[EngineOp, ...] # TLContext 가 평평한 OpSpec 으로 펼침
RECIPE_DESCRIPTORS["softmax_merge"] = RecipeDescriptor(
operands={"s": "R", "m": "RW", "l": "RW", "O": "RW"},
primary_out=PrimaryOutSpec(
from_shape="s", from_dtype="s", transform="identity",
),
tile_alignment="single_shot",
internal_scratch_bytes_fn=lambda G, TILE, d, bpe: bpe * (
G + G + G + G * TILE + G # m_loc + m_new + corr + P + l_loc
),
engine_seq=(
EngineOp("MATH", "rmax", src="s", dst="m_loc", reduce_axis=-1),
EngineOp("MATH", "max_elem", src_a="m", src_b="m_loc", dst="m_new"),
EngineOp("MATH", "exp_diff", src_a="m", src_b="m_new", dst="corr"),
EngineOp("MATH", "exp_diff", src_a="s", src_b="m_new", dst="P", bcast_axis=0),
EngineOp("MATH", "rsum", src="P", dst="l_loc", reduce_axis=-1),
EngineOp("MATH", "fma", src_a="l", src_b="corr", src_c="l_loc", dst="l"),
EngineOp("MATH", "mul_bcast", src_a="O", src_b="corr", dst="O", bcast_axis=1),
EngineOp("MATH", "copy", src="m_new", dst="m"),
),
)
```
TLContext 가 `tl.composite(prologue=[{"op": "softmax_merge", ...}],
op="gemm", ...)` 에서:
1. operand R/RW 를 `RECIPE_DESCRIPTORS["softmax_merge"]` 와 검증.
2. scratch (m_loc, m_new, corr, P, l_loc) 를 scratch_scope helper 로 할당 (ADR-0063).
3. primary output P 의 shape 을 `s.shape` 에서 derive (identity).
4. `engine_seq` 를 8 개의 평평한 MATH OpSpec 으로 펼침 (모든 주소/크기 채움).
각 OpSpec 의 `scope = KERNEL`.
5. **Auto-bind (충돌 검사, D6.6).** head GEMM 이 `operands["a"]` 를
*명시하지 않은* 경우에만 `operands["a"] = P_handle` auto-bind.
커널이 `a` 를 명시적으로 제공했다면 validation error — 모호한
primary-output 대체.
6. `CompositeCmd(ops=(8 MATH + 1 GEMM + epilogue), rw_handles=(m, l, O))` emit.
**RECIPE_DESCRIPTORS 는 HW 경로 어디에도 안 나타남**. PE_SCHEDULER 는
평평한 ops 리스트만 봄.
### D6. Invariants
1. **GEMM 수.** `0 ≤ count(op.kind == "gemm" in cmd.ops) ≤ 1`. GEMM 0 개면
MATH-only composite.
2. **softmax_merge 가 V operand 없음.** `RECIPE_DESCRIPTORS[
"softmax_merge"].operands` 에 HBM-resident handle 없음 → prologue 동안
DMA 삽입 없음 → V DMA 는 GEMM head 시작 시에만 발생 → decode opt2 의
#1 / #2 간 K-before-V DMA priority 자연 강제.
3. **Cross-composite strict FIFO.** PE_SCHEDULER 가 in-flight `rw_handles`
추적. 신규 composite 의 `rw_handles` (또는 op 입력) 가 in-flight 와
교차하면 모든 이전 composite 완료까지 대기 — out-of-order reorder 없음.
4. **legacy caller 후방 호환.** 사용자 API `tl.composite(op="gemm", a, b,
epilogue=[...])` 보존. TLContext 가 내부적으로 평평한 ops `CompositeCmd`
로 lowering.
5. **PE_SCHEDULER recipe-free.** PE_SCHEDULER 는 RECIPE_DESCRIPTORS import
안 함, "softmax_merge" 의 존재 모름, recipe 별 코드 분기 없음.
6. **암묵 operand 대체 없음 (auto-bind 충돌).** prologue recipe 가
`primary_out` 선언하고 head op 의 auto-bind 대상 operand (예: GEMM
`a`) 도 *커널이 명시적* 으로 제공하면, TLContext 가 emit 시 validation
error. 커널은 recipe 가 바인딩하게 두거나 (operand 생략) **혹은**
명시적으로 지정 (prologue 생략, 또는 `primary_out` 없는 recipe 사용)
해야 함. "어느 값이 이기느냐" 의 모호함 방지.
7. **Prologue MATH operand TCM-only.** D4 의 재기술: *prologue recipe*
(비 GEMM) OpSpec 의 모든 operand 는 TCM 상주. head op (gemm 또는 math)
은 예외 — HBM 스트리밍 허용. Phase 1 에서 TLContext emit 시 강제.
### D7. Boundary 요약 (compiler vs scheduler vs engine)
```
HOST DEVICE
TLContext (compiler analog) PE_SCHEDULER (scheduler/dispatcher)
- RECIPE_DESC 와 검증 - 평평한 ops 리스트 스캔
- primary_out shape derive - GEMM 식별 (≤ 1)
- scratch 할당 - position-기반 stage 배치
- recipe → 평평한 OpSpec 펼침 - space 보고 DMA stage 자동 삽입
- rw_handles 계산 - strict-FIFO RW 해저드 추적
- CompositeCmd emit - Stage 리스트 emit
imports: RECIPE_DESCRIPTORS imports: recipe 관련 없음
ENGINES (PE_MATH / PE_GEMM / PE_DMA)
- Stage.params 읽음 (op_kind, addresses, n_elements)
- 기존 latency 모델 변화 없음
imports: recipe 관련 없음
```
### D8. Composite 출력 handle + 출력-space DMA
`tl.composite(...)` 는 **출력 `TensorHandle`** 을 반환(단순 `CompletionHandle`
아님) — composite 결과를 `tl.dot` 결과처럼 downstream op 에 먹일 수 있음.
handle 의 `pending` 필드가 composite 완료를 참조; downstream consumer 가
auto-await(`_await_pending`, ADR-0062 lazy 패턴), `tl.wait(handle)` 동작.
**`out` 은 항상 `TensorHandle`** (raw 주소 아님) — 위치는 커널의 선택, 추측 안 함:
- **HBM 출력** → `out=tl.ref(addr, shape)`. `tl.ref` 가 `space="hbm"` handle
반환 → composite 가 DMA_WRITE 로 write-back. STORE + DMA_WRITE 는 composite
파이프라인 **내부**에 유지 — fused op 의 일부이지 별도 `tl.store` 아님.
- **In-place TCM** → `out=<기존 TCM handle>` (예: recipe 누적기 `O`).
`space="tcm"` → STORE 만, 결과 TCM 에 남음 (chainable).
- **미지정** → TLContext 가 **TCM scratch 자동 할당** (`tl.dot` 과 동일) 후 반환.
편의상 `out_ptr: int` 는 **HBM 출력 shorthand** 로 허용 —
`out=tl.ref(out_ptr, <출력 shape>)` 와 동치(컴파일러가 `space="hbm"` handle 로
감쌈). 출력 미지정 시 scratch 주소는 커널이 아니라 컴파일러가 소유.
(`tl.ref` 는 `space="hbm"` 반환 — HBM 데이터 참조이므로. operand-**입력** DMA
결정은 D4 대로 `pinned` 기반이라 이 라벨이 입력 스트리밍에 영향 없음.)
출력 handle 의 **`space` 가 타일 루프의 write-back 을 결정** (D4 의 operand
`pinned` 규칙의 출력판):
| `out.space` | 타일 루프 write-back |
| --- | --- |
| `"tcm"` | STORE 만 — TCM 에 남음 (chainable; **DMA_WRITE 없음**) |
| `"hbm"` | STORE + DMA_WRITE — 최종 HBM 출력 |
이로써 출력이 TCM 에 남는 composite(예: recipe 누적기)가 합법이 되고,
TCM-scratch 주소가 DMA 엔진의 PA 디코더(고비트 scratch 주소를 거부)로 가는
일을 막음. 기존 벤치는 HBM 출력을 `tl.ref` 로 감싸므로 write-back
불변(byte-equal).
## Alternatives
### A1. tile 당 3 composite (prologue 개념 없이)
#2 를 `softmax_merge` + `gemm(P·V)` + `add(O)` 의 3 개로 분할. 장점:
flat-ops 재구조화 없음 (`add` 가 기존 epilogue 경로 재사용). 단점: tile
당 dispatch 3 회 (대신 2) — ADR-0060 §8 item 4 sizing ("**두 개** 로
분할") 과 어긋남; ~32 ns × N_tiles 의 fixed-cost 절감 소실. **기각.**
### A2. Mixed-engine recipe (`flash_pv_merge`)
MATH + GEMM + MATH 흡수하는 단일 recipe (ADR-0060 §5.6 원안 표현). 장점:
가장 적은 dispatch (1 composite). 단점: PE_SCHEDULER 가 engine-경계
recipe 알아야 함, "엔진은 Stage.op_kind 만" boundary 위반, recipe 표가
ISA-급이 되어 HW 인터페이스 오염. **기각.**
### A3. Generic op-list DAG (`tl.ex_composite([...])`)
사용자 정의 임의 op 시퀀스 + 명명 슬롯 + 데이터플로 분석. 장점: 최대
유연성. 단점: 단일 user (softmax_merge) 가 generic DAG 언어를 정당화 못함;
PE_SCHEDULER 에 데이터플로 분석 필요; ADR-0060 §8 item 4 의 "inner loop
흡수" 트랩 재등장. **기각 (Simplicity First).**
### A4. RW-aware scheduler reorder (strict FIFO 가 아님)
PE_SCHEDULER 가 `rw_handles` 교차 없으면 후속 composite 가 in-flight 를
추월 허용 (예: 다음 tile 의 #1 가 현재 tile 의 #2 와 겹침 — #1 은
m/l/O 안 만짐). 장점: 더 많은 overlap. 단점: scheduler state 증가, 이득
한정 (다음 #1 은 어차피 K DMA 대기). **연기** — strict FIFO 먼저;
RW-aware reorder 는 후속 ADR 가능.
### A5. softmax_merge 를 #1 (Q·Kᵀ) 의 epilogue 로 embedding
오답: softmax_merge 의 출력 `P` 는 #2 의 P·V GEMM 의 입력, #1 의
epilogue *이후* 실행. embedding 하면 #2 가 자기 입력을 #1 의 epilogue
scratch 에서 읽는 형태 — 순환 tile-loop 의존성. **기각 (incorrect).**
## Consequences
### Positive
- ADR-0060 §5.6 / §8 item 4 carve-out 정확 구현: opt2 = tile 당 2
composite, K-before-V DMA priority 자연.
- composite contract 균일화: scope-기반 배치의 평평한 ops 리스트. HW cmd
에서 "head" 와 "epilogue" 구조 구분 없음.
- DMA 가 operand `space` 에서 자동 추론; 커널 표면 단순화.
- HW 인터페이스 (`CompositeCmd` 모양) 가 recipe 증가에도 작게 유지 —
recipe 펼침이 host-side.
- 새 fused 패턴 (linear+gelu+norm, rmsnorm+linear 등) 은 `RECIPE_DESCRIPTORS`
entry 추가만으로 가능; PE_SCHEDULER 변경 없음.
### Negative
- CompositeCmd 구조 변경 — 모든 현재 caller 의 refactor. 의미 보존
(D6.4); ADR-0064 Revision 2 land 후 기존 골든 불변.
- `OpSpec.operands` 가 positional `tuple[Any, ...]` 에서
`dict[str, TensorHandle]` 로 변경 — 기존 epilogue lowering 접촉.
- 새 TLContext 코드: recipe 펼침 + scratch slot 할당 + primary_out
auto-bind. ~80–120 LOC.
- PE_SCHEDULER 의 `_generate_plan` 에 position-기반 스캔 + DMA 자동 삽입
+ strict-FIFO RW tracker. ~60–100 LOC.
- **strict-FIFO RW 해저드 추적이 안전하지만 보수적.** 두 composite 가
`rw_handle` 을 공유할 때 신규는 모든 이전 겹치는 composite 가 *완전히
완료* 될 때까지 대기 — 원칙적으로 overlap 가능한 경우에도
(예: 다음-tile #1 = Q·Kᵀ 가 m/l/O 안 만지므로 현재-tile 의 #2 와
overlap 가능하지만, 같은 핸들을 만지는 후속 path 도 FIFO 로 직렬화).
더 똑똑한 scheduler 대비 **composite-level overlap 을 과소 노출**.
RW-aware reorder (A4) 가 연기된 개선.
## Open review items
1. **identity 를 넘는 recipe shape derivation 룰.** 미래 recipe 가
non-identity transform (예: 전치된 primary_out) 필요할 수 있음.
`PrimaryOutSpec.transform` 필드가 forward-compat; 필요 시 새 transform 추가.
2. **Recipe scratch allocator** 의 ADR-0063 `scratch_scope` 통합. 초기
배선: TLContext 가 활성 `scratch_scope` 가 있으면 그 안에서 할당,
없으면 kernel-scope persistent slot.
3. **`tl_recipes.py` 위치.** `triton_emu/` 아래 새 모듈. recipe 지식을
compiler analog (TLContext) 와 함께 위치.
4. **GEMM `out=TensorHandle`.** 기존 `out_ptr: int` 와 함께 새 형태.
둘 다 허용; handle 형태가 신규 코드에 권장 (handle identity 로
strict-FIFO RW 추적 가능).
## Test Requirements
1. **CompositeCmd 평평화 refactor — 의미 보존.** 기존 모든 bench 의
op_log 가 refactor 전후로 byte-equal (legacy API `tl.composite(
op="gemm", a, b, epilogue=[...])` 가 같은 Stage 시퀀스로 lowering).
2. **softmax_merge recipe lowering.** TLContext 호출 `tl.composite(
prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}],
op="gemm", b=V_ref, out=O, epilogue=[{"op": "add", "other": O}])` 가
정확히 10 ops (8 MATH + 1 GEMM + 1 MATH(add)) + `rw_handles == (m, l, O)`
의 `CompositeCmd` 생성.
3. **K-before-V DMA priority invariant.** decode opt2 의 op_log 에서
#2 의 MATH prologue 동안 V 관련 DMA 없음 (#2 의 V DMA 는 GEMM head
시작 시에만).
4. **Strict-FIFO RW 직렬화.** `rw_handle` 공유하는 두 연속 composite 가
dispatch 순서대로 완료; stage 가 interleave 안 함.
5. **`pinned` 에서 DMA 자동 삽입.** not-`pinned` operand (예: `b=tl.ref(...)`)
의 GEMM composite 가 DMA_READ Stage emit; `pinned` operand (예: recipe
의 TCM scratch / primary-out, 또는 `tl.load` 결과) 는 emit 안 함.
(as-built D4 노트: DMA 결정은 `space` 태그가 아니라 기존 `pinned`
플래그 사용.)
6. **opt2 가 opt3 와 수치 동등.** data mode 에서 opt2 의 최종
`(m, l, O)` 가 opt3 와 fp tolerance 안.
7. **opt2 dispatch ratio (ADR-0064 Rev2 이후) — robust.** opt3 vs opt2
per-tile PE_CPU dispatch cycles 가 `opt3 > 2 × opt2` 만족
(정성적 gate, calibration-독립). Default-calibration 모델 기대치는
≈ 4.0× (FIXED=40 cycles, R=0.0625 cycles/byte) — DDD-0065 §11 의
모델-유도 숫자 참조; 테스트 gate 는 느슨한 `> 2×` 한계라
calibration 변경에도 살아남음. Ratio 가 FIXED-dominated —
command-count 감소가 1차 신호임을 반영.
8. **GEMM-count invariant.** GEMM OpSpec 두 개 가진 composite 가
TLContext emit 시 validation error.
9. **Composite 크기 cap 안에 (ADR-0064 D7).** Decode opt2 의 `#2`
composite (10 ops, ~322 logical bytes — operand 핸들의 per-op 합산)
가 default `MAX_COMPOSITE_LOGICAL_BYTES=1024` 안에 편안히 —
segmentation 없음. recipe lowering 테스트가 `cmd.logical_bytes
< 1024` 단언.
10. **Auto-bind 충돌.** `tl.composite(op="gemm", a=A_handle,
prologue=[{"op": "softmax_merge", ...}], ...)` 처럼 softmax_merge 가
`primary_out` 선언하는데 `a` 도 명시되면 emit 시 validation error
(D6.6).
11. **MATH operand TCM invariant.** `tl.composite(op="math",
a=tl.ref(addr, shape), ...)` (HBM-resident 핸들을 MATH operand 로
전달) 가 emit 시 validation error (D6.7).
## Dependencies
- **ADR-0060** §5.6, §8 item 4 — 본 ADR 이 구현하는 carve-out.
- **ADR-0064 Revision 2** — dispatch cost model; 먼저 land, opt2 의
fewer-issues win 을 측정 가능하게 함.
- **ADR-0063** `scratch_scope` — recipe intermediate 가 안에 할당;
`m, l, O` 는 persistent arena 로 밖에 할당 (DDD-0060 §6.2 참조).
- **ADR-0046** `tl_context` contract — `prologue=[...]` kwarg,
`out=TensorHandle`, RECIPE_DESCRIPTORS lookup 으로 확장.
- **ADR-0042** tile plan generators — (재설계 아닌) 확장 — 평평한 ops
리스트 처리 + DMA 자동 삽입.
- **ADR-0014** PE pipeline — boundary 보존: scheduler 가 recipe-free,
엔진이 op_kind-opaque.
## Migration
> **구현 순서 노트.** 아래의 ADR-0064-우선 순서가 원래 계획이었음.
> 실제 구현은 **1↔2 단계를 뒤집어** — ADR-0065 Phase 1 (평평화
> `CompositeCmd`) 이 ADR-0064 Rev2 *보다 먼저* land. 이유: `logical_bytes`
> (ADR-0064 D2) 가 평평화 형태 위에서 자연스럽게 정의되므로, refactor 를
> 먼저 하면 버려지는 transitional `logical_bytes` 를 피함. 각 단계는
> 독립적으로 안전: P1 은 기존 cost 경로 하에서 의미 보존 refactor
> (op_log byte-equal); ADR-0064 Rev2 가 이미 평평화된 형태 위에 구조적
> 공식을 교체. 최종 상태는 어느 쪽이든 동일.
Land 순서 (원래 계획; as-built 1↔2 뒤집힘은 위 노트 참조):
1. **ADR-0064 Revision 2** (별도 PR): 구조적 dispatch cost +
`logical_bytes` property + topology config override + 일회성 골든 재생성.
2. **ADR-0065 Phase 1** (본 ADR, PR a): `CompositeCmd` 평평화 refactor +
legacy-API lowering. refactor only; 골든 불변.
3. **ADR-0065 Phase 2** (본 ADR, PR b): `tl_recipes.py` + `softmax_merge`
recipe + PE_SCHEDULER position-스캔 + DMA 자동 삽입 + strict-FIFO
RW tracker.
4. **ADR-0065 Phase 3** (본 ADR, PR c): `_gqa_decode_long.py` opt2
variant + 수치 동등 테스트 + opt3 대비 dispatch-ratio 측정.
상세 구현 계획: **DDD-0065** 참조.
docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm-ko.md
-747
View File
@@ -1,747 +0,0 @@
# ADR-0060: AHBM GQA Fused Attention 커널 (Llama3-70B)
## Status
Proposed
**Context model:** Llama3-70B.
**Decision drivers:** agentic workload → 낮은 batch, 긴 context;
KV-load-bound decode; long-context prefill용 sequence-parallel (Ring KV).
**Supersedes / extends:** 기존 mesh-native 어텐션 커널
`_attention_mesh_kv`(prefill)와 `_attention_mesh_mlo`(decode), 그리고
`milestone-gqa-llama70b` eval bench. *§A. 기존 kernbench 작업과의 관계* 참조 —
그 코드가 본 ADR이 진짜 GQA·causal·long-context 커널로 업그레이드하는
baseline이다.
**Supporting ADRs** (efficiency / scale enabler — *GQA blocker 아님*; §8 정정
참조): **ADR-0063** `tl.scratch_scope`(per-tile scratch 재활용 — 현실적 context
길이에 필요), **ADR-0062** lazy `tl.load`(첫 사용 시점 auto-wait를 갖는
non-blocking load → load/compute 오버랩), **ADR-0061** `tl.broadcast`(선택적
mask/범용 편의). 두 GEMM(Q·Kᵀ, P·V)은 scheduler가 관리하는 `tl.composite`
커맨드로 발행된다(기존 `CompositeCmd`; 새 커맨드 종류 없음); 진짜 GQA 자체는
커널 재구조화만 필요(§5.2).
**Algorithm lineage.** 이 커널은 **FlashAttention**(tiling + online/streaming
softmax, P·V fused — full score matrix 미생성)이다. §4의 KV-parallel
split-and-combine는 **FlashDecoding**(split-KV + log-sum-exp merge). §5.5의 Ring
경로는 **Ring Attention**(KV 블록을 mesh 주위로 회전, 같은 online softmax로
fold). 새 수학은 도입하지 않으며, 본 ADR은 이 알려진 알고리즘을 kernbench의
**greenlet `tl` 프로그래밍 모델**(ADR-0020, ADR-0046)과 **IPCQ** PE↔PE
collective(ADR-0023/0025)에 매핑한다.
---
## TL;DR — 최종 GQA 커널 (composite hybrid) pseudocode
결정(§1)을 한 곳에: **GEMM → `tl.composite`; softmax 머지 + tree reduction →
커널; `tl.load`은 lazy.** KV head별로, `G`를 matmul M 차원에 fold(진짜 GQA,
broadcast 없음). 이것이 reference 형태이며, 산문 섹션이 각 부분을 부연한다.
```python
def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr,
counter, start_pe, N, q_block, scale, *, tl):
# ---- geometry (kernel arithmetic, §2) ----
pe_id = tl.program_id(axis=rank_axis)
G = H_q // H_kv # query heads per KV head (=8)
causal = q_block.is_prefill
for kv in range(H_kv): # one KV head per iteration (§5.2)
# Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062):
# issue now, auto-wait at first use inside the first composite.
q_g = tl.load(q_base(kv), (G * q_block.T_q, d)) # [G·T_q, d]
my_len = valid_len(counter, start_pe, pe_id, N) if not causal else S_kv
n_tiles = ceil(my_len / TILE)
# persistent arena (outside scratch_scope, ADR-0063): -inf, 0, zeros
m, l, O = init_running(G * q_block.T_q, d)
for j in range(n_tiles):
if causal and tile_all_future(j, q_block):
continue # causal skip (kernel if)
with tl.scratch_scope(): # per-tile temporaries (ADR-0063)
# --- GEMM #1 on the scheduler: Q·Kⱼᵀ; K streamed by composite ---
Sj = tl.composite("gemm", a=q_g,
b=tl.ref(k_tile(kv, j), (d, TILE))) * scale # Kᵀ pre-stored [d,TILE]
if causal and tile_partial(j, q_block):
Sj = Sj + causal_mask(j, q_block) # additive boundary mask
# --- online softmax merge in the kernel (MATH ops) ---
m_new = tl.maximum(m, tl.max(Sj, axis=-1))
P = tl.exp(Sj - m_new)
corr = tl.exp(m - m_new)
l = l * corr + tl.sum(P, axis=-1)
# --- GEMM #2 on the scheduler: P·Vⱼ; V streamed by composite ---
Oj = tl.composite("gemm", a=P,
b=tl.ref(v_tile(kv, j), (TILE, d)))
O = O * corr + Oj # running merge
m = m_new
# ---- cross-PE combine (§4): log-sum-exp tree to root, or store ----
if N == 1:
tl.store(o_base(kv), O / l)
else:
tree_reduce_and_store(m, l, O, pe_id, N, o_base(kv)) # tl.send/tl.recv
```
> **이 형태인 이유:** 두 `tl.composite` 호출이 tiling + K/V DMA 스트리밍 +
> cross-tile 파이프라이닝을 PE_SCHEDULER에 offload한다(CPU는 coarse descriptor를
> 내고 앞서 나가 → 엔진이 saturate 유지, §1); softmax 머지와 IPCQ tree reduction은
> 커널에 남는다(기존 `CompositeCmd`가 cross-tile 상태를 못 들기 때문). `K`는
> reshape-not-transpose caveat를 피하려 transpose된 `[d, TILE]`로 미리 저장된다
> (§3, §B). 전용 "flash-composite" 커맨드 종류는 도입하지 **않는다**(§8 항목 4).
---
## A. 기존 kernbench 작업과의 관계 (먼저 읽을 것)
kernbench는 **오늘날 이미 IPCQ 상에서 online-softmax `(m, ℓ, O)` 머지로
FlashAttention을 돌린다.** 두 커널이 존재한다:
| File | Role | Mechanism |
|---|---|---|
| `src/kernbench/benches/_attention_mesh_kv.py` | prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold |
| `src/kernbench/benches/_attention_mesh_mlo.py` | decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge |
둘 다 `milestone-gqa-llama70b`(4 패널:
`{single,multi}_user × {prefill,decode}`,
`src/kernbench/benches/milestone_gqa_llama70b.py`)이 구동하며
`tests/attention/test_milestone_gqa_llama70b.py`에서 테스트된다.
**이들은 greenlet `tl` API로 작성되었다:** `tl.load`, `tl.dot`,
`tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, 그리고
`TensorHandle`에 대한 Python `-`/`*`/`/`(각각 `MathCmd` emit) — GEMM은 composite가
아니라 **blocking `tl.dot`**로. 실행 중 `(m, ℓ, O)`는 루프를 관통하는 Python
`TensorHandle`일 뿐이다. 본 ADR은 실행 상태와 softmax 머지는 커널에 유지하되
**두 GEMM은 scheduler가 관리하는 `tl.composite` 경로로 옮긴다**(§1 참조) — 이것이
**설계에 중요**하다.
**baseline의 의도된 세 한계** — 효율적 GQA 커널이 정확히 들어내야 하는 것:
1. **GQA 재사용 없음.** `h_q == h_kv == 1`
(`test_milestone_gqa_llama70b.py:137-142`). 테스트는 이를 *broadcast view*에
대한 MemoryStore byte-conservation 실패로 돌리지만, 그 실패는 baseline의
**head-packing 핵**의 속성이다(`_view(K, (h_q·d, S_kv))`는 모든 head를 하나의
matmul 차원에 뭉치고 `h_q == h_kv`일 때만 byte를 보존). 올바른 수정은
broadcast op이 아니라 **커널 재구조화**다: **한 번에 한 KV head**를 처리하고
`G` group 행을 matmul **M** 차원에 fold(§5.2). 이는 byte-보존 reshape만 쓰므로
진짜 GQA(`h_q = G·h_kv`)가 **신규 primitive 없이** 돈다 — §8 참조.
2. **O(N) reduction.** baseline은 all-to-all bidirectional fan-out을 하여 *모든*
rank가 full 답을 갖는다(`n_ranks − 1` 단계). 어텐션은 query owner에서만 `O`가
필요 → root로의 **tree reduction**은 `⌈log₂ N⌉` 단계(§4).
3. **검증 스케일만.** `S = 16`인 이유는 1 MiB scratch bump allocator가 per-tile
임시값을 누수하고(`test_milestone_gqa_llama70b.py:123-148`) causal masking /
tiling이 없기 때문 → **ADR-0063**(재활용) + §5(tiling, causal skip) + composite
K/V 스트리밍(§3) + **ADR-0062**(lazy load 오버랩)으로 해결.
**문서 부채(범위 밖이나 기록):** baseline은 ADR-0055/0056/0057/0058/0059를
인용하나 **파일로 존재하지 않는다** — ghost 참조다. 본 ADR은 이를 소급
작성하지 않으며; 권고는 Detailed Design Document의 *Open Decisions* 참조.
---
## 0. 참조 차원 (Llama3-70B)
| Symbol | Meaning | Value |
|---|---|---|
| `H_q` | query heads | 64 |
| `H_kv` | KV heads | 8 |
| `G` | GQA group size = `H_q / H_kv` | 8 |
| `d` | head dim | 128 |
| `L` | layers | 80 |
| `D` | model dim | 8192 |
하드웨어 요약:
- **AHBM (chip)** = **CUBE**(메모리 cube, 각각 **PE**를 담은 logic die) 집합 +
**IO die**(ADR-0003).
- **IPCQ**: PE↔PE 큐, PE당 4 mesh-방향 queue-pair
(`N/S/E/W`, ADR-0023 D3; inter-SIP용 `global_*`, ADR-0032). 커널
API: `tl.send(dir, src)` / `tl.recv(dir, shape, dtype)`
(`tl_context.py:402-499`).
- **Composite command** (`CompositeCmd`, `pe_commands.py:144-162`): 단일
GEMM(또는 MATH) *head* + element-wise *epilogue* 단계
(`bias/relu/scale/add/...`), PE_SCHEDULER에 **non-blocking**으로 발행되며,
scheduler가 tile plan을 생성하고 타일당 DMA→GEMM→write를 스트리밍한다
(ADR-0014 D6; `pe_scheduler.py:104-143`). 이는 일반적 multi-op DAG가 **아니다**:
두 GEMM을 chain할 수 없고, 인스턴스 간 register 상태를 못 들며, IPCQ를
pop/wait 못 한다. 따라서 본 ADR은 두 어텐션 GEMM(Q·Kᵀ, P·V)을 **각각** 자체
composite로 발행하고 cross-GEMM softmax 머지 + IPCQ reduction은 커널에 유지한다
— 새 "flash-composite" 커맨드 종류는 **필요 없다**(§1, §8 참조).
- **Allocation policy (SP):** multi-user 패널에서 CUBE당 query head 하나; KV
group의 `G` query head는 한 CUBE의 PE 안에 매핑.
**직교하는 두 매핑 레이어** (round-robin placement):
- **Layer 1 (KV-parallel, intra-request):** KV token `i`는 PE
`(start_pe + i) mod N`에 안착 — round-robin이라 한 긴 request가 `N` PE에 걸쳐
분할, ≤1 token으로 균형.
- **Layer 2 (inter-request):** `start_pe = request_id mod N`이 request별 "PE-0
역할"을 회전.
- 결합: `pe = (request_id + global_token_idx) mod N`.
`N` = 하나의 (query head, request)에 대한 KV-parallel reduction group의 PE 수.
Reduction은 **intra-CUBE** 유지(query head는 CUBE를 가로지르지 않음 — 명시적
non-goal).
---
## 0.5 커널 경계, 전제조건, I/O 계약
### 0.5.1 디코더 레이어 내 위치
```
1. RMSNorm
2. QKV projection (GEMM) ─┐ qkv_rope kernel (SEPARATE, upstream)
3. RoPE on Q and K ─┤
4. write new K,V → KV cache ─┘
5. ===== THIS KERNEL: FlashAttention ===== (post-RoPE Q, K-cache, V-cache → O)
6. Output projection (GEMM) out_proj kernel (SEPARATE, downstream)
7. residual add → FFN ...
```
**전제조건 (upstream `qkv_rope`, 본 커널 아님):**
- **P1.** Q는 이미 RoPE-회전됨. 여기서 회전 없음.
- **P2.** K-cache는 **post-RoPE** K 저장(어텐션 시점에 재회전 안 함 — post-RoPE
캐싱의 이유).
- **P3.** decode의 경우, 새 step의 K 행은 RoPE-회전되어 그 소유 PE의 K-cache
슬롯에 **본 커널 launch 전에 `qkv_rope`가** 추가한다. ⇒ 본 커널은 KV cache에
대해 **pure read**.
- **P4.** V는 회전 안 함; V-cache는 raw projected V 보유.
- **P5.** upstream RoPE 위치는 token의 **절대 global 위치**
`global_idx = local_slot·N + ((pe_id − start_pe) mod N)`(§2.1). Round-robin
placement ≠ RoPE 위치.
### 0.5.2 Shape 기호
`T_q` = 이번 launch의 query 길이(decode: 1; prefill/chunk: chunk 폭).
`S` = 전체 context 길이. `S_pe` = 이 PE 소유 key ≈ `⌈S/N⌉`.
저장 `bf16`(numpy proxy `f16`, `memory_store.py:16`); `m,ℓ,O` 누산기는 정밀도가
중요한 곳에서 `f32`.
### 0.5.3 INPUTS (커널 launch당)
| Input | Shape (per KV head) | Location | Notes |
|---|---|---|---|
| `Q` | `[G, T_q, d]` | per-PE HBM (loaded to TCM) | post-RoPE (P1). group의 `G` query 행 batched. |
| `K_cache` | `[S_pe, d]` | per-PE HBM, base `K_base[pe]`, contiguous | post-RoPE (P2). Read-only. 로컬은 dense, global index는 strided (§2.1). |
| `V_cache` | `[S_pe, d]` | per-PE HBM, base `V_base[pe]` | raw V (P4). Read-only. |
| `global_token_counter` | scalar | launch arg | 커널이 `S_pe`, slot↔global, causal bound 도출. |
| `start_pe` (= `request_id mod N`) | scalar | launch arg | Layer-2 회전. |
| `pe_id`, `N` | scalars | launch / `tl.program_id` | reduction-group geometry. |
| `q_block_meta` `{q_start, T_q}` | launch arg | prefill/SP causal masking & skip. |
| `O_base` | address | launch arg | 최종 O가 쓰이는 곳 (root PE만). |
| `softmax_scale` | scalar | launch arg | `1/√d`. |
**Ring Attention (§5.5)**에는 ring step마다 추가: IPCQ로 들어오는 `K_block,
V_block`을 ping-pong 버퍼에(post-RoPE), 그리고 causal step-skip용
`step_kv_global_range`.
### 0.5.4 OUTPUTS
| Output | Shape | Location | Notes |
|---|---|---|---|
| `O` (final) | `[G, T_q, d]` | `O_base`, **root PE만** | 정규화 `O_acc / ℓ_acc`, 저장 dtype으로 cast. |
**중간값 (non-root PE, IPCQ 상 — 커널-가시 출력 아님):**
`(m_i, ℓ_i, O_i)` (`m,ℓ`: `[G, T_q]`; `O_i`: `[G, T_q, d]` 비정규화)이
`tl.send`로 `tree_parent`에 push(§4). `O_i`가 무거운 부분.
**No-SP (`N=1`):** IPCQ partial 없음; 단일 PE의 실행 중 `(m,ℓ,O)`를 제자리에서
정규화하여 `O_base`에 기록.
**명시적 비-출력:** KV-cache write(upstream, P3), output projection(downstream),
score `S` 및 prob `P`(미생성).
---
## 1. 결정 (메커니즘)
**커널을 *하이브리드*로 구현한다: 두 GEMM(Q·Kᵀ, P·V)을 scheduler가 관리하는
`tl.composite(op="gemm")` 커맨드로 발행하고, online-softmax 머지와 cross-PE
reduction은 커널 수준 `tl` op으로 유지한다.** `tl.load`은 **lazy**다(non-blocking;
wait는 load된 데이터의 첫 사용 시점에 자동 삽입 — ADR-0062). 따라서 명시적 HBM
load가 뒤따르는 compute와 오버랩된다. kernbench의 실행 + latency 모델에 근거한
근거:
- **GEMM tiling이 PE_SCHEDULER에 offload된다.** `CompositeCmd`는 non-blocking
(`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`): 커널이 **하나의 coarse
descriptor**(M = `G·T_q`, PE당 전체 tile sweep)를 push하면 scheduler가 tile
plan을 생성하고 타일당 DMA→GEMM→write를 스트리밍한다(ADR-0014 D6). K/V는
scheduler가 HBM에서 스트리밍하는 `tl.ref` operand이므로, 타일당 **K/V prefetch는
scheduler의 일**이다 — 명시적 prefetch op 없음. CPU(greenlet)는 현재 composite가
도는 동안 **다음** composite를 낼 수 있게 풀려나, scheduler가 타일을 가로질러
GEMM 엔진을 saturate 유지한다.
- **이는 하드웨어를 반영**하며 CPU issue-rate를 execution-rate로부터 decouple한다.
대조적으로 blocking per-op `tl.dot` 경로는 매 GEMM마다 CPU를 멈추고, 끼어드는
softmax MATH op 동안 GEMM 엔진 **버블**을 남긴다; CPU가 fine-grained per-tile
발행을 따라잡을 수 있을 때만 현실적이다.
- **실행 중 `(m, ℓ, O)` flash 상태는 Python `TensorHandle`로** 루프를 관통한다
(baseline이 이미 그렇게 함); softmax 머지(max/exp/sum/rescale)는 두 GEMM
composite **사이**의 커널 수준 `tl` MATH다. 기존 `CompositeCmd`는 두 GEMM을
chain하거나 cross-tile register 상태를 못 들므로(§0), 머지는 필연적으로 커널에
산다 — 이것이 하이브리드 분할이지, 우회한 한계가 아니다.
- **Cross-PE 결합**은 P·V 후 `(m, ℓ, O)`에 대한 log-sum-exp **tree**로, 커널 수준
`tl.send`/`tl.recv`(§4)를 통해 — 불변.
그래서 per-tile 내부 파이프라인은:
```
q_g = tl.load(Q group) # lazy; auto-wait at first use
per tile j:
Sⱼ = tl.composite("gemm", a=q_g, b=tl.ref(Kⱼ)) → Sⱼ # scheduler streams Kⱼ DMA + GEMM
Sⱼ += maskⱼ # kernel MATH, boundary tile only
online-softmax: mⱼ, m_new, P, corr, ℓ # kernel MATH
Oⱼ = tl.composite("gemm", a=P, b=tl.ref(Vⱼ)) → Oⱼ # scheduler streams Vⱼ DMA + GEMM
O = O*corr + Oⱼ; m = m_new # kernel MATH (running merge)
```
각 타일의 MATH 임시값을 `tl.scratch_scope`(ADR-0063)로 감싸 scratch를 O(1)로
유지하고, 다음 타일의 composite를 현재 타일 결과를 wait하기 전에 발행
(non-blocking 핸들)하여 scheduler가 타일을 가로질러 파이프라인되게 한다.
제어 흐름(tile skip, mask 생성, reduction 스케줄링, 주소 산술)은 **커널**에 산다
(greenlet 본문의 평범한 Python `if`/산술). 이는 kernbench의 greenlet 모델이 이미
허용하는 바 그대로다(`kernel_runner.py`, ADR-0020 D3).
> **이것이 무엇을 대체하나.** 이전 iteration은 "composite는 latency 이점 없음"이라는
> 논거로 순수 greenlet primitive 경로(전부 `tl.dot`, composite 없음)를 제안했다.
> 그것은 **오직** 시뮬레이터가 현재 per-op CPU 발행 비용을 **0**으로 청구하기
> 때문에만(`dispatch_cycles=0`, `pe_cpu.py`) 성립한다 — descriptor offload가 숨기려
> 존재하는 바로 그 CPU issue-rate / DMA-program 비용을 모델이 지워버린 것이다.
> 하이브리드가 효율 커널의 충실한 표현이다. 이점의 **측정 가능한** 크기(CPU가 많은
> 타일에 대해 엔진을 saturate할 수 있는가?)는 op 종류별 발행 비용 모델링에
> 좌우되며, future work로 추적된다(cost model; §9). `dispatch_cycles=0`에서도
> non-blocking composite 경로는 blocking `tl.dot` 경로가 남기는 GEMM 엔진 버블을
> 채운다.
>
> **이것이 효율적 선택인 이유.** GEMM tiling + DMA 스트리밍 + cross-tile
> 파이프라이닝은 검증된 `CompositeCmd` scheduler 경로에 offload되고; softmax 머지와
> 검증된 IPCQ collective는 커널에 남는다. 진짜 새 기계장치는 작고 범용적인 두
> primitive뿐이다(ADR-0062 lazy `tl.load`, ADR-0063 `tl.scratch_scope`); reduction은
> `tl.send`/`tl.recv`를 재사용한다. 전용 "flash-composite" 커맨드(softmax 머지 +
> carried register 상태 + IPCQ-push epilogue를 내재화하는 한 종류)는 만들지
> **않는다** — 크고 특수목적이며, 하이브리드 대비 유일한 delta(완전 softmax
> offload)가 현재 모델링 충실도에서 정당화되지 않는다; §8 참조.
---
## 2. 메모리 레이아웃 & 드라이버 책임
### 2.1 KV cache 할당
- per-PE KV 버퍼는 per-PE 최대 context `⌈max_context / N⌉ × d × dtype`로 K와 V를
각각 sizing. kernbench에서는 sequence 차원에 대해 `pe`(또는 `cube`) 축이
`row_wise`인 `DPPolicy`로 배치된다(`policy/placement/dp.py`; baseline은 K/V에
`DPPolicy(pe="row_wise")`, Q에 `replicate`).
- PE 안에서 할당된 token은 **연속 append**된다(slot 0,1,2,…). global index는
strided(`i, i+N, …`)이나 per-PE 버퍼는 dense ⇒ DMA가 연속 유지.
- Slot → global: `global_idx = local_slot·N + ((pe_id − start_pe) mod N)`.
### 2.2 드라이버 launch당 의무 (최소)
드라이버는 launch마다 base + counter + rotation을 공급하고; 커널이 나머지를
도출한다:
| Launch arg | Purpose |
|---|---|
| `K_base[pe]`, `V_base[pe]` | per-PE KV 버퍼 base (tensor VA에서) |
| `O_base` | 어텐션 출력 목적지 |
| `global_token_counter` | 현재 sequence 위치 |
| `start_pe = request_id mod N` | Layer-2 회전 |
| `pe_id` (`tl.program_id`), `N` | geometry |
| `q_block_meta` | prefill/SP query block 시작 + 길이 |
커널 도출(평범한 산술):
- **이번 step에 내가 쓸 차례:** `(start_pe + counter) mod N == pe_id`
- **내 valid 길이:** `base = counter // N; rem = counter % N;
my_len = base + (1 if ((pe_id − start_pe) mod N) < rem else 0)`
- **read 범위:** tile `0 .. ⌈my_len / TILE⌉`.
- **causal bound / per-tile skip:** query block과 각 tile의 global 위치에서.
드라이버 = base + counter + rotation. 단일 공식
`(request_id + token_idx) mod N`이 placement policy 전부다.
---
## 3. Per-tile op 시퀀스 (greenlet `tl`)
한 iteration = 한 PE의 한 KV tile. 두 GEMM은 `tl.composite(op="gemm")`
(scheduler 관리 tiling + K/V DMA 스트리밍); softmax 머지는 그 사이의 커널 `tl`
MATH다. 실제 `tl` 이름(`tl_context.py`), lazy `tl.load`(ADR-0062),
`tl.scratch_scope`(ADR-0063) 사용:
```python
# running state (persistent arena — allocated once, outside the scope)
# m: [G, T_q] l: [G, T_q] O: [G, T_q, d]
q_g = tl.load(Q_group_ptr, (G*T_q, d)) # lazy; auto-wait at first use (ADR-0062)
with tl.scratch_scope(): # per-tile MATH temporaries recycled
Sj = tl.composite("gemm", a=q_g, # [G·T_q, TILE]; scheduler streams Kⱼ DMA
b=tl.ref(K_base + j*TILE*d, (TILE, d))) * softmax_scale
if mask_j is not None:
Sj = Sj + mask_j # additive causal mask (boundary tile)
m_j = tl.max(Sj, axis=-1)
m_new = tl.maximum(m, m_j)
P = tl.exp(Sj - m_new) # no full-matrix softmax; streaming
corr = tl.exp(m - m_new) # rescale factor for old accumulators
l = l * corr + tl.sum(P, axis=-1)
Oj = tl.composite("gemm", a=P, # [G·T_q, d]; scheduler streams Vⱼ DMA
b=tl.ref(V_base + j*TILE*d, (TILE, d)))
O = O * corr + Oj # running merge (kernel MATH)
m = m_new
```
Notes:
- `q_g`는 `[G·T_q, d]`로 reshape된 GQA-batched query(`G` group 행을 matmul M
차원에 fold; byte-보존). 한 K/V tile이 모든 `G·T_q` 행을 서비스 — GQA 재사용
레버 — broadcast 없이.
- **K/V는 `tl.ref` operand**로 composite scheduler가 HBM에서 타일당 스트리밍한다
(`pe_scheduler.py:104-143`): 그것이 *바로* prefetch/파이프라인이므로 명시적
prefetch op이 없다. 타일 `j+1`의 composite를 타일 `j`를 wait하기 전에 발행
(non-blocking 핸들)하면 scheduler가 타일을 가로질러 파이프라인되고 GEMM 엔진이
saturate 유지된다.
- `tl.trans`는 kernbench에서 **메타데이터 전용**이고(`tl_context.py:390`)
`MemoryStore.read`는 transpose가 아니라 *reshape*한다(`memory_store.py:73`).
zero/structural 실행엔 무해하나; 비자명 수치 데이터에선 reshape-not-transpose가
되어 transpose된 K로의 Q·Kᵀ는 주의가 필요하다(§11) — K를 transpose해 `[d, TILE]`로
미리 저장하거나, 실제 `tl.transpose`를 추가(후보 primitive, 시뮬레이터의
성능-모델링 목적상 아마 불필요).
- Masking: 커널이 query/KV global offset에서 boundary-tile mask를 만들어 더한다
(`Sj + mask_j`); full-past tile은 `None` 전달; full-future tile은 **skip**(`if`이
enqueue 안 함).
- **새** "flash-composite" 커맨드 종류(softmax 머지 + carried `(m,ℓ,O)`를
내재화하는 것)는 쓰지 **않는다**; 기존 `CompositeCmd`가 각 GEMM을 담당하고 머지는
커널에 남는다(§1, §8 항목 4).
---
## 4. Reduction (KV-parallel / SP combine) — root로의 tree
tile sweep 후 각 PE는 `(m_i, ℓ_i, O_i)`(비정규화)를 가진다. 결합/교환 가능한
log-sum-exp 머지로 결합(baseline의 fold와 동일 수학, `_attention_mesh_mlo.py:117-122`):
```python
def merge(m_a, l_a, O_a, m_b, l_b, O_b):
m = tl.maximum(m_a, m_b)
sa, sb = tl.exp(m_a - m), tl.exp(m_b - m)
return m, l_a*sa + l_b*sb, O_a*sa + O_b*sb
# final: O = O_root / l_root # normalise once at the tree root
```
- **타이밍:** 교환은 **P·V 후**(각 PE가 최종 `O_i`를 가짐).
- **토폴로지:** mesh tree, depth `⌈log₂ N⌉`, 물리 이웃 상. `N=8`의 경우: level-0
`(0↔1,2↔3,4↔5,6↔7)`, level-1 `(1↔3,5↔7)`, level-2 `(3↔7)`, root = `7`. 각 tree
pair는 물리 `N/S/E/W` 이웃이어야 `tl.send(dir)`/`tl.recv(dir)`이 실제 link를
쓴다(튜닝 항목 §9; 이웃을 wiring하는 SFR install은
`configure_sfr_intracube_pe_ring`, baseline이 사용).
- **왜 baseline fan-out이 아니라 tree인가:** 어텐션은 한 곳(`O_base`를 쓰는 PE)에서
`O`가 필요하다. tree는 `⌈log₂ N⌉` 단계(`N=8`에 3) vs baseline all-to-all의 `N−1`
(7). per-PE sweep이 짧은 decode에서 reduction이 지배적 비용이므로 실제 이득이다.
(baseline fan-out은 답을 모든 rank에 복제 — 여기선 불필요.)
- **Payload:** `O_i`(`d`-vector)가 무겁고; `m,ℓ`은 `(g, T_q)`당 scalar. 별도
`tl.send`로 전송(baseline은 triplet을 세 send로 — 같은 패턴).
커널 구조(greenlet):
```python
# leaf / internal node: fold children that send to me, then forward up
for child_dir in tree_children_dirs(pe_id, N):
m_c = tl.recv(child_dir, m.shape); l_c = tl.recv(child_dir, l.shape)
O_c = tl.recv(child_dir, O.shape)
m, l, O = merge(m, l, O, m_c, l_c, O_c)
if not is_root(pe_id):
tl.send(parent_dir(pe_id), m); tl.send(parent_dir(pe_id), l); tl.send(parent_dir(pe_id), O)
else:
tl.store(O_base, O / l)
```
`tl.recv`는 데이터 도착까지 blocking(`tl_context.py:446-499`); 머지 순서가 정적
tree로 고정되므로 data-dependent `pop`이 불필요 — **하드웨어/composite
`pop`-as-dependency 변경이 필요 없는** 이유(§6).
---
## 5. Case별 커널
네 case 모두 §3(tile sweep) + §4(reduction)를 공유하며; `N`, query-block 폭,
어느 `tile_*` predicate가 발동하는지에서만 다르다.
### 5.1 공통 skeleton
```python
def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N,
q_block, softmax_scale, *, tl):
pe_id = tl.program_id(axis=rank_axis)
my_len = valid_len(counter, start_pe, pe_id, N)
n_tiles = ceil(my_len / TILE)
q_g = load_Q_group(q_ptr) # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast)
m, l, O = init_running() # persistent arena: -inf, 0, zeros
for j in range(n_tiles): # K/V streamed per tile by the composite scheduler (§3)
if tile_all_future(j, q_block): # causal skip (kernel if)
continue
run_tile(j, mask_or_null(j, q_block)) # §3: 2 composites + softmax MATH, in tl.scratch_scope
if N == 1:
tl.store(o_ptr, O / l) # no reduction
else:
tree_reduce_and_store(pe_id, N, o_ptr) # §4
```
### 5.2 DECODE, no SP (`N=1`, 한 PE가 head의 KV 소유)
- `T_q = 1`, 모든 과거 KV에 attend ⇒ future tile 없음, 마지막 ragged tile만 mask.
- **GQA 재사용이 전부**(decode는 KV-load-bound): KV head의 `G=8` query 행을 matmul
**M** 차원에 fold(`q_g`를 `[G, T_q, d] → [G·T_q, d]`로 reshape, byte-보존).
그러면 `Q·Kᵀ`는 `composite([G·T_q, d], Kᵀ[d, TILE]) → [G·T_q, TILE]`이고
`P·V`는 `composite([G·T_q, TILE], V[TILE, d]) → [G·T_q, d]`. KV tile(`[TILE, d]`)이
공유 `K`/`V` operand — **한 번 스트리밍되어 모든 `G·T_q` 행이 자동 재사용** —
그것들이 GEMM의 M 행이기 때문. K/V broadcast 불필요; composite의 tile plan 내
`m = G·T_q`가 timing이 모든 `G` 행의 작업을 올바르게 세게 한다(leading batch
축은 세어지지 *않음* — §8 참조).
- `S_pe`가 scratch에 맞으면(작은/중간 context) 이것은 **one-shot** partial
attention으로 축약(Q·Kᵀ용 composite 하나, softmax 하나, P·V용 composite 하나) —
바로 baseline `_attention_mesh_mlo`의 `_partial_attention`, 단지 GQA-batched에
composite 경로. Tiling(§3)은 `S_pe`가 scratch scope의 tile 예산을 초과할 때만
발동.
### 5.3 DECODE, with SP / KV-parallel (`N=8`)
- 한 request의 KV가 8 PE에 round-robin; 각각 ≈`my_len` token 소유. 각 PE가 자기
로컬 tile에 대해 `G=8` query 행을 GQA-batch.
- `T_q=1` ⇒ 짧은 per-PE sweep ⇒ **reduction 지배** ⇒ §4 tree(3단계)가 구조적으로
중요한 부분. Reduction latency는 batch될 때 다른 동시 decode token으로, 혹은 긴
single-stream context의 긴 per-PE tile sweep으로 숨겨진다.
### 5.4 PREFILL, no SP
- 전체 prompt 상주; query는 `T_q` token 블록(chunk).
- causality가 실재, query block `[qs, qe)` vs KV tile `[ks, ke)`:
- `ke ≤ qs` → `tile_all_past` → `mask=None`, full compute.
- `ks ≥ qe` → `tile_all_future` → **skip**(커널 `if`).
- overlap → `tile_partial` → 커널이 삼각 additive mask 생성, tile op 시퀀스가
더함(`Sj + mask_j`).
- GQA는 group의 query head를 decode와 같은 축으로 batch.
### 5.5 PREFILL, with SP (Ring KV)
- KV가 `N` PE(ring)에 sharded; 각 step이 peer의 KV 블록을 IPCQ로 전달. 실행 중
`(m,ℓ,O)`는 **ring step을 가로질러** carry(Python 핸들, `_attention_mesh_kv`가
오늘 하는 그대로).
- IPCQ가 다음 step의 KV 수신을 현재 step의 compute와 `tl.recv_async`/`tl.wait`로
오버랩(`tl_context.py:543-560` — 이미 존재).
- **Causal ring skip:** 들어오는 step의 KV 블록이 전부 로컬 query block *이후*면 그
compute를 skip(recv-consume 안 함 / fold 안 함) — causal 어텐션에서 약 절반 step
제거.
- 수신 버퍼는 ping-pong(persistent arena, 재활용 안 함).
- ring 후, `(m,ℓ,O)`는 query-소유 PE별 최종; KV-parallel 분할이 남으면 §4가
tree-merge한 뒤 정규화 + store.
baseline `_attention_mesh_kv`가 이미 ring fold를 구현한다; 본 ADR은 GQA 재사용,
causal step-skip, `recv_async` 오버랩을 추가한다.
---
## 6. 왜 하드웨어 / composite 변경이 불필요한가
- HW/composite `pop`-as-dependency가 도움 될 유일한 곳은 poll을 숨길 다른 작업이
없는 외로운 reduction — 즉 batch=1 **그리고** 짧은 context. 타깃은 **agentic =
낮은 batch, 긴 context** ⇒ 각 PE가 많은 KV tile을 가짐; §4 tree의 `tl.recv`
blocking은 동시 token / 긴 context의 sweep 작업으로 가려진다.
- §4 reduction은 **정적** tree를 쓰므로 recv 순서가 컴파일 타임에 고정 —
data-dependent pop 없음. `tl.recv`(blocking)로 충분.
- **결정: greenlet `tl.send`/`tl.recv` collective로 출시.** 짧은-context,
single-stream, latency-critical 타깃이 나타날 때만 HW `pop` 재검토.
---
## 7. 제어 vs 실행 분할 (load-bearing 원칙)
| Concern | Owner |
|---|---|
| tile skip (future), mask 생성, causal bound | **커널** (greenlet 본문의 Python `if` + 산술) |
| 주소 / offset / valid-length 산술 | **커널** (counter에서) |
| reduction 스케줄링, IPCQ send/recv 순서 | **커널** (정적 tree) |
| Q·Kᵀ, P·V (per-tile K/V DMA 스트리밍 + tiling 포함) | **`tl.composite`** → PE_SCHEDULER |
| Q load, mask add, softmax math, 실행 중 `(m,ℓ,O)` 머지 | **`tl` op** PE 엔진 상 (커널 발행) |
커널이 결정하고; GEMM은 composite로 scheduler에 offload되며; 나머지 `tl` op은
이미 결정된 작업을 실행한다. 이는 kernbench가 이미 지원하는 greenlet + composite
모델 그대로 — 새 제어 추상화 없음.
---
## 8. 필요한 kernbench 변경
**설계 iteration의 정정:** 진짜 GQA(`h_q > h_kv`)는 **신규 primitive 불필요** —
§5.2의 커널 재구조화(KV head별, `G`를 M에 fold, byte-보존 reshape)만 필요. 보조
ADR들은 *efficiency / scale* enabler이지 GQA blocker가 아니다.
**커널 내 알고리즘 작업 (신규 primitive 없음; 기존 `tl` API):**
- **GQA Q축 batching**(재사용 레버) — KV head별로 `G·T_q`를 matmul M 차원에 fold
(§5.2); `_view` 식 byte-보존 reshape; GEMM은 M = `G·T_q`인
`tl.composite(op="gemm")`. 오늘 timing/data 모드 모두 동작.
- **composite를 통한 GEMM**(§1/§3) — Q·Kᵀ와 P·V를 각각 non-blocking
`tl.composite(op="gemm")`로 발행; PE_SCHEDULER가 tiling하고 `tl.ref` K/V
operand의 DMA를 스트리밍(기존 `CompositeCmd`; 새 커맨드 종류 없음).
- root로의 tree reduction(§4)이 baseline all-to-all fan-out 대체 — `tl.send`/
`tl.recv` 상의 순수 커널 제어 흐름.
- causal tile skip + additive boundary mask(§3/§5.4) — 커널 `if` + `+`로 더한
mask 텐서.
- round-robin KV placement / valid-length 산술(§2) — launch-arg 산술 +
`DPPolicy(pe="row_wise")`.
**efficiency / scale을 위한 신규 primitive (각각 보조 ADR 있음):**
1. **per-tile scratch 재활용** — **ADR-0063**(`tl.scratch_scope`). *스케일에 필요*:
`S=16` 상한(1 MiB bump allocator)을 제거해 현실적 context 길이가 돌게 함. 셋 중
최고 가치.
2. **Lazy `tl.load`** — **ADR-0062**(non-blocking load + 첫 사용 시점 auto-wait; API
표면 불변). *효율*: 명시적 load(Q group, 비-composite 커널)를 뒤따르는 compute와
오버랩. 타일당 **K/V** prefetch는 composite scheduler가 처리(§1)하므로, 이것은
나머지 명시적 load를 담당. 전역 시맨틱 변경 → 기존 골든 재생성(ADR-0062 D3).
3. **GQA head / mask broadcast** — **ADR-0061**(`tl.broadcast`). *선택적 편의*, GQA
blocker 아님(위 정정 참조). `G·T_q` 행에 걸친 additive-mask 구성과 범용 커널에
유용; `np.matmul`이 data 모드에서 이미 broadcast하므로 정확성엔 불필요. 최저
우선순위.
**명시적 REJECTED (효율적 대안 선택):**
4. ~~내부 루프 전체를 내재화하는 전용 "flash-composite" 커맨드 종류 —
DMA→MM→VEC→DMA→MM→VEC + carried `(m,ℓ,O)` register 상태 + 꼬리 IPCQ push.~~ 두
GEMM은 기존 `CompositeCmd`를 **쓴다**(§1/§3) — 그것이 scheduler 관리 tiling, K/V
DMA 스트리밍, cross-tile 파이프라이닝을 준다. 기각되는 것은 softmax 머지 +
cross-tile register 수명 + IPCQ-push epilogue까지 흡수하는 **새** 커맨드 종류다:
크고 특수목적이며, 하이브리드 대비 유일한 delta(완전 softmax offload)가 현재
모델링 충실도(per-op CPU 발행 비용 = 0; §1 참조)에서 정당화되지 않는다. cost
model(§9)이 완전 offload를 측정 가능하게 가치 있게 만들면 재검토.
5. ~~하드웨어 `pop`-as-dependency.~~ 범위 밖(§6).
6. ~~본 커널 내 RoPE / QKV projection / KV-cache write.~~ Upstream `qkv_rope`
(P1–P5). RoPE를 fold하면 매 decode step마다 과거 tile을 재회전해야 하고 Ring
Attention의 post-RoPE pass-through를 깬다.
---
## 9. Open 튜닝 항목 (kernbench에서 측정, blocking 아님)
1. **Composite tile-pipeline depth** — KV-load-bound에 지배적; 커널이 wait 전에
non-blocking composite를 얼마나 앞서 발행하는지, 그리고 scheduler의 타일당
스트리밍 depth.
2. **reduction tree의 PE↔mesh-이웃 매핑** — 각 depth-`⌈log₂ N⌉` pair가 물리
`N/S/E/W` 이웃인지 보장; 나쁜 매핑은 hop 추가. SFR install 대조 검증.
3. **TILE 크기** — scratch 상주(`S/P tile + O_acc + G-way GQA`)와 DMA 효율 균형;
ADR-0063 및 scheduler의 `TILE_M/K/N`(`pe_scheduler.py`)과 상호작용.
4. §5.5의 **Ring 버퍼 ping-pong vs `recv_async` depth**.
5. **decode의 one-shot vs tiled 교차점**(§5.2) — tiling이 단일 composite를 이기는
`S_pe` 임계값.
6. **per-op CPU 발행 비용 (cost model)** — 현재 `dispatch_cycles=0`(`pe_cpu.py`)이라
composite-vs-primitive 발행 오버헤드가 보이지 않음. op 종류별 차등 발행 비용
(`tl.composite` descriptor push ≫ primitive op)이 하이브리드의 CPU-saturation
이점을 **측정 가능**하게 한다(§1). **ADR-0064**에 명세; 별도 future work로 추적.
---
## 10. Coverage 요약
| Case | KV placement | Inner loop | Reduction | Masking |
|---|---|---|---|---|
| Decode, no SP | 1 PE, all KV | one-shot or tiled, GQA-batched | none | last tile only |
| Decode, SP | round-robin `N` PEs | tiled, GQA-batched | §4 tree (`⌈log₂ N⌉`) | last tile only |
| Prefill, no SP | resident | tiled per q-block | none | skip future / triangular boundary |
| Prefill, SP (Ring) | ring-sharded | per ring step (recv_async overlap) | running state + §4 tree | causal step-skip + boundary |
**Case별 I/O** (전체 계약 §0.5):
| Case | Inputs | Output | IPCQ partials |
|---|---|---|---|
| Decode, no SP | `Q[G,1,d]`, full `K/V[S,d]` on 1 PE | `O[G,1,d]` at `O_base` | none |
| Decode, SP | `Q[G,1,d]`, per-PE `K/V[S_pe,d]` | `O[G,1,d]` (root) | `(m,ℓ,O_i)` per PE → tree |
| Prefill, no SP | `Q[G,T_q,d]`, `K/V[≤end,d]` | `O[G,T_q,d]` at `O_base` | none |
| Prefill, SP (Ring) | `Q[G,T_q,d]`, ring-delivered `K/V` blocks | `O[G,T_q,d]` (root) | running `(m,ℓ,O)` + tree |
모든 case에서: KV-cache write와 RoPE는 **upstream**에서; output projection은
**downstream**에서; score `S`와 prob `P`는 미생성.
---
## 11. 검증 계획 (Phase 1 테스트 개요)
SPEC/ADR coverage: R5 (PE↔PE IPCQ, PE↔HBM), R2 (traversal에 의한 latency),
ADR-0023/0025 (IPCQ), ADR-0046 (`tl` 계약), ADR-0054 (eval bench).
시뮬레이터의 계약은 bit-exact 수치가 아니라 **traversal에 의한 latency + 결정성 +
구조적 정확성**(SPEC §0, §0.1)이다 — Phase 2 데이터는 주로 data path를
exercise하려 존재하고, `tl.trans`는 reshape-not-transpose이며, `bf16`은 `f16`으로
모델된다. 따라서 검증은 **구조/타이밍 우선**, 수치 parity는 제한된 2차 체크.
1. **data 모드 실행(`enable_data=True`):** GQA 커널(`h_q = G·h_kv`)이 baseline
head-packing이 부딪히는 byte-conservation 에러 없이 — 네 case 모두 — 완료. (이는
§5.2 재구조화가 필요하지, 신규 primitive가 *아님*.)
**수치 parity(2차):** reshape-as-transpose가 정확한 symmetric/identity 입력에
대해, 커널 `O`가 numpy FlashAttention reference와 fp tolerance 내 일치. 완전
asymmetric parity는 실제 `tl.transpose`에 좌우됨(범위 밖; 플래그됨).
2. **GQA 재사용:** `h_q = G·h_kv`에서 K/V `dma_read_count`가 `G`와 무관(타일당 한
load, group에 걸쳐 재사용), 반면 GEMM 작업은 `G`로 스케일. 레버가 실제 발동함을
assert.
3. **tree reduction 단계 수:** decode-SP가 `⌈log₂ N⌉` reduction 라운드(`N−1`이
아님)를 발행; op_log `ipcq_send`/`recv` 수가 tree와 일치.
4. **Causal skip:** prefill이 모든 `tile_all_future` tile을 skip — GEMM 수가 full
grid가 아니라 하삼각 tile 수와 일치.
5. **Long context (ADR-0063):** scope 없이 1 MiB를 넘기는 `S` sweep이 완료하고
reference와 일치.
6. **Load/compute 오버랩:** tiled sweep의 end-to-end latency가 직렬
`Σ(load+compute)`보다 낮음 — composite scheduler의 타일당 K/V 스트리밍(§1/§3)과
Q load를 오버랩하는 lazy `tl.load`(ADR-0062)에서. (오버랩은 실제 모델 동시성,
빼기가 아님.)
7. **Composite GEMM offload (구조적):** 각 tile의 Q·Kᵀ와 P·V가 blocking `tl.dot`이
아니라 `CompositeCmd`(non-blocking)를 PE_SCHEDULER에 emit; op_log가 composite
tile plan을 보이고 커널이 다음 tile의 composite를 wait 전에 발행(cross-tile
파이프라이닝).
8. **결정성:** 동일 입력 → 동일 op_log + latency (SPEC §0.1).
---
## B. 하이브리드 전환에서 나온 Open 설계 항목 (추후 검토)
이들은 결정이 순수 greenlet primitive 경로에서 **composite hybrid + lazy
`tl.load`**(이번 개정)로 옮겨갈 때 생겼다. 설계를 막는 것은 없으며; 각각 구현 중
검증 패스가 필요하다. (묻지 않고) 작업 합의에 따라 여기 기록한다 — 권고는 내가
예측한 기본값이며; 검토 시 수정.
1. **DDD-0060이 아직 미동기화.** Detailed Design Document는 여전히 옛
`tl.load_async` 더블버퍼 경로와 primitive `tl.dot` 내부 루프를 기술한다(그
§4.3/§5/§10). 하이브리드(composite GEMM, lazy load, K 사전 transpose)로
갱신해야 한다. DDD가 파생 how-to이고 큰 재작성이라 *검토용으로 남김*; ADR이 이제
권위 있는 기록이다. **권고:** 구현 시작 전 후속으로 DDD 동기화.
2. **composite GEMM의 K operand 방향.** Q·Kᵀ는 `b = [d, TILE]`가 필요하나 KV
cache는 K를 `[S_pe, d]`로 저장한다. `tl.trans`는 메타데이터 전용이고
`MemoryStore.read`는 transpose가 아니라 reshape(`memory_store.py:73`) — 비자명
데이터엔 런타임 transpose가 틀리다. **권고:** cache에 K를 **사전 transpose**
`[d, S_pe]`로 저장(pseudocode와 §3이 이를 가정)하여 `tl.ref(k_tile, (d, TILE))`이
연속 slice가 되게. upstream `qkv_rope` write 레이아웃이 이를 지원하는지 검증,
아니면 실제 `tl.transpose` 추가(더 무거움; 연기).
3. **composite 출력 버퍼 vs `tl.scratch_scope`.** 각 Q·Kᵀ composite는 `Sj`를
`out_addr`에 쓰고; 커널이 softmax MATH용으로 그것을 읽는다. 그 출력 버퍼와
in-flight composite의 타깃은 per-tile `scratch_scope`(ADR-0063)가 소비 전에
재활용하지 **않는** 곳에 살아야 한다 — in-flight lazy load와 같은 규율(ADR-0062
D-Negative). **권고:** *현재* tile의 composite 출력은 scoped arena에(같은
iteration에 소비); persistent `(m,ℓ,O)`는 바깥에. 다음 tile의 composite를 일찍
발행할 때(cross-tile 파이프라이닝) use-after-recycle이 없음을 검증.
4. **composite 스트리밍 하의 GQA `dma_read_count` 레버.** 레버(§11.2: K/V
`dma_read_count`가 `G`와 무관)는 composite가 모든 `G·T_q` M-행에 재사용되는
**하나의** K/V tile DMA를 emit한다고 가정한다. scheduler의 `generate_gemm_plan`은
`TILE_M/K/N`(`pe_scheduler.py:35-37`, 32/64/32)으로 tiling한다 — `G·T_q`에 대한
M-tiling이 M-tile마다 공유 K/V tile DMA를 재발행하지 **않음**을 확인(즉 operand
DMA가 M-tile에 걸쳐 공유되거나, 아니면 레버가 약해짐). **권고:** levers 테스트에서
assert; 위반 시 GQA 이득은 DMA가 아니라 compute에만 — 여전히 정확하나 헤드라인이
바뀜.
5. **커널 TILE vs scheduler `TILE_M/K/N`.** 커널은 논리적 KV `TILE`을 다루고;
scheduler는 고정 `TILE_M/K/N`으로 내부 재tiling한다. 두 tiling 레이어가
상호작용(scratch 상주, 파이프라인 depth). **권고:** 커널 TILE을 K/V 스트리밍
granularity로 보고 scheduler가 GEMM을 sub-tile하게; DDD에 관계를 문서화하고 둘 다
sweep(§9 항목 1, 3).
6. **cost model은 별도 ADR.** 하이브리드의 CPU-saturation 이점은
`dispatch_cycles=0`인 동안 보이지 않는다. op 종류별 issue-cost 모델은
**ADR-0064**에 명세; 본 ADR의 §1/§9는 *측정 가능한*(구조적뿐 아닌) 이점을 위해
그것에 의존. **권고:** eval에서 하이브리드 latency 이점을 주장하기 전에 ADR-0064
모델을 도입.
7. **Ring 경로(§5.5) GEMM.** §5.5는 여전히 primitive op + `recv_async`로 ring
fold를 기술한다. 일관성을 위해 ring의 per-step Q·Kᵀ / P·V도 composite여야 하고;
IPCQ `recv_async` 오버랩은 직교하며 유지. **권고:** 구현 중 ring step에 같은
하이브리드 형태 적용; 낮은 리스크, §3 미러.
docs/adr-proposed/ADR-0060-algo-gqa-fused-attention-ahbm.md
-817
View File
@@ -1,817 +0,0 @@
# ADR-0060: AHBM GQA Fused Attention Kernel (Llama3-70B)
## Status
Proposed
**Context model:** Llama3-70B.
**Decision drivers:** agentic workload → low batch, long context;
KV-load-bound decode; sequence-parallel (Ring KV) for long-context prefill.
**Supersedes / extends:** the existing mesh-native attention kernels
`_attention_mesh_kv` (prefill) and `_attention_mesh_mlo` (decode) and the
`milestone-gqa-llama70b` eval bench. See *§A. Relationship to existing
kernbench work* — that code is the baseline this ADR upgrades to a real
GQA, causal, long-context kernel.
**Supporting ADRs** (efficiency / scale enablers — *not* GQA blockers;
see §8 correction): **ADR-0063** `tl.scratch_scope` (per-tile scratch
recycling — required for realistic context length), **ADR-0062** lazy
`tl.load` (non-blocking load with auto-wait on first use → load/compute
overlap), **ADR-0061** `tl.broadcast` (optional mask/general
convenience). The two GEMMs (Q·Kᵀ, P·V) are issued as scheduler-managed
`tl.composite` commands (the existing `CompositeCmd`; no new command
kind); real GQA itself needs only kernel restructuring (§5.2).
**Algorithm lineage.** This kernel is **FlashAttention** (tiling +
online/streaming softmax with fused P·V — no full score matrix
materialised). The KV-parallel split-and-combine in §4 is
**FlashDecoding** (split-KV with log-sum-exp merge). The Ring path in
§5.5 is **Ring Attention** (KV blocks rotated around the mesh, folded by
the same online softmax). No new math is introduced; this ADR maps those
known algorithms onto the kernbench **greenlet `tl` programming model**
(ADR-0020, ADR-0046) and the **IPCQ** PE↔PE collective (ADR-0023/0025).
---
## TL;DR — final GQA kernel (composite hybrid) pseudocode
The decision (§1) in one place: **GEMMs → `tl.composite`; softmax merge +
tree reduction → kernel; `tl.load` is lazy.** Per-KV-head, `G` folded into
the matmul M dim (real GQA, no broadcast). This is the reference shape; the
prose sections elaborate each piece.
```python
def gqa_attention(q_ptr, k_ptr, v_ptr, o_ptr,
counter, start_pe, N, q_block, scale, *, tl):
# ---- geometry (kernel arithmetic, §2) ----
pe_id = tl.program_id(axis=rank_axis)
G = H_q // H_kv # query heads per KV head (=8)
causal = q_block.is_prefill
for kv in range(H_kv): # one KV head per iteration (§5.2)
# Q group: G rows folded into M → [G·T_q, d]. Lazy load (ADR-0062):
# issue now, auto-wait at first use inside the first composite.
q_g = tl.load(q_base(kv), (G * q_block.T_q, d)) # [G·T_q, d]
my_len = valid_len(counter, start_pe, pe_id, N) if not causal else S_kv
n_tiles = ceil(my_len / TILE)
# persistent arena (outside scratch_scope, ADR-0063): -inf, 0, zeros
m, l, O = init_running(G * q_block.T_q, d)
for j in range(n_tiles):
if causal and tile_all_future(j, q_block):
continue # causal skip (kernel if)
with tl.scratch_scope(): # per-tile temporaries (ADR-0063)
# --- GEMM #1 on the scheduler: Q·Kⱼᵀ; K streamed by composite ---
Sj = tl.composite("gemm", a=q_g,
b=tl.ref(k_tile(kv, j), (d, TILE))) * scale # Kᵀ pre-stored [d,TILE]
if causal and tile_partial(j, q_block):
Sj = Sj + causal_mask(j, q_block) # additive boundary mask
# --- online softmax merge in the kernel (MATH ops) ---
m_new = tl.maximum(m, tl.max(Sj, axis=-1))
P = tl.exp(Sj - m_new)
corr = tl.exp(m - m_new)
l = l * corr + tl.sum(P, axis=-1)
# --- GEMM #2 on the scheduler: P·Vⱼ; V streamed by composite ---
Oj = tl.composite("gemm", a=P,
b=tl.ref(v_tile(kv, j), (TILE, d)))
O = O * corr + Oj # running merge
m = m_new
# ---- cross-PE combine (§4): log-sum-exp tree to root, or store ----
if N == 1:
tl.store(o_base(kv), O / l)
else:
tree_reduce_and_store(m, l, O, pe_id, N, o_base(kv)) # tl.send/tl.recv
```
> **Why this shape:** the two `tl.composite` calls offload tiling + K/V DMA
> streaming + cross-tile pipelining to PE_SCHEDULER (CPU issues coarse
> descriptors and runs ahead → engines stay saturated, §1); the softmax
> merge and the IPCQ tree reduction stay in the kernel because the existing
> `CompositeCmd` cannot carry cross-tile state. `K` is pre-stored
> transposed `[d, TILE]` to sidestep the reshape-not-transpose caveat
> (§3, §B). A bespoke "flash-composite" command kind is **not** introduced
> (§8 item 4).
---
## A. Relationship to existing kernbench work (read first)
kernbench **already runs FlashAttention with an online-softmax `(m, ℓ, O)`
merge over IPCQ today.** Two kernels exist:
| File | Role | Mechanism |
|---|---|---|
| `src/kernbench/benches/_attention_mesh_kv.py` | prefill (Ring K/V) | per-rank partial attention, bidirectional K/V fan-out, online-softmax fold |
| `src/kernbench/benches/_attention_mesh_mlo.py` | decode (split-KV) | per-rank one-shot partial attention, bidirectional `(m,ℓ,O)` fan-out, log-sum-exp merge |
Both are driven by `milestone-gqa-llama70b` (4 panels:
`{single,multi}_user × {prefill,decode}`,
`src/kernbench/benches/milestone_gqa_llama70b.py`) and tested in
`tests/attention/test_milestone_gqa_llama70b.py`.
**They are written in the greenlet `tl` API:** `tl.load`, `tl.dot`,
`tl.softmax`/`tl.max`/`tl.sum`/`tl.exp`, `tl.send`/`tl.recv`, and Python
`-`/`*`/`/` on `TensorHandle` (each emits a `MathCmd`) — the GEMMs as
**blocking `tl.dot`**, not composites. The running `(m, ℓ, O)` is just
Python `TensorHandle`s threaded through the loop. This ADR keeps the
running state and the softmax merge in the kernel but **moves the two
GEMMs onto the scheduler-managed `tl.composite` path** (see §1) — this
**matters for the design**.
**Three deliberate limitations of the baseline** — exactly what an
*efficient GQA* kernel must lift:
1. **No GQA reuse.** `h_q == h_kv == 1`
(`test_milestone_gqa_llama70b.py:137-142`). The test attributes this to
a MemoryStore byte-conservation failure on a *broadcast view*, but that
failure is a property of the baseline's **head-packing hack**
(`_view(K, (h_q·d, S_kv))`, which conflates all heads into one matmul
dim and only conserves bytes when `h_q == h_kv`). The correct fix is
**kernel restructuring**, not a broadcast op: process **one KV head at
a time** and fold the `G` group rows into the matmul **M** dimension
(§5.2). That uses only byte-conserving reshapes, so real GQA
(`h_q = G·h_kv`) runs with **no new primitive** — see §8.
2. **O(N) reduction.** The baseline does an all-to-all bidirectional
fan-out so *every* rank ends with the full answer (`n_ranks − 1`
steps). Attention only needs `O` at the query owner → a **tree
reduction** to root is `⌈log₂ N⌉` steps (§4).
3. **Validation scale only.** `S = 16` because the 1 MiB scratch bump
allocator leaks per-tile temporaries
(`test_milestone_gqa_llama70b.py:123-148`) and there is no causal
masking / tiling → fixed by **ADR-0063** (recycling) + §5 (tiling,
causal skip) + composite K/V streaming (§3) + **ADR-0062** (lazy load
overlap).
**Documentation debt (out of scope but recorded):** the baseline cites
ADR-0055/0056/0057/0058/0059, **none of which exist as files** — they are
ghost references. This ADR does not retro-write them; see the Detailed
Design Document's *Open Decisions* for the recommendation.
---
## 0. Reference dimensions (Llama3-70B)
| Symbol | Meaning | Value |
|---|---|---|
| `H_q` | query heads | 64 |
| `H_kv` | KV heads | 8 |
| `G` | GQA group size = `H_q / H_kv` | 8 |
| `d` | head dim | 128 |
| `L` | layers | 80 |
| `D` | model dim | 8192 |
Hardware recap:
- **AHBM (chip)** = set of **CUBE**s (memory cubes, each with a logic die
containing **PE**s) + **IO die** (ADR-0003).
- **IPCQ**: PE↔PE queues, 4 mesh-direction queue-pairs per PE
(`N/S/E/W`, ADR-0023 D3; `global_*` for inter-SIP, ADR-0032). Kernel
API: `tl.send(dir, src)` / `tl.recv(dir, shape, dtype)`
(`tl_context.py:402-499`).
- **Composite command** (`CompositeCmd`, `pe_commands.py:144-162`): a
single GEMM (or MATH) *head* plus element-wise *epilogue* stages
(`bias/relu/scale/add/...`), issued **non-blocking** to PE_SCHEDULER,
which generates a tile plan and streams DMA→GEMM→write per tile
(ADR-0014 D6; `pe_scheduler.py:104-143`). It is **not** a general
multi-op DAG: it cannot chain two GEMMs, cannot carry register state
across instances, and cannot pop/wait on IPCQ. This ADR therefore issues
**each** of the two attention GEMMs (Q·Kᵀ, P·V) as its own composite and
keeps the cross-GEMM softmax merge + the IPCQ reduction in the kernel —
it does **not** need a new "flash-composite" command kind (see §1, §8).
- **Allocation policy (SP):** one query head per CUBE in the multi-user
panels; the `G` query heads of a KV group map within a CUBE's PEs.
**Two orthogonal mapping layers** (round-robin placement):
- **Layer 1 (KV-parallel, intra-request):** KV token `i` lands on PE
`(start_pe + i) mod N` — round-robin so one long request is split
across `N` PEs, balanced to ≤1 token.
- **Layer 2 (inter-request):** `start_pe = request_id mod N` rotates the
"PE-0 role" per request.
- Combined: `pe = (request_id + global_token_idx) mod N`.
`N` = number of PEs in the KV-parallel reduction group for one
(query head, request). Reduction stays **intra-CUBE** (a query head never
spans CUBEs — explicit non-goal).
---
## 0.5 Kernel boundary, preconditions, I/O contract
### 0.5.1 Position in the decoder layer
```
1. RMSNorm
2. QKV projection (GEMM) ─┐ qkv_rope kernel (SEPARATE, upstream)
3. RoPE on Q and K ─┤
4. write new K,V → KV cache ─┘
5. ===== THIS KERNEL: FlashAttention ===== (post-RoPE Q, K-cache, V-cache → O)
6. Output projection (GEMM) out_proj kernel (SEPARATE, downstream)
7. residual add → FFN ...
```
**Preconditions (upstream `qkv_rope`, NOT this kernel):**
- **P1.** Q is already RoPE-rotated. No rotation here.
- **P2.** K-cache stores **post-RoPE** K (never re-rotated at attention
time — the reason for post-RoPE caching).
- **P3.** For decode, the new step's K row is RoPE-rotated and appended
to its owning PE's K-cache slot **by `qkv_rope` before this kernel
launches.** ⇒ this kernel is **pure read** on the KV cache.
- **P4.** V is not rotated; V-cache holds raw projected V.
- **P5.** RoPE position upstream is the token's **absolute global
position** `global_idx = local_slot·N + ((pe_id − start_pe) mod N)`
(§2.1). Round-robin placement ≠ RoPE position.
### 0.5.2 Shape symbols
`T_q` = query length this launch (decode: 1; prefill/chunk: chunk width).
`S` = total context length. `S_pe` = keys owned by this PE ≈ `⌈S/N⌉`.
Storage `bf16` (numpy proxy `f16`, `memory_store.py:16`); `m,ℓ,O`
accumulators `f32` where precision matters.
### 0.5.3 INPUTS (per kernel launch)
| Input | Shape (per KV head) | Location | Notes |
|---|---|---|---|
| `Q` | `[G, T_q, d]` | per-PE HBM (loaded to TCM) | post-RoPE (P1). `G` query rows of the group batched. |
| `K_cache` | `[S_pe, d]` | per-PE HBM, base `K_base[pe]`, contiguous | post-RoPE (P2). Read-only. Dense locally, strided in global index (§2.1). |
| `V_cache` | `[S_pe, d]` | per-PE HBM, base `V_base[pe]` | raw V (P4). Read-only. |
| `global_token_counter` | scalar | launch arg | kernel derives `S_pe`, slot↔global, causal bounds. |
| `start_pe` (= `request_id mod N`) | scalar | launch arg | Layer-2 rotation. |
| `pe_id`, `N` | scalars | launch / `tl.program_id` | reduction-group geometry. |
| `q_block_meta` `{q_start, T_q}` | launch arg | prefill/SP causal masking & skip. |
| `O_base` | address | launch arg | where final O is written (root PE only). |
| `softmax_scale` | scalar | launch arg | `1/√d`. |
For **Ring Attention (§5.5)** add per ring step: incoming `K_block,
V_block` via IPCQ into ping-pong buffers (post-RoPE), plus
`step_kv_global_range` for the causal step-skip.
### 0.5.4 OUTPUTS
| Output | Shape | Location | Notes |
|---|---|---|---|
| `O` (final) | `[G, T_q, d]` | `O_base`, **root PE only** | normalised `O_acc / ℓ_acc`, cast to storage dtype. |
**Intermediate (non-root PE, on IPCQ — not a kernel-visible output):**
`(m_i, ℓ_i, O_i)` (`m,ℓ`: `[G, T_q]`; `O_i`: `[G, T_q, d]` unnormalised)
pushed to `tree_parent` via `tl.send` (§4). `O_i` is the heavy part.
**No-SP (`N=1`):** no IPCQ partials; the single PE's running `(m,ℓ,O)` is
normalised in place and written to `O_base`.
**Explicitly NOT outputs:** KV-cache writes (upstream, P3), output
projection (downstream), score `S` and probs `P` (never materialised).
---
## 1. Decision (mechanism)
**Implement the kernel as a *hybrid*: issue the two GEMMs (Q·Kᵀ and P·V)
as scheduler-managed `tl.composite(op="gemm")` commands, and keep the
online-softmax merge and the cross-PE reduction as kernel-level `tl`
ops.** `tl.load` is **lazy** (non-blocking; the wait is auto-inserted at
first use of the loaded data — ADR-0062), so explicit HBM loads overlap
the compute that follows. Rationale grounded in kernbench's execution +
latency model:
- **GEMM tiling is offloaded to PE_SCHEDULER.** A `CompositeCmd` is
non-blocking (`kernel_runner.py:182-191`, `pe_scheduler.py:104-121`):
the kernel pushes **one coarse descriptor** (M = `G·T_q`, the whole
per-PE tile sweep) and the scheduler generates the tile plan and streams
DMA→GEMM→write per tile (ADR-0014 D6). K/V are `tl.ref` operands the
scheduler streams from HBM, so per-tile **K/V prefetch is the
scheduler's job** — no explicit prefetch op. The CPU (greenlet) is freed
to issue the **next** composite while the current one runs, so the
scheduler keeps the GEMM engine saturated across tiles.
- **This reflects the hardware** and decouples CPU issue-rate from
execution-rate. The blocking per-op `tl.dot` path, by contrast, stalls
the CPU on every GEMM and leaves GEMM-engine **bubbles** during the
interleaved softmax MATH ops; it is realistic only if the CPU can keep
up with fine-grained per-tile issue.
- **The running `(m, ℓ, O)` flash state stays Python `TensorHandle`s**
threaded through the loop (the baseline already does this); the softmax
merge (max/exp/sum/rescale) is kernel-level `tl` MATH **between** the two
GEMM composites. The existing `CompositeCmd` cannot chain two GEMMs or
carry cross-tile register state (§0), so the merge necessarily lives in
the kernel — this is the hybrid split, not a limitation worked around.
- **Cross-PE combination** is a log-sum-exp **tree** over `(m, ℓ, O)`
after P·V, via kernel-level `tl.send`/`tl.recv` (§4) — unchanged.
So the per-tile inner pipeline is:
```
q_g = tl.load(Q group) # lazy; auto-wait at first use
per tile j:
Sⱼ = tl.composite("gemm", a=q_g, b=tl.ref(Kⱼ)) → Sⱼ # scheduler streams Kⱼ DMA + GEMM
Sⱼ += maskⱼ # kernel MATH, boundary tile only
online-softmax: mⱼ, m_new, P, corr, ℓ # kernel MATH
Oⱼ = tl.composite("gemm", a=P, b=tl.ref(Vⱼ)) → Oⱼ # scheduler streams Vⱼ DMA + GEMM
O = O*corr + Oⱼ; m = m_new # kernel MATH (running merge)
```
with each tile's MATH temporaries wrapped in `tl.scratch_scope`
(ADR-0063) so scratch stays O(1), and the next tile's composites issued
before the current tile's results are waited on (non-blocking handles) so
the scheduler pipelines across tiles.
Control flow (tile skip, mask generation, reduction scheduling, address
arithmetic) lives in the **kernel** (plain Python `if`/arithmetic in the
greenlet body). This is exactly what kernbench's greenlet model already
permits (`kernel_runner.py`, ADR-0020 D3).
> **What this supersedes.** An earlier iteration proposed a pure greenlet
> primitive path (all `tl.dot`, no composite) on the argument that
> "composite yields no latency benefit." That holds **only because** the
> simulator currently charges **zero** per-op CPU issue cost
> (`dispatch_cycles=0`, `pe_cpu.py`) — it models away exactly the CPU
> issue-rate / DMA-program cost that descriptor offload exists to hide.
> The hybrid is the faithful representation of an efficient kernel. The
> **measurable** size of the win (can the CPU saturate the engines for
> many tiles?) is gated on modelling an op-type-differentiated issue cost,
> tracked as future work (cost model; §9). Even at `dispatch_cycles=0` the
> non-blocking composite path fills the GEMM-engine bubbles the blocking
> `tl.dot` path leaves.
>
> **Why this is the efficient choice.** GEMM tiling + DMA streaming +
> cross-tile pipelining are offloaded to the proven `CompositeCmd`
> scheduler path; the softmax merge and the proven IPCQ collective stay in
> the kernel. The only genuinely new machinery is two small, general
> primitives (ADR-0062 lazy `tl.load`, ADR-0063 `tl.scratch_scope`); the
> reduction reuses `tl.send`/`tl.recv`. A bespoke "flash-composite"
> command (one kind internalising the softmax merge + carried register
> state + an IPCQ-push epilogue) is **not** built — large, special-purpose,
> and its only delta over this hybrid (full softmax offload) is not
> justified at the current modelling fidelity; see §8.
---
## 2. Memory layout & driver responsibilities
### 2.1 KV cache allocation
- Per-PE KV buffers sized to per-PE max context `⌈max_context / N⌉ × d ×
dtype`, for K and V separately. In kernbench these are deployed with a
`DPPolicy` whose `pe` (or `cube`) axis is `row_wise` over the sequence
dimension (`policy/placement/dp.py`; the baseline uses
`DPPolicy(pe="row_wise")` for K/V and `replicate` for Q).
- Within a PE, assigned tokens are **appended contiguously** (slot
0,1,2,…). Global indices are strided (`i, i+N, …`) but the per-PE
buffer is dense ⇒ DMA stays contiguous.
- Slot → global: `global_idx = local_slot·N + ((pe_id − start_pe) mod N)`.
### 2.2 Driver per-launch duties (minimal)
The driver supplies, per launch, the bases + counter + rotation; the
kernel derives the rest:
| Launch arg | Purpose |
|---|---|
| `K_base[pe]`, `V_base[pe]` | per-PE KV buffer bases (from tensor VA) |
| `O_base` | attention output destination |
| `global_token_counter` | current sequence position |
| `start_pe = request_id mod N` | Layer-2 rotation |
| `pe_id` (`tl.program_id`), `N` | geometry |
| `q_block_meta` | prefill/SP query block start + length |
Kernel-derived (plain arithmetic):
- **my turn to write this step:** `(start_pe + counter) mod N == pe_id`
- **my valid length:** `base = counter // N; rem = counter % N;
my_len = base + (1 if ((pe_id − start_pe) mod N) < rem else 0)`
- **read range:** tiles `0 .. ⌈my_len / TILE⌉`.
- **causal bounds / per-tile skip:** from global positions of the query
block vs each tile.
Driver = bases + counter + rotation. The single formula
`(request_id + token_idx) mod N` is the entire placement policy.
---
## 3. Per-tile op sequence (greenlet `tl`)
One iteration = one KV tile on one PE. The two GEMMs are
`tl.composite(op="gemm")` (scheduler-managed tiling + K/V DMA streaming);
the softmax merge is kernel `tl` MATH between them. Real `tl` names
(`tl_context.py`), with lazy `tl.load` (ADR-0062), `tl.scratch_scope`
(ADR-0063):
```python
# running state (persistent arena — allocated once, outside the scope)
# m: [G, T_q] l: [G, T_q] O: [G, T_q, d]
q_g = tl.load(Q_group_ptr, (G*T_q, d)) # lazy; auto-wait at first use (ADR-0062)
with tl.scratch_scope(): # per-tile MATH temporaries recycled
Sj = tl.composite("gemm", a=q_g, # [G·T_q, TILE]; scheduler streams Kⱼ DMA
b=tl.ref(K_base + j*TILE*d, (TILE, d))) * softmax_scale
if mask_j is not None:
Sj = Sj + mask_j # additive causal mask (boundary tile)
m_j = tl.max(Sj, axis=-1)
m_new = tl.maximum(m, m_j)
P = tl.exp(Sj - m_new) # no full-matrix softmax; streaming
corr = tl.exp(m - m_new) # rescale factor for old accumulators
l = l * corr + tl.sum(P, axis=-1)
Oj = tl.composite("gemm", a=P, # [G·T_q, d]; scheduler streams Vⱼ DMA
b=tl.ref(V_base + j*TILE*d, (TILE, d)))
O = O * corr + Oj # running merge (kernel MATH)
m = m_new
```
Notes:
- `q_g` is the GQA-batched query reshaped to `[G·T_q, d]` (the `G` group
rows folded into the matmul M dim; byte-conserving). One K/V tile serves
all `G·T_q` rows — the GQA reuse lever — with no broadcast.
- **K/V are `tl.ref` operands** the composite scheduler streams from HBM
per tile (`pe_scheduler.py:104-143`): that *is* the prefetch/pipeline,
so there is no explicit prefetch op. Issuing tile `j+1`'s composites
before waiting on tile `j` (non-blocking handles) keeps the scheduler
pipelined across tiles and the GEMM engine saturated.
- `tl.trans` is **metadata-only** in kernbench (`tl_context.py:390`) and
`MemoryStore.read` *reshapes* rather than transposes
(`memory_store.py:73`). For zero/structural runs this is harmless; for
non-trivial numeric data it yields a reshape-not-transpose, so Q·Kᵀ via
a transposed K needs care (§11) — store K pre-transposed `[d, TILE]`, or
add a real `tl.transpose` (a candidate further primitive, likely
unnecessary given the simulator's performance-modeling purpose).
- Masking: the kernel builds the boundary-tile mask from query/KV global
offsets and adds it (`Sj + mask_j`); full-past tiles pass `None`;
full-future tiles are **skipped** (the `if` never enqueues them).
- A **new** "flash-composite" command kind (one that internalises the
softmax merge + carried `(m,ℓ,O)`) is **not** used; the existing
`CompositeCmd` covers each GEMM and the merge stays in the kernel
(§1, §8 item 4).
---
## 4. Reduction (KV-parallel / SP combine) — tree to root
After its tile sweep each PE holds `(m_i, ℓ_i, O_i)` (unnormalised).
Combine via the associative/commutative log-sum-exp merge (identical math
to the baseline's fold, `_attention_mesh_mlo.py:117-122`):
```python
def merge(m_a, l_a, O_a, m_b, l_b, O_b):
m = tl.maximum(m_a, m_b)
sa, sb = tl.exp(m_a - m), tl.exp(m_b - m)
return m, l_a*sa + l_b*sb, O_a*sa + O_b*sb
# final: O = O_root / l_root # normalise once at the tree root
```
- **Timing:** exchange is **after P·V** (each PE has its final `O_i`).
- **Topology:** mesh tree, depth `⌈log₂ N⌉`, over physical neighbours.
For `N=8`: level-0 `(0↔1,2↔3,4↔5,6↔7)`, level-1 `(1↔3,5↔7)`, level-2
`(3↔7)`, root = `7`. Each tree pair must be a physical `N/S/E/W`
neighbour so `tl.send(dir)`/`tl.recv(dir)` use a real link
(tuning item §9; the SFR install that wires neighbours is
`configure_sfr_intracube_pe_ring`, used by the baseline).
- **Why a tree, not the baseline fan-out:** attention needs `O` at one
place (the PE that writes `O_base`). A tree is `⌈log₂ N⌉` steps
(3 for `N=8`) vs the baseline all-to-all's `N−1` (7). For decode, where
the per-PE sweep is short, the reduction is the dominant cost, so this
is a real win. (The baseline's fan-out replicates the answer to all
ranks — unnecessary here.)
- **Payload:** `O_i` (`d`-vector) is heavy; `m,ℓ` are scalars per `(g,
T_q)`. Sent as separate `tl.send`s (the baseline sends the triplet as
three sends — same pattern).
Kernel structure (greenlet):
```python
# leaf / internal node: fold children that send to me, then forward up
for child_dir in tree_children_dirs(pe_id, N):
m_c = tl.recv(child_dir, m.shape); l_c = tl.recv(child_dir, l.shape)
O_c = tl.recv(child_dir, O.shape)
m, l, O = merge(m, l, O, m_c, l_c, O_c)
if not is_root(pe_id):
tl.send(parent_dir(pe_id), m); tl.send(parent_dir(pe_id), l); tl.send(parent_dir(pe_id), O)
else:
tl.store(O_base, O / l)
```
`tl.recv` blocks until data arrives (`tl_context.py:446-499`); the merge
order is fixed by the static tree, so no data-dependent `pop` is needed —
which is why **no hardware/composite `pop`-as-dependency change is
required** (§6).
---
## 5. Per-case kernels
All four cases share §3 (tile sweep) + §4 (reduction); they differ only
in `N`, query-block width, and which `tile_*` predicates fire.
### 5.1 Common skeleton
```python
def attn_kernel(q_ptr, k_ptr, v_ptr, o_ptr, counter, start_pe, N,
q_block, softmax_scale, *, tl):
pe_id = tl.program_id(axis=rank_axis)
my_len = valid_len(counter, start_pe, pe_id, N)
n_tiles = ceil(my_len / TILE)
q_g = load_Q_group(q_ptr) # [G·T_q, d]; lazy tl.load, G folded into M (no broadcast)
m, l, O = init_running() # persistent arena: -inf, 0, zeros
for j in range(n_tiles): # K/V streamed per tile by the composite scheduler (§3)
if tile_all_future(j, q_block): # causal skip (kernel if)
continue
run_tile(j, mask_or_null(j, q_block)) # §3: 2 composites + softmax MATH, in tl.scratch_scope
if N == 1:
tl.store(o_ptr, O / l) # no reduction
else:
tree_reduce_and_store(pe_id, N, o_ptr) # §4
```
### 5.2 DECODE, no SP (`N=1`, one PE owns the head's KV)
- `T_q = 1`, attends to all past KV ⇒ no future tiles, mask only on the
final ragged tile.
- **GQA reuse is the whole game** (decode is KV-load-bound): the `G=8`
query rows of the KV head are folded into the matmul **M** dimension
(`q_g` reshaped `[G, T_q, d] → [G·T_q, d]`, a byte-conserving reshape).
Then `Q·Kᵀ` is `composite([G·T_q, d], Kᵀ[d, TILE]) → [G·T_q, TILE]` and
`P·V` is `composite([G·T_q, TILE], V[TILE, d]) → [G·T_q, d]`. The KV tile
(`[TILE, d]`) is the shared `K`/`V` operand — **streamed once, reused by
all `G·T_q` rows automatically** because they are the M rows of the
GEMM. No broadcast of K/V is needed; `m = G·T_q` in the composite's tile
plan also makes the timing count all `G` rows' work correctly (a leading
batch axis would *not* be counted — see §8).
- If `S_pe` fits in scratch (small/medium context) this degenerates to a
**one-shot** partial attention (one composite for Q·Kᵀ, one softmax, one
composite for P·V) — exactly the baseline `_attention_mesh_mlo`
`_partial_attention`, just GQA-batched and on the composite path. Tiling
(§3) only kicks in when `S_pe` exceeds the scratch scope's tile budget.
### 5.3 DECODE, with SP / KV-parallel (`N=8`)
- One request's KV is round-robin across 8 PEs; each owns ≈`my_len`
tokens. Each PE GQA-batches its `G=8` query rows over its local tiles.
- `T_q=1` ⇒ short per-PE sweep ⇒ **reduction dominates** ⇒ the §4 tree
(3 steps) is the structurally important part. Reduction latency is
hidden by other concurrent decode tokens when batched, or by the long
per-PE tile sweep for long single-stream context.
### 5.4 PREFILL, no SP
- Whole prompt resident; query is a block of `T_q` tokens (chunked).
- Causality is real, per query block `[qs, qe)` vs KV tile `[ks, ke)`:
- `ke ≤ qs` → `tile_all_past` → `mask=None`, full compute.
- `ks ≥ qe` → `tile_all_future` → **skip** (kernel `if`).
- overlap → `tile_partial` → kernel builds a triangular additive mask,
the tile op sequence adds it (`Sj + mask_j`).
- GQA batches the group's query heads along the same axis as decode.
### 5.5 PREFILL, with SP (Ring KV)
- KV sharded across `N` PEs (the ring); each step delivers a peer's KV
block over IPCQ. Running `(m,ℓ,O)` is carried **across ring steps**
(Python handles, exactly as `_attention_mesh_kv` does today).
- IPCQ overlaps next step's KV receive with current step's compute via
`tl.recv_async`/`tl.wait` (`tl_context.py:543-560` — already exists).
- **Causal ring skip:** if the incoming step's KV block is entirely
*after* the local query block, skip its compute (don't recv-consume /
don't fold) — eliminates ≈half the steps for causal attention.
- Receive buffers ping-pong (persistent arena, not recycled).
- After the ring, `(m,ℓ,O)` is final per query-owning PE; if KV-parallel
splits remain, §4 tree-merges them, then normalise + store.
The baseline `_attention_mesh_kv` already implements the ring fold; this
ADR adds GQA reuse, causal step-skip, and `recv_async` overlap.
---
## 6. Why no hardware / composite change is needed
- The only place a HW/composite `pop`-as-dependency would help is a lone
reduction with no other work to hide a poll — i.e. batch=1 **and** short
context. The target is **agentic = low batch, long context** ⇒ each PE
has many KV tiles; the §4 tree's `tl.recv` blocks are covered by the
sweep work of concurrent tokens / long context.
- The §4 reduction uses a **static** tree, so the recv order is fixed at
compile time — no data-dependent pop. `tl.recv` (blocking) suffices.
- **Decision: ship with the greenlet `tl.send`/`tl.recv` collective.**
Revisit a HW `pop` only if a short-context, single-stream,
latency-critical target emerges.
---
## 7. Control vs execution split (the load-bearing principle)
| Concern | Owner |
|---|---|
| tile skip (future), mask generation, causal bounds | **kernel** (Python `if` + arithmetic in greenlet body) |
| address / offset / valid-length arithmetic | **kernel** (from counter) |
| reduction scheduling, IPCQ send/recv ordering | **kernel** (static tree) |
| Q·Kᵀ, P·V (incl. per-tile K/V DMA streaming + tiling) | **`tl.composite`** → PE_SCHEDULER |
| Q load, mask add, softmax math, running `(m,ℓ,O)` merge | **`tl` ops** on PE engines (kernel-issued) |
The kernel decides; the GEMMs are offloaded to the scheduler as
composites; the remaining `tl` ops execute already-decided work. This is
exactly the greenlet + composite model kernbench already supports — no new
control abstraction.
---
## 8. Required kernbench changes
**Correction from design iteration:** real GQA (`h_q > h_kv`) needs **no
new primitive** — only the kernel restructuring in §5.2 (per KV head,
`G` folded into M, byte-conserving reshapes). The supporting ADRs are
*efficiency / scale* enablers, not GQA blockers.
**Algorithm work in the kernel (no new primitive; existing `tl` API):**
- **GQA Q-axis batching** (the reuse lever) — fold `G·T_q` into the matmul
M dim per KV head (§5.2); `_view`-style byte-conserving reshape; the
GEMM is a `tl.composite(op="gemm")` with M = `G·T_q`. Runs today in both
timing and data mode.
- **GEMMs via composite** (§1/§3) — Q·Kᵀ and P·V each issued as a
non-blocking `tl.composite(op="gemm")`; PE_SCHEDULER tiles them and
streams the `tl.ref` K/V operands' DMA (existing `CompositeCmd`; no new
command kind).
- Tree reduction to root (§4) replacing the baseline all-to-all fan-out —
pure kernel control flow over `tl.send`/`tl.recv`.
- Causal tile skip + additive boundary mask (§3/§5.4) — kernel `if` +
a mask tensor added with `+`.
- Round-robin KV placement / valid-length arithmetic (§2) — launch-arg
arithmetic + `DPPolicy(pe="row_wise")`.
**New primitives for efficiency / scale (each has a supporting ADR):**
1. **Per-tile scratch recycling** — **ADR-0063** (`tl.scratch_scope`).
*Required for scale*: removes the `S=16` ceiling (1 MiB bump
allocator) so realistic context lengths run. Highest-value of the
three.
2. **Lazy `tl.load`** — **ADR-0062** (non-blocking load + auto-wait on
first use; API surface unchanged). *Efficiency*: overlaps explicit
loads (the Q group, non-composite kernels) with following compute. The
per-tile **K/V** prefetch is handled by the composite scheduler (§1),
so this covers the remaining explicit loads. Global semantics change →
existing goldens regenerate (ADR-0062 D3).
3. **GQA head / mask broadcast** — **ADR-0061** (`tl.broadcast`).
*Optional convenience*, not a GQA blocker (see correction above).
Useful for additive-mask construction across the `G·T_q` rows and for
general kernels; `np.matmul` already broadcasts in data mode, so it is
not needed for correctness. Lowest priority.
**Explicitly REJECTED (efficient alternative chosen):**
4. ~~A bespoke "flash-composite" command kind that internalises the whole
inner loop — DMA→MM→VEC→DMA→MM→VEC with carried `(m,ℓ,O)` register
state and a tail IPCQ push.~~ The two GEMMs **do** use the existing
`CompositeCmd` (§1/§3) — that gives scheduler-managed tiling, K/V DMA
streaming, and cross-tile pipelining. What is rejected is a **new**
command kind that also absorbs the softmax merge + cross-tile register
lifetime + an IPCQ-push epilogue: it is large and special-purpose, and
its only delta over the hybrid (full softmax offload) is not justified
at the current modelling fidelity (per-op CPU issue cost = 0; see §1).
Revisit if the cost model (§9) makes full offload measurably worthwhile.
5. ~~Hardware `pop`-as-dependency.~~ Out of scope (§6).
6. ~~RoPE / QKV projection / KV-cache write inside this kernel.~~ Upstream
`qkv_rope` (P1–P5). Folding RoPE in would force re-rotating past tiles
every decode step and break Ring Attention's post-RoPE pass-through.
---
## 9. Open tuning items (measured in kernbench, not blocking)
1. **Composite tile-pipeline depth** — dominant for KV-load-bound; how far
ahead the kernel issues non-blocking composites before waiting, and the
scheduler's per-tile streaming depth.
2. **PE↔mesh-neighbour mapping for the reduction tree** — ensure each
depth-`⌈log₂ N⌉` pair is a physical `N/S/E/W` neighbour; bad mapping
adds hops. Verify against the SFR install.
3. **TILE size** — balance scratch residency (`S/P tiles + O_acc + G-way
GQA`) against DMA efficiency; interacts with ADR-0063 and the
scheduler's `TILE_M/K/N` (`pe_scheduler.py`).
4. **Ring buffer ping-pong vs `recv_async` depth** in §5.5.
5. **One-shot vs tiled crossover for decode** (§5.2) — the `S_pe`
threshold where tiling beats a single composite.
6. **Per-op CPU issue cost (cost model)** — currently `dispatch_cycles=0`
(`pe_cpu.py`), so composite-vs-primitive issue overhead is invisible.
An op-type-differentiated issue cost (a `tl.composite` descriptor push
≫ a primitive op) is what makes the hybrid's CPU-saturation win
**measurable** (§1). Specified in **ADR-0064**; tracked as separate
future work.
---
## 10. Coverage summary
| Case | KV placement | Inner loop | Reduction | Masking |
|---|---|---|---|---|
| Decode, no SP | 1 PE, all KV | one-shot or tiled, GQA-batched | none | last tile only |
| Decode, SP | round-robin `N` PEs | tiled, GQA-batched | §4 tree (`⌈log₂ N⌉`) | last tile only |
| Prefill, no SP | resident | tiled per q-block | none | skip future / triangular boundary |
| Prefill, SP (Ring) | ring-sharded | per ring step (recv_async overlap) | running state + §4 tree | causal step-skip + boundary |
**I/O per case** (full contract §0.5):
| Case | Inputs | Output | IPCQ partials |
|---|---|---|---|
| Decode, no SP | `Q[G,1,d]`, full `K/V[S,d]` on 1 PE | `O[G,1,d]` at `O_base` | none |
| Decode, SP | `Q[G,1,d]`, per-PE `K/V[S_pe,d]` | `O[G,1,d]` (root) | `(m,ℓ,O_i)` per PE → tree |
| Prefill, no SP | `Q[G,T_q,d]`, `K/V[≤end,d]` | `O[G,T_q,d]` at `O_base` | none |
| Prefill, SP (Ring) | `Q[G,T_q,d]`, ring-delivered `K/V` blocks | `O[G,T_q,d]` (root) | running `(m,ℓ,O)` + tree |
In all cases: KV-cache writes and RoPE happen **upstream**; output
projection **downstream**; score `S` and probs `P` are never materialised.
---
## 11. Verification plan (Phase 1 test outline)
SPEC/ADR coverage: R5 (PE↔PE IPCQ, PE↔HBM), R2 (latency by traversal),
ADR-0023/0025 (IPCQ), ADR-0046 (`tl` contract), ADR-0054 (eval bench).
The simulator's contract is **latency by traversal + determinism +
structural correctness** (SPEC §0, §0.1), not bit-exact numerics — Phase 2
data exists mainly to exercise the data path, `tl.trans` is
reshape-not-transpose, and `bf16` is modelled as `f16`. Verification is
therefore **structural/timing-first**, with numeric parity as a bounded
secondary check.
1. **Runs in data mode (`enable_data=True`):** the GQA kernel
(`h_q = G·h_kv`) completes without the byte-conservation error that the
baseline head-packing hits — for all four cases. (This needs the §5.2
restructuring, *not* a new primitive.)
**Numeric parity (secondary):** for symmetric/identity inputs where
reshape-as-transpose is exact, kernel `O` matches a numpy
FlashAttention reference within fp tolerance. Full asymmetric parity is
gated on a real `tl.transpose` (out of scope; flagged).
2. **GQA reuse:** with `h_q = G·h_kv`, the K/V `dma_read_count` is
independent of `G` (one load per tile, reused across the group), while
GEMM work scales with `G`. Asserts the lever actually fires.
3. **Tree reduction step count:** decode-SP issues `⌈log₂ N⌉` reduction
rounds (not `N−1`); op_log `ipcq_send`/`recv` counts match the tree.
4. **Causal skip:** prefill skips all `tile_all_future` tiles — GEMM
count equals the lower-triangular tile count, not the full grid.
5. **Long context (ADR-0063):** a sweep at `S` that overflows 1 MiB
without scopes completes and matches the reference.
6. **Load/compute overlap:** end-to-end latency of the tiled sweep is
below the serial `Σ(load+compute)` — from the composite scheduler
streaming K/V per tile (§1/§3) and lazy `tl.load` (ADR-0062) overlapping
the Q load. (Overlap is real modelled concurrency, not a subtraction.)
7. **Composite GEMM offload (structural):** each tile's Q·Kᵀ and P·V emit a
`CompositeCmd` (non-blocking) to PE_SCHEDULER, not a blocking `tl.dot`;
op_log shows the composite tile plan and the kernel issues the next
tile's composites before waiting (cross-tile pipelining).
8. **Determinism:** identical inputs → identical op_log + latency
(SPEC §0.1).
---
## B. Open design items from the hybrid pivot (review later)
These arose when the decision moved from a pure greenlet primitive path to
the **composite hybrid + lazy `tl.load`** (this revision). None blocks the
design; each needs a verification pass during implementation. Recorded here
(rather than asked) per the working agreement — the recommendation is my
predicted default; revise on review.
1. **DDD-0060 is not yet synced.** The Detailed Design Document still
describes the old `tl.load_async` double-buffer path and primitive
`tl.dot` inner loop (its §4.3/§5/§10). It must be updated to the hybrid
(composite GEMMs, lazy load, K pre-transposed). *Left for review*
because the DDD is a derived how-to and a large rewrite; the ADR is now
the authoritative record. **Recommend:** sync DDD as a follow-up before
implementation starts.
2. **K operand orientation for the composite GEMM.** Q·Kᵀ needs `b =
[d, TILE]`, but the KV cache stores K as `[S_pe, d]`. `tl.trans` is
metadata-only and `MemoryStore.read` reshapes, not transposes
(`memory_store.py:73`) — so a runtime transpose is wrong for non-trivial
data. **Recommend:** store K **pre-transposed** `[d, S_pe]` in the cache
(the pseudocode and §3 assume this), making `tl.ref(k_tile, (d, TILE))`
a contiguous slice. Verify the upstream `qkv_rope` write layout supports
this, or add a real `tl.transpose` (heavier; deferred).
3. **Composite output buffer vs `tl.scratch_scope`.** Each Q·Kᵀ composite
writes `Sj` to an `out_addr`; the kernel then reads it for the softmax
MATH. That output buffer, and the in-flight composites' targets, must
live where the per-tile `scratch_scope` (ADR-0063) will **not** recycle
them before they are consumed — same discipline as in-flight lazy loads
(ADR-0062 D-Negative). **Recommend:** composite outputs for the *current*
tile live in the scoped arena (consumed same iteration); the persistent
`(m,ℓ,O)` stays outside. Verify no use-after-recycle when the next
tile's composites are issued early (cross-tile pipelining).
4. **GQA `dma_read_count` lever under composite streaming.** The lever
(§11.2: K/V `dma_read_count` independent of `G`) assumes the composite
emits **one** K/V tile DMA reused across all `G·T_q` M-rows. The
scheduler's `generate_gemm_plan` tiles by `TILE_M/K/N`
(`pe_scheduler.py:35-37`, 32/64/32) — confirm the M-tiling over `G·T_q`
does **not** re-issue the shared K/V tile DMA per M-tile (i.e. operand
DMA is shared across M-tiles, or the lever weakens). **Recommend:**
assert it in the levers test; if violated, the GQA win is in compute
only, not DMA — still correct, but the headline changes.
5. **Kernel TILE vs scheduler `TILE_M/K/N`.** The kernel reasons about a
logical KV `TILE`; the scheduler re-tiles internally at fixed
`TILE_M/K/N`. Two tiling layers interact (scratch residency, pipeline
depth). **Recommend:** treat the kernel TILE as the K/V streaming
granularity and let the scheduler sub-tile the GEMM; document the
relationship in the DDD and sweep both (§9 items 1, 3).
6. **Cost model is a separate ADR.** The hybrid's CPU-saturation benefit is
invisible while `dispatch_cycles=0`. The per-op-type issue-cost model is
specified in **ADR-0064**; this ADR's §1/§9 depend on it for the
*measurable* (not just structural) win. **Recommend:** land ADR-0064's
model before claiming hybrid latency wins in the eval.
7. **Ring path (§5.5) GEMMs.** §5.5 still describes the ring fold with
primitive ops + `recv_async`. For consistency the ring's per-step Q·Kᵀ /
P·V should also be composites; the IPCQ `recv_async` overlap is
orthogonal and stays. **Recommend:** apply the same hybrid shape to the
ring step during implementation; low risk, mirrors §3.
docs/adr-proposed/ADR-0061-prog-tl-broadcast.md
+4 -6
View File
@@ -32,9 +32,8 @@ the group is the decode-time efficiency win (ADR-0060 §5.2).
### Why it does not work today ### Why it does not work today
The existing mesh kernels reshape tensors with a metadata-only helper Pre-ADR-0060 mesh kernels reshaped tensors with a metadata-only `_view`
`_view` (`src/kernbench/benches/_attention_mesh_kv.py:25-36`, helper that rewrote `shape` but kept the original `nbytes`:
`_attention_mesh_mlo.py:25-36`):
```python ```python
def _view(handle, new_shape): def _view(handle, new_shape):
@@ -62,8 +61,7 @@ A reshape conserves bytes; a **broadcast does not** (`[S_kv, 1, d]` →
kernel tries to view a `h_kv`-headed K as if it had `h_q` heads, the kernel tries to view a `h_kv`-headed K as if it had `h_q` heads, the
stored array has `1/G` of the requested bytes and `read` raises. stored array has `1/G` of the requested bytes and `read` raises.
The current test suite documents this as a deliberate limitation The pre-ADR-0060 baseline documented this as a deliberate limitation:
(`tests/attention/test_milestone_gqa_llama70b.py:137-142`):
> "v1 uses `h_q == h_kv == 1` to avoid … GQA broadcast view (which is > "v1 uses `h_q == h_kv == 1` to avoid … GQA broadcast view (which is
> symbolic and does not survive MemoryStore's nbytes check under > symbolic and does not survive MemoryStore's nbytes check under
@@ -196,7 +194,7 @@ entire point of GQA (it `G×`'s KV-cache HBM footprint and bandwidth).
2. **Phase 2 data**: broadcasting a known array yields the numpy 2. **Phase 2 data**: broadcasting a known array yields the numpy
`broadcast_to(...).copy()` result; a downstream `tl.dot` consuming it `broadcast_to(...).copy()` result; a downstream `tl.dot` consuming it
passes the MemoryStore nbytes check (the exact case that fails today). passes the MemoryStore nbytes check (the exact case that fails today).
3. **GQA end-to-end**: re-run a reduced `milestone-gqa-llama70b` panel 3. **GQA end-to-end**: re-run a reduced `milestone-gqa-headline` panel
with `h_q=8, h_kv=1` under `enable_data=True`; it must complete (today with `h_q=8, h_kv=1` under `enable_data=True`; it must complete (today
it raises `ValueError: Shape mismatch`). it raises `ValueError: Shape mismatch`).
4. **Broadcast-incompatible shape** (e.g. axis size 3 → 8) raises a 4. **Broadcast-incompatible shape** (e.g. axis size 3 → 8) raises a
docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model-ko.md
-166
View File
@@ -1,166 +0,0 @@
# ADR-0064: op 종류별 CPU 발행 비용 모델 (command construct + dispatch)
## Status
Proposed
> **ADR-0060**(AHBM GQA Fused Attention)의 보조 ADR. 거기서의 하이브리드
> 결정(GEMM은 `tl.composite`, softmax 머지는 커널)이 이기는 이유는
> **타일링을 PE_SCHEDULER에 offload하여 CPU가 coarse descriptor를 내리고
> 앞서 나가 엔진을 saturate**하기 때문이다. 그 이점이 현재 시뮬레이터에서는
> per-op CPU 발행 비용이 0이라 **보이지 않는다**. 본 ADR은 발행 비용을
> 실재화하고 **op 종류별로 차등**화하여, composite-vs-primitive 트레이드오프
> (그리고 CPU saturation)가 측정 가능하도록 한다.
## Context
### 현재 상태
- 모든 `tl.*` op은 커맨드를 emit하기 전에 `_emit_dispatch_overhead()`를
호출하며(`tl_context.py:123-125, 190, 227, 235, 612, …`), 이는
`dispatch_cycles > 0`일 때**만** `PeCpuOverheadCmd(cycles=dispatch_cycles)`를
emit한다. 즉 발행 비용은 `tl.load`, `tl.dot`, MATH op, `tl.composite`에
동일하게 적용되는 **단일 균일 노브**다.
- 그 노브는 두 실행 경로 모두에서 **0**으로 하드코딩되어 있다
(`pe_cpu.py:101` greenlet runner, `:195` legacy replay). ⇒ 커맨드 발행
— *descriptor를 construct하고 scheduler 큐에 push* — 이 현재 PE_CPU에서
**0 ns**다.
- `PeCpuOverheadCmd`는 PE_CPU(`kernel_runner.py:131-132`)와
PE_SCHEDULER(`pe_scheduler.py:97-100`)에서 `yield env.timeout(cmd.cycles)`로
소비된다. 수동 `tl.cycles(n)`도 존재한다(`tl_context.py:695`).
### 본 ADR을 규정하는 두 가지 발견 (ADR-0060 검토에서)
- **Q1 — composite 발행 비용.** `tl.composite`를 construct + push 하는 비용은
CPU 시간 기준 **~40 ns** 정도(descriptor 빌드 + 큐 push)로 예상된다. 현재는
**0**이다. 훅은 존재하며, 값(그리고 op 종류별 차등)만 빠져 있다.
- **Q2 — scheduler dispatch vs DMA latency.** PE_SCHEDULER의 composite
dispatch는 **non-blocking**이다: `_dispatch_composite`는 tile plan을 생성하고
feeder에 enqueue한 뒤 즉시 반환한다(`pe_scheduler.py:104-121`). 실제 **DMA
latency는 타일이 흐를 때 PE_DMA에서 부과**되며(`drain_ns =
compute_drain_ns(path, nbytes)`, `pe_dma.py:89`), scheduler dispatch에
덩어리로 들어가지 **않는다** ⇒ **중복 계상 없음**, DMA는 modeled component에
머문다(SPEC §0.1). scheduler의 plan 생성 자체는 현재 sim 시간 **0**이다.
### 균일-그리고-0이 하이브리드에 틀린 이유
하이브리드의 논거 전체는 **하나의** composite descriptor가 `N_tiles`만큼의
GEMM 타일링을 offload하므로, CPU가 `O(N_tiles × ops/tile)`개의 fine 커맨드
대신 `O(1)`개의 coarse 커맨드를 낸다는 것이다. 발행 비용 = 0(그리고 균일)이면
모델은 다음을 보여줄 수 없다:
- CPU가 충분히 빨리 못 밀어넣을 때 primitive 경로가 엔진을 **saturate 못 할 수
있음**(핵심 ADR-0060 §1 주장), 그리고
- composite가 단일 primitive보다 construct 비용이 **더 크지만** 그것이 대체하는
다수의 primitive보다는 훨씬 작음.
단일 균일 `dispatch_cycles`로는 이를 표현할 수 없다: 실제 construct 비용에서
`tl.load` ≪ `tl.composite`이기 때문이다.
## Decision
### D1. PE_CPU의 op 종류별 발행 비용 테이블
단일 `dispatch_cycles` 스칼라를 **커맨드 종류별 비용 테이블**로 교체하고,
발행 시점(커맨드가 scheduler로 dispatch되기 전)에 PE_CPU에서 부과한다.
greenlet이 이 시간을 지불하며, 이것이 바로 CPU가 얼마나 빨리 work를 밀어넣을
수 있는지를 gate한다 — saturation 레버.
시작 추정치(추후 calibrate — 검토 항목 참조):
| 발행 op | CPU 발행 비용 (construct + push) |
|---|---|
| `tl.composite` (GEMM/MATH descriptor) | ~40 ns (Q1 추정) |
| `tl.load` / `tl.store` (DMA descriptor) | 작음 (수 ns) |
| `tl.dot` / MATH / IPCQ send/recv | 작음 (수 ns) |
요점은 **비율**이다: composite construct ≫ primitive 발행, 그러나 그것이
대체하는 primitive 발행들의 합 ≪. 정확한 ns는 설정 가능하다.
### D2. 실행 latency는 엔진에 유지 (Q2 — 변경 없음)
DMA/GEMM/MATH 실행 latency는 오늘처럼 PE_DMA / PE_GEMM / PE_MATH에서 부과된다.
발행 비용(D1)은 **CPU 측에만** 추가되며, `drain_ns`를 건드리지 않으므로 중복
계상이 없다. 이는 SPEC §0.1(latency는 modeled component에서)을 보존한다.
### D3. scheduler plan 생성 비용 — 0에서 시작, 재검토
PE_SCHEDULER의 tile-plan 생성은 오늘 sim 시간 0이다. 일단 유지한다(지배적
레버는 CPU 발행 비용 D1); scheduler 측 비용이 유의미하다고 판명되면 기존
`overhead_ns` node attr로 노출한다. 여기서 결정하지 않고 검토 항목으로 둔다.
### D4. 설정 가능한 값; 골든 재생성
비용 테이블은 설정 가능하다(per-topology / node attrs, 기본값은 D1 추정치).
발행 비용을 0이 아니게 하면 **모든** bench latency가 바뀌므로, 골든 latency를
한 번 **재생성**한다 — ADR-0062의 전역 lazy-load 회귀와 동일한 자세의 모델
충실도 개선이다.
## Alternatives
### A1. 단일 균일 `dispatch_cycles > 0` 유지
기각: op construct 비용은 자릿수 단위로 다르다(`tl.load` vs `tl.composite`).
균일값은 primitive를 과대 청구하거나 composite를 과소 청구하며, 어느 쪽이든
하이브리드가 의존하는 composite-vs-primitive 트레이드오프를 왜곡한다.
### A2. 발행 비용을 PE_CPU 대신 PE_SCHEDULER에 부과
기각: saturation 질문은 *"CPU가 엔진을 바쁘게 유지할 만큼 descriptor를 빨리
밀어넣을 수 있는가?"*이며 — 이는 **PE_CPU**의 issue-bandwidth 속성이다.
scheduler에 부과하면 CPU back-pressure를 모델하지 못한다.
### A3. DMA program/setup 시간을 별도의 고정 per-descriptor 비용으로 모델
실제 HW는 전송 시간과 별개로 DMA descriptor program 비용을 지불한다. 연기:
초기에는 descriptor-program 비용을 **발행 op**의 D1 비용에 합친다(DMA를
유발하는 `tl.load` / composite). calibration이 유의미함을 보이면 PE_DMA 고정
setup으로 분리한다. 검토 항목.
## Consequences
### Positive
- 하이브리드의 CPU-offload / saturation 이점(ADR-0060 §1)이 구조적일 뿐
아니라 **측정 가능**해진다.
- composite-vs-primitive 및 타일링 granularity 트레이드오프가 latency에
보여 eval 서사를 가능케 한다.
- 하드웨어에 더 충실(issue bandwidth는 실제 병목).
### Negative
- **모든** bench 골든이 이동 → 한 번의 재생성(D4); CI 골든 fixture 갱신.
- op 종류별 값은 calibration 필요; ~40 ns composite 수치는 추정이고
primitive는 미지정 — 결과는 숫자만큼만 정확하다(calibrate 전까지 절대
latency를 과대 주장 금지).
- 발행 경로에 비용 테이블 lookup 추가(런타임 영향 미미).
## Open review items (자율 결정; 검토 시 수정)
1. **op 종류별 값의 calibration 출처.** composite ≈ 40 ns는 작업 추정치;
primitive 발행 비용은 placeholder. *권고:* 문서화된 가정(instruction-issue +
큐 push)에서 값을 택하고, 실제 reference가 생기기 전까지 절대 latency를
잠정으로 취급; **비율**을 방어 가능하게 유지.
2. **scheduler plan-gen 비용 (D3).** *권고:* 초기에는 0 유지; scheduler-bound
동작을 보이는 workload가 나오면 `overhead_ns`로 노출.
3. **DMA program 시간 (A3).** *권고:* 먼저 발행 op의 비용에 합치고; 필요 시
PE_DMA setup으로 분리.
4. **테이블 위치.** *권고:* 흩어진 node attr이 아니라, 커맨드 종류로 키잉되고
per-topology 오버라이드 가능한 작은 중앙 cost 모듈 — 값을 한 곳에서 검토
가능하도록.
5. **회귀 롤아웃.** *권고:* lazy-load(ADR-0062)와 본 ADR의 cost model을 한
번의 골든 재생성 패스로 함께 도입하여 두 번의 churn을 피함.
6. **legacy replay 경로와의 상호작용**(`pe_cpu.py:_execute_legacy`) — greenlet과
replay 경로가 동일한 cost 테이블을 읽어 결과가 일치하도록 보장. 검증 요.
## Test Requirements
1. **composite는 한 번, primitive는 op마다 청구.** `N_tiles`에 걸쳐 하나의
`tl.composite`를 내는 커널은 PE_CPU에서 composite 발행 비용을 한 번 청구하고;
동등한 `N_tiles × ops` primitive 커널은 op당 비용을 `N_tiles × ops`번 청구.
PE_CPU busy time이 그에 따라 다름을 assert.
2. **DMA latency 불변 (Q2).** 고정된 전송에 대해 PE_DMA `drain_ns`는 ADR 이전
값과 동일 — 발행 비용은 DMA에 합쳐지지 않고 PE_CPU에 가산(중복 계상 없음).
3. **saturation 관측 가능.** per-op 발행 비용이 0이 아니면, many-tile primitive
sweep는 GEMM 엔진 idle(CPU-bound 발행)을 보이고 composite sweep는 바쁘게
유지 — ADR-0060 §1 레버.
4. **결정성:** 동일 입력 → 동일 op_log + latency (SPEC §0.1).
5. **경로 일치:** greenlet과 legacy-replay 경로가 동일 커널에 대해 동일한 발행
비용 회계를 산출.
docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model.md
-179
View File
@@ -1,179 +0,0 @@
# ADR-0064: Per-op-type CPU issue cost model (command construct + dispatch)
## Status
Proposed
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). The hybrid
> decision there (GEMMs via `tl.composite`, softmax merge in the kernel)
> wins by **offloading tiling to PE_SCHEDULER so the CPU issues coarse
> descriptors and runs ahead, keeping the engines saturated**. That win is
> currently **invisible in the simulator** because per-op CPU issue cost is
> zero. This ADR makes the issue cost real and **op-type-differentiated**,
> so composite-vs-primitive trade-offs (and CPU saturation) are measurable.
## Context
### What exists today
- Every `tl.*` op calls `_emit_dispatch_overhead()` before emitting its
command (`tl_context.py:123-125, 190, 227, 235, 612, …`), which emits
`PeCpuOverheadCmd(cycles=dispatch_cycles)` **only if** `dispatch_cycles
> 0`. So the issue cost is a **single uniform knob** applied identically
to `tl.load`, `tl.dot`, a MATH op, and `tl.composite`.
- That knob is hardcoded to **0** in both live execution paths
(`pe_cpu.py:101` greenlet runner, `:195` legacy replay). ⇒ issuing a
command — *constructing the descriptor and pushing it to the scheduler
queue* — currently costs **0 ns** on PE_CPU.
- `PeCpuOverheadCmd` is consumed as `yield env.timeout(cmd.cycles)` on
PE_CPU (`kernel_runner.py:131-132`) and on PE_SCHEDULER
(`pe_scheduler.py:97-100`); a manual `tl.cycles(n)` also exists
(`tl_context.py:695`).
### Two findings that frame this ADR (from ADR-0060 review)
- **Q1 — composite issue cost.** Constructing + pushing a `tl.composite`
is expected to cost on the order of **~40 ns** of CPU time (descriptor
build + queue push). Today it is **0**. The hook exists; only the value
(and its per-op-type differentiation) is missing.
- **Q2 — scheduler dispatch vs DMA latency.** `PE_SCHEDULER`'s composite
dispatch is **non-blocking**: `_dispatch_composite` generates the tile
plan and enqueues to the feeder, returning immediately
(`pe_scheduler.py:104-121`). The actual **DMA latency is charged on
PE_DMA** as each tile flows through (`drain_ns = compute_drain_ns(path,
nbytes)`, `pe_dma.py:89`), **not** lumped into the scheduler dispatch ⇒
**no double-counting**, and DMA stays on its modelled component
(SPEC §0.1). The scheduler's own plan-generation currently costs **0**
sim time.
### Why uniform-and-zero is wrong for the hybrid
The hybrid's whole argument is that **one** composite descriptor offloads
`N_tiles` worth of GEMM tiling, so the CPU issues `O(1)` coarse commands
instead of `O(N_tiles × ops/tile)` fine ones. With issue cost = 0 (and
uniform), the model cannot show:
- that the primitive path may **fail to saturate** the engines when the
CPU cannot push fast enough (the core ADR-0060 §1 claim), nor
- that a composite **costs more to construct** than a single primitive but
far less than the many primitives it replaces.
A single uniform `dispatch_cycles` cannot express this: `tl.load` ≪
`tl.composite` in real construct cost.
## Decision
### D1. Per-op-type issue-cost table on PE_CPU
Replace the single `dispatch_cycles` scalar with a **cost table keyed by
command type**, charged on PE_CPU at issue (before the command is
dispatched to the scheduler). The greenlet pays this time, which is
exactly what gates how fast the CPU can push work — the saturation lever.
Starting estimates (calibrate later — see review items):
| Issued op | CPU issue cost (construct + push) |
|---|---|
| `tl.composite` (GEMM/MATH descriptor) | ~40 ns (Q1 estimate) |
| `tl.load` / `tl.store` (DMA descriptor) | small (a few ns) |
| `tl.dot` / MATH / IPCQ send/recv | small (a few ns) |
The point is the **ratio**: a composite construct ≫ a primitive issue, but
≪ the sum of the primitive issues it replaces. Exact ns are configurable.
### D2. Execution latency stays on the engines (Q2 — no change)
DMA/GEMM/MATH execution latency remains charged on PE_DMA / PE_GEMM /
PE_MATH as today. The issue cost (D1) is **CPU-side only** and additive; it
does **not** touch `drain_ns`, so there is no double-count. This preserves
SPEC §0.1 (latency on modelled components).
### D3. Scheduler plan-generation cost — start at 0, revisit
`PE_SCHEDULER`'s tile-plan generation costs 0 sim time today. Keep that for
now (the dominant lever is CPU issue cost, D1); expose it later via the
existing `overhead_ns` node attr if a scheduler-side cost proves material.
Marked as a review item, not decided here.
### D4. Configurable values; goldens regenerate
The cost table is configurable (per-topology / node attrs, default to the
D1 estimates). Turning issue cost non-zero changes **every** bench's
latency, so golden latencies are **regenerated** once — a model-fidelity
improvement, same posture as ADR-0062's global lazy-load regression.
## Alternatives
### A1. Keep a single uniform `dispatch_cycles > 0`
Rejected: op construct costs differ by an order of magnitude (`tl.load` vs
`tl.composite`); a uniform value either over-charges primitives or
under-charges composites, and in both cases misrepresents the
composite-vs-primitive trade-off the hybrid depends on.
### A2. Charge the issue cost on PE_SCHEDULER instead of PE_CPU
Rejected: the saturation question is *"can the CPU push descriptors fast
enough to keep the engines busy?"* — that is a **PE_CPU** issue-bandwidth
property. Charging it on the scheduler would not model CPU back-pressure.
### A3. Model DMA program/setup time as a separate fixed per-descriptor cost
Real HW pays a DMA descriptor program cost distinct from transfer time.
Deferred: initially fold the descriptor-program cost into the **issuing
op's** D1 cost (a `tl.load` / composite that triggers DMA). Split it out to
a PE_DMA fixed setup only if calibration shows it matters. Review item.
## Consequences
### Positive
- The hybrid's CPU-offload / saturation win (ADR-0060 §1) becomes
**measurable**, not just structural.
- Composite-vs-primitive and tiling-granularity trade-offs are visible in
latency, enabling the eval narrative.
- More faithful to hardware (issue bandwidth is a real bottleneck).
### Negative
- **All** bench goldens shift → one-time regeneration (D4); CI golden
fixtures update.
- Per-op-type values need calibration; the ~40 ns composite figure is an
estimate, primitives are unspecified — results are only as good as the
numbers (do not over-claim absolute latencies until calibrated).
- Adds a cost-table lookup on the issue path (negligible runtime).
## Open review items (decided autonomously; revise on review)
1. **Calibration source for the per-op-type values.** Composite ≈ 40 ns is
a working estimate; primitive issue costs are placeholders. *Recommend:*
pick values from a documented assumption (instruction-issue + queue
push) and treat absolute latency as provisional until a real reference
exists; keep the **ratios** defensible.
2. **Scheduler plan-gen cost (D3).** *Recommend:* keep 0 initially; expose
via `overhead_ns` if a workload shows scheduler-bound behaviour.
3. **DMA program time (A3).** *Recommend:* fold into the issuing op's cost
first; split to PE_DMA setup only if needed.
4. **Where the table lives.** *Recommend:* a small central cost module
keyed by command type, overridable per topology — not scattered node
attrs — so the values are reviewable in one place.
5. **Regression rollout.** *Recommend:* land lazy-load (ADR-0062) and this
ADR's cost model in one golden-regeneration pass to avoid two churns.
6. **Interaction with the legacy replay path** (`pe_cpu.py:_execute_legacy`)
— ensure both the greenlet and replay paths read the same cost table so
results match. Verify.
## Test Requirements
1. **Composite charges once; primitives charge per-op.** A kernel issuing
one `tl.composite` over `N_tiles` charges one composite issue cost on
PE_CPU; the equivalent `N_tiles × ops` primitive kernel charges the
per-op cost `N_tiles × ops` times. Assert the PE_CPU busy time differs
accordingly.
2. **DMA latency unchanged (Q2).** For a fixed transfer, PE_DMA `drain_ns`
is identical to the pre-ADR value — the issue cost is additive on
PE_CPU, not folded into DMA (no double-count).
3. **Saturation is observable.** With non-zero per-op issue cost, a
many-tile primitive sweep shows GEMM-engine idle (CPU-bound issue)
whereas the composite sweep keeps it busy — the ADR-0060 §1 lever.
4. **Determinism:** identical inputs → identical op_log + latency
(SPEC §0.1).
5. **Path parity:** greenlet and legacy-replay paths produce identical
issue-cost accounting for the same kernel.
docs/adr-proposed/ADR-0070-algo-gqa-short-context-attention.md
+243
View File
@@ -0,0 +1,243 @@
# ADR-0070: GQA Short-Context Attention — Unified A1/A2/A4/B Mapping for Prefill and Decode
## Status
Proposed — short-context (single tile to a handful of KV tiles) GQA
attention benchmark covering **both prefill and decode** in **four
mapping modes** (A1/A2/A4/B) and **three composite tiers** (without
composite, GEMM-only composite, composite + softmax_merge fused).
Compute-only attention — weight DMA, ITL batching, and end-to-end
LLM driver scope are out of scope (separate ADRs).
## Context
A GQA layer at the LLaMA-3.1-70B headline shape (`h_q=64`, `h_kv=8`,
`d_head=128`, GQA group `G = h_q/h_kv = 8`) on a 4×4 cube SIP can be
mapped to AHBM in several ways that trade KV-cache memory, HBM read
volume per cube, and intra-cube IPCQ traffic against PE parallelism.
> **Bench-vs-headline shape.** The kernels and tests use a
> `d_head = 64` proxy (`tests/attention/test_gqa_short_context.py`
> `D_HEAD = 64`) so the per-tile working set fits comfortably in
> scratch across all four modes. The mapping decisions are
> `d_head`-independent — only the per-PE GEMM tile bytes scale.
The four candidate mappings (per *GQA-mapping-on-AHBM-short-context* note §1):
```
Mode kv_per_cube C group_size Reduce / broadcast topology
---- ----------- ------- ---------- ----------------------------------
A1 1 h_kv P (=8) row chain + col bridge (2×4 mesh)
A2 2 h_kv/2 P/2 (=4) row chain only
A4 4 h_kv/4 P/4 (=2) single intra_W / intra_E hop
B 8 1 1 no broadcast / no reduce (single PE)
```
Within each cube the 8 PEs split into `kv_per_cube` groups of
`group_size = P/kv_per_cube` PEs; each group owns one KV head.
The kernel must be able to compare these mappings at multiple context
lengths (8K..64K KV tokens) and across the three composite tiers so
the team can size the per-cube HBM / IPCQ / GEMM tradeoff empirically.
## Decision
### 1. Mapping (unified A1/A2/A4/B)
Two unified kernels select mode at launch via `kv_per_cube ∈ {1,2,4,8}`:
- **Prefill** (`gqa_attention_prefill_short_kernel`):
- Q-tile split — `T_q` rows floor-balanced across the group's
`group_size` PEs.
- FA2 head fusion — `G` Q heads fused into the M dimension of one
batched GEMM per PE per tile.
- **IPCQ KV broadcast** — group root (`pe_in_group == 0`) loads K/V
from HBM and IPCQ-broadcasts each tile to the rest of its group;
q-tiles are independent so no intra-group reduce.
- **Decode** (`gqa_attention_decode_short_kernel`):
- Sequence-shard — each PE owns `S_local = S_kv/group_size` tokens
of the group's KV head.
- FA2 head fusion — same as prefill.
- **IPCQ chain reduce** — per-PE partial `(m, ℓ, O)` chain-reduce
up to group root (PE 0), which normalizes and stores.
Geometry (2×4 mesh):
```
group_size = 8 : 2 row × 4 col (A1)
group_size = 4 : 1 row × 4 col (A2)
group_size = 2 : 1 row × 2 col (A4)
group_size = 1 : single PE (B)
```
### 2. Shard addressing (ADR-0011 D-VA1 contract)
Deploy is per-(sip, cube, pe) but `tl.load` receives a single global
VA per tensor. The kernel computes its own shard base offset from
`tl.program_id(axis=0)` (PE id) and `tl.program_id(axis=1)` (cube id):
```python
# Decode shape (pe=row_wise): K is split across PEs by sequence shard.
cube_K_base = cube_id * kv_per_cube * K_HEAD_BYTES
head_K_base = group_id_in_cube * K_HEAD_BYTES
pe_K_seq_offset = pe_in_group * n_tiles_per_pe * K_TILE_BYTES
k_shard_base = k_ptr + cube_K_base + head_K_base + pe_K_seq_offset
```
For **prefill** the K/V dp is `pe=replicate` (group root reads and
broadcasts) so the `pe_in_group` term drops out:
`k_head_shard_base = k_ptr + cube_K_base + head_K_base`.
Skipping `cube_id` collapses all cubes onto cube 0's HBM region
(observed pre-fix as an 11.5× per-cube DMA imbalance).
### 3. Three composite tiers
The same mapping is exercised against three GEMM/MATH dispatch styles
so the contribution of the composite API vs the recipe-driven fusion
can be isolated:
```
(1) without composite primitives only (tl.dot, tl.exp, …)
(2) with composite (GEMM-only) tl.composite(op="gemm")
(3) with composite + softmax_merge tl.composite(prologue=[softmax_merge],
op="gemm", out=O,
epilogue=[add])
```
File layout (`src/kernbench/benches/gqa_helpers/short_ctx/`):
```
_gqa_attention_{prefill,decode}_short.py (1)
_gqa_attention_{prefill,decode}_short_composite.py (2)
_gqa_attention_{prefill,decode}_short_composite_fused.py (3)
```
Tier (2) is **GEMM-only by definition** — no recipe-driven fusion,
so both prefill's and decode's P·V stay `tl.dot`. Recipe-driven
fusion is exactly what tier (3) adds.
Tier (3) on **multi-cube modes (A1/A2/A4) of prefill** relies on the
D4 supplement: any composite operand that is an IPCQ recv'd slot
(non-root PE's K_T in Q·Kᵀ, V in P·V) is pinned and read in place
instead of being DMA-streamed from HBM. Without the supplement every
recv slot fed to a composite — both Q·Kᵀ's `b=K_T` and P·V's `b=V`
— PageFaults on PA decode (PE scratch addresses overflow the 51-bit
PA range).
### 4. Caller contract — `_validate_config`
Each kernel module exposes a single `_validate_config(...)` helper
called by the bench wrapper before launch. The kernel itself is lean
(no inline `if` guards): the contract block is enforced caller-side,
sim cost zero, but catches every silent-shard-corruption foot-gun
(`kv_per_cube=3 → group_size=2` via integer division, non-integer
GQA group, mismatched cube count, partial-tile `S_kv`, …) before any
address arithmetic runs.
### 5. Tensor layouts (host-side, mode-invariant byte totals)
- **Q**: `(kv_per_cube·T_q, h_q·d_head/kv_per_cube)`
`dp=(cube=column_wise, pe=replicate)` over `C = h_kv/kv_per_cube`.
- **K**: `(h_kv·S_kv·d_head/TILE_S_KV, TILE_S_KV)` tile-major.
Decode: `dp=(cube=row_wise, pe=row_wise)`.
Prefill: `dp=(cube=row_wise, pe=replicate)` (broadcast model).
- **V**: `(h_kv·S_kv, d_head)` native. Same dp as K.
- **O**: same dp as Q.
Caller pre-scales Q by `1/√d_head`.
## Verification
### Smoke + regression (`tests/attention/test_gqa_short_context.py`)
- Prefill / decode smoke (all 4 modes, single-tile S_kv).
- Multi-tile coverage (all 4 modes, S_kv chosen so each PE owns
≥2 tiles).
- Op-count invariants — prefill `dma_writes = group_size·kv_per_cube·C`,
`dma_reads = P·C + 2·kv_per_cube·C·n_tiles`; decode
`dma_writes = kv_per_cube·C`.
- IPCQ topology — short-context kernel never emits inter-CUBE E/W
IPCQ.
- **ADR-0011 D-VA1 regression** — `per_cube_disjoint_src_addrs`,
`per_cube_dma_balanced` (max/min DMA busy ratio < 1.1×) per mode
per phase.
### Composite smoke
- `test_{prefill,decode}_composite_smoke[A1/A2/A4/B]` — tier (2),
4 modes × 2 phases.
- `test_{prefill,decode}_composite_fused_smoke[A1/A2/A4/B]` — tier
(3), 4 modes × 2 phases.
- `test_decode_composite_fused_multitile[A1/A2/A4/B]` —
`n_tiles_per_pe == 2` per mode so the recipe-fused tile loop
actually executes.
### Mode × context sweep (`tests/attention/test_gqa_short_context_sweep_*.py`)
Six sweep files (3 tiers × {prefill, decode}) measure each
`(mode, S_kv)` cell and dump the same metric set to CSV:
```
wall_us, n_pe, gemm_util, math_count, math_pipeline_us, hbm_bw_util,
hbm_read_mb, hbm_write_kb, ipcq_kb, kv_cache_per_cube_mb
```
S_kv ∈ {8K, 16K, 32K, 64K}; outputs under `docs/sweeps/`.
Composite/fused sweeps count both `gemm_f16` (non-pipeline) **and**
`TileToken/GEMM` (pipeline composite path); fused additionally counts
`TileToken/MATH`. Baseline sweeps count `gemm_f16` only.
## Known limitations
1. **Causal mask** — kernel is non-causal. Adding causal masking is
orthogonal to mapping and is scheduled as a separate ADR.
2. **f16 online-softmax accumulator** — `(m, ℓ, O)` are f16 throughout.
Long-context numerical drift will need an f32 accumulator before
correctness-grade use.
3. **End-to-end numeric validation** — current tests are op-count /
topology / per-cube-DMA invariants. No full-kernel
`np.allclose` against a torch/numpy reference yet.
4. **Skinny-M GEMM underfill** — decode (`T_q=1`) gives `M = G = 8`,
well below `PE_SCHEDULER` supertile `TILE_M = 32`. The two GEMM
paths handle this differently:
- **Primitive `tl.dot`** dispatches the GEMM at the actual
`m = 8` (no padding); GEMM time reflects the real 8-row work.
- **`tl.composite(op="gemm")`** tiles at `TILE_M = 32`, padding
`M` 4× with zeros; GEMM time reflects the padded 32-row work.
Consequence: when sweep CSVs show tier (2)/(3) `gemm_util`
higher than tier (1), most of that gap is supertile padding
overhead, **not** extra useful work or fusion savings — read it as
"composite tile shape doesn't match decode-skinny shape" rather
than "composite/fusion is more compute-intensive." This is the
**intended decode shape** for this benchmark; lifting it requires
batched-M (multi-request inference, `B_sys > 1`), tracked as a
separate batched variant — see "Future work".
5. **Multi-cube prefill composite without the ADR-0065 D4 supplement**
PageFaults. The supplement — IPCQ recv'd slots are pinned and read
in place as composite operands — ships alongside this ADR.
## Future work
- Batched-M variant (`B_sys ∈ {1,2,4,8,16}`) so composite/fused
pipeline overlap shows up in wall_clock rather than only in
per-engine utilization.
- Long-context (S_kv ≥ 128K) sweep extension for the headline LLaMA
decode target.
- Causal-mask + f32 accumulator promotion.
- 3-variant comparison plots from the sweep CSVs (wall, gemm util,
hbm bw util, ipcq, kv cache) — generated under `docs/diagrams/`.
## References
- ADR-0011 — PhysAddr + VA contract (D-VA1: kernel-computed shard
offset).
- ADR-0060 — GQA fused attention on AHBM (precursor of the unified
mapping).
- ADR-0064 — CPU issue-cost model (composite supertile motivation).
- ADR-0065 — Flat ops, composite, softmax_merge recipe (with the D4
supplement applied here).
- ADR-0070 supersedes ADR-0060 §B.split.2 short-prefill clause.
- *GQA-mapping-on-AHBM-short-context* note (research design rationale).
docs/adr-proposed/DDD-0060-gqa-fused-attention-detailed-design.md
+225 -402
View File
@@ -1,39 +1,42 @@
# Detailed Design Document — AHBM GQA Fused Attention (ADR-0060) # Detailed Design Document — AHBM GQA Fused Attention (ADR-0060)
**Status:** Draft (companion to ADR-0060, Proposed) **Status:** Draft (companion to ADR-0060, Proposed)
**Scope:** implementation-ready design for an efficient GQA FlashAttention **Scope:** implementation-ready plan for the **two** GQA FlashAttention
kernel on AHBM in the kernbench simulator. kernels (decode=reduce, prefill=ring) on AHBM in kernbench.
**Audience:** reviewer resuming with *"GQA 검토 시작하자"*. Read ADR-0060 **Audience:** reviewer resuming with *"GQA 검토 시작하자"*. Read **ADR-0060
first (decision record), then this (the how), then the *Open Decisions* first** — it is now the authoritative design record (TL;DR has full
section (§10) for the choices made autonomously that need your sign-off. pseudocode for both kernels). This DDD is the *how to build it*: file plan,
phase plan, helper signatures, tests. It does **not** re-derive the design;
it points to ADR-0060 sections.
> **How this document was produced.** It is the result of several > **Alignment note (this revision).** ADR-0060 moved to a **composite
> design↔evaluation iterations against the real kernbench source > hybrid + hierarchical CUBE-Group SP** design with **two kernels**. This
> (`src/kernbench/triton_emu/tl_context.py`, `sim_engine/*`, > DDD is rewritten to match. The earlier DDD (greenlet-primitive,
> `components/builtin/pe_*`, the `_attention_mesh_*` kernels, and the GQA > `tl.load_async`, single unified kernel) is superseded.
> tests). Every claim about what exists is grounded in a file path; every
> claim about what is new is tied to a supporting ADR (0061/0062/0063).
--- ---
## 1. Goal and success criteria ## 1. Goal and success criteria
**Goal:** an *efficient* GQA fused attention kernel that runs on AHBM in **Goal:** two *efficient* GQA fused-attention kernels on AHBM in kernbench —
kernbench, across the four operating points decode/prefill × **decode+SP** (head-replicated, KV static shard, 2-level reduce) and
single-/multi-user, at realistic (not just validation) scale, in both **prefill+SP** (1 Q head per CUBE, Ring KV) — at realistic scale, in both
timing (`enable_data=False`) and data (`enable_data=True`) modes. timing (`enable_data=False`) and data (`enable_data=True`) modes.
**Success criteria:** **Success criteria:**
1. Correct: kernel output matches a numpy FlashAttention reference within 1. Correct: kernel output matches a numpy FlashAttention reference within fp
fp tolerance, with **real GQA** (`H_q=64, H_kv=8, G=8`). tolerance, with **real GQA** (`H_q=64, H_kv=8, G=8`).
2. Efficient: the three levers actually fire and are observable in op_log: 2. Efficient — levers observable in op_log (ADR-0060 §11):
GQA K/V-load amortisation, KV prefetch overlap, `⌈log₂ N⌉` reduction. GQA K/V-load amortisation; decode 2-level reduce (`⌈log₂P⌉` + center-mesh
over `C`, not `C·P−1`); prefill ring (no reduce, KV rotates); load/compute
overlap (lazy `tl.load` + composite streaming); composite GEMM offload.
3. Scales: context length not capped by the 1 MiB scratch bump allocator. 3. Scales: context length not capped by the 1 MiB scratch bump allocator.
4. Deterministic and SPEC-compliant (all latency from modelled events). 4. Deterministic and SPEC-compliant (all latency from modelled events).
**Non-goals:** RoPE / QKV projection / KV-cache writes (upstream **Non-goals:** RoPE / QKV projection / KV-cache writes (upstream
`qkv_rope`); output projection (downstream); cross-CUBE query-head split; `qkv_rope`); output projection (downstream); cross-**SIP** head split
cycle-accurate microarchitecture (SPEC §5). (a head stays in one SIP, ADR-0060 §0); cycle-accurate microarchitecture
(SPEC §5); a bespoke "flash-composite" command kind (ADR-0060 §8 item 4).
--- ---
@@ -41,31 +44,33 @@ cycle-accurate microarchitecture (SPEC §5).
``` ```
runtime_api (host) runtime_api (host)
RuntimeContext.zeros/empty/launch context.py (tensor deploy, DPPolicy, kernel launch) RuntimeContext.zeros/empty/launch context.py (tensor deploy, DPPolicy(+cube_start), launch)
│ KernelLaunchMsg │ KernelLaunchMsg
▼ ▼
sim_engine sim_engine
GraphEngine(enable_data=…) engine.py GraphEngine(enable_data=…) engine.py
DataExecutor (Phase 2) data_executor.py ← ADR-0061 adds _execute_broadcast DataExecutor (Phase 2) data_executor.py ← _execute_broadcast (ADR-0061, optional)
MemoryStore (nbytes check) memory_store.py ← the GQA blocker lives here MemoryStore memory_store.py
│ │
▼ ▼
components / PE pipeline (per PE) components / PE pipeline (per PE)
PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,FETCH_STORE,TCM,IPCQ} PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,IPCQ,TCM}
greenlet kernel ↔ SimPy kernel_runner.py, pe_cpu.py greenlet kernel ↔ SimPy kernel_runner.py, pe_cpu.py
CompositeCmd → tile plan → engines pe_scheduler.py (ADR-0014 D6)
│ │
▼ ▼
tl programming model tl_context.py tl programming model tl_context.py
load/store/dot/softmax/max/sum/exp/maximum/send/recv/trans/… load (→ lazy, ADR-0062) / store / composite / dot / max/sum/exp/maximum
← ADR-0061 adds broadcast/repeat send/recv/recv_async / scratch_scope (ADR-0063) / broadcast (ADR-0061, opt)
← ADR-0062 adds load_async (+ wait arm)
← ADR-0063 adds scratch_scope
``` ```
The attention kernels are **bench-side Python** (`src/kernbench/benches/`) Both kernels are **bench-side Python** (`src/kernbench/benches/`) run as
invoked via `ctx.launch(name, kernel_fn, *args)`; they run as greenlet greenlet kernels emitting `tl.*` ops; the **GEMMs are `tl.composite`**
kernels that emit `tl.*` ops. No component (device-side) code changes are (scheduler-managed tiling + K/V DMA streaming, existing `CompositeCmd` — no
needed for the algorithm — only the three `tl`/sim_engine primitives. new command kind), the softmax merge + reduction/ring stay kernel-level
(ADR-0060 §1). The remote work already added building blocks:
`DPPolicy.cube_start`, `_attention_mesh_mlo_2d.py` (2D cube reduce —
currently AllReduce, to become reduce-to-root), `topologies/llama70b_4sip.yaml`.
--- ---
@@ -75,424 +80,241 @@ needed for the algorithm — only the three `tl`/sim_engine primitives.
| File | Change | ADR | | File | Change | ADR |
|---|---|---| |---|---|---|
| `src/kernbench/triton_emu/tl_context.py` | add `broadcast`/`repeat`; `load_async` + `wait` arm; `scratch_scope` ctx-mgr | 0061/0062/0063 | | `src/kernbench/triton_emu/tl_context.py` | `tl.load` → **lazy** (non-blocking + auto-wait on first use); `scratch_scope` ctx-mgr; `broadcast` (optional) | 0062/0063/0061 |
| `src/kernbench/common/pe_commands.py` | add `BroadcastCmd`; `DmaReadCmd.blocking` flag (or `DmaReadAsyncCmd`) | 0061/0062 | | `src/kernbench/components/builtin/pe_dma.py` | non-blocking read path: post DMA, resolve on a scheduled event (the `recv_async` pattern for loads) | 0062 |
| `src/kernbench/sim_engine/data_executor.py` | add `_execute_broadcast`; handle async dma_read resolution | 0061/0062 | | `src/kernbench/triton_emu/kernel_runner.py` | auto-wait: yield a handle's pending load event at first consume; lazy-load dispatch arm | 0062 |
| `src/kernbench/components/builtin/pe_dma.py` | non-blocking read path returns a future-resolving event | 0062 | | `src/kernbench/sim_engine/data_executor.py` | `_execute_broadcast` (optional) | 0061 |
| `src/kernbench/sim_engine/event_log.py` / op_log | recognise `broadcast` op_name (math) | 0061 | | `src/kernbench/components/builtin/pe_cpu.py` + cost table | per-op-type CPU issue cost (replaces uniform `dispatch_cycles`) | **0064** (separate) |
### 3.2 New kernel + bench files No new GEMM command kind: the two GEMMs use the **existing** `CompositeCmd`.
### 3.2 Kernel + bench files
| File | Role | | File | Role |
|---|---| |---|---|
| `src/kernbench/benches/_gqa_attention.py` | the unified GQA kernel (4 cases via flags), with a shared `_flash_tile_sweep` helper and `_tree_reduce` helper | | `src/kernbench/benches/_gqa_decode.py` | decode+SP kernel — evolve `_attention_mesh_mlo`/`_mlo_2d` to composite-hybrid + **2-level reduce-to-root** (Level-2 PE tree + Level-1 center-mesh; currently 2D AllReduce) |
| `src/kernbench/benches/_gqa_helpers.py` | `valid_len`, `tile_*` predicates, mask builders, tree topology (children/parent dirs), GQA-batched Q load | | `src/kernbench/benches/_gqa_prefill.py` | prefill+SP kernel — evolve `_attention_mesh_kv` to **head-parallel (1 Q head/CUBE) Ring KV** + composite-hybrid |
| extend `src/kernbench/benches/milestone_gqa_llama70b.py` | add headline-scale panels with real `G`, prefetch, tiling | | `src/kernbench/benches/_gqa_helpers.py` | `head_of_group`, `valid_len_2level`, `init_running`, tree/center-mesh topology, `hierarchical_reduce_and_store`, ring helpers, causal `block_*` predicates + mask builders |
| extend `src/kernbench/benches/milestone_gqa_llama70b.py` | headline panels: real `G`, `C=G=8`, contiguous KV, 4-SIP |
| `topologies/llama70b_4sip.yaml` | **exists** (remote) — 4 SIP × 16-CUBE 4×4 × 8 PE |
### 3.3 New tests ### 3.3 New tests
| File | Covers | | File | Covers |
|---|---| |---|---|
| `tests/attention/test_gqa_correctness.py` | numpy reference parity, all 4 cases, real GQA | | `tests/attention/test_gqa_decode.py` | decode reduce: GQA reuse (dma_read ⟂ `G`), 2-level reduce step count, root-only output, long context |
| `tests/attention/test_gqa_levers.py` | GQA load amortisation, tree step count, causal skip, prefetch overlap | | `tests/attention/test_gqa_prefill.py` | prefill ring: no `(m,ℓ,O)` reduce, `C` KV rotations, causal step-skip, per-CUBE distributed output |
| `tests/test_tl_broadcast.py` | ADR-0061 unit | | `tests/test_tl_lazy_load.py` | ADR-0062 unit (overlap, auto-wait correctness, op_log parity) |
| `tests/test_tl_load_async.py` | ADR-0062 unit |
| `tests/test_tl_scratch_scope.py` | ADR-0063 unit | | `tests/test_tl_scratch_scope.py` | ADR-0063 unit |
| `tests/test_tl_broadcast.py` | ADR-0061 unit (optional) |
| `tests/test_dppolicy_cube_start.py` | **exists** (remote) — sub-mesh placement |
--- ---
## 4. Data model and launch contract ## 4. Data model and launch contract
### 4.1 Tensor placement (DPPolicy) ### 4.1 Tensor placement — CUBE Group, 2-level SP, contiguous KV
Using `DPPolicy(cube, pe, num_cubes, num_pes)` (`policy/placement/dp.py`): A `CUBE Group` is `C` CUBEs in one SIP owning one KV head (ADR-0060 §0).
Placement via `DPPolicy(..., cube_start=...)`:
| Tensor | single-user (PE ring in 1 cube) | multi-user (cube ring) | | Tensor | Decode+SP | Prefill+SP (`C=G`) |
|---|---|---| |---|---|---|
| `Q` (group, replicated) | `cube=replicate, pe=replicate, num_pes=N` | `cube=replicate, pe=replicate, num_cubes=N` | | `Q` | `replicate` (all `G` heads, M-folded) over `C×P` ranks | **head-parallel**: CUBE `i` holds Q head `i` |
| `K`, `V` (sequence-sharded) | `cube=replicate, pe=row_wise, num_pes=N` | `cube=row_wise, pe=replicate, num_cubes=N` | | `K`,`V` | sequence-sharded **two axes** (`cube` row_wise over `C`, `pe` row_wise over `P`) | sequence-sharded over `C` CUBEs (the ring blocks) |
| `O` (group, replicated) | as `Q` | as `Q` | | `O` | reduced to CUBE-Group root (all `G` heads) | per-CUBE (one head each), distributed |
This matches the baseline `_PANEL_DISPATCH` **KV cache is *contiguous* `C×P` position blocks** (rank `r` = `[r·B,(r+1)·B)`),
(`milestone_gqa_llama70b.py:64-81`). `row_wise` over the sequence (M) axis **shared by both kernels** — prefill writes it (via `qkv_rope`), decode reads
realises the round-robin shard (§2.1 of ADR-0060); the contiguous per-PE + extends it, no reshard. Contiguous is required for prefill causal skip
buffer is the row-wise slice. (ADR-0060 §2.1). `cube_start` offsets the group's CUBE sub-mesh within the
4×4 SIP mesh.
`rank_axis` selects which `tl.program_id` carries the ring rank: `0` for **Driver SP-enable threshold (fallback).** Contiguous sharding under-utilizes
single-user (PE id in cube), `1` for multi-user (cube level; only PE 0 of ranks for short/early decode (only frontier ranks hold data). So the driver
each cube participates) — kept from the baseline so the SFR installs **enables SP only past a context-length threshold** where KV-sweep time
(`configure_sfr_intracube_pe_ring` / `configure_sfr_intercube_multisip`) dominates reduction + placement imbalance; below it, fall back to a smaller
apply unchanged. `C` (or `C=P=1`, single rank, no reduce). The threshold is a sweep item
(P7 / §9), not a fixed constant.
### 4.2 Launch signature `--device` enumerates **SIPs**; one SIP-device runs its 2 CUBE Groups
(2 KV heads); the head is picked by CUBE coordinate (`head_of_group`), not an
in-kernel loop (ADR-0060 §0).
### 4.2 Launch signatures
```python ```python
ctx.launch( # decode+SP
f"{panel}_gqa", gqa_attention_kernel, ctx.launch(f"{panel}_gqa_decode", gqa_decode_sp,
q, k, v, o, # Tensors (TensorArg) q, k, v, o, counter, start_pe, start_cube, C, P, softmax_scale)
S_q, S_kv_per_rank, H_q, H_kv, D_HEAD, n_ranks, # scalars # prefill+SP
rank_axis, mode, counter, start_pe, # scalars (mode: 0=decode,1=prefill) ctx.launch(f"{panel}_gqa_prefill", gqa_prefill_sp,
softmax_scale, # float scalar q, k, v, o, T_q, S_kv_local, d, C, softmax_scale, q_block, cube_start)
_auto_dim_remap=False, # keep d_head from colliding (baseline note)
)
``` ```
`H_q > H_kv` is the change from the baseline's `H_q=H_kv=1`; it requires Full per-kernel I/O contract: ADR-0060 §0.5.3/§0.5.4 (note the **output head
ADR-0061 to survive Phase 2. distribution differs** — decode root-gathered, prefill distributed).
### 4.3 GQA-batched Q load — fold `G` into the matmul M dim (no broadcast) ### 4.3 GQA reuse — fold `G` into the matmul M dim (decode)
For a KV head `kv`, its group is query heads `[kv·G, (kv+1)·G)`. Load Decode replicates Q and **M-folds** the `G` query heads into the GEMM row
those `G` rows and **reshape the group into the matmul M dimension** — dim so one `Q·Kᵀ` serves all `G` heads sharing one `K` (ADR-0060 §0, §5.2):
this is the key to GQA reuse and it needs only a byte-conserving reshape, `Q` group `[G, T_q, d]` → `[G·T_q, d]` (byte-conserving `_view`); the
**not** a broadcast: composite carries `m = G·T_q` (timing counts all `G` rows). Prefill does
**not** M-fold (1 head/CUBE). Caveat: `tl.trans` is reshape-not-transpose →
```python store `K` pre-transposed `[d, S/(C·P)]` (ADR-0060 §3, §B).
def load_q_group(q_ptr, kv, G, T_q, d, tl):
# Q laid out [H_q, T_q, d]; the group is contiguous in the head axis.
qg = tl.load(q_ptr + kv*G*T_q*d, shape=(G, T_q, d), dtype="f16")
return _view(qg, (G*T_q, d)) # byte-conserving reshape → M = G*T_q
```
Then per KV tile (`Kj`, `Vj` are `[TILE, d]`, one KV head, loaded once):
```python
Sj = tl.dot(q_g, tl.trans(Kj)) * scale # [G*T_q, d]·[d, TILE] -> [G*T_q, TILE]
Oj = tl.dot(P, Vj) # [G*T_q, TILE]·[TILE, d] -> [G*T_q, d]
```
K/V are the **shared operands** (the matmul's `b`), loaded once and reused
across all `G·T_q` M-rows automatically. The emitted `GemmCmd` carries
`m = G·T_q`, so Phase-1 timing counts all `G` rows' work correctly.
> **Why not a leading batch axis** (`[G,T_q,d]·[G,d,TILE]`)? `tl.dot`
> derives `m,k,n` from the **last two dims only** (`tl_context.py:220`)
> and `GemmCmd` carries no batch dim — so a leading `G` axis is computed
> correctly by `np.matmul` in Phase 2 (it broadcasts) but **under-counted
> by `G×` in Phase-1 timing**. Folding into M fixes both. This supersedes
> the earlier batched-vs-loop question.
> **Transpose caveat.** `tl.trans` is metadata-only and `MemoryStore.read`
> *reshapes* (not transposes) (`memory_store.py:73`,
> `data_executor.py:153-168` does `np.matmul` on the reshaped operands).
> Fine for timing/structure and for the baseline's zero data; for
> non-trivial numeric parity, store K pre-transposed as `[d, S_pe]` or add
> a real `tl.transpose` (§10.10).
--- ---
## 5. Kernel design (concrete, real `tl`) ## 5. Kernel design
### 5.1 Shared tile sweep helper **The full pseudocode for both kernels (and the 3 decode CPU-pipelining
variants) lives in ADR-0060 TL;DR + §5.** Implementation notes only here.
```python ### 5.1 Decode (reduce) — `_gqa_decode.py`
def flash_tile_sweep(q_g, k_base, v_base, n_tiles, TILE, d, scale, - Inner tile = §3 composite-hybrid: `Sj = composite(q_g, Kⱼ)·scale` → softmax
q_block, causal, tl, PREFETCH=2): MATH → `Oj = composite(P, Vⱼ)` → running merge (ADR-0060 §3).
"""One PE's FlashAttention sweep over its local KV tiles. - **Ship opt3** (software pipelining: issue next tile's `Q·Kᵀ` before this
Returns running (m, l, O) in the persistent arena.""" tile's softmax; `Sj` in a persistent double buffer) — removes the
m, l, O = _init_running(q_g.shape, d, tl) # persistent arena GEMM-engine bubble, no new command kind (ADR-0060 §5.6). opt1 is the naïve
f_k, f_v = {}, {} baseline; opt2 (`ex_composite`) waits on ADR-0064.
for j in range(min(PREFETCH, n_tiles)): # prime (ADR-0062) - Combine = **`hierarchical_reduce_and_store`** (ADR-0060 §4): Level-2 PE
if causal and tile_all_future(j, q_block): continue reduce-to-root tree (intra-CUBE, the cube's KV slice further split across
f_k[j] = tl.load_async(k_base + j*TILE*d, (TILE, d)) its `P` PEs — decision (a)) → Level-1 center-root CUBE-mesh reduce
f_v[j] = tl.load_async(v_base + j*TILE*d, (TILE, d)) (intra-CUBE-Group). Data-driven, level-pipelined, root-only.
for j in range(n_tiles):
if causal and tile_all_future(j, q_block): continue
with tl.scratch_scope(): # recycle temporaries (ADR-0063)
Kj = tl.wait(f_k[j])
Sj = tl.dot(q_g, tl.trans(Kj)) * scale
mask = make_mask_tile(j, q_block) if (causal and tile_partial(j, q_block)) else None
if mask is not None: Sj = Sj + mask
m_j = tl.max(Sj, axis=-1)
m_new = tl.maximum(m, m_j)
P = tl.exp(Sj - m_new)
corr = tl.exp(m - m_new)
l = l * corr + tl.sum(P, axis=-1)
Vj = tl.wait(f_v[j])
O = O * corr + tl.dot(P, Vj)
m = m_new
nxt = j + PREFETCH # refill pipeline
if nxt < n_tiles and not (causal and tile_all_future(nxt, q_block)):
f_k[nxt] = tl.load_async(k_base + nxt*TILE*d, (TILE, d))
f_v[nxt] = tl.load_async(v_base + nxt*TILE*d, (TILE, d))
return m, l, O
```
> Note: `m, l, O` are reassigned inside the `with` but their **storage** ### 5.2 Prefill (ring) — `_gqa_prefill.py`
> must be the persistent arena, not the recycled scope. Concretely the - Head-parallel: CUBE `i` owns Q head `i` + KV slice `i`. **No reduce** (each
> merge ops write to pre-allocated persistent handles (double-buffered); CUBE outputs a distinct head). Ring rotates KV blocks; GQA reuse via the
> ADR-0063 §D3 describes the persistent-vs-scoped arena split. This is the rotation (ADR-0060 §5.5).
> one subtlety the implementer must get right — see §10.5. - **K/V are `tl.load`'d TCM handles** (not `tl.ref`) so the ring can
`tl.send`/`recv_async` them; `K` pre-transposed `[d, S/C]`. Receive buffers
**ping-pong** (persistent arena) — recv into a *separate* buffer while
computing/sending the current one (do not clobber).
- Causal step-skip + boundary mask; `recv_async` overlaps next block's
receive with current compute.
- Within a CUBE: tile the Q head's `T_q` rows across `P` PEs (disjoint output
rows → no intra-CUBE reduce); fall back to KV-block split only if `T_q < P`
(ADR-0060 §B decode-split item 3).
### 5.2 Tree reduction helper ### 5.3 Shared helpers — `_gqa_helpers.py`
`head_of_group(cube_id) = cube_id // C`; `init_running(...)` → persistent
```python `(m=-inf, ℓ=0, O=0)`; `valid_len_2level(...)` (contiguous block extent);
def tree_reduce_and_store(m, l, O, pe_id, N, o_ptr, tl): `hierarchical_reduce_and_store(...)`; tree/center-mesh child/parent dirs;
for cdir in tree_children_dirs(pe_id, N): # static tree (§4 ADR-0060) `block_all_future`/`block_partial` + `causal_mask`.
m_c = tl.recv(cdir, m.shape); l_c = tl.recv(cdir, l.shape)
O_c = tl.recv(cdir, O.shape)
mn = tl.maximum(m, m_c); sa = tl.exp(m - mn); sb = tl.exp(m_c - mn)
l = l*sa + l_c*sb; O = O*sa + O_c*sb; m = mn
if is_root(pe_id, N):
tl.store(o_ptr, O / l)
else:
pdir = parent_dir(pe_id, N)
tl.send(pdir, m); tl.send(pdir, l); tl.send(pdir, O)
```
`tree_children_dirs` / `parent_dir` map the binary tree (level pairs in
ADR-0060 §4) onto physical `N/S/E/W` neighbours — this mapping is the
tuning item ADR-0060 §9.2 and an Open Decision (§10.2).
### 5.3 Top-level kernel (all four cases)
```python
def gqa_attention_kernel(q_ptr,k_ptr,v_ptr,o_ptr,
S_q,S_kv_per_rank,H_q,H_kv,d,n_ranks,
rank_axis,mode,counter,start_pe,scale,*,tl):
if rank_axis != 0 and tl.program_id(axis=0) != 0: # multi-user: only PE0 per cube
return
pe_id = tl.program_id(axis=rank_axis)
G = H_q // H_kv
causal = (mode == 1) # prefill
q_block = (counter, S_q) if causal else (counter, 1)
# decode small-S fast path = one tile; otherwise tiled (TILE from config)
for kv in range(H_kv): # one KV head per iter (or PE-mapped)
q_g = load_q_group(q_ptr, kv, G, S_q, d, tl)
my_len = valid_len(counter, start_pe, pe_id, n_ranks) if mode==0 else S_kv_per_rank
n_tiles = ceil(my_len / TILE)
m,l,O = flash_tile_sweep(q_g, k_base(kv), v_base(kv), n_tiles, TILE, d,
scale, q_block, causal, tl)
if n_ranks == 1:
tl.store(o_ptr_for(kv), O / l)
elif mode == 1 and is_ring(...): # prefill ring already folded across steps
ring_fold_then_store(...) # §5.4
else:
tree_reduce_and_store(m,l,O, pe_id, n_ranks, o_ptr_for(kv), tl)
```
This collapses the baseline's two separate kernels (`_attention_mesh_kv`,
`_attention_mesh_mlo`) into one parameterised kernel; the `mode`/`n_ranks`
flags select the case. (Keeping them separate is also fine — Open
Decision §10.6.)
### 5.4 Ring (prefill SP) specifics
The Ring path keeps the baseline's structure (`_attention_mesh_kv.py`)
with three upgrades:
- `recv_async` for the next step's K/V (overlap) — `tl_context.py:543`.
- causal step-skip: `if step_all_future(step, q_block): continue` before
folding (skips ≈half the steps).
- GQA: `q_g` is the `G`-row group; `tl.broadcast` for the KV head.
The fold across ring steps is the same online-softmax merge as §5.2's
inner update — running `(m,l,O)` carried across steps as Python handles
(exactly as the baseline does today).
--- ---
## 6. The three new primitives — integration points ## 6. Supporting primitives — integration points
### 6.1 ADR-0061 `tl.broadcast` (OPTIONAL — not on the GQA critical path) ### 6.1 ADR-0062 lazy `tl.load` (efficiency — overlap)
GQA reuse is achieved by folding `G` into M (§4.3), so broadcast is **not** - `tl.load` issues `DmaReadCmd` non-blocking, returns a handle with a pending
required to unblock `h_q > h_kv`. It remains a clean primitive for additive event; the runtime **auto-inserts the wait at first consume** (generalises
causal-mask construction across the `G·T_q` rows and for general kernels. the `recv_async`/wait pattern, `kernel_runner.py:248-285`).
- `tl_context.py`: new `broadcast(x, new_shape)` → emits `BroadcastCmd`, - `pe_dma.py`: non-blocking read; DMA occupies the (capacity-1) read channel,
output handle from `_make_compute_out` (so `nbytes` is correct). overlaps compute on PE_GEMM/PE_MATH.
- `data_executor.py`: `_execute_broadcast` = `np.broadcast_to(src, - **Global** semantics change → existing goldens regenerate (ADR-0062 D3).
shape).copy()` → `store.write`. - op_log unchanged (`memory/dma_read`) ⇒ `dma_read_count` stable.
- Latency: `pe_math._compute_ns(prod(new_shape))`; logs `math/broadcast`,
**no** `dma_read`.
- Lowest priority; deliver after the efficiency/scale features land.
### 6.2 ADR-0062 `tl.load_async` ### 6.2 ADR-0063 `tl.scratch_scope` (scale)
- `tl_context.py`: `load_async(ptr, shape, dtype) → LoadFuture`; extend - ctx-mgr saving/restoring `_scratch_cursor`. Persistent arena for `(m,ℓ,O)`
`wait` to resolve `LoadFuture`. (+ decode opt3 `Sj` double buffer, + prefill ring ping-pong buffers, +
- `pe_dma.py`: non-blocking read path; transfer occupies the read channel in-flight lazy-load buffers) lives **outside** the scope. Removes `S=16`.
in parallel; completion is a scheduled event.
- op_log unchanged (`memory/dma_read`) ⇒ `dma_read_count` metric stable.
### 6.3 ADR-0063 `tl.scratch_scope` ### 6.3 ADR-0061 `tl.broadcast` (optional)
- `tl_context.py`: context manager saving/restoring `_scratch_cursor`. - Not on the GQA critical path (M-fold gives reuse). Convenience for additive
- Persistent arena: `m,l,O` (+ double-buffer) and live `LoadFuture` mask construction. Lowest priority.
buffers allocated outside the scope.
- Removes the `S=16` cap (`test_milestone_gqa_llama70b.py:123-148`). ### 6.4 ADR-0064 per-op-type CPU issue cost (separate ADR)
- Replaces uniform `dispatch_cycles=0`. Makes the hybrid's CPU-saturation
win (and the decode opt2 `ex_composite` fewer-issues win) **measurable**.
Land before claiming hybrid latency wins (ADR-0060 §9).
--- ---
## 7. Phased implementation plan (test-first per CLAUDE.md) ## 7. Phased implementation plan (test-first per CLAUDE.md)
Each phase is an independent Phase-1→Phase-2 cycle. Ordered so each builds Each phase is an independent Phase-1→Phase-2 cycle; each keeps the baseline
on a green predecessor; every phase keeps the existing baseline green. green.
| Phase | Deliverable | Gate | | Phase | Deliverable | Gate |
|---|---|---| |---|---|---|
| **P1** | Real GQA by restructuring (per KV head, `G`→M fold); one-shot, no tiling. **No new primitive.** | data-mode run completes for `h_q=G·h_kv`; K/V `dma_read_count` independent of `G`; `gemm` m=`G·T_q` | | **P1** | Real GQA, decode, composite-hybrid M-fold; one-shot, no tiling. **No new primitive** (existing `CompositeCmd`). | data-mode completes `h_q=G·h_kv`; K/V `dma_read_count` ⟂ `G`; composite tile plan `m=G·T_q` |
| **P2** | Tree reduction replacing fan-out (decode SP) | `⌈log₂ N⌉` reduction rounds; structure preserved | | **P2** | Decode **2-level reduce-to-root** (Level-2 PE tree + Level-1 center-mesh), replacing the 2D AllReduce | reduce rounds = `⌈log₂P⌉`+center-mesh; **root-only** output; Level-1 on CUBE NOC, Level-2 on PE IPCQ |
| **P3** | ADR-0063 `tl.scratch_scope` + tiled sweep | long-`S` decode completes (today caps at S=16) | | **P3** | **ADR-0063** `scratch_scope` + tiled sweep | long-`S` decode completes (today caps S=16); scratch O(1) |
| **P4** | ADR-0062 `tl.load_async` + prefetch in sweep | latency < serial Σ(load+compute) | | **P4** | **ADR-0062** lazy `tl.load` + composite K/V streaming | tiled-sweep latency < serial `Σ(load+compute)`; goldens regenerated |
| **P5** | Causal masking + tile/step skip (prefill, ring) | GEMM count = triangular tile count | | **P5** | Decode **opt3** software pipelining | GEMM engine non-idle across tiles (structural / op_log) |
| **P6** | Headline-scale milestone panels (real `G`, tiling, prefetch) + figures | sweep.json headline rows; figures render | | **P6** | Prefill **head-parallel Ring KV** + causal step-skip | `C` KV rotations, no `(m,ℓ,O)` reduce, per-CUBE distributed `O`, causal skip ≈½ |
| **P7** *(opt)* | ADR-0061 `tl.broadcast` (mask convenience); optional `tl.transpose` for numeric parity | broadcast unit test; asymmetric numeric parity | | **P7** | Contiguous shared KV layout + 4-SIP topology + headline milestone panels (real `G`, `C=G`) | shared layout (no reshard); headline sweep.json rows |
| **P8** *(opt)* | **ADR-0064** cost model + decode **opt2** `ex_composite`; **ADR-0061** broadcast | cost-model goldens; opt2 fewer-issues measurable |
**P1 alone delivers "GQA actually runs" — and it needs no new feature**, P1 delivers "GQA runs" on the composite path. P3 is the key *scale* feature;
only the §4.3/§5 restructuring (the iteration-2 finding). P3 is the most P4 the overlap lever; P2/P6 the SP algorithms.
important *new* feature (removes the S=16 scale ceiling). P4 is the
efficiency lever. P2/P5 are algorithmic. P7 is optional polish.
--- ---
## 8. Verification plan (concrete) ## 8. Verification plan (concrete)
Grounded in SPEC R2/R5, ADR-0023/0025 (IPCQ), ADR-0046 (`tl`), Mirrors ADR-0060 §11; grounded in SPEC R2/R5, ADR-0023/0025, ADR-0046, ADR-0054.
ADR-0054 (eval bench). Mirrors ADR-0060 §11.
**Runs + numeric (Phase 2, `enable_data=True`):** **Runs + numeric (`enable_data=True`):**
- `test_gqa_correctness.py`: for `(G∈{1,8}, T_q∈{1,16}, S, N∈{1,8})` the - decode & prefill complete for `(G∈{1,8}, T_q∈{1,16}, S, C∈{1,8}, P∈{1,8})`
GQA kernel **completes** (no byte-conservation error) — the `G=8` cases without byte-conservation error (the `G=8` cases baseline cannot express).
that the baseline head-packing cannot express. This is gated on the
§4.3 restructuring, **not** on a new primitive.
- *Numeric parity (secondary):* `O ≈ numpy FlashAttention` for - *Numeric parity (secondary):* `O ≈ numpy FlashAttention` for
symmetric/identity inputs where reshape-as-transpose is exact; full symmetric/identity inputs (reshape-as-transpose exact); full asymmetric
asymmetric parity is gated on an optional `tl.transpose` (§10.10). parity gated on a real `tl.transpose` (deferred).
**Levers (op_log assertions):** **Levers (op_log):**
- GQA amortisation: K/V `dma_read_count` constant as `G` grows 1→8; - GQA amortisation: K/V `dma_read_count` ⟂ `G`; compute scales with `G`.
`gemm` work grows with `G`. - Decode reduce: rounds = `⌈log₂P⌉` + center-mesh (not `C·P−1`); result at
- Tree: decode-SP `ipcq_send`/`recv` counts = tree edges, reduction one rank; Level-1=CUBE NOC, Level-2=PE IPCQ.
rounds = `⌈log₂ N⌉` (not `N−1`). - Prefill ring: `C` KV rotations, **no** reduce, each CUBE writes a distinct
- Causal skip: prefill `gemm_count` = lower-triangular tile count. head.
- Prefetch overlap: tiled-sweep latency strictly < `Σ(load+compute)`. - Causal skip: GEMM/step count = lower-triangular.
- Scratch: peak `_scratch_cursor` bounded by one tile, independent of - Overlap: tiled-sweep latency < `Σ(load+compute)`.
`n_tiles`. - Scratch: peak cursor bounded by one tile.
**Invariants:** **Invariants:** determinism (identical op_log + latency); every routed
- Determinism: identical inputs → identical op_log + latency (SPEC §0.1). request latency > 0 (SPEC §0.1).
- Every routed request latency > 0 (SPEC §0.1).
--- ---
## 9. Performance model (what we expect to see) ## 9. Performance model (expected)
Rough per-PE decode cost (KV-load-bound), with prefetch overlap:
Per-rank decode (KV-load-bound), composite-streamed + overlapped:
``` ```
t_decode_pe ≈ max( DMA(K+V over S_pe tiles), compute(QKᵀ + PV over S_pe, ×G amortised) ) t_decode_rank ≈ max( DMA(K+V over S/(C·P) tiles), compute(QKᵀ+PV, ×G amortised) )
+ t_tree_reduce(⌈log₂ N⌉ hops, payload O[G,d]) + t_reduce( ⌈log₂P⌉ intra-CUBE + center-mesh over C inter-CUBE )
``` ```
Prefill (ring): `t ≈ C · max( recv(KV block), compute(QKᵀ+PV over block) )`
with causal step-skip removing ≈half the steps; no reduce term.
Levers and their expected signature in op_log / latency: Levers: GQA reuse → KV bytes `H_kv·S·d` not `H_q·S·d`; overlap → `max` not
- **GQA reuse** → KV bytes moved are `H_kv·S·d` not `H_q·S·d` (8× less), sum; decode reduce `⌈log₂P⌉+center-mesh` vs baseline `C·P−1`; prefill ring
with compute still `H_q`-scaled. Decode being KV-load-bound, this is the trades the reduce for KV rotation (good when `O` is big). The CPU-saturation
dominant decode win. win (composite offload) is *measurable* only with ADR-0064.
- **Prefetch** → DMA hidden behind compute ⇒ `t ≈ max(...)` not sum.
- **Tree** → reduction `⌈log₂ N⌉` (3 for N=8) vs baseline `N−1` (7).
- **Causal skip** (prefill) → ≈½ the tile/step work.
These are the quantities the milestone bench's op_log summary + a new
latency column should report (Open Decision §10.3 on adding latency to the
sweep JSON).
--- ---
## 10. Open Decisions / Assumptions made autonomously (REVIEW THESE) ## 10. Open items (status)
Choices I made to keep progressing without your input. Each is reversible; Most original "Open Decisions" are now **resolved** in ADR-0060; the live
none is yet implemented in production code (docs only). Marked review items live in **ADR-0060 §B** (three groups: hybrid pivot,
**[recommend]** = my default if you don't object. hierarchical CUBE-Group SP, decode/prefill split). Key resolved choices:
### 10.1 Greenlet `tl` model over composites **[recommend: greenlet]** | Was open | Now |
The biggest design call. ADR-0060 §1/§8 argues the greenlet path is |---|---|
equally efficient in kernbench's per-op latency model and far smaller to | greenlet vs composite | **composite hybrid** (GEMMs→composite, merge→kernel) |
build than a flash-composite. **Risk:** if you specifically want the | one kernel vs two | **two kernels** (decode reduce / prefill ring) |
*composite* abstraction exercised for the eval narrative (e.g. to study | GQA matmul shape | **M-fold** (decode); head-parallel (prefill) |
scheduler-managed tiling), we'd instead invest in a flash-composite | reduction topology | decode **2-level reduce-to-root**; prefill **no reduce (ring)** |
command kind. I judged that out of proportion to the goal. *Confirm the | `tl.load_async` | **lazy `tl.load`** (ADR-0062, redefined) |
greenlet direction.* | KV placement | **contiguous `C×P`**, shared prefill/decode (ADR-0060 §2.1) |
### 10.2 Tree reduction topology / neighbour mapping **[recommend: binary tree on N/S/E/W]** Still to confirm (ADR-0060 §B): DDD↔impl reconcile (`_attention_mesh_mlo_2d`
ADR-0060 §4 fixes a binary tree (N=8: root=7). The mapping of tree edges AllReduce → reduce-to-root); `head_of_group` sub-mesh partition; `C=G`
to physical mesh directions depends on the SFR install's neighbour coupling; short-context KV balance; ADR-0064 calibration. Pre-existing:
wiring. I assumed `configure_sfr_intracube_pe_ring` gives a ring on which ghost ADRs 0055–0059 (separate backfill); `bf16→f16` proxy.
the tree pairs are 1-hop neighbours. **Needs verification** against the
actual installed neighbour table; if pairs aren't 1-hop, either relabel
PEs or accept multi-hop sends. *Confirm or hand me the intended mapping.*
### 10.3 Keep the baseline all-to-all as an option? **[recommend: replace with tree]**
The baseline fan-out gives the answer on *every* rank; the tree gives it
on root only. Attention needs root-only, so tree wins. If any downstream
consumer needs O replicated (none found), we'd keep fan-out. *Assuming
root-only is fine.*
### 10.4 GQA matmul shape — RESOLVED in iteration 2 **[fold G into M]**
Earlier this asked "batched leading axis vs per-`g` loop." Resolved: do
**neither** — fold `G·T_q` into the matmul M dimension (§4.3). A leading
`G` axis would be computed by `np.matmul` in data mode but **under-counted
`G×` in Phase-1 timing** (`GemmCmd` carries no batch dim,
`tl_context.py:220`). Folding into M gives correct timing *and* correct
data with a single byte-conserving reshape and one shared K/V operand. No
new primitive, no per-`g` loop.
### 10.5 Persistent-vs-scoped arena split **[recommend: two explicit arenas]**
ADR-0063 requires `m,l,O` to live outside the recycled scope. The cleanest
implementation reserves a small fixed persistent region at kernel start
and double-buffers the merge. I assumed this is acceptable extra
bookkeeping inside the shared helper. *Alternative:* make `scratch_scope`
take an explicit "keep-alive" handle list. I picked the two-arena model
for clarity.
### 10.6 One unified kernel vs two (decode/prefill) **[recommend: unify]**
§5.3 unifies into one parameterised kernel; the baseline has two. Unifying
reduces duplication but makes the kernel branch-heavy. *Alternative:* keep
`_gqa_decode` and `_gqa_prefill` separate, sharing helpers. Low stakes;
either is fine.
### 10.7 Ghost ADRs 0055–0059 **[recommend: backfill as Draft, separate task]**
The baseline cites ADR-0055/0056/0057/0058/0059 which **do not exist**
(`grep` finds refs in code/tests, no files; no git history of them). This
is pre-existing documentation debt, not created by this work. I did **not**
backfill them (out of scope, and reverse-engineering intent is risky).
*Recommendation:* a separate `Draft` ADR-backfill task to document the
existing mesh kernels + milestone bench + SFR install, so ADR-0060 has a
real predecessor to "supersede". *Confirm you want this as follow-up.*
### 10.8 Numbering **[recommend: as chosen]**
GQA ADR = **0060**; supporting = **0061/0062/0063**. I avoided 0055–0059
(claimed by ghost refs) to prevent collisions. The file was renamed from
`ADR-XXXX-…` to `ADR-0060-algo-…` (title was erroneously "ADR-001").
### 10.9 bf16 modelled as f16 **[assumption, pre-existing]**
`memory_store.py:16` maps `bf16→float16` (numpy has no bf16). Numeric
parity tests must use tolerances that accommodate this, and we should not
claim bf16 *accuracy* results — only timing/structure. Pre-existing
simulator limitation; flagged so results aren't over-interpreted.
### 10.10 `tl.trans` is reshape-not-transpose **[decision: structural-first verify]**
Discovered in iteration 2: `tl.trans` only swaps shape metadata
(`tl_context.py:390`) and `MemoryStore.read` *reshapes* the bytes
(`memory_store.py:73`); `data_executor._execute_gemm` then `np.matmul`s
the reshaped operands (`data_executor.py:153-168`). So for **non-trivial
numeric data**, `Q·Kᵀ` via `tl.trans(K)` is mathematically a reshape, not
a transpose — wrong numbers (the baseline never notices: it runs on
all-zeros and only asserts op_log structure). **Decision:** treat the
simulator as a *performance* model (SPEC §0) — verify structure + timing +
determinism first; treat numeric parity as a bounded secondary check
(symmetric/identity inputs, or store K pre-transposed). A real
data-materialising `tl.transpose` is a *candidate* further primitive if
true numeric validation is ever required — **recommend deferring it**; it
is not needed for the performance-modeling goal. *Confirm you're content
with structural-first verification.*
### 10.11 GQA reuse = fold `G` into M, NOT a broadcast op **[finding, iteration 2]**
The headline mechanism finding. The deferred-GQA limitation
(`test_milestone_gqa_llama70b.py:137-142`) was blamed on a missing
broadcast / the MemoryStore nbytes check. Iteration 2 shows the real
cause is the baseline's head-packing reshape, and the fix is kernel
restructuring (§4.3), which needs **no new primitive**. This demoted
ADR-0061 from "the blocker" to "optional convenience" and reordered the
phase plan (§7). Recorded explicitly because it changes the project's
critical path: *GQA itself is not gated on any new feature; scale
(ADR-0063) and efficiency (ADR-0062) are the features that matter.*
--- ---
@@ -500,33 +322,34 @@ critical path: *GQA itself is not gated on any new feature; scale
| Risk | Likelihood | Mitigation | | Risk | Likelihood | Mitigation |
|---|---|---| |---|---|---|
| Numeric parity blocked by trans-as-reshape | med | structural-first verify (§10.10); pre-transpose K or defer `tl.transpose` | | Numeric parity blocked by trans-as-reshape | med | structural-first verify; pre-transpose K (ADR-0060 §B item 2) |
| Tree pairs not 1-hop on installed mesh | med | verify SFR neighbour table (§10.2); relabel PEs | | 2-level reduce pairs not 1-hop on the mesh | med | verify SFR neighbour table; center-root sub-mesh partition (§B) |
| scratch_scope use-after-scope bug | med | two-arena discipline in shared helper; debug poison | | scratch_scope use-after-scope (Sj / ring / lazy-load buffers) | med | persistent-arena discipline; ping-pong buffers (§5.2) |
| prefetch buffer recycled while in flight | med | live `LoadFuture` buffers in persistent arena (§10.5) | | prefill ring buffer clobbered while in flight | med | recv into the *other* ping-pong buffer (§5.2) |
| GQA broadcast `.copy()` memory blow-up | low | bounded by TILE in tiled design | | decode opt2 `acc` read before composites drain | med | `tl.wait()` after the loop (ADR-0060 TL;DR opt2) |
| Effort underestimated (6 phases) | med | P0–P1 alone deliver the core unblock; later phases independent | | impl drift (`_attention_mesh_mlo_2d` is AllReduce, Q-replicated) | med | reconcile to reduce-to-root; add prefill ring (§B) |
| CPU-saturation win invisible | exp | ADR-0064 cost model (P8) |
--- ---
## 12. Glossary & references ## 12. Glossary & references
- **GQA / G** — grouped-query attention; `G = H_q/H_kv` query heads share - **GQA / G** — `G = H_q/H_kv` query heads share one KV head. Llama3-70B: `G=8`.
one KV head. Llama3-70B: `G=8`. - **CUBE Group** — `C` CUBEs in one SIP owning one KV head (ADR-0060 §0).
- **FlashAttention / FlashDecoding / Ring Attention** — see ADR-0060 - **2-level SP** — Level-1 inter-CUBE (over `C`) × Level-2 intra-CUBE PE
lineage note. (over `P`); ranks = `C·P`.
- **IPCQ** — inter-PE queues; `tl.send`/`tl.recv` over `N/S/E/W` - **Composite hybrid** — GEMMs via `tl.composite` (scheduler), softmax merge
(ADR-0023/0025/0032). + reduction in the kernel (ADR-0060 §1).
- **Greenlet `tl` model** — kernel interleaves with SimPy; ops emit - **M-fold** — stack the `G` query heads into the GEMM M (row) dim so one
commands scheduled on PE engines (ADR-0020, ADR-0046). `Q·Kᵀ` does all heads sharing one `K`.
- **Persistent vs scoped arena** — ADR-0063 scratch split. - **FlashAttention / FlashDecoding / Ring Attention** — ADR-0060 lineage.
Key source anchors: `tl_context.py` (tl API), `memory_store.py:67-73` Source anchors: `tl_context.py` (tl API), `pe_scheduler.py:104-143`
(nbytes check), `data_executor.py` (Phase 2), `pe_dma.py`/`pe_gemm.py`/ (composite tile plan), `pe_dma.py:45,89` (read channel, drain), `pe_cpu.py`
`pe_math.py` (latency), `_attention_mesh_kv.py` / `_attention_mesh_mlo.py` (dispatch_cycles), `lrab_hierarchical_allreduce.py` (center-root pattern),
(baseline kernels), `milestone_gqa_llama70b.py` (eval bench), `_attention_mesh_mlo_2d.py` / `_attention_mesh_kv.py` (impl to evolve),
`tests/attention/test_milestone_gqa_llama70b.py` (current limits). `DPPolicy.cube_start`, `topologies/llama70b_4sip.yaml`.
ADRs: **0060** (this design), **0061** (broadcast), **0062** (async load), ADRs: **0060** (this design), **0062** (lazy load), **0063** (scratch scope),
**0063** (scratch scope); related accepted **0014/0020/0023/0025/0042/ **0061** (broadcast, optional), **0064** (CPU issue cost model); related
0045/0046/0052/0054**. accepted **0014/0017/0020/0023/0025/0046/0054**.
docs/adr/ADR-0014-dev-pe-pipeline-execution-model.md
+9 -2
View File
@@ -184,8 +184,15 @@ The empty-`ops` form is the legacy single-op path.
- `DMA_READ`: `simpy.Resource(capacity=1)`. - `DMA_READ`: `simpy.Resource(capacity=1)`.
- `DMA_WRITE`: `simpy.Resource(capacity=1)`. - `DMA_WRITE`: `simpy.Resource(capacity=1)`.
- Both channels run concurrently (READ ∥ WRITE allowed). - Both channels run concurrently (READ ∥ WRITE allowed).
- Within a channel, requests serialize (READ ∥ READ disallowed; same - Within a channel, requests serialize **at the issue path** (READ ∥ READ
for WRITE). issued out-of-order is disallowed; same for WRITE). The channel is
held only until the request is enqueued onto the next hop (router),
then released — it does **not** block until the HBM round-trip
completes. Multiple in-flight requests are therefore allowed, and
HBM-level serialization is the responsibility of the HBM controller's
per-PC `available_at` timestamps (ADR-0033 D1). This is what allows
back-to-back small-tile DMAs to amortize the per-request head latency
through the fabric.
- `vc_comm` is an orthogonal channel for IPCQ traffic defined in - `vc_comm` is an orthogonal channel for IPCQ traffic defined in
ADR-0023 D8 — out of scope for this ADR. ADR-0023 D8 — out of scope for this ADR.
docs/adr/ADR-0026-par-dppolicy-intra-device.md
+51 -2
View File
@@ -2,7 +2,7 @@
## Status ## Status
Accepted (Revision 5 — Phase 2 landed 2026-04-14, 523 passed + 1 strict xfail) Accepted (Revision 6 — cube_start added 2026-06-04; Revision 5 landed 2026-04-14)
## Context ## Context
@@ -31,9 +31,11 @@ class DPPolicy:
pe: Literal["replicate", "column_wise", "row_wise"] = "replicate" pe: Literal["replicate", "column_wise", "row_wise"] = "replicate"
num_pes: int | None = None num_pes: int | None = None
num_cubes: int | None = None num_cubes: int | None = None
cube_start: int = 0 # Revision 6 — first cube in SIP; see D8
``` ```
Removed fields: `sip`, `num_sips`. Removed fields: `sip`, `num_sips`.
Added fields (Revision 6): `cube_start` — see D8.
### D2. `ShardSpec` — structural (sip, cube, pe) coordinates, `pe_index` fully removed ### D2. `ShardSpec` — structural (sip, cube, pe) coordinates, `pe_index` fully removed
@@ -120,7 +122,7 @@ def resolve_dp_policy(
for ls in local_shards: for ls in local_shards:
all_shards.append(ShardSpec( all_shards.append(ShardSpec(
sip=target_sip, # from caller (current_device) sip=target_sip, # from caller (current_device)
cube=cube_id, # local within SIP cube=policy.cube_start + cube_id, # Rev 6: shifted by D8 cube_start; default 0 = local within SIP
pe=ls.local_pe, # local within cube (explicit name) pe=ls.local_pe, # local within cube (explicit name)
offset_bytes=cube_offset + ls.offset_bytes, offset_bytes=cube_offset + ls.offset_bytes,
nbytes=ls.nbytes, nbytes=ls.nbytes,
@@ -239,6 +241,53 @@ property: code that expected a global flat could otherwise silently
receive a SIP-local result and index incorrectly — that possibility is receive a SIP-local result and index incorrectly — that possibility is
eliminated. eliminated.
### D8. Add `cube_start: int = 0` for disjoint cube sub-meshes within a SIP
**Revision 6 addition (2026-06-04)**.
**Purpose**: address a disjoint cube sub-mesh within one SIP (e.g.
cubes 8..15 alongside cubes 0..7). Required by the GQA Llama-70B
8-KV-group headline target — two 2×4 KV-groups per SIP × 4 SIPs = 64
cubes — where the second KV-group on each SIP must land on cubes 8..15
instead of the default 0..7.
**Semantics**: `resolve_dp_policy` returns `ShardSpec` with
`cube = policy.cube_start + cube_id`, where `cube_id` iterates
`0..num_cubes-1` within the launch. Selected cubes lie in
`[cube_start, cube_start + num_cubes)` within the target SIP.
**Default**: `cube_start = 0` preserves every existing call site
bit-for-bit (`ShardSpec.cube ∈ [0, num_cubes)` as before). The CCL
milestone bench's full-SIP `DPPolicy(num_cubes=16)` continues to
produce cubes 0..15.
**Constraint**: intra-device invariant preserved.
`cube_start ∈ [0, cubes_per_sip)` and
`cube_start + num_cubes ≤ cubes_per_sip`. SIP boundary crossing remains
the job of `ahbm.set_device(rank)` (ADR-0024).
**Why scalar** (not 2D `cube_mesh_origin` or arbitrary `cube_ids` list):
consumer kernels (e.g. `_attention_mesh_mlo_2d`) assume row-major
contiguous cubes; scalar `cube_start` pairs naturally with `num_cubes`
(range = `[start, start + count)`); zero migration churn at default 0.
More general designs can be added on top later if a non-contiguous use
case appears.
**Backward compatibility**: additive change with default value. No
existing call site is forced to change. The "breaking change" stance
in D7 applies only to the original sip/num_sips removal — `cube_start`
does NOT break existing callers.
**Kernel-side note**: kernbench's `tl.program_id(axis=1)` returns the
physical cube id (ADR-0022), not a launch-local rank. Kernels that
derive ring positions from `program_id(axis=1)` must subtract
`cube_start` to recover launch-local rank when `cube_start > 0` —
otherwise they compute out-of-bounds sub-mesh positions and deadlock.
See `_attention_mesh_mlo_2d.py` for the reference pattern. A future
revision of ADR-0022 could move `program_id(axis=1)` to launch-local
rank semantics; that change would let kernels drop the explicit
subtraction.
## Dependencies ## Dependencies
- **ADR-0024** (launcher): `set_device(rank)` and current-device scoping - **ADR-0024** (launcher): `set_device(rank)` and current-device scoping
docs/adr/ADR-0060-algo-gqa-fused-attention-ahbm.md
+1299
View File
File diff suppressed because it is too large Load Diff
docs/adr-proposed/ADR-0062-prog-tl-async-load.md → docs/adr/ADR-0062-prog-tl-async-load.md
+1 -1
View File
@@ -2,7 +2,7 @@
## Status ## Status
Proposed Accepted
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Decode and > Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Decode and
> long-context attention are **KV-load-bound**; load/compute overlap is a > long-context attention are **KV-load-bound**; load/compute overlap is a
docs/adr-proposed/ADR-0063-prog-tl-scratch-scope.md → docs/adr/ADR-0063-prog-tl-scratch-scope.md
+63 -1
View File
@@ -2,7 +2,7 @@
## Status ## Status
Proposed Accepted
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Long-context > Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Long-context
> attention sweeps many K/V tiles; each tile's intermediates > attention sweeps many K/V tiles; each tile's intermediates
@@ -102,6 +102,68 @@ kernel keeps two arenas:
This mirrors how real flash-attention SRAM budgeting works: a small This mirrors how real flash-attention SRAM budgeting works: a small
persistent accumulator region + a recycled tile working set. persistent accumulator region + a recycled tile working set.
#### D3.1 The persistent-arena write mechanism: `tl.copy_to(dst, src)`
The merge ops (`tl.maximum`, `tl.exp`, binary `*` / `+`) all call
`_make_compute_out(...)` which allocates from the bump cursor (D1).
Inside a `scratch_scope`, their result handles therefore live **inside**
the scope and vanish on `__exit__`. To realise D3's two-arena split, the
kernel needs a way to **write a scoped result's bytes back to a
persistent address**.
The primitive that closes this gap:
```python
def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None:
"""Copy ``src``'s bytes into ``dst``'s address (both TCM).
Shapes and dtypes must match. ``dst`` is typically a handle
allocated outside any active ``scratch_scope`` (the persistent
arena); ``src`` is a scoped handle whose bytes must outlive
scope ``__exit__``.
"""
```
Symmetric to `tl.store` (the HBM-side byte copy), kept TCM-only here so
the running-state writeback doesn't pollute op_log with spurious DMA.
**Mechanics:**
- New `CopyCmd(src, dst, nbytes, data_op=True)` command.
- op_log: `op_kind="math"`, `op_name="copy"` — runs on the vector engine.
- Latency: `pe_math._compute_ns(prod(shape))` — models on-chip register
writeback, not HBM transfer.
- Emit-time validation: `dst.shape == src.shape`, `dst.dtype == src.dtype`,
`dst.space == "tcm"`, `src.space == "tcm"`. Authoring errors surface in
Phase 1, not deep in Phase 2 data execution.
**Call-site pattern:**
```python
m, l, O = init_running(...) # persistent (outside scope)
for j in range(n_tiles):
with tl.scratch_scope():
... # per-tile work (recycled)
m_new = tl.maximum(m, mj) # scoped scratch
l_new = l * scale_old + l_step * scale_step
O_new = O * scale_old + O_step * scale_step
tl.copy_to(m, m_new) # ← persist new running state
tl.copy_to(l, l_new)
tl.copy_to(O, O_new)
# exit: scoped m_new/l_new/O_new gone; their bytes live in persistent m/l/O
```
The copy happens **before** `__exit__`, so the read of `src` (scoped) is
valid; after exit only `dst` (persistent) is read, satisfying D2's
no-read-after-exit safety contract.
**Why a dedicated primitive rather than `dst=` kwargs on every math op**
(considered, rejected): adding `dst=` to `tl.maximum`, `tl.exp`,
`_binary_math`, `_unary_math`, `_reduction` is ~25 LOC across 5
op-families and breaks the uniform "call returns a fresh handle"
pattern. `tl.copy_to` is one primitive, one command, one executor
handler — minimal surface area for the same effect.
### D4. Nesting ### D4. Nesting
Scopes nest (stack of save-points). Inner scope exit rewinds to the inner Scopes nest (stack of save-points). Inner scope exit rewinds to the inner
docs/adr/ADR-0064-perf-cpu-issue-cost-model.md
+471
View File
@@ -0,0 +1,471 @@
# ADR-0064: Structural CPU dispatch cost model (`logical_bytes` + FIXED + R)
## Status
Accepted (Revision 2)
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention) and **ADR-0065**
> (flat-ops composite + first stateful recipe). The hybrid decision there
> (GEMMs via `tl.composite`, softmax merge in the kernel) wins by
> **offloading tiling to PE_SCHEDULER so the CPU issues coarse descriptors
> and runs ahead, keeping the engines saturated**. That win is currently
> **invisible in the simulator** because per-op CPU issue cost is zero.
>
> Revision 2 replaces the **op-type calibration table** (the original
> proposal) with a **structural formula** derived from each command's
> `logical_bytes` — no per-op-type calibration needed; new op kinds are
> covered automatically.
## Context
### What exists today
- Every `tl.*` op calls `_emit_dispatch_overhead()` before emitting its
command (`tl_context.py:196-212`), which emits
`PeCpuOverheadCmd(cycles=dispatch_cycles)` **only if** `dispatch_cycles
> 0`. The knob is **uniform** across op kinds and hardcoded to **0**
in both live execution paths (`pe_cpu.py:101` greenlet, `:195` replay).
- ⇒ issuing a command — *constructing the descriptor and pushing it to
the scheduler queue* — currently costs **0 ns** on PE_CPU.
- `PeCpuOverheadCmd` is consumed as `yield env.timeout(cmd.cycles)` on
PE_CPU (`kernel_runner.py:131-132`).
### Why uniform-and-zero is wrong for the hybrid
ADR-0060 §1's argument is that **one** composite descriptor offloads
`N_tiles` worth of GEMM tiling, so the CPU issues `O(1)` coarse commands
instead of `O(N_tiles × ops/tile)` fine ones. With issue cost = 0, the
model cannot show:
- that the primitive path may **fail to saturate** the engines when the
CPU cannot push fast enough, nor
- that a composite **costs more to construct** than a single primitive
but far less than the many primitives it replaces.
### Why per-op-type calibration (Revision 1) was over-shaped
The original proposal had a `cost_table[kind]` keyed by op kind
(`composite`, `load`, `dot`, `math`, …). That required:
- a value per kind (calibration cost ≥ |kinds|),
- a new entry every time a new kind appears,
- and yet the *ratio* it tried to capture — "composite ≫ primitive,
but ≪ the primitives it replaces" — is structurally a function of
**how many fields the command carries**, not of the op kind.
A command's *byte footprint* is the natural proxy: a composite carrying
N OpSpecs has ~N× the bytes of a primitive op with one OpSpec. The
*fixed* part (queue head update, completion register, MMIO-class
latency) is per-command. The two together compose: `FIXED + bytes × R`.
### What this model actually exposes
The **primary** signal is **command-count reduction** through the
per-command FIXED cost. The **byte** term is a secondary refinement
that prevents pathologically-large composites from looking free. With
realistic on-die queue bandwidth (16 B/cycle, D3), FIXED accounts for
≥85% of the dispatch cost differential between opt3 (≈10 cmds/tile)
and opt2 (≈2 cmds/tile) for decode opt2.
This framing also bounds the model: if a single composite were allowed
to grow unboundedly large, the byte term alone would not stop the
formula from rewarding ever-bigger fused commands beyond what real HW
supports — hence the descriptor-size cap in **D7**.
## Decision
### D1. Structural dispatch cost formula
Each PE command going to PE_SCHEDULER incurs PE_CPU dispatch cycles:
```
dispatch_cycles(cmd) = FIXED_PER_CMD + cmd.logical_bytes × R
```
where:
- `FIXED_PER_CMD` (cycles per command) models queue-tail update, MMIO-class
RTT, completion-event registration — fixed per command regardless of size.
- `R` (cycles per byte) models the queue-write bandwidth — bytes of the
command serialized into the scheduler queue.
- `cmd.logical_bytes` (int) is each command's *HW-logical* byte size,
computed from D2 below — not Python's `sys.getsizeof`.
PE_CPU emits `PeCpuOverheadCmd(cycles=dispatch_cycles(cmd))` before
dispatching, exactly as the existing hook (`tl_context.py:_emit_dispatch_
overhead`) — only the cycle value changes.
### D2. `logical_bytes` rule
Each PE command dataclass exposes `logical_bytes: int` (property). The
counting rule (HW-friendly, ignores Python overhead):
| Field kind | Bytes |
|---|---|
| Command framing (cmd-type discriminator + completion id ref) | 4 |
| Opcode (op kind enum) | 1 |
| Enum (scope, etc.) | 1 |
| `TensorHandle` reference (address only — shape/dtype assumed in descriptor table) | 8 |
| Scalar (int/float) | 4 |
| Tuple length marker | 1 |
`CompositeCmd` recursively sums its `ops` and `rw_handles`:
```python
@property
def logical_bytes(self) -> int:
return (
4 # framing
+ 1 + sum(op.logical_bytes for op in self.ops)
+ 1 + 8 * len(self.rw_handles)
)
```
`OpSpec`:
```python
@property
def logical_bytes(self) -> int:
return (
1 + 1 # opcode + scope
+ 1 + 8 * len(self.operands) # named operand handles
+ (8 if self.out is not None else 0) # out handle
+ 1 + sum(_extra_bytes(v) for v in self.extra.values())
)
def _extra_bytes(v) -> int:
"""Type-aware byte count for OpSpec.extra values."""
if isinstance(v, bool): return 1
if isinstance(v, (int, float)): return 4
if isinstance(v, (tuple, list)): return 1 + 4 * len(v) # shape, axes, …
if isinstance(v, str): return 1 # opcode-like tag
return 4 # default scalar
# Example types in extra:
# m, k, n int → 4 each
# reduce_axis int → 4
# shape=(64, 64) tuple → 1 + 8 = 9
# factor=1.0 float → 4
```
(Identical rule for `DmaReadCmd`, `MathCmd`, etc. — one property per
dataclass, ~3 lines each.)
**Counting rule — per-op summation, no deduplication.** Each
`TensorHandle` reference inside an `OpSpec` is counted independently,
even when the same handle appears in multiple OpSpecs or in
`rw_handles`. The `rw_handles` block is metadata for cross-composite
hazard tracking (ADR-0065 D6.3) and is counted separately from
operand references in `ops`. There is no deduplication. This matches
HW reality: the descriptor encodes each operand slot as an independent
address field, and the dispatcher tracks `rw_handles` as a distinct
metadata block. Example: `OpSpec(kind="mul_bcast", operands={"src_a":
O, "src_b": corr}, out=O)` counts handle `O` *twice* (once for
`src_a`, once for `out`); if `O` is also in the enclosing
CompositeCmd's `rw_handles`, it is counted a third time.
### D3. Defaults — anchored on a typical composite ≈ 43 ns
Anchor: a **single-OpSpec composite for a DMA-staged GEMM path** —
one `OpSpec(kind="gemm", ...)`; DMA stages are auto-inserted by
PE_SCHEDULER from operand `space` (ADR-0065 D4) and **do not appear in
`logical_bytes`** (the kernel does not issue them as separate cmds).
Breakdown:
```
framing 4
ops tuple length 1
GEMM OpSpec 40 (opcode 1 + scope 1 + 1 + 2 handles 16
+ out 8 + 1 + extra m/k/n 12)
rw_handles tuple length 1
rw_handles content 8 (one RW handle for the output)
─────────────────────────────
total ~54 bytes
```
Target dispatch = ~43 ns. On-die producer→consumer queue at 16 bytes/cycle
(typical on-die descriptor queue width).
```
FIXED_PER_CMD = 40 cycles
R = 0.0625 cycles/byte (= 16 bytes/cycle)
```
Verification: `40 + 54 × 0.0625 = 43.375 cycles ≈ 43 ns` ✓
Cycle→ns conversion uses the PE node's existing `clock_freq_ghz` attr
(the same one used by PE_MATH `_compute_ns`). The cost-model knobs are
**cycle-domain only** — they do not duplicate the clock setting.
### D4. Topology config override
Defaults are baked into `pe_cpu.py`. Topology yaml may override under a
`pe_cost_model:` section at the PE node attrs (cycle-domain knobs only;
clock comes from the PE's existing `clock_freq_ghz`):
```yaml
pe:
attrs:
clock_freq_ghz: 1.0 # existing, used for cycle→ns
pe_cost_model:
fixed_per_cmd_cycles: 40
byte_cycles_recip: 0.0625 # = 16 bytes/cycle
max_composite_logical_bytes: 1024 # D7 — descriptor size cap
```
Missing keys fall back to defaults. The dispatch formula reads from
`node.attrs["pe_cost_model"]` at PE_CPU init.
### D5. Scope — what does and does not pay
| Path | Pays dispatch cost? |
|---|---|
| PE_CPU → PE_SCHEDULER for any `PeCommand` | **Yes** |
| `PeCpuOverheadCmd` itself (already cycles-explicit) | **No** (formula bypass) |
| Stages auto-generated by PE_SCHEDULER (DMA_READ/WRITE/FETCH/STORE) | **No** (PE_SCHEDULER-internal) |
| Engine compute latency (DMA `drain_ns`, GEMM/MATH `_compute_ns`) | **No change** — stays on engines (SPEC §0.1) |
This preserves the "latency on modelled components" invariant — dispatch
cost is *additional* CPU-side time, not folded into engine times.
### D6. Configurable values; goldens regenerate
Turning issue cost non-zero changes **every** bench's latency. Golden
latencies are **regenerated once** when this ADR lands — same posture as
ADR-0062 D3 lazy-load. After regeneration, the same calibration is in
effect for ADR-0065 opt2 measurement.
### D7. Composite size cap (hard limit — validation error)
Each `CompositeCmd`'s `logical_bytes` is capped at
**`max_composite_logical_bytes`** (default **1024 bytes**, overridable
per D4). A composite whose `logical_bytes` exceeds the cap is **rejected
at emit time** by the host-side TLContext with a `ValueError` — there is
**no automatic segmentation**. The kernel author must restructure the
recipe (e.g., split it into multiple smaller composites explicitly) so
each `CompositeCmd` fits within the descriptor capacity.
**Why a cap is needed.** Real hardware imposes hard limits — descriptor
queue entry size, scheduler parser buffer, command SRAM, firmware
input. Without a cap, the `FIXED + bytes × R` model would reward
arbitrarily large fused composites beyond what hardware accepts (e.g.,
fusing 100 primitive ops into one composite, paying one `FIXED`).
**Why 1024 bytes specifically.** This is a **safe engineering limit**,
not a measured HW number — intentionally chosen to be well above all
currently known composites (decode opt2's `#2` at ~322 bytes is the
largest in the kernbench codebase) while still representing a *finite*
descriptor capacity that future recipes must respect. The number is
overridable per topology (D4); when a real HW reference appears, the
value should be recalibrated. The role of this default is to make the
cap *exist as a discipline*, not to fit a specific HW.
**Rationale for a hard error over auto-segmentation.** Auto-splitting an
oversized composite (the original Revision 2 proposal) added emitter
complexity — inter-segment ordering, shared `rw_handles`, completion
chaining — to paper over what is, in practice, a kernel that asked for a
descriptor larger than the hardware allows. Surfacing it as an explicit
error keeps the emitter simple and makes the HW constraint visible to the
kernel author, who is best placed to decide how to split the work.
Decode opt2's `#2` composite (10 ops, ~322 bytes) sits comfortably
inside the 1024 cap — no error for the GQA workload.
### D8. Single-op-cmd fast-path (per-cmd-type FIXED)
Every **single-op** command — the DMA descriptors (`DmaReadCmd`,
`DmaWriteCmd`) *and* the single-op compute dispatches (`GemmCmd`,
`MathCmd`, `CopyCmd`) — carries a separate, lighter FIXED than the
general control-path cost: **`fixed_per_single_op_cmd_cycles`** (default
**8 cycles**) replaces `fixed_per_cmd_cycles` for all of them.
**`CompositeCmd` is the only command that stays on the general
control-path cost** `fixed_per_cmd_cycles` (= 40 cycles).
Concretely, the D1 formula becomes
```
dispatch_cycles(cmd) = FIXED(type(cmd)) + cmd.logical_bytes × R
FIXED(CompositeCmd) = fixed_per_cmd_cycles (= 40)
FIXED(everything else, i.e. every single-op cmd:
DmaReadCmd / DmaWriteCmd / GemmCmd / MathCmd / CopyCmd)
= fixed_per_single_op_cmd_cycles (= 8)
```
**Why a separate, lighter FIXED for single-op cmds.** D1 treated every
emitted command uniformly because that is the simplest defensible
model. In practice an single-op command — whether a DMA descriptor or a
single GEMM / elementwise / copy dispatch to an engine — is *much*
cheaper than the generic control-path cost the 40-cycle FIXED was
modeling. It is a single descriptor or instruction pushed to an engine
queue, with no scheduler-side plan to build:
- NVIDIA Hopper TMA: a single PTX instruction (`cp.async.bulk`) initiates
a bulk async DMA — ~1 ISA cycle on the SM side, with the TMA engine
doing fan-out and chunking itself.
- NVIDIA Ampere `cp.async`: ~1–2 cycles per warp-level async copy.
- A single MMA / tensor-core issue (`wgmma`, `mma.sync`): one
instruction to the math pipe, not a control-path round trip.
- AMD AQL packet write to a HSA queue: ~5–15 cycles.
- Generic descriptor-ring-push designs (Synopsys/Xilinx-style DMA IP):
~5–20 cycles for the MMIO writes.
- Tenstorrent tile descriptor emission, Habana Gaudi TPC descriptor RAM
+ start register: single-digit cycles per descriptor.
The 40-cycle FIXED is justified for `CompositeCmd` specifically, because
a composite is *not* a single dispatch: the scheduler must build a
tile-feeder plan, track per-tile read/write hazards across its internal
`DMA_READ → FETCH → GEMM → STORE → DMA_WRITE` stages, and wire a
completion handle. An single-op `GemmCmd` is one engine issue with none of
that machinery; charging it the same 40 cycles as a composite conflates
a single instruction with a whole scheduled plan.
**Why this matters for kernel evaluation.** A uniform 40-cycle FIXED
inflates the dispatch overhead of user-orchestrated single-op-primitive
kernels. A chunked-prefetching async GEMM at `K = N_K · TILE_K` emits,
per K-tile, a `DmaReadCmd` (B-chunk load) **and** a `GemmCmd` (the
per-chunk `tl.dot`) **and** a `MathCmd` (the running `tl.add`
accumulate) — so the per-command FIXED dominates its dispatch budget.
The single-op fast-path keeps the *structural* command-count signal that
D1 was designed to capture (a recipe that emits `N` single-op commands
*does* pay more than a recipe that emits one composite covering `N`
tiles internally) without conflating it with a modeling overcharge that
would be specific to a 40-cycle generic control path.
**Why 8 cycles specifically.** 8 sits at the mid-range of the survey
above (TMA ~1 to descriptor-ring ~15). It is intentionally a *modeling
default*, not a measured HW number — the role of this default is to
distinguish the single-op dispatch path from the composite control path in
qualitative behaviour. The value is overridable per topology (D4).
**Effect on composite numbers.** Composite tile-internal stages
(`DMA_READ`, `FETCH`, `GEMM`, `STORE`, `DMA_WRITE` per HW tile) are
emitted by the scheduler's tile-feeder loop, not by the host-side
TLContext, so they do **not** go through `_charge_dispatch`. The single
`CompositeCmd` itself still pays the 40-cycle control-path FIXED. D8
only changes the FIXED charged to user-side single-op primitives
(`tl.load` / `tl.store` / `tl.dot` / `tl.add` / …). Composite
measurements remain stable.
## Alternatives
### A1. Keep Revision 1's op-type calibration table
Rejected: calibration cost scales with |kinds|, and the *ratio* the
table tried to capture is structurally a function of cmd size. The
structural formula reaches the same qualitative behaviour with two
calibratable numbers instead of N.
### A2. Byte-only formula (no FIXED term)
Rejected. With FIXED = 0, opt2 (Option Y per ADR-0065) does **not** win
over opt3 — the total *bytes* dispatched per tile are similar (opt3 ≈ 232,
opt2 ≈ 380); the win is entirely in *fewer per-cmd fixed costs*. A
byte-only formula erases the very signal the model needs to expose.
### A3. Charge dispatch on PE_SCHEDULER instead of PE_CPU
Rejected: the saturation question is *"can the CPU push descriptors fast
enough to keep the engines busy?"* — that is a **PE_CPU** issue-bandwidth
property. Charging on the scheduler would not model CPU back-pressure.
### A4. Model DMA program/setup time as a separate fixed per-descriptor cost
Deferred: initially fold the descriptor-program cost into the **issuing
op's** dispatch cost. Split it out to a PE_DMA fixed setup only if
calibration shows it matters.
## Consequences
### Positive
- Hybrid's CPU-offload / saturation win (ADR-0060 §1) becomes
**measurable**, with a structurally honest model (no calibration table).
- Adding new op kinds (e.g., ADR-0065's `softmax_merge` 8-step recipe)
costs zero — they fit the same formula automatically.
- More faithful to hardware (queue-head MMIO RTT + queue-write bandwidth).
### Negative
- **All** bench goldens shift → one-time regeneration (D6); CI golden
fixtures update.
- Two calibration knobs (FIXED, R) need values; defaults are anchored on
a documented assumption — treat absolute latencies as provisional until
a reference exists; keep the **ratios** defensible.
- Adds a small `logical_bytes` property to each PE command dataclass.
## Open review items
1. **Calibration source for FIXED and R.** Defaults from "typical
composite ≈ 43 ns + on-die queue 16 bytes/cycle"; reasonable for an
on-die descriptor queue. Revisit when a HW reference appears.
2. **Scheduler plan-gen cost vs large composites.** Stays 0 — D5 keeps
PE_SCHEDULER's plan-generation outside the dispatch formula. The
D7 cap (1024 bytes ≈ 30–35 OpSpecs) bounds the worst case, but a
composite near the cap still costs the scheduler the same as a 1-op
composite under the current zero-cost model. If a stress test
(large-composite microbench) shows scheduler-bound behaviour,
expose `overhead_ns` per-op-count.
3. **Where the override lives.** `pe_cost_model:` block under PE node
attrs in topology yaml — keeps all knobs in one place, reviewable.
Clock comes from the PE's existing `clock_freq_ghz` attr, not
duplicated here.
4. **Path parity.** Both greenlet (`_execute_legacy` and `kernel_runner`)
and replay paths must read the same cost model. Verify.
5. **Sensitivity of conclusions to R.** opt2 < opt3 must hold across a
reasonable range of `R` (queue-bandwidth assumption). Sensitivity
sweep is part of Test Requirements (#9).
## Test Requirements
Tests are written against the **formula** (D1), not against specific
numeric anchors, so they remain valid when calibration changes or when
OpSpec/CompositeCmd fields are added.
1. **Formula preservation.** For any `CompositeCmd` `c`, PE_CPU's
recorded dispatch overhead equals `FIXED_PER_CMD + c.logical_bytes
× R` (within ±1 cycle for floor/round-off). Parametrized over
several composites: a 1-OpSpec GEMM composite, a 5-op MATH chain,
and a 10-op recipe composite. Default-calibration numbers (anchor
≈43 ns for the 1-OpSpec composite) are informative reference, not
the test gate — the test gate is the formula equality.
2. **Qualitative ratio (robust).** opt3 per-tile PE_CPU dispatch
strictly exceeds opt2 per-tile dispatch by at least a 2× margin —
`opt3 > 2 × opt2`. The default-calibration model predicts ≈4×;
the gate is the loose 2× bound so the test does not break when
calibration is moved (e.g., when a HW reference replaces the
default). Informative numbers — see ADR-0065 §verification and
DDD-0065 §11 for the model expectation.
3. **Override path.** Topology yaml `pe_cost_model:` block changes the
per-PE dispatch cost; default is recovered when block is missing.
The formula identity from #1 must hold with the override values.
4. **`PeCpuOverheadCmd` bypass.** Manual `tl.cycles(n)` issues exactly `n`
cycles, not `n + dispatch_cycles(...)`.
5. **No double-count.** PE_DMA `drain_ns`, PE_GEMM/MATH `_compute_ns`
identical to pre-ADR values.
6. **Determinism.** Identical inputs → identical op_log + latency
(SPEC §0.1).
7. **Path parity.** Greenlet and replay paths produce identical
dispatch-cycle accounting for the same kernel.
8. **Composite size cap (D7).** A composite that would emit `logical_bytes
> max_composite_logical_bytes` raises a `ValueError` at emit time (no
segmentation). A composite within the cap emits normally.
9. **Sensitivity (qualitative).** At `R ∈ {0.25, 0.0625, 0.03125}`
cycles/byte, `opt3 > opt2` at *all* three points. Direction (ratio
monotonically increases as R decreases) is also asserted, but
absolute ratio values are *not* required.
## Migration
ADR-0064 Revision 2 lands as a single PR with:
- `logical_bytes` property on each `PeCommand` dataclass (type-aware
extra-field counting per D2)
- formula application in `pe_cpu.py` dispatch path
- `pe_cost_model:` override read at PE_CPU init (cycle-domain knobs +
`max_composite_logical_bytes`)
- **composite size cap (D7)** — TLContext-side hard cap with
`max_composite_logical_bytes` default 1024: a composite over the cap
raises a `ValueError` at emit time (no auto-segmentation). Not
triggered by any existing bench (largest current composite is ~322
bytes), but the check lands so the descriptor-capacity limit is
enforced as recipes grow.
- one-time goldens regeneration
After this lands, ADR-0065 builds on top with no further goldens churn
in existing benches (ADR-0065 is a meaning-preserving refactor of
`CompositeCmd` for the existing path; only opt2 is a new bench).
docs/adr/ADR-0065-prog-flat-ops-composite-softmax-merge.md
+488
View File
@@ -0,0 +1,488 @@
# ADR-0065: Flat-ops `CompositeCmd` + first stateful recipe (`softmax_merge`)
## Status
Accepted
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). Implements
> the §5.6 / §8 item 4 carve-out exactly: opt2 of decode = (#1 existing
> GEMM composite for Q·Kᵀ) + (#2 a single composite carrying softmax +
> P·V + online-softmax merge). The "ex_composite" / "flash_pv_merge" name
> referenced in ADR-0060 §5.6 is *retired* — this ADR fixes the canonical
> shape using the existing `tl.composite` entry with two structural
> additions: (a) a flat `ops` tuple, and (b) a first stateful recipe
> `softmax_merge` that lowers to a sequence of MATH micro-ops.
## Context
### What ADR-0060 left to this ADR
ADR-0060 §5.6 shipped opt3 (software pipelining) as the recommended now,
and explicitly carved out opt2 — fewer per-tile dispatches, K-before-V
DMA priority, **measurable only after ADR-0064**'s cost model — as the
follow-up.
§8 item 4 sized the carve-out: "if revisited, split it into **two**
composites — `#1` = Q·Kᵀ (existing composite + `scale`), `#2` = softmax +
P·V + online-softmax accumulator merge". This ADR implements `#2`.
### Why composite needs a structural change
Today's `CompositeCmd` has the shape `(op, a, b, out_addr, ops=(head,
*epi))` with implicit conventions:
- `op` selects the head engine (gemm | math)
- `ops[0]` is the head OpSpec, `ops[1:]` are epilogue OpSpecs that run
*after* the head per-output-tile or per-K-tile (scope-determined)
For decode opt2 we need MATH ops to run **before** the head GEMM (the
online-softmax `m, l, O` update + emit `P` for the GEMM input). The
current shape has no slot for that.
Two failed paths considered earlier:
- Embedding softmax_merge as an *epilogue* of #1 (Q·Kᵀ) — wrong: it
needs to read `Sj` (= #1's output) and writes a fresh `P` consumed by
the *next* GEMM (#2's P·V); circular.
- A single mixed-engine recipe `flash_pv_merge` absorbing MATH + GEMM +
MATH — violates the engine boundary, complicates PE_SCHEDULER.
The clean shape is: **drop the head/epilogue distinction from the
command format**. Make `ops` a flat ordered tuple; let each op's
`scope` + position determine its tile-loop placement. The "prologue"
concept stays in the *user-facing* `tl.composite(...)` API as
ergonomics — but the **HW command never sees it**.
### Why a recipe (not a generic op list)
Decode opt2's `#2` MATH chain is 8 steps with a fixed structure (online
softmax merge). Letting the kernel author wire those 8 steps would be
both verbose and a hazard surface (correctness depends on step order).
A *recipe* — a single named op (`softmax_merge`) that TLContext expands
to the 8 steps with addresses pre-filled — keeps the kernel concise and
the structure under one source of truth. The recipe table lives in
TLContext (the compiler analog); PE_SCHEDULER never reads it. See
boundary in D5.
## Decision
### D1. `CompositeCmd` is a flat ordered op list
```python
@dataclass(frozen=True)
class CompositeCmd:
completion: CompletionHandle
ops: tuple[OpSpec, ...] # ordered MATH/GEMM ops
rw_handles: tuple[TensorHandle, ...] = () # cross-composite hazard
data_op: bool = True
@property
def logical_bytes(self) -> int: ... # ADR-0064 D2
```
The previous `(op, a, b, out_addr, out_nbytes, math_op)` fields are
absorbed into `ops`. The first OpSpec in `ops` (if `kind == "gemm"`) is
the head GEMM; preceding OpSpecs are pre-loop ops, following OpSpecs
are post-loop / per-tile epilogue ops depending on `scope`.
### D2. `OpSpec` carries named operands
```python
@dataclass(frozen=True)
class OpSpec:
kind: str # "gemm" | "rmax" | ...
scope: Scope # KERNEL | K_TILE | OUTPUT_TILE
operands: dict[str, TensorHandle] # named — see D5
extra: dict[str, Any] # scalars, axes, m/k/n, etc.
out: TensorHandle | None = None # explicit write-back handle
@property
def logical_bytes(self) -> int: ...
```
Named operands let PE_SCHEDULER route inputs to engine ports (e.g., GEMM
takes `operands["a"]`, `operands["b"]`) without positional convention.
### D3. Position + scope determines tile-loop placement
**Semantics split.** *Position* determines the **phase** (where in the
tile-loop lifecycle the op fires); *scope* determines the **repetition**
within that phase. `Scope.KERNEL` means "**once per `CompositeCmd`
invocation**", not "once per kernel launch" — its single firing happens
at the op's position in the sequence. The same KERNEL value carries
different *phases* depending on whether the OpSpec sits before, at, or
after the GEMM op.
PE_SCHEDULER scans `cmd.ops` for the GEMM op (≤1 per CompositeCmd, see
D6). Calling its index `g`:
| OpSpec position | scope | Phase | Placement in plan |
|---|---|---|---|
| `0 .. g-1` | KERNEL | pre-tile-loop | run once before tile loop |
| `g` | KERNEL | head GEMM | drives tile loop with `extra["m"], extra["k"], extra["n"]` |
| `g+1 .. ` | K_TILE | in-accumulation | per K-tile epilogue (inside K accumulation loop) |
| `g+1 .. ` | OUTPUT_TILE | per-output-tile | per (m, n) tile epilogue (after K accumulation) |
| `g+1 .. ` | KERNEL | post-tile-loop | run once after tile loop |
A composite with no GEMM op (e.g., MATH-only composite) treats all ops
as KERNEL-scope sequential — phase collapses to "single-shot in order".
### D4. PE_SCHEDULER auto-inserts DMAs from the operand `pinned` flag
PE_SCHEDULER scans the GEMM's `operands`. For each:
- **not `pinned`** (data lives in HBM, must be streamed) → emit DMA_READ
Stage before the FETCH/GEMM Stages
- **`pinned`** (already staged in TCM via a prior `tl.load`, or a recipe's
TCM scratch / primary-out `P`) → no DMA Stage, consumed in place
The kernel never emits explicit DMA cmds for composite operands;
PE_SCHEDULER infers them from handles. `tl.ref` operands are not pinned
(DMA_READ); `tl.load` results and recipe scratch / primary-out are pinned
(in place).
> **As-built note (P3).** An earlier draft of D4 derived the DMA decision
> from a `space` tag (`hbm`→DMA, `tcm`→in place) and would have flipped the
> current `space` values (today `tl.load`=`hbm`, `tl.ref`=`tcm`). The
> implementation **keeps the existing `pinned` flag** as the DMA signal
> instead — `pinned` already encodes "already in TCM, skip DMA", flipping
> `space` would change existing bench Stage sequences for no functional
> gain. Unifying `space` and retiring `pinned` is a deferred low-priority
> cleanup.
**Prologue recipe ops are TCM-only (Phase 1 limit, enforced at TLContext
emit).** Every operand of a *prologue recipe* (non-GEMM) OpSpec **must** be
TCM-resident; passing an HBM handle as a recipe operand raises a validation
error at emit time. This bounds DMA staging to the head op's operands in
Phase 1. The **head op itself** (gemm *or* math) keeps the existing
DMA-staged-from-HBM behavior — D4's TCM-only rule does **not** apply to it.
This invariant is also restated in D6 #7.
### D5. RECIPE_DESCRIPTORS — TLContext-internal, NOT in `pe_commands.py`
Recipes live in a new TLContext-adjacent module
(`src/kernbench/triton_emu/tl_recipes.py`). PE_SCHEDULER does **not**
import it. The first recipe:
```python
@dataclass(frozen=True)
class RecipeDescriptor:
# Operand contract: name → "R" (read-only) | "RW" (read-write)
operands: dict[str, str]
# Type / shape rule for the implicit primary output (auto-bound to
# head's first matrix operand). None means recipe emits no value.
primary_out: PrimaryOutSpec | None
# Single-shot prologue (full reduction) vs tile-aligned execution.
tile_alignment: Literal["single_shot", "tile_aligned"]
# Internal scratch budget (bytes), function of operand dims.
internal_scratch_bytes_fn: Callable[..., int]
# Engine micro-op sequence: TLContext expands this into flat OpSpecs.
engine_seq: tuple[EngineOp, ...]
RECIPE_DESCRIPTORS["softmax_merge"] = RecipeDescriptor(
operands={"s": "R", "m": "RW", "l": "RW", "O": "RW"},
primary_out=PrimaryOutSpec(
from_shape="s", from_dtype="s", transform="identity",
),
tile_alignment="single_shot",
internal_scratch_bytes_fn=lambda G, TILE, d, bpe: bpe * (
G + G + G + G * TILE + G # m_loc + m_new + corr + P + l_loc
),
engine_seq=(
EngineOp("MATH", "rmax", src="s", dst="m_loc", reduce_axis=-1),
EngineOp("MATH", "max_elem", src_a="m", src_b="m_loc", dst="m_new"),
EngineOp("MATH", "exp_diff", src_a="m", src_b="m_new", dst="corr"),
EngineOp("MATH", "exp_diff", src_a="s", src_b="m_new", dst="P", bcast_axis=0),
EngineOp("MATH", "rsum", src="P", dst="l_loc", reduce_axis=-1),
EngineOp("MATH", "fma", src_a="l", src_b="corr", src_c="l_loc", dst="l"),
EngineOp("MATH", "mul_bcast", src_a="O", src_b="corr", dst="O", bcast_axis=1),
EngineOp("MATH", "copy", src="m_new", dst="m"),
),
)
```
TLContext, on `tl.composite(prologue=[{"op": "softmax_merge", ...}],
op="gemm", ...)`:
1. Validates operand R/RW against `RECIPE_DESCRIPTORS["softmax_merge"]`.
2. Allocates scratch (m_loc, m_new, corr, P, l_loc) via the scratch_scope
helper (ADR-0063).
3. Derives the primary output P's shape from `s.shape` (identity).
4. Expands `engine_seq` into 8 flat MATH OpSpecs with all addresses /
sizes filled. Each OpSpec gets `scope=KERNEL`.
5. **Auto-binds (with conflict check, D6.6).** If the head GEMM does
*not* explicitly provide `operands["a"]`, auto-bind
`operands["a"] = P_handle`. If the kernel *did* provide `a`
explicitly, emit a validation error — ambiguous primary-output
replacement.
6. Emits `CompositeCmd(ops=(8 MATH + 1 GEMM + epilogue), rw_handles=
(m, l, O))`.
**RECIPE_DESCRIPTORS appears nowhere in the HW path**. PE_SCHEDULER
sees only the flat ops list.
### D6. Invariants
1. **GEMM count.** `0 ≤ count(op.kind == "gemm" in cmd.ops) ≤ 1`. A
composite with 0 GEMMs is MATH-only.
2. **softmax_merge has no V operand.** `RECIPE_DESCRIPTORS["softmax_merge"
].operands` does not include any HBM-resident handle → no DMA inserted
during the prologue → V DMA can only happen during the GEMM head →
natural K-before-V DMA priority across #1 / #2 of decode opt2.
3. **Strict FIFO across composites.** PE_SCHEDULER tracks in-flight
`rw_handles`. A new composite whose `rw_handles` (or any of its op
inputs) intersect with an in-flight composite's `rw_handles` waits
for all earlier composites to complete — no out-of-order reorder.
4. **No backwards-incompatibility for legacy callers.** The user-facing
API `tl.composite(op="gemm", a, b, epilogue=[...])` is preserved.
TLContext internally lowers it to a flat-ops `CompositeCmd`.
5. **PE_SCHEDULER recipe-free.** PE_SCHEDULER does not import
RECIPE_DESCRIPTORS, does not know that "softmax_merge" exists, and
has no recipe-specific code branches.
6. **No implicit operand replacement (auto-bind conflict).** If a
prologue recipe declares a `primary_out` AND the head op's auto-bind
target operand (e.g., GEMM `a`) is *also* provided by the kernel
explicitly, TLContext raises a validation error at emit time. The
kernel must either let the recipe bind (omit the operand) **or**
specify it explicitly (omit the prologue, or use a recipe without
`primary_out`). This prevents the ambiguity of "which value wins".
7. **Prologue MATH operands TCM-only.** Restatement of D4: every operand
of a *prologue recipe* (non-GEMM) OpSpec must be TCM-resident. The head
op (gemm or math) is exempt — it may stream from HBM. Phase 1 enforced
at TLContext emit.
### D7. Boundary summary (compiler vs scheduler vs engine)
```
HOST DEVICE
TLContext (compiler analog) PE_SCHEDULER (scheduler/dispatcher)
- validate against RECIPE_DESC. - scan flat ops list
- derive primary_out shape - identify GEMM (≤1)
- allocate scratch - position-based stage placement
- expand recipe → flat OpSpecs - auto-insert DMA stages from space
- compute rw_handles - strict-FIFO RW hazard tracking
- emit CompositeCmd - emit Stage list
imports: RECIPE_DESCRIPTORS imports: nothing recipe-related
ENGINES (PE_MATH / PE_GEMM / PE_DMA)
- read Stage.params (op_kind, addresses, n_elements)
- existing latency models unchanged
imports: nothing recipe-related
```
### D8. Composite output handle + output-space DMA
`tl.composite(...)` returns the **output `TensorHandle`** (not a bare
`CompletionHandle`), so a composite's result can feed a downstream op the
same way `tl.dot`'s result does. The handle's `pending` field references
the composite's completion; downstream consumers auto-await it
(`_await_pending`, ADR-0062 lazy pattern) and `tl.wait(handle)` works.
**`out` is always a `TensorHandle`** (never a raw address) — its location is
the kernel's choice, never guessed:
- **HBM output** → `out=tl.ref(addr, shape)`. `tl.ref` returns a handle with
`space="hbm"`, so the composite writes it back via DMA_WRITE. The
STORE + DMA_WRITE stay **inside** the composite pipeline — they are part of
the fused op, *not* a separate `tl.store`.
- **In-place TCM** → `out=<existing TCM handle>` (e.g. the recipe accumulator
`O`). `space="tcm"` → STORE only, result stays in TCM (chainable).
- **Omitted** → TLContext **auto-allocates a TCM scratch** (like `tl.dot`)
and returns it.
For ergonomics, `out_ptr: int` is accepted as **shorthand for an HBM
output** — equivalent to `out=tl.ref(out_ptr, <output shape>)` (the compiler
wraps it into an `space="hbm"` handle). When no output is specified the
compiler owns the scratch address, never the kernel. (`tl.ref` returns
`space="hbm"` — it references HBM data. The operand-**input** DMA decision
stays `pinned`-based per D4, so this label does not affect input streaming.)
The output handle's **`space` drives the tile loop's write-back** (the
output analog of the operand `pinned` rule in D4):
| `out.space` | tile-loop write-back |
| --- | --- |
| `"tcm"` | STORE only — result stays in TCM (chainable; **no DMA_WRITE**) |
| `"hbm"` | STORE + DMA_WRITE — final HBM output |
This makes a TCM-resident composite output (e.g. a recipe accumulator) legal
**and** keeps a TCM-scratch address out of the DMA engine's PA decoder
(which would otherwise reject the high-bit scratch address). Existing benches
wrap their HBM output via `tl.ref` → write-back unchanged (byte-equal).
## Alternatives
### A1. Three composites per tile (no prologue concept)
Split #2 further into `softmax_merge` + `gemm(P·V)` + `add(O)` as three
separate CompositeCmds. Pros: no flat-ops restructuring (`add` reuses
existing epilogue path). Cons: 3 dispatches per tile instead of 2 —
diverges from ADR-0060 §8 item 4 sizing ("split into **two**"); loses
~32 ns × N_tiles of fixed-cost savings. **Rejected.**
### A2. Mixed-engine recipe (`flash_pv_merge`)
One recipe absorbing MATH + GEMM + MATH (ADR-0060 §5.6 original
phrasing). Pros: tightest dispatch (1 composite). Cons: PE_SCHEDULER
must know engine-boundary recipes, violates the "engines see only
Stage.op_kind" boundary, recipe table becomes ISA-like and pollutes
HW interface. **Rejected.**
### A3. Generic op-list DAG (`tl.ex_composite([...])`)
User-defined sequence of any ops with named slots and dataflow analysis.
Pros: maximally flexible. Cons: a single user (softmax_merge) does not
justify a generic DAG language; PE_SCHEDULER would need dataflow
analysis; reintroduces ADR-0060 §8 item 4's "absorb the inner loop"
trap. **Rejected (Simplicity First).**
### A4. RW-aware scheduler reorder (not strict FIFO)
PE_SCHEDULER could allow a later CompositeCmd to overtake an in-flight
one when their `rw_handles` don't intersect (e.g., next-tile #1 could
overlap with this-tile #2 since #1 doesn't touch m/l/O). Pros: more
overlap. Cons: more scheduler state, gains are bounded (next #1 waits
for K DMA anyway). **Deferred** — strict FIFO ships first; RW-aware
reorder may be a future follow-up ADR.
### A5. Embedding softmax_merge as an epilogue of #1 (Q·Kᵀ)
Wrong: softmax_merge's output `P` is the input of #2's P·V GEMM, which
runs *after* #1's epilogue. Embedding it would force #2 to read its own
input from #1's epilogue scratch — circular tile-loop dependency.
**Rejected (incorrect).**
## Consequences
### Positive
- ADR-0060 §5.6 / §8 item 4 carve-out implemented exactly: opt2 = 2
composites per tile, K-before-V DMA priority natural.
- Composite contract becomes uniform: flat ops list with scope-based
placement. No structural distinction between "head" and "epilogue"
in the HW cmd.
- DMAs auto-inferred from operand `space`; kernel surface simpler.
- HW interface (`CompositeCmd` shape) stays small even as recipes grow
— recipe expansion happens host-side.
- New fused patterns (linear+gelu+norm, rmsnorm+linear, …) can be added
by inserting one `RECIPE_DESCRIPTORS` entry; no PE_SCHEDULER changes.
### Negative
- CompositeCmd's struct shape changes — a refactor for all current
callers. Semantics preserved (D6.4); existing goldens unchanged once
ADR-0064 Revision 2 lands.
- `OpSpec.operands` changes from positional `tuple[Any, ...]` to
`dict[str, TensorHandle]` — touches the existing epilogue lowering.
- New TLContext code: recipe expansion + scratch-slot allocation +
primary_out auto-bind. ~80–120 LOC.
- PE_SCHEDULER's `_generate_plan` gets a position-based scan + DMA
auto-insertion + strict-FIFO RW tracker. ~60–100 LOC.
- **Strict-FIFO RW hazard tracking is safe but conservative.** When
two composites share any `rw_handle`, the new one waits for *all*
prior overlapping composites to fully complete — even when their
compute could in principle overlap (e.g., next-tile #1 = Q·Kᵀ does
not touch m/l/O and could overlap with this-tile's #2, but FIFO
still serializes any path that touches the same handles down the
line). This **may under-expose composite-level overlap** relative to
a smarter scheduler. RW-aware reorder (A4) is the deferred
improvement.
## Open review items
1. **Recipe shape derivation rule beyond identity.** Future recipes may
need non-identity transforms (e.g., transposed primary_out). The
`PrimaryOutSpec.transform` field is forward-compat; add new
transforms as needed.
2. **Recipe scratch allocator** integration with ADR-0063's
`scratch_scope`. Initial wiring: TLContext allocates inside the
current `scratch_scope` context if active, else in a kernel-scope
persistent slot.
3. **`tl_recipes.py` location.** New module under `triton_emu/`. Keeps
recipe knowledge co-located with the compiler analog (TLContext).
4. **GEMM `out=TensorHandle`.** New form alongside legacy `out_ptr:
int`. Both accepted; handle form preferred for new code (enables
strict-FIFO RW tracking by handle identity).
## Test Requirements
1. **CompositeCmd flat-ops refactor — meaning preservation.** Every
existing bench's op_log is byte-equal pre/post refactor (legacy API
`tl.composite(op="gemm", a, b, epilogue=[...])` lowers to the same
Stage sequence).
2. **softmax_merge recipe lowering.** TLContext call `tl.composite(
prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}],
op="gemm", b=V_ref, out=O, epilogue=[{"op": "add", "other": O}])`
produces a `CompositeCmd` with exactly 10 ops in the expected order
(8 MATH + 1 GEMM + 1 MATH(add)) and `rw_handles == (m, l, O)`.
3. **K-before-V DMA priority invariant.** op_log of decode opt2 shows
no V-related DMA during #2's MATH prologue (#2's V DMA begins only
when the GEMM head starts).
4. **Strict-FIFO RW serialization.** Two consecutive composites sharing
any `rw_handle` complete in dispatch order; their stages do not
interleave.
5. **DMA auto-insertion from `pinned`.** A GEMM composite with a
not-`pinned` operand (e.g. `b=tl.ref(...)`) emits a DMA_READ Stage; a
`pinned` operand (e.g. a recipe's TCM scratch / primary-out, or a
`tl.load` result) does not. (As-built D4 note: the DMA decision uses
the existing `pinned` flag, not a `space` tag.)
6. **opt2 numeric parity with opt3.** In data mode, opt2's final
`(m, l, O)` matches opt3's within fp tolerance.
7. **opt2 dispatch ratio (after ADR-0064 Rev2) — robust.** opt3 vs
opt2 per-tile PE_CPU dispatch cycles satisfies `opt3 > 2 × opt2`
(qualitative gate; calibration-independent). The default-calibration
model expectation is ≈ 4.0× (FIXED=40 cycles, R=0.0625 cycles/byte)
— see DDD-0065 §11 for the model-derived numbers; only the loose
`> 2×` bound is the test gate so it survives calibration changes.
Ratio is FIXED-dominated, reflecting command-count reduction as the
primary signal.
8. **GEMM-count invariant.** A composite with two GEMM OpSpecs raises a
validation error at TLContext emit.
9. **Within composite size cap (ADR-0064 D7).** Decode opt2's `#2`
composite (10 ops, ~322 logical bytes — per-op summation of operand
handles) fits comfortably under the default
`MAX_COMPOSITE_LOGICAL_BYTES=1024` — no segmentation. The recipe
lowering test asserts `cmd.logical_bytes < 1024`.
10. **Auto-bind conflict.** `tl.composite(op="gemm", a=A_handle,
prologue=[{"op": "softmax_merge", ...}], ...)` where softmax_merge
declares `primary_out` raises a validation error at emit time
(D6.6).
11. **MATH operand TCM invariant.** `tl.composite(op="math",
a=tl.ref(addr, shape), ...)` (passing an HBM-resident handle as a
MATH operand) raises a validation error at emit time (D6.7).
## Dependencies
- **ADR-0060** §5.6, §8 item 4 — the carve-out this ADR implements.
- **ADR-0064 Revision 2** — dispatch cost model; lands first, makes
opt2's fewer-issues win measurable.
- **ADR-0063** `scratch_scope` — recipe intermediates allocated inside;
`m, l, O` allocated outside as persistent arena (per DDD-0060 §6.2).
- **ADR-0046** `tl_context` contract — extended by `prologue=[...]`
kwarg, `out=TensorHandle`, RECIPE_DESCRIPTORS lookup.
- **ADR-0042** tile plan generators — extended (not redesigned) to
process the flat ops list and auto-insert DMAs.
- **ADR-0014** PE pipeline — boundary preserved: scheduler stays
recipe-free, engines stay op_kind-opaque.
## Migration
> **Implementation ordering note.** The ADR-0064-first ordering below was
> the original plan. The actual implementation **reversed steps 1 and 2** —
> ADR-0065 Phase 1 (flat-ops `CompositeCmd`) landed *before* ADR-0064 Rev2.
> Reason: `logical_bytes` (ADR-0064 D2) is naturally defined on the flat-ops
> shape, so doing the refactor first avoids a throwaway transitional
> `logical_bytes`. Each step is independently safe: P1 is a meaning-
> preserving refactor (op_log byte-equal) under the pre-existing cost path;
> ADR-0064 Rev2 then swaps in the structural formula on the already-flat
> shape. The end state is identical either way.
Land order (original plan; see note above for the as-built reversal of 1↔2):
1. **ADR-0064 Revision 2** (separate PR): structural dispatch cost +
`logical_bytes` property + topology config override + one-time
goldens regeneration.
2. **ADR-0065 Phase 1** (this ADR, PR a): `CompositeCmd` flat-ops
restructure + legacy-API lowering. Refactor only; goldens unchanged.
3. **ADR-0065 Phase 2** (this ADR, PR b): `tl_recipes.py` +
`softmax_merge` recipe + PE_SCHEDULER position-scan + DMA
auto-insertion + strict-FIFO RW tracker.
4. **ADR-0065 Phase 3** (this ADR, PR c): `_gqa_decode_long.py` opt2
variant + numeric parity test + dispatch-ratio measurement against
opt3 baseline.
Detailed implementation plan: see **DDD-0065**.
docs/ddd/DDD-0065-flat-ops-composite-softmax-merge-detailed-design.md
+699
View File
@@ -0,0 +1,699 @@
# Detailed Design Document — Flat-ops `CompositeCmd` + `softmax_merge` recipe (ADR-0065)
**Status:** Accepted (companion to ADR-0065, Accepted) — reconciled to
the as-built implementation.
**Scope:** the `CompositeCmd` flat-ops restructure, the `softmax_merge`
recipe, the decode opt2 variant in `_gqa_attention_decode_opt2.py`, and the
data-mode numeric path that makes opt2's recipe compute.
**Audience:** reviewer / implementer. Read **ADR-0065** and **ADR-0064
Revision 2** first.
---
## 0. As-built reconciliation (read first)
This DDD was authored in the design round; the implementation diverged in
several places. The sections below have been corrected, but the headline
deltas are collected here. (ADR-0065 itself was kept in sync; see commits
`47e2c78`,`184f654`,`55f025c`,`17fb940`,`35453cc`,`4a55ae5`,`e8d6c28`,
`4089e18`,`3d4a3d4`,`5bbe293`,`e9d62b9`,`f0c5b9c`.)
| Topic | Original DDD | As-built |
| --- | --- | --- |
| **Phase order** | P0 (ADR-0064 Rev2) lands first | **P1 (flat-ops) first**, then P0 — `logical_bytes` is naturally defined on the flat shape, so the refactor goes first |
| **Operand-input DMA** | from `handle.space` (`space=="hbm"`→DMA_READ) | **from the `pinned` flag** (already-in-TCM→skip). The `space` flip was avoided (it would break byte-equality); `tl.ref`→hbm, but the *decision* stays pinned-based |
| **Output DMA / handle** | (not addressed; `out=O`) | **ADR-0065 D8** — `tl.composite` returns the output `TensorHandle`; `out`=handle (HBM via `tl.ref`) \| `out_ptr` (HBM shorthand) \| omitted→auto-TCM; **DMA_WRITE gated on `out.space=="hbm"`** (TCM output stays on-chip, chainable) |
| **Size cap (D7)** | deterministic segmentation | **hard cap → `ValueError`** at emit (no segmentation) |
| **Prologue execution** | feeder yields `prologue_stages` first | feeder feeds each prologue stage as a **standalone 1-stage tile** (folding broke PE_MATH→PE_DMA routing) |
| **D6.7 (MATH TCM-only)** | every non-GEMM op | **prologue recipe ops only**; the head op (gemm *or* math) keeps DMA-staged-from-HBM |
| **Numeric parity** | "just works" via the existing data executor | composites did **not** compute in data mode at all — built in **N1–N4** (below): composite-as-gemm record, recipe MATH compute, GEMM accumulate, opt2 ordering |
| **Dispatch ratio** | model ≈ 4.0× | **measured 5.22×** per-tile (gate `>2×`); R-sweep 4.20×/5.22×/5.45× |
| **opt2↔opt3 parity** | exact in data mode | exact recipe math verified by N2/N3 (isolated, non-trivial); e2e opt2↔opt3 holds in the **near-uniform regime** (kernbench K reshape-as-transpose caveat, shared by both kernels) |
**Data-mode numerics (N1–N4), absent from the original DDD — see §15.**
---
## 1. Goal and success criteria
**Goal:** make decode opt2 *runnable*, *numerically equivalent to opt3*,
and *measurably cheaper in PE_CPU dispatch* under the ADR-0064 Rev2 cost
model. Existing benches keep their (post-ADR-0064 Rev2) goldens
unchanged.
**Success criteria:**
1. **Refactor safety.** `CompositeCmd` flat-ops restructure preserves the
op_log byte-for-byte for every existing bench
(`tests/<existing>` all green, no goldens diff).
2. **Recipe correctness.** opt2 final `(m, l, O)` matches opt3 within fp
tolerance in data mode (`enable_data=True`).
3. **Boundary preservation.** PE_SCHEDULER imports nothing recipe-related;
PE_MATH/PE_GEMM/PE_DMA see only `Stage.params["op_kind"]` (no new
engine code).
4. **Dispatch ratio (robust).** Per-tile PE_CPU dispatch cycles satisfy
`opt3 > 2 × opt2` (qualitative invariant, calibration-independent).
Default ADR-0064 Rev2 calibration (FIXED=40 cycles, R=0.0625
cycles/byte) gives model-expected ≈ 4.0× — informative only, not the
test gate.
5. **K-before-V invariant.** op_log of decode opt2 shows zero V-related
DMA during #2's MATH prologue.
**Non-goals:** prefill (causal `if` in kernel control flow, separate
work); RW-aware scheduler reorder (strict FIFO ships); user-defined
recipes; new engine micro-ops in PE_MATH (it stays op_kind-opaque).
---
## 2. Where the design sits in the current code
```
runtime_api (host)
RuntimeContext.zeros/empty/launch context.py
│ KernelLaunchMsg
▼
sim_engine
GraphEngine(enable_data=…) engine.py
MemoryStore memory_store.py
│
▼
components / PE pipeline (per PE)
PE_CPU → PE_SCHEDULER → PE_{DMA,GEMM,MATH,IPCQ,TCM}
greenlet kernel ↔ SimPy kernel_runner.py, pe_cpu.py
CompositeCmd → tile plan → engines pe_scheduler.py (ADR-0014 D6)
│
▼
tl programming model tl_context.py
composite / dot / load / store / scratch_scope
── ADR-0065 NEW ──────────────────────────
tl_recipes.py RECIPE_DESCRIPTORS, EngineOp,
RecipeDescriptor, PrimaryOutSpec
```
The new module `tl_recipes.py` sits beside `tl_context.py`. PE_SCHEDULER
(`pe_scheduler.py`) and `tiling.py` extend to consume the flat ops list
without importing `tl_recipes`.
---
## 3. File plan
### 3.1 Modified production files
| File | Change | ADR ref |
|---|---|---|
| `src/kernbench/common/pe_commands.py` | `CompositeCmd` flat-ops dataclass; `OpSpec.operands: dict`; new `Scope` value semantics (KERNEL position-based); `logical_bytes` property on every `PeCommand`. | 0065 D1/D2, 0064 D2 |
| `src/kernbench/triton_emu/tl_context.py` | `composite(prologue=[...], out=TensorHandle, …)`; recipe lowering pass (RECIPE_DESCRIPTORS lookup, scratch alloc, OpSpec expansion, primary_out auto-bind); legacy API preserved (`op="gemm", a, b, epilogue=[...]` still works, internally lowers to flat ops). | 0065 D1/D5 |
| `src/kernbench/triton_emu/tl_recipes.py` (NEW) | `RecipeDescriptor`, `PrimaryOutSpec`, `EngineOp` dataclasses; `RECIPE_DESCRIPTORS["softmax_merge"]` entry; expansion helper. | 0065 D5 |
| `src/kernbench/components/builtin/pe_scheduler.py` | `_generate_plan` rewritten for flat ops: scan for GEMM (≤1), position-based KERNEL placement, OUTPUT_TILE/K_TILE epilogue, strict-FIFO RW tracker keyed by handle id. | 0065 D3/D6 |
| `src/kernbench/components/builtin/tiling.py` | `generate_plan_from_ops(ops, tile_*, bpe, pe_prefix)` wraps `generate_gemm_plan`/`generate_math_plan`. Input DMA_READ from the **`pinned`** flag; **output DMA_WRITE gated on `out.space=="hbm"`** (D8). `_math_stage` carries operand+output addrs on recipe prologue stages (N2). | 0065 D3/D4/D8 |
| `src/kernbench/common/pe_cost_model.py` (NEW) | `PeCostModel` (FIXED/R/cap) + `from_node_attrs`; replaces the removed Rev1 `cpu_issue_cost.py`. | 0064 D1/D4 |
| `src/kernbench/components/builtin/pe_cpu.py` | reads `pe_cost_model:` + `clock_freq_ghz` from node attrs, threads to TLContext (legacy + greenlet KernelRunner). | 0064 D1/D4 |
| `src/kernbench/sim_engine/op_log.py` | `_extract_op_info(CompositeCmd)` scans for the GEMM op + `accumulate` flag (N1/N3); `record_end` promotes an address-carrying MATH stage to op_kind="math" (N2). | 0065 N1-N3 |
| `src/kernbench/sim_engine/data_executor.py` | `_compute_math` +5 recipe ops; `_execute_gemm` best-effort skip + accumulate (N1-N3). | 0065 N1-N3 |
| `src/kernbench/benches/_gqa_attention_decode_opt2.py` (NEW) | opt2 variant — tile 0 establishes (m,l,O) with primitives; the merge tile uses `#1 scores = tl.composite(op="gemm", a=Q, b=K_T1)` + `#2 tl.composite(prologue=[softmax_merge], op="gemm", b=V_ref, out=O, epilogue=[add])`. | 0065 §Migration P3 |
### 3.2 New tests
| File | Covers |
|---|---|
| `tests/test_composite_flat_ops_refactor.py` | P1 flat-ops lowering shape (legacy + math + epilogue) |
| `tests/test_pe_cost_model.py` | ADR-0064 Rev2 — `logical_bytes`/formula, override, `PeCpuOverheadCmd` bypass, D7 cap→error, live wiring/path-parity |
| `tests/test_tl_recipe_softmax_merge.py` | recipe lowering — 10 ops in order + `rw_handles == (m, l, O)` + size cap |
| `tests/test_tl_composite_validation.py` | auto-bind conflict (D6.6); prologue MATH operand TCM-only (D6.7) |
| `tests/test_pe_scheduler_flat_ops_plan.py` | position-scan; pinned-based DMA; prologue stages; MATH-only |
| `tests/test_pe_scheduler_strict_fifo.py` | `_RwHazardTracker` admit/retire (RAW/WAW/disjoint/legacy) |
| `tests/test_composite_output_handle.py` | D8 — composite returns handle (auto-TCM); chaining; output-space DMA; `tl.ref`→hbm |
| `tests/test_recipe_data_mode.py` | **N2/N3** — recipe MATH compute m/l/O-rescale + full accumulate (O = O·corr + P@V) vs numpy |
| `tests/test_e2e_pipeline.py` | **N1** — composite GEMM computes a@b in data mode |
| `tests/attention/test_gqa_decode_opt2.py` | dispatch ratio (>2×) + R-sweep; K-before-V; **N4** opt2↔opt3 parity |
---
## 4. Data model — `CompositeCmd` flat-ops
### 4.1 New shape
```python
@dataclass(frozen=True)
class CompositeCmd:
completion: CompletionHandle
ops: tuple[OpSpec, ...] # ordered MATH/GEMM ops
rw_handles: tuple[TensorHandle, ...] = () # cross-composite hazard
data_op: bool = True
@property
def logical_bytes(self) -> int:
return (
4 # framing
+ 1 + sum(op.logical_bytes for op in self.ops)
+ 1 + 8 * len(self.rw_handles)
)
```
The previous `op`, `a`, `b`, `out_addr`, `out_nbytes`, `math_op` fields
are **removed**. All info lives in `ops[i]`.
### 4.2 `OpSpec` extended
```python
@dataclass(frozen=True)
class OpSpec:
kind: str # "gemm" | "rmax" | "exp_diff" | ...
scope: Scope # KERNEL | K_TILE | OUTPUT_TILE
operands: dict[str, TensorHandle]
extra: dict[str, Any] = field(default_factory=dict)
out: TensorHandle | None = None
@property
def logical_bytes(self) -> int:
return (
1 + 1 # opcode + scope enum
+ 1 + 8 * len(self.operands) # named handles
+ (8 if self.out is not None else 0)
+ 1 + 4 * len(self.extra) # scalars
)
```
Backward compat: existing `operands: tuple[Any, ...]` use sites
(epilogue `bias`, `scale`, `add`, etc.) migrate to `operands={"other":
h}` / `operands={"src": h}` per a small lookup in
`pe_commands.EPILOGUE_OPS`.
### 4.3 `Scope` semantics — phase vs repetition (ADR-0065 D3)
`Scope` keeps its three values (`KERNEL`, `K_TILE`, `OUTPUT_TILE`).
What changes is the **semantics split**:
- **Position** determines the *phase* (where in the tile-loop lifecycle
the op fires — pre-loop / in-loop / post-loop).
- **Scope** determines the *repetition* within the phase
(KERNEL = once per `CompositeCmd` invocation; K_TILE = per K-tile
iteration; OUTPUT_TILE = per output tile).
`KERNEL` does *not* mean "once per kernel launch" or "once anywhere".
It means "once per `CompositeCmd` invocation, at this op's position".
A composite with multiple KERNEL-scope ops (e.g., softmax_merge's 8
prologue ops) executes each exactly once, in sequence order, before
the GEMM tile loop begins.
Concretely:
- KERNEL ops *before* GEMM → pre-tile-loop (single shot, in order)
- KERNEL ops *after* GEMM → post-tile-loop (single shot, in order)
- K_TILE / OUTPUT_TILE ops *after* GEMM → inside tile loop as today
---
## 5. Recipe: `softmax_merge`
### 5.1 `tl_recipes.py` definitions
```python
@dataclass(frozen=True)
class PrimaryOutSpec:
from_shape: str # operand name to copy shape from
from_dtype: str # operand name to copy dtype from
transform: str # "identity" | "trans" (forward-compat)
@dataclass(frozen=True)
class EngineOp:
engine: Literal["MATH", "GEMM", "DMA"]
op_kind: str # PE_MATH op_kind label (opaque to engine)
# Slot names: any of "src", "src_a", "src_b", "src_c", "dst" — these
# are RecipeDescriptor.operands names or internal scratch names.
src: str | None = None
src_a: str | None = None
src_b: str | None = None
src_c: str | None = None
dst: str | None = None
reduce_axis: int | None = None
bcast_axis: int | None = None
@dataclass(frozen=True)
class RecipeDescriptor:
operands: dict[str, str] # name → "R" | "RW"
primary_out: PrimaryOutSpec | None # implicit output spec
tile_alignment: Literal["single_shot", "tile_aligned"]
internal_scratch_bytes_fn: Callable[..., int]
engine_seq: tuple[EngineOp, ...]
RECIPE_DESCRIPTORS: dict[str, RecipeDescriptor] = { ... }
```
### 5.2 Engine sequence (8 MATH ops)
| # | engine | op_kind | reads (slots) | writes (slots) | extras |
|---|---|---|---|---|---|
| 1 | MATH | `rmax` | s | m_loc | `reduce_axis=-1` |
| 2 | MATH | `max_elem` | m, m_loc | m_new | — |
| 3 | MATH | `exp_diff` | m, m_new | corr | — |
| 4 | MATH | `exp_diff` | s, m_new | P (primary_out)| `bcast_axis=0` |
| 5 | MATH | `rsum` | P | l_loc | `reduce_axis=-1` |
| 6 | MATH | `fma` | l, corr, l_loc | l | — |
| 7 | MATH | `mul_bcast` | O, corr | O | `bcast_axis=1` |
| 8 | MATH | `copy` | m_new | m | — |
Intermediates `m_loc, m_new, corr, P, l_loc` are scratch slots allocated
by TLContext before lowering. `P` is the recipe's primary output; auto-
bound to the head GEMM's `operands["a"]`.
### 5.3 PE_MATH op_kind additions
PE_MATH is op_kind-opaque ([pe_math.py:90-99](src/kernbench/components/builtin/pe_math.py#L90-L99)).
Latency = `ceil(num_elements / vector_width) / clock`. Adding new
op_kind labels (rmax, max_elem, exp_diff, rsum, fma, mul_bcast, copy)
requires **no PE_MATH code changes** — they pass through as string
labels in `Stage.params["op_kind"]`. Latency stays consistent for the
same `num_elements`.
(Reductions like `rmax` / `rsum` have log-depth final stages in real HW
not modelled by `ceil(num_elements / vector_width)`; this is an existing
approximation, not introduced by ADR-0065.)
---
## 6. TLContext lowering pipeline
Input (user-facing API, preserved):
```python
tl.composite(
op="gemm", b=tl.ref(v_tile(j), (TILE, d)), out=O,
prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}],
epilogue=[{"op": "add", "other": O}],
)
```
Lowering steps in `TLContext.composite()`:
1. **Validate prologue items.** For each, lookup
`RECIPE_DESCRIPTORS[item["op"]]`. Check all operands present, type
tags (R / RW) match handle states (e.g., RW must not be `tl.ref`).
2. **Allocate internal scratch.** Call
`recipe.internal_scratch_bytes_fn(G, TILE, d, bpe)`; allocate via
`scratch_scope` helper (ADR-0063). Get slot addresses for `m_loc,
m_new, corr, P, l_loc`.
3. **Derive primary_out handle.** From `recipe.primary_out`: `P` handle =
`TensorHandle(shape=Sj.shape, dtype=Sj.dtype, addr=<P slot>,
space="tcm")`.
4. **Expand `engine_seq` → flat OpSpecs.** For each `EngineOp`:
- Resolve slot names (operands or scratch) to `TensorHandle`s.
- Build `OpSpec(kind=op_kind, scope=Scope.KERNEL, operands={...},
extra={...}, out=...)`.
- `out` is the handle for `dst` slot; `operands` is `{src*: handle}`.
5. **Auto-bind head GEMM (with conflict check, ADR-0065 D6.6).** GEMM
OpSpec built from `op="gemm"`, `operands={"a": P_handle, "b":
V_handle}`, `out=<output handle>`, `extra={"m": G, "k": TILE, "n":
d}`. **Validation:** if the kernel passed `a=<some handle>`
explicitly AND there is a prologue recipe with `primary_out` that
would auto-bind into `a`, raise a validation error (auto-bind
conflict). The kernel must omit `a` (let recipe bind) or omit the
prologue / use a `primary_out=None` recipe (no auto-bind).
6. **Build epilogue OpSpecs.** Each `epilogue=[{...}]` entry maps to
`OpSpec(kind=..., scope=Scope.OUTPUT_TILE, operands={...}, ...)`.
The `add` epilogue specifically gets `operands={"other": O_handle}`
so the Stage later carries `other_handle` (GAP 9 fix).
7. **Compute `rw_handles`.** Collect all handles tagged RW in prologue
recipes + the head GEMM's `out` (if accumulating) + any RW epilogue
targets. For softmax_merge: `rw_handles = (m, l, O)`.
8. **Operand validation (ADR-0065 D6.7, as-built scope).** Checked on the
**prologue recipe ops only**: every operand handle must be
`space == "tcm"` (HBM recipe operand → `ValueError`). The head op (gemm
*or* math) is **exempt** — it may stream from HBM. softmax_merge's 8 MATH
ops only reference TCM-resident handles → the invariant naturally holds.
Also: GEMM-count ≤ 1 (D6.1).
9. **Check composite size cap (ADR-0064 D7, as-built).** If
`cmd.logical_bytes > max_composite_logical_bytes` (default 1024), raise
`ValueError` at emit — **hard cap, no auto-segmentation**. softmax_merge's
10-op composite (~322 bytes) is well under the cap.
10. **Resolve the output handle (ADR-0065 D8) + emit.** `out`=handle (HBM
`tl.ref` \| in-place TCM) \| `out_ptr` (HBM shorthand) \| omitted→auto-TCM
scratch. Emit `CompositeCmd(ops=(8 MATH + 1 GEMM + 1 MATH(add)),
rw_handles=(m,l,O))`, attach a `CompositeFuture` (carrying the
completion) to the output handle, and **return that handle** so the
result chains like `tl.dot`'s (downstream auto-awaits via
`_await_pending`).
**Legacy lowering.** Pre-existing
`tl.composite(op="gemm", a=A, b=B, epilogue=[{"op":"bias","bias":h}])`
goes through the same emit path: no prologue, head OpSpec from `op="gemm"`,
epilogue OpSpecs built from `epilogue=[...]`. Output `CompositeCmd` has
flat ops list; semantics preserved.
---
## 7. PE_SCHEDULER tile plan generation (flat-ops)
```python
def _generate_plan(cmd: CompositeCmd) -> PipelinePlan:
pp = self._pe_prefix
bpe = 2
# 1. Identify GEMM op (≤1).
gemm_idx = next(
(i for i, o in enumerate(cmd.ops) if o.kind == "gemm"),
None,
)
# 2. MATH-only composite (no GEMM): single tile of all ops in order.
if gemm_idx is None:
return _math_only_plan(cmd.ops, pp)
gemm = cmd.ops[gemm_idx]
M, K, N = gemm.extra["m"], gemm.extra["k"], gemm.extra["n"]
# 3. Partition ops by position.
pre_loop_ops = cmd.ops[:gemm_idx] # KERNEL scope
after_ops = cmd.ops[gemm_idx+1:] # mixed scopes
# 4. Pre-loop MATH Stages (single-shot, before tile loop).
pre_stages = [_math_stage_from_op(op, pp) for op in pre_loop_ops]
# 5. Tile loop (existing generate_gemm_plan pattern, but:
# - DMA_READ inserted only when operand.space == "hbm"
# - K_TILE epilogue from after_ops with scope == K_TILE
# - OUTPUT_TILE epilogue from after_ops with scope == OUTPUT_TILE
# - DMA_WRITE inserted only when out.space == "hbm")
tile_plans = _generate_tile_loop(
gemm, M, K, N, after_ops, pp, bpe,
)
# 6. Post-loop KERNEL Stages (after tile loop completes).
post_loop_ops = [op for op in after_ops if op.scope == Scope.KERNEL]
post_stages = [_math_stage_from_op(op, pp) for op in post_loop_ops]
return PipelinePlan(
prologue_stages=tuple(pre_stages),
tiles=tile_plans,
epilogue_stages=tuple(post_stages),
m_tiles=..., k_tiles=..., n_tiles=...,
)
```
`PipelinePlan` gets two new fields:
```python
@dataclass(frozen=True)
class PipelinePlan:
tiles: list[TilePlan]
m_tiles: int
k_tiles: int
n_tiles: int
prologue_stages: tuple[Stage, ...] = ()
epilogue_stages: tuple[Stage, ...] = ()
```
Feeder change (`_feed_loop`): yield `prologue_stages` Stages first, then
tile loop as today, then `epilogue_stages`. For the single-Stage case
(no prologue/epilogue), behaviour identical to today.
---
## 8. Strict-FIFO RW hazard tracker
```python
class _RwHazardTracker:
def __init__(self):
self._inflight: list[tuple[CompletionHandle, set[str]]] = []
# (completion, set of handle ids written or RW)
def admit(self, env, cmd: CompositeCmd) -> Generator:
"""Block until no in-flight composite has overlapping RW."""
rw_ids = {h.id for h in cmd.rw_handles}
# Strict FIFO: wait for ALL prior composites whose rw_set
# intersects, in order.
while True:
blocking = [c for c, s in self._inflight if s & rw_ids]
if not blocking:
break
# Wait on the *first* blocking composite's completion.
yield blocking[0].done_event
self._inflight.append((cmd.completion, rw_ids))
def retire(self, completion: CompletionHandle) -> None:
self._inflight = [
(c, s) for c, s in self._inflight if c != completion
]
```
`PE_SCHEDULER` invokes `admit()` before queueing a CompositeCmd's tile
plan and `retire()` upon `done_event.succeed()`. Strict FIFO: even when
overlap is partial, the new composite waits for all prior overlapping
composites to *fully* drain — simple, easy to reason about, easy to
test. (RW-aware reorder deferred per ADR-0065 A4.)
**Conservatism trade-off (ADR-0065 Consequences/Negative).** This rule
is safe but may under-expose overlap. Concretely for decode opt2:
next-tile's #1 (Q·Kᵀ) does not touch `(m, l, O)` and *could* in
principle run while this-tile's #2 is still computing on those handles.
Under strict FIFO, next-tile's #1 still gets dispatched (no rw overlap
with prior in-flight) — so #1 cross-tile pipelining is preserved.
However, if a future bench composed multiple #2-like recipes that share
handles, FIFO would serialize them even when their compute could
overlap. Acceptable for Phase 1; revisit when a real workload shows
the gap.
---
## 9. Phased implementation plan (test-first per CLAUDE.md)
Each phase is an independent Phase-1 → Phase-2 cycle per CLAUDE.md Part 1.
| Phase | Deliverable | Gate |
|---|---|---|
| **P0 (= ADR-0064 Rev2)** | `logical_bytes` property + dispatch formula + topology override + size cap; one-time goldens regeneration | typical 1-OpSpec composite ≈ 43 ns; dispatch ratio test infrastructure in place |
| **P1** | `CompositeCmd` flat-ops dataclass refactor + `OpSpec.operands: dict` + legacy `tl.composite` lowering preserved | every existing bench's op_log byte-equal pre/post |
| **P2** | `tl_recipes.py` + `RECIPE_DESCRIPTORS["softmax_merge"]` + TLContext lowering for `prologue=[...]` | recipe lowering test (10 ops, rw_handles, addresses) passes |
| **P3** | PE_SCHEDULER `_generate_plan` flat-ops + `PipelinePlan.prologue_stages` + DMA auto-insertion + `_feed_loop` extension | all existing bench Stage sequences unchanged; new MATH chain visible in op_log for opt2 mock |
| **P4** | Strict-FIFO RW hazard tracker + `_RwHazardTracker` integration | two-composite RW conflict test serializes correctly |
| **P5** | `_gqa_decode_long.py` opt2 variant; data-mode numeric parity check; K-before-V invariant check | opt2 matches opt3 within fp tolerance; no V DMA during prologue |
| **P6** | Dispatch-ratio measurement: opt3 vs opt2 per-tile PE_CPU cycles + R sensitivity sweep | `opt3 > 2 × opt2` at default calibration (gate); model-expected ≈ 4.0× (informative); `opt2 < opt3` at all `R ∈ {0.25, 0.0625, 0.03125}` (ADR-0064 Test #9) |
**As-built ordering:** the implementation reversed P0↔P1 — **P1 (flat-ops
refactor) landed first**, then P0 (ADR-0064 Rev2 cost model), because
`logical_bytes` is naturally defined on the flat shape (doing P1 first avoids
a throwaway transitional definition). P1 is a pure, byte-equal refactor; P0
then swaps in the structural formula and regenerates goldens (in practice no
golden churn occurred). P5 was split: **P5(B)** shipped the opt2 kernel +
op_log dispatch measurement first (numeric parity deferred), then **D8 +
N1–N4** (§15) built the data-mode numeric path. P6 closed the measurement.
---
## 10. Verification plan (concrete)
Mirrors ADR-0065 §Test Requirements; grounded in SPEC R2/R5, ADR-0023/
0025, ADR-0046, ADR-0054.
**Refactor safety (P1):**
- Parametrize over `tests/<existing bench>.py` baseline op_logs.
- Pre/post refactor diff must be empty for *every* bench.
- Failure → revert P1 atomically; do not land partial.
**Recipe lowering (P2):**
- Given `Sj=(8,64), m=(8,), l=(8,), O=(8,128)`, lower
`tl.composite(prologue=[{"op":"softmax_merge", "s":Sj,...}], op="gemm",
b=V_ref, out=O, epilogue=[{"op":"add", "other":O}])`.
- Assert `len(cmd.ops) == 10`, `cmd.ops[0..7].kind == [rmax, max_elem,
exp_diff, exp_diff, rsum, fma, mul_bcast, copy]`,
`cmd.ops[8].kind == "gemm"`, `cmd.ops[9].kind == "add"`.
- Assert `cmd.rw_handles == (m, l, O)`.
- Assert all OpSpec addresses are concrete (no None).
**Scheduler plan (P3):**
- For the above CompositeCmd, assert PipelinePlan has 8 prologue_stages
(MATH), tile loop with FETCH/GEMM/MATH(add), no epilogue_stages.
- Assert no DMA_READ for prologue MATH stages (operands TCM-resident).
- Assert one DMA_READ for V (`b.space == "hbm"`).
- Assert no DMA_WRITE (`out.space == "tcm"`).
**RW hazard (P4):**
- Submit composite A (rw_handles=(O,)) then composite B (operands
reading O). Assert B's Stages start only after A's done_event.
**Numeric parity (P5):**
- In data mode, run opt2 and opt3 for `(G=8, T_q=1, S_kv=64, C=1, P=1)`.
- Assert `O_opt2 ≈ O_opt3` within `fp16` tolerance.
**K-before-V invariant (P5):**
- op_log inspection: for each tile's #2 CompositeCmd, no
`memory/dma_read` events between #2 start and the GEMM compute begin
that are not the V tile DMA itself.
**Dispatch ratio (P6):**
- For `S_kv=64, n_tiles=16`, measure total PE_CPU dispatch cycles for
opt3 vs opt2 paths.
- Assert `opt3 > 2 × opt2` (robust gate — calibration-independent).
- Record observed ratio; the default-calibration model expects ≈ 4.0×
(DDD §11). The recorded number is informative for performance
tracking, not a test gate.
- **Sensitivity sweep.** Repeat with `R ∈ {0.25, 0.0625, 0.03125}`
cycles/byte; assert opt2 < opt3 in all three. (ADR-0064 Test #9.)
**Invariant guards (P2):**
- Auto-bind conflict: `tl.composite(op="gemm", a=A_handle,
prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}],
...)` → `pytest.raises(ValidationError, match="auto-bind conflict")`.
- MATH operand TCM-only: `tl.composite(op="math", a=tl.ref(addr,
shape), ...)` → `pytest.raises(ValidationError, match="MATH operand
must be TCM")`.
- Per-op logical_bytes summation: a synthetic composite with the same
handle `O` in 3 OpSpecs + in `rw_handles` reports
`logical_bytes` higher than a deduplicated equivalent by exactly
`3 × 8` (operand references) + `8` (rw_handles entry) = 32 bytes.
**Determinism (all phases):**
- Identical inputs → identical op_log + latency (SPEC §0.1).
---
## 11. Performance model (expected)
Per-tile PE_CPU dispatch cycles, ADR-0064 Rev2 default calibration
(FIXED=40 cycles, R=0.0625 cycles/byte, 1 cycle = 1 ns at 1 GHz):
```
opt3: 10 cmds × 40 + total_bytes(≈232) × 0.0625 = 400 + 14.5 ≈ 414.5 ns
opt2: 2 cmds × 40 + total_bytes(≈380) × 0.0625 = 80 + 23.75 ≈ 103.75 ns
ratio = 414.5 / 103.75 ≈ 4.0×
```
The FIXED term dominates (~96% of the differential), reflecting
ADR-0064's framing: the primary signal is command-count reduction.
The byte term is a secondary refinement (visible: opt2 carries more
total bytes than opt3, so its byte cost is higher; this is more than
offset by 8× fewer FIXED payments).
Per-tile engine timing dominates total latency; dispatch is one
component. The ratio above is the *dispatch-only* relative cost. Total
latency improvement depends on whether dispatch was on the critical
path — measured per workload in P6.
**R sensitivity** (ADR-0064 Test #9 — sanity at neighbouring R):
| R (cycles/byte) | opt3 ns | opt2 ns | ratio |
|---|---|---|---|
| 0.25 (4 B/cy) | 458 | 175 | 2.6× |
| 0.0625 (16 B/cy) ✓| 414.5 | 103.75 | 4.0× |
| 0.03125 (32 B/cy) | 407.25 | 91.875 | 4.4× |
`opt2 < opt3` at all three; ratio is monotonic in `1/R` (more
FIXED-dominated as R decreases).
**As-built (measured, `test_gqa_decode_opt2.py`).** The per-tile emitters
give opt3=960 ns, opt2=184 ns → **5.22×** at the default calibration (higher
than the model's ≈4.0× because the opt3 emitter includes the full online
merge); R-sweep ratios **4.20× / 5.22× / 5.45×** for R∈{0.25,0.0625,0.03125},
monotonically increasing as R decreases. The gate is the calibration-robust
`opt3 > 2 × opt2`.
---
## 12. Open items (status)
Most ADR-0065 §Open review items map to specific implementation choices
already documented in §§5–8. Live remaining items:
| Item | Resolution venue |
|---|---|
| Non-identity `PrimaryOutSpec.transform` (e.g., `"trans"`) | Future recipe addition; current code uses string dispatch in step 3 of §6 |
| `_gqa_decode_long.py` opt2 — TILE size sweep | Open (existing §B-3, §9 of ADR-0060) |
| Larger composite descriptors approaching ADR-0064 D7 cap (1024 bytes) | Hard cap → `ValueError` at emit (no segmentation); kernel must split the recipe |
| `scratch_scope` budget vs softmax_merge intermediates (G·TILE for P) | ADR-0063 sizing check at S_kv = 256K |
| Scheduler plan-gen cost for large composites (ADR-0064 Open Item 2) | If P6 measurements show scheduler-bound behaviour near the cap, expose per-op-count `overhead_ns` on PE_SCHEDULER |
---
## 13. Risks
| Risk | Likelihood | Mitigation |
|---|---|---|
| `CompositeCmd` refactor (P1) changes op_log for some bench | med | Parametrized regression over all benches; revert atomically |
| `OpSpec.operands` migration (positional → dict) breaks epilogue site | med | Per-EPILOGUE_OPS migration table; existing epilogue tests cover |
| `_RwHazardTracker` false-positive serializes too aggressively | low | Strict FIFO is the safe direction; RW-aware reorder is a future ADR |
| ADR-0064 Rev2 and ADR-0065 P1 land in wrong order → stale goldens | med | PR ordering enforced; CI gate requires both to be green |
| `softmax_merge` scratch budget exceeds `scratch_scope` at long context | med | ADR-0063 sizing check; if needed, tile_alignment promotion (future) |
| PE_MATH op_kind label collision with new strings | low | Labels are opaque strings — no enum conflict; just add to allowed set if validated |
---
## 14. Glossary & references
- **flat ops list** — `CompositeCmd.ops: tuple[OpSpec, ...]` with no
head/epilogue field; placement determined by op position + scope.
- **prologue (user API)** — `tl.composite(prologue=[...])` kwarg for
recipe ops that lower into pre-GEMM KERNEL-scope OpSpecs. *HW command
does not carry this concept.*
- **recipe** — a named macro op (e.g., `softmax_merge`) that TLContext
expands to multiple flat OpSpecs at emit time. Registry =
`RECIPE_DESCRIPTORS`.
- **primary_out** — recipe's implicit output (e.g., `P` for
softmax_merge). Auto-bound to head GEMM's first matrix operand.
- **handle-identity strict FIFO** — cross-composite serialization rule:
a composite waits for all prior composites whose `rw_handles`
intersect with its own to complete.
- **logical_bytes** — HW-logical byte size of a PE command, derived
from a structural rule (ADR-0064 D2), not from Python's
`sys.getsizeof`.
Source anchors: `tl_context.py` (tl API), `pe_scheduler.py:104-143`
(composite tile plan), `pe_dma.py:45,89` (read channel), `pe_math.py`
(op_kind opaque latency), `pe_cpu.py` (dispatch hook),
`pe_commands.py` (CompositeCmd, OpSpec, EPILOGUE_OPS).
ADRs: **0065** (this design), **0064 Revision 2** (dispatch cost),
**0060** (GQA fused attention, §5.6 / §8 item 4 carve-out), **0063**
(scratch_scope), **0046** (tl_context contract), **0042** (tile plan
generators), **0014** (PE pipeline execution model).
---
## 15. Data-mode numerics (D8 + N1–N4) — as-built
The original DDD assumed opt2's numeric parity "just works" via the existing
DataExecutor. It does not: a composite in the builtin tiled path is recorded
**only as latency TileTokens** (`op_kind="unknown"`, no addresses) — its
matmul/MATH is never computed; the output address is never written. Making
opt2's recipe compute required building the composite/recipe data-mode path.
**D8 — output handle + output-space DMA (commit `e8d6c28`).** Prerequisite:
a composite's output must be a real, addressable, chainable handle, and a
TCM-resident output must NOT be DMA_WRITTEN (its high-bit scratch address
crashes the DMA PA decoder). `tl.composite` returns the output `TensorHandle`
(auto-TCM if `out` omitted); the tile loop's DMA_WRITE is gated on
`out.space=="hbm"`. This is what lets opt2 *run* in data mode at all.
**N1 — composite GEMM record (commit `4089e18`).** `op_log._extract_op_info
(CompositeCmd)` scans `ops` for the GEMM (not `ops[0]` — a recipe's GEMM sits
after the prologue) and emits one `composite_gemm` record (a/b/out
addrs+spaces). `PE_SCHEDULER` records the composite's numeric op (no-op
without an op_logger). `DataExecutor._execute_gemm` is best-effort: skips on
(KeyError, ValueError) reading inputs (torch.empty / sharded operands).
**N2 — recipe MATH ops (commit `3d4a3d4`).** `_compute_math` gains
`rmax`/`rsum` (keepdims) + `max_elem`/`exp_diff`/`mul_bcast` (binary; numpy
broadcasting covers `bcast_axis`). `tiling._math_stage` carries operand+output
addrs/shapes/spaces+axis on the **recipe prologue** stages only;
`op_log.record_end` promotes a MATH stage to `op_kind="math"` *only* when it
carries `input_addrs` (legacy epilogue stages stay opaque → byte-equal). A
latent P2 bug was fixed: reductions keep the reduced axis as size 1 (keepdims)
so they broadcast against `m=(G,1)`. Verified: the 8 MATH ops produce m, l
(full online-softmax update) and O (rescaled by corr) vs a numpy reference.
**N3 — GEMM-epilogue accumulate (commit `5bbe293`).** The composite GEMM
record is marked `accumulate=True` when an `add` epilogue targets the GEMM's
own output; `_execute_gemm` then does `O = O_existing + a@b`. The composite's
numeric op is recorded at **completion** (`_retire_on_done`), not dispatch, so
its `t_start` sorts AFTER its prologue MATH ops → the GEMM reads the computed
P and the rescaled O. Verified: full recipe → `O = O_old·corr + P@V` vs numpy.
**N4 — opt2 parity + ordering (commit `e9d62b9`).** `composite()` awaits its
prologue recipe operands (the cross-composite RAW: #2's `s` is #1's score
tile) so #2 waits for #1. With N1–N3, opt2 computes a correct attention output
e2e in data mode and matches opt3 within fp tolerance. **Scope:** the e2e
opt2↔opt3 parity holds in the near-uniform-attention regime — kernbench's K
reshape-as-transpose (opt3 docstring: correct for zero/symmetric inputs)
makes opt2's per-sub-tile K interpretation differ from opt3's single-tile one
for high-contrast inputs (a pre-existing kernbench limitation shared by both
kernels, not the recipe). The **exact** recipe math is verified non-trivially
by N2/N3 (`test_recipe_data_mode.py`).
docs/diagrams/allreduce_latency_plots/AllReduce_LRAB_2DMesh_6SiP_2x3.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 38 KiB

After

Width:  |  Height:  |  Size: 38 KiB

docs/diagrams/allreduce_latency_plots/AllReduce_LRAB_2Dtorus_6SiP_2x3.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 36 KiB

After

Width:  |  Height:  |  Size: 36 KiB

docs/diagrams/allreduce_latency_plots/AllReduce_LRAB_2Dtorus_6SiP_2x3_with_TCM_SRAM_HBM.csv
+12 -12
View File
@@ -1,13 +1,13 @@
buffer_kind,sip_topology,n_sips,n_elem,bytes_per_pe,latency_ns buffer_kind,sip_topology,n_sips,n_elem,bytes_per_pe,latency_ns
hbm,torus_2d,6,128,256,2120.040000000012 hbm,torus_2d,6,128,256,2345.040000000012
hbm,torus_2d,6,1024,2048,2717.2783333333473 hbm,torus_2d,6,1024,2048,2942.2783333333473
hbm,torus_2d,6,8192,16384,7315.184999999989 hbm,torus_2d,6,8192,16384,7540.184999999989
hbm,torus_2d,6,32768,65536,23081.26500000037 hbm,torus_2d,6,32768,65536,23306.26500000037
sram,torus_2d,6,128,256,2060.040000000012 sram,torus_2d,6,128,256,2285.040000000012
sram,torus_2d,6,1024,2048,2909.2783333333473 sram,torus_2d,6,1024,2048,3134.2783333333527
sram,torus_2d,6,8192,16384,9523.184999999869 sram,torus_2d,6,8192,16384,9748.184999999869
sram,torus_2d,6,32768,65536,32201.265000000385 sram,torus_2d,6,32768,65536,32426.265000000385
tcm,torus_2d,6,128,256,1964.040000000012 tcm,torus_2d,6,128,256,2189.040000000012
tcm,torus_2d,6,1024,2048,2477.2783333333473 tcm,torus_2d,6,1024,2048,2702.2783333333473
tcm,torus_2d,6,8192,16384,6403.185000000109 tcm,torus_2d,6,8192,16384,6628.18500000005
tcm,torus_2d,6,32768,65536,19865.265000000378 tcm,torus_2d,6,32768,65536,20090.265000000378
1 buffer_kind sip_topology n_sips n_elem bytes_per_pe latency_ns
2 hbm torus_2d 6 128 256 2120.040000000012 2345.040000000012
3 hbm torus_2d 6 1024 2048 2717.2783333333473 2942.2783333333473
4 hbm torus_2d 6 8192 16384 7315.184999999989 7540.184999999989
5 hbm torus_2d 6 32768 65536 23081.26500000037 23306.26500000037
6 sram torus_2d 6 128 256 2060.040000000012 2285.040000000012
7 sram torus_2d 6 1024 2048 2909.2783333333473 3134.2783333333527
8 sram torus_2d 6 8192 16384 9523.184999999869 9748.184999999869
9 sram torus_2d 6 32768 65536 32201.265000000385 32426.265000000385
10 tcm torus_2d 6 128 256 1964.040000000012 2189.040000000012
11 tcm torus_2d 6 1024 2048 2477.2783333333473 2702.2783333333473
12 tcm torus_2d 6 8192 16384 6403.185000000109 6628.18500000005
13 tcm torus_2d 6 32768 65536 19865.265000000378 20090.265000000378
docs/diagrams/allreduce_latency_plots/AllReduce_LRAB_2Dtorus_6SiP_2x3_with_TCM_SRAM_HBM.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 75 KiB

After

Width:  |  Height:  |  Size: 77 KiB

docs/diagrams/allreduce_latency_plots/AllReduce_LRAB_Ring1D_6SiP_1x6.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 37 KiB

After

Width:  |  Height:  |  Size: 37 KiB

docs/diagrams/allreduce_latency_plots/comparison_mesh_vs_ring_vs_2DTorus_vs_theoretical_vs_fsim.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 86 KiB

After

Width:  |  Height:  |  Size: 92 KiB

docs/diagrams/allreduce_latency_plots/summary.csv
+36 -36
View File
@@ -1,37 +1,37 @@
algorithm,sip_topology,n_sips,n_elem,bytes_per_pe,bytes_per_sip,latency_ns algorithm,sip_topology,n_sips,n_elem,bytes_per_pe,bytes_per_sip,latency_ns
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,8,16,256,2666.552500000015 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,8,16,256,2918.5525000000157
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,32,64,1024,2747.7400000000152 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,32,64,1024,2999.740000000016
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,64,128,2048,2855.990000000018 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,64,128,2048,3107.990000000019
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,128,256,4096,3072.490000000019 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,128,256,4096,3324.4900000000207
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,512,1024,16384,3337.1133333333582 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,512,1024,16384,3589.1133333333582
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,1024,2048,32768,3708.0333333333692 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,1024,2048,32768,3960.0333333333692
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,2048,4096,65536,4449.873333333393 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,2048,4096,65536,4701.873333333393
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,4096,8192,131072,5933.020000000124 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,4096,8192,131072,6185.020000000124
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,8192,16384,262144,8900.379999999863 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,8192,16384,262144,9152.379999999861
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,16384,32768,524288,14835.099999999224 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,16384,32768,524288,15087.099999999224
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,32768,65536,1048576,26704.540000000765 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,32768,65536,1048576,26956.540000000765
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,49152,98304,1572864,38573.97999999701 lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,49152,98304,1572864,38825.97999999701
lrab_hierarchical_allreduce,ring_1d,6,8,16,256,2365.255833333347 lrab_hierarchical_allreduce,ring_1d,6,8,16,256,2628.2558333333477
lrab_hierarchical_allreduce,ring_1d,6,32,64,1024,2436.9433333333473 lrab_hierarchical_allreduce,ring_1d,6,32,64,1024,2699.943333333348
lrab_hierarchical_allreduce,ring_1d,6,64,128,2048,2532.526666666683 lrab_hierarchical_allreduce,ring_1d,6,64,128,2048,2795.526666666683
lrab_hierarchical_allreduce,ring_1d,6,128,256,4096,2723.693333333349 lrab_hierarchical_allreduce,ring_1d,6,128,256,4096,2986.693333333351
lrab_hierarchical_allreduce,ring_1d,6,512,1024,16384,3048.635000000021 lrab_hierarchical_allreduce,ring_1d,6,512,1024,16384,3311.635000000021
lrab_hierarchical_allreduce,ring_1d,6,1024,2048,32768,3393.4016666666957 lrab_hierarchical_allreduce,ring_1d,6,1024,2048,32768,3656.4016666666957
lrab_hierarchical_allreduce,ring_1d,6,2048,4096,65536,4082.401666666714 lrab_hierarchical_allreduce,ring_1d,6,2048,4096,65536,4345.401666666714
lrab_hierarchical_allreduce,ring_1d,6,4096,8192,131072,5458.80166666677 lrab_hierarchical_allreduce,ring_1d,6,4096,8192,131072,5721.801666666768
lrab_hierarchical_allreduce,ring_1d,6,8192,16384,262144,8216.934999999943 lrab_hierarchical_allreduce,ring_1d,6,8192,16384,262144,8479.934999999887
lrab_hierarchical_allreduce,ring_1d,6,16384,32768,524288,13733.201666665835 lrab_hierarchical_allreduce,ring_1d,6,16384,32768,524288,13996.201666665835
lrab_hierarchical_allreduce,ring_1d,6,32768,65536,1048576,24765.73500000064 lrab_hierarchical_allreduce,ring_1d,6,32768,65536,1048576,25028.73500000064
lrab_hierarchical_allreduce,ring_1d,6,49152,98304,1572864,35798.268333331536 lrab_hierarchical_allreduce,ring_1d,6,49152,98304,1572864,36061.26833333154
lrab_hierarchical_allreduce,torus_2d,6,8,16,256,1700.6025000000095 lrab_hierarchical_allreduce,torus_2d,6,8,16,256,1925.6025000000104
lrab_hierarchical_allreduce,torus_2d,6,32,64,1024,1753.2900000000102 lrab_hierarchical_allreduce,torus_2d,6,32,64,1024,1978.290000000011
lrab_hierarchical_allreduce,torus_2d,6,64,128,2048,1823.540000000012 lrab_hierarchical_allreduce,torus_2d,6,64,128,2048,2048.540000000012
lrab_hierarchical_allreduce,torus_2d,6,128,256,4096,1964.040000000012 lrab_hierarchical_allreduce,torus_2d,6,128,256,4096,2189.040000000012
lrab_hierarchical_allreduce,torus_2d,6,512,1024,16384,2196.8183333333463 lrab_hierarchical_allreduce,torus_2d,6,512,1024,16384,2421.8183333333463
lrab_hierarchical_allreduce,torus_2d,6,1024,2048,32768,2477.2783333333473 lrab_hierarchical_allreduce,torus_2d,6,1024,2048,32768,2702.2783333333473
lrab_hierarchical_allreduce,torus_2d,6,2048,4096,65536,3038.1983333333583 lrab_hierarchical_allreduce,torus_2d,6,2048,4096,65536,3263.1983333333583
lrab_hierarchical_allreduce,torus_2d,6,4096,8192,131072,4159.5050000000665 lrab_hierarchical_allreduce,torus_2d,6,4096,8192,131072,4384.5050000000665
lrab_hierarchical_allreduce,torus_2d,6,8192,16384,262144,6403.185000000109 lrab_hierarchical_allreduce,torus_2d,6,8192,16384,262144,6628.18500000005
lrab_hierarchical_allreduce,torus_2d,6,16384,32768,524288,10890.5449999995 lrab_hierarchical_allreduce,torus_2d,6,16384,32768,524288,11115.5449999995
lrab_hierarchical_allreduce,torus_2d,6,32768,65536,1048576,19865.265000000378 lrab_hierarchical_allreduce,torus_2d,6,32768,65536,1048576,20090.265000000378
lrab_hierarchical_allreduce,torus_2d,6,49152,98304,1572864,28839.98500000059 lrab_hierarchical_allreduce,torus_2d,6,49152,98304,1572864,29064.985000000597
1 algorithm sip_topology n_sips n_elem bytes_per_pe bytes_per_sip latency_ns
2 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 8 16 256 2666.552500000015 2918.5525000000157
3 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 32 64 1024 2747.7400000000152 2999.740000000016
4 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 64 128 2048 2855.990000000018 3107.990000000019
5 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 128 256 4096 3072.490000000019 3324.4900000000207
6 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 512 1024 16384 3337.1133333333582 3589.1133333333582
7 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 1024 2048 32768 3708.0333333333692 3960.0333333333692
8 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 2048 4096 65536 4449.873333333393 4701.873333333393
9 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 4096 8192 131072 5933.020000000124 6185.020000000124
10 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 8192 16384 262144 8900.379999999863 9152.379999999861
11 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 16384 32768 524288 14835.099999999224 15087.099999999224
12 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 32768 65536 1048576 26704.540000000765 26956.540000000765
13 lrab_hierarchical_allreduce mesh_2d_no_wrap 6 49152 98304 1572864 38573.97999999701 38825.97999999701
14 lrab_hierarchical_allreduce ring_1d 6 8 16 256 2365.255833333347 2628.2558333333477
15 lrab_hierarchical_allreduce ring_1d 6 32 64 1024 2436.9433333333473 2699.943333333348
16 lrab_hierarchical_allreduce ring_1d 6 64 128 2048 2532.526666666683 2795.526666666683
17 lrab_hierarchical_allreduce ring_1d 6 128 256 4096 2723.693333333349 2986.693333333351
18 lrab_hierarchical_allreduce ring_1d 6 512 1024 16384 3048.635000000021 3311.635000000021
19 lrab_hierarchical_allreduce ring_1d 6 1024 2048 32768 3393.4016666666957 3656.4016666666957
20 lrab_hierarchical_allreduce ring_1d 6 2048 4096 65536 4082.401666666714 4345.401666666714
21 lrab_hierarchical_allreduce ring_1d 6 4096 8192 131072 5458.80166666677 5721.801666666768
22 lrab_hierarchical_allreduce ring_1d 6 8192 16384 262144 8216.934999999943 8479.934999999887
23 lrab_hierarchical_allreduce ring_1d 6 16384 32768 524288 13733.201666665835 13996.201666665835
24 lrab_hierarchical_allreduce ring_1d 6 32768 65536 1048576 24765.73500000064 25028.73500000064
25 lrab_hierarchical_allreduce ring_1d 6 49152 98304 1572864 35798.268333331536 36061.26833333154
26 lrab_hierarchical_allreduce torus_2d 6 8 16 256 1700.6025000000095 1925.6025000000104
27 lrab_hierarchical_allreduce torus_2d 6 32 64 1024 1753.2900000000102 1978.290000000011
28 lrab_hierarchical_allreduce torus_2d 6 64 128 2048 1823.540000000012 2048.540000000012
29 lrab_hierarchical_allreduce torus_2d 6 128 256 4096 1964.040000000012 2189.040000000012
30 lrab_hierarchical_allreduce torus_2d 6 512 1024 16384 2196.8183333333463 2421.8183333333463
31 lrab_hierarchical_allreduce torus_2d 6 1024 2048 32768 2477.2783333333473 2702.2783333333473
32 lrab_hierarchical_allreduce torus_2d 6 2048 4096 65536 3038.1983333333583 3263.1983333333583
33 lrab_hierarchical_allreduce torus_2d 6 4096 8192 131072 4159.5050000000665 4384.5050000000665
34 lrab_hierarchical_allreduce torus_2d 6 8192 16384 262144 6403.185000000109 6628.18500000005
35 lrab_hierarchical_allreduce torus_2d 6 16384 32768 524288 10890.5449999995 11115.5449999995
36 lrab_hierarchical_allreduce torus_2d 6 32768 65536 1048576 19865.265000000378 20090.265000000378
37 lrab_hierarchical_allreduce torus_2d 6 49152 98304 1572864 28839.98500000059 29064.985000000597
docs/diagrams/allreduce_latency_plots/topology.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 194 KiB

After

Width:  |  Height:  |  Size: 194 KiB

docs/diagrams/gemm_plots/gemm_mac_utilization_measured.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 40 KiB

After

Width:  |  Height:  |  Size: 29 KiB

docs/diagrams/gemm_plots/gemm_mac_utilization_theoretical_vs_measured.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 46 KiB

After

Width:  |  Height:  |  Size: 32 KiB

docs/diagrams/gemm_plots/gemm_stage_breakdown.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 42 KiB

After

Width:  |  Height:  |  Size: 31 KiB

docs/diagrams/pe2pe_latency_plots/latency_intercube_C0PE0_to_C1PE0_horizontal.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 53 KiB

After

Width:  |  Height:  |  Size: 52 KiB

docs/diagrams/pe2pe_latency_plots/latency_intercube_C0PE0_to_C4PE0_vertical.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 52 KiB

After

Width:  |  Height:  |  Size: 45 KiB

docs/diagrams/pe2pe_latency_plots/latency_intracube_PE0_to_PE1_horizontal.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 51 KiB

After

Width:  |  Height:  |  Size: 53 KiB

docs/diagrams/pe2pe_latency_plots/latency_intracube_PE0_to_PE4_vertical.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 52 KiB

After

Width:  |  Height:  |  Size: 47 KiB

docs/diagrams/pe2pe_latency_plots/overview.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 137 KiB

After

Width:  |  Height:  |  Size: 137 KiB

docs/diagrams/pe2pe_latency_plots/summary.csv
+80 -80
View File
@@ -1,81 +1,81 @@
hop,label,size_bytes,path,total_ns hop,label,size_bytes,path,total_ns
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),128,ipcq,24.88749999999891 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),128,ipcq,42.88749999999891
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),128,raw,33.57999999999811 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),128,raw,51.57999999999811
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),256,ipcq,28.13749999999891 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),256,ipcq,46.13749999999891
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),256,raw,36.07999999999811 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),256,raw,54.07999999999811
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),384,ipcq,29.88749999999891 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),384,ipcq,47.88749999999891
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),384,raw,37.07999999999811 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),384,raw,55.07999999999811
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),512,ipcq,31.63749999999891 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),512,ipcq,49.63749999999891
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),512,raw,38.07999999999811 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),512,raw,56.07999999999811
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),768,ipcq,35.13749999999891 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),768,ipcq,53.13749999999891
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),768,raw,40.07999999999811 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),768,raw,58.07999999999811
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),1024,ipcq,38.63749999999891 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),1024,ipcq,56.63749999999891
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),1024,raw,42.07999999999811 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),1024,raw,60.07999999999811
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),2048,ipcq,52.63749999999891 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),2048,ipcq,70.63749999999891
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),2048,raw,50.07999999999811 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),2048,raw,68.07999999999811
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),4096,ipcq,80.63750000000073 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),4096,ipcq,98.63750000000073
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),4096,raw,66.08000000000175 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),4096,raw,84.08000000000175
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),8192,ipcq,136.63750000000073 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),8192,ipcq,154.63750000000073
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),8192,raw,98.08000000000175 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),8192,raw,116.08000000000175
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),10240,ipcq,164.63750000000073 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),10240,ipcq,182.63750000000073
latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),10240,raw,114.08000000000175 latency_intracube_PE0_to_PE1_horizontal,Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal),10240,raw,132.08000000000175
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),128,ipcq,38.49749999999585 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),128,ipcq,56.49749999999585
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),128,raw,47.18999999999505 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),128,raw,65.18999999999505
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),256,ipcq,43.24749999999585 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),256,ipcq,61.24749999999585
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),256,raw,51.18999999999505 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),256,raw,69.18999999999505
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),384,ipcq,44.99749999999585 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),384,ipcq,62.99749999999585
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),384,raw,52.18999999999505 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),384,raw,70.18999999999505
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),512,ipcq,46.74749999999585 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),512,ipcq,64.74749999999585
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),512,raw,53.18999999999505 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),512,raw,71.18999999999505
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),768,ipcq,50.24749999999585 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),768,ipcq,68.24749999999585
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),768,raw,55.18999999999505 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),768,raw,73.18999999999505
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),1024,ipcq,53.74749999999585 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),1024,ipcq,71.74749999999585
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),1024,raw,57.18999999999505 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),1024,raw,75.18999999999505
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),2048,ipcq,67.74749999999585 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),2048,ipcq,85.74749999999585
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),2048,raw,65.18999999999505 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),2048,raw,83.18999999999505
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),4096,ipcq,95.74750000000131 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),4096,ipcq,113.74750000000131
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),4096,raw,81.19000000000233 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),4096,raw,99.19000000000233
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),8192,ipcq,151.7475000000013 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),8192,ipcq,169.7475000000013
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),8192,raw,113.19000000000233 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),8192,raw,131.19000000000233
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),10240,ipcq,179.7475000000013 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),10240,ipcq,197.7475000000013
latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),10240,raw,129.19000000000233 latency_intracube_PE0_to_PE4_vertical,Intra-cube PE-to-PE latency: PE0 → PE4 (vertical),10240,raw,147.19000000000233
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),128,ipcq,81.15999999999804 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),128,ipcq,99.15999999999804
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),128,raw,89.28999999999724 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),128,raw,107.28999999999724
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),256,ipcq,88.65999999999804 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),256,ipcq,106.65999999999804
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),256,raw,95.53999999999724 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),256,raw,113.53999999999724
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),384,ipcq,90.90999999999804 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),384,ipcq,108.90999999999804
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),384,raw,96.53999999999724 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),384,raw,114.53999999999724
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),512,ipcq,93.15999999999804 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),512,ipcq,111.15999999999804
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),512,raw,97.53999999999724 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),512,raw,115.53999999999724
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),768,ipcq,97.65999999999804 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),768,ipcq,115.65999999999804
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),768,raw,99.53999999999724 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),768,raw,117.53999999999724
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),1024,ipcq,103.15999999999804 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),1024,ipcq,121.15999999999804
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),1024,raw,102.53999999999724 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),1024,raw,120.53999999999724
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),2048,ipcq,125.15999999999804 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),2048,ipcq,143.15999999999804
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),2048,raw,114.53999999999724 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),2048,raw,132.53999999999724
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),4096,ipcq,169.15999999999985 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),4096,ipcq,187.15999999999985
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),4096,raw,138.54000000000087 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),4096,raw,156.54000000000087
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),8192,ipcq,257.15999999999985 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),8192,ipcq,275.15999999999985
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),8192,raw,186.54000000000087 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),8192,raw,204.54000000000087
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),10240,ipcq,301.15999999999985 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),10240,ipcq,319.15999999999985
latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),10240,raw,210.54000000000087 latency_intercube_C0PE0_to_C1PE0_horizontal,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal),10240,raw,228.54000000000087
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),128,ipcq,103.15999999999804 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),128,ipcq,121.15999999999804
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),128,raw,111.28999999999724 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),128,raw,129.28999999999724
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),256,ipcq,112.65999999999804 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),256,ipcq,130.65999999999804
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),256,raw,119.53999999999724 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),256,raw,137.53999999999724
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),384,ipcq,114.90999999999804 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),384,ipcq,132.90999999999804
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),384,raw,120.53999999999724 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),384,raw,138.53999999999724
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),512,ipcq,117.15999999999804 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),512,ipcq,135.15999999999804
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),512,raw,121.53999999999724 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),512,raw,139.53999999999724
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),768,ipcq,121.65999999999804 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),768,ipcq,139.65999999999804
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),768,raw,123.53999999999724 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),768,raw,141.53999999999724
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),1024,ipcq,127.15999999999804 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),1024,ipcq,145.15999999999804
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),1024,raw,126.53999999999724 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),1024,raw,144.53999999999724
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),2048,ipcq,149.15999999999804 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),2048,ipcq,167.15999999999804
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),2048,raw,138.53999999999724 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),2048,raw,156.53999999999724
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),4096,ipcq,193.15999999999985 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),4096,ipcq,211.15999999999985
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),4096,raw,162.54000000000087 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),4096,raw,180.54000000000087
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),8192,ipcq,281.15999999999985 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),8192,ipcq,299.15999999999985
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),8192,raw,210.54000000000087 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),8192,raw,228.54000000000087
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),10240,ipcq,325.15999999999985 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),10240,ipcq,343.15999999999985
latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),10240,raw,234.54000000000087 latency_intercube_C0PE0_to_C4PE0_vertical,Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical),10240,raw,252.54000000000087
1 hop label size_bytes path total_ns
2 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 128 ipcq 24.88749999999891 42.88749999999891
3 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 128 raw 33.57999999999811 51.57999999999811
4 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 256 ipcq 28.13749999999891 46.13749999999891
5 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 256 raw 36.07999999999811 54.07999999999811
6 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 384 ipcq 29.88749999999891 47.88749999999891
7 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 384 raw 37.07999999999811 55.07999999999811
8 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 512 ipcq 31.63749999999891 49.63749999999891
9 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 512 raw 38.07999999999811 56.07999999999811
10 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 768 ipcq 35.13749999999891 53.13749999999891
11 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 768 raw 40.07999999999811 58.07999999999811
12 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 1024 ipcq 38.63749999999891 56.63749999999891
13 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 1024 raw 42.07999999999811 60.07999999999811
14 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 2048 ipcq 52.63749999999891 70.63749999999891
15 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 2048 raw 50.07999999999811 68.07999999999811
16 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 4096 ipcq 80.63750000000073 98.63750000000073
17 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 4096 raw 66.08000000000175 84.08000000000175
18 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 8192 ipcq 136.63750000000073 154.63750000000073
19 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 8192 raw 98.08000000000175 116.08000000000175
20 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 10240 ipcq 164.63750000000073 182.63750000000073
21 latency_intracube_PE0_to_PE1_horizontal Intra-cube PE-to-PE latency: PE0 → PE1 (horizontal) 10240 raw 114.08000000000175 132.08000000000175
22 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 128 ipcq 38.49749999999585 56.49749999999585
23 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 128 raw 47.18999999999505 65.18999999999505
24 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 256 ipcq 43.24749999999585 61.24749999999585
25 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 256 raw 51.18999999999505 69.18999999999505
26 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 384 ipcq 44.99749999999585 62.99749999999585
27 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 384 raw 52.18999999999505 70.18999999999505
28 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 512 ipcq 46.74749999999585 64.74749999999585
29 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 512 raw 53.18999999999505 71.18999999999505
30 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 768 ipcq 50.24749999999585 68.24749999999585
31 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 768 raw 55.18999999999505 73.18999999999505
32 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 1024 ipcq 53.74749999999585 71.74749999999585
33 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 1024 raw 57.18999999999505 75.18999999999505
34 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 2048 ipcq 67.74749999999585 85.74749999999585
35 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 2048 raw 65.18999999999505 83.18999999999505
36 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 4096 ipcq 95.74750000000131 113.74750000000131
37 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 4096 raw 81.19000000000233 99.19000000000233
38 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 8192 ipcq 151.7475000000013 169.7475000000013
39 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 8192 raw 113.19000000000233 131.19000000000233
40 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 10240 ipcq 179.7475000000013 197.7475000000013
41 latency_intracube_PE0_to_PE4_vertical Intra-cube PE-to-PE latency: PE0 → PE4 (vertical) 10240 raw 129.19000000000233 147.19000000000233
42 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 128 ipcq 81.15999999999804 99.15999999999804
43 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 128 raw 89.28999999999724 107.28999999999724
44 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 256 ipcq 88.65999999999804 106.65999999999804
45 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 256 raw 95.53999999999724 113.53999999999724
46 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 384 ipcq 90.90999999999804 108.90999999999804
47 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 384 raw 96.53999999999724 114.53999999999724
48 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 512 ipcq 93.15999999999804 111.15999999999804
49 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 512 raw 97.53999999999724 115.53999999999724
50 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 768 ipcq 97.65999999999804 115.65999999999804
51 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 768 raw 99.53999999999724 117.53999999999724
52 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 1024 ipcq 103.15999999999804 121.15999999999804
53 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 1024 raw 102.53999999999724 120.53999999999724
54 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 2048 ipcq 125.15999999999804 143.15999999999804
55 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 2048 raw 114.53999999999724 132.53999999999724
56 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 4096 ipcq 169.15999999999985 187.15999999999985
57 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 4096 raw 138.54000000000087 156.54000000000087
58 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 8192 ipcq 257.15999999999985 275.15999999999985
59 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 8192 raw 186.54000000000087 204.54000000000087
60 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 10240 ipcq 301.15999999999985 319.15999999999985
61 latency_intercube_C0PE0_to_C1PE0_horizontal Inter-cube PE-to-PE latency: Cube0.PE0 → Cube1.PE0 (horizontal) 10240 raw 210.54000000000087 228.54000000000087
62 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 128 ipcq 103.15999999999804 121.15999999999804
63 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 128 raw 111.28999999999724 129.28999999999724
64 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 256 ipcq 112.65999999999804 130.65999999999804
65 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 256 raw 119.53999999999724 137.53999999999724
66 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 384 ipcq 114.90999999999804 132.90999999999804
67 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 384 raw 120.53999999999724 138.53999999999724
68 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 512 ipcq 117.15999999999804 135.15999999999804
69 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 512 raw 121.53999999999724 139.53999999999724
70 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 768 ipcq 121.65999999999804 139.65999999999804
71 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 768 raw 123.53999999999724 141.53999999999724
72 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 1024 ipcq 127.15999999999804 145.15999999999804
73 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 1024 raw 126.53999999999724 144.53999999999724
74 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 2048 ipcq 149.15999999999804 167.15999999999804
75 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 2048 raw 138.53999999999724 156.53999999999724
76 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 4096 ipcq 193.15999999999985 211.15999999999985
77 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 4096 raw 162.54000000000087 180.54000000000087
78 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 8192 ipcq 281.15999999999985 299.15999999999985
79 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 8192 raw 210.54000000000087 228.54000000000087
80 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 10240 ipcq 325.15999999999985 343.15999999999985
81 latency_intercube_C0PE0_to_C4PE0_vertical Inter-cube PE-to-PE latency: Cube0.PE0 → Cube4.PE0 (vertical) 10240 raw 234.54000000000087 252.54000000000087
docs/report/1H-codesign-paper/.gitignore
+2
View File
@@ -0,0 +1,2 @@
# Downloaded Tectonic toolchain binary (large, machine-local; re-fetched by /paper)
build/.tools/
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
+26909
View File
File diff suppressed because one or more lines are too long
docs/report/1H-codesign-paper/figures/allreduce_buffer_kind.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 77 KiB

docs/report/1H-codesign-paper/figures/allreduce_comparison.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 92 KiB

docs/report/1H-codesign-paper/figures/allreduce_topology.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 194 KiB

docs/report/1H-codesign-paper/figures/cube_architecture.pdf
BIN
View File
Binary file not shown.
docs/report/1H-codesign-paper/figures/cube_architecture.svg
+312
View File
@@ -0,0 +1,312 @@
<svg xmlns="http://www.w3.org/2000/svg" width="860" height="798" viewBox="55 2 860 798">
<title>cube</title>
<rect width="970" height="900" fill="#ffffff"/>
<text x="485" y="22" text-anchor="middle" font-family="monospace" font-size="18" font-weight="bold" fill="#1f2937">CUBE TOPOLOGY — 17.0×14.0mm | 6×6 Router Mesh | n_to_one mode | 64 pseudo-ch</text>
<text x="485" y="40" text-anchor="middle" font-family="monospace" font-size="15" fill="#ffffff">Per-PE: 8 ch × 32.0 GB/s = 256.0 GB/s | Cube total: 64 × 32.0 = 2048.0 GB/s</text>
<rect x="60" y="60" width="850.0" height="700.0" rx="6" fill="none" stroke="#94a3b8" stroke-width="2" stroke-dasharray="8,4"/>
<rect x="260" y="285" width="450" height="250" rx="6" fill="#ecfdf5" stroke="#047857" stroke-width="2" opacity="0.6"/>
<text x="485" y="395" text-anchor="middle" font-family="monospace" font-size="16" font-weight="bold" fill="#047857">HBM_CTRL | 64 pseudo channels</text>
<text x="485" y="412" text-anchor="middle" font-family="monospace" font-size="14" fill="#059669">Total BW: 2048 GB/s</text>
<rect x="270.0" y="289" width="12.9" height="8" rx="1" fill="#3b82f6" opacity="0.8"/>
<rect x="283.4" y="289" width="12.9" height="8" rx="1" fill="#3b82f6" opacity="0.8"/>
<rect x="296.9" y="289" width="12.9" height="8" rx="1" fill="#3b82f6" opacity="0.8"/>
<rect x="310.3" y="289" width="12.9" height="8" rx="1" fill="#3b82f6" opacity="0.8"/>
<rect x="323.8" y="289" width="12.9" height="8" rx="1" fill="#3b82f6" opacity="0.8"/>
<rect x="337.2" y="289" width="12.9" height="8" rx="1" fill="#3b82f6" opacity="0.8"/>
<rect x="350.6" y="289" width="12.9" height="8" rx="1" fill="#3b82f6" opacity="0.8"/>
<rect x="364.1" y="289" width="12.9" height="8" rx="1" fill="#3b82f6" opacity="0.8"/>
<rect x="377.5" y="289" width="12.9" height="8" rx="1" fill="#60a5fa" opacity="0.8"/>
<rect x="390.9" y="289" width="12.9" height="8" rx="1" fill="#60a5fa" opacity="0.8"/>
<rect x="404.4" y="289" width="12.9" height="8" rx="1" fill="#60a5fa" opacity="0.8"/>
<rect x="417.8" y="289" width="12.9" height="8" rx="1" fill="#60a5fa" opacity="0.8"/>
<rect x="431.2" y="289" width="12.9" height="8" rx="1" fill="#60a5fa" opacity="0.8"/>
<rect x="444.7" y="289" width="12.9" height="8" rx="1" fill="#60a5fa" opacity="0.8"/>
<rect x="458.1" y="289" width="12.9" height="8" rx="1" fill="#60a5fa" opacity="0.8"/>
<rect x="471.6" y="289" width="12.9" height="8" rx="1" fill="#60a5fa" opacity="0.8"/>
<rect x="485.0" y="289" width="12.9" height="8" rx="1" fill="#475569" opacity="0.8"/>
<rect x="498.4" y="289" width="12.9" height="8" rx="1" fill="#475569" opacity="0.8"/>
<rect x="511.9" y="289" width="12.9" height="8" rx="1" fill="#475569" opacity="0.8"/>
<rect x="525.3" y="289" width="12.9" height="8" rx="1" fill="#475569" opacity="0.8"/>
<rect x="538.8" y="289" width="12.9" height="8" rx="1" fill="#475569" opacity="0.8"/>
<rect x="552.2" y="289" width="12.9" height="8" rx="1" fill="#475569" opacity="0.8"/>
<rect x="565.6" y="289" width="12.9" height="8" rx="1" fill="#475569" opacity="0.8"/>
<rect x="579.1" y="289" width="12.9" height="8" rx="1" fill="#475569" opacity="0.8"/>
<rect x="592.5" y="289" width="12.9" height="8" rx="1" fill="#94a3b8" opacity="0.8"/>
<rect x="605.9" y="289" width="12.9" height="8" rx="1" fill="#94a3b8" opacity="0.8"/>
<rect x="619.4" y="289" width="12.9" height="8" rx="1" fill="#94a3b8" opacity="0.8"/>
<rect x="632.8" y="289" width="12.9" height="8" rx="1" fill="#94a3b8" opacity="0.8"/>
<rect x="646.2" y="289" width="12.9" height="8" rx="1" fill="#94a3b8" opacity="0.8"/>
<rect x="659.7" y="289" width="12.9" height="8" rx="1" fill="#94a3b8" opacity="0.8"/>
<rect x="673.1" y="289" width="12.9" height="8" rx="1" fill="#94a3b8" opacity="0.8"/>
<rect x="686.6" y="289" width="12.9" height="8" rx="1" fill="#94a3b8" opacity="0.8"/>
<text x="324" y="286" text-anchor="middle" font-family="monospace" font-size="12" fill="#3b82f6">PE0×8ch</text>
<text x="431" y="286" text-anchor="middle" font-family="monospace" font-size="12" fill="#60a5fa">PE1×8ch</text>
<text x="539" y="286" text-anchor="middle" font-family="monospace" font-size="12" fill="#475569">PE2×8ch</text>
<text x="646" y="286" text-anchor="middle" font-family="monospace" font-size="12" fill="#94a3b8">PE3×8ch</text>
<rect x="270.0" y="523" width="12.9" height="8" rx="1" fill="#f59e0b" opacity="0.8"/>
<rect x="283.4" y="523" width="12.9" height="8" rx="1" fill="#f59e0b" opacity="0.8"/>
<rect x="296.9" y="523" width="12.9" height="8" rx="1" fill="#f59e0b" opacity="0.8"/>
<rect x="310.3" y="523" width="12.9" height="8" rx="1" fill="#f59e0b" opacity="0.8"/>
<rect x="323.8" y="523" width="12.9" height="8" rx="1" fill="#f59e0b" opacity="0.8"/>
<rect x="337.2" y="523" width="12.9" height="8" rx="1" fill="#f59e0b" opacity="0.8"/>
<rect x="350.6" y="523" width="12.9" height="8" rx="1" fill="#f59e0b" opacity="0.8"/>
<rect x="364.1" y="523" width="12.9" height="8" rx="1" fill="#f59e0b" opacity="0.8"/>
<rect x="377.5" y="523" width="12.9" height="8" rx="1" fill="#fbbf24" opacity="0.8"/>
<rect x="390.9" y="523" width="12.9" height="8" rx="1" fill="#fbbf24" opacity="0.8"/>
<rect x="404.4" y="523" width="12.9" height="8" rx="1" fill="#fbbf24" opacity="0.8"/>
<rect x="417.8" y="523" width="12.9" height="8" rx="1" fill="#fbbf24" opacity="0.8"/>
<rect x="431.2" y="523" width="12.9" height="8" rx="1" fill="#fbbf24" opacity="0.8"/>
<rect x="444.7" y="523" width="12.9" height="8" rx="1" fill="#fbbf24" opacity="0.8"/>
<rect x="458.1" y="523" width="12.9" height="8" rx="1" fill="#fbbf24" opacity="0.8"/>
<rect x="471.6" y="523" width="12.9" height="8" rx="1" fill="#fbbf24" opacity="0.8"/>
<rect x="485.0" y="523" width="12.9" height="8" rx="1" fill="#ef4444" opacity="0.8"/>
<rect x="498.4" y="523" width="12.9" height="8" rx="1" fill="#ef4444" opacity="0.8"/>
<rect x="511.9" y="523" width="12.9" height="8" rx="1" fill="#ef4444" opacity="0.8"/>
<rect x="525.3" y="523" width="12.9" height="8" rx="1" fill="#ef4444" opacity="0.8"/>
<rect x="538.8" y="523" width="12.9" height="8" rx="1" fill="#ef4444" opacity="0.8"/>
<rect x="552.2" y="523" width="12.9" height="8" rx="1" fill="#ef4444" opacity="0.8"/>
<rect x="565.6" y="523" width="12.9" height="8" rx="1" fill="#ef4444" opacity="0.8"/>
<rect x="579.1" y="523" width="12.9" height="8" rx="1" fill="#ef4444" opacity="0.8"/>
<rect x="592.5" y="523" width="12.9" height="8" rx="1" fill="#f87171" opacity="0.8"/>
<rect x="605.9" y="523" width="12.9" height="8" rx="1" fill="#f87171" opacity="0.8"/>
<rect x="619.4" y="523" width="12.9" height="8" rx="1" fill="#f87171" opacity="0.8"/>
<rect x="632.8" y="523" width="12.9" height="8" rx="1" fill="#f87171" opacity="0.8"/>
<rect x="646.2" y="523" width="12.9" height="8" rx="1" fill="#f87171" opacity="0.8"/>
<rect x="659.7" y="523" width="12.9" height="8" rx="1" fill="#f87171" opacity="0.8"/>
<rect x="673.1" y="523" width="12.9" height="8" rx="1" fill="#f87171" opacity="0.8"/>
<rect x="686.6" y="523" width="12.9" height="8" rx="1" fill="#f87171" opacity="0.8"/>
<text x="324" y="539" text-anchor="middle" font-family="monospace" font-size="12" fill="#f59e0b">PE4×8ch</text>
<text x="431" y="539" text-anchor="middle" font-family="monospace" font-size="12" fill="#fbbf24">PE5×8ch</text>
<text x="539" y="539" text-anchor="middle" font-family="monospace" font-size="12" fill="#ef4444">PE6×8ch</text>
<text x="646" y="539" text-anchor="middle" font-family="monospace" font-size="12" fill="#f87171">PE7×8ch</text>
<line x1="135" y1="135" x2="285" y2="135" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="135" y1="135" x2="135" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="135" x2="435" y2="135" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="135" x2="285" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="435" y1="135" x2="585" y2="135" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="435" y1="135" x2="435" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="585" y1="135" x2="685" y2="135" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="585" y1="135" x2="585" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="135" x2="835" y2="135" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="135" x2="685" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="835" y1="135" x2="835" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="135" y1="260" x2="285" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="135" y1="260" x2="135" y2="335" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="260" x2="435" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="260" x2="285" y2="335" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="435" y1="260" x2="585" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="435" y1="260" x2="435" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="585" y1="260" x2="685" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="585" y1="260" x2="585" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="260" x2="835" y2="260" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="260" x2="685" y2="335" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="835" y1="260" x2="835" y2="335" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="135" y1="335" x2="285" y2="335" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="135" y1="335" x2="135" y2="485" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="335" x2="685" y2="335" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="335" x2="285" y2="485" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="335" x2="835" y2="335" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="335" x2="685" y2="485" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="835" y1="335" x2="835" y2="485" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="135" y1="485" x2="285" y2="485" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="135" y1="485" x2="135" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="485" x2="685" y2="485" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="485" x2="285" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="485" x2="835" y2="485" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="485" x2="685" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="835" y1="485" x2="835" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="135" y1="560" x2="285" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="135" y1="560" x2="135" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="560" x2="435" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="560" x2="285" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="435" y1="560" x2="585" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="435" y1="560" x2="435" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="585" y1="560" x2="685" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="585" y1="560" x2="585" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="560" x2="835" y2="560" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="560" x2="685" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="835" y1="560" x2="835" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="135" y1="685" x2="285" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="285" y1="685" x2="435" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="435" y1="685" x2="585" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="585" y1="685" x2="685" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<line x1="685" y1="685" x2="835" y2="685" stroke="#94a3b8" stroke-width="1" opacity="0.4"/>
<circle cx="135" cy="135" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="135" y="138" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r0c0</text>
<rect x="119" y="81" width="32" height="16" rx="3" fill="#ffffff" stroke="#a855f7" stroke-width="1"/>
<text x="135" y="92" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#a855f7">PE0</text>
<line x1="135" y1="127" x2="149" y2="97" stroke="#a855f7" stroke-width="1" opacity="0.6"/>
<circle cx="285" cy="135" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="285" y="138" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r0c1</text>
<rect x="269" y="81" width="32" height="16" rx="3" fill="#ffffff" stroke="#a855f7" stroke-width="1"/>
<text x="285" y="92" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#a855f7">PE1</text>
<line x1="285" y1="127" x2="299" y2="97" stroke="#a855f7" stroke-width="1" opacity="0.6"/>
<circle cx="435" cy="135" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="435" y="138" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r0c2</text>
<circle cx="585" cy="135" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="585" y="138" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r0c3</text>
<circle cx="685" cy="135" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="685" y="138" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r0c4</text>
<circle cx="835" cy="135" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="835" y="138" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r0c5</text>
<circle cx="135" cy="260" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="135" y="263" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r1c0</text>
<circle cx="285" cy="260" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="285" y="263" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r1c1</text>
<circle cx="435" cy="260" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="435" y="263" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r1c2</text>
<rect x="419" y="206" width="32" height="16" rx="3" fill="#ffffff" stroke="#f59e0b" stroke-width="1"/>
<text x="435" y="217" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#f59e0b">M_CPU</text>
<line x1="435" y1="252" x2="449" y2="222" stroke="#f59e0b" stroke-width="1" opacity="0.6"/>
<circle cx="585" cy="260" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="585" y="263" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r1c3</text>
<circle cx="685" cy="260" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="685" y="263" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r1c4</text>
<rect x="669" y="206" width="32" height="16" rx="3" fill="#ffffff" stroke="#a855f7" stroke-width="1"/>
<text x="685" y="217" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#a855f7">PE2</text>
<line x1="685" y1="252" x2="699" y2="222" stroke="#a855f7" stroke-width="1" opacity="0.6"/>
<circle cx="835" cy="260" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="835" y="263" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r1c5</text>
<rect x="819" y="206" width="32" height="16" rx="3" fill="#ffffff" stroke="#a855f7" stroke-width="1"/>
<text x="835" y="217" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#a855f7">PE3</text>
<line x1="835" y1="252" x2="849" y2="222" stroke="#a855f7" stroke-width="1" opacity="0.6"/>
<circle cx="135" cy="335" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="135" y="338" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r2c0</text>
<circle cx="285" cy="335" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="285" y="338" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r2c1</text>
<circle cx="685" cy="335" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="685" y="338" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r2c4</text>
<circle cx="835" cy="335" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="835" y="338" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r2c5</text>
<circle cx="135" cy="485" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="135" y="488" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r3c0</text>
<rect x="119" y="523" width="32" height="16" rx="3" fill="#ffffff" stroke="#d97706" stroke-width="1"/>
<text x="135" y="534" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#d97706">SRAM</text>
<line x1="135" y1="493" x2="149" y2="523" stroke="#d97706" stroke-width="1" opacity="0.6"/>
<circle cx="285" cy="485" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="285" y="488" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r3c1</text>
<circle cx="685" cy="485" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="685" y="488" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r3c4</text>
<circle cx="835" cy="485" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="835" y="488" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r3c5</text>
<circle cx="135" cy="560" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="135" y="563" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r4c0</text>
<rect x="119" y="598" width="32" height="16" rx="3" fill="#ffffff" stroke="#a855f7" stroke-width="1"/>
<text x="135" y="609" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#a855f7">PE4</text>
<line x1="135" y1="568" x2="149" y2="598" stroke="#a855f7" stroke-width="1" opacity="0.6"/>
<circle cx="285" cy="560" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="285" y="563" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r4c1</text>
<rect x="269" y="598" width="32" height="16" rx="3" fill="#ffffff" stroke="#a855f7" stroke-width="1"/>
<text x="285" y="609" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#a855f7">PE5</text>
<line x1="285" y1="568" x2="299" y2="598" stroke="#a855f7" stroke-width="1" opacity="0.6"/>
<circle cx="435" cy="560" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="435" y="563" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r4c2</text>
<circle cx="585" cy="560" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="585" y="563" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r4c3</text>
<circle cx="685" cy="560" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="685" y="563" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r4c4</text>
<circle cx="835" cy="560" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="835" y="563" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r4c5</text>
<circle cx="135" cy="685" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="135" y="688" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r5c0</text>
<circle cx="285" cy="685" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="285" y="688" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r5c1</text>
<circle cx="435" cy="685" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="435" y="688" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r5c2</text>
<circle cx="585" cy="685" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="585" y="688" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r5c3</text>
<circle cx="685" cy="685" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="685" y="688" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r5c4</text>
<rect x="669" y="723" width="32" height="16" rx="3" fill="#ffffff" stroke="#a855f7" stroke-width="1"/>
<text x="685" y="734" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#a855f7">PE6</text>
<line x1="685" y1="693" x2="699" y2="723" stroke="#a855f7" stroke-width="1" opacity="0.6"/>
<circle cx="835" cy="685" r="17" fill="#ffffff" stroke="#94a3b8" stroke-width="1"/>
<text x="835" y="688" text-anchor="middle" font-family="monospace" font-size="12" fill="#1f2937">r5c5</text>
<rect x="819" y="723" width="32" height="16" rx="3" fill="#ffffff" stroke="#a855f7" stroke-width="1"/>
<text x="835" y="734" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#a855f7">PE7</text>
<line x1="835" y1="693" x2="849" y2="723" stroke="#a855f7" stroke-width="1" opacity="0.6"/>
<polyline points="135,143 208,216 251,216 324,289" fill="none" stroke="#10b981" stroke-width="1.5" opacity="0.6" stroke-dasharray="4,3"/>
<text x="239" y="216" font-family="monospace" font-size="12" fill="#047857">256GB/s</text>
<polyline points="285,143 358,216 358,216 431,289" fill="none" stroke="#10b981" stroke-width="1.5" opacity="0.6" stroke-dasharray="4,3"/>
<text x="368" y="216" font-family="monospace" font-size="12" fill="#047857">256GB/s</text>
<polyline points="685,268 674,278 549,278 539,289" fill="none" stroke="#10b981" stroke-width="1.5" opacity="0.6" stroke-dasharray="4,3"/>
<text x="622" y="278" font-family="monospace" font-size="12" fill="#047857">256GB/s</text>
<polyline points="835,268 824,278 657,278 646,289" fill="none" stroke="#10b981" stroke-width="1.5" opacity="0.6" stroke-dasharray="4,3"/>
<text x="751" y="278" font-family="monospace" font-size="12" fill="#047857">256GB/s</text>
<polyline points="135,552 146,542 313,542 324,531" fill="none" stroke="#10b981" stroke-width="1.5" opacity="0.6" stroke-dasharray="4,3"/>
<text x="239" y="542" font-family="monospace" font-size="12" fill="#047857">256GB/s</text>
<polyline points="285,552 296,542 421,542 431,531" fill="none" stroke="#10b981" stroke-width="1.5" opacity="0.6" stroke-dasharray="4,3"/>
<text x="368" y="542" font-family="monospace" font-size="12" fill="#047857">256GB/s</text>
<polyline points="685,677 612,604 612,604 539,531" fill="none" stroke="#10b981" stroke-width="1.5" opacity="0.6" stroke-dasharray="4,3"/>
<text x="622" y="604" font-family="monospace" font-size="12" fill="#047857">256GB/s</text>
<polyline points="835,677 762,604 719,604 646,531" fill="none" stroke="#10b981" stroke-width="1.5" opacity="0.6" stroke-dasharray="4,3"/>
<text x="751" y="604" font-family="monospace" font-size="12" fill="#047857">256GB/s</text>
<rect x="65" y="360" width="50" height="100" rx="3" fill="#ffffff" stroke="#475569" stroke-width="1.5" opacity="0.9"/>
<text x="90" y="357" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#475569">UCIe-W</text>
<rect x="67" y="362" width="46" height="23" rx="2" fill="#cbd5e1" opacity="0.7"/>
<text x="90" y="376" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c0</text>
<polyline points="127,135 120,142 120,366 113,374" fill="none" stroke="#cbd5e1" stroke-width="1" opacity="0.5"/>
<rect x="67" y="386" width="46" height="23" rx="2" fill="#94a3b8" opacity="0.7"/>
<text x="90" y="400" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c1</text>
<polyline points="127,260 120,267 120,390 113,398" fill="none" stroke="#94a3b8" stroke-width="1" opacity="0.5"/>
<rect x="67" y="410" width="46" height="23" rx="2" fill="#64748b" opacity="0.7"/>
<text x="90" y="424" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c2</text>
<polyline points="127,560 120,553 120,428 113,422" fill="none" stroke="#64748b" stroke-width="1" opacity="0.5"/>
<rect x="67" y="434" width="46" height="23" rx="2" fill="#374151" opacity="0.7"/>
<text x="90" y="448" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c3</text>
<polyline points="127,685 120,678 120,452 113,446" fill="none" stroke="#374151" stroke-width="1" opacity="0.5"/>
<rect x="435" y="65" width="100" height="50" rx="3" fill="#ffffff" stroke="#475569" stroke-width="1.5" opacity="0.9"/>
<text x="485" y="62" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#475569">UCIe-N</text>
<rect x="437" y="67" width="23" height="46" rx="2" fill="#cbd5e1" opacity="0.7"/>
<text x="448" y="93" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c0</text>
<polyline points="135,127 142,120 442,120 448,113" fill="none" stroke="#cbd5e1" stroke-width="1" opacity="0.5"/>
<rect x="461" y="67" width="23" height="46" rx="2" fill="#94a3b8" opacity="0.7"/>
<text x="472" y="93" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c1</text>
<polyline points="285,127 292,120 466,120 472,113" fill="none" stroke="#94a3b8" stroke-width="1" opacity="0.5"/>
<rect x="485" y="67" width="23" height="46" rx="2" fill="#64748b" opacity="0.7"/>
<text x="496" y="93" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c2</text>
<polyline points="685,127 678,120 504,120 496,113" fill="none" stroke="#64748b" stroke-width="1" opacity="0.5"/>
<rect x="509" y="67" width="23" height="46" rx="2" fill="#374151" opacity="0.7"/>
<text x="520" y="93" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c3</text>
<polyline points="835,127 828,120 528,120 520,113" fill="none" stroke="#374151" stroke-width="1" opacity="0.5"/>
<rect x="855" y="360" width="50" height="100" rx="3" fill="#ffffff" stroke="#475569" stroke-width="1.5" opacity="0.9"/>
<text x="880" y="357" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#475569">UCIe-E</text>
<rect x="857" y="362" width="46" height="23" rx="2" fill="#cbd5e1" opacity="0.7"/>
<text x="880" y="376" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c0</text>
<polyline points="843,135 850,142 850,367 857,374" fill="none" stroke="#cbd5e1" stroke-width="1" opacity="0.5"/>
<rect x="857" y="386" width="46" height="23" rx="2" fill="#94a3b8" opacity="0.7"/>
<text x="880" y="400" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c1</text>
<polyline points="843,260 850,267 850,391 857,398" fill="none" stroke="#94a3b8" stroke-width="1" opacity="0.5"/>
<rect x="857" y="410" width="46" height="23" rx="2" fill="#64748b" opacity="0.7"/>
<text x="880" y="424" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c2</text>
<polyline points="843,560 850,553 850,428 857,422" fill="none" stroke="#64748b" stroke-width="1" opacity="0.5"/>
<rect x="857" y="434" width="46" height="23" rx="2" fill="#374151" opacity="0.7"/>
<text x="880" y="448" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c3</text>
<polyline points="843,685 850,678 850,452 857,446" fill="none" stroke="#374151" stroke-width="1" opacity="0.5"/>
<rect x="435" y="705" width="100" height="50" rx="3" fill="#ffffff" stroke="#475569" stroke-width="1.5" opacity="0.9"/>
<text x="485" y="702" text-anchor="middle" font-family="monospace" font-size="14" font-weight="bold" fill="#475569">UCIe-S</text>
<rect x="437" y="707" width="23" height="46" rx="2" fill="#cbd5e1" opacity="0.7"/>
<text x="448" y="733" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c0</text>
<polyline points="135,693 142,700 442,700 448,707" fill="none" stroke="#cbd5e1" stroke-width="1" opacity="0.5"/>
<rect x="461" y="707" width="23" height="46" rx="2" fill="#94a3b8" opacity="0.7"/>
<text x="472" y="733" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c1</text>
<polyline points="285,693 292,700 466,700 472,707" fill="none" stroke="#94a3b8" stroke-width="1" opacity="0.5"/>
<rect x="485" y="707" width="23" height="46" rx="2" fill="#64748b" opacity="0.7"/>
<text x="496" y="733" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c2</text>
<polyline points="685,693 678,700 504,700 496,707" fill="none" stroke="#64748b" stroke-width="1" opacity="0.5"/>
<rect x="509" y="707" width="23" height="46" rx="2" fill="#374151" opacity="0.7"/>
<text x="520" y="733" text-anchor="middle" font-family="monospace" font-size="10" fill="#1f2937">c3</text>
<polyline points="835,693 828,700 528,700 520,707" fill="none" stroke="#374151" stroke-width="1" opacity="0.5"/>
<rect x="60" y="775" width="10" height="10" rx="2" fill="#3b82f6" stroke="#94a3b8" stroke-width="0.5"/>
<text x="74" y="784" font-family="monospace" font-size="13" fill="#1f2937">PE Router</text>
<rect x="147" y="775" width="10" height="10" rx="2" fill="#f59e0b" stroke="#94a3b8" stroke-width="0.5"/>
<text x="161" y="784" font-family="monospace" font-size="13" fill="#1f2937">M_CPU / SRAM</text>
<rect x="255" y="775" width="10" height="10" rx="2" fill="#475569" stroke="#94a3b8" stroke-width="0.5"/>
<text x="269" y="784" font-family="monospace" font-size="13" fill="#1f2937">UCIe</text>
<rect x="307" y="775" width="10" height="10" rx="2" fill="#1f2937" stroke="#94a3b8" stroke-width="0.5"/>
<text x="321" y="784" font-family="monospace" font-size="13" fill="#1f2937">Relay</text>
<rect x="366" y="775" width="10" height="10" rx="2" fill="#10b981" stroke="#94a3b8" stroke-width="0.5"/>
<text x="380" y="784" font-family="monospace" font-size="13" fill="#1f2937">HBM Link</text>
<rect x="446" y="775" width="10" height="10" rx="2" fill="#1f2937" stroke="#94a3b8" stroke-width="0.5"/>
<text x="460" y="784" font-family="monospace" font-size="13" fill="#1f2937">Mesh Link</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 30 KiB

docs/report/1H-codesign-paper/figures/gemm_composite_vs_async_tflops.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 75 KiB

docs/report/1H-codesign-paper/figures/gemm_hbm_bw_util.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 32 KiB

docs/report/1H-codesign-paper/figures/gemm_mac_utilization_measured.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 29 KiB

docs/report/1H-codesign-paper/figures/gemm_mac_utilization_theoretical_vs_measured.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 32 KiB

docs/report/1H-codesign-paper/figures/gemm_per_pe_tflops.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 29 KiB

docs/report/1H-codesign-paper/figures/gemm_stage_breakdown.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 34 KiB

docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_6cases_latency.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 58 KiB

docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_6cases_memory.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 51 KiB

docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_6cases_parallelism.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 74 KiB

docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_6cases_traffic.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 45 KiB

docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 149 KiB

docs/report/1H-codesign-paper/figures/gqa_decode_long_ctx_composite_breakdown.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 73 KiB

docs/report/1H-codesign-paper/figures/gqa_decode_models.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 169 KiB

docs/report/1H-codesign-paper/figures/gqa_latency.json
+97
View File
@@ -0,0 +1,97 @@
{
"version": 1,
"panels": [
"single_user_prefill_gqa",
"multi_user_prefill_gqa",
"single_user_decode_gqa",
"multi_user_decode_gqa"
],
"rows": [
{
"panel": "single_user_prefill_gqa",
"kind": "prefill",
"C": 1,
"S_kv": 16,
"latency_ns": 445.1180000000004,
"op_log_summary": {
"gemm_count": 2,
"ipcq_copy_count": 0,
"dma_read_count": 3,
"dma_write_count": 1
},
"engine_occupancy_ns": {
"pe_gemm": 2.048000000000002,
"pe_math": 5.0,
"pe_dma": 72.0,
"pe_fetch_store": 0,
"pe_ipcq": 0,
"pe_cpu": 0
}
},
{
"panel": "multi_user_prefill_gqa",
"kind": "prefill",
"C": 4,
"S_kv": 16,
"latency_ns": 4630.408000000019,
"op_log_summary": {
"gemm_count": 32,
"ipcq_copy_count": 24,
"dma_read_count": 12,
"dma_write_count": 4
},
"engine_occupancy_ns": {
"pe_gemm": 8.192000000000917,
"pe_math": 236.0,
"pe_dma": 1446.945000000003,
"pe_fetch_store": 0,
"pe_ipcq": 0,
"pe_cpu": 0
}
},
{
"panel": "single_user_decode_gqa",
"kind": "decode",
"C": 1,
"P": 8,
"S_kv": 64,
"latency_ns": 3631.730500000015,
"op_log_summary": {
"gemm_count": 16,
"ipcq_copy_count": 21,
"dma_read_count": 24,
"dma_write_count": 1
},
"engine_occupancy_ns": {
"pe_gemm": 16.383999999998196,
"pe_math": 160.0,
"pe_dma": 1450.1150000000016,
"pe_fetch_store": 0,
"pe_ipcq": 0,
"pe_cpu": 0
}
},
{
"panel": "multi_user_decode_gqa",
"kind": "decode",
"C": 4,
"P": 8,
"S_kv": 128,
"latency_ns": 6692.578999999866,
"op_log_summary": {
"gemm_count": 64,
"ipcq_copy_count": 93,
"dma_read_count": 96,
"dma_write_count": 1
},
"engine_occupancy_ns": {
"pe_gemm": 32.76799999998184,
"pe_math": 688.0000000000009,
"pe_dma": 15919.751500000013,
"pe_fetch_store": 0,
"pe_ipcq": 0,
"pe_cpu": 0
}
}
]
}
docs/report/1H-codesign-paper/figures/gqa_latency_by_panel.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 38 KiB

docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_hbm_budget.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 47 KiB

docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_kv_sharding_diagram.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 135 KiB

docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_kv_sharding_table.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 123 KiB

docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_memory_comm_analytical.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 132 KiB

docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_memory_comm_paired.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 183 KiB

docs/report/1H-codesign-paper/figures/gqa_long_ctx_6cases_summary.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 182 KiB

docs/report/1H-codesign-paper/figures/gqa_op_engine_breakdown.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 62 KiB

docs/report/1H-codesign-paper/figures/gqa_prefill_compute_bound.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 145 KiB

docs/report/1H-codesign-paper/figures/gqa_short_context/batch_scaling.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 72 KiB

docs/report/1H-codesign-paper/figures/gqa_short_context/gemm_util_a1_variant_ablation.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 96 KiB

docs/report/1H-codesign-paper/figures/gqa_short_context/per_cube_tradeoff_mode_compare.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 133 KiB

docs/report/1H-codesign-paper/figures/gqa_short_context/wall_a1_variant_compare.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 70 KiB

docs/report/1H-codesign-paper/figures/gqa_short_context/wall_variant1_mode_compare.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 73 KiB

docs/report/1H-codesign-paper/figures/gqa_short_context/wall_variant2_mode_compare.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 75 KiB

docs/report/1H-codesign-paper/figures/gqa_short_context/wall_variant3_mode_compare.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 84 KiB

docs/report/1H-codesign-paper/figures/ipcq_alternatives_architecture.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 253 KiB

docs/report/1H-codesign-paper/figures/ipcq_alternatives_architecture_flow.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 305 KiB

docs/report/1H-codesign-paper/figures/ipcq_alternatives_architecture_stacked.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 381 KiB

docs/report/1H-codesign-paper/figures/ipcq_alternatives_decision_matrix.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 189 KiB

docs/report/1H-codesign-paper/figures/ipcq_alternatives_latency.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 160 KiB

docs/report/1H-codesign-paper/figures/latency_model.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 61 KiB

docs/report/1H-codesign-paper/figures/pe_architecture.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1014 KiB

docs/report/1H-codesign-paper/figures/probe_pe_dma_summary.txt
+57
View File
@@ -0,0 +1,57 @@
=== H2D Write Latency (IO->HBM, data=32768B) ===
Case Target Hops Actual Ovhd Drain Wire Ovhd% Drain% Eff.BW BN.BW Util%
-------------------------------------------------------------------------------------------------------------------
h2d-1hop cube0.pe0 1 289.53 50.0 256.0 0.18 17.3% 88.4% 113.17 128.0 88.4%
h2d-2hop cube4.pe0 2 326.04 78.0 256.0 0.20 23.9% 78.5% 100.50 128.0 78.5%
h2d-3hop cube8.pe0 3 362.56 106.0 256.0 0.20 29.2% 70.6% 90.38 128.0 70.6%
h2d-4hop cube12.pe0 4 399.06 134.0 256.0 0.21 33.6% 64.1% 82.11 128.0 64.1%
[v] Monotonic increase: PASS
BW Saturation (Util% by data size):
Case 4KB 16KB 64KB 256KB 1MB
------------------------------------------------------------------
h2d-1hop 38.9% 71.8% 91.1% 97.6% 99.4%
h2d-2hop 29.0% 62.1% 86.8% 96.3% 99.1%
h2d-3hop 23.2% 54.7% 82.8% 95.1% 98.7%
h2d-4hop 19.3% 48.8% 79.2% 93.8% 98.4%
=== D2H Read Latency (HBM->IO, data=32768B) ===
Case Source Hops Actual Ovhd Drain Wire Ovhd% Drain% Eff.BW BN.BW Util%
-------------------------------------------------------------------------------------------------------------------
d2h-1hop cube0.pe0 1 571.20 23.0 256.0 0.04 4.0% 44.8% 57.37 128.0 44.8%
d2h-2hop cube4.pe0 2 635.72 51.0 256.0 0.05 8.0% 40.3% 51.54 128.0 40.3%
d2h-3hop cube8.pe0 3 700.24 79.0 256.0 0.06 11.3% 36.6% 46.80 128.0 36.6%
d2h-4hop cube12.pe0 4 764.76 107.0 256.0 0.07 14.0% 33.5% 42.85 128.0 33.5%
[v] Monotonic increase: PASS
BW Saturation (Util% by data size):
Case 4KB 16KB 64KB 256KB 1MB
------------------------------------------------------------------
d2h-1hop 9.2% 28.9% 61.9% 86.7% 96.3%
d2h-2hop 7.8% 25.2% 57.4% 84.4% 95.6%
d2h-3hop 6.7% 22.4% 53.5% 82.2% 94.9%
d2h-4hop 5.9% 20.1% 50.2% 80.1% 94.2%
[v] D2H >= H2D (reverse data path): PASS
=== PE DMA Latency (pe_dma -> router -> HBM, data=32768B) ===
Case Target Actual Ovhd Drain Wire Ovhd% Drain% Eff.BW BN.BW Util%
----------------------------------------------------------------------------------------------------------------------------
pe-local-hbm c0.pe0->c0.slice0 141.00 2.0 128.0 0.00 1.4% 90.8% 232.40 256.0 90.8%
pe-same-half-hbm c0.pe0->c0.slice1 147.87 4.0 128.0 0.03 2.7% 86.6% 221.60 256.0 86.6%
pe-cross-half-hbm c0.pe0->c0.slice4 161.17 10.0 128.0 0.09 6.2% 79.4% 203.31 256.0 79.4%
pe-cross-cube-hbm-best c0.pe0->c1.slice0 330.52 30.0 256.0 0.01 9.1% 77.5% 99.14 128.0 77.5%
pe-cross-cube-hbm-worst c0.pe0->c15.slice0 677.12 180.0 256.0 0.06 26.6% 37.8% 48.39 128.0 37.8%
* Local BN: 256.0 GB/s, Remote BN: 256.0 GB/s
[v] Cross-cube best < worst: PASS (330.52ns < 677.12ns)
BW Saturation (Util% by data size):
Case 4KB 16KB 64KB 256KB 1MB
------------------------------------------------------------------
pe-local-hbm 88.9% 97.0% 99.2% 99.8% 100.0%
pe-same-half-hbm 79.9% 94.1% 98.5% 99.6% 99.9%
pe-cross-half-hbm 61.3% 86.4% 96.2% 99.0% 99.8%
pe-cross-cube-hbm-best 51.6% 81.0% 94.5% 98.6% 99.6%
pe-cross-cube-hbm-worst 15.1% 41.6% 74.0% 91.9% 97.8%
============================================================
docs/report/1H-codesign-paper/figures/roofline_long_context.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 98 KiB

docs/report/1H-codesign-paper/figures/roofline_short_context.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 106 KiB

docs/report/1H-codesign-paper/figures/sip_architecture.pdf
BIN
View File
Binary file not shown.
docs/report/1H-codesign-paper/figures/sip_architecture.svg
+94
View File
@@ -0,0 +1,94 @@
<svg xmlns="http://www.w3.org/2000/svg" width="508" height="495" viewBox="70 40 508 495">
<title>sip</title>
<rect width="648" height="648" fill="#ffffff"/>
<line x1="108.0" y1="144.0" x2="252.0" y2="144.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="180.0" y="140.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="108.0" y1="144.0" x2="108.0" y2="264.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="108.0" y="200.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="252.0" y1="144.0" x2="396.0" y2="144.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="324.0" y="140.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="252.0" y1="144.0" x2="252.0" y2="264.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="252.0" y="200.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="396.0" y1="144.0" x2="540.0" y2="144.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="468.0" y="140.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="396.0" y1="144.0" x2="396.0" y2="264.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="396.0" y="200.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="540.0" y1="144.0" x2="540.0" y2="264.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="540.0" y="200.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="108.0" y1="264.0" x2="252.0" y2="264.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="180.0" y="260.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="108.0" y1="264.0" x2="108.0" y2="384.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="108.0" y="320.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="252.0" y1="264.0" x2="396.0" y2="264.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="324.0" y="260.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="252.0" y1="264.0" x2="252.0" y2="384.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="252.0" y="320.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="396.0" y1="264.0" x2="540.0" y2="264.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="468.0" y="260.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="396.0" y1="264.0" x2="396.0" y2="384.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="396.0" y="320.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="540.0" y1="264.0" x2="540.0" y2="384.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="540.0" y="320.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="108.0" y1="384.0" x2="252.0" y2="384.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="180.0" y="380.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="108.0" y1="384.0" x2="108.0" y2="504.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="108.0" y="440.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="252.0" y1="384.0" x2="396.0" y2="384.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="324.0" y="380.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="252.0" y1="384.0" x2="252.0" y2="504.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="252.0" y="440.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="396.0" y1="384.0" x2="540.0" y2="384.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="468.0" y="380.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="396.0" y1="384.0" x2="396.0" y2="504.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="396.0" y="440.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="540.0" y1="384.0" x2="540.0" y2="504.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="540.0" y="440.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="108.0" y1="504.0" x2="252.0" y2="504.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="180.0" y="500.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="252.0" y1="504.0" x2="396.0" y2="504.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="324.0" y="500.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<line x1="396.0" y1="504.0" x2="540.0" y2="504.0" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="468.0" y="500.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">1.0mm 512GB/s</text>
<polyline points="324.0,56.0 108.0,56.0 108.0,144.0" fill="none" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="216.0" y="96.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">2.5mm 512GB/s</text>
<polyline points="324.0,56.0 252.0,56.0 252.0,144.0" fill="none" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="288.0" y="96.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">2.5mm 512GB/s</text>
<polyline points="324.0,56.0 396.0,56.0 396.0,144.0" fill="none" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="360.0" y="96.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">2.5mm 512GB/s</text>
<polyline points="324.0,56.0 540.0,56.0 540.0,144.0" fill="none" stroke="#000000" stroke-width="1" opacity="0.8"/>
<text x="432.0" y="96.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#000000">2.5mm 512GB/s</text>
<rect x="84.0" y="128.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="108.0" y="148.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (0,0)</text>
<rect x="228.0" y="128.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="252.0" y="148.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (1,0)</text>
<rect x="372.0" y="128.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="396.0" y="148.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (2,0)</text>
<rect x="516.0" y="128.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="540.0" y="148.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (3,0)</text>
<rect x="84.0" y="248.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="108.0" y="268.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (0,1)</text>
<rect x="228.0" y="248.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="252.0" y="268.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (1,1)</text>
<rect x="372.0" y="248.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="396.0" y="268.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (2,1)</text>
<rect x="516.0" y="248.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="540.0" y="268.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (3,1)</text>
<rect x="84.0" y="368.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="108.0" y="388.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (0,2)</text>
<rect x="228.0" y="368.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="252.0" y="388.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (1,2)</text>
<rect x="372.0" y="368.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="396.0" y="388.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (2,2)</text>
<rect x="516.0" y="368.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="540.0" y="388.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (3,2)</text>
<rect x="84.0" y="488.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="108.0" y="508.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (0,3)</text>
<rect x="228.0" y="488.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="252.0" y="508.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (1,3)</text>
<rect x="372.0" y="488.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="396.0" y="508.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (2,3)</text>
<rect x="516.0" y="488.0" width="48.0" height="32.0" rx="4" fill="#cbd5e1" stroke="#000000" stroke-width="1"/>
<text x="540.0" y="508.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#1e293b">CUBE (3,3)</text>
<rect x="308.0" y="50.0" width="32.0" height="12.0" rx="4" fill="#0ea5e9" stroke="#000000" stroke-width="1"/>
<text x="324.0" y="60.0" text-anchor="middle" font-family="monospace" font-size="10" fill="#ffffff">IO io0</text>
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 10 KiB

docs/report/1H-codesign-paper/main.tex
+55
View File
@@ -0,0 +1,55 @@
\documentclass[10pt,twocolumn]{article}
\usepackage[margin=0.75in]{geometry}
\usepackage{graphicx}
\usepackage{booktabs}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{siunitx}
\DeclareSIUnit\flop{FLOP}
\DeclareSIUnit\cycle{cycle}
\DeclareSIUnit\byte{B}
\usepackage{xcolor}
\usepackage{hyperref}
\hypersetup{colorlinks=true,linkcolor=blue!50!black,citecolor=blue!50!black,urlcolor=blue!50!black}
\usepackage{caption}
\usepackage{subcaption}
\captionsetup{font=small,labelfont=bf}
\usepackage{lmodern} % scalable Latin Modern fonts (required by microtype expansion)
\usepackage{microtype}
\usepackage{tikz}
\usetikzlibrary{arrows.meta,positioning,calc,fit}
\graphicspath{{figures/}}
\title{\textbf{Hardware--Software Co-Design for Grouped-Query Attention on AHBM}}
\author{Mukesh Garg \and Jiyoon Lee \and Yangwook Kang}
\date{
\small
AGI Computing Lab, System Technology Group\\
2026 Q1-Q3 Report
}
\begin{document}
\maketitle
\input{sections/00-exec-summary}
\input{sections/01-introduction}
\input{sections/02-platform}
\input{sections/03-gemm}
\input{sections/04-allreduce}
\input{sections/05-gqa} % section header + intro
\input{sections/05a-roofline} % 5.1 Roofline Analysis
\input{sections/05b-capacity-planning} % 5.2 Capacity Planning
\input{sections/05c-parallelism-selection} % 5.3 Parallelism Selection
\input{sections/05x-fused-kernel} % 5.4-5.7 Placement + short/long ctx + composite
\input{sections/05z-summary} % 5.8 Summary
\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/00-exec-summary.tex
+36
View File
@@ -0,0 +1,36 @@
\section*{Executive Summary}
\addcontentsline{toc}{section}{Executive Summary}
Grouped-Query Attention (GQA) has largely replaced GPT-3-style
multi-head attention in modern decoder-only LLMs (e.g., Llama 3 and
Mistral) due to its reduced memory and bandwidth requirements. Mapping
GQA 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: GQA-aware placement
of KV caches and weights across TCM, SRAM, and HBM; PE\_IPCQ,
an efficient on-device collective communication primitive; and a
composite-command GEMM pipeline that tightly pipelines memory
and compute operations within each PE under PE\_SCHEDULER control.
Kernbench-based evaluation show up to \SI{38}{\percent} lower collective
communication overhead, \SIrange{25}{35}{\percent} lower all-reduce
latency than a mesh-based interconnect, and up to \SI{78}{\percent} of
peak MAC efficiency for compute-intensive GEMM kernels. These results
demonstrate that the fused GQA kernel can efficiently exploit AHBM's HBM bandwidth.
While GQA serves as the motivating workload in this study, the
resulting architectural mechanisms are broadly applicable across the
AHBM software stack. PE\_IPCQ provides a scalable communication
substrate for collective operations and communication-intensive workloads, while
composite-command execution enables efficient fusion of memory
movement, GEMM kernels, normalization, and other element-wise
operations. Together with the hierarchical data-placement framework,
these capabilities form a reusable foundation for future AI kernels
and communication libraries on AHBM. Looking ahead, these same mechanisms provide
the basis for optimizing FFN-intensive workloads such as Mixture-of-Experts (MoE)
layers and for developing integrated optimization strategies spanning both attention
and MoE components within next-generation AI models.
docs/report/1H-codesign-paper/sections/01-introduction.tex
+86
View File
@@ -0,0 +1,86 @@
\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. GQA
dominates inference-time memory traffic and KV-cache capacity in modern
LLM serving, making it the primary bandwidth bottleneck on memory-centric
architectures such as AHBM. 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.
docs/report/1H-codesign-paper/sections/01-introduction_ko.tex
+73
View File
@@ -0,0 +1,73 @@
\section{서론}
\label{sec:intro}
AHBM은 연산 유닛을 HBM 스택 내부에 직접 통합한 메모리 중심 가속기
아키텍처이다. 각 처리 요소(PE)는 전용 HBM 슬라이스와 짝지어지며,
TCM과 SRAM을 거쳐 메모리와 MAC 어레이 사이의 데이터를 단계적으로
이동시킨다. 이러한 메모리 중심 구조에서 커널 성능은 연산 처리량만이
아니라, 메모리 계층과 PE 사이에 데이터를 얼마나 효율적으로 배치하고,
이동시키며, 공유하는가에 의해 결정된다. 따라서 AHBM에서의 AI 커널
최적화는 커널 알고리즘과 아키텍처 메커니즘을 함께 발전시키는
하드웨어--소프트웨어 코디자인을 필요로 한다.
상세한 성능 분석과 신속한 설계 탐색을 위해, 우리는
\textbf{KernBench}---AHBM을 위한 소스 수준 이산 사건 시뮬레이션
플랫폼---을 개발하였다. KernBench는 메모리 시스템 지연, PE 실행
모델, PE 간 통신, 호스트 측 오케스트레이션을 포함한 AHBM 실행 모델을
구현하면서, 커널과 호스트 소프트웨어를 소스 코드로부터 직접 실행한다.
이를 통해 실행 거동을 세밀하게 관찰할 수 있을 뿐만 아니라, 컴파일러나
런타임과 같은 상위 소프트웨어 스택과 독립적으로 하드웨어--소프트웨어
코디자인 선택을 체계적으로 평가할 수 있다. 본 보고서에 제시된 모든
결과는 KernBench를 사용하여 얻은 것이며, KernBench는 본 연구에서
다루는 모든 메커니즘과 커널의 공통 평가 플랫폼 역할을 한다.
본 보고서는 LLM 추론에서 가장 성능 및 대역폭 임계 구성요소 중 하나인
Grouped-Query Attention(GQA)에 초점을 맞춘다. Llama~3, Mistral과
같은 최신 디코더 전용 모델은 GPT-3 형태의 다중-헤드 어텐션(MHA)으로
부터 GQA---여러 질의 헤드가 하나의 KV 헤드를 공유함으로써 KV 캐시
용량과 메모리 대역폭 요구량을 줄이는 방식---로 대거 전환되었다. GQA는
모델 수준에서 시스템 효율을 개선하지만, 이를 AHBM에 효율적으로
매핑하기 위해서는 세 가지 아키텍처 요구사항이 발생한다: PE 간 통신을
최소화하기 위한 KV 캐시와 가중치의 최적 배치, 불가피한 PE 간 트래픽에
대한 저오버헤드 지원, 그리고 각 PE 내에서 메모리 접근과 연산의
효율적인 파이프라이닝이다.
이러한 요구사항을 충족하기 위해, 본 보고서는 세 가지 하드웨어--
소프트웨어 코디자인 메커니즘을 제안한다. 첫째, GQA-aware 데이터
배치는 TCM/SRAM/HBM 계층에 KV 캐시와 가중치를 분산 배치하여 통신
오버헤드를 줄이고 데이터 지역성을 향상시킨다. 둘째, PE\_IPCQ는 리덕션
등 통신 집약적 연산을 위한 효율적인 온디바이스 집합 통신 프리미티브를
제공한다. 셋째, composite-command GEMM 파이프라인은 PE\_SCHEDULER의
제어 하에 각 PE 내부에서 메모리 이동과 연산을 긴밀하게 파이프라이닝
하여, 명령 오버헤드를 줄이면서도 MAC 어레이의 가동률을 효율적으로
유지한다.
연산 측 인에이블러(composite-command GEMM 파이프라인,
\S\ref{sec:gemm})와 통신 측 인에이블러(PE\_IPCQ,
\S\ref{sec:allreduce})는 우선 독립적으로 개발 및 평가된다. 이후 융합
GQA 커널(\S\ref{sec:gqa})은 이들을 GQA-aware 데이터 배치와 결합하여
AHBM 상에서 종단간 어텐션 구현을 시연한다. 대응 관계는 직접적이다:
어텐션의 $QK^{\top}$와 $PV$ 곱은 정확히 composite-command 파이프라인이
이득을 주는 GEMM이며, KV 리덕션은 정확히 PE\_IPCQ가 이득을 주는 집합
연산이다. 이러한 메커니즘들이 결합되어, 융합 GQA 커널은 AHBM의 HBM
대역폭을 효율적으로 활용할 수 있게 된다.
본 연구에서 GQA가 주된 동기 부여 워크로드 역할을 하지만, 도출된
메커니즘들은 훨씬 광범위한 AI 커널 분류군에서 재사용 가능한 구성
요소로 의도되었다. PE\_IPCQ는 분산 및 통신 집약적 워크로드 전반에
걸쳐 집합 통신을 지원할 수 있으며, composite-command 실행은 GEMM
기반 커널, 피드포워드 네트워크(FFN), 정규화, 그 외 융합 연산자
파이프라인에 적용 가능하다. 계층적 데이터 배치 프레임워크와 함께 이
메커니즘들은 AHBM에서의 향후 AI 커널과 통신 라이브러리의 기반을
형성한다. 2026년 하반기에는 이 기반을 FFN 및 MoE 주도 워크로드와
완전한 LLM 실행의 종단간 최적화로 확장할 예정이다.
본 보고서의 나머지 부분은 다음과 같이 구성된다.
\S\ref{sec:platform}에서는 KernBench 플랫폼과 본 연구 전반에 걸쳐
사용된 AHBM 구성을 기술한다. \S\ref{sec:gemm},
\S\ref{sec:allreduce}, \S\ref{sec:gqa}에서는 각각 composite-command
GEMM 파이프라인, PE\_IPCQ 집합 통신, 융합 GQA 커널을 다룬다.
\S\ref{sec:discussion}에서는 이러한 결과의 광범위한 아키텍처적 함의를
논의한다. 마지막으로 \S\ref{sec:conclusion}과 \S\ref{sec:future}에서는
주요 결과를 요약하고 FFN, MoE, 전체 모델 최적화에 관한 향후 과제를
제시한다.
docs/report/1H-codesign-paper/sections/02-platform.tex
+465
View File
@@ -0,0 +1,465 @@
\begin{figure*}[t]
\centering
\begin{subfigure}[b]{0.495\textwidth}
\centering
\includegraphics[width=\linewidth,height=0.7\linewidth,keepaspectratio]{sip_architecture.pdf}
\caption{SIP Architecture}
\label{fig:sip-arch}
\end{subfigure}\hfill
\begin{subfigure}[b]{0.42\textwidth}
\centering
\includegraphics[width=\linewidth,height=0.84\linewidth]{cube_architecture.pdf}
\caption{CUBE Architecture}
\label{fig:cube-arch}
\end{subfigure}
\vspace{0.6em}
\begin{subfigure}[t]{\textwidth}
\centering
\includegraphics[width=\linewidth]{pe_architecture.png}
\caption{PE level (zoom-in of one PE in (b)): \textsf{PE\_CPU}
dispatches commands to \textsf{PE\_SCHED}, which routes tile-token
streams through \textsf{PE\_DMA}, \textsf{PE\_FETCH\_STORE}, and
\textsf{GEMM}/\textsf{MATH} engines; \textsf{PE\_IPCQ} is the
on-device collective control plane.}
\label{fig:pe-arch}
\end{subfigure}
\caption{Modeled hardware graph at the SIP, CUBE, and PE levels.
This figure is an \emph{illustrative topology} chosen for
readability; the \emph{experimental configuration} used throughout
this report (port counts, slice counts, and link parameters) is
specified in \S\ref{sec:hw} / Table~\ref{tab:hw}, and numbers in the
two may differ. KernBench is not tied to this particular
arrangement: each box is a modeled component node, each line a
directed link with bandwidth and propagation attributes, and any
topology that respects those attributes is supported.}
\label{fig:hw-arch}
\end{figure*}
\begin{figure*}[t]
\centering
\includegraphics[width=\textwidth]{latency_model.png}
\caption{Conceptual schematic of the latency model. Two source nodes
(Requester A/B) inject flits through a chain of routers into a
destination node; on the shared edge between routers, flits from the
two transactions are interleaved flit-by-flit by the wire's FIFO
arrival order. End-to-end latency is the sum of four contributions:
\textbf{per-node overhead} (the switch's fixed processing cost,
shown in light yellow---the same colour as the destination's
processing-logic block), \textbf{per-edge transmission}
($\textit{flit\_size}/\textit{BW}$ on each wire),
\textbf{drain} (per-flit service occupancy at the destination's
channel), and \textbf{queuing delay} (waiting in a FIFO when a
shared resource is busy). The places where queuing actually
accumulates are highlighted in green: the router output queue and
the destination's input queue. This is the model, not a
measurement; specific bandwidths and overheads are listed in
\S\ref{sec:hw} (Table~\ref{tab:hw}).}
\label{fig:latency-model}
\end{figure*}
\section{The KernBench Platform}
\label{sec:platform}
All results in this report are produced on \emph{KernBench}, a
system-level discrete-event simulator for LLM kernels running on AHBM.
This section explains why the platform exists, the device and
execution model it presents to a kernel writer, how its latency model
turns that execution into a number, and the concrete hardware
configuration used for every experiment that follows.
\subsection{Why KernBench}
\label{sec:why}
In a production end-to-end (E2E) stack, kernel performance is entangled
with every layer above the hardware: the compiler's tiling and
scheduling choices, the framework's operator dispatch, the collective
library, and the runtime. Good E2E numbers require \emph{all} of those
layers to be co-optimized, which makes it hard to answer a narrower but
more fundamental question: \emph{given the hardware, how fast can a
well-written kernel be, and which hardware features actually make it
faster?}
KernBench is built to answer exactly that question. Kernels are written
and executed at the \emph{source level}---as algorithmic descriptions
in a small tile-oriented kernel API---with no dependency on a compiler
or any other software-stack layer. The simulator takes the kernel and a
hardware topology and reports the latency the modeled hardware would
deliver. This isolation is deliberate: it lets us study algorithm-level
optimizations (how to tile a GEMM, how to schedule a collective, how to
fuse an attention kernel) and the hardware features that support them,
without the confound of compiler maturity or framework overhead. The
cost is that KernBench numbers are \emph{not} E2E latencies; they are
the achievable-kernel latencies an ideal software stack would expose.
\subsection{Device and execution model}
\label{sec:exec}
KernBench models AHBM as a hierarchy of SIPs, CUBEs, and processing
elements (PEs). At the system level, multiple SIPs are connected
through inter-package links, while each SIP contains a collection of
CUBEs joined by an on-package interconnect
(Fig.~\ref{fig:hw-arch}\subref{fig:sip-arch}).
Each CUBE contains eight PEs, shared SRAM, HBM controllers, an
\textsf{M\_CPU} control processor, and an intra-CUBE router mesh
(Fig.~\ref{fig:hw-arch}\subref{fig:cube-arch}). Together these
components form the execution substrate for all kernels evaluated in
this report. This organization reflects the memory-centric nature of
AHBM: each PE is paired with a dedicated slice of HBM bandwidth and
local on-PE storage (TCM), so compute lives next to the data it
consumes rather than fetching it through a far-away memory
controller. The platform's performance question is therefore not
``how many FLOPs can the chip do'' but ``how well can a kernel keep
each compute engine fed from its locally-attached memory while
moving the unavoidable traffic between PEs efficiently.''
Within a PE, commands are dispatched by \textsf{PE\_CPU} to
\textsf{PE\_SCHED}, which routes work to specialized execution
engines---DMA, FETCH/STORE, GEMM, vector-math, and IPCQ
(Fig.~\ref{fig:hw-arch}\subref{fig:pe-arch}). Commands come in two
flavours. \emph{Atomic} commands target a single engine---a plain
DMA read, a GEMM tile, a vector-math op, an IPCQ send/receive---and
are the natural unit for short or special-purpose work; the
\textsf{PE\_CPU} itself also runs control-plane work directly.
\emph{Composite} commands, by contrast, carry an ordered pipeline of
operations across multiple engines (\textsf{DMA\_READ} $\rightarrow$
\textsf{FETCH} $\rightarrow$ \textsf{GEMM}/\textsf{MATH} $\rightarrow$
\textsf{STORE} $\rightarrow$ \textsf{DMA\_WRITE}) that the
\textsf{PE\_SCHED} tiles and streams without per-tile redispatch. The
composite form is the substrate for the GEMM optimization
(\S\ref{sec:gemm}) and, combined with on-PE collectives, for fused
attention (\S\ref{sec:gqa}).
KernBench is layered along the flow of a request:
\begin{itemize}
\item The \textbf{runtime API} is host-facing and
topology-agnostic---it deploys tensors and launches kernels but knows
nothing about routing or interconnect.
\item The \textbf{simulation engine} schedules discrete events and
routes every request through the modeled graph.
\item The \textbf{components} are device-side nodes that model
hardware behaviour: the per-PE blocks (scheduler, DMA, GEMM and
vector-math engines, TCM, IPCQ), the NoC routers, the HBM
controllers, and the inter-chiplet links.
\end{itemize}
Data and timing are handled in two passes, so that a kernel's
numeric results and its latency are computed consistently but
independently. \textbf{Pass~1 (timing)} runs the kernel under
the discrete-event engine: memory ops (\textsf{tl.load},
\textsf{tl.store}) execute against a host-side \textsf{MemoryStore}
and return real tensor data, while compute ops (GEMM, vector-math)
only emit records into an op-log carrying their operands, shapes,
dtypes, and scheduled time. \textbf{Pass~2 (data)} replays that
op-log offline in numpy, producing the actual numeric outputs and
comparing them against a reference computation under per-dtype
tolerances (e.g.\ \texttt{rtol}/\texttt{atol} $=10^{-3}$ for f16).
Pass~2 is optional---runs that need only latency skip it---but when
enabled it guarantees that every reported timing number corresponds
to a kernel whose numeric output has been verified end-to-end.
\subsection{Latency model: graph traversal and contention}
\label{sec:latency}
The modeled hardware hierarchy described above is represented
internally as a directed graph. Nodes correspond to hardware
components---PEs, routers, memory controllers, SRAM blocks, IO
chiplets---while edges represent communication links with associated
bandwidth (\si{\giga\byte\per\second}) and propagation
(\si{\nano\second}) attributes. Every operation in KernBench---DMA
transfers, remote memory accesses, collective communication, command
dispatch---is modelled as a traversal through this graph. End-to-end
latency is decomposed into four contributions accumulated along the
traversal path (Fig.~\ref{fig:latency-model}): \textbf{per-node
overhead} at each component, \textbf{per-edge transmission} on each
wire, \textbf{drain} (per-flit service occupancy) at the destination,
and \textbf{queuing delay} at the shared FIFOs that the wire and the
destination share between concurrent transactions.
\paragraph{The hardware as a graph.} The topology is compiled once at
configuration time into this graph and is never mutated during a run.
There are no hidden shortcuts, implicit bypasses, or magic paths: if a
request reaches its destination, the path it took is explicit in the
graph, and the latency it incurred is the sum of the per-node and
per-edge costs paid along that path. The same graph representation
applies recursively at every hierarchy level---system, SIP, CUBE, and
PE (Fig.~\ref{fig:hw-arch}).
\paragraph{From graph to discrete-event simulation.} The graph is
driven by a discrete-event engine. Two kinds of events advance
simulation time: \emph{node events} (component switching overhead,
service completions such as an HBM channel commit or a GEMM tile
finish) and \emph{edge events} (the flit-by-flit serialization of a
payload across a bandwidth-limited link). The engine maintains a
priority queue of pending events ordered by their scheduled time and
fires them one at a time, with ties broken under a deterministic
policy so that the same kernel on the same topology always yields the
same trace. Per-request correlation IDs are stamped at injection and
carried through every hop, so the path from injection to completion
is fully traceable. Every modeled latency contribution
corresponds to exactly one of these events on exactly one node or
edge---there is no slack in the budget.
\paragraph{Latency contributions.} Three kinds of latency accumulate
along a traversal: (i) \emph{per-node fixed overhead}---each component
carries a small switching cost (router decode, controller pickup,
scheduler handoff); (ii) \emph{per-edge transfer time}---each link's
payload is decomposed into fixed-size flits (default
\SI{256}{\byte}), and each flit arrives at
$\text{prop}+\text{flit\_bytes}/\text{bw}$ after the previous one,
giving wormhole semantics across multi-hop paths; and (iii)
\emph{per-service occupancy}---memory controllers, GEMM stages, and
collective engines hold the request for their service time before
releasing it downstream. Each of these is attached to a specific node
or edge in the graph; together they make up the entire latency budget.
Beyond data movement and execution latency, KernBench
also models the control-plane cost required to issue work
to accelerator engines. Since every DMA, GEMM, vector-
math, and IPCQ operation is initiated through the
\textsf{PE\_CPU} $\rightarrow$ \textsf{PE\_SCHED} path,
command dispatch latency contributes to the end-to-end
execution time of all kernels.
\paragraph{Command dispatch overhead model.} The cost incurred by
the \textsf{PE\_CPU} when it dispatches a command to one of the
accelerator engines is modelled structurally rather than with a
per-operation calibration table. The \textsf{PE\_CPU} charges, per
command,
\[
d_{\text{cmd}} = \textsf{FIXED} + b_{\text{logical}} \cdot R,
\]
This linear-in-size form reflects the serialization cost of writing a
command descriptor into the scheduler queue: a fixed per-command
bookkeeping cost plus a byte-wise descriptor-transfer cost.
Concretely, $b_{\text{logical}}$ is the command's hardware-logical byte
size, \textsf{FIXED} captures the fixed per-command cost (queue-tail
update, completion registration) and $R$ captures the per-byte cost of
serializing the command descriptor into the scheduler queue. This
\textsf{FIXED} is distinct from Table~\ref{tab:hw}'s
\SI{2}{\nano\second} / \SI{1}{\nano\second} \textsf{PE\_CPU} /
\textsf{PE\_SCHED} fixed costs: the latter are per-component
traversal overheads each command pays when it transits those nodes,
whereas the \textsf{FIXED} term here is the command-descriptor
issue cost. \textsf{FIXED} depends on the command class: a single-op
command --- one engine operation, i.e.\ a single DMA descriptor, GEMM,
or elementwise issue --- carries a lighter 8-cycle \textsf{FIXED},
while the composite command, which the scheduler expands into a
multi-stage tile-feeder plan, carries 40 cycles
(\S\ref{sec:gemm-vs-async}). With the composite \textsf{FIXED} $= 40$
cycles and $R = 0.0625$ cycles/byte (i.e.\ \SI{16}{\byte\per\cycle},
at \SI{1}{\giga\hertz}) a typical composite lands at roughly
\SI{43}{\nano\second}, and a hard cap on a composite's descriptor size
prevents the model from rewarding arbitrarily large fused commands
beyond what real descriptor queues accept. In the configurations measured here, command issue is
not the bottleneck---data movement is---so this term stays small
relative to DMA and collective time.
\begin{table*}[t]
\centering
\caption{PE\,$\to$\,HBM DMA latency probe at varying hop distances
(\SI{32}{\kibi\byte} transfer; output captured from \texttt{kernbench
probe}). \emph{Util\%} is the achieved bandwidth as a fraction of the
path's bottleneck-edge bandwidth.}
\label{tab:probe-pe-dma}
\small
\begin{tabular}{@{}lrrr@{}}
\toprule
Case & Latency~(\si{\nano\second}) & Util\% (\,32\,KiB) & Util\% (\,1\,MiB) \\
\midrule
PE\,$\to$\,local HBM & 141.0 & 90.8 & 100.0 \\
PE\,$\to$\,same-half HBM & 147.9 & 86.6 & ~99.9 \\
PE\,$\to$\,cross-half HBM & 161.2 & 79.4 & ~99.8 \\
PE\,$\to$\,cross-CUBE (best) & 330.5 & 77.5 & ~99.6 \\
PE\,$\to$\,cross-CUBE (worst) & 677.1 & 37.8 & ~97.8 \\
\bottomrule
\end{tabular}
\end{table*}
\subsubsection{Congestion and contention modeling}
\label{sec:congestion}
The simulator's
sharpness comes from how it models contention for those nodes and
edges. \emph{Every directed edge has a FIFO}: an arriving flit takes
its bandwidth-limited transfer time on top of whatever earlier flits
are still being served, so a busy link queues later traffic behind
earlier traffic rather than transferring everything at peak BW.
\emph{HBM is modelled with per-pseudo-channel parallelism}: a
stateless array of channel-availability timestamps with address-based
channel selection captures the bank-level concurrency that real HBM
exposes, so 64 channels per CUBE deliver real parallelism on uniform
addresses but a hot channel surfaces as the bottleneck on skewed ones.
\emph{Every component has a serial worker}: a router carrying two
heavy streams interleaves them at flit granularity in arrival order
rather than fanning out for free, so two concurrent collectives sharing
a link share its bandwidth, not double it. Without these mechanisms
the simulator would simply re-confirm the peak-BW roofline; with them,
it reveals where the real bottlenecks form and which hardware levers
actually relieve them---which is exactly the question the codesign
work in this report turns on.
\subsection{Accuracy}
\label{sec:accuracy}
The model is precise about the effects that
dominate kernel latency on this class of hardware: per-edge bandwidth
occupancy and flit-level serialization, HBM pseudo-channel parallelism,
and per-component switching overhead. Two independent cross-checks
drawn from the experiments in this report confirm that this precision
translates into physically reasonable kernel latencies.
First, in the GEMM study (\S\ref{sec:gemm}), simulator-measured MAC
efficiency tracks an analytic ideal-pipeline model within
\SI{1.4}{ppt} across every swept shape---from a single-tile
$M{=}K{=}N{=}32$ at $\sim\SI{7.7}{\percent}$ up to the deep-$K$
$K{=}3072$ case at $\sim\SI{90}{\percent}$. The residual gap is
fill/tail overhead the closed-form pipeline model omits.
Second, in the all-reduce study (\S\ref{sec:allreduce},
Fig.~\ref{fig:allreduce-cmp}), simulator latency for a 2D-torus over
six devices follows the expected startup-plus-per-packet shape across
the entire payload sweep---tight at small payloads where startup
dominates, and within a single-digit multiplicative factor at the
largest payloads, where the residual gap is explained by per-router
switching the analytic shape elides.
Third, the simulator's per-traversal behaviour can be probed
directly. Running \texttt{kernbench probe} on the modelled topology
issues a sequence of PE-to-HBM DMA reads at progressively greater
hop distances and reports the per-component overhead, per-edge
serialization, and per-PC drain that the model charges
(Table~\ref{tab:probe-pe-dma}). Three properties stand out. First,
the reported latency increases \emph{monotonically} with hop
count---from \SI{141}{\nano\second} at the local HBM slice to
\SI{677}{\nano\second} at the worst-case remote-CUBE slice---with the
increment per added hop matching the per-router overhead and the
UCIe-link cost in Table~\ref{tab:hw}. Second, the bandwidth-saturation
curves track the wire's serialization model: at \SI{1}{\mebi\byte}
transfers the local PE DMA reaches \SI{100}{\percent} of the per-edge
bandwidth limit, while the longest cross-CUBE path saturates at
\SI{97.8}{\percent}---exactly the gap that the per-edge propagation
and per-router overhead model would predict. Third, the simulator
passes the internal-consistency invariants the probe builds in
(monotonic hop progression; D2H reads at least as long as the
equivalent H2D writes because of the reverse-path acknowledgement;
cross-CUBE best less than cross-CUBE worst). These are properties
that hold \emph{only} when the underlying graph traversal, FIFO
contention, and propagation models are internally
self-consistent---a model error in any of them would surface as a
non-monotonic or under-utilising curve here long before it polluted a
kernel-level measurement.
The known simplifications---FIFO router arbitration (instead of
round-robin), HBM scheduler without write-buffer reordering, no bank
conflict, no refresh or thermal effects, and no upstream
backpressure---are the price of a deterministic, inspectable model.
These simplifications bound the absolute accuracy, but the agreement
with both analytic models and the probe's internal-consistency
checks above indicates that KernBench is sufficiently accurate for
evaluating the \emph{relative} hardware--software design trade-offs
(tiling A vs.\ B, topology X vs.\ Y, with vs.\ without composite
command, mesh vs.\ torus) that are the primary objective of this
work.
\subsection{Modeled hardware configuration}
\label{sec:hw}
Table~\ref{tab:hw} summarizes the hardware configuration
used throughout this report. The intent of this
configuration is not to model a specific product, but to
represent a realistic memory-centric accelerator and to
provide a consistent baseline for evaluating hardware--
software co-design mechanisms.
The compute capability within each CUBE is provisioned
such that inference workloads can effectively saturate
the available HBM bandwidth. The aggregate compute
throughput is therefore balanced against memory-system
bandwidth rather than being intentionally over- or
under-provisioned. Concretely, 8 PEs $\times$
\SI{8}{\tera\flop\per\second} give roughly
\SI{64}{\tera\flop\per\second} of f16 compute per CUBE
against \SI{2048}{\giga\byte\per\second} of HBM bandwidth,
which places the balanced point at about
$\sim$31~FLOP/byte; inference decode kernels typically
operate well below that, so HBM bandwidth---not compute---is
the natural ceiling on this configuration. For reference, the
GQA decode kernels evaluated in \S\ref{sec:gqa} operate at
arithmetic intensity well below this balance point, so their
ceiling is set by HBM and inter-PE traffic rather than by GEMM
throughput. HBM bandwidth
is distributed evenly across the PEs within a CUBE, with
each PE responsible for servicing approximately one-eighth
of the CUBE's memory bandwidth through its dedicated HBM
channels.
The on-chip interconnect is configured using bandwidth
and latency parameters representative of commercially
available mesh-network IPs. The goal is not to study a
particular NoC implementation, but rather to evaluate
kernel behavior under a realistic baseline communication
fabric.
For die-to-die communication, UCIe-A bandwidth
characteristics are used as the reference point for
inter-CUBE links. Inter-SIP communication is modeled
using PCIe-class links. Together, these assumptions
provide a representative communication hierarchy for
evaluating collective communication and distributed
kernel execution.
Unless otherwise stated, all experiments use this
configuration. Workload-specific parameters such as
matrix dimensions, sequence lengths, and collective
payload sizes are introduced in their respective sections.
\begin{table}[t]
\centering
\caption{Modeled hardware configuration}
\label{tab:hw}
\small
\begin{tabular}{@{}ll@{}}
\toprule
\textbf{Parameter} & \textbf{Value} \\
\midrule
\multicolumn{2}{@{}l}{\emph{Hierarchy}} \\
SIPs & 2 (1D ring) \\
CUBEs per SIP & 16 ($4\times4$ mesh) \\
PEs per CUBE & 8 (4 corners $\times$ 2) \\
PEs total & 256 \\
\midrule
\multicolumn{2}{@{}l}{\emph{Processing element (PE)}} \\
GEMM engine peak & \SI{8}{\tera\flop\per\second} (f16) \\
TCM (on-PE) & \SI{16}{\mega\byte}, \SI{512}{\giga\byte\per\second} R/W \\
\quad kernel scratch & \SI{1}{\mega\byte} \\
DMA engines & 1 read + 1 write \\
\textsf{PE\_CPU} fixed cost & \SI{2}{\nano\second} \\
\textsf{PE\_SCHED} fixed cost & \SI{1}{\nano\second} \\
\midrule
\multicolumn{2}{@{}l}{\emph{Memory (per CUBE)}} \\
HBM capacity & \SI{48}{\giga\byte} (8 slices) \\
HBM aggregate BW & \SI{2048}{\giga\byte\per\second} \\
HBM pseudo-channels & 64 (8 per PE), \SI{32}{\giga\byte\per\second} each \\
SRAM (shared) & \SI{32}{\mega\byte}, \SI{128}{\giga\byte\per\second} link \\
HBM burst & \SI{256}{\byte} \\
\midrule
\multicolumn{2}{@{}l}{\emph{Interconnect}} \\
Intra-CUBE NoC link & \SI{256}{\giga\byte\per\second}, \SI{2}{\nano\second}/router \\
Inter-CUBE (UCIe PHY) & \SI{512}{\giga\byte\per\second}, \SI{8}{\nano\second}, XY routing \\
Inter-SIP (PCIe) & \SI{768}{\giga\byte\per\second} per endpoint \\
\midrule
\multicolumn{2}{@{}l}{\emph{Command-issue cost model (defaults)}} \\
FIXED per single-op command & 8 cycles \\
FIXED per composite command & 40 cycles \\
per-byte rate $R$ & 0.0625 cycles/byte (\SI{16}{\byte\per\cycle}) \\
composite size cap & \SI{1024}{\byte} \\
\bottomrule
\end{tabular}
\end{table}
docs/report/1H-codesign-paper/sections/03-gemm.tex
+401
View File
@@ -0,0 +1,401 @@
\section{GEMM Acceleration via the Composite Command}
\label{sec:gemm}
GEMM is the compute core of every transformer block---the QKV
projections, the attention score and context products, and the
feed-forward matrices are all matrix multiplications. On a tiled
accelerator a single logical GEMM expands into many hardware tiles, and
each tile must be read from HBM, fetched into the register file, computed,
stored, and written back. If every one of those tile-stage steps were an
independently issued command, two costs would dominate. First, the host
and the PE control processor would pay a per-command issue overhead
$O(\text{tiles}\times\text{stages})$ times, which for a deep-$K$ reduction
is hundreds to thousands of issues. Second, with stages dispatched
one-at-a-time the scheduler cannot overlap the DMA of the next tile with
the compute of the current one---the pipeline never fills, and the MAC
array sits idle waiting for data. The hardware question is therefore: what
issue mechanism lets a single GEMM saturate the MAC array without drowning
in command overhead?
\subsection{Design}
The answer is the \emph{composite command}. A single command carries
the ordered five-stage tile pipeline (\textsf{DMA\_READ},
\textsf{FETCH}, \textsf{GEMM}, \textsf{STORE}, \textsf{DMA\_WRITE}),
and the PE scheduler splits the payload into hardware tiles
(here $32\times64\times32$), emitting one tile token per tile.
Subsequent stages are reached by \emph{token self-routing} between
the on-PE engines, so a tile flows through the chain without
returning to the scheduler between stages. Because the whole
pipeline is described by one command, the issue cost is paid once
per GEMM rather than once per tile-stage, and the scheduler is free
to keep every stage busy on different tiles simultaneously---tile
$i$'s GEMM overlaps tile $i{+}1$'s DMA read. A multi-operation composite
additionally lets an epilogue (for example a vector-math step) ride the
same tile loop, firing per $K$-tile, per output tile, or once per kernel
according to its declared scope; this is what later allows an attention
kernel to fuse its softmax work into the GEMM pipeline
(\S\ref{sec:gqa}).
\subsection{Results}
We sweep eight GEMM shapes spanning square, tall, wide, and deep-$K$
geometries under two operand-staging variants that bracket the
realistic LLM cases. In \textsf{load\_ref}, the activation $A$ is
pre-staged on chip and only the weight $W$ streams from HBM during
the composite (the ``activation in TCM, weights from HBM'' case
typical of decoding with a small batch). In \textsf{ref\_ref}, both
operands stream from HBM during the composite (the ``cold both
sides'' case, where the activation does not fit on-chip or is itself
the output of a producer composite that wrote back to HBM). Both
variants run the same composite command; only the up-front staging
of $A$ differs.
We evaluate the composite GEMM along two axes that together describe
how well it lands on the hardware: \emph{HBM bandwidth utilization}
(does the workload saturate the per-PE HBM bandwidth, i.e.\ is it
memory-bound on this configuration?) and \emph{achieved GEMM
throughput} (what fraction of the \SI{8}{\tera\flop\per\second}
per-PE peak does the kernel actually deliver). Both metrics use the
composite window as the denominator, so the up-front \textsf{tl.load}
of $A$ in \textsf{load\_ref} is excluded — only HBM traffic and
compute that happen \emph{inside} the composite count.
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_hbm_bw_util.png}
\caption{Per-PE HBM bandwidth utilization during the composite window,
for the two staging variants. The ceiling is the per-PE HBM bandwidth
(\SI{256}{\giga\byte\per\second}; dashed line at \SI{100}{\percent}).
\textsf{ref\_ref} streams both $A$ and $W$ and so always sits at a
higher BW-utilization than \textsf{load\_ref} on the same shape;
both variants saturate ($>\SI{85}{\percent}$) once the kernel
has enough total bytes to keep the link busy (deep-$K$ $K{=}3072$,
or low-reuse output-dominated $M{=}128,K{=}8,N{=}128$). Small or
single-tile shapes do not have enough total traffic to saturate the
link and are bottlenecked by pipeline fill rather than HBM.}
\label{fig:gemm-bw}
\end{figure}
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_per_pe_tflops.png}
\caption{Per-PE GEMM throughput delivered during the composite window
(reference line at the \SI{8}{\tera\flop\per\second} per-PE engine
peak). At the deep-$K$ corner $M{=}32,K{=}3072,N{=}32$ the
activation-pre-staged kernel reaches
\SI{7.18}{\tera\flop\per\second} ($\sim\SI{90}{\percent}$ of peak);
the both-from-HBM variant drops to
\SI{3.83}{\tera\flop\per\second} ($\sim\SI{48}{\percent}$) at the
same shape because the second operand doubles HBM pressure and the
kernel hits the BW ceiling shown in Fig.~\ref{fig:gemm-bw}.
Under-tile shapes ($M{=}K{=}N{=}32$, $M{=}8$, $K{=}8$) are
hard-capped by tile-fill at \SI{50}{\percent}, \SI{25}{\percent},
\SI{12.5}{\percent} of peak regardless of staging.}
\label{fig:gemm-tflops}
\end{figure}
The composite reaches the hardware roofline in two distinct regimes.
\emph{Memory-bound}: $K{=}3072$ deep-$K$ saturates HBM (\SI{88}{\percent}
load\_ref, \SI{94}{\percent} ref\_ref) and the low-reuse
$M{=}128,K{=}8,N{=}128$ corner saturates even harder
($>\SI{96}{\percent}$ in both), because back-to-back output
writes dominate. \emph{Compute-rich}: at the same deep-$K$ shape,
\textsf{load\_ref} delivers \SI{7.18}{\tera\flop\per\second}
($\sim\SI{90}{\percent}$ of the per-PE peak) — the configuration's
clean win. The \emph{distance between the two variants} at the same
shape is the cost of \emph{not} pre-staging $A$: at $K{=}3072$ the
throughput halves (\SI{7.18}{\tera\flop\per\second} $\to$
\SI{3.83}{\tera\flop\per\second}) because the second operand doubles
HBM pressure
and the kernel hits the BW ceiling. This is the operational
take-away — when the activation can live on-chip the composite
delivers near-peak GEMM; when it cannot the BW ceiling dominates and
throughput halves.
Figure~\ref{fig:gemm-util} validates that simulator-measured
efficiency tracks an analytic ideal-pipeline model (within
\SI{1.4}{ppt} across every swept shape), and
Figure~\ref{fig:gemm-stages} resolves the per-stage engine wall-clock
that backs the above interpretation.
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_mac_utilization_theoretical_vs_measured.png}
\caption{GEMM MAC utilization and efficiency, theoretical vs.\ measured
(\textsf{load\_ref} staging). Tile-fill sets the ceiling: under-tile
shapes (marked $\ast$) such as $M{=}K{=}N{=}32$, $M{=}8$, and $K{=}8$
cannot fill the MAC tile and cap at \SI{50}{\percent},
\SI{25}{\percent}, \SI{12.5}{\percent}. For tile-filling shapes,
efficiency climbs with tile count---from
\textasciitilde\SI{15}{\percent} at one tile to \textasciitilde%
\SI{90}{\percent} measured at 48 tiles ($K{=}3072$)---and the
measured bars track the analytic ideal-pipeline prediction within
\SI{1.4}{ppt} across every shape.}
\label{fig:gemm-util}
\end{figure}
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{gemm_stage_breakdown.png}
\caption{Per-stage engine wall-clock (DMA in, Fetch, GEMM, DMA out)
under \textsf{load\_ref} staging. For the deep-$K$ shape ($K{=}3072$,
48 tiles) DMA in, Fetch, and GEMM are all close to
\textasciitilde\SI{785}{\nano\second}, running concurrently in a
tightly pipelined regime while DMA-out is negligible. For the
low-reuse shape ($M{=}128,K{=}8,N{=}128$, 16 output tiles) DMA-out
grows to \textasciitilde\SI{336}{\nano\second} while GEMM compute is
\textasciitilde\SI{262}{\nano\second}---a data-movement-bound regime
where the output write becomes the largest stage.}
\label{fig:gemm-stages}
\end{figure}
\subsection{Analysis and meaning}
The composite command does not manufacture bandwidth or MAC
throughput; it removes the two software-shaped obstacles between a
GEMM and its hardware roofline. By paying the issue cost once per
GEMM rather than once per tile-stage, and by chaining tile-stages
through token self-routing, it lets the scheduler keep every stage
busy on a different tile, so compute-rich shapes actually reach the
efficiency their arithmetic intensity allows
($\sim$\SI{90}{\percent} of GEMM peak at the deep-$K$
\textsf{load\_ref} corner) and data-bound shapes actually reach
their HBM-BW ceiling instead of stalling on command overhead.
The split between \textsf{load\_ref} and \textsf{ref\_ref} also makes
the hardware-software boundary explicit: when the activation fits in
on-chip storage the composite is BW-headroom-rich and saturates the
GEMM engine; when it does not, both operands compete for the same
HBM port and the kernel is BW-bound. This is the lever the GQA
kernel of \S\ref{sec:gqa} reaches for next — keeping the right
working set on-chip so the composite pipeline lands in the
compute-rich regime rather than the BW-bound one.
\subsection{Why composite, and not kernel-orchestrated async loading?}
\label{sec:gemm-vs-async}
A reader familiar with double-buffered GEMM kernels on conventional
hardware may ask: why a hardware-side composite command at all? Why
isn't the obvious kernel-level pattern --- async-load each operand,
overlap with compute, accumulate --- sufficient?
To answer this concretely we contrast composite against two
kernel-orchestrated baselines that have access to the same single-op
primitives the platform exposes (\textsf{tl.load} for async DMA into
TCM, \textsf{tl.dot} for a single-op GEMM command on TCM-resident
operands, \textsf{tl.store} for a DMA write-back). Both baselines
pre-stage activation $A$ identically to \textsf{load\_ref}, so the
window measured is the engine-pipeline window with $A$'s up-front DMA
excluded for all three kernels.
\paragraph{Async-full (naive).} A single \textsf{tl.load(A)} followed by
a single \textsf{tl.load(B)} (async, queued behind $A$),
\textsf{tl.dot(A, B)}, and \textsf{tl.store(out)}. The kernel issues
four commands total. The decisive constraint is that
\textsf{tl.dot}'s \textsf{\_await\_pending(b)} blocks the GEMM
command until the \emph{entire} $B$ has landed in TCM --- there is no
way at the runtime API surface to express ``start computing on
tile 0 of $B$ while tile 1 is still in flight.'' Load-of-$B$ and
GEMM therefore serialize.
\paragraph{Async-tiled (chunked prefetch).} The kernel-level workaround is to
split $B$ along $K$ into \textsf{TILE\_K}-sized chunks, issue async
\textsf{tl.load}s for those chunks, issue one \textsf{tl.dot} per
chunk (each blocking only on its own $b_i$), and accumulate via
\textsf{out = out + tl.dot(...)}. This is the standard
double-buffered-GEMM pattern transcribed to the single-op primitives.
We measure two versions of this kernel that differ only in
\emph{prefetch depth}:
\textbf{depth-2 (TCM-bounded)} keeps at most two B-chunks in flight
at any time by issuing the next \textsf{tl.load} just before each
\textsf{tl.dot}; \textbf{depth-$\infty$ (queue-all)} issues all $N_K$
B-chunk loads up front so the DMA engine has the deepest possible
request queue. For a $K{=}3072$ shape both versions emit
48 $A$-chunk loads + 48 $B$-chunk loads + 48 \textsf{tl.dot}s + 47
elementwise adds + 1 store = 192 host-side commands; the only
difference is the temporal interleaving of B-load and dot dispatches.
\paragraph{Why two depths.} The depth distinction matters because the
async-full kernel and the depth-$\infty$ async-tiled kernel both pin the
\emph{entire} $B$ in TCM simultaneously --- $K \cdot N \cdot 2$
bytes. The on-PE scratch is capped at \SI{1}{\mebi\byte}
(\texttt{topology.yaml: pe\_tcm.kernel\_scratch\_mb=1}), so an LLM-scale
attention $B = K_{\text{KV}} \times d_{\text{head}}$ at
$K_{\text{KV}}=4096, d_{\text{head}}=128$ already needs
\SI{1}{\mebi\byte}, and at $K_{\text{KV}}=8192$ it needs \SI{2}{\mebi\byte}
--- past the cap. async-full and queue-all async-tiled are therefore
not just slower than composite but \emph{architecturally infeasible}
at LLM context length. The depth-2 async-tiled kernel is the only
kernel-level
variant whose peak TCM footprint stays
$O(2 \cdot \textsf{TILE\_K} \cdot N)$ regardless of $K$, the same
order as composite's per-tile streaming buffer. It is the apples-to-apples
comparison.
\paragraph{Per-PE throughput.} Figure~\ref{fig:gemm-async} reports
per-PE TFLOP/s for all four kernels side-by-side. The single-op
fast-path in the dispatch cost model (\S\ref{sec:congestion}) is
enabled, so every single-op command the async kernels emit --- DMA
descriptors and \textsf{tl.dot}/\textsf{tl.add} alike --- is charged the
light 8-cycle \textsf{FIXED}, not the 40-cycle composite control-path
cost; neither async-tiled variant carries an inflated per-command cost.
\paragraph{The $K{=}3072$ corner, concretely.} We take this shape
($M{=}32, K{=}3072, N{=}32$) as the running example throughout the
mechanism discussion because it is the regime where the four
kernels spread the widest --- the deepest $K$ in the sweep maps
onto $K / \textsf{TILE\_K} = 48$ hardware tiles, so per-tile costs
amplify into the largest measurable gap. The work content is
identical for all four kernels: $\sim$6.3 M f16 MACs and
$\sim$386 KiB of $B$ traffic from HBM, which together require
$\sim$786 ns of GEMM-engine compute and $\sim$781 ns of DMA on a
saturated per-PE link. What differs is the number of host commands
the same work is decomposed into --- 2 for composite (one
\textsf{tl.load(A)} plus one composite), 4 for async-full, and 192
for either async-tiled variant (48 $A$-loads + 48 $B$-loads + 48
\textsf{tl.dot}s + 47 elementwise adds + 1 store). The
engine-pipeline-window throughput tracks that decomposition closely:
composite reaches \SI{7.18}{\tera\flop\per\second} (post-overlap
limit, only \SI{10}{\percent} below the \SI{8}{\tera\flop\per\second}
per-PE GEMM peak), async-full \SI{3.91}{\tera\flop\per\second}
(DMA and compute serialize on a single big dot), and both async-tiled
variants $\sim$\SI{2.53}{\tera\flop\per\second} (192 commands' worth of
structural dispatch cost --- even at the light per-command rate ---
accumulates on the wall). The next paragraph attributes those gaps to
specific simulator mechanisms.
\begin{figure*}[t]
\centering
\includegraphics[width=\linewidth]{gemm_composite_vs_async_tflops.png}
\caption{Per-PE achieved TFLOP/s for the same shape sweep run under
four issuance patterns: composite (one command, scheduler streams
per-tile internally), async-full (one \textsf{tl.dot} on
fully-loaded $B$), async-tiled with depth-2 double-buffer
(TCM-bounded; the only kernel-level variant that scales to LLM
context length), and async-tiled with depth-$\infty$
(all B-tiles queued up front; included as a sanity check that the
prefetch depth is \emph{not} what separates the async-tiled kernel from
composite). All curves exclude $A$'s up-front DMA from the
measurement window. Composite wins at every full-tile shape; the
depth-2 and depth-$\infty$ async-tiled kernels deliver
\emph{indistinguishable} throughput (e.g. \SI{2.522}{} vs.\
\SI{2.544}{\tera\flop\per\second} at $K{=}3072$), confirming that
prefetch depth is not the lever --- the structural per-command
dispatch cost is. The gap between composite and the async-tiled
kernels grows with $K_{\text{useful}}$
(\textbf{$\sim$$2.8\times$} at $K{=}3072$, where the async-tiled
kernels emit 192 host commands while composite emits one).
The under-tile corner $M{=}128,K{=}8,N{=}128$ inverts: composite's
per-tile orchestration overhead exceeds the per-tile useful work,
and all three async kernels beat it.}
\label{fig:gemm-async}
\end{figure*}
\paragraph{Decomposing the gap.} Three structural mechanisms separate
composite from the kernel-level baselines, and they layer.
\emph{1. Inter-engine token routing happens below the host-side
dispatch path.} The composite encodes the full
\textsf{DMA\_READ}$\to$\textsf{FETCH}$\to$\textsf{GEMM}$\to$\textsf{STORE}$\to$\textsf{DMA\_WRITE}
pipeline once. The scheduler's tile-feeder loop then emits one
\emph{tile token} per HW tile inside that one composite, and each
token self-routes between engines after each stage finishes. The
per-tile token routing is a scheduler-internal event, not a fresh
host command, so it does not pay the structural CPU dispatch cost.
At $K{=}3072$ the composite emits 48 tile tokens that flow
fully-pipelined through five stages each --- 240 inter-engine
hand-offs total --- behind a single command from the host's point of
view.
\emph{2. \textsf{tl.dot} cannot replicate that per-tile pipeline at the
kernel level.} A single-op GEMM command is handled on the GEMM engine as
a single monolithic compute timeout for the supplied $M{\times}K{\times}N$;
there is no internal token loop that would let a streaming DMA of
$B[i{+}1]$ overlap with the GEMM of $B[i]$ inside one
\textsf{tl.dot}. The user can only recover inter-tile overlap by
emitting one \textsf{tl.dot} per chunk --- which is exactly what the
async-tiled baseline does, at the price of $N$ host-side commands.
\emph{3. The host-side dispatch cost the async-tiled baseline pays is
structural, not modelling slack.} KernBench charges every host-emitted
command a structural CPU dispatch cost $d_{\text{cmd}} = \textsf{FIXED} +
b_{\text{logical}} \cdot R$ (\S\ref{sec:congestion}). The cost model is
deliberately charitable to the async kernels here: only a
\textsf{CompositeCmd} pays the 40-cycle control-path \textsf{FIXED} (it
alone drives a scheduler-built tile-feeder plan), while \emph{every}
single-op command --- the 96 DMA descriptors, the 48 \textsf{tl.dot}s,
and the 47 elementwise adds alike --- pays only the light 8-cycle
\textsf{FIXED}, calibrated to descriptor-ring-push / single-instruction
issue patterns in modern accelerators (NVIDIA Hopper TMA $\sim$1 ISA
cycle, a single \textsf{mma.sync} one instruction, AMD AQL packet writes
$\sim$5--15 cycles). So the async-tiled baseline is \emph{not} penalized
by an inflated per-command cost on \emph{any} of its operations. Yet it
still emits 192 host commands against composite's one, so $\sim$192
light dispatches accumulate to $\sim$\SI{1.5}{\micro\second} of
structural PE\_CPU time that composite never pays --- composite hides
its 48 tiles' 240 inter-engine hand-offs as scheduler-internal events
(mechanism 1). Layered on top, the dependency chain on the running
accumulator serializes the 47 adds on the math engine. The net result
is that the async-tiled kernel runs $\sim$2.8$\times$ slower than
composite at $K{=}3072$ despite getting the inter-chunk overlap right
--- down from $\sim$6.3$\times$ under the earlier uniform-40 cost model,
because D8 removed the per-command overcharge, but \emph{not} closed:
the residual gap is the command-count structure (one composite vs.\ 192
single-ops), not a modelling artifact.
\emph{4. Prefetch depth is not the lever, command count is.} The
depth-2 and depth-$\infty$ async-tiled kernels land within \SI{1}{\percent}
of each other at every measured shape (Figure~\ref{fig:gemm-async}).
This is the diagnostic against a natural objection: ``surely the
async-tiled kernel was just under-prefetching; deepen the queue and the
DMA$\to$GEMM overlap recovers.'' Deepening the prefetch queue
changes \emph{when} DMA descriptors hit the engine but not their
total number or the structural dispatch cost they each pay. The
$\sim$2.8$\times$ gap to composite is not a prefetch-depth gap; it
is the gap between
``one composite command with internal per-tile token routing'' and
``$N_K$ host-side dot/add commands, each charged separately.''
The depth-2 kernel additionally constrains peak TCM occupancy to
$O(2 \cdot \textsf{TILE\_K} \cdot N)$, matching composite's per-tile
streaming buffer; the depth-$\infty$ kernel needs the full $B$ in
TCM, which makes it infeasible at LLM context length even if the
throughput were competitive.
\paragraph{Where the composite advantage doesn't apply.} The shape
$M{=}128,K{=}8,N{=}128$ inverts the picture: composite delivers
\SI{0.66}{\tera\flop\per\second} and both async kernels reach
\SI{1.21}{\tera\flop\per\second}. The reason is consistent with the
analysis above and explicit in the simulator state. Composite emits
16 tile tokens (one per output tile) for this shape, each carrying
$K_{\text{useful}}{=}8$ MACs across a TILE\_K$=64$ pipeline ---
$12.5\%$ of the hardware tile's MAC slots are useful, the rest is
K-padding. The per-tile inter-engine hand-off cost stays the same
regardless. When the per-tile useful work is small enough that
hand-off overhead exceeds the GEMM work itself, a single-op
\textsf{tl.dot} --- which submits one monolithic GEMM command with no
per-tile orchestration --- wins. The take-away is the bound on
composite's value: it amortizes useful per-tile compute, not
padding-dominated under-tile shapes. Real kernels at this corner are
better served by reshape-into-batched-GEMM transforms that move
under-tile $K$ into a tile-filling dimension before reaching the GEMM
engine, which is exactly what the GQA decode kernel of
\S\ref{sec:gqa} does for its $K_{\text{useful}}=\text{head\_dim}=128$
inner reduction.
\paragraph{Summary of the comparison.} Composite is the only one of
the four kernels to combine (a) macro-command dispatch at the host
boundary (amortizing the structural CPU cost across all the work a
single GEMM does), (b) scheduler-internal per-HW-tile streaming of
DMA$\rightleftarrows$compute, and (c) TCM-bounded streaming buffer.
Kernel-orchestrated async kernels can have any two of those, not all
three: async-full pays one host dispatch (a) but forfeits per-tile
overlap (b) and pins all of $B$ in TCM (c); depth-$\infty$ async-tiled
achieves inter-chunk overlap but at $N_K$ host dispatches and
full-$B$ TCM occupancy; depth-2 async-tiled fixes the TCM
footprint (c) but still pays $N_K$ host dispatches. The two corners
where async catches up
(small-$K$ where composite has nothing useful to amortize;
under-tile shapes where per-tile useful work is sub-token) are
diagnostic of \emph{where composite is the wrong tool}, not of a
slack the user kernel could close.
docs/report/1H-codesign-paper/sections/04-allreduce.tex
+288
View File
@@ -0,0 +1,288 @@
\section{PE\_IPCQ and Collective Communication}
\label{sec:allreduce}
Distributing a transformer across devices turns every tensor-parallel
layer into a collective: partial results computed on different PEs, CUBEs,
and SIPs must be summed and redistributed with an all-reduce. Underneath
the algorithm this is fundamentally a PE-to-PE problem---many short
messages flowing between neighbors as the reduction proceeds. The natural
software realization is a per-direction ring buffer whose head and tail
pointers the producer and consumer update atomically and poll, but our
H2 2025 report measured this scheme end-to-end and found that the
atomic-pointer traffic together with the consumer's polling loop dominate
the per-message cost: the queue itself becomes the bottleneck well before
the link runs out of bandwidth, and a pure-SW collective spends most of
its time on metadata rather than on actually moving partials. A second,
orthogonal problem is link sharing---if the collective rides the kernel's
generic DMA path it competes with the GEMM's compute traffic on the same
wires, so a large tile transfer head-of-line-blocks a pending reduction.
This section asks how a dedicated hardware primitive can lift queue
management off the software path entirely and, in the same design, stop
collective traffic from stalling behind compute traffic, so the all-reduce
runs overlapped with compute and at the interconnect's physical limit.
\subsection{Design}
The proposed block is \textbf{PE\_IPCQ} (inter-PE communication queue),
a small controller dropped into every PE next to PE\_DMA, PE\_GEMM,
PE\_MATH, and PE\_TCM. It is a control-plane block---it holds no
payload data---and consists of three pieces: a per-direction
\emph{QPair register file} (\textasciitilde\SI{576}{\byte} of flip-flops
covering up to eight directions), a combinational slot-address
generator and backpressure comparator, and a credit injector/receiver
wired to the NoC. Each QPair holds the local pointers
(\texttt{my\_head}, \texttt{my\_tail}), shadowed views of the peer's
(\texttt{peer\_head\_cache}, \texttt{peer\_tail\_cache}), the local and
peer rx-buffer physical bases, ring depth and slot size (both
power-of-two), and the peer's credit-target address. The ring data
itself lives in a reserved slot region of TCM (or PE-local HBM, or
cube-shared SRAM), addressed by the QPair registers; PE\_IPCQ never
touches the bytes, only the pointers. The PE\_CPU sees the controller
as an MMIO peripheral and the N/S/E/W direction labels as logical
ports---one kernel image runs across a 1D ring, 2D mesh, or 2D torus
because the topology only changes which peers the QPair registers
point at.
\emph{Initialization.} The host CCL backend brings the whole machine
to a usable state before any kernel runs.
\texttt{init\_process\_group(backend="ahbm")} loads \texttt{ccl.yaml},
resolves the algorithm + topology + buffer\_kind + slot configuration,
allocates an rx ring-buffer region on every participating PE so every
rank now knows every other rank's \texttt{rx\_base\_pa}, and fans out
an \texttt{IpcqInitMsg} that writes the QPair register file on each
PE\_IPCQ over MMIO. The same fan-out wires the per-direction
credit-return channel: each PE\_IPCQ records its peer's credit-target
address and a back-pointer that the credit receiver will use to update
the right QPair. After init, every register the runtime needs is
preloaded; nothing in the kernel path allocates memory, walks a table,
or talks to the host.
\emph{Send.} When the kernel executes \texttt{tl.send(dir="E",
src\_addr, nbytes)}, the PE\_CPU performs a single MMIO write into
PE\_IPCQ describing the request (direction, source address/space,
length, sender handle). The controller evaluates backpressure in one
combinational compare---\texttt{(my\_head $-$ peer\_tail\_cache) $<$
n\_slots}---and, on a hit, the slot-address generator returns
\texttt{dst = peer\_rx\_base\_pa + (my\_head \% n\_slots) $\times$
slot\_size} in one to two cycles. PE\_IPCQ then emits one
\texttt{IpcqDmaToken} on its dedicated port to PE\_DMA's
\texttt{vc\_comm} channel, carrying the data descriptor (\texttt{src,
dst, nbytes}) plus a small piggyback header (\texttt{sender\_seq =
my\_head, src\_coord, direction}). \texttt{my\_head} is incremented in
the same cycle---a local flip-flop bump, not a cross-PE atomic---and
\texttt{tl.send} returns fire-and-forget. If the backpressure compare
fails, the controller stalls the CPU in either of two modes selected
at init: \texttt{poll} (CPU re-reads a status CSR) or \texttt{sleep}
(controller asserts a wake event when a credit arrives). Both modes
are benchmarked in the results.
\emph{Transport.} PE\_DMA is the only block that touches the fabric,
and it has been extended for IPCQ in two ways. First, it now exposes
two virtual channels---\texttt{vc\_compute} for GEMM/Math
\texttt{TileToken}s and \texttt{vc\_comm} for \texttt{IpcqDmaToken}s
---with independent state machines and a chunk-level
(\SI{256}{\byte}) weighted round-robin arbiter on the shared physical
link. A large GEMM tile DMA can no longer monopolize the link against
a pending IPCQ send, enabling compute and communication to make
progress independently as an architectural property of the design. Second, on the sender side PE\_DMA packs the piggyback
header into the same flit train as the data, and on the receiver side
it runs the I6 \emph{atomic terminal handler}: pay the per-flit
bottleneck-BW drain, then, in one indivisible block, (i) write the
payload into \texttt{MemoryStore} at \texttt{dst\_addr} (the
receiver's ring slot) and (ii) forward an \texttt{IpcqMetaArrival}
$\{\texttt{sender\_seq, dst\_addr}\}$ to the local PE\_IPCQ over a
PE-internal wire. No yield, no PE\_CPU interaction, no cache-coherence
round trip---the bytes land in TCM and the metadata reaches the
controller in the same cycle.
\emph{Receive and credit return.} PE\_IPCQ's Meta Extractor
range-matches the incoming \texttt{dst\_addr} against each direction's
$[\texttt{rx\_base}, \texttt{rx\_base} + \texttt{n\_slots} \times
\texttt{slot\_size})$ window---unambiguous even when two directions
share a peer, as in a 2-rank bidirectional ring---and updates
\texttt{peer\_head\_cache[d] := max(prev, sender\_seq + 1)}, releasing
any \texttt{tl.recv} blocked on direction $d$. When the kernel
eventually consumes the slot, the controller increments
\texttt{my\_tail} and emits a \SI{16}{\byte} credit packet on the
dedicated fast-path channel preloaded at init; the latency is the
real fabric path (per-node overhead + edge propagation + 16 B /
bottleneck BW), so an in-cube credit returns faster than a cross-SIP
credit and the model carries no magic constants. The sender's PE\_IPCQ
absorbs the credit, advances \texttt{peer\_tail\_cache}, and
de-asserts backpressure if it was stalled. The pointer-synchronization
problem the software baseline solved with explicit atomic RMWs and
polling loops has been split in two and dissolved into existing
traffic: \emph{head} updates ride on the DMA payload itself, with the
receiver's PE\_DMA performing the data + metadata write in a single
atomic step, so the sender never blocks waiting for the receiver to
see it; \emph{tail} updates ride on a dedicated 16 B credit on a
side-channel, with no software in the loop. Send becomes a single
MMIO write, receive becomes a flip-flop read, backpressure becomes
one 64-bit subtract, and the per-message cost in IPCQ is dominated by
the link traversal rather than by metadata bookkeeping---which is
what the H2 2025 measurement showed the pure-software queue could
not achieve. On top of this substrate the collective runs a
hierarchical local-reduce / global all-reduce-broadcast schedule
across whatever inter-device topology the configuration specifies,
with the IPCQ ring buffer placed in on-PE TCM, PE-local HBM, or
cube-shared SRAM---the third knob the results section sweeps.
\subsection{Design alternatives}
\label{sec:ipcq-alternatives}
PE\_IPCQ is one point in a small space of hardware mechanisms for
moving a short message from one PE to a neighbor and signalling its
arrival. Three established alternatives anchor the space, each the
HW realization of a familiar host-networking idea: a \emph{doorbell +
polling} scheme (the classic MMIO doorbell---write the payload by DMA,
write a doorbell, let the peer poll or take an interrupt); a
\emph{hardware message queue} (HMQ, the NVLink-style descriptor engine
that pushes a queue entry to the peer, with large payloads still
riding a second DMA); and a \emph{completion-queue} design (RDMA-CQ,
the InfiniBand/RoCE pattern where a DMA write auto-posts a completion
entry the peer's CQ polls). PE\_IPCQ is the fourth: a hardware ring
with credit return, splitting the control plane into PE\_IPCQ and the
data plane into PE\_DMA, with head updates riding the payload and tail
updates riding a 16\,B side-channel credit (\S\ref{sec:allreduce}).
\begin{figure*}[t]
\centering
\includegraphics[width=0.78\linewidth]{ipcq_alternatives_architecture_flow.png}
\caption{Per-send data and control flow for the four PE-to-PE
signalling mechanisms (sender\,$\rightarrow$\,NoC\,$\rightarrow$\,receiver).
Doorbell and RDMA-CQ each issue two fabric transactions (payload then
doorbell / completion) and leave the peer polling or taking an
interrupt; HMQ adds a dedicated descriptor engine but still moves large
payloads on a second DMA; PE\_IPCQ folds head-pointer signalling into
the payload flit train and returns the tail credit on a side channel,
so a send is one MMIO write and a receive is a flip-flop read. This is
a \emph{design schematic}, not a measured comparison.}
\label{fig:ipcq-arch}
\end{figure*}
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{ipcq_alternatives_decision_matrix.png}
\caption{Why the ring+credit design was chosen, across five criteria:
single-send latency, whether the host CPU sits on the critical path,
whether the receiver must poll or take a wake-up interrupt, whether the
control and data datapaths are duplicated, and whether the mechanism is
right-sized for single-owner PE-to-PE traffic (rather than a
multi-tenant fabric). PE\_IPCQ is the only design that clears every
criterion. The accompanying per-send step-count tally
($\sim$28 control events for IPCQ versus $\sim$38 for HMQ, $\sim$53 for
RDMA-CQ, and $\sim$56 for doorbell+polling) is an \emph{illustrative}
order-of-magnitude comparator over hand-counted pipeline steps---not a
simulator measurement. The measured, simulator-grounded results follow
in the next subsection.}
\label{fig:ipcq-decision}
\end{figure}
The qualitative comparison motivates the design but is not a
quantitative claim: the cycle-step tallies above are hand-counted
control events, deliberately separated from the measured latencies that
follow. Everything in the results subsection runs on the PE\_IPCQ
substrate and is simulator-grounded.
\subsection{Results}
All measurements in this section run on the PE\_IPCQ substrate
described above; the topology sweep is intended to characterize how
effectively the proposed mechanism exposes the underlying interconnect's
properties, not to compare PE\_IPCQ against an alternative
communication primitive. Following the milestone-evaluation convention,
the collective sweep builds its own six-device (six-SIP, $2\times3$)
configurations---distinct from the two-SIP default of
Table~\ref{tab:hw}---and measures all-reduce latency as a function of
payload size for three inter-device topologies: a 1D ring, a 2D mesh
(no wrap), and a 2D torus (Figure~\ref{fig:allreduce-topo}).
Table~\ref{tab:allreduce} and Figure~\ref{fig:allreduce-cmp} report the
result.
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{allreduce_topology.png}
\caption{The three six-device ($2\times3$) inter-device topologies the
collective sweep runs over, and the hierarchical local-reduce /
global all-reduce-broadcast schedule mapped onto each: a 1D ring, a 2D
mesh (no wrap-around), and a 2D torus (wrap-around links on both axes).
The torus's wrap links shorten the worst-case reduction path, which is
what the latency sweep below rewards.}
\label{fig:allreduce-topo}
\end{figure}
\begin{table}[t]
\centering
\caption{All-reduce latency (ns) across six devices, by topology and
per-PE payload. Lower is better; the torus wins at every size.}
\label{tab:allreduce}
\small
\begin{tabular}{@{}rrrr@{}}
\toprule
\textbf{Bytes/PE} & \textbf{2D mesh} & \textbf{Ring 1D} & \textbf{2D torus} \\
\midrule
256 & 4189 & 3883 & 2957 \\
4{,}096 & 5566 & 5240 & 4031 \\
16{,}384 & 10016 & 9376 & 7396 \\
65{,}536 & 27821 & 25925 & 20858 \\
98{,}304 & 39690 & 36957 & 29833 \\
\bottomrule
\end{tabular}
\end{table}
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{allreduce_comparison.png}
\caption{All-reduce latency vs.\ per-PE payload for the three PE\_IPCQ
topologies, against the analytic torus model. The isolated point in
the top panel (\SI{366}{\micro\second}) is the only data point
available from the H2 2025 software-queue measurement campaign---an
intra-device all-reduce across 16 CUBEs on a single device. A true
6-device inter-SIP measurement under the same software queue was not
collected, so this point in fact \emph{understates} the SW-queue cost
for the multidevice configuration KernBench measures here; even so the
proposed PE\_IPCQ inter-device curves sit roughly an order of
magnitude below it, which is the headline SW-vs-HW comparison this
section makes. The measured torus tracks the analytic curve within a
small constant factor, the gap reflecting real link serialization that
the analytic model idealizes away.}
\label{fig:allreduce-cmp}
\end{figure}
\begin{figure}[t]
\centering
\includegraphics[width=\linewidth]{allreduce_buffer_kind.png}
\caption{Effect of IPCQ staging-buffer placement (2D torus). At
\SI{64}{\kibi\byte}/PE, TCM staging (\SI{20858}{\nano\second}) beats HBM
(\SI{24074}{\nano\second}) by \textasciitilde\SI{13}{\percent} and SRAM
(\SI{33194}{\nano\second}) by \textasciitilde\SI{37}{\percent}; at small
payloads the three are indistinguishable.}
\label{fig:allreduce-buf}
\end{figure}
Three findings follow. First, topology matters and the effect grows with
size: at \SI{96}{\kilo\byte}/PE the 2D torus is \SI{25}{\percent} faster
than the mesh and \SI{19}{\percent} faster than the ring, because its
wrap-around links shorten the worst-case reduction path. Second, the
measured curves grow smoothly with payload and stay within a small
constant factor of the analytic torus model, with the gap reflecting real
link serialization that the analytic model idealizes away. Third, where
the IPCQ staging buffer lives is a first-order knob: keeping it in on-PE
TCM is materially faster than HBM or shared SRAM at large payloads
(Figure~\ref{fig:allreduce-buf}).
\subsection{Analysis and meaning}
PE\_IPCQ turns the collective from an off-device, contention-prone
operation into an on-device primitive whose latency tracks the
interconnect's physical limits. The control/data split keeps the engine
small while reusing PE\_DMA for movement, and the virtual-channel split is
what lets a reduction proceed without stalling the compute DMA that feeds
the very GEMMs producing the partials---the same property the fused
attention kernel relies on when it interleaves KV reduction with score
computation (\S\ref{sec:gqa}). For the hardware roadmap the results argue
two things: provisioning wrap-around (torus) inter-device links is worth
roughly a \SI{20}{}--\SI{25}{\percent} collective-latency reduction at
scale, and giving the IPCQ a fast on-PE staging buffer (TCM) rather than
forcing it through HBM or SRAM is worth another \SI{14}{}--%
\SI{38}{\percent} at large payloads.
docs/report/1H-codesign-paper/sections/05-gqa.tex
+26
View File
@@ -0,0 +1,26 @@
\section{Fused Grouped-Query Attention}
\label{sec:gqa}
Attention is the 1H focus, and it is where the two preceding optimizations
have to come together. Grouped-Query Attention (GQA) shrinks the KV cache
by sharing each KV head across a group of query heads (here $h_q=8$ query
heads to $h_{kv}=1$ KV head, a group factor $G=8$), which makes decoding
feasible at long context but also makes it acutely memory-bound: a decode
step processes a single query position ($T_q=1$) against the entire KV
history, so its arithmetic intensity is low and its time is dominated by
streaming the KV cache out of HBM. FlashAttention-style tiling with an
online-softmax merge avoids ever materializing the full score matrix, but
realizing it as a fast \emph{fused} kernel needs both building blocks from
this report: efficient GEMM issue (\S\ref{sec:gemm}) for the
$Q\!\cdot\!K^{\top}$ and $P\!\cdot\!V$ products, and an efficient on-device
reduction (\S\ref{sec:allreduce}) for the multi-user and
sequence-parallel KV reductions. \emph{Fused} here is meant in the
FlashAttention sense---$Q\!\cdot\!K^{\top}$, the online softmax, and
$P\!\cdot\!V$ collapse into a single kernel that never materializes the
score matrix---and, beyond that, the cross-device KV reduction is absorbed
into the same kernel (on PE\_IPCQ) rather than issued as a separate
all-reduce. This section is the capstone: the fused
kernel that uses the composite command and PE\_IPCQ at the same time.
Multi-head attention (MHA) was studied in prior work and serves here as
the established baseline rather than being re-derived.
docs/report/1H-codesign-paper/sections/05a-roofline.tex
+85
View File
@@ -0,0 +1,85 @@
\subsection{Roofline Analysis: Batch and Context Regimes}
\label{sec:roofline}
Before deploying any sharding strategy, the workload's position on the
decode roofline sets what \emph{can} be improved and what \emph{cannot}.
Two quantities matter: the machine's \textbf{critical batch}
$B^\ast = C \cdot b / (2 W)$ (where $C$ is per-PE FLOPs, $W$ is per-PE
HBM bandwidth, $b$ is bytes per parameter), and its \textbf{balance
context} $L^\ast = 2 N / (\mathrm{AI} \cdot \mathrm{kv\_bpt})$ (where
$N$ is active parameters and $\mathrm{AI} = C/W$ is arithmetic
intensity). Above $B^\ast$ the deployment leaves the memory-bound
regime; above $L^\ast$ per-user KV streaming dominates and no amount
of batching hides it. For Llama-3.1-70B on the default AHBM machine
($C \!=\! \SI{8}{TFLOPS}$/PE, $W\!=\! \SI{256}{GB/s}$/PE), we get
$B^\ast \!\approx\! 31$ and $L^\ast \!\approx\! \SI{13.4}{K}$ tokens.
\paragraph{Short-context regime ($S_{kv} \!<\! L^\ast$).}
Figure~\ref{fig:roofline-short} shows one decode step at
$S_{kv}\!=\! \SI{8}{K}$. The step-latency panel (left) makes it
obvious that weight fetch dominates at low $B$: at $B\!=\!1$, one
step spends $\sim\!\SI{535}{ms}$ streaming weights and only
$\SI{28}{ms}$ on per-sequence compute and KV combined. The
cost-per-token panel (right) is where the batch story lives: dividing
step latency by $B$ shrinks the weight term as $1/B$, so per-token
cost drops from $\SI{562}{ms}$ at $B\!=\!1$ to $\SI{29.7}{ms}$ at
$B\!=\!256$ --- a $19\times$ reduction. This is the memory-bound-to-
compute-bound crossover: past $B^\ast$, weight cost is fully
amortized and per-token time asymptotes to the compute + KV floor
($\sim\!\SI{27}{ms}$). \textbf{At short context, batching directly
lowers cost per token.}
\begin{figure*}[h]
\centering
\includegraphics[width=\textwidth]{roofline_short_context.png}
\caption{Roofline decomposition for Llama-3.1-70B decode at
$S_{kv}\!=\! \SI{8}{K}$ (short, well below $L^\ast\!\approx\! \SI{13.4}{K}$).
\textbf{Left}: step latency vs.\ batch $B$; weight fetch is flat
(one HBM sweep per step), compute and KV grow linearly with $B$.
\textbf{Right}: per-token cost ($=$ step $\div B$); the $1/B$
weight-fetch term amortizes rapidly, dropping total per-token time
from \SI{562}{ms} at $B\!=\!1$ to \SI{29.7}{ms} at $B\!=\!256$
--- a $19\times$ throughput gain from batching alone.}
\label{fig:roofline-short}
\end{figure*}
\paragraph{Long-context regime ($S_{kv} \!\gg\! L^\ast$).}
Figure~\ref{fig:roofline-long} shows the same decomposition at
$S_{kv}\!=\! \SI{1}{M}$ ($\sim\!78 \times L^\ast$). The step-latency
panel shows KV fetch has swelled by two orders of magnitude:
per-token KV cost is now $\sim\!\SI{1342}{ms}$, dwarfing both
weight ($\SI{535}{ms}$ at $B\!=\!1$) and compute ($\SI{17}{ms}$).
The cost-per-token panel is the punchline: increasing $B$ still
shrinks the weight term but leaves the giant KV floor untouched,
because \textbf{KV fetch is per-sequence} --- adding another user
adds a full extra copy of their KV read. Total per-token cost falls
only from \SI{1894}{ms} at $B\!=\!1$ to \SI{1361}{ms} at $B\!=\!256$
--- a mere $28\%$ reduction despite $256\times$ the batch.
\textbf{At long context, batching stops paying because per-user KV
streaming dominates.}
\begin{figure*}[h]
\centering
\includegraphics[width=\textwidth]{roofline_long_context.png}
\caption{Same decomposition at $S_{kv}\!=\! \SI{1}{M}$ (long,
$\sim\!78\times L^\ast$). KV fetch has grown to \SI{1342}{ms}/token
and is now the dominant term at every $B$. Batching only shrinks
the weight component; the KV floor is unmovable because each new
user brings a full per-sequence KV read. Total per-token cost falls
just $28\%$ from $B\!=\!1$ to $B\!=\!256$, versus $19\times$ at
short context.}
\label{fig:roofline-long}
\end{figure*}
\paragraph{Implication for deployment.} The two regimes call for
opposite strategies. In the short-context regime, the operator packs
$B$ as high as HBM allows to sit on the compute floor --- this is
where cost-per-token is minimized and hardware utilization is
highest. In the long-context regime, per-user KV is the binding
resource; batching offers little benefit, so the operator instead
shrinks $\mathrm{kv\_bpt}$ (GQA / MQA / MLA, INT4 KV, sparse
attention) and shards the sequence dimension itself (CP), routing
long-context requests to a dedicated pool with more chips per user.
The next two subsections (\S\ref{sec:capacity-planning},
\S\ref{sec:parallelism-selection}) turn these regime observations
into concrete sizing and sharding rules.
docs/report/1H-codesign-paper/sections/05b-capacity-planning.tex
+118
View File
@@ -0,0 +1,118 @@
\subsection{Capacity Planning: LLM Serving on AHBM}
\label{sec:capacity-planning}
Once a fused kernel is chosen, the deployment question is orthogonal:
\emph{how many cubes (or SIPs) does a given model, context length, and
latency SLO require?} Hyperscalers decide this along three axes that
must all be satisfied simultaneously. The total PE count is
$\max(A, B, C) \times N_{\text{replicas}}$, where $A$ is the capacity
floor, $B$ the KV-cache headroom, and $C$ the throughput SLO
(Table~\ref{tab:capacity-axes}).
\begin{table}[h]
\centering
\small
\caption{Three-axis sizing. $N$ = active params, $b$ = bytes per param,
$\mathrm{HBM}_{\mathrm{PE}}$ = per-PE HBM budget, $S_{kv}$ = context
length, $\mathrm{kv\_bpt}$ = KV cache bytes per token per user
($=2 \cdot h_{kv} \cdot d_{\text{head}} \cdot b \cdot L$).}
\label{tab:capacity-axes}
\begin{tabular}{@{}l l@{}}
\toprule
Axis & Formula \\
\midrule
A. Capacity floor & $\lceil N b / \mathrm{HBM}_{\mathrm{PE}} \rceil$ \\
B. KV headroom & $\lceil (N b + u \cdot S_{kv} \cdot \mathrm{kv\_bpt}) / \mathrm{HBM}_{\mathrm{PE}} \rceil$ \\
C. Throughput SLO & $N_{\text{replicas}} = \lceil n_{\text{users}} / B_{\mathrm{SLO}} \rceil$ \\
\bottomrule
\end{tabular}
\end{table}
\paragraph{SLO targets.} The per-token latency budget $B_{\mathrm{SLO}}$
follows the use case. TTFT (Time To First Token) is dominated by
prefill; TPOT (Time Per Output Token) by one decode step
(Table~\ref{tab:slo-targets}).
\begin{table}[h]
\centering
\small
\caption{Typical SLO targets. Voice needs a tight per-token cadence;
chat balances both; batch is priced on throughput not latency.}
\label{tab:slo-targets}
\begin{tabular}{@{}l l l@{}}
\toprule
Workload & TTFT & TPOT \\
\midrule
Voice / real-time & 200--300 ms & 10--25 ms \\
Interactive chat & 300 ms -- 1 s & 20--50 ms \\
Batch / offline & seconds -- minutes & N/A \\
\bottomrule
\end{tabular}
\end{table}
\paragraph{Regime-aware rules of thumb.} Context length relative to the
machine's balance point $L^*=2N/(\mathrm{AI}\cdot\mathrm{kv\_bpt})$
determines which strategy applies. Below $L^*$ the deployment is
compute-bound and batch scales to $\sim\!2B^*$; above, KV streaming
dominates and $B$ must shrink (Table~\ref{tab:regime-rules}).
\begin{table*}[t]
\centering
\small
\caption{Deployment strategy by context regime.}
\label{tab:regime-rules}
\begin{tabular}{@{}l l l@{}}
\toprule
Regime & Batch & Cost/token \\
\midrule
Short ($S_{kv} < L^*$) & $B \approx 2 B^*$ & low \\
Long ($S_{kv} > L^*$) & smaller $B$, more CP & higher \\
Extreme ($S_{kv} \gg L^*$) & dedicated pool, disagg.\ prefill & much higher \\
\bottomrule
\end{tabular}
\end{table*}
\paragraph{Sample deployment templates.} A single base model is
typically routed to one of several sharding tiers depending on the
request's expected context length; the API gateway dispatches to the
appropriate replica pool (Table~\ref{tab:deploy-templates}).
\begin{table*}[t]
\centering
\small
\caption{Sample deployment tiers for one base model. TP is fixed to
match $h_{kv}$; CP grows with context; $B$ shrinks to hold TPOT.}
\label{tab:deploy-templates}
\begin{tabular}{@{}l c c c c@{}}
\toprule
Tier & CP$\times$TP & Cubes/repl. & Max $S_{kv}$ & Typical $B$ \\
\midrule
Small & $1 \times 8$ & 1 & 32k & 64 \\
Medium & $4 \times 8$ & 4 & 128k & 32 \\
Large & $32 \times 8$ & 32 & 1M & 4 \\
\bottomrule
\end{tabular}
\end{table*}
\paragraph{What to do when each axis binds.} The dominant axis dictates
the first lever to pull (Table~\ref{tab:axis-playbook}). Axes $A$ and
$C$ scale nearly linearly with hardware; axis $B$ is the one where
algorithmic techniques (GQA, MLA, INT4 KV, sparse attention) buy the
most because they attack $\mathrm{kv\_bpt}$ directly rather than adding
chips.
\begin{table*}[t]
\centering
\small
\caption{Playbook per binding axis.}
\label{tab:axis-playbook}
\begin{tabular}{@{}l l l@{}}
\toprule
Binding & First lever & Second lever \\
\midrule
A. Capacity & $\uparrow$ TP / PP & bigger-HBM chip; FP8 / INT4 weights \\
B. KV & $\uparrow$ CP; $\downarrow B$ & GQA / MQA / MLA; INT4 KV; sparse attn \\
C. Throughput & $\uparrow$ replicas (DP) & loosen SLO; smaller model; spec.\ decoding \\
\bottomrule
\end{tabular}
\end{table*}
docs/report/1H-codesign-paper/sections/05c-parallelism-selection.tex
+91
View File
@@ -0,0 +1,91 @@
\subsection{Parallelism Selection: TP $\times$ CP $\times$ PP $\times$ DP $\times$ EP}
\label{sec:parallelism-selection}
Given a capacity-feasible layout (\S\ref{sec:capacity-planning}), the
remaining decision is \emph{how to shard}: which axes to enable and at
what degrees. Memory is the feasibility filter; once you fit, the
design problem is minimising \emph{exposed} communication (communication
that could not be overlapped with compute). The core principle is to
\textbf{rank techniques by how much they communicate per unit of
compute, and assign the chattiest ones to the fastest links.}
\begin{table*}[h]
\centering
\small
\caption{Parallelism axes at a glance. TP AllReduce volume does not
shrink with degree — only compute does — which sets its practical
ceiling near a fast-link (NVLink / intra-SIP) domain.}
\label{tab:parallelism-compare}
\begin{tabular}{@{}l l l l l@{}}
\toprule
Axis & Shards & Comm pattern & Frequency & Hard ceiling \\
\midrule
DP & optimizer state (w/ ZeRO) & AllReduce (grads) & 1$\times$/step & critical batch \\
TP & weights + acts + KV heads & AllReduce & 2$\times$/layer & fast-link domain; $\lesssim h_{kv}$ \\
PP & weights + KV by layer & P2P activation hand-off & per stage boundary & $n_{\text{layers}}$ \\
CP & activations + KV by position & Ring P2P per hop & per layer (overlappable) & $\mathrm{seq\_len} / \mathrm{block}$ \\
EP & expert weights (MoE only) & All-to-all & 2$\times$/MoE layer & $n_{\text{experts}}$ \\
\bottomrule
\end{tabular}
\end{table*}
\paragraph{Symptom-driven axis selection.} The dominant problem dictates
the first knob to try (Table~\ref{tab:parallelism-problem}).
\begin{table}[h]
\centering
\small
\caption{Which axis to reach for first, by dominant problem.}
\label{tab:parallelism-problem}
\begin{tabular}{@{}l l@{}}
\toprule
Dominant problem & First axis \\
\midrule
Weight memory doesn't fit & TP within a fast domain \\
KV of ONE long sequence & CP \\
KV for MANY sequences & DP replicas \\
Single-request TPOT & TP ($\lesssim 8$) \\
Aggregate throughput & Outer DP \\
Model spans multiple nodes & TP intra-node, PP inter-node \\
MoE expert weight memory & EP \\
\bottomrule
\end{tabular}
\end{table}
\paragraph{When to add, when to stop.} Every axis has a natural
saturation point past which further degree hurts more than it helps
(Table~\ref{tab:parallelism-criteria}). Sizing decisions are almost
always constrained by two axes simultaneously (e.g.\ TP by NVLink
domain and $h_{kv}$; CP by ring-hop hiding and per-rank block size).
\begin{table*}[t]
\centering
\small
\caption{Add / stop criteria per axis.}
\label{tab:parallelism-criteria}
\begin{tabular}{@{}l l l@{}}
\toprule
Axis & Add when & Stop when \\
\midrule
DP & more independent requests & global batch $>$ critical \\
TP & weights need sharding or TPOT tight & collectives dominate; local GEMMs shrink \\
PP & depth too large; TP would cross slow links & bubble $>$ $\sim$10\% \\
CP & one sequence needs position sharding & ring hops can't overlap compute \\
EP & MoE expert memory & expert GEMMs too small; routing imbalance \\
\bottomrule
\end{tabular}
\end{table*}
\paragraph{Common misconceptions.} Three widely repeated rules survive
contact with real workloads only partially:
\emph{(i)} memory is not just a feasibility filter but a continuous
performance variable --- more free HBM enables larger $B$, more prefix
cache, higher throughput even after weights fit;
\emph{(ii)} the ``TP $\leq h_{kv}$'' ceiling is an efficiency
preference, not a correctness constraint --- vLLM and others support
KV-head replication and head-dim splitting for TP $>$ $h_{kv}$;
\emph{(iii)} ``always start at TP=8'' is disproven by public
disclosures (DeepSeek-V3 inference uses TP=4, some MLPs at TP=1;
DeepSeek-V3 training uses PP=16 + EP=64 with no TP), which shows the
right starting point is the smallest TP that fits memory and meets
latency, followed by measurement.

Some files were not shown because too many files have changed in this diff Show More