Compare commits
102 Commits
77eece81c4
..
master
| Author | SHA1 | Date | |
|---|---|---|---|
| fc4747668e | |||
| bfe3e0e8d1 | |||
| 55653cbb3e | |||
| 5ce7a5b30b | |||
| bc5e704572 | |||
| c39bbb16aa | |||
| 77abc95d78 | |||
| f54822d247 | |||
| b7c4c680aa | |||
| 6410c0843c | |||
| 97e765b58f | |||
| bda76cbd66 | |||
| a1c3c28f62 | |||
| 9b4aee83da | |||
| 221097ef08 | |||
| c99a238826 | |||
| 04c7d247f0 | |||
| 1ddd1baa7a | |||
| 5e92b89821 | |||
| 3360be2488 | |||
| adc14c84af | |||
| 9c5062bdfb | |||
| 4ded67a443 | |||
| 47c4f40f4d | |||
| 50aab8d591 | |||
| fc1dcdde24 | |||
| fd45bd2408 | |||
| 674e194dc9 | |||
| 770aed6291 | |||
| 1bce080d79 | |||
| 5d7ef48b14 | |||
| 2a0bc26db0 | |||
| 6f61657ba4 | |||
| da1909fdf2 | |||
| ed28eea156 | |||
| f1cf0257d3 | |||
| afe35ee0b5 | |||
| 84abf847c4 | |||
| a7f39ade7e | |||
| c2b42999d8 | |||
| 9afeb6bc16 | |||
| bccc90db82 | |||
| 1a6d52c58a | |||
| 09721040ca | |||
| bf5b659a3d | |||
| b8e1a3322f | |||
| 91b63eb9f6 | |||
| 55c20bd1a6 | |||
| 6ef09dd6f5 | |||
| 2834383700 | |||
| 85f6716fc1 | |||
| 09e9331026 | |||
| f9d0077472 | |||
| 65f358fdfa | |||
| 9fdde44922 | |||
| ae131aa2af | |||
| baf22f01ab | |||
| 13f5cbc317 | |||
| 8c108154a7 | |||
| 711a9a257f | |||
| cb6d21f145 | |||
| d88e6cc72e | |||
| edb30326ce | |||
| d0db9ff40e | |||
| 23a2bc3fb3 | |||
| eec61838b5 | |||
| f8caf2e16c | |||
| 09e70bd516 | |||
| e20e1f7362 | |||
| 3437d7e079 | |||
| 3da116d40d | |||
| f7aa8134b9 | |||
| a455ae61b4 | |||
| 1971ac731c | |||
| a98e93110b | |||
| f08dda3bd4 | |||
| 7d90d53d4d | |||
| faf011dae5 | |||
| 1dade267ff | |||
| 24c705419e | |||
| cd6f0ed91d | |||
| e9a5c438e3 | |||
| bb668a3ac3 | |||
| 0662c76fa7 | |||
| 37dc48fd27 | |||
| 7c45047e93 | |||
| 73e0b315fe | |||
| 1552028c25 | |||
| 56c51b184a | |||
| 8102ddbe30 | |||
| 84bb418e1e | |||
| dd3337f2e4 | |||
| b6315c3c90 | |||
| 821bbf26a2 | |||
| 2d8271c981 | |||
| 0c6ca0aaed | |||
| b64c43b947 | |||
| e8b7f4f064 | |||
| 23992548f7 | |||
| 92b9221533 | |||
| 694e0cc9b9 | |||
| 6b6e29968a |
@@ -252,6 +252,74 @@ composite (가장 큰 것이 decode opt2 의 `#2` ~322 bytes) 보다 훨씬 위
|
||||
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 표 유지
|
||||
|
||||
@@ -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).
|
||||
@@ -184,8 +184,15 @@ The empty-`ops` form is the legacy single-op path.
|
||||
- `DMA_READ`: `simpy.Resource(capacity=1)`.
|
||||
- `DMA_WRITE`: `simpy.Resource(capacity=1)`.
|
||||
- Both channels run concurrently (READ ∥ WRITE allowed).
|
||||
- Within a channel, requests serialize (READ ∥ READ disallowed; same
|
||||
for WRITE).
|
||||
- Within a channel, requests serialize **at the issue path** (READ ∥ READ
|
||||
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
|
||||
ADR-0023 D8 — out of scope for this ADR.
|
||||
|
||||
|
||||
@@ -270,6 +270,80 @@ 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
|
||||
|
||||
|
Before Width: | Height: | Size: 38 KiB After Width: | Height: | Size: 38 KiB |
|
Before Width: | Height: | Size: 36 KiB After Width: | Height: | Size: 36 KiB |
@@ -1,13 +1,13 @@
|
||||
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,1024,2048,2717.2783333333473
|
||||
hbm,torus_2d,6,8192,16384,7315.184999999989
|
||||
hbm,torus_2d,6,32768,65536,23081.26500000037
|
||||
sram,torus_2d,6,128,256,2060.040000000012
|
||||
sram,torus_2d,6,1024,2048,2909.2783333333473
|
||||
sram,torus_2d,6,8192,16384,9523.184999999869
|
||||
sram,torus_2d,6,32768,65536,32201.265000000385
|
||||
tcm,torus_2d,6,128,256,1964.040000000012
|
||||
tcm,torus_2d,6,1024,2048,2477.2783333333473
|
||||
tcm,torus_2d,6,8192,16384,6403.185000000109
|
||||
tcm,torus_2d,6,32768,65536,19865.265000000378
|
||||
hbm,torus_2d,6,128,256,2345.040000000012
|
||||
hbm,torus_2d,6,1024,2048,2942.2783333333473
|
||||
hbm,torus_2d,6,8192,16384,7540.184999999989
|
||||
hbm,torus_2d,6,32768,65536,23306.26500000037
|
||||
sram,torus_2d,6,128,256,2285.040000000012
|
||||
sram,torus_2d,6,1024,2048,3134.2783333333527
|
||||
sram,torus_2d,6,8192,16384,9748.184999999869
|
||||
sram,torus_2d,6,32768,65536,32426.265000000385
|
||||
tcm,torus_2d,6,128,256,2189.040000000012
|
||||
tcm,torus_2d,6,1024,2048,2702.2783333333473
|
||||
tcm,torus_2d,6,8192,16384,6628.18500000005
|
||||
tcm,torus_2d,6,32768,65536,20090.265000000378
|
||||
|
||||
|
|
Before Width: | Height: | Size: 75 KiB After Width: | Height: | Size: 77 KiB |
|
Before Width: | Height: | Size: 37 KiB After Width: | Height: | Size: 37 KiB |
|
Before Width: | Height: | Size: 86 KiB After Width: | Height: | Size: 92 KiB |
@@ -1,37 +1,37 @@
|
||||
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,32,64,1024,2747.7400000000152
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,64,128,2048,2855.990000000018
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,128,256,4096,3072.490000000019
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,512,1024,16384,3337.1133333333582
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,1024,2048,32768,3708.0333333333692
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,2048,4096,65536,4449.873333333393
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,4096,8192,131072,5933.020000000124
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,8192,16384,262144,8900.379999999863
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,16384,32768,524288,14835.099999999224
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,32768,65536,1048576,26704.540000000765
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,49152,98304,1572864,38573.97999999701
|
||||
lrab_hierarchical_allreduce,ring_1d,6,8,16,256,2365.255833333347
|
||||
lrab_hierarchical_allreduce,ring_1d,6,32,64,1024,2436.9433333333473
|
||||
lrab_hierarchical_allreduce,ring_1d,6,64,128,2048,2532.526666666683
|
||||
lrab_hierarchical_allreduce,ring_1d,6,128,256,4096,2723.693333333349
|
||||
lrab_hierarchical_allreduce,ring_1d,6,512,1024,16384,3048.635000000021
|
||||
lrab_hierarchical_allreduce,ring_1d,6,1024,2048,32768,3393.4016666666957
|
||||
lrab_hierarchical_allreduce,ring_1d,6,2048,4096,65536,4082.401666666714
|
||||
lrab_hierarchical_allreduce,ring_1d,6,4096,8192,131072,5458.80166666677
|
||||
lrab_hierarchical_allreduce,ring_1d,6,8192,16384,262144,8216.934999999943
|
||||
lrab_hierarchical_allreduce,ring_1d,6,16384,32768,524288,13733.201666665835
|
||||
lrab_hierarchical_allreduce,ring_1d,6,32768,65536,1048576,24765.73500000064
|
||||
lrab_hierarchical_allreduce,ring_1d,6,49152,98304,1572864,35798.268333331536
|
||||
lrab_hierarchical_allreduce,torus_2d,6,8,16,256,1700.6025000000095
|
||||
lrab_hierarchical_allreduce,torus_2d,6,32,64,1024,1753.2900000000102
|
||||
lrab_hierarchical_allreduce,torus_2d,6,64,128,2048,1823.540000000012
|
||||
lrab_hierarchical_allreduce,torus_2d,6,128,256,4096,1964.040000000012
|
||||
lrab_hierarchical_allreduce,torus_2d,6,512,1024,16384,2196.8183333333463
|
||||
lrab_hierarchical_allreduce,torus_2d,6,1024,2048,32768,2477.2783333333473
|
||||
lrab_hierarchical_allreduce,torus_2d,6,2048,4096,65536,3038.1983333333583
|
||||
lrab_hierarchical_allreduce,torus_2d,6,4096,8192,131072,4159.5050000000665
|
||||
lrab_hierarchical_allreduce,torus_2d,6,8192,16384,262144,6403.185000000109
|
||||
lrab_hierarchical_allreduce,torus_2d,6,16384,32768,524288,10890.5449999995
|
||||
lrab_hierarchical_allreduce,torus_2d,6,32768,65536,1048576,19865.265000000378
|
||||
lrab_hierarchical_allreduce,torus_2d,6,49152,98304,1572864,28839.98500000059
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,8,16,256,2918.5525000000157
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,32,64,1024,2999.740000000016
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,64,128,2048,3107.990000000019
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,128,256,4096,3324.4900000000207
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,512,1024,16384,3589.1133333333582
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,1024,2048,32768,3960.0333333333692
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,2048,4096,65536,4701.873333333393
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,4096,8192,131072,6185.020000000124
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,8192,16384,262144,9152.379999999861
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,16384,32768,524288,15087.099999999224
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,32768,65536,1048576,26956.540000000765
|
||||
lrab_hierarchical_allreduce,mesh_2d_no_wrap,6,49152,98304,1572864,38825.97999999701
|
||||
lrab_hierarchical_allreduce,ring_1d,6,8,16,256,2628.2558333333477
|
||||
lrab_hierarchical_allreduce,ring_1d,6,32,64,1024,2699.943333333348
|
||||
lrab_hierarchical_allreduce,ring_1d,6,64,128,2048,2795.526666666683
|
||||
lrab_hierarchical_allreduce,ring_1d,6,128,256,4096,2986.693333333351
|
||||
lrab_hierarchical_allreduce,ring_1d,6,512,1024,16384,3311.635000000021
|
||||
lrab_hierarchical_allreduce,ring_1d,6,1024,2048,32768,3656.4016666666957
|
||||
lrab_hierarchical_allreduce,ring_1d,6,2048,4096,65536,4345.401666666714
|
||||
lrab_hierarchical_allreduce,ring_1d,6,4096,8192,131072,5721.801666666768
|
||||
lrab_hierarchical_allreduce,ring_1d,6,8192,16384,262144,8479.934999999887
|
||||
lrab_hierarchical_allreduce,ring_1d,6,16384,32768,524288,13996.201666665835
|
||||
lrab_hierarchical_allreduce,ring_1d,6,32768,65536,1048576,25028.73500000064
|
||||
lrab_hierarchical_allreduce,ring_1d,6,49152,98304,1572864,36061.26833333154
|
||||
lrab_hierarchical_allreduce,torus_2d,6,8,16,256,1925.6025000000104
|
||||
lrab_hierarchical_allreduce,torus_2d,6,32,64,1024,1978.290000000011
|
||||
lrab_hierarchical_allreduce,torus_2d,6,64,128,2048,2048.540000000012
|
||||
lrab_hierarchical_allreduce,torus_2d,6,128,256,4096,2189.040000000012
|
||||
lrab_hierarchical_allreduce,torus_2d,6,512,1024,16384,2421.8183333333463
|
||||
lrab_hierarchical_allreduce,torus_2d,6,1024,2048,32768,2702.2783333333473
|
||||
lrab_hierarchical_allreduce,torus_2d,6,2048,4096,65536,3263.1983333333583
|
||||
lrab_hierarchical_allreduce,torus_2d,6,4096,8192,131072,4384.5050000000665
|
||||
lrab_hierarchical_allreduce,torus_2d,6,8192,16384,262144,6628.18500000005
|
||||
lrab_hierarchical_allreduce,torus_2d,6,16384,32768,524288,11115.5449999995
|
||||
lrab_hierarchical_allreduce,torus_2d,6,32768,65536,1048576,20090.265000000378
|
||||
lrab_hierarchical_allreduce,torus_2d,6,49152,98304,1572864,29064.985000000597
|
||||
|
||||
|
|
Before Width: | Height: | Size: 194 KiB After Width: | Height: | Size: 194 KiB |
|
Before Width: | Height: | Size: 40 KiB After Width: | Height: | Size: 29 KiB |
|
Before Width: | Height: | Size: 46 KiB After Width: | Height: | Size: 32 KiB |
|
Before Width: | Height: | Size: 42 KiB After Width: | Height: | Size: 31 KiB |
|
Before Width: | Height: | Size: 53 KiB After Width: | Height: | Size: 52 KiB |
|
Before Width: | Height: | Size: 52 KiB After Width: | Height: | Size: 45 KiB |
|
Before Width: | Height: | Size: 51 KiB After Width: | Height: | Size: 53 KiB |
|
Before Width: | Height: | Size: 52 KiB After Width: | Height: | Size: 47 KiB |
|
Before Width: | Height: | Size: 137 KiB After Width: | Height: | Size: 137 KiB |
@@ -1,81 +1,81 @@
|
||||
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,raw,33.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,raw,36.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,raw,37.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,raw,38.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,raw,40.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,raw,42.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,raw,50.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,raw,66.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,raw,98.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,raw,114.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,raw,47.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,raw,51.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,raw,52.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,raw,53.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,raw,55.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,raw,57.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,raw,65.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,raw,81.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,raw,113.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,raw,129.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,raw,89.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,raw,95.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,raw,96.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,raw,97.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,raw,99.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,raw,102.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,raw,114.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,raw,138.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,raw,186.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,raw,210.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,raw,111.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,raw,119.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,raw,120.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,raw,121.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,raw,123.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,raw,126.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,raw,138.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,raw,162.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,raw,210.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,raw,234.54000000000087
|
||||
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,51.57999999999811
|
||||
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,54.07999999999811
|
||||
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,55.07999999999811
|
||||
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,56.07999999999811
|
||||
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,58.07999999999811
|
||||
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,60.07999999999811
|
||||
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,68.07999999999811
|
||||
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,84.08000000000175
|
||||
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,116.08000000000175
|
||||
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,132.08000000000175
|
||||
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,65.18999999999505
|
||||
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,69.18999999999505
|
||||
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,70.18999999999505
|
||||
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,71.18999999999505
|
||||
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,73.18999999999505
|
||||
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,75.18999999999505
|
||||
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,83.18999999999505
|
||||
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,99.19000000000233
|
||||
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,131.19000000000233
|
||||
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,147.19000000000233
|
||||
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,107.28999999999724
|
||||
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,113.53999999999724
|
||||
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,114.53999999999724
|
||||
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,115.53999999999724
|
||||
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,117.53999999999724
|
||||
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,120.53999999999724
|
||||
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,132.53999999999724
|
||||
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,156.54000000000087
|
||||
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,204.54000000000087
|
||||
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,228.54000000000087
|
||||
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,129.28999999999724
|
||||
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,137.53999999999724
|
||||
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,138.53999999999724
|
||||
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,139.53999999999724
|
||||
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,141.53999999999724
|
||||
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,144.53999999999724
|
||||
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,156.53999999999724
|
||||
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,180.54000000000087
|
||||
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,228.54000000000087
|
||||
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,252.54000000000087
|
||||
|
||||
|
@@ -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.
|
||||
|
Before Width: | Height: | Size: 86 KiB After Width: | Height: | Size: 92 KiB |
@@ -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>
|
||||
|
After Width: | Height: | Size: 30 KiB |
|
After Width: | Height: | Size: 75 KiB |
|
After Width: | Height: | Size: 32 KiB |
|
Before Width: | Height: | Size: 40 KiB After Width: | Height: | Size: 29 KiB |
|
Before Width: | Height: | Size: 46 KiB After Width: | Height: | Size: 32 KiB |
|
After Width: | Height: | Size: 29 KiB |
|
Before Width: | Height: | Size: 42 KiB After Width: | Height: | Size: 34 KiB |
|
After Width: | Height: | Size: 58 KiB |
|
After Width: | Height: | Size: 51 KiB |
|
After Width: | Height: | Size: 74 KiB |
|
After Width: | Height: | Size: 45 KiB |
|
After Width: | Height: | Size: 149 KiB |
|
After Width: | Height: | Size: 73 KiB |
|
After Width: | Height: | Size: 169 KiB |
|
After Width: | Height: | Size: 47 KiB |
|
After Width: | Height: | Size: 135 KiB |
|
After Width: | Height: | Size: 123 KiB |
|
After Width: | Height: | Size: 132 KiB |
|
After Width: | Height: | Size: 183 KiB |
|
After Width: | Height: | Size: 182 KiB |
|
After Width: | Height: | Size: 145 KiB |
|
After Width: | Height: | Size: 72 KiB |
|
After Width: | Height: | Size: 96 KiB |
|
After Width: | Height: | Size: 133 KiB |
|
After Width: | Height: | Size: 70 KiB |
|
After Width: | Height: | Size: 73 KiB |
|
After Width: | Height: | Size: 75 KiB |
|
After Width: | Height: | Size: 84 KiB |
|
After Width: | Height: | Size: 253 KiB |
|
After Width: | Height: | Size: 305 KiB |
|
After Width: | Height: | Size: 381 KiB |
|
After Width: | Height: | Size: 189 KiB |
|
After Width: | Height: | Size: 160 KiB |
|
After Width: | Height: | Size: 98 KiB |
|
After Width: | Height: | Size: 106 KiB |
@@ -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>
|
||||
|
After Width: | Height: | Size: 10 KiB |
@@ -15,6 +15,7 @@
|
||||
\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}
|
||||
@@ -28,7 +29,7 @@
|
||||
\date{
|
||||
\small
|
||||
AGI Computing Lab, System Technology Group\\
|
||||
2026 H1 Report
|
||||
2026 Q1-Q3 Report
|
||||
}
|
||||
|
||||
\begin{document}
|
||||
@@ -39,9 +40,16 @@ AGI Computing Lab, System Technology Group\\
|
||||
\input{sections/02-platform}
|
||||
\input{sections/03-gemm}
|
||||
\input{sections/04-allreduce}
|
||||
\input{sections/05-gqa}
|
||||
\input{sections/06-discussion}
|
||||
\input{sections/07-conclusion}
|
||||
\input{sections/08-future-work}
|
||||
\input{sections/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}
|
||||
|
||||
@@ -25,9 +25,12 @@ 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. 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
|
||||
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
|
||||
|
||||
@@ -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, 전체 모델 최적화에 관한 향후 과제를
|
||||
제시한다.
|
||||
@@ -1,16 +1,15 @@
|
||||
\begin{figure*}[t]
|
||||
\centering
|
||||
\begin{subfigure}[b]{0.42\textwidth}
|
||||
\begin{subfigure}[b]{0.495\textwidth}
|
||||
\centering
|
||||
\includegraphics[width=\linewidth,height=0.6\linewidth,keepaspectratio]{sip_architecture.pdf}
|
||||
\caption{SIP level: $4{\times}4$ CUBE mesh + IO chiplet.}
|
||||
\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.6\linewidth,keepaspectratio]{cube_architecture.pdf}
|
||||
\caption{CUBE level (zoom-in of one CUBE in (a)): 8 PEs, HBM
|
||||
channels, M\_CPU, SRAM, NoC router mesh, UCIe links.}
|
||||
\includegraphics[width=\linewidth,height=0.84\linewidth]{cube_architecture.pdf}
|
||||
\caption{CUBE Architecture}
|
||||
\label{fig:cube-arch}
|
||||
\end{subfigure}
|
||||
|
||||
@@ -26,9 +25,12 @@
|
||||
on-device collective control plane.}
|
||||
\label{fig:pe-arch}
|
||||
\end{subfigure}
|
||||
\caption{Modeled hardware graph at the SIP, CUBE, and PE levels
|
||||
(one example configuration; specific parameters in \S\ref{sec:hw} /
|
||||
Table~\ref{tab:hw}). KernBench is not tied to this particular
|
||||
\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.}
|
||||
@@ -144,8 +146,20 @@ KernBench is layered along the flow of a request:
|
||||
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.
|
||||
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}
|
||||
@@ -185,7 +199,7 @@ 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 nanosecond in a reported latency
|
||||
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.
|
||||
|
||||
@@ -202,6 +216,69 @@ 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}
|
||||
|
||||
@@ -225,26 +302,6 @@ 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.
|
||||
|
||||
\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,
|
||||
\]
|
||||
where $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. The
|
||||
default anchoring (\textsf{FIXED} $= 40$ cycles, $R = 0.0625$
|
||||
cycles/byte, i.e.\ \SI{16}{\byte\per\cycle}, at \SI{1}{\giga\hertz})
|
||||
places a typical composite 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.
|
||||
|
||||
\subsection{Accuracy}
|
||||
\label{sec:accuracy}
|
||||
@@ -254,41 +311,30 @@ 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 roughly
|
||||
\SIrange{10}{20}{\percent} across a wide range of tile counts; the
|
||||
residual gap is attributable to pipeline-fill and DMA effects the
|
||||
analytic 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. A
|
||||
single-device point from an external full-system simulator (FSIM) at
|
||||
the largest payload sits an order of magnitude above the KernBench
|
||||
multi-device torus, illustrating the well-known gap between an
|
||||
achievable-kernel number and a full end-to-end-stack number rather
|
||||
than a model error. 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
|
||||
external full-system simulator data 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.
|
||||
translates into physically reasonable kernel latencies.
|
||||
|
||||
A third source of confidence comes from directly probing the
|
||||
simulator's per-traversal behaviour. 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
|
||||
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
|
||||
@@ -308,26 +354,19 @@ 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.
|
||||
|
||||
\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*}
|
||||
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}
|
||||
@@ -344,10 +383,22 @@ 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. 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.
|
||||
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
|
||||
@@ -371,7 +422,7 @@ payload sizes are introduced in their respective sections.
|
||||
|
||||
\begin{table}[t]
|
||||
\centering
|
||||
\caption{Modeled hardware configuration (shared by all experiments).}
|
||||
\caption{Modeled hardware configuration}
|
||||
\label{tab:hw}
|
||||
\small
|
||||
\begin{tabular}{@{}ll@{}}
|
||||
@@ -389,11 +440,12 @@ 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 \\
|
||||
CPU / scheduler overhead & \SI{2}{\nano\second} / \SI{1}{\nano\second} \\
|
||||
\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{1024}{\giga\byte\per\second} \\
|
||||
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} \\
|
||||
@@ -404,7 +456,8 @@ Inter-CUBE (UCIe PHY) & \SI{512}{\giga\byte\per\second}, \SI{8}{\nano\second}, X
|
||||
Inter-SIP (PCIe) & \SI{768}{\giga\byte\per\second} per endpoint \\
|
||||
\midrule
|
||||
\multicolumn{2}{@{}l}{\emph{Command-issue cost model (defaults)}} \\
|
||||
FIXED per command & 40 cycles \\
|
||||
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
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
\section{GEMM Acceleration via the Composite Command}
|
||||
\label{sec:gemm}
|
||||
|
||||
\subsection{Why it is needed}
|
||||
|
||||
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
|
||||
@@ -21,20 +19,17 @@ in command overhead?
|
||||
|
||||
\subsection{Design}
|
||||
|
||||
The answer is the \emph{composite command}. A single command carries the
|
||||
ordered tile pipeline
|
||||
\[
|
||||
\textsf{DMA\_READ}\rightarrow\textsf{FETCH}\rightarrow\textsf{GEMM}
|
||||
\rightarrow\textsf{STORE}\rightarrow\textsf{DMA\_WRITE},
|
||||
\]
|
||||
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 DMA\,$\rightarrow$\,fetch\,$\rightarrow$\,GEMM\,$%
|
||||
\rightarrow$\,store 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
|
||||
(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
|
||||
@@ -45,66 +40,362 @@ kernel to fuse its softmax work into the GEMM pipeline
|
||||
\subsection{Results}
|
||||
|
||||
We sweep eight GEMM shapes spanning square, tall, wide, and deep-$K$
|
||||
geometries, under three operand-staging variants
|
||||
(\textsf{ref\_ref}, both operands streamed from HBM; \textsf{load\_ref},
|
||||
one operand resident in TCM; \textsf{load\_load}, both resident).
|
||||
Figure~\ref{fig:gemm-util} reports MAC utilization and efficiency, and
|
||||
Figure~\ref{fig:gemm-stages} breaks the kernel into per-stage engine
|
||||
busy time.
|
||||
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.
|
||||
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{23}{\percent} at one tile to \textasciitilde%
|
||||
\SI{78}{\percent} measured at 48 tiles ($K{=}3072$)---and the measured
|
||||
bars track the analytic prediction within
|
||||
\SIrange{10}{20}{\percent}.}
|
||||
\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)
|
||||
the Fetch and GEMM stages are large and comparable
|
||||
(\textasciitilde\SI{770}{} and \SI{785}{\nano\second}) while DMA-out is
|
||||
negligible---a compute-rich, well-pipelined regime. For the low-reuse
|
||||
shape ($M{=}128,K{=}8,N{=}128$) DMA-out grows to
|
||||
\textasciitilde\SI{350}{\nano\second} and compute is small---a
|
||||
data-movement-bound regime.}
|
||||
\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}
|
||||
|
||||
Two regularities stand out. (i) Utilization is governed by how completely
|
||||
the problem fills the MAC tile: the three under-tile shapes are hard-capped
|
||||
well below \SI{100}{\percent}, independent of how the kernel is issued.
|
||||
(ii) Among tile-filling shapes, efficiency is governed by tile count---more
|
||||
tiles amortize the one-time pipeline fill and issue cost, so the deep-$K$
|
||||
shape reaches \textasciitilde\SI{78}{\percent} of peak while a single-tile
|
||||
shape reaches only \textasciitilde\SI{23}{\percent}. The stage breakdown
|
||||
explains why: with 48 tiles the GEMM and Fetch stages overlap and stay
|
||||
busy, whereas the low-reuse shape spends most of its time moving data in
|
||||
and out.
|
||||
|
||||
\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 issue cost once per GEMM and chaining tile-stages
|
||||
through token self-routing, it lets the scheduler keep the pipeline full,
|
||||
so compute-rich shapes actually reach the efficiency their arithmetic
|
||||
intensity allows ($\sim$\SI{78}{\percent} measured at 48 tiles), and
|
||||
data-bound shapes actually reach their DMA bound instead of stalling on
|
||||
command overhead. The close agreement between measured and theoretical
|
||||
efficiency (Figure~\ref{fig:gemm-util}) is also the report's primary
|
||||
validation that KernBench's latency model is faithful in the regime that
|
||||
matters. The hardware implication is concrete: a single-command,
|
||||
self-routing tile pipeline is the issue mechanism that makes the MAC array
|
||||
usable, and it is a prerequisite---not a luxury---for the fused attention
|
||||
kernel of \S\ref{sec:gqa}.
|
||||
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.
|
||||
|
||||
@@ -1,46 +1,215 @@
|
||||
\section{PE\_IPCQ and Collective Communication}
|
||||
\label{sec:allreduce}
|
||||
|
||||
\subsection{Why it is needed}
|
||||
|
||||
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. If that
|
||||
collective is handled by the host or by a generic DMA path, three problems
|
||||
appear. The reduction traffic competes with the kernel's own compute DMA
|
||||
on the same links, causing head-of-line blocking; there is no efficient
|
||||
peer-to-peer ring primitive, so data takes extra hops; and the ordering is
|
||||
hard to make deterministic. The hardware question is how to perform
|
||||
collectives \emph{on the device}, overlapped with compute and reproducible
|
||||
run-to-run.
|
||||
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}
|
||||
|
||||
KernBench models a dedicated per-PE collective engine, \textbf{PE\_IPCQ}
|
||||
(inter-PE communication queue). It is a control-plane block: it owns the
|
||||
ring-buffer address arithmetic, head/tail pointers, peer-pointer caches,
|
||||
backpressure, and the four-direction (N/S/E/W) neighbor map, with eight
|
||||
ring buffers per PE (four directions $\times$ \{tx, rx\}). Crucially,
|
||||
PE\_IPCQ does \emph{not} move data itself---it delegates the actual
|
||||
transfer to PE\_DMA, keeping a clean control/data split. To stop
|
||||
collective traffic from blocking compute, PE\_DMA is extended into a
|
||||
two-channel virtual-channel model: \texttt{vc\_compute} carries tile
|
||||
load/store for GEMM and vector math, \texttt{vc\_comm} carries IPCQ sends,
|
||||
each with an independent state machine. The same physical link is shared
|
||||
but progresses in chunks (\SI{256}{\byte}), so a large GEMM DMA does not
|
||||
lock the link end-to-end against a pending reduction. 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.
|
||||
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}
|
||||
|
||||
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. Table~\ref{tab:allreduce}
|
||||
and Figure~\ref{fig:allreduce-cmp} report the result.
|
||||
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
|
||||
@@ -64,13 +233,19 @@ per-PE payload. Lower is better; the torus wins at every size.}
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{allreduce_comparison.png}
|
||||
\caption{All-reduce latency vs.\ per-PE payload for the three topologies,
|
||||
against the analytic torus model and an external full-system simulator
|
||||
(FSIM) reference. The measured torus tracks the analytic curve within a
|
||||
small constant factor; the FSIM single-device point
|
||||
(\SI{366}{\micro\second}) sits an order of magnitude above the
|
||||
KernBench algorithmic latency, illustrating the difference between an
|
||||
achievable-kernel number and a full end-to-end-stack number.}
|
||||
\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}
|
||||
|
||||
|
||||
@@ -1,8 +1,6 @@
|
||||
\section{Fused Grouped-Query Attention}
|
||||
\label{sec:gqa}
|
||||
|
||||
\subsection{Why it is needed}
|
||||
|
||||
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
|
||||
@@ -16,120 +14,13 @@ 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. This section is the capstone: the fused
|
||||
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.
|
||||
|
||||
\subsection{Design}
|
||||
|
||||
The fused GQA kernel issues its matrix products as scheduler-managed
|
||||
composite commands and keeps the online-softmax merge and the cross-device
|
||||
KV reduction inside the kernel, on PE\_IPCQ. Two kernel families cover the
|
||||
two phases. The \emph{prefill} kernel is head-parallel and rotates the KV
|
||||
shards around an inter-CUBE ring (``Ring KV''). The \emph{decode} kernel
|
||||
is head-replicated with a statically sharded KV cache and reduces partial
|
||||
attention outputs through an M-fold intra-CUBE chain and, for multiple
|
||||
users, a two-level reduce-to-root. Two further primitives make long
|
||||
context practical: a \emph{lazy load} that issues the KV \textsf{DMA\_READ}
|
||||
and returns immediately, auto-waiting only at first use so that KV load
|
||||
overlaps score computation; and per-tile \emph{scratch recycling} that
|
||||
keeps the running softmax accumulators ($m,\ell,O$) in a persistent arena
|
||||
while freeing per-tile temporaries, so the kernel fits the
|
||||
\SI{1}{\mebi\byte} scratch budget across many tiles. A further refinement
|
||||
that restructures the decode step into two stateful composites (a named
|
||||
\textsf{softmax\_merge} recipe) is designed but not yet wired into the
|
||||
measured path; results below reflect the implemented kernel only.
|
||||
|
||||
\subsection{Results: long-context decode and parallelism strategies}
|
||||
|
||||
The four headline panels above stress the kernel at moderate context
|
||||
lengths. Long-context decode---the regime where KV cache size, not
|
||||
attention compute, sets serving cost---turns the choice of how to
|
||||
parallelize across cubes and PEs into a first-order design knob. We
|
||||
compare four strategies on the LLaMA-3.1-70B single-KV-head-group
|
||||
target (8 CUBEs $\times$ 8 PEs, one KV-head group):
|
||||
|
||||
\begin{itemize}\setlength\itemsep{1pt}
|
||||
\item \textbf{Case 1} (Cube-SP $\times$ PE-TP): KV split by $S_{kv}$
|
||||
across CUBEs; PEs tensor-parallel on the batch dimension
|
||||
(wastes PE-TP work at $B{=}1$).
|
||||
\item \textbf{Case 2} (Cube-Repl $\times$ PE-TP): full KV
|
||||
replicated to every CUBE; PEs tensor-parallel on batch.
|
||||
\item \textbf{Case 3} (Cube-Repl $\times$ PE-SP): full KV
|
||||
replicated; PEs sequence-parallel on $S_{kv}$ with an
|
||||
intra-CUBE all-reduce.
|
||||
\item \textbf{Case 4} ($\star$, Cube-SP $\times$ PE-SP): KV split
|
||||
64-way (across both CUBEs and PEs) with a two-phase all-reduce
|
||||
on the running softmax state $(m, \ell, O)$.
|
||||
\end{itemize}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_latency.png}
|
||||
\caption{End-to-end decode latency per parallelism strategy
|
||||
(LLaMA-3.1-70B single-KV-head group, 8 CUBEs $\times$ 8 PEs).
|
||||
Replication into CUBEs (Cases 2/3) wins the latency race
|
||||
(\SI{20.2}{\micro\second} for Case 3), but Case~4 ($\star$, KV split
|
||||
64-way) finishes within \SI{14}{\micro\second} of the leader while
|
||||
paying a different cost---visible in Figure~\ref{fig:gqa-4cases-mem}.}
|
||||
\label{fig:gqa-4cases-lat}
|
||||
\end{figure}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_memory.png}
|
||||
\caption{Per-CUBE KV memory footprint for the four cases. Cases 1
|
||||
and 4---both with the KV cache split across cubes (Cube-SP)---hold
|
||||
only \SI{0.5}{\mebi\byte} of KV state per CUBE; Cases 2 and 3, which
|
||||
replicate the full KV, hold \SI{4}{\mebi\byte} per CUBE, an
|
||||
\textbf{8$\times$} blowup at this configuration that scales linearly
|
||||
with context length.}
|
||||
\label{fig:gqa-4cases-mem}
|
||||
\end{figure}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_4cases_traffic.png}
|
||||
\caption{Per-case op-count breakdown. The replicated-KV PE-TP design
|
||||
(Case 2) avoids almost all on-device communication
|
||||
(\textasciitilde0 IPCQ copies), but at the cost of KV memory.
|
||||
Case~4's two-phase reduce charges \textasciitilde190 IPCQ copies and
|
||||
\textasciitilde190 DMA reads---this is the traffic that PE\_IPCQ
|
||||
(\S\ref{sec:allreduce}) is built to absorb at on-device speed.}
|
||||
\label{fig:gqa-4cases-traffic}
|
||||
\end{figure}
|
||||
|
||||
Three things stand out. First, the fastest case in pure latency
|
||||
(Case~3, \SI{20.2}{\micro\second}) is also the most memory-hungry,
|
||||
requiring the full KV state on every CUBE---an option that fails to
|
||||
scale once context length blows past the per-CUBE budget. Second,
|
||||
Case~4's KV-split design gives back roughly \SI{14}{\micro\second}
|
||||
versus Case~3 in exchange for an \textbf{8$\times$} KV-memory
|
||||
reduction; for practical long-context serving where KV capacity is
|
||||
the binding constraint, this is the trade the design chooses
|
||||
(marked $\star$). Third, Case~4 pays its way in
|
||||
\emph{communication}: the op-count panel shows \textasciitilde190
|
||||
IPCQ copies and \textasciitilde190 DMA reads, precisely the
|
||||
on-device collective traffic that PE\_IPCQ and the torus links of
|
||||
\S\ref{sec:allreduce} are provisioned to move quickly---so the
|
||||
``slower'' strategy is in fact the one that fully cashes in the
|
||||
communication-side codesign work of this report.
|
||||
|
||||
\subsection{Analysis and meaning}
|
||||
|
||||
These panels are the clearest statement of the codesign thesis in the
|
||||
report. Because the composite command keeps GEMM issue cheap and the MAC
|
||||
array barely occupied, the fused attention kernel's latency is set almost
|
||||
entirely by data movement: streaming the KV cache and reducing partials
|
||||
across devices. That is precisely the cost that the communication-side
|
||||
work targets---PE\_IPCQ for the on-device reduction, the lazy load for
|
||||
load/compute overlap, fast TCM staging and torus links for the reduction
|
||||
itself. In other words, the two enablers are not independent features that
|
||||
happen to appear in the same kernel; the GEMM optimization is what
|
||||
\emph{exposes} the data-movement bottleneck (by removing the compute and
|
||||
issue overhead that would otherwise hide it), and the communication
|
||||
optimization is what \emph{attacks} it. For an attention-dominated decoder
|
||||
the meaningful hardware investments are therefore the ones that move data
|
||||
faster and reduce it on-device---not additional MAC throughput, which this
|
||||
workload cannot use.
|
||||
|
||||
@@ -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.
|
||||
@@ -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*}
|
||||
@@ -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.
|
||||
@@ -0,0 +1,515 @@
|
||||
\subsection{Data Placement Policy}
|
||||
\label{sec:gqa-placement}
|
||||
|
||||
Long-context decode is bound by the KV cache, so the first-order design
|
||||
question is how to place that cache---and the running softmax state it
|
||||
feeds---across the two hardware axes the machine exposes: the CUBEs and,
|
||||
within each CUBE, the PEs (here $C{=}8$ CUBEs $\times$ $P{=}8$ PEs, for
|
||||
$C\!\cdot\!P{=}64$ attention engines over one KV-head group on the
|
||||
LLaMA-3.1-70B target). Each axis can \emph{replicate} the KV cache or
|
||||
\emph{shard} it, and a shard can run along the sequence dimension
|
||||
$S_{kv}$ or the head dimension $d_{\text{head}}$. The cross product is a
|
||||
small, enumerable taxonomy; six placements span its meaningful corners
|
||||
(Figure~\ref{fig:gqa-kv-sharding}).
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_kv_sharding_diagram.png}
|
||||
\caption{The six KV-placement strategies, drawn on the
|
||||
$S_{kv}\!\times\!d_{\text{head}}$ KV tensor (rows = sequence, columns =
|
||||
head dimension). Cube colour bands and dashed PE dividers show which
|
||||
axis each level shards. Cases~1--3 either replicate the cache or shard
|
||||
it on a single axis (8-way at most); Cases~4--6 reach a full 64-way
|
||||
split, three different ways: Case~4 splits $S_{kv}$ across CUBEs and
|
||||
$d_{\text{head}}$ across PEs, Case~5 the mirror, and Case~6~$\star$
|
||||
splits $S_{kv}$ on \emph{both} axes.}
|
||||
\label{fig:gqa-kv-sharding}
|
||||
\end{figure}
|
||||
|
||||
Two quantities decide which placement is viable, and they pull against
|
||||
each other. The first is \textbf{per-PE KV memory}. With a per-PE HBM
|
||||
budget of \SI{6.0}{\giga\byte} and \SI{1.76}{\giga\byte} of attention
|
||||
weights resident, the KV headroom is \SI{4.24}{\giga\byte} per PE. At a
|
||||
production context of $S_{kv}{=}1\,\text{M}$ tokens the unsharded cache
|
||||
is \SI{40}{\giga\byte}/PE (Case~1), an 8-way shard is \SI{5}{\giga\byte}
|
||||
(Cases~2--3)---both \emph{over} the headroom---while only the 64-way
|
||||
placements bring it to \SI{640}{\mega\byte}/PE (Cases~4--6), comfortably
|
||||
inside budget. Memory alone therefore eliminates Cases~1--3 at long
|
||||
context. The second quantity is \textbf{communication per token}, and it
|
||||
is what separates the three survivors. Sharding $d_{\text{head}}$
|
||||
(Cases~4--5) makes each PE hold only a slice of every head, so the
|
||||
$Q\!\cdot\!K^{\top}$ score is \emph{partial} and must be all-reduced
|
||||
across the slice owners on every token---a reduction whose volume scales
|
||||
with $S_{kv}$ ($\sim$\SI{166}{\mega\byte}/token analytically, intra-CUBE
|
||||
on the NoC for Case~4, inter-CUBE on UCIe for Case~5). Case~6~$\star$
|
||||
instead shards $S_{kv}$ on both axes, so every PE computes a
|
||||
\emph{complete} score over its own token range and only the small running
|
||||
softmax state $(m,\ell,O)$ is merged across PEs
|
||||
($\sim$\SI{6.2}{\mega\byte}/token)---a $\sim$27$\times$ lighter collective
|
||||
than the $d_{\text{head}}$-split designs.
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_long_ctx_6cases_summary.png}
|
||||
\caption{Long-context placement analysis at $S_{kv}{=}1\,\text{M}$ tokens.
|
||||
\emph{Left:} the per-PE HBM budget---\SI{1.76}{\giga\byte} of attention
|
||||
weights leave \SI{4.24}{\giga\byte} of KV headroom (red line).
|
||||
\emph{Middle:} per-PE KV memory per case (log scale); only the 64-way
|
||||
placements (Cases~4--6, \SI{640}{\mega\byte}) clear the headroom, while
|
||||
the unsharded (\SI{40}{\giga\byte}) and 8-way (\SI{5}{\giga\byte}) cases
|
||||
overflow. \emph{Right:} analytical communication per token (log scale);
|
||||
the $d_{\text{head}}$-split Cases~4--5 pay a partial-score all-reduce
|
||||
($\sim$\SI{166}{\mega\byte}/token) that the both-axes-$S_{kv}$ split of
|
||||
Case~6~$\star$ avoids ($\sim$\SI{6.2}{\mega\byte}/token, merging only the
|
||||
softmax state). On these two axes Case~6 (marked $\star$ in the figure)
|
||||
is the single placement that lands both inside the memory budget and at
|
||||
low per-token communication; whether that combination is the right one to
|
||||
pick is a regime-specific question taken up next.}
|
||||
\label{fig:gqa-budget}
|
||||
\end{figure}
|
||||
|
||||
These two costs---per-PE KV memory and per-token communication---are
|
||||
intrinsic properties of each placement, fixed by how it shards the cache
|
||||
and independent of the workload regime. They do not by themselves name a
|
||||
winner: a short prompt where the whole cache fits on one PE values low
|
||||
communication and tolerates replication, whereas a million-token decode
|
||||
is bound by the memory wall and will pay communication to escape it.
|
||||
Which placement is appropriate is therefore a per-regime question, which
|
||||
the short- and long-context subsections that follow answer by running the
|
||||
options on the simulator and reading off latency, traffic, and the
|
||||
redundant compute each one induces.
|
||||
|
||||
\subsection{Inference with Short-Context Length}
|
||||
\label{sec:gqa-short}
|
||||
|
||||
The fused GQA kernel issues its matrix products as scheduler-managed
|
||||
composite commands and keeps the online-softmax merge and the cross-device
|
||||
KV reduction inside the kernel, on PE\_IPCQ. Two kernel families cover the
|
||||
two phases. The \emph{prefill} kernel splits the query tile across the PEs
|
||||
of a group and broadcasts each KV tile from the group's root PE over
|
||||
intra-CUBE IPCQ, so the HBM K/V read is paid once per group instead of
|
||||
once per PE. The \emph{decode} kernel sequence-shards the KV cache across
|
||||
the same group of PEs, runs per-PE local attention, then chain-reduces the
|
||||
partial $(m,\ell,O)$ triples back up to the group root. Two further
|
||||
primitives make long context practical: a \emph{lazy load} that issues the
|
||||
KV \textsf{DMA\_READ} and returns immediately, auto-waiting only at first
|
||||
use so KV load overlaps score computation; and per-tile \emph{scratch
|
||||
recycling} that keeps the running accumulators in a persistent arena while
|
||||
freeing per-tile temporaries, so the kernel fits the
|
||||
\SI{1}{\mebi\byte} scratch budget across many tiles. A further refinement
|
||||
restructures the decode step so the per-tile matrix products and the
|
||||
online-softmax merge are issued as \emph{composite} commands rather than
|
||||
hand-tiled primitives; \S\ref{sec:gqa-composite} measures it against the
|
||||
primitive baseline at long context.
|
||||
|
||||
This is a different decomposition from the six-case placement taxonomy of
|
||||
\S\ref{sec:gqa-placement}: that taxonomy spreads a single KV-head group
|
||||
across all 64 PEs and is the right lens for long context
|
||||
(\S\ref{sec:gqa-long}), whereas short context---where the whole cache fits
|
||||
comfortably---turns instead on how the eight KV heads are distributed
|
||||
across the eight CUBEs. The design question for short context is therefore
|
||||
not whether to fuse softmax---that is settled---but how to distribute the
|
||||
eight KV heads across the eight CUBEs of a SIP. Four mappings cover the spectrum, named
|
||||
by KV-head count per CUBE: \textsf{1-kv-per-cube} dedicates one whole
|
||||
CUBE to each head and uses all eight PEs of that CUBE on the head's
|
||||
sequence shard; \textsf{2-kv-per-cube} packs two heads per CUBE with
|
||||
group size four; \textsf{4-kv-per-cube} packs four heads with group size
|
||||
two; \textsf{8-kv-per-cube} loads all heads onto a single CUBE with one
|
||||
PE per head. The mapping fixes a trade-off: as KV heads per CUBE grows, the per-CUBE KV
|
||||
footprint grows linearly and the mapping uses proportionally fewer CUBEs
|
||||
(eight down to one), but the intra-CUBE IPCQ broadcast/reduce cost shrinks
|
||||
to zero (8-kv-per-cube uses a single PE per head and has no group
|
||||
communication).
|
||||
|
||||
A second axis is how each GEMM tile is issued. We isolate this with three
|
||||
\emph{composite tiers} applied to the same mapping: \textsf{without
|
||||
composite} uses primitive \textsf{tl.dot} and unfused softmax;
|
||||
\textsf{with composite (GEMM-only)} issues the GEMMs as
|
||||
\textsf{tl.composite(op="gemm")} commands so the scheduler can pipeline
|
||||
the tile stages but leaves softmax as primitives; \textsf{with composite
|
||||
+ softmax\_merge} additionally folds the softmax fold into the $P\!\cdot\!V$
|
||||
composite through a named \textsf{softmax\_merge} prologue recipe.
|
||||
|
||||
We sweep all four mappings $\times$ three composite tiers $\times$
|
||||
$S_{kv}\in\{8\text{K},16\text{K},32\text{K},64\text{K}\}$ for both
|
||||
phases (prefill uses $T_q{=}8$ sliced tiles; decode uses $T_q{=}1$). The
|
||||
headline latency comparison is in Figure~\ref{fig:gqa-short-wall-baseline}.
|
||||
The mappings separate cleanly, and in proportion to how many PEs each
|
||||
dedicates to a head: \textsf{1-kv-per-cube} spreads one head across all
|
||||
eight PEs of its CUBE, \textsf{8-kv-per-cube} gives each head a single PE,
|
||||
and the four mappings form a $\{64,32,16,8\}$-active-PE ladder. Decode
|
||||
latency tracks that ladder inversely---at $S_{kv}{=}64$K the
|
||||
\textsf{8-kv-per-cube}/\textsf{1-kv-per-cube} ratio is $7.4\times$
|
||||
($\SI{87.0}{}$ vs.\ $\SI{11.8}{\micro\second}$), and even at $8$K it is
|
||||
$4.8\times$ ($\SI{10.8}{}$ vs.\ $\SI{2.2}{\micro\second}$); prefill shows
|
||||
the same ordering more gently ($2.3\times$ at $64$K). The reason is that
|
||||
decode here is \emph{bandwidth-bound}: each active PE streams its KV shard
|
||||
at \SI{46}{}--\SI{76}{\percent} of the \SI{256}{\giga\byte\per\second}
|
||||
per-PE ceiling (rising with context), so a mapping's speed is set by how
|
||||
many PEs it puts to work---more PEs, more aggregate HBM bandwidth, lower
|
||||
latency. The $M{=}G{=}8$ score GEMM is skinny and leaves the MAC array
|
||||
lightly loaded throughout; the lever here is data movement, not compute.
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_short_context/wall_variant1_mode_compare.png}
|
||||
\caption{Wall-clock latency of the four short-context mappings, baseline
|
||||
kernel (no composite; log scale). The mappings separate along the
|
||||
$\{64,32,16,8\}$-active-PE ladder set by heads-per-CUBE:
|
||||
\textsf{1-kv-per-cube} (64 PEs) is fastest and \textsf{8-kv-per-cube}
|
||||
(8 PEs) slowest, by $4.8\times$ at $8$K growing to $7.4\times$ at $64$K in
|
||||
decode ($2.3\times$ at $64$K in prefill). Decode is bandwidth-bound, so
|
||||
latency scales inversely with the number of active PEs. Composite tiers
|
||||
track the baseline at this skinny shape
|
||||
(Figure~\ref{fig:gqa-short-a1-variant}).}
|
||||
\label{fig:gqa-short-wall-baseline}
|
||||
\end{figure}
|
||||
|
||||
Figure~\ref{fig:gqa-short-tradeoff} shows what the mappings trade for that
|
||||
latency. The differentiator is \emph{density}: \textsf{8-kv-per-cube}
|
||||
concentrates all eight heads onto one CUBE, so its per-CUBE KV footprint is
|
||||
$8\times$ that of \textsf{1-kv-per-cube} (one CUBE per head), and it
|
||||
occupies a single one of the SIP's sixteen CUBEs against
|
||||
\textsf{1-kv-per-cube}'s eight. Per-PE HBM \emph{utilization}, by contrast,
|
||||
is similar across mappings---every active PE runs near the same
|
||||
\SI{46}{}--\SI{76}{\percent} of the per-PE ceiling---so the latency ladder
|
||||
comes from the \emph{count} of active PEs, not from any per-PE
|
||||
concentration. IPCQ traffic runs opposite to density:
|
||||
\textsf{1-kv-per-cube} pays the eight-PE chain-reduce while
|
||||
\textsf{8-kv-per-cube} (one PE per head) pays none, but absolute IPCQ volume
|
||||
stays under \SI{120}{\kibi\byte} per run, two orders of magnitude below the
|
||||
HBM traffic. The result is a genuine trade-off rather than a single winner:
|
||||
\textsf{1-kv-per-cube} minimizes per-request latency by spending the most
|
||||
hardware (eight CUBEs, 64 PEs) on one request, while denser mappings leave
|
||||
CUBEs free for other requests---the batched-serving axis taken up below.
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_short_context/per_cube_tradeoff_mode_compare.png}
|
||||
\caption{Per-CUBE trade-off across the four mappings (baseline kernel;
|
||||
log scale). Top row: per-CUBE KV footprint---\textsf{8-kv-per-cube} packs
|
||||
all eight heads onto one CUBE, an $8\times$ larger per-CUBE KV than
|
||||
\textsf{1-kv-per-cube} and one CUBE used against eight (the $8\times$ ratio
|
||||
is context-invariant; the absolute footprint grows with $S_{kv}$). Bottom
|
||||
row: IPCQ traffic per run---\textsf{1-kv-per-cube} pays the eight-PE
|
||||
chain-reduce while \textsf{8-kv-per-cube} (single PE per head) pays zero.
|
||||
The density (top) is what sets how many concurrent requests a SIP can hold.}
|
||||
\label{fig:gqa-short-tradeoff}
|
||||
\end{figure}
|
||||
|
||||
The composite ablation in Figure~\ref{fig:gqa-short-a1-variant} shows the
|
||||
three tiers within \textsf{1-kv-per-cube}: wall-clock is nearly flat, with
|
||||
the GEMM-only composite tier a few percent \emph{slower}
|
||||
(\SI{14.0}{} vs.\ \SI{11.8}{\micro\second} at $64$K decode) and the
|
||||
\textsf{softmax\_merge} tier matching the baseline. This is the expected
|
||||
outcome for the decode-skinny shape---$M{=}G{=}8$ is well below the
|
||||
scheduler's $\textsf{TILE\_M}{=}32$ supertile (\S\ref{sec:gemm}), so the
|
||||
composite path pads $M$ by $4\times$ with zeros and the fusion has no slack
|
||||
to win back; the padding is pure overhead here. It shows up in
|
||||
Figure~\ref{fig:gqa-short-gemm-util} as inflated GEMM-engine busy time on
|
||||
the composite tiers---the padded-$M$ GEMM is charged in full against a
|
||||
sub-\SI{15}{\micro\second} decode wall, driving the accounted utilization
|
||||
well above the baseline's \SI{12}{}--\SI{18}{\percent}. Fusion benefit
|
||||
requires lifting $M$, either by the long-context Q-tile split
|
||||
(\S\ref{sec:gqa-composite}) or by batching multiple users---which is also
|
||||
the lever that converts a mapping's density into throughput.
|
||||
|
||||
\paragraph{Latency versus batched throughput.} The latency ladder and the
|
||||
density trade-off pull in opposite directions once a SIP serves more than
|
||||
one request. Decode users are independent---each attends its own KV
|
||||
cache---and these mappings generate no cross-CUBE traffic, so a SIP runs
|
||||
several users at once on disjoint CUBE groups, up to $16/C$ users: two for
|
||||
\textsf{1-kv-per-cube}, sixteen for \textsf{8-kv-per-cube}. We measure this
|
||||
directly, launching $B$ concurrent users and reading aggregate latency off
|
||||
the shared clock (Figure~\ref{fig:gqa-batch}). The overlap is nearly
|
||||
perfect: aggregate latency stays within \SI{3}{}--\SI{5}{\percent} of the
|
||||
single-user latency all the way to a full SIP (\textsf{8-kv-per-cube} at
|
||||
$16$ users is \SI{11.4}{} vs.\ \SI{10.8}{\micro\second}), so that small
|
||||
residual is the only cross-user interference the shared fabric adds.
|
||||
Aggregate throughput at $8$K therefore rises almost linearly with $B$, from
|
||||
\SI{0.86}{}~requests\,/\,\si{\micro\second} (\textsf{1-kv-per-cube}, capped
|
||||
at two users) to \SI{1.41}{} (\textsf{8-kv-per-cube}, sixteen users): the
|
||||
\emph{dense} mapping wins on throughput even though it is the \emph{slowest}
|
||||
per request.
|
||||
|
||||
The reason is worth spelling out, because both mappings fill the same
|
||||
hardware: at capacity \textsf{1-kv-per-cube} runs two users
|
||||
$\times$ 64 PEs and \textsf{8-kv-per-cube} sixteen users $\times$ 8 PEs, so
|
||||
each puts all 128 PEs of the SIP to work. What differs is per-PE
|
||||
\emph{efficiency}. Splitting one head across eight PEs
|
||||
(\textsf{1-kv-per-cube}) leaves each PE only $S_{kv}/8$ tokens---a single
|
||||
tile---so its fixed overheads (pipeline fill on the lazy KV load, and the
|
||||
eight-PE online-softmax chain-reduce) dominate the little streaming work it
|
||||
does, and it sustains just \SI{46}{\percent} of the per-PE HBM ceiling.
|
||||
Giving each head a whole PE (\textsf{8-kv-per-cube}) streams the full
|
||||
$S_{kv}$-token cache contiguously with no reduce, amortizing those overheads
|
||||
over $8\times$ more work and reaching \SI{76}{\percent}. The measured
|
||||
throughput ratio ($1.41/0.86 = 1.6\times$) matches the utilization ratio
|
||||
($0.76/0.46 = 1.7\times$) almost exactly: because decode is
|
||||
bandwidth-bound, aggregate throughput is set by total HBM efficiency, not
|
||||
by PE count. \textsf{1-kv-per-cube} is really spending hardware on
|
||||
\emph{latency}---eight PEs per head buy only a $4.8\times$ speedup, a $60\%$
|
||||
strong-scaling efficiency---and that same $40\%$ overhead is what caps its
|
||||
throughput once the SIP is full. The design conclusion is regime-dependent: \textsf{1-kv-per-cube} is the
|
||||
latency-optimal choice for single-stream or low-batch decode, while denser
|
||||
mappings are the throughput-optimal choice for high-batch short-context
|
||||
serving. At long context the per-request latencies already scale as $1/C$,
|
||||
so the same accounting predicts the throughputs converge; a measured
|
||||
long-context batch sweep is 2H work (\S\ref{sec:future}).
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_short_context/batch_scaling.png}
|
||||
\caption{Batch scaling at $S_{kv}{=}8$K: aggregate throughput (requests per
|
||||
\si{\micro\second}) versus the number of concurrent decode users on one SIP,
|
||||
each user on a disjoint CUBE group. Each mapping scales almost linearly with
|
||||
$B$ up to its SIP capacity $16/C$---\textsf{1-kv-per-cube} saturates at two
|
||||
users, \textsf{8-kv-per-cube} at sixteen. Solid segments are measured
|
||||
concurrent runs; the dashed extensions are the saturated-throughput ceiling
|
||||
beyond capacity, where further users run in waves at that sustained rate
|
||||
(so the SIP is full, not idle). The dense mapping reaches the highest
|
||||
ceiling; the spread mapping is latency-optimal but capacity-limited.
|
||||
Concurrent users overlap to within \SI{3}{}--\SI{5}{\percent} of the
|
||||
single-user latency.}
|
||||
\label{fig:gqa-batch}
|
||||
\end{figure}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_short_context/wall_a1_variant_compare.png}
|
||||
\caption{Composite-tier ablation on \textsf{1-kv-per-cube} (log scale):
|
||||
wall-clock is nearly flat across tiers---the GEMM-only composite runs a few
|
||||
percent slower (padding overhead) and \textsf{softmax\_merge} matches the
|
||||
baseline. With $M{=}G{=}8 < \textsf{TILE\_M}{=}32$ the composite scheduler
|
||||
pads $M$ by $4\times$, leaving no fusion slack to recover at this shape.
|
||||
Fusion pays off only when $M$ fills the supertile---long-context Q-tile
|
||||
splitting or batched-$M$ inference.}
|
||||
\label{fig:gqa-short-a1-variant}
|
||||
\end{figure}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_short_context/gemm_util_a1_variant_ablation.png}
|
||||
\caption{GEMM-engine busy fraction on \textsf{1-kv-per-cube} across the
|
||||
three composite tiers. The baseline runs at \SI{12}{}--\SI{18}{\percent};
|
||||
the composite tiers read much higher---up to and past
|
||||
\SI{100}{\percent} of the short decode wall---because the padded
|
||||
$8\!\to\!32$ $M$ dimension is charged as GEMM time, bookkeeping overhead
|
||||
rather than useful work. The engine is never the bottleneck at this shape;
|
||||
decode is KV-bandwidth-bound.}
|
||||
\label{fig:gqa-short-gemm-util}
|
||||
\end{figure}
|
||||
|
||||
\subsection{Inference with Long-Context Length}
|
||||
\label{sec:gqa-long}
|
||||
|
||||
% TODO: prefill long-context kernel implementation description
|
||||
% (Sequence-Parallel partition of S_kv, per-case mechanics).
|
||||
% Bench code: src/kernbench/benches/gqa_helpers/long_ctx/
|
||||
|
||||
% TODO: prefill long-context performance figure.
|
||||
|
||||
Long-context decode is the regime where the KV cache, not attention
|
||||
compute, sets serving cost, so the placement question of
|
||||
\S\ref{sec:gqa-placement} becomes decisive here. To pick the right
|
||||
placement for this regime we run each of the six options as a fused
|
||||
decode kernel on the simulator (one decode step on the LLaMA-3.1-70B
|
||||
single-KV-head-group target, $C{=}8$ CUBEs $\times$ $P{=}8$ PEs) at a
|
||||
tractable $S_{kv}{=}8192$, and read off end-to-end latency, on-device op
|
||||
traffic, and the redundant compute each one induces. The swept context is
|
||||
small enough that all six fit in memory at $S_{kv}{=}8192$; the
|
||||
placements the long-context memory budget rules out (Cases~1--3) are
|
||||
drawn in red, run here only to expose their issue and communication
|
||||
structure.
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_latency.png}
|
||||
\caption{Measured end-to-end decode latency per placement
|
||||
($S_{kv}{=}8192$; red = ruled out by the long-context memory budget,
|
||||
blue = the predicted Pareto choice). The fastest raw latency belongs to
|
||||
Case~3 (\SI{17.8}{\micro\second})---but Case~3 replicates the full KV
|
||||
cache into every CUBE, so it overflows the per-PE budget at production
|
||||
context and wastes 8$\times$ the compute
|
||||
(Figure~\ref{fig:gqa-6cases-par}). Among the placements that actually fit
|
||||
1\,M-token memory (Cases~4--6), Case~6~$\star$ is the fastest
|
||||
(\SI{30.6}{\micro\second}, versus \SI{31.4}{} and \SI{34.5}{\micro\second}
|
||||
for the $d_{\text{head}}$-split Cases~4 and~5)---making it the placement
|
||||
of choice for long-context decode.}
|
||||
\label{fig:gqa-6cases-lat}
|
||||
\end{figure}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_parallelism.png}
|
||||
\caption{Redundant compute per placement, measured as active-PE
|
||||
$\times$ $S_{\text{local}}$ (PE-tokens; lower means less wasted work).
|
||||
The minimum is \num{8192} PE-tokens---one pass over the sequence.
|
||||
Case~3 inflates this 8$\times$ to \num{65536} by replicating the KV
|
||||
cache across all eight CUBEs so every CUBE redundantly re-attends the
|
||||
whole sequence; the $d_{\text{head}}$-split Cases~4--5 likewise carry
|
||||
\num{65536} because each token is processed across eight head slices.
|
||||
Case~6~$\star$ achieves the full 64-way split at the minimal
|
||||
\num{8192} PE-tokens---fully parallel, no replication.}
|
||||
\label{fig:gqa-6cases-par}
|
||||
\end{figure}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_6cases_traffic.png}
|
||||
\caption{Measured on-device op traffic per placement. The unsharded
|
||||
Case~1 issues no IPCQ copies (each PE has the full cache, nothing to
|
||||
reduce); the single-axis Case~3 charges 168. Among the 64-way splits,
|
||||
the $d_{\text{head}}$-split Cases~4--5 charge the most---280 IPCQ copies
|
||||
and 8 DMA writes each, the partial-score all-reduce that head-slicing
|
||||
forces---while Case~6~$\star$ needs only 189 IPCQ copies and a single DMA
|
||||
write, because it merges just the running softmax state $(m,\ell,O)$
|
||||
rather than partial scores. This is the on-device collective traffic that
|
||||
PE\_IPCQ and the torus links of \S\ref{sec:allreduce} are provisioned to
|
||||
absorb at link speed.}
|
||||
\label{fig:gqa-6cases-traffic}
|
||||
\end{figure}
|
||||
|
||||
Taken together the three panels select the placement. The only
|
||||
memory-feasible family at production context is the 64-way splits
|
||||
(Cases~4--6), and within it Case~6~$\star$ is both the fastest and the
|
||||
lightest-communicating, because it merges only the running softmax state
|
||||
rather than the partial scores that the $d_{\text{head}}$-split
|
||||
Cases~4--5 must all-reduce. Case~3's lower raw latency is beside the
|
||||
point---it replicates the full cache into every CUBE, overflowing the
|
||||
per-PE budget and wasting $8\times$ the compute. For long-context decode
|
||||
the placement of choice is therefore the both-axes sequence shard, and the
|
||||
cross-PE softmax reduction it does pay is precisely the traffic the
|
||||
communication-side codesign of this report is built to move quickly.
|
||||
|
||||
\subsection{Use of Composite Commands}
|
||||
\label{sec:gqa-composite}
|
||||
|
||||
The decode kernel of \S\ref{sec:gqa-long} issues its local attention as
|
||||
primitive operations: it walks each PE's $S_{\text{local}}$ token slice in
|
||||
\SI{1024}{}-token tiles, and for every tile issues a $Q\!\cdot\!K^{\top}$
|
||||
\textsf{dot}, the online-softmax primitives, and a $P\!\cdot\!V$
|
||||
\textsf{dot}, merging the running $(m,\ell,O)$ state by hand. The number
|
||||
of PE\_CPU commands this costs grows with the context. At a production
|
||||
context of $S_{kv}{=}1\,\text{M}$ tokens the Case-6 64-way split gives
|
||||
each PE $S_{\text{local}}{=}16384$ tokens, so its local attention is a
|
||||
$Q\!\cdot\!K^{\top}$ of $(8,128)\!\cdot\!(128,16384)$ and a
|
||||
$P\!\cdot\!V$ of $(8,16384)\!\cdot\!(16384,128)$---sixteen hand-issued
|
||||
tiles, each a fresh batch of CPU commands.
|
||||
|
||||
The composite command lets the kernel hand that tiling to PE\_SCHEDULER.
|
||||
We compare three command forms of the \emph{same} Case-6 kernel---identical
|
||||
placement and identical $(m,\ell,O)$ reduce, differing only in how the
|
||||
local attention is issued:
|
||||
\begin{itemize}
|
||||
\item \textbf{primitive}---the hand-tiled \textsf{dot}/softmax kernel
|
||||
above (the baseline of \S\ref{sec:gqa-long}).
|
||||
\item \textbf{composite}---each matrix product is one coarse
|
||||
\textsf{composite} GEMM over the \emph{whole} $S_{\text{local}}$, with
|
||||
$K$ and $V$ passed as HBM references so PE\_SCHEDULER streams and
|
||||
tiles them on the fixed $32\!\times\!64\!\times\!32$ MAC tile; the
|
||||
softmax stays primitive.
|
||||
\item \textbf{composite\,+\,softmax\_merge}---additionally folds the
|
||||
online-softmax merge and $P\!\cdot\!V$ into a single stateful
|
||||
\textsf{softmax\_merge} recipe composite.
|
||||
\end{itemize}
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_decode_long_ctx_composite.png}
|
||||
\caption{Three command forms of the Case-6 decode kernel, swept over
|
||||
context length ($S_{\text{local}}{=}S_{kv}/64$ per PE). \emph{Right:}
|
||||
PE\_CPU commands issued. The hand-tiled primitive kernel rises
|
||||
$O(n_{\text{tiles}})$---from 96 commands at one tile to 426 at the
|
||||
1\,M-token, sixteen-tile production point---while both composite forms
|
||||
issue a context-\emph{independent} $O(1)$ count (94 and 98) that
|
||||
\emph{saturates}: one coarse descriptor offloads the entire per-tile
|
||||
fan-out. \emph{Left:} the consequence for wall-clock latency is none---all
|
||||
three land on the same curve (\SI{30.6}{}, \SI{231}{},
|
||||
\SI{461}{\micro\second} at 8\,K\,/\,64\,K\,/\,128\,K), because decode is
|
||||
bound by streaming the KV cache, not by issue. Command-count is measured
|
||||
at emit time (exact, to 1\,M); latency on the data-mode engine over the
|
||||
tractable range.}
|
||||
\label{fig:gqa-composite}
|
||||
\end{figure}
|
||||
|
||||
Figure~\ref{fig:gqa-composite} reads off the two quantities that matter,
|
||||
and they point in opposite directions. The PE\_CPU command count (right)
|
||||
collapses from a context-growing $O(n_{\text{tiles}})$ to a flat $O(1)$:
|
||||
at 1\,M tokens the composite form issues \num{94} commands against the
|
||||
primitive kernel's \num{426}, a $4.5\times$ reduction
|
||||
($4.2\times$ in modeled dispatch cost), and---crucially---that number no
|
||||
longer grows with context. The wall-clock latency (left), by contrast,
|
||||
is unchanged across all three forms: decode is bound by streaming the KV
|
||||
cache out of HBM, so the command form does not move the critical path.
|
||||
|
||||
That juxtaposition is the point. The composite command is not a latency
|
||||
optimization for this memory-bound decode; it is a \emph{CPU-issue}
|
||||
optimization. Its value is removing the per-tile dispatch work that would
|
||||
otherwise grow without bound as context grows, freeing PE\_CPU to run
|
||||
ahead and keep the engines fed---which is exactly what lets the
|
||||
data-movement cost analyzed next show through as the true bottleneck
|
||||
rather than being masked by issue overhead. The \textsf{softmax\_merge}
|
||||
recipe folds the online merge into the same descriptor; on this
|
||||
memory-bound path its marginal cost over the plain GEMM composite is small
|
||||
(98 vs.\ 94 commands), and like the plain composite it keeps the issued
|
||||
count flat as context scales.
|
||||
|
||||
\paragraph{The compute-bound mirror: prefill.} Decode's verdict---command
|
||||
form is latency-neutral---is a property of its regime, not of the
|
||||
composite command. A decode step has $T_q{=}1$, so its score and context
|
||||
products are skinny ($M{=}G\,T_q{=}8$): the MAC array is barely fed and
|
||||
the kernel is bound by streaming the KV cache. Prefill is the opposite
|
||||
corner. It processes a block of query positions at once, so $M{=}G\,T_q$
|
||||
is large and tile-filling, the GEMMs carry real arithmetic intensity
|
||||
($\sim$$M$ flops/byte, well above the roofline ridge), and the kernel is
|
||||
\emph{compute-bound}. This is the regime the composite command was built
|
||||
for (\S\ref{sec:gemm}): it streams the per-HW-tile
|
||||
DMA$\rightleftarrows$compute pipeline so the MAC array stays fed, whereas
|
||||
the primitive kernel's blocking \textsf{tl.dot} serializes each tile's
|
||||
load and compute and starves the array between tiles. We run the same
|
||||
three command forms on a single-rank compute-bound prefill (FlashAttention
|
||||
$Q$-block $\times$ $S_{kv}$-tile, online softmax) and sweep the context
|
||||
length (Figure~\ref{fig:gqa-prefill-cb}).
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_prefill_compute_bound.png}
|
||||
\caption{Compute-bound prefill, three command forms, swept over context
|
||||
length ($M{=}8\,T_q$ tile-filling). \emph{Left:} end-to-end latency.
|
||||
\emph{Right:} MAC utilization (achieved $\div$ the
|
||||
\SI{8}{\tera\flop\per\second} per-PE peak). The hand-tiled primitive sits
|
||||
flat at $\sim$\SI{68}{\percent}---its serial load$\to$dot path leaves the
|
||||
MAC array idle between tiles regardless of context. The composite forms
|
||||
climb with context (\SI{67}{}$\to$\SI{80}{\percent} plain,
|
||||
\SI{67}{}$\to$\SI{83}{\percent} with the recipe) because a deeper $P\!\cdot
|
||||
\!V$ reduction gives more HW tiles to pipeline, and they convert that into
|
||||
wall-clock: at \num{1024} the recipe form is \SI{646.9}{} vs.\
|
||||
\SI{794.1}{\micro\second} (\SI{19}{\percent} faster). The margin
|
||||
\emph{grows} with context---the compute-bound mirror of the GEMM result of
|
||||
\S\ref{sec:gemm}.}
|
||||
\label{fig:gqa-prefill-cb}
|
||||
\end{figure}
|
||||
|
||||
The two studies together state the composite command's value precisely. It
|
||||
has two distinct benefits, and which one matters is set by the workload's
|
||||
roofline position. The first is \emph{host-issue offload}: one macro
|
||||
command in place of $O(n_{\text{tiles}})$ fine ones, which removes
|
||||
PE\_CPU dispatch work and is regime-independent (it shows in the decode
|
||||
command count). The second is \emph{MAC-array feeding}: the
|
||||
scheduler-internal per-tile DMA$\rightleftarrows$compute pipeline, which
|
||||
only converts to latency when the workload is compute-bound enough to have
|
||||
a MAC array worth keeping busy (it shows in the prefill utilization). A
|
||||
memory-bound decode exercises only the first; a compute-bound prefill
|
||||
exercises both. The composite command is the single mechanism that
|
||||
delivers each where it applies.
|
||||
|
||||
% Summary moved to sections/05z-summary.tex so the three new
|
||||
% subsections (roofline, capacity planning, parallelism selection)
|
||||
% can appear between the technical results and the closing summary.
|
||||
|
||||
% TODO: cross-regime DP (data parallelism) applicability:
|
||||
% - Does Case-4 long-context placement compose with batch-level DP
|
||||
% without further changes?
|
||||
% - Does the short-context placement compose the same way?
|
||||
% - Implications for multi-user serving (single vs. mixed regimes).
|
||||
@@ -0,0 +1,22 @@
|
||||
\subsection{Summary}
|
||||
\label{sec:gqa-analysis}
|
||||
|
||||
The section's results line up into one picture of the fused kernel.
|
||||
Placement is decided by memory first and communication second: only the
|
||||
64-way splits fit production context, and among them the both-axes
|
||||
sequence shard (Case~6) minimizes the per-token collective by merging
|
||||
softmax state rather than partial scores. Command form is decided by
|
||||
roofline position: the composite command removes the per-tile issue work
|
||||
in both regimes, but converts to wall-clock only in compute-bound prefill,
|
||||
while memory-bound decode stays bound by KV streaming regardless of command
|
||||
form. The thread connecting the two is that, once composite issue makes
|
||||
GEMM cheap and leaves the MAC array idle, the fused kernel's latency is set
|
||||
almost entirely by data movement---streaming the KV cache and reducing
|
||||
partials across PEs---which is exactly the cost the communication-side
|
||||
work (on-device reduction, lazy load/compute overlap, fast TCM staging and
|
||||
torus links) is built to attack. Roofline framing
|
||||
(\S\ref{sec:roofline}) turns this observation into a sizing and sharding
|
||||
programme in the next two subsections: capacity planning
|
||||
(\S\ref{sec:capacity-planning}) and parallelism selection
|
||||
(\S\ref{sec:parallelism-selection}). What this implies for hardware
|
||||
investment across the report as a whole is taken up in the discussion.
|
||||
@@ -0,0 +1,309 @@
|
||||
\section{Supporting Agentic Workloads}
|
||||
\label{sec:agentic}
|
||||
|
||||
The fused GQA kernel of \S\ref{sec:gqa} solved the efficient execution of a
|
||||
\emph{single} logical attention stream: one query stream against one KV
|
||||
cache, tiled and merged with an online softmax. Agentic inference
|
||||
introduces a different execution model, in which one request dynamically
|
||||
\emph{forks} into several cooperating reasoning branches and later
|
||||
\emph{joins} them. The purpose of this section is to show that this
|
||||
multi-stream model needs no new attention algorithm---the expensive
|
||||
hardware machinery built for \S\ref{sec:gqa} is reused unchanged---and to
|
||||
work out the resulting design across three implementation layers. The
|
||||
central claim, stated once here and defended throughout, is:
|
||||
|
||||
\begin{quote}
|
||||
\emph{The proposed agentic execution does not change the semantics of
|
||||
transformer attention; it changes only how independent query rows are
|
||||
scheduled over a shared KV cache.}
|
||||
\end{quote}
|
||||
|
||||
\noindent Attention stays identical; only the scheduling differs. What
|
||||
follows is a \emph{design} built on the measured \S\ref{sec:gqa} kernel; it
|
||||
is not yet a KernBench measurement, and points that go beyond the
|
||||
implemented path are flagged as such.
|
||||
|
||||
\subsection{Motivation: from one attention stream to many}
|
||||
\label{sec:agentic-why}
|
||||
|
||||
Agentic execution has two recurring shapes: a \emph{loop} (an agent that
|
||||
repeatedly generates, calls a tool, and continues) and a \emph{fan-out /
|
||||
fan-in} (a parent agent forks several specialised sub-agents and later
|
||||
synthesises their results). The loop is ordinary autoregressive decoding
|
||||
and needs nothing new. The fan-out is the interesting case, because the
|
||||
branches are \emph{structurally related} rather than independent:
|
||||
\[
|
||||
\begin{aligned}
|
||||
\text{context}_i \;=\;& \underbrace{\text{shared prefix}}_{\text{common}} \\
|
||||
&+\; \underbrace{\text{private role}_i + \text{private suffix}_i}_{\text{branch-specific}} .
|
||||
\end{aligned}
|
||||
\]
|
||||
Ordinary batching groups unrelated requests that merely arrive together;
|
||||
agentic fan-out groups requests that share an identical prefix KV cache and
|
||||
differ only in a short private suffix. That structure creates two levers
|
||||
that unrelated-request batching cannot pull:
|
||||
|
||||
\begin{enumerate}
|
||||
\item \textbf{Shared-prefix KV reuse.} All branches attend to the same
|
||||
prefix KV, so it is read (and stored) once, not once per branch.
|
||||
\item \textbf{Query-row batching.} The new query rows of many sub-agents
|
||||
read the \emph{same} shared KV, so they can be concatenated into one
|
||||
taller GEMM.
|
||||
\end{enumerate}
|
||||
|
||||
\noindent Both levers land directly on the \S\ref{sec:gqa} design, which
|
||||
already multicasts a (small) query against a stationary KV cache and merges
|
||||
the result hierarchically. Fan-out simply makes the query taller, and
|
||||
fan-in adds a join step; neither touches the attention math.
|
||||
|
||||
\subsection{Architecture: three execution layers}
|
||||
\label{sec:agentic-arch}
|
||||
|
||||
Before the mechanics, we fix the layering, because the recurring reader
|
||||
question is ``whose job is this?''. Agentic execution divides cleanly along
|
||||
the request-flow layers already used in this report: the \textbf{agentic
|
||||
framework} decides \emph{what} to run and how to combine it, the
|
||||
\textbf{runtime} decides \emph{how} to map it onto the hardware without
|
||||
copying shared state, and the \textbf{kernel} does the actual math and
|
||||
reduction on the PEs.
|
||||
|
||||
\[
|
||||
\text{Framework} \;\longrightarrow\; \text{Runtime}
|
||||
\;\longrightarrow\; \text{Kernel}
|
||||
\]
|
||||
|
||||
\noindent Keeping the split this way preserves the topology-agnostic
|
||||
runtime boundary: the framework never sees the SIP/CUBE/PE hierarchy, and
|
||||
the kernel never sees the agent tree. The remainder of the section walks
|
||||
the fan-out $\rightarrow$ runtime $\rightarrow$ kernel $\rightarrow$ fan-in
|
||||
path through exactly these three layers.
|
||||
Table~\ref{tab:agentic-levels} states each layer's responsibilities; the
|
||||
kernel row is unchanged from \S\ref{sec:gqa}.
|
||||
|
||||
\begin{table*}[t]
|
||||
\centering
|
||||
\caption{Division of responsibilities for agentic attention. Only the
|
||||
framework and runtime rows are new work; the kernel is the
|
||||
\S\ref{sec:gqa} fused GQA kernel, unchanged in its math and reduction.}
|
||||
\label{tab:agentic-levels}
|
||||
\small
|
||||
\begin{tabular}{@{}p{0.16\textwidth}p{0.40\textwidth}p{0.36\textwidth}@{}}
|
||||
\toprule
|
||||
\textbf{Level} & \textbf{Responsibilities} & \textbf{Explicitly not its job} \\
|
||||
\midrule
|
||||
\textbf{Agentic framework}
|
||||
& Manage the agent tree: fork sub-agents and assign roles; identify the
|
||||
shared prefix; mark which branches are co-schedulable (same shared prefix).
|
||||
Fan-in: enforce schema-constrained results, deduplicate/rank findings,
|
||||
hold evidence behind pointers, insert hierarchical reducers, and drive the
|
||||
main agent's final synthesis.
|
||||
& No topology, routing, KV placement, or scheduling. Sees agents and text,
|
||||
not CUBEs or PEs. \\
|
||||
\addlinespace
|
||||
\textbf{Runtime}
|
||||
& Fork the logical context by page table (shared-prefix pages $+$ private
|
||||
suffix pages) with \emph{no copy} of shared KV. Concatenate co-scheduled
|
||||
sub-agent query rows into $Q_{\text{cmb}}$. Select the execution policy
|
||||
(Stationary-KV vs.\ Distributed-Q, \S\ref{sec:agentic-choice}) from the
|
||||
cost model. Launch the
|
||||
fused GQA kernel (composite command $+$ PE\_IPCQ) and hand the batched work
|
||||
and policy to the simulation engine.
|
||||
& No attention math and no per-hop routing (delegated to the engine and
|
||||
policy); no agent-level semantics. \\
|
||||
\addlinespace
|
||||
\textbf{Kernel}
|
||||
& Multicast $Q_{\text{cmb}}$ to all PEs; per-PE local attention over the
|
||||
stationary KV shard via composite-command $Q\!\cdot\!K^{\top}$ and
|
||||
$P\!\cdot\!V$; produce per-row online-softmax state; merge that state
|
||||
hierarchically over PE\_IPCQ by matching row index (row-parallel, not
|
||||
serialised).
|
||||
& No knowledge of agents, page tables, or policy selection. Identical to
|
||||
\S\ref{sec:gqa}; a taller $Q$ is the only difference. \\
|
||||
\bottomrule
|
||||
\end{tabular}
|
||||
\end{table*}
|
||||
|
||||
\subsection{Stationary-KV Execution Policy}
|
||||
\label{sec:agentic-policyA}
|
||||
|
||||
The core attention operation is $Q\!\cdot\!K^{\top}$, and under fan-out
|
||||
multiple sub-agents contribute different query rows while reading a common
|
||||
$K$. Rather than issuing $Q_A\!\cdot\!K_{\text{shared}}$,
|
||||
$Q_B\!\cdot\!K_{\text{shared}}$, \dots\ as separate small GEMMs, the runtime
|
||||
concatenates the rows,
|
||||
\[
|
||||
Q_{\text{cmb}} = \begin{bmatrix} Q_A \\ Q_B \\ \vdots \end{bmatrix},
|
||||
\qquad
|
||||
Q_{\text{cmb}}\!\cdot\!K_{\text{shared}},
|
||||
\]
|
||||
and issues one taller GEMM. The arithmetic is unchanged---each row is still
|
||||
independent---but the $M$ dimension grows. For a representative fan-out of
|
||||
$8$ sub-agents at $20$ new tokens each, $M = 8\times20 = 160$; with a per-PE
|
||||
KV shard of $256$ sequence positions this is a $160\times d$ by $d\times256$
|
||||
local product---an $M{=}160$, $N{=}256$ shape that sits well above the
|
||||
per-command issue overhead the composite command (\S\ref{sec:gemm}) is built
|
||||
to amortise. Fan-out is therefore especially valuable during \emph{prefill}
|
||||
of the private suffixes, where the combined $M$ is large; during decode the
|
||||
query stays short and the workload remains, as \S\ref{sec:gqa} found,
|
||||
KV-bandwidth bound.
|
||||
|
||||
The Stationary-KV policy maps onto the \S\ref{sec:gqa} placement
|
||||
essentially unchanged: logically replicated $Q$ over a stationary,
|
||||
sequence-sharded KV cache. One
|
||||
KV head spans four CUBEs and 32 PEs, with each PE owning a different KV
|
||||
\emph{sequence} shard $K_p[S_p,d]$, $V_p[S_p,d_v]$. The combined $Q$ is
|
||||
multicast to all of them; each PE forms a complete local score
|
||||
$Q_{\text{cmb}}\!\cdot\!K_p^{\top}$ over its own shard, and the per-row
|
||||
online-softmax state is reduced hierarchically---8-PE merge inside each
|
||||
CUBE, then a 4-CUBE merge---using the same PE\_IPCQ merge primitive from
|
||||
\S\ref{sec:allreduce}.
|
||||
|
||||
\paragraph{The reduction algorithm does not change.} This is the point to
|
||||
emphasise. Agentic batching does \emph{not} require a new reduction
|
||||
algorithm: the hierarchical online-softmax reduction of \S\ref{sec:gqa}
|
||||
remains exactly as-is, and only the query batch becomes larger. The
|
||||
reduction is always \emph{same row index across sequence shards}, never
|
||||
\emph{different query rows against each other}, so $M$ batched rows do
|
||||
\emph{not} become $M$ serialised communication rounds; the $(m,\ell,O)$ row
|
||||
states are exchanged as vectors/tiles and merged in parallel. Because the
|
||||
merge is row-indexed, a taller $Q$ widens each payload but adds no rounds.
|
||||
This is what makes agentic support a reuse of \S\ref{sec:gqa} rather than a
|
||||
redesign.
|
||||
|
||||
\paragraph{Logical vs.\ physical $Q$ replication.} ``Replicated $Q$'' is a
|
||||
\emph{logical} statement. Because the shard axis is the KV \emph{sequence}
|
||||
dimension $S$, every PE must form the full score $Q\!\cdot\!K_p^{\top}$
|
||||
against its local $K_p$ and therefore needs $Q$'s \emph{entire} hidden
|
||||
dimension $d$; what is partitioned across PEs is $K$/$V$ along $S$, never
|
||||
$Q$ along its columns. Splitting $Q$ (and $K$) on the hidden dimension
|
||||
would instead make each PE's product \emph{partial} and force a pre-softmax
|
||||
hidden-dimension reduction ($QK^{\top}=\sum_i Q_iK_i^{\top}$)---that is
|
||||
tensor-/head-parallel attention, a different structure from the
|
||||
sequence-parallel one assumed here, and one that cannot coexist with using
|
||||
the PE axis for sequence shards. Logical replication also does not mean 32
|
||||
physical copies: $Q$ can be multicast once into a CUBE-local shared buffer
|
||||
(shared SRAM) that all PEs in the CUBE read, and a large $Q$ can further be
|
||||
\emph{row}-tiled in time ($Q[0{:}16,:],\,Q[16{:}32,:],\dots$)---row tiling
|
||||
splits the $M$ dimension, not the hidden-dimension columns. In short:
|
||||
the Stationary-KV policy uses logically replicated $Q$ across
|
||||
sequence-parallel PEs while $K$ and $V$ are partitioned along the sequence
|
||||
dimension; $Q$ may be
|
||||
temporally row-tiled or physically shared through multicast buffers, but it
|
||||
is not partitioned along the hidden-dimension columns.
|
||||
|
||||
\paragraph{Replication is not $32\times$ the compute work (attention
|
||||
FLOPs).} Multicasting $Q$ to 32 PEs does not multiply attention FLOPs,
|
||||
because each PE computes against a different KV sequence shard rather than
|
||||
the same one. Let the KV sequence length be $S$; with sequence parallelism
|
||||
over 32 PEs, each PE owns $S/32$ positions. The score GEMM
|
||||
$Q[M,d]\!\cdot\!K^{\top}[d,S]$ costs $\propto M\,d\,S$, so each PE performs
|
||||
$M\,d\,(S/32)$ and the 32 shards sum to
|
||||
\[
|
||||
32 \cdot M\,d\,\tfrac{S}{32} \;=\; M\,d\,S,
|
||||
\]
|
||||
identical to attention over one undivided sequence. Replication therefore
|
||||
changes $Q$ distribution, reduction traffic, buffering, and scheduling---not
|
||||
the total attention FLOPs.
|
||||
|
||||
\subsection{Distributed-Q Execution Policy (alternative)}
|
||||
\label{sec:agentic-policyB}
|
||||
|
||||
The natural alternative is to partition (distribute) the query rows across
|
||||
PE groups (\emph{the Distributed-Q policy}) rather than replicate them. It
|
||||
is not automatically
|
||||
better. Because the 32 PEs already shard the KV \emph{sequence}, every query
|
||||
row must still attend to \emph{all} shards; partitioning $Q$ across PE
|
||||
groups therefore forces each group to reach every KV shard, which requires
|
||||
one of: regrouping KV shards per $Q$ group, replicating KV across groups, or
|
||||
reading remote KV through symmetric memory. Each of these adds
|
||||
memory-system complexity that the Stationary-KV policy avoids entirely.
|
||||
Time-multiplexing the same PEs over $Q$ groups is the fourth option, but
|
||||
that is temporal tiling---already available under the Stationary-KV policy
|
||||
as row tiling---not true spatial $Q$ partitioning. The Distributed-Q policy
|
||||
is thus a proposed adaptive extension, not the
|
||||
baseline, and is only worth its complexity when $Q$ grows large enough that
|
||||
multicast and reduction traffic dominate remote/regrouped-KV cost.
|
||||
|
||||
\subsection{Why the Stationary-KV policy is the baseline}
|
||||
\label{sec:agentic-choice}
|
||||
|
||||
The choice reduces to a cost comparison the runtime can estimate. The
|
||||
Stationary-KV policy pays for $Q$ multicast and a hierarchical reduction;
|
||||
the Distributed-Q policy pays for moving or replicating KV plus extra
|
||||
scheduling:
|
||||
\[
|
||||
T_{\text{SK}} = T_{Q\text{-mcast}} + T_{\text{local GEMM}} + T_{\text{hier.\ reduce}},
|
||||
\]
|
||||
\[
|
||||
\begin{aligned}
|
||||
T_{\text{DQ}} = {}& T_{Q\text{-part}} + T_{\text{remote/repl.\ KV}} + T_{\text{local GEMM}} \\
|
||||
&+ T_{\text{grp.\ reduce}} + T_{\text{remap}}.
|
||||
\end{aligned}
|
||||
\]
|
||||
For the assumed mapping (one KV head $=$ 4 CUBEs $=$ 32 PEs), the KV cache
|
||||
is large and stationary while $Q$ is comparatively small, so $T_{\text{SK}}$
|
||||
is the lower cost and the Stationary-KV policy is the recommended baseline.
|
||||
A future runtime can compute both estimates per launch---from the current
|
||||
agent count, $Q$ size, KV-head mapping, and interconnect state---and switch
|
||||
to the Distributed-Q policy only in the regime where a very large $Q$ batch
|
||||
makes multicast and reduction traffic outweigh the cost of remote or
|
||||
regrouped KV. Until that regime is measured, the Distributed-Q policy
|
||||
remains a designed, not-yet-implemented option.
|
||||
|
||||
\subsection{Fan-in: joining sub-agent branches}
|
||||
\label{sec:agentic-fanin}
|
||||
|
||||
Fan-out is only half of the pattern; after it, each branch produces a
|
||||
private continuation and the parent must synthesise them. We treat fan-in
|
||||
as a runtime/framework design problem with a clear optimization ladder.
|
||||
|
||||
\paragraph{Problem.} The branch KV caches \emph{cannot} be concatenated.
|
||||
Each token's K/V depends on its full causal history, so stacking several
|
||||
branches' private KV does not form the cache of any single valid
|
||||
sequence. Join must therefore happen at the token/text level, and its cost
|
||||
is dominated by the number of join-input tokens, because the new main-agent
|
||||
KV grows in proportion to them.
|
||||
|
||||
\paragraph{Baseline.} A naive full-text gather concatenates every branch's
|
||||
raw output: $8$ agents $\times\ 1000$ tokens $=\ 8000$ tokens pushed through
|
||||
every layer during continuation prefill---inflating prefill work, KV
|
||||
allocation, and later decode-time KV reads. This is the cost the ladder
|
||||
below drives down.
|
||||
|
||||
\paragraph{Optimization 1 --- schema-constrained results.} Constrain each
|
||||
sub-agent to emit a compact structured result (claim, confidence, evidence
|
||||
handles) instead of free text, cutting join input by an order of magnitude
|
||||
($8\times1000 \to 8\times50$).
|
||||
|
||||
\paragraph{Optimization 2 --- deduplication and ranking.} Overlapping
|
||||
findings across branches are merged and ranked in the framework before they
|
||||
reach the main context. This is preprocessing, not reasoning, and shrinks
|
||||
the input further without involving the main model.
|
||||
|
||||
\paragraph{Optimization 3 --- pointer-based evidence.} Detailed evidence
|
||||
stays outside the main context behind handles, materialised only when the
|
||||
main agent actually requests it, so the effective input is summaries plus
|
||||
only the evidence truly used.
|
||||
|
||||
\paragraph{Optimization 4 --- hierarchical reducers.} For wide fan-out,
|
||||
intermediate reducer agents summarise groups of branches, shrinking the
|
||||
final join prompt and organising fan-in traffic hierarchically---mirroring
|
||||
how the hierarchical online-softmax merge organises attention reduction.
|
||||
|
||||
\paragraph{Future --- latent-state join.} A more aggressive step replaces
|
||||
text outputs with a few learned latent tokens, further cutting join prefill
|
||||
and KV growth. This requires training the main model to consume latent
|
||||
tokens and aligning branch representations; it is a model--system co-design
|
||||
direction, not a drop-in runtime optimization, and is out of scope here.
|
||||
|
||||
\medskip
|
||||
\noindent Throughout, the shared prefix KV is reused by page-table
|
||||
reference and only the join suffix is new, so shortening join input is the
|
||||
primary lever on continuation-prefill cost, KV growth, and later
|
||||
decode-time KV reads. Taken together, \S\ref{sec:agentic-policyA}--\ref{sec:agentic-fanin}
|
||||
show why this workload is a natural extension of the 1H design rather than a
|
||||
new one: the decisive hardware levers of this report---cheap composite
|
||||
issue and a fast on-device reduction path---are exactly what make agentic
|
||||
fan-out efficient, and no part of the attention math is altered to get
|
||||
there. Quantifying the fan-out speedup and the fan-in join savings on
|
||||
KernBench is left as measured 2H work.
|
||||
@@ -0,0 +1,67 @@
|
||||
\section{Hardware Performance-Spec Search for GQA}
|
||||
\label{sec:hwspec}
|
||||
|
||||
\emph{Work in progress --- this section states the study's intent and
|
||||
method; experimental data, figures, and conclusions are not yet
|
||||
available and will be added in a later revision.}
|
||||
|
||||
The studies so far fixed the modeled hardware (\S\ref{sec:hw}) and
|
||||
varied the \emph{software}: the composite command, the reduction path, and
|
||||
the KV placement. This section inverts the question. Given the fused GQA
|
||||
kernel of \S\ref{sec:gqa} as the target workload, \emph{which hardware
|
||||
specification best serves it}, and where does spending more silicon stop
|
||||
paying off? The discussion of \S\ref{sec:discussion} already gives a
|
||||
strong prior---attention is data-movement bound, so raw MAC throughput is
|
||||
not the binding constraint---but that conclusion was drawn at one operating
|
||||
point. A systematic sweep is needed to turn it into a defensible
|
||||
performance-spec recommendation across context lengths, agent counts, and
|
||||
sharding regimes.
|
||||
|
||||
\subsection{Sweep axes}
|
||||
\label{sec:hwspec-axes}
|
||||
|
||||
Four hardware knobs are varied, spanning the compute side and both levels
|
||||
of the interconnect so that the compute/communication balance can be
|
||||
located rather than assumed. The KernBench cost model already exposes each
|
||||
as a parameter, so the sweep reuses the same deterministic engine as the
|
||||
rest of the report; no production change is required.
|
||||
|
||||
\begin{table}[h]
|
||||
\centering
|
||||
\caption{Hardware knobs varied in the performance-spec search. Ranges and
|
||||
step counts are to be finalised with the experiment.}
|
||||
\label{tab:hwspec-axes}
|
||||
\small
|
||||
\begin{tabular}{@{}p{0.30\textwidth}p{0.14\textwidth}@{}}
|
||||
\toprule
|
||||
\textbf{Knob} & \textbf{Axis} \\
|
||||
\midrule
|
||||
GEMM throughput (TFLOPS) & compute \\
|
||||
MATH-engine ALU count & compute \\
|
||||
CUBE-to-CUBE (die-to-die) bandwidth & interconnect \\
|
||||
SIP-to-SIP (card-to-card) bandwidth & interconnect \\
|
||||
\bottomrule
|
||||
\end{tabular}
|
||||
\end{table}
|
||||
|
||||
The first two axes scale the on-PE compute engines; the latter two scale
|
||||
the two inter-device links that carry the KV reduction and cross-device
|
||||
softmax merge. Sweeping them jointly---rather than one at a time---is what
|
||||
exposes the interactions: for example, whether added die-to-die bandwidth
|
||||
only helps once card-to-card bandwidth is also raised, or whether GEMM
|
||||
throughput is genuinely inert for attention once issue overhead is removed.
|
||||
|
||||
\subsection{Method and target output}
|
||||
\label{sec:hwspec-method}
|
||||
|
||||
The intended procedure is a joint sweep over the four axes, running the
|
||||
fused GQA kernel at representative decode and prefill configurations
|
||||
(including the agentic fan-out shapes of \S\ref{sec:agentic}), and reading
|
||||
latency and per-engine busy time from the engine's completion timestamps
|
||||
and op log---the same measurement path used in \S\ref{sec:gqa}. The target
|
||||
deliverables are: (i) the sensitivity of GQA latency to each knob in
|
||||
isolation, (ii) the Pareto frontier of latency against a simple
|
||||
area/cost proxy, and (iii) a recommended balanced specification---the knob
|
||||
combination past which further investment does not move GQA latency. These
|
||||
results are the natural quantitative counterpart to the qualitative
|
||||
ranking in \S\ref{sec:discussion}, and completing them is a 2H objective.
|
||||
@@ -40,7 +40,10 @@ hardware levers for this workload:
|
||||
GQA panels leave the GEMM and vector-math engines two to three orders of
|
||||
magnitude below the DMA engine in busy time. Adding MAC area would not move
|
||||
decode latency; the workload cannot use it. This is the single most
|
||||
actionable finding for an attention-dominated roadmap.
|
||||
actionable finding for an attention-dominated roadmap. The implication is
|
||||
that future hardware investment should prioritize communication and
|
||||
memory-system efficiency over additional compute throughput for
|
||||
attention-dominated inference workloads.
|
||||
|
||||
\paragraph{Caveats.} These conclusions are achievable-kernel results from a
|
||||
deterministic model, not E2E measurements; absolute numbers carry the
|
||||
@@ -1,4 +1,4 @@
|
||||
\section{Future Work --- 2H}
|
||||
\section{Future Work}
|
||||
\label{sec:future}
|
||||
|
||||
The 1H work covered attention end to end. The natural next step is to
|
||||
@@ -30,11 +30,22 @@ vs.\ sequence parallelism---under a single, software-stack-independent
|
||||
model, so the interconnect and memory implications of each choice are
|
||||
measured rather than assumed.
|
||||
|
||||
\paragraph{Agentic workloads.} Agentic inference interleaves many
|
||||
short, bursty decode requests with tool use and long shared contexts,
|
||||
which stresses the system differently from a single long generation:
|
||||
context reuse across requests, dynamic batching, and uneven expert load all
|
||||
change how compute and data should be dispersed. Characterizing how total
|
||||
compute and data movement distribute across the SIP/CUBE/PE hierarchy under
|
||||
such workloads---and which of the 1H hardware levers still dominate when the
|
||||
workload is this irregular---is the broader 2H agenda.
|
||||
\paragraph{Agentic workloads: from design to implementation.}
|
||||
Section~\ref{sec:agentic} established \emph{how} the fused GQA design
|
||||
extends to agentic fan-out/fan-in and \emph{which} responsibilities fall to
|
||||
the framework, the runtime, and the kernel. The 2H step is no longer to
|
||||
analyse the workload but to \emph{build} that support: the detailed design
|
||||
and implementation of all three levels. This is necessary work rather than
|
||||
optional, because today's agentic frameworks are effectively all
|
||||
closed-source---there is no open substrate that exposes shared-prefix KV
|
||||
reuse, cross-agent query batching, and schema/pointer-based fan-in down to
|
||||
the hardware. Concretely, 2H targets: a \textbf{framework} layer that
|
||||
manages the agent tree, identifies co-schedulable branches, and enforces
|
||||
compact fan-in; a \textbf{runtime} layer that forks logical contexts by
|
||||
page table without copying shared KV, batches sub-agent query rows, and
|
||||
selects the execution policy; and a \textbf{kernel} layer that runs the
|
||||
batched replicated-Q attention and hierarchical online-softmax merge over
|
||||
PE\_IPCQ. Bringing these up on KernBench turns the
|
||||
Section~\ref{sec:agentic} design into a measured, end-to-end agentic path
|
||||
and lets us confirm which of the 1H hardware levers still dominate once the
|
||||
full framework/runtime/kernel stack is in place.
|
||||
@@ -45,17 +45,28 @@
|
||||
6. **Fused Grouped Query Attention** (composite + PE_IPCQ) — `sections/05-gqa.tex`
|
||||
necessity · design · results · analysis. (+ GQA seq/head/user configs)
|
||||
|
||||
7. **Discussion** — `sections/06-discussion.tex`
|
||||
7. **Supporting Agentic Workloads** — `sections/06-agentic.tex`
|
||||
How the §6 GQA design extends to agentic fan-out/fan-in: shared-prefix KV
|
||||
reuse, batched replicated-Q attention, schema/pointer-based fan-in.
|
||||
Implementation split across three levels (agentic framework / runtime /
|
||||
kernel) with each component's responsibilities. Design, not yet measured.
|
||||
|
||||
8. **Hardware Performance-Spec Search for GQA** — `sections/07-hw-spec-search.tex`
|
||||
*Work in progress.* Joint sweep of GEMM TFLOPS, MATH-engine ALU count,
|
||||
CUBE↔CUBE (die-to-die) BW, SIP↔SIP (card-to-card) BW to find the balanced
|
||||
spec for GQA. Intent + method only; data/figures/conclusions deferred.
|
||||
|
||||
9. **Discussion** — `sections/08-discussion.tex`
|
||||
Which HW changes are meaningful, and under what regimes (cross-cutting).
|
||||
|
||||
8. **Conclusion** — `sections/07-conclusion.tex`
|
||||
The codesign thesis, stated plainly, supported by the measured results.
|
||||
10. **Conclusion** — `sections/09-conclusion.tex`
|
||||
The codesign thesis, stated plainly, supported by the measured results.
|
||||
|
||||
9. **Future Work — 2H** — `sections/08-future-work.tex`
|
||||
Add the FFN/MoE layer toward full LLM decoding; how compute & data should
|
||||
be distributed for agentic / MoE workloads.
|
||||
11. **Future Work — 2H** — `sections/10-future-work.tex`
|
||||
Add the FFN/MoE layer toward full LLM decoding; how compute & data should
|
||||
be distributed for agentic / MoE workloads.
|
||||
|
||||
10. **References** *(optional)* — external literature only (FlashAttention,
|
||||
12. **References** *(optional)* — external literature only (FlashAttention,
|
||||
Megatron-LM, GPT-3, Llama 3). No ADR/SPEC entries.
|
||||
|
||||
## Per-section structure (§4/§5/§6)
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
mode,kv_per_cube,C,S_kv,B,capacity,agg_latency_us,mean_user_latency_us,throughput_users_per_us
|
||||
A1,1,8,8192,1,2,2.2471645000005376,2.2471645000005376,0.44500524994932983
|
||||
A1,1,8,8192,2,2,2.3131845000012543,2.2471645000012357,0.864608940617973
|
||||
A2,2,4,8192,1,4,3.3555127500005475,3.3555127500005475,0.2980170467240325
|
||||
A2,2,4,8192,2,4,3.3935227500010514,3.355512750001042,0.5893580645656141
|
||||
A2,2,4,8192,4,4,3.4695427500010703,3.3555127500011586,1.1528896711241752
|
||||
A4,4,2,8192,1,8,5.586135500000557,5.586135500000557,0.17901463364071643
|
||||
A4,4,2,8192,2,8,5.602295500000823,5.586135500000673,0.35699652044411195
|
||||
A4,4,2,8192,8,8,5.783285499986028,5.586135499984957,1.3832967436968704
|
||||
B,8,1,8192,1,16,10.816091000000947,10.816091000000947,0.09245484343649775
|
||||
B,8,1,8192,2,16,10.848411000001535,10.816091000001236,0.18435879687815265
|
||||
B,8,1,8192,16,16,11.38492099986691,10.816090999974403,1.4053676788962384
|
||||
|
@@ -0,0 +1,69 @@
|
||||
"""Batch-scaling figure for concurrent decode users on one SIP.
|
||||
|
||||
Reads ``docs/sweeps/decode_batch_sweep.csv`` (from
|
||||
``tests/attention/test_gqa_decode_batch_sweep.py``) and plots aggregate
|
||||
throughput (requests / us) versus batch size B for the four KV mappings,
|
||||
one panel per context length. Each mapping's curve runs up to its SIP
|
||||
capacity 16/C (A1=2 ... B=16 users).
|
||||
"""
|
||||
from pathlib import Path
|
||||
|
||||
import csv
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
CSV = ROOT / "docs" / "sweeps" / "decode_batch_sweep.csv"
|
||||
OUT = (ROOT / "docs" / "report" / "1H-codesign-paper"
|
||||
/ "figures" / "gqa_short_context" / "batch_scaling.png")
|
||||
|
||||
MODE_LABEL = {"A1": "1-kv-per-cube", "A2": "2-kv-per-cube",
|
||||
"A4": "4-kv-per-cube", "B": "8-kv-per-cube"}
|
||||
MODE_COLOR = {"A1": "#1f77b4", "A2": "#2ca02c",
|
||||
"A4": "#ffd000", "B": "#d62728"}
|
||||
MODES = ["A1", "A2", "A4", "B"]
|
||||
|
||||
|
||||
def main():
|
||||
rows = list(csv.DictReader(CSV.open()))
|
||||
contexts = sorted({int(r["S_kv"]) for r in rows})
|
||||
fig, axes = plt.subplots(1, len(contexts), figsize=(6.5 * len(contexts), 4.5),
|
||||
squeeze=False)
|
||||
xmax = max(int(r["B"]) for r in rows) # largest SIP capacity (= 16)
|
||||
for col, S in enumerate(contexts):
|
||||
ax = axes[0][col]
|
||||
for m in MODES:
|
||||
pts = sorted(
|
||||
((int(r["B"]), float(r["throughput_users_per_us"]))
|
||||
for r in rows if r["mode"] == m and int(r["S_kv"]) == S),
|
||||
key=lambda t: t[0],
|
||||
)
|
||||
if not pts:
|
||||
continue
|
||||
xs, ys = zip(*pts)
|
||||
# Measured: concurrent users up to the SIP capacity 16/C.
|
||||
ax.plot(xs, ys, marker="o", color=MODE_COLOR[m],
|
||||
label=MODE_LABEL[m])
|
||||
# Beyond capacity the SIP is full: extra users run in waves at
|
||||
# the saturated rate, so throughput plateaus. Draw that ceiling
|
||||
# as a dashed extension (projected, not concurrently measured).
|
||||
if xs[-1] < xmax:
|
||||
ax.plot([xs[-1], xmax], [ys[-1], ys[-1]],
|
||||
linestyle="--", color=MODE_COLOR[m], alpha=0.55)
|
||||
ax.annotate(f"{ys[-1]:.2f}", (xs[-1], ys[-1]),
|
||||
textcoords="offset points", xytext=(4, 4), fontsize=8)
|
||||
ax.set_title(f"S_kv = {S // 1024}K")
|
||||
ax.set_xlabel("batch size B (concurrent users)")
|
||||
ax.set_ylabel("throughput (requests / µs)")
|
||||
ax.grid(True, alpha=0.3)
|
||||
ax.legend(fontsize=9)
|
||||
fig.suptitle("Batch scaling: aggregate throughput vs concurrent users "
|
||||
"(one SIP)", fontsize=12, fontweight="bold")
|
||||
fig.tight_layout()
|
||||
OUT.parent.mkdir(parents=True, exist_ok=True)
|
||||
fig.savefig(OUT, dpi=140, bbox_inches="tight")
|
||||
print(f" ✓ {OUT}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,231 @@
|
||||
"""Generate 6 comparison figures for the short-context GQA attention sweep.
|
||||
|
||||
Sweep matrix: 3 variants × 2 phases × 4 modes × 4 contexts (8K/16K/32K/64K).
|
||||
|
||||
Figures
|
||||
Wall-clock latency:
|
||||
1. wall_variant1_mode_compare.png — (1) without composite
|
||||
2. wall_variant2_mode_compare.png — (2) with composite (GEMM-only)
|
||||
3. wall_variant3_mode_compare.png — (3) with composite + softmax_merge
|
||||
4. wall_a1_variant_compare.png — 1-kv-per-cube: variant comparison
|
||||
|
||||
Mode trade-off + composite ablation:
|
||||
5. per_cube_tradeoff_mode_compare.png — HBM BW + IPCQ (variant-invariant)
|
||||
6. gemm_util_a1_variant_ablation.png — 1-kv-per-cube GEMM util
|
||||
|
||||
Output: docs/report/1H-codesign-paper/figures/gqa_short_context/
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[3]
|
||||
SWEEPS = ROOT / "docs" / "sweeps"
|
||||
OUT = (ROOT / "docs" / "report" / "1H-codesign-paper"
|
||||
/ "figures" / "gqa_short_context")
|
||||
OUT.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
CONTEXTS = [8192, 16384, 32768, 65536]
|
||||
CONTEXT_LABELS = ["8K", "16K", "32K", "64K"]
|
||||
MODES = ["A1", "A2", "A4", "B"]
|
||||
MODE_LABEL = {
|
||||
"A1": "1-kv-per-cube",
|
||||
"A2": "2-kv-per-cube",
|
||||
"A4": "4-kv-per-cube",
|
||||
"B": "8-kv-per-cube",
|
||||
}
|
||||
MODE_COLOR = {
|
||||
MODE_LABEL["A1"]: "#1f77b4",
|
||||
MODE_LABEL["A2"]: "#2ca02c",
|
||||
MODE_LABEL["A4"]: "#ffd000",
|
||||
MODE_LABEL["B"]: "#d62728",
|
||||
}
|
||||
VARIANTS = [
|
||||
("baseline", "(1) without composite"),
|
||||
("composite", "(2) with composite (GEMM-only)"),
|
||||
("composite_fused", "(3) with composite + softmax_merge"),
|
||||
]
|
||||
VARIANT_COLOR = {
|
||||
"baseline": "#1f77b4",
|
||||
"composite": "#2ca02c",
|
||||
"composite_fused": "#d62728",
|
||||
}
|
||||
|
||||
CSV_MAP = {
|
||||
("decode", "baseline"): SWEEPS / "short_context_decode_sweep.csv",
|
||||
("decode", "composite"): SWEEPS / "short_context_decode_composite_sweep.csv",
|
||||
("decode", "composite_fused"): SWEEPS / "short_context_decode_composite_fused_sweep.csv",
|
||||
("prefill", "baseline"): SWEEPS / "short_context_prefill_sweep.csv",
|
||||
("prefill", "composite"): SWEEPS / "short_context_prefill_composite_sweep.csv",
|
||||
("prefill", "composite_fused"): SWEEPS / "short_context_prefill_composite_fused_sweep.csv",
|
||||
}
|
||||
|
||||
|
||||
def load(phase, variant):
|
||||
"""Return {(mode, S_kv): row dict}."""
|
||||
rows = {}
|
||||
with CSV_MAP[(phase, variant)].open() as f:
|
||||
for r in csv.DictReader(f):
|
||||
rows[(r["mode"], int(r["S_kv"]))] = r
|
||||
return rows
|
||||
|
||||
|
||||
def grouped_bars(ax, *, x_labels, groups, group_color, ylabel, title,
|
||||
y_log=False, value_fmt="{:.3g}"):
|
||||
"""Grouped bars: x = x_labels, groups = list of (label, values)."""
|
||||
n = len(groups)
|
||||
nx = len(x_labels)
|
||||
width = 0.8 / n
|
||||
x = np.arange(nx)
|
||||
for i, (label, values) in enumerate(groups):
|
||||
offset = (i - (n - 1) / 2) * width
|
||||
bars = ax.bar(x + offset, values, width, label=label,
|
||||
color=group_color[label], edgecolor="black",
|
||||
linewidth=0.4)
|
||||
for bar, v in zip(bars, values):
|
||||
if v <= 0 and y_log:
|
||||
continue
|
||||
ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height(),
|
||||
value_fmt.format(v), ha="center", va="bottom",
|
||||
fontsize=7)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(x_labels)
|
||||
ax.set_xlabel("S_kv")
|
||||
ax.set_ylabel(ylabel)
|
||||
ax.set_title(title)
|
||||
if y_log:
|
||||
ax.set_yscale("log")
|
||||
ax.grid(True, axis="y", alpha=0.3, which="both")
|
||||
ax.legend(fontsize=8, loc="upper left")
|
||||
|
||||
|
||||
# ── 1-3. Wall mode comparison per variant ──────────────────────────
|
||||
|
||||
def fig_wall_mode_compare(variant_key, suptitle, fname):
|
||||
fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
|
||||
for ax, phase in zip(axes, ("prefill", "decode")):
|
||||
rows = load(phase, variant_key)
|
||||
groups = [
|
||||
(MODE_LABEL[m], [float(rows[(m, s)]["wall_us"]) for s in CONTEXTS])
|
||||
for m in MODES
|
||||
]
|
||||
grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
|
||||
group_color=MODE_COLOR,
|
||||
ylabel="Wall-clock latency (μs)",
|
||||
title=phase.capitalize(),
|
||||
y_log=True, value_fmt="{:.0f}")
|
||||
fig.suptitle(suptitle, fontsize=12, fontweight="bold")
|
||||
fig.tight_layout()
|
||||
out = OUT / fname
|
||||
fig.savefig(out, dpi=140, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f" ✓ {out}")
|
||||
|
||||
|
||||
# ── 4. Wall 1-kv-per-cube variant comparison ───────────────────────
|
||||
|
||||
def fig_wall_a1_variant_compare(fname):
|
||||
fig, axes = plt.subplots(1, 2, figsize=(13, 4.8))
|
||||
for ax, phase in zip(axes, ("prefill", "decode")):
|
||||
groups = []
|
||||
for variant_key, _ in VARIANTS:
|
||||
rows = load(phase, variant_key)
|
||||
values = [float(rows[("A1", s)]["wall_us"]) for s in CONTEXTS]
|
||||
groups.append((variant_key, values))
|
||||
grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
|
||||
group_color=VARIANT_COLOR,
|
||||
ylabel="Wall-clock latency (μs)",
|
||||
title=phase.capitalize(),
|
||||
y_log=True, value_fmt="{:.0f}")
|
||||
ax.legend([v[1] for v in VARIANTS], fontsize=8, loc="upper left")
|
||||
fig.suptitle("1-kv-per-cube: variant comparison",
|
||||
fontsize=12, fontweight="bold")
|
||||
fig.tight_layout()
|
||||
out = OUT / fname
|
||||
fig.savefig(out, dpi=140, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f" ✓ {out}")
|
||||
|
||||
|
||||
# ── 5. Per-cube trade-off (HBM BW + IPCQ, mode compare) ────────────
|
||||
|
||||
def fig_per_cube_tradeoff(fname):
|
||||
fig, axes = plt.subplots(2, 2, figsize=(13, 9))
|
||||
metrics = [
|
||||
("kv_cache_per_cube_mb", "Per-cube KV footprint (MB)", True, "{:.0f}"),
|
||||
("ipcq_kb", "IPCQ traffic (KB)", True, "{:.1f}"),
|
||||
]
|
||||
for row, (key, ylabel, log, fmt) in enumerate(metrics):
|
||||
for col, phase in enumerate(("prefill", "decode")):
|
||||
ax = axes[row][col]
|
||||
rows = load(phase, "baseline")
|
||||
groups = [
|
||||
(MODE_LABEL[m], [float(rows[(m, s)][key]) for s in CONTEXTS])
|
||||
for m in MODES
|
||||
]
|
||||
grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
|
||||
group_color=MODE_COLOR, ylabel=ylabel,
|
||||
title=phase.capitalize(),
|
||||
y_log=log, value_fmt=fmt)
|
||||
fig.suptitle("Per-cube trade-off: KV footprint + IPCQ mode comparison",
|
||||
fontsize=12, fontweight="bold")
|
||||
fig.tight_layout()
|
||||
out = OUT / fname
|
||||
fig.savefig(out, dpi=140, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f" ✓ {out}")
|
||||
|
||||
|
||||
# ── 6. GEMM util ablation (1-kv-per-cube, variant compare) ─────────
|
||||
|
||||
def fig_gemm_util_a1_ablation(fname):
|
||||
fig, axes = plt.subplots(1, 2, figsize=(13, 5.4))
|
||||
for ax, phase in zip(axes, ("prefill", "decode")):
|
||||
groups = []
|
||||
for variant_key, _ in VARIANTS:
|
||||
rows = load(phase, variant_key)
|
||||
values = [float(rows[("A1", s)]["gemm_util"]) for s in CONTEXTS]
|
||||
groups.append((variant_key, values))
|
||||
grouped_bars(ax, x_labels=CONTEXT_LABELS, groups=groups,
|
||||
group_color=VARIANT_COLOR,
|
||||
ylabel="GEMM engine utilization (per PE)",
|
||||
title=phase.capitalize(),
|
||||
y_log=False, value_fmt="{:.4f}")
|
||||
ax.legend([v[1] for v in VARIANTS], fontsize=8, loc="upper left")
|
||||
fig.suptitle("1-kv-per-cube: GEMM utilization variant comparison",
|
||||
fontsize=12, fontweight="bold")
|
||||
fig.text(
|
||||
0.5, -0.04,
|
||||
"⚠ Higher gemm_util on composite/fused is largely supertile "
|
||||
"padding overhead (M=G=8 padded to TILE_M=32), not pure fusion "
|
||||
"gain. See ADR-0070 limitation #4.",
|
||||
ha="center", fontsize=9, style="italic", color="#555")
|
||||
fig.tight_layout()
|
||||
out = OUT / fname
|
||||
fig.savefig(out, dpi=140, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f" ✓ {out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(f"Output: {OUT}")
|
||||
fig_wall_mode_compare(
|
||||
"baseline",
|
||||
"(1) Without composite: latency mode comparison",
|
||||
"wall_variant1_mode_compare.png")
|
||||
fig_wall_mode_compare(
|
||||
"composite",
|
||||
"(2) With composite (GEMM-only): latency mode comparison",
|
||||
"wall_variant2_mode_compare.png")
|
||||
fig_wall_mode_compare(
|
||||
"composite_fused",
|
||||
"(3) With composite + softmax_merge: latency mode comparison",
|
||||
"wall_variant3_mode_compare.png")
|
||||
fig_wall_a1_variant_compare("wall_a1_variant_compare.png")
|
||||
fig_per_cube_tradeoff("per_cube_tradeoff_mode_compare.png")
|
||||
fig_gemm_util_a1_ablation("gemm_util_a1_variant_ablation.png")
|
||||
print("\nDone.")
|
||||
@@ -0,0 +1,17 @@
|
||||
variant,mode,kv_per_cube,C,S_kv,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
|
||||
composite_fused,A1,1,8,8192,2.2471645000005376,64,0.11665545624270428,656,0.0,0.457688,16.0625,8.0,57.75,2.0
|
||||
composite_fused,A1,1,8,16384,4.090972500001895,64,0.4907507933676371,1296,13.903615999992937,0.501714,32.0625,8.0,57.75,4.0
|
||||
composite_fused,A1,1,8,32768,6.65601250000007,64,0.8261186408347434,2576,41.71084799996763,0.61606,64.0625,8.0,57.75,8.0
|
||||
composite_fused,A1,1,8,65536,11.786092499984429,64,1.05893212704674,5136,97.32531200143696,0.695438,128.0625,8.0,57.75,16.0
|
||||
composite_fused,A2,2,4,8192,3.8395887500001407,32,0.5228809986460942,624,6.951807999995537,0.534693,16.03125,8.0,24.75,4.0
|
||||
composite_fused,A2,2,4,16384,6.404628750001605,32,0.858544064722257,1264,20.855423999989405,0.640318,32.03125,8.0,24.75,8.0
|
||||
composite_fused,A2,2,4,32768,11.534708749999874,32,1.0820101547616419,2544,48.66265599996224,0.710638,64.03125,8.0,24.75,16.0
|
||||
composite_fused,A2,2,4,65536,21.794868749984772,32,1.21334541192213,5104,104.27712000153959,0.751966,128.03125,8.0,24.75,32.0
|
||||
composite_fused,A4,4,2,8192,5.915787499999919,16,0.9294884239792724,608,10.427711999993306,0.693399,16.015625,8.0,8.25,8.0
|
||||
composite_fused,A4,4,2,16384,11.045867500001041,16,1.1298951395303998,1248,24.33132799998764,0.742178,32.015625,8.0,8.25,16.0
|
||||
composite_fused,A4,4,2,32768,21.306027499999384,16,1.2411841672219757,2528,52.13855999995954,0.769266,64.015625,8.0,8.25,32.0
|
||||
composite_fused,A4,4,2,65536,41.82634749998455,16,1.2999645260103314,5088,107.7530240015909,0.783573,128.015625,8.0,8.25,64.0
|
||||
composite_fused,B,8,1,8192,10.856904999999854,8,1.149560763397397,600,12.16566399999219,0.75528,16.0078125,8.0,0.0,16.0
|
||||
composite_fused,B,8,1,16384,21.117065000001226,8,1.2522906947682175,1240,26.069279999986758,0.776244,32.0078125,8.0,0.0,32.0
|
||||
composite_fused,B,8,1,32768,41.637384999999774,8,1.3058641410537761,2520,53.8765119999582,0.787177,64.0078125,8.0,0.0,64.0
|
||||
composite_fused,B,8,1,65536,82.67802499997313,8,1.33323087973183,5080,109.49097600161657,0.792762,128.0078125,8.0,0.0,128.0
|
||||
|
@@ -0,0 +1,17 @@
|
||||
variant,mode,kv_per_cube,C,S_kv,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
|
||||
composite,A1,1,8,8192,2.532214500000584,64,0.33402541530150626,656,0.0,0.406166,16.0625,8.0,57.75,2.0
|
||||
composite,A1,1,8,16384,4.176996500001463,64,0.40499148132408613,1360,0.0,0.491382,32.0625,8.0,57.75,4.0
|
||||
composite,A1,1,8,32768,7.466560500002,64,0.4531264428807111,2768,0.0,0.549182,64.0625,8.0,57.75,8.0
|
||||
composite,A1,1,8,65536,14.045688499950572,64,0.4817558071541797,5584,0.0,0.58356,128.0625,8.0,57.75,16.0
|
||||
composite,A2,2,4,8192,3.9256127500006404,32,0.4309258471789165,656,0.0,0.522976,16.03125,8.0,24.75,4.0
|
||||
composite,A2,2,4,16384,7.215176750001498,32,0.4689138072801237,1360,0.0,0.568385,32.03125,8.0,24.75,8.0
|
||||
composite,A2,2,4,32768,13.794304750002688,32,0.4905351971318266,2768,0.0,0.594231,64.03125,8.0,24.75,16.0
|
||||
composite,A2,2,4,65536,26.95256074991636,32,0.5021112511805295,5584,0.0,0.608068,128.03125,8.0,24.75,32.0
|
||||
composite,A4,4,2,8192,6.726335500000743,16,0.5029924540606681,656,0.0,0.609842,16.015625,8.0,8.25,8.0
|
||||
composite,A4,4,2,16384,13.305463500001585,16,0.5085574057667106,1360,0.0,0.616138,32.015625,8.0,8.25,16.0
|
||||
composite,A4,4,2,32768,26.463719500003965,16,0.5113863151276257,2768,0.0,0.619338,64.015625,8.0,8.25,32.0
|
||||
composite,A4,4,2,65536,52.78023149984703,16,0.5128126048745716,5584,0.0,0.620952,128.015625,8.0,8.25,64.0
|
||||
composite,B,8,1,8192,13.096491000001318,8,0.5166721375947838,656,0.0,0.626122,16.0078125,8.0,0.0,16.0
|
||||
composite,B,8,1,16384,26.254747000003057,8,0.5154566524738308,1360,0.0,0.624344,32.0078125,8.0,0.0,32.0
|
||||
composite,B,8,1,32768,52.57125900000788,8,0.5148510519664782,2768,0.0,0.623459,64.0078125,8.0,0.0,64.0
|
||||
composite,B,8,1,65536,105.20428299969761,8,0.514548785079354,5584,0.0,0.623016,128.0078125,8.0,0.0,128.0
|
||||
|
@@ -0,0 +1,17 @@
|
||||
mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
|
||||
A1,1,8,8192,149.2685345000007,64,0.001756190618994587,656,0.013780533230866485,16.0625,8.0,57.75,2.0
|
||||
A1,1,8,16384,282.2252665000025,64,0.0018576933472394743,1360,0.014545118695104373,32.0625,8.0,57.75,4.0
|
||||
A1,1,8,32768,548.138730500006,64,0.0019129755692392889,2768,0.014961540835691614,64.0625,8.0,57.75,8.0
|
||||
A1,1,8,65536,1080.073058500367,64,0.0019416760593127277,5584,0.01517767698303756,128.0625,8.0,57.75,16.0
|
||||
A2,2,4,8192,142.16883275000066,32,0.0036877843748065785,656,0.028881154333033524,16.03125,8.0,24.75,4.0
|
||||
A2,2,4,16384,277.0102967500018,32,0.0037853322143696724,1360,0.029609007665885423,32.03125,8.0,24.75,8.0
|
||||
A2,2,4,32768,546.8006247504073,32,0.0038353138330044275,2768,0.02998167752182655,64.03125,8.0,24.75,16.0
|
||||
A2,2,4,65536,1087.0962007518285,32,0.003858263874988177,5584,0.03015188534127058,128.03125,8.0,24.75,32.0
|
||||
A4,4,2,8192,142.21417550000066,16,0.007373217165681769,656,0.05768763888097753,16.015625,8.0,8.25,8.0
|
||||
A4,4,2,16384,280.8788035001758,16,0.007466394665122732,1360,0.058373931374247456,32.015625,8.0,8.25,16.0
|
||||
A4,4,2,32768,558.5655195010398,16,0.007509063580845671,2768,0.058686042828569145,64.015625,8.0,8.25,32.0
|
||||
A4,4,2,65536,1114.319951501857,16,0.007528006645388841,5584,0.05882332081702009,128.015625,8.0,8.25,64.0
|
||||
B,8,1,8192,147.93485600006025,8,0.014176185766516818,656,0.11085960701508589,16.0078125,8.0,0.0,16.0
|
||||
B,8,1,16384,294.3171420004266,8,0.014250967413897551,1360,0.1113900460475132,32.0078125,8.0,0.0,32.0
|
||||
B,8,1,32768,587.2722140010417,8,0.01428401991446495,2768,0.11162115018757507,64.0078125,8.0,0.0,64.0
|
||||
B,8,1,65536,1173.1823580017838,8,0.014300603725891686,5584,0.11173710472707321,128.0078125,8.0,0.0,128.0
|
||||
|
@@ -0,0 +1,17 @@
|
||||
mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
|
||||
A1,1,8,8192,2.2471645000005376,64,0.11665545624270428,656,0.457688,16.0625,8.0,57.75,2.0
|
||||
A1,1,8,16384,3.6068965000013704,64,0.14535709577464812,1360,0.569049,32.0625,8.0,57.75,4.0
|
||||
A1,1,8,32768,6.326360500001814,64,0.1657471147905734,2768,0.648161,64.0625,8.0,57.75,8.0
|
||||
A1,1,8,65536,11.765288499959512,64,0.17824909265982342,5584,0.696668,128.0625,8.0,57.75,16.0
|
||||
A2,2,4,8192,3.3555127500005475,32,0.15624676139283236,656,0.611829,16.03125,8.0,24.75,4.0
|
||||
A2,2,4,16384,6.074976750001311,32,0.17260576347056114,1360,0.675064,32.03125,8.0,24.75,8.0
|
||||
A2,2,4,32768,11.513904750002315,32,0.18214081543451005,2768,0.711922,64.03125,8.0,24.75,16.0
|
||||
A2,2,4,65536,22.391760749934242,32,0.18731461303283142,5584,0.731921,128.03125,8.0,24.75,32.0
|
||||
A4,4,2,8192,5.586135500000557,16,0.1877104484844272,656,0.734318,16.015625,8.0,8.25,8.0
|
||||
A4,4,2,16384,11.025063500001211,16,0.19021677290108474,1360,0.743578,32.015625,8.0,8.25,16.0
|
||||
A4,4,2,32768,21.902919500003218,16,0.19149520227204342,2768,0.748302,64.015625,8.0,8.25,32.0
|
||||
A4,4,2,65536,43.65863149988279,16,0.19214088284049563,5584,0.750688,128.015625,8.0,8.25,64.0
|
||||
B,8,1,8192,10.816091000000947,8,0.1938918598225168,656,0.75813,16.0078125,8.0,0.0,16.0
|
||||
B,8,1,16384,21.693947000002314,8,0.1933398288471476,1360,0.755602,32.0078125,8.0,0.0,32.0
|
||||
B,8,1,32768,43.449659000006385,8,0.19306499045255035,2768,0.754344,64.0078125,8.0,0.0,64.0
|
||||
B,8,1,65536,86.96108299976913,8,0.19292786406575965,5584,0.753716,128.0078125,8.0,0.0,128.0
|
||||
|
@@ -0,0 +1,17 @@
|
||||
variant,mode,kv_per_cube,C,S_kv,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
|
||||
composite_fused,A1,1,8,8192,35.53158299944201,64,0.3512557265116541,4800,97.32531200143696,0.0290446,16.0625,64.0,114688.0,2.0
|
||||
composite_fused,A1,1,8,16384,63.02314799463935,64,0.4196030322618738,9920,208.55424000307917,0.0326229,32.0625,64.0,229376.0,4.0
|
||||
composite_fused,A1,1,8,32768,118.00299802135117,64,0.4607744625054221,20160,431.01209600543973,0.0347788,64.0625,64.0,458752.0,8.0
|
||||
composite_fused,A1,1,8,65536,227.95666800683736,64,0.4835519703723695,40640,875.9278079059123,0.0359717,128.0625,64.0,917504.0,16.0
|
||||
composite_fused,A2,2,4,8192,35.27508350025257,32,0.5747630902258046,2400,96.40870400077105,0.0585116,16.0625,64.0,49152.0,4.0
|
||||
composite_fused,A2,2,4,16384,65.58160349091561,32,0.653336388872743,4960,206.59008000165224,0.0627005,32.0625,64.0,98304.0,8.0
|
||||
composite_fused,A2,2,4,32768,126.19361352364719,32,0.69726913692688,10080,426.9528320015669,0.0650429,64.0625,64.0,196608.0,16.0
|
||||
composite_fused,A2,2,4,65536,247.42072338639758,32,0.7205501525262498,20320,867.678335950613,0.0662839,128.0625,64.0,393216.0,32.0
|
||||
composite_fused,A4,4,2,8192,45.62825600103196,16,0.7872118539828873,1200,92.28038400042057,0.0904703,16.0625,64.0,16384.0,8.0
|
||||
composite_fused,A4,4,2,16384,87.81637548340065,16,0.8628360665552867,2480,197.74368000090124,0.09365,32.0625,64.0,32768.0,16.0
|
||||
composite_fused,A4,4,2,32768,172.19397554247547,16,0.9029073142513345,5040,408.6702720000148,0.0953343,64.0625,64.0,65536.0,32.0
|
||||
composite_fused,A4,4,2,65536,340.9491752808783,16,0.9235491703853554,10160,830.5234559737444,0.096202,128.0625,64.0,131072.0,64.0
|
||||
composite_fused,B,8,1,8192,68.33786600000411,8,2.036048828365075,712,121.6138560003899,0.120811,16.0625,64.0,0.0,16.0
|
||||
composite_fused,B,8,1,16384,133.87386600000133,8,2.209234414765302,1480,260.6011200008355,0.122862,32.0625,64.0,0.0,32.0
|
||||
composite_fused,B,8,1,32768,264.9458659999771,8,2.2985744263323977,3016,538.575647996068,0.12392,64.0625,64.0,0.0,64.0
|
||||
composite_fused,B,8,1,65536,527.0898659999762,8,2.3439567929361727,6088,1094.5247039788662,0.124457,128.0625,64.0,0.0,128.0
|
||||
|
@@ -0,0 +1,17 @@
|
||||
variant,mode,kv_per_cube,C,S_kv,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
|
||||
composite,A1,1,8,8192,35.312546999435405,64,0.19162004939598534,5248,0.0,0.0292247,16.0625,64.0,114688.0,2.0
|
||||
composite,A1,1,8,16384,62.80411199463298,64,0.21548245124172694,10880,0.0,0.0327367,32.0625,64.0,229376.0,4.0
|
||||
composite,A1,1,8,32768,117.78396202135924,64,0.2297967187462748,22144,0.0,0.0348435,64.0625,64.0,458752.0,8.0
|
||||
composite,A1,1,8,65536,227.7376320068799,64,0.23769780826259174,44672,0.0,0.0360063,128.0625,64.0,917504.0,16.0
|
||||
composite,A2,2,4,8192,35.873466499721395,32,0.21785371648964683,2624,0.0,0.0575356,16.0625,64.0,49152.0,4.0
|
||||
composite,A2,2,4,16384,67.67155399355386,32,0.23097350479268527,5440,0.0,0.0607641,32.0625,64.0,98304.0,8.0
|
||||
composite,A2,2,4,32768,129.35661752446276,32,0.24166271963282046,11072,0.0,0.0634525,64.0625,64.0,196608.0,16.0
|
||||
composite,A2,2,4,65536,254.34501247150266,32,0.24581313146006306,22336,0.0,0.0644793,128.0625,64.0,393216.0,32.0
|
||||
composite,A4,4,2,8192,46.88723350110906,16,0.2114076532174909,1312,0.0,0.088041,16.0625,64.0,16384.0,8.0
|
||||
composite,A4,4,2,16384,89.10037348339473,16,0.22249783278807841,2720,0.0,0.0923004,32.0625,64.0,32768.0,16.0
|
||||
composite,A4,4,2,32768,173.47797354247328,16,0.22855512537871095,5536,0.0,0.0946287,64.0625,64.0,65536.0,32.0
|
||||
composite,A4,4,2,65536,342.2399932809714,16,0.23170453934016594,11168,0.0,0.0958392,128.0625,64.0,131072.0,64.0
|
||||
composite,B,8,1,8192,68.73804599999497,8,0.3341725483628459,656,0.0,0.120108,16.0625,64.0,0.0,16.0
|
||||
composite,B,8,1,16384,134.27404599999218,8,0.3421415930417734,1360,0.0,0.122496,32.0625,64.0,0.0,32.0
|
||||
composite,B,8,1,32768,265.34604599999824,8,0.3462703641501576,2768,0.0,0.123733,64.0625,64.0,0.0,64.0
|
||||
composite,B,8,1,65536,527.4900460000299,8,0.3483723443539211,5584,0.0,0.124363,128.0625,64.0,0.0,128.0
|
||||
|
@@ -0,0 +1,17 @@
|
||||
mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
|
||||
A1,1,8,8192,1103.359903999566,64,0.0019006962210579498,5248,0.001870649814732448,16.0625,64.0,114688.0,2.0
|
||||
A1,1,8,16384,2182.457823968016,64,0.0019218259129421755,10880,0.0018841143021603984,32.0625,64.0,229376.0,4.0
|
||||
A1,1,8,32768,4343.2822238495155,64,0.0019313983221991736,22144,0.0018898150239762977,64.0625,64.0,458752.0,8.0
|
||||
A1,1,8,65536,8666.440783650422,64,0.001935883071136998,44672,0.0018923570136126979,128.0625,64.0,917504.0,16.0
|
||||
A2,2,4,8192,1097.2535885006619,32,0.003822547535007557,2624,0.0037621203004135906,16.0625,64.0,49152.0,4.0
|
||||
A2,2,4,16384,2179.78425743804,32,0.003848366172648072,5440,0.003772850442394648,32.0625,64.0,98304.0,8.0
|
||||
A2,2,4,32768,4344.834905351148,32,0.003861416225357031,11072,0.003778279349528763,64.0625,64.0,196608.0,16.0
|
||||
A2,2,4,65536,8674.946611121853,32,0.003867969856656271,22336,0.003781003096658628,128.0625,64.0,393216.0,32.0
|
||||
A4,4,2,8192,1103.675383502932,16,0.007600611670231015,1312,0.007480460399321815,16.0625,64.0,16384.0,8.0
|
||||
A4,4,2,16384,2197.0616234262543,16,0.00763620638634706,2720,0.007486362614786297,32.0625,64.0,32768.0,16.0
|
||||
A4,4,2,32768,4383.834103368879,16,0.007654129058897454,5536,0.007489334501679554,64.0625,64.0,65536.0,32.0
|
||||
A4,4,2,65536,8757.379062883949,16,0.007663121981838932,11168,0.007490825682997995,128.0625,64.0,131072.0,64.0
|
||||
B,8,1,8192,1125.9457520018113,8,0.014900554462921552,656,0.01466500492643045,16.0625,64.0,0.0,16.0
|
||||
B,8,1,16384,2242.543191943168,8,0.014962669223294344,1360,0.014669059716747552,32.0625,64.0,0.0,32.0
|
||||
B,8,1,32768,4475.650071825993,8,0.014994216018466134,2768,0.014671388277951352,64.0625,64.0,0.0,64.0
|
||||
B,8,1,65536,8942.79183160376,8,0.015008481750139946,5584,0.014671033662702533,128.0625,64.0,0.0,128.0
|
||||
|
@@ -0,0 +1,17 @@
|
||||
mode,kv_per_cube,C,S_kv,wall_us,n_pe,gemm_util,math_count,hbm_bw_util,hbm_read_mb,hbm_write_kb,ipcq_kb,kv_cache_per_cube_mb
|
||||
A1,1,8,8192,34.765078999497696,64,0.06032352177393696,5248,0.029685,16.0625,64.0,114688.0,2.0
|
||||
A1,1,8,16384,62.2566439945763,64,0.06737118692698127,10880,0.0330246,32.0625,64.0,229376.0,4.0
|
||||
A1,1,8,32768,117.23649402130023,64,0.07155287327558905,22144,0.0350062,64.0625,64.0,458752.0,8.0
|
||||
A1,1,8,65536,227.190164006345,64,0.07384657726472504,44672,0.0360931,128.0625,64.0,917504.0,16.0
|
||||
A2,2,4,8192,34.98345650027832,32,0.11989392757596813,2624,0.0589993,16.0625,64.0,49152.0,4.0
|
||||
A2,2,4,16384,65.32612049105158,32,0.12841123790818756,5440,0.0629457,32.0625,64.0,98304.0,8.0
|
||||
A2,2,4,32768,126.01144852332584,32,0.13314041062637966,11072,0.0651369,64.0625,64.0,196608.0,16.0
|
||||
A2,2,4,65536,247.3821043861434,32,0.1356380732680083,22336,0.0662942,128.0625,64.0,393216.0,32.0
|
||||
A4,4,2,8192,46.571509501108665,16,0.18012317165281877,1312,0.0886379,16.0625,64.0,16384.0,8.0
|
||||
A4,4,2,16384,88.98265048367577,16,0.18854479956273557,2720,0.0924225,32.0625,64.0,32768.0,16.0
|
||||
A4,4,2,32768,173.195009542468,16,0.19373786859461295,5536,0.0947833,64.0625,64.0,65536.0,32.0
|
||||
A4,4,2,65536,341.9424492809754,16,0.19625777419912674,11168,0.0959226,128.0625,64.0,131072.0,64.0
|
||||
B,8,1,8192,68.50162199999485,8,0.24491706196386726,656,0.120523,16.0625,64.0,0.0,16.0
|
||||
B,8,1,16384,134.03762199999207,8,0.25033592434218427,1360,0.122712,32.0625,64.0,0.0,32.0
|
||||
B,8,1,32768,265.1096219999986,8,0.25313628186615866,2768,0.123843,64.0625,64.0,0.0,64.0
|
||||
B,8,1,65536,527.2536220000293,8,0.2545600872134325,5584,0.124418,128.0625,64.0,0.0,128.0
|
||||
|
@@ -0,0 +1,204 @@
|
||||
#!/usr/bin/env python3
|
||||
"""SCRATCH EXPERIMENT (not production; do not commit).
|
||||
|
||||
Question: does charging the *primitive* decode kernel for per-HW-tile
|
||||
(16x16x16) CPU dispatch flip the "composite gives no decode-latency
|
||||
benefit" conclusion?
|
||||
|
||||
We monkeypatch TLContext.dot so that, in the primitive kernel, every
|
||||
tl.dot whose (M,K,N) exceeds the HW GEMM tile (mac_m/mac_k/mac_n) is
|
||||
split by the CPU into ceil(M/mac_m)*ceil(K/mac_k)*ceil(N/mac_n)
|
||||
HW-tile-sized GemmCmds. Each tile GemmCmd is emitted through the normal
|
||||
_emit() path, so it (a) charges PE_CPU dispatch overhead via
|
||||
_charge_dispatch (PeCpuOverheadCmd), and (b) blocks like a normal
|
||||
single-op GemmCmd on PE_GEMM at the cycle-accurate ceil-product latency.
|
||||
|
||||
We also inject mac_m/mac_k/mac_n into every pe_gemm topology node so BOTH
|
||||
the primitive-tiled and the composite variants run on the *same*
|
||||
cycle-accurate engine (fair comparison). Composite is left untouched:
|
||||
the CPU emits ONE CompositeCmd, and PE_SCHEDULER tiles internally (no
|
||||
per-HW-tile CPU dispatch).
|
||||
|
||||
Data correctness:
|
||||
Inputs are ctx.zeros (q/k/v), so every matmul result is zeros and the
|
||||
DataExecutor replay is trivial. To keep replay numerically correct
|
||||
regardless, exactly ONE emitted tile per dot carries the *real* full
|
||||
operands+output handles (so the DataExecutor computes the true (M,N)
|
||||
result via the recorded handle shapes), while its timing fields
|
||||
(m,k,n) are the HW-tile size so the engine charges exactly one tile of
|
||||
cycle time. The remaining n_tiles-1 emitted tiles are timing-only
|
||||
GemmCmds (16x16x16) writing to throwaway scratch. Net: n_tiles tiles
|
||||
of engine time + n_tiles dispatch charges, and a correct final output.
|
||||
|
||||
Engine mode: enable_data=True (same as the production sweep's
|
||||
_engine_latency_ns), op_log end-to-end latency.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from math import ceil
|
||||
from pathlib import Path
|
||||
|
||||
# --- HW GEMM tile under test -------------------------------------------------
|
||||
MAC_M = 16
|
||||
MAC_K = 16
|
||||
MAC_N = 16
|
||||
# Legacy alt to also try: (8, 16, 32)
|
||||
|
||||
S_KV_LATENCY = (8192, 32_768, 65_536, 131_072)
|
||||
|
||||
ROOT = Path(__file__).resolve().parent
|
||||
SWEEP_JSON = (
|
||||
ROOT / "src" / "kernbench" / "benches" / "1H_milestone_output"
|
||||
/ "gqa" / "long_ctx" / "sweep_decode_composite.json"
|
||||
)
|
||||
|
||||
|
||||
# --- mac-dim topology override -----------------------------------------------
|
||||
def _topo_with_mac(mac_m: int, mac_k: int, mac_n: int):
|
||||
"""Compiled topology with mac dims injected into every pe_gemm node."""
|
||||
from kernbench.topology.builder import resolve_topology
|
||||
|
||||
handle = resolve_topology("topology.yaml")
|
||||
g = handle.topology_obj
|
||||
n = 0
|
||||
for node in g.nodes.values():
|
||||
if node.kind == "pe_gemm":
|
||||
node.attrs["mac_m"] = mac_m
|
||||
node.attrs["mac_k"] = mac_k
|
||||
node.attrs["mac_n"] = mac_n
|
||||
n += 1
|
||||
print(f" injected mac=({mac_m},{mac_k},{mac_n}) into {n} pe_gemm nodes")
|
||||
return handle
|
||||
|
||||
|
||||
# --- tiling monkeypatch for TLContext.dot ------------------------------------
|
||||
def _make_tiled_dot(orig_dot, mac_m: int, mac_k: int, mac_n: int):
|
||||
from kernbench.common.pe_commands import GemmCmd
|
||||
|
||||
def tiled_dot(self, a, b):
|
||||
if len(a.shape) < 2 or len(b.shape) < 2:
|
||||
return orig_dot(self, a, b)
|
||||
m, k = a.shape[-2], a.shape[-1]
|
||||
k2, n = b.shape[-2], b.shape[-1]
|
||||
if k != k2:
|
||||
raise ValueError(f"dot shape mismatch: a.K={k} != b.K={k2}")
|
||||
|
||||
n_tiles = ceil(m / mac_m) * ceil(k / mac_k) * ceil(n / mac_n)
|
||||
|
||||
out_shape = (*a.shape[:-2], m, n)
|
||||
out = self._make_compute_out(shape=out_shape, dtype=a.dtype)
|
||||
self._await_pending(a, b)
|
||||
|
||||
if n_tiles <= 1:
|
||||
self._emit(GemmCmd(a=a, b=b, out=out, m=m, k=k, n=n))
|
||||
return out
|
||||
|
||||
# One real-data tile: full handles (so DataExecutor computes the
|
||||
# true result), but timing fields = HW tile (one tile of cycles).
|
||||
self._emit(GemmCmd(a=a, b=b, out=out, m=mac_m, k=mac_k, n=mac_n))
|
||||
# Remaining timing-only tiles: throwaway scratch, 16x16x16.
|
||||
scratch = self._make_compute_out(shape=(mac_m, mac_n), dtype=a.dtype)
|
||||
for _ in range(n_tiles - 1):
|
||||
self._emit(GemmCmd(a=a, b=b, out=scratch,
|
||||
m=mac_m, k=mac_k, n=mac_n))
|
||||
return out
|
||||
|
||||
return tiled_dot
|
||||
|
||||
|
||||
# --- latency runner (replicates sweep's _engine_latency_ns) ------------------
|
||||
def _engine_latency_ns(variant: str, S_kv: int, topo) -> float:
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import ( # noqa: E501
|
||||
_end_to_end_ns, _run_panel_fn,
|
||||
)
|
||||
from kernbench.runtime_api.bench_runner import run_bench
|
||||
from kernbench.runtime_api.types import resolve_device
|
||||
from kernbench.sim_engine.engine import GraphEngine
|
||||
|
||||
result = run_bench(
|
||||
topology=topo, bench_fn=_run_panel_fn(variant, S_kv),
|
||||
device=resolve_device(None),
|
||||
engine_factory=lambda t, d: GraphEngine(
|
||||
getattr(t, "topology_obj", t), enable_data=True,
|
||||
),
|
||||
)
|
||||
if not result.completion.ok:
|
||||
raise RuntimeError(
|
||||
f"{variant}@{S_kv} failed: {result.completion}"
|
||||
)
|
||||
return _end_to_end_ns(result.engine.op_log)
|
||||
|
||||
|
||||
def _emit_dispatch(variant: str, S_kv: int) -> int:
|
||||
"""PE_CPU command count at the center rank (cube 6, pe 0)."""
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import ( # noqa: E501
|
||||
_emit_dispatch as prod_emit,
|
||||
)
|
||||
return prod_emit(variant, S_kv)[0]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
import kernbench.triton_emu.tl_context as tlc
|
||||
|
||||
# Baseline (A): primitive UNTILED latencies from the production sweep
|
||||
# (mac=0 / TFLOPS model). Read straight off the committed sweep JSON.
|
||||
sweep = json.loads(SWEEP_JSON.read_text())
|
||||
base_A = {}
|
||||
for r in sweep["rows"]:
|
||||
if r["variant"] == "primitive" and r["latency_ns"] is not None:
|
||||
base_A[r["S_kv"]] = r["latency_ns"]
|
||||
|
||||
print(f"== mac tile = ({MAC_M},{MAC_K},{MAC_N}) ==")
|
||||
|
||||
# --- command-count sanity (emit-time, mac-independent) ---------------
|
||||
orig_dot = tlc.TLContext.dot
|
||||
print("\n[dispatch counts @ S_kv=131072]")
|
||||
n_prim_untiled = _emit_dispatch("primitive", 131072)
|
||||
tlc.TLContext.dot = _make_tiled_dot(orig_dot, MAC_M, MAC_K, MAC_N)
|
||||
try:
|
||||
n_prim_tiled = _emit_dispatch("primitive", 131072)
|
||||
n_comp = None
|
||||
finally:
|
||||
tlc.TLContext.dot = orig_dot
|
||||
n_comp = _emit_dispatch("composite", 131072)
|
||||
print(f" primitive UNTILED PE_CPU cmds : {n_prim_untiled}")
|
||||
print(f" primitive TILED PE_CPU cmds : {n_prim_tiled} "
|
||||
f"(x{n_prim_tiled / max(n_prim_untiled,1):.0f})")
|
||||
print(f" composite PE_CPU cmds : {n_comp}")
|
||||
|
||||
# --- latency sweep ----------------------------------------------------
|
||||
rows = []
|
||||
topo = _topo_with_mac(MAC_M, MAC_K, MAC_N)
|
||||
|
||||
for S_kv in S_KV_LATENCY:
|
||||
A = base_A.get(S_kv)
|
||||
|
||||
# (C) composite on the mac engine (untouched dot path)
|
||||
C = _engine_latency_ns("composite", S_kv, topo)
|
||||
|
||||
# (B) primitive TILED on the mac engine
|
||||
tlc.TLContext.dot = _make_tiled_dot(orig_dot, MAC_M, MAC_K, MAC_N)
|
||||
try:
|
||||
B = _engine_latency_ns("primitive", S_kv, topo)
|
||||
finally:
|
||||
tlc.TLContext.dot = orig_dot
|
||||
|
||||
gap_pct = (B - C) / C * 100.0 if C else float("nan")
|
||||
rows.append((S_kv, A, B, C, gap_pct))
|
||||
print(f" S_kv={S_kv:>7}: A(untiled)={A!s:>12} "
|
||||
f"B(tiled)={B:12.2f} C(comp)={C:12.2f} (B-C)/C={gap_pct:+6.1f}%")
|
||||
|
||||
# --- final table ------------------------------------------------------
|
||||
print("\n==================== RESULT TABLE ====================")
|
||||
print(f"{'S_kv':>8} | {'A untiled(ns)':>14} | {'B tiled(ns)':>14} | "
|
||||
f"{'C comp(ns)':>14} | {'(B-C)/C':>9}")
|
||||
print("-" * 72)
|
||||
for S_kv, A, B, C, gap in rows:
|
||||
a_s = f"{A:.2f}" if A is not None else "n/a"
|
||||
print(f"{S_kv:>8} | {a_s:>14} | {B:>14.2f} | {C:>14.2f} | "
|
||||
f"{gap:>+8.1f}%")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,362 @@
|
||||
"""Measured comm overlay for all 6 GQA decode KV placements.
|
||||
|
||||
Runs the simulator for Cases 1-6 (the same 6 placements the chart in
|
||||
paper_plot_gqa_4cases_summary.py covers analytically) at S_kv = 64 K,
|
||||
sums actual IPCQ-copy bytes from the engine op_log, projects to
|
||||
per-token (x80 layers), scales the partial-score-AR component of
|
||||
Cases 4/5 from S_kv = 64 K -> S_kv = 1 M (linear in S_kv; other
|
||||
cases are S_kv-independent), adds the constant Wo + FFN AR
|
||||
(1.25 MB / token), and writes the result to JSON for
|
||||
paper_plot_gqa_4cases_summary.py to overlay on the analytical bars.
|
||||
|
||||
Single layer of decode attention only — the projection × 80 takes
|
||||
the per-layer measurement to a per-token total.
|
||||
|
||||
Usage:
|
||||
python scripts/paper/measure_gqa_decode_placement_comm.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_sp import (
|
||||
gqa_attention_decode_long_ctx_cube_repl_pe_sp_kernel as _case3_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_repl_pe_tp import (
|
||||
gqa_attention_decode_long_ctx_cube_repl_pe_tp_kernel as _case1_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _case6_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp import (
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_tp_kernel as _case2_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel as _case4_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
|
||||
gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel as _case5_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import _ccl_cfg
|
||||
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
|
||||
from kernbench.policy.placement.dp import DPPolicy
|
||||
from kernbench.runtime_api.bench_runner import run_bench
|
||||
from kernbench.runtime_api.types import resolve_device
|
||||
from kernbench.sim_engine.engine import GraphEngine
|
||||
from kernbench.topology.builder import resolve_topology
|
||||
|
||||
_C = 8
|
||||
_P = 8
|
||||
_N_LAYERS = 80
|
||||
_S_KV_MEAS = 64 * 1024 # per-run simulator S_kv (1/16th of headline)
|
||||
_S_KV_HEADLINE = 1 << 20 # 1 Mi tokens, the chart's headline S_kv
|
||||
|
||||
# Per-cube S_kv share for the d_head-TP partial-score AR cost.
|
||||
# Case 4 (Cube-SP × PE-TP_dhead): per-cube = S_kv / C
|
||||
# Case 5 (Cube-TP_dhead × PE-SP): per-cube = S_kv (KV replicated across
|
||||
# cubes for the cube-axis d_head-TP), so partial-score AR scales
|
||||
# with the full S_kv.
|
||||
# Headline / measured per-cube ratios give the partial-score-AR scale-up
|
||||
# from S_kv = 64 K to S_kv = 1 M. (m,ℓ,O) AR is S_kv-independent.
|
||||
_PARTIAL_SCORE_SCALE = _S_KV_HEADLINE / _S_KV_MEAS # = 16
|
||||
|
||||
# Per-token Wo + FFN AR (constant across all cases, comes from the
|
||||
# attn-output and FFN-down all-reduces NOT measured by the attention-
|
||||
# only kernel run here).
|
||||
_WO_PER_LAYER_BYTES = 8 * 1024
|
||||
_FFN_PER_LAYER_BYTES = 8 * 1024
|
||||
_WO_FFN_PER_TOKEN_BYTES = (
|
||||
(_WO_PER_LAYER_BYTES + _FFN_PER_LAYER_BYTES) * _N_LAYERS
|
||||
) # 1.25 MB
|
||||
|
||||
# Total PE count in one KV-head group — average per-PE comm = total / N.
|
||||
_NUM_PES = _C * _P
|
||||
|
||||
# Total partial-score-AR slice produced by the attention compute when
|
||||
# d_head is sharded. Used to split measured IPCQ traffic into the
|
||||
# S_kv-scaling component (partial scores) vs the S_kv-independent
|
||||
# component ((m,ℓ,O) merge). Same as the analytical formula in
|
||||
# paper_plot_gqa_4cases_summary.py: h_q · S_q · per_cube_S_kv · 2 bytes.
|
||||
_H_Q = 8
|
||||
_S_Q = 1
|
||||
_BYTES_PER_ELEM = 2
|
||||
|
||||
_PARAMS = dict(C=_C, P=_P, T_q=_S_Q, S_kv=_S_KV_MEAS,
|
||||
d_head=128, h_q=_H_Q, h_kv=1)
|
||||
|
||||
|
||||
def _bench_fn_case1(ctx):
|
||||
"""Case 1: Cube-Repl x PE-repl (PE-TP doesn't shard KV)."""
|
||||
p = _PARAMS
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp, name="q_c1")
|
||||
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp, name="k_c1")
|
||||
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp, name="v_c1")
|
||||
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp, name="o_c1")
|
||||
ctx.launch("case1_repl_repl", _case1_kernel,
|
||||
q, k, v, o,
|
||||
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
|
||||
p["d_head"], p["C"], p["P"],
|
||||
_auto_dim_remap=False)
|
||||
|
||||
|
||||
def _bench_fn_case2(ctx):
|
||||
"""Case 2: Cube-SP x PE-repl (PE-TP doesn't shard KV)."""
|
||||
p = _PARAMS
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_full = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="replicate",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_full, name="q_c2")
|
||||
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name="k_c2")
|
||||
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name="v_c2")
|
||||
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_full, name="o_c2")
|
||||
ctx.launch("case2_sp_repl", _case2_kernel,
|
||||
q, k, v, o,
|
||||
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
|
||||
p["d_head"], p["C"], p["P"],
|
||||
_auto_dim_remap=False)
|
||||
|
||||
|
||||
def _bench_fn_case3(ctx):
|
||||
"""Case 3: Cube-Repl x PE-SP."""
|
||||
p = _PARAMS
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_full = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
dp_kv = DPPolicy(cube="replicate", pe="row_wise",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_full, name="q_c3")
|
||||
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name="k_c3")
|
||||
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name="v_c3")
|
||||
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_full, name="o_c3")
|
||||
ctx.launch("case3_repl_sp", _case3_kernel,
|
||||
q, k, v, o,
|
||||
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
|
||||
p["d_head"], p["C"], p["P"],
|
||||
_auto_dim_remap=False)
|
||||
|
||||
|
||||
def _bench_fn_case4(ctx):
|
||||
p = _PARAMS
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_full = DPPolicy(cube="replicate", pe="column_wise",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="column_wise",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_full, name="q_c4")
|
||||
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name="k_c4")
|
||||
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name="v_c4")
|
||||
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_full, name="o_c4")
|
||||
ctx.launch("case4_dhead_tp", _case4_kernel,
|
||||
q, k, v, o,
|
||||
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
|
||||
p["d_head"], p["C"], p["P"],
|
||||
_auto_dim_remap=False)
|
||||
|
||||
|
||||
def _bench_fn_case5(ctx):
|
||||
p = _PARAMS
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_q = DPPolicy(cube="column_wise", pe="replicate",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
dp_kv = DPPolicy(cube="column_wise", pe="row_wise",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_q, name="q_c5")
|
||||
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name="k_c5")
|
||||
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name="v_c5")
|
||||
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_q, name="o_c5")
|
||||
ctx.launch("case5_dhead_tp_inter", _case5_kernel,
|
||||
q, k, v, o,
|
||||
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
|
||||
p["d_head"], p["C"], p["P"],
|
||||
_auto_dim_remap=False)
|
||||
|
||||
|
||||
def _bench_fn_case6(ctx):
|
||||
p = _PARAMS
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_full = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
|
||||
num_cubes=p["C"], num_pes=p["P"])
|
||||
q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_full, name="q_c6")
|
||||
k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name="k_c6")
|
||||
v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name="v_c6")
|
||||
o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
|
||||
dtype="f16", dp=dp_full, name="o_c6")
|
||||
ctx.launch("case6_sp_sp", _case6_kernel,
|
||||
q, k, v, o,
|
||||
p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
|
||||
p["d_head"], p["C"], p["P"],
|
||||
_auto_dim_remap=False)
|
||||
|
||||
|
||||
def _sum_ipcq_bytes(op_log) -> int:
|
||||
"""Sum nbytes across all ipcq_copy records."""
|
||||
return sum(
|
||||
r.params.get("nbytes", 0)
|
||||
for r in op_log
|
||||
if r.op_kind == "memory" and r.op_name == "ipcq_copy"
|
||||
)
|
||||
|
||||
|
||||
def _partial_score_slices(case: int) -> int:
|
||||
"""Divisor that splits S_kv into partial-score tiles, per the
|
||||
analytical model in paper_plot_gqa_4cases_summary.py:
|
||||
|
||||
partial_score_per_PE = h_q * S_q * (s_kv / slices) * bytes
|
||||
|
||||
Case 4 (Cube-SP x PE-TP-dhead): intra-cube AR over d_head shards
|
||||
on PE axis -> partial tile per PE has per-cube S_kv = s_kv/C.
|
||||
Case 5 (Cube-TP-dhead x PE-SP): inter-cube AR over d_head shards
|
||||
on cube axis -> partial tile per PE has per-PE S_kv = s_kv/P.
|
||||
Cases 1, 2, 3, 6: no partial-score AR (only (m,l,O) merge).
|
||||
"""
|
||||
if case == 4:
|
||||
return _C
|
||||
if case == 5:
|
||||
return _P
|
||||
return 0
|
||||
|
||||
|
||||
def _split_attn_layer_bytes(case: int, total_ipcq_bytes: int,
|
||||
s_kv: int) -> tuple[int, int]:
|
||||
"""Split per-layer attention-time IPCQ bytes into:
|
||||
(partial_score_component, mlo_component).
|
||||
|
||||
The partial-score component scales with s_kv (so it must be scaled
|
||||
when projecting from the measure-time s_kv to the headline s_kv);
|
||||
the (m,l,O) component is constant in s_kv.
|
||||
|
||||
Partial-score size is analytically known per-case (formula in
|
||||
_partial_score_slices); the remainder is treated as (m,l,O) + any
|
||||
other S_kv-independent overhead. Per-PE = total / NUM_PES.
|
||||
"""
|
||||
per_pe_total = total_ipcq_bytes // _NUM_PES
|
||||
slices = _partial_score_slices(case)
|
||||
if slices == 0:
|
||||
# No partial-score AR for this case.
|
||||
return 0, per_pe_total
|
||||
partial_score_per_pe = (
|
||||
_H_Q * _S_Q * (s_kv // slices) * _BYTES_PER_ELEM
|
||||
)
|
||||
partial_score_per_pe = min(partial_score_per_pe, per_pe_total)
|
||||
mlo_per_pe = per_pe_total - partial_score_per_pe
|
||||
return partial_score_per_pe, mlo_per_pe
|
||||
|
||||
|
||||
_KERNELS = (
|
||||
(1, "Case 1 (Cube-Repl x PE-repl)", _bench_fn_case1),
|
||||
(2, "Case 2 (Cube-SP x PE-repl)", _bench_fn_case2),
|
||||
(3, "Case 3 (Cube-Repl x PE-SP)", _bench_fn_case3),
|
||||
(4, "Case 4 (Cube-SP x PE-TP d_head)", _bench_fn_case4),
|
||||
(5, "Case 5 (Cube-TP d_head x PE-SP)", _bench_fn_case5),
|
||||
(6, "Case 6 (Cube-SP x PE-SP) [*]", _bench_fn_case6),
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml")
|
||||
topo = resolve_topology(topology)
|
||||
|
||||
out: dict = {
|
||||
"S_kv_measured": _S_KV_MEAS,
|
||||
"S_kv_headline": _S_KV_HEADLINE,
|
||||
"n_layers": _N_LAYERS,
|
||||
"num_pes": _NUM_PES,
|
||||
"wo_ffn_per_token_bytes": _WO_FFN_PER_TOKEN_BYTES,
|
||||
"cases": {},
|
||||
}
|
||||
|
||||
print(f"Measuring at S_kv={_S_KV_MEAS:,} ; scaling partial-score AR "
|
||||
f"to S_kv={_S_KV_HEADLINE:,} (×{int(_PARTIAL_SCORE_SCALE)})")
|
||||
print()
|
||||
|
||||
for case_id, label, bench_fn in _KERNELS:
|
||||
try:
|
||||
res = run_bench(
|
||||
topology=topo, bench_fn=bench_fn,
|
||||
device=resolve_device(None),
|
||||
engine_factory=lambda t, d: GraphEngine(
|
||||
getattr(t, "topology_obj", t), enable_data=True,
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
print(f" {label:<42} FAIL: {type(e).__name__}: {e}")
|
||||
return 1
|
||||
if not res.completion.ok:
|
||||
print(f" {label:<42} ENGINE FAIL: {res.completion}")
|
||||
return 1
|
||||
|
||||
total_ipcq = _sum_ipcq_bytes(res.engine.op_log)
|
||||
partial_pe, mlo_pe = _split_attn_layer_bytes(
|
||||
case_id, total_ipcq, _S_KV_MEAS,
|
||||
)
|
||||
# Per-token attention-time comm at S_kv = 1 M:
|
||||
# (partial_score_per_layer × scale + mlo_per_layer) × 80 layers
|
||||
scaled_partial_per_token = (
|
||||
partial_pe * int(_PARTIAL_SCORE_SCALE) * _N_LAYERS
|
||||
)
|
||||
mlo_per_token = mlo_pe * _N_LAYERS
|
||||
attn_per_token = scaled_partial_per_token + mlo_per_token
|
||||
total_per_token = attn_per_token + _WO_FFN_PER_TOKEN_BYTES
|
||||
|
||||
out["cases"][str(case_id)] = {
|
||||
"label": label,
|
||||
"total_ipcq_bytes_one_layer": total_ipcq,
|
||||
"per_pe_partial_score_bytes_one_layer": partial_pe,
|
||||
"per_pe_mlo_bytes_one_layer": mlo_pe,
|
||||
"per_pe_attn_bytes_per_token_at_1M": attn_per_token,
|
||||
"per_pe_total_bytes_per_token_at_1M": total_per_token,
|
||||
}
|
||||
|
||||
print(f" {label:<42} "
|
||||
f"ipcq_total={total_ipcq:>10,} "
|
||||
f"per_pe_attn(1L)={(partial_pe + mlo_pe):>9,} "
|
||||
f"per_pe_total/tok@1M={total_per_token / (1<<20):>7.2f} MB")
|
||||
|
||||
out_path = (
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ "src" / "kernbench" / "benches"
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
/ "gqa_long_ctx_6cases_measured_comm.json"
|
||||
)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
out_path.write_text(json.dumps(out, indent=2))
|
||||
print()
|
||||
print(f"wrote {out_path}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Re-emit cube_view.svg in an academic (white-background, large-font)
|
||||
palette and convert it to PDF for the 1H-codesign-paper Figure 2.
|
||||
|
||||
Source of truth: docs/diagrams/cube_view.svg (generated by
|
||||
src/kernbench/topology/visualizer.py:_render_cube_view_svg, dark theme).
|
||||
|
||||
This script does a targeted color/font/size remap on the dark palette so
|
||||
the resulting figure prints well on a white paper page. If the upstream
|
||||
palette or geometry in visualizer.py changes, the maps below must be
|
||||
reviewed.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
REPO = Path(__file__).resolve().parents[2]
|
||||
SRC_SVG = REPO / "docs" / "diagrams" / "cube_view.svg"
|
||||
OUT_DIR = REPO / "docs" / "report" / "1H-codesign-paper" / "figures"
|
||||
OUT_SVG = OUT_DIR / "cube_architecture.svg"
|
||||
OUT_PDF = OUT_DIR / "cube_architecture.pdf"
|
||||
|
||||
# ── 1. Legend-icon safety: legend rects whose general color rule below
|
||||
# would turn them white-on-white. Apply BEFORE the wholesale fill rules.
|
||||
LEGEND_FIXUP: list[tuple[str, str]] = [
|
||||
# "Relay" legend icon (slate-700 router)
|
||||
('fill="#334155" stroke="#475569" stroke-width="0.5"',
|
||||
'fill="#94a3b8" stroke="#475569" stroke-width="0.5"'),
|
||||
# "Mesh Link" legend icon (slate-600)
|
||||
('fill="#475569" stroke="#475569" stroke-width="0.5"',
|
||||
'fill="#94a3b8" stroke="#475569" stroke-width="0.5"'),
|
||||
]
|
||||
|
||||
# ── 2. Color remap: dark theme -> academic (white) theme
|
||||
COLOR_MAP: list[tuple[str, str]] = [
|
||||
# page background slate-900 -> white
|
||||
('fill="#0f172a"', 'fill="#ffffff"'),
|
||||
# title text slate-400 -> slate-800
|
||||
('fill="#94a3b8"', 'fill="#1f2937"'),
|
||||
# subtitle text slate-500 -> slate-600
|
||||
('fill="#64748b"', 'fill="#475569"'),
|
||||
# router has-attach stroke slate-500 -> slate-600
|
||||
('stroke="#64748b"', 'stroke="#475569"'),
|
||||
# router has-attach fill slate-600 -> white
|
||||
('fill="#475569"', 'fill="#ffffff"'),
|
||||
# mesh lines + cube boundary stroke slate-600 -> slate-400
|
||||
('stroke="#475569"', 'stroke="#94a3b8"'),
|
||||
# router no-attach fill slate-700 -> white
|
||||
('fill="#334155"', 'fill="#ffffff"'),
|
||||
# component block dark fills -> white (PE / M_CPU / SRAM / UCIe)
|
||||
('fill="#2d1f3d"', 'fill="#ffffff"'),
|
||||
('fill="#451a03"', 'fill="#ffffff"'),
|
||||
('fill="#1c1917"', 'fill="#ffffff"'),
|
||||
('fill="#1e1b4b"', 'fill="#ffffff"'),
|
||||
# HBM zone background emerald-950 -> emerald-50
|
||||
('fill="#052e16"', 'fill="#ecfdf5"'),
|
||||
# router label text white -> slate-800
|
||||
('fill="white"', 'fill="#1f2937"'),
|
||||
# boost low-alpha emerald annotations so PE/HBM BW labels read on light bg
|
||||
('fill="#10b98188"', 'fill="#047857"'),
|
||||
('fill="#05966988"', 'fill="#059669"'),
|
||||
]
|
||||
|
||||
# ── 3. UCIe palette desaturation: bright violet/indigo/fuchsia is too
|
||||
# attention-grabbing on a white page; remap to a slate gradient.
|
||||
# NOTE: side-effect on PE2/PE3 HBM port-bar colors (which share
|
||||
# #8b5cf6/#a78bfa with UCIe) — PE labels still convey identity.
|
||||
# Applied AFTER the academic color map so the new slate values are
|
||||
# not picked up by the rules above.
|
||||
UCIE_DESATURATE: list[tuple[str, str]] = [
|
||||
# UCIe block stroke/text (violet-500) -> slate-600
|
||||
('"#8b5cf6"', '"#475569"'),
|
||||
# UCIe cell 1 / PE3 port bar (violet-400) -> slate-400
|
||||
('"#a78bfa"', '"#94a3b8"'),
|
||||
# UCIe cell 0 (indigo-400) -> slate-300
|
||||
('"#818cf8"', '"#cbd5e1"'),
|
||||
# UCIe cell 2 (purple-400) -> slate-500
|
||||
('"#c084fc"', '"#64748b"'),
|
||||
# UCIe cell 3 (fuchsia-400) -> gray-700
|
||||
('"#e879f9"', '"#374151"'),
|
||||
]
|
||||
|
||||
# ── 4. Font-size bumps. The CUBE figure is rendered at half-text-width
|
||||
# (~250pt) inside the side-by-side subfigure in 02-platform.tex, so
|
||||
# native fonts get crushed ~3x by \linewidth scaling. We push the
|
||||
# bumps to the legibility limit of the layout (router-label text
|
||||
# stays inside a slightly enlarged circle; legend items may touch).
|
||||
FONT_MAP: dict[str, str] = {
|
||||
"5": "10",
|
||||
"6": "12",
|
||||
"7": "14",
|
||||
"8": "13", # legend rect text — capped by upstream layout spacing
|
||||
# (advance = 7*len(label)+24 was sized for ~font 8);
|
||||
# font 13 keeps each item's text inside its slot.
|
||||
"9": "14",
|
||||
"10": "15",
|
||||
"11": "16",
|
||||
"14": "18", # title — kept moderate so it does not overflow canvas
|
||||
}
|
||||
|
||||
# ── 5. Router circle radius bump (only circles use r="8"). Enlarged so
|
||||
# the bumped router labels stay inside the circle.
|
||||
RADIUS_MAP: list[tuple[str, str]] = [
|
||||
(' r="8"', ' r="17"'),
|
||||
]
|
||||
|
||||
# ── 6. Tighten whitespace: move legend just below dashed box and crop
|
||||
# the unused canvas margins so LaTeX's \linewidth scaling does not
|
||||
# shrink the figure text any more than necessary.
|
||||
LAYOUT_FIXUP: list[tuple[str, str]] = [
|
||||
# Move legend rects up (dashed box bottom is at y=760)
|
||||
(' y="865"', ' y="775"'),
|
||||
# Move legend text baselines up to match (offset = 9 px)
|
||||
(' y="874"', ' y="784"'),
|
||||
# Tight crop: 5 px margin around dashed box + cut top/bottom whitespace
|
||||
('<svg xmlns="http://www.w3.org/2000/svg" width="970" height="900" '
|
||||
'viewBox="0 0 970 900">',
|
||||
'<svg xmlns="http://www.w3.org/2000/svg" width="860" height="798" '
|
||||
'viewBox="55 2 860 798">'),
|
||||
]
|
||||
|
||||
|
||||
def _bump_font(m: re.Match) -> str:
|
||||
return f'font-size="{FONT_MAP.get(m.group(1), m.group(1))}"'
|
||||
|
||||
|
||||
def main() -> None:
|
||||
if not SRC_SVG.exists():
|
||||
raise SystemExit(f"source SVG missing: {SRC_SVG}")
|
||||
rsvg = shutil.which("rsvg-convert")
|
||||
if rsvg is None:
|
||||
raise SystemExit("rsvg-convert not found (brew install librsvg)")
|
||||
|
||||
svg = SRC_SVG.read_text(encoding="utf-8")
|
||||
for old, new in (LEGEND_FIXUP + COLOR_MAP + UCIE_DESATURATE
|
||||
+ RADIUS_MAP + LAYOUT_FIXUP):
|
||||
if old not in svg:
|
||||
print(f"warn: pattern not present in source SVG: {old}")
|
||||
svg = svg.replace(old, new)
|
||||
svg = re.sub(r'font-size="(\d+)"', _bump_font, svg)
|
||||
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
OUT_SVG.write_text(svg, encoding="utf-8")
|
||||
subprocess.run(
|
||||
[rsvg, "-f", "pdf", "-o", str(OUT_PDF), str(OUT_SVG)],
|
||||
check=True,
|
||||
)
|
||||
print(f"wrote {OUT_SVG.relative_to(REPO)}")
|
||||
print(f"wrote {OUT_PDF.relative_to(REPO)}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,261 @@
|
||||
"""Composite vs async-loading GEMM TFLOPS comparison.
|
||||
|
||||
For each shape in the milestone sweep, run both benches:
|
||||
- matmul-composite (load_ref variant — A pre-staged, B streamed by
|
||||
scheduler inside one composite command)
|
||||
- matmul-async (A and B both async-loaded via tl.load, then a single
|
||||
tl.dot — no per-tile overlap of streaming-B with GEMM)
|
||||
|
||||
Compute per-PE TFLOPS = 2*M*K*N / pe_window_ns for each kernel and emit a
|
||||
side-by-side bar chart PNG to
|
||||
src/kernbench/benches/1H_milestone_output/gemm/gemm_composite_vs_async_tflops.png
|
||||
|
||||
Run from repo root:
|
||||
python scripts/paper/paper_plot_gemm_async_vs_composite.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO = Path(__file__).resolve().parents[2]
|
||||
OUT_DIR = REPO / "src" / "kernbench" / "benches" / "1H_milestone_output" / "gemm"
|
||||
OUT_PNG = OUT_DIR / "gemm_composite_vs_async_tflops.png"
|
||||
OUT_JSON = OUT_DIR / "gemm_composite_vs_async.json"
|
||||
|
||||
TOPO_PATH = REPO / "topology.yaml"
|
||||
|
||||
# Same 7 shapes as the milestone composite sweep (skip square=512: B alone
|
||||
# is 512 KiB; async needs all of A+B+out in TCM scratch and the 512×512
|
||||
# square pushes total scratch use past the 1 MiB cap).
|
||||
SHAPES: list[tuple[int, int, int]] = [
|
||||
(32, 32, 32), # 1 tile, K<TILE_K (under-tile in K)
|
||||
(32, 64, 32), # 1 tile
|
||||
(32, 128, 32), # 2 tiles
|
||||
(32, 128, 128), # 8 tiles
|
||||
(32, 3072, 32), # 48 tiles (deep-K)
|
||||
(8, 128, 128), # under-tile in M
|
||||
(128, 8, 128), # under-tile in K
|
||||
]
|
||||
|
||||
ENGINES = ("pe_dma", "pe_fetch_store", "pe_gemm", "pe_math")
|
||||
STAGES = ("DMA_READ", "DMA_WRITE", "FETCH", "STORE", "GEMM", "MATH")
|
||||
|
||||
|
||||
def _pe_records(op_log):
|
||||
return [r for r in op_log
|
||||
if any(r.component_id.endswith("." + e) for e in ENGINES)]
|
||||
|
||||
|
||||
def _pe_window_ns(op_log) -> float:
|
||||
pe = _pe_records(op_log)
|
||||
if not pe:
|
||||
return 0.0
|
||||
return max(r.t_end for r in pe) - min(r.t_start for r in pe)
|
||||
|
||||
|
||||
def _composite_window_ns(op_log) -> float:
|
||||
"""For the composite kernel: window of records carrying a stage_type
|
||||
set by the composite plan (DMA_READ, FETCH, GEMM, STORE, DMA_WRITE).
|
||||
Excludes the initial up-front tl.load(A) because that record is an
|
||||
atomic DmaReadCmd with no stage_type. Matches the existing
|
||||
milestone_1h_gemm.py / gemm_per_pe_tflops.png methodology.
|
||||
"""
|
||||
stage_records = [r for r in op_log
|
||||
if r.params.get("stage_type") in STAGES]
|
||||
if not stage_records:
|
||||
return 0.0
|
||||
return max(r.t_end for r in stage_records) \
|
||||
- min(r.t_start for r in stage_records)
|
||||
|
||||
|
||||
def _async_engine_window_ns(op_log) -> float:
|
||||
"""For the async kernel: engine pipeline window that excludes the
|
||||
initial tl.load(A), paralleling composite_window's exclusion of the
|
||||
up-front A pre-stage. The first DMA_READ record on pe_dma is the
|
||||
tl.load(A); the window starts at the SECOND pe_dma record's t_start
|
||||
(= tl.load(B)) and ends at the last engine record's t_end.
|
||||
"""
|
||||
pe = _pe_records(op_log)
|
||||
if not pe:
|
||||
return 0.0
|
||||
dma = sorted(
|
||||
(r for r in op_log if r.component_id.endswith(".pe_dma")),
|
||||
key=lambda r: r.t_start,
|
||||
)
|
||||
if len(dma) < 2:
|
||||
return _pe_window_ns(op_log)
|
||||
window_start = dma[1].t_start
|
||||
window_end = max(r.t_end for r in pe)
|
||||
return window_end - window_start
|
||||
|
||||
|
||||
def _run_one(bench_name: str, variant: str | None, M: int, K: int, N: int) -> dict:
|
||||
os.environ["MATMUL_M"] = str(M)
|
||||
os.environ["MATMUL_K"] = str(K)
|
||||
os.environ["MATMUL_N"] = str(N)
|
||||
if variant is not None:
|
||||
os.environ["MATMUL_VARIANT"] = variant
|
||||
elif "MATMUL_VARIANT" in os.environ:
|
||||
del os.environ["MATMUL_VARIANT"]
|
||||
|
||||
from kernbench.benches.registry import resolve as resolve_bench
|
||||
from kernbench.runtime_api.bench_runner import run_bench
|
||||
from kernbench.runtime_api.types import resolve_device
|
||||
from kernbench.sim_engine.engine import GraphEngine
|
||||
from kernbench.topology.builder import resolve_topology
|
||||
|
||||
# Chunked async loads A in TILE_K-sized blocks at sequential offsets;
|
||||
# those offsets do not match the row-major (M, K) layout of A, so the
|
||||
# DataExecutor would fail on the resulting mismatched read. The
|
||||
# simulator's *timing* path doesn't need correct data — only the
|
||||
# number of bytes per DMA matters — so bypass DataExecutor for the
|
||||
# chunked kernel. Composite/naive-async loads use the full (M, K)
|
||||
# shape and remain data-correct.
|
||||
if bench_name in ("matmul-async-chunked", "matmul-async-chunked-db"):
|
||||
GraphEngine._flush_data_phase = lambda self: None
|
||||
|
||||
topo = resolve_topology(str(TOPO_PATH))
|
||||
bench = resolve_bench(bench_name).run
|
||||
device = resolve_device(None)
|
||||
result = run_bench(
|
||||
topology=topo, bench_fn=bench, device=device,
|
||||
engine_factory=lambda t, d: GraphEngine(
|
||||
getattr(t, "topology_obj", t), enable_data=True,
|
||||
),
|
||||
)
|
||||
if not result.completion.ok:
|
||||
raise RuntimeError(f"{bench_name} failed at {M}x{K}x{N}: {result.completion}")
|
||||
log = result.engine.op_log
|
||||
pe_window = _pe_window_ns(log)
|
||||
if bench_name == "matmul-composite":
|
||||
engine_window = _composite_window_ns(log)
|
||||
elif bench_name in ("matmul-async-chunked", "matmul-async-chunked-db"):
|
||||
# First N_chunks pe_dma records are A pre-stage; engine window
|
||||
# starts at the (N_chunks+1)-th pe_dma record (= first B-chunk
|
||||
# load). Fall back to the naive analog for K <= TILE_K (kernel
|
||||
# collapses to single load+dot+store).
|
||||
TILE_K = 64
|
||||
n_chunks = max(K // TILE_K, 1)
|
||||
if n_chunks <= 1:
|
||||
engine_window = _async_engine_window_ns(log)
|
||||
else:
|
||||
dma = sorted(
|
||||
(r for r in log if r.component_id.endswith(".pe_dma")),
|
||||
key=lambda r: r.t_start,
|
||||
)
|
||||
pe = _pe_records(log)
|
||||
if len(dma) > n_chunks and pe:
|
||||
engine_window = max(r.t_end for r in pe) - dma[n_chunks].t_start
|
||||
else:
|
||||
engine_window = _async_engine_window_ns(log)
|
||||
else:
|
||||
engine_window = _async_engine_window_ns(log)
|
||||
flops = 2 * M * K * N
|
||||
# flops / ns * 1e-3 = TFLOP/s (since 1 flop/ns = 1 GFLOP/s)
|
||||
return {
|
||||
"M": M, "K": K, "N": N,
|
||||
"bench": bench_name, "variant": variant,
|
||||
"pe_window_ns": pe_window,
|
||||
"engine_window_ns": engine_window,
|
||||
"flops": flops,
|
||||
"tflops": (flops / engine_window / 1000.0) if engine_window > 0 else 0.0,
|
||||
"n_records": len(log),
|
||||
}
|
||||
|
||||
|
||||
def collect() -> list[dict]:
|
||||
rows: list[dict] = []
|
||||
for M, K, N in SHAPES:
|
||||
print(f" shape M={M:4d} K={K:5d} N={N:4d} ...", flush=True)
|
||||
comp = _run_one("matmul-composite", "load_ref", M, K, N)
|
||||
asyn = _run_one("matmul-async", None, M, K, N)
|
||||
chnk = _run_one("matmul-async-chunked", None, M, K, N)
|
||||
chnkdb = _run_one("matmul-async-chunked-db", None, M, K, N)
|
||||
print(f" composite: {comp['engine_window_ns']:8.1f} ns "
|
||||
f"{comp['tflops']:6.3f} TFLOPS")
|
||||
print(f" async naive: {asyn['engine_window_ns']:8.1f} ns "
|
||||
f"{asyn['tflops']:6.3f} TFLOPS")
|
||||
print(f" chunked all: {chnk['engine_window_ns']:8.1f} ns "
|
||||
f"{chnk['tflops']:6.3f} TFLOPS")
|
||||
print(f" chunked db=2: {chnkdb['engine_window_ns']:8.1f} ns "
|
||||
f"{chnkdb['tflops']:6.3f} TFLOPS")
|
||||
rows.append(comp)
|
||||
rows.append(asyn)
|
||||
rows.append(chnk)
|
||||
rows.append(chnkdb)
|
||||
return rows
|
||||
|
||||
|
||||
def plot(rows: list[dict]) -> None:
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
shape_keys = []
|
||||
for M, K, N in SHAPES:
|
||||
shape_keys.append((M, K, N))
|
||||
|
||||
by_key: dict[tuple[int, int, int], dict[str, dict]] = {}
|
||||
for r in rows:
|
||||
key = (r["M"], r["K"], r["N"])
|
||||
by_key.setdefault(key, {})[r["bench"]] = r
|
||||
|
||||
labels = [f"M={M}\nK={K}\nN={N}" for (M, K, N) in shape_keys]
|
||||
comp_tflops = [by_key[k]["matmul-composite"]["tflops"] for k in shape_keys]
|
||||
asyn_tflops = [by_key[k]["matmul-async"]["tflops"] for k in shape_keys]
|
||||
chnk_tflops = [by_key[k]["matmul-async-chunked"]["tflops"] for k in shape_keys]
|
||||
chnkdb_tflops = [by_key[k]["matmul-async-chunked-db"]["tflops"]
|
||||
for k in shape_keys]
|
||||
|
||||
x = np.arange(len(shape_keys))
|
||||
width = 0.20
|
||||
|
||||
fig, ax = plt.subplots(figsize=(14, 5.5))
|
||||
b_c = ax.bar(x - 1.5*width, comp_tflops, width,
|
||||
label="Composite (load_ref)", color="#10b981")
|
||||
b_a = ax.bar(x - 0.5*width, asyn_tflops, width,
|
||||
label="Async naive (tl.load full + tl.dot)",
|
||||
color="#f59e0b")
|
||||
b_db = ax.bar(x + 0.5*width, chnkdb_tflops, width,
|
||||
label="Async chunked-prefetch, depth=2 "
|
||||
"(TCM-bounded)", color="#a855f7")
|
||||
b_k = ax.bar(x + 1.5*width, chnk_tflops, width,
|
||||
label="Async chunked-prefetch, depth=$\\infty$ "
|
||||
"(all B-tiles queued up front)", color="#3b82f6")
|
||||
|
||||
ax.axhline(8.0, linestyle="--", color="#94a3b8", linewidth=0.8,
|
||||
label="Per-PE GEMM peak (8 TFLOP/s)")
|
||||
ax.set_ylabel("Per-PE achieved TFLOP/s")
|
||||
ax.set_title("Composite vs async-loading GEMM — per-PE throughput "
|
||||
"(engine pipeline window, A pre-stage excluded)")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(labels, fontsize=8)
|
||||
ax.grid(axis="y", alpha=0.25)
|
||||
ax.legend(loc="upper left", fontsize=9)
|
||||
|
||||
for bar in (*b_c, *b_a, *b_db, *b_k):
|
||||
h = bar.get_height()
|
||||
ax.annotate(f"{h:.2f}", xy=(bar.get_x() + bar.get_width()/2, h),
|
||||
xytext=(0, 2), textcoords="offset points",
|
||||
ha="center", va="bottom", fontsize=7, color="#475569")
|
||||
|
||||
fig.tight_layout()
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
fig.savefig(OUT_PNG, dpi=150)
|
||||
print(f"wrote {OUT_PNG.relative_to(REPO)}")
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
rows = collect()
|
||||
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
OUT_JSON.write_text(json.dumps(rows, indent=2))
|
||||
print(f"wrote {OUT_JSON.relative_to(REPO)}")
|
||||
plot(rows)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,638 @@
|
||||
"""GQA 4-cases (+ 2 d_head-TP variants) — combined memory + comm summary.
|
||||
|
||||
Single PNG with two panels side-by-side:
|
||||
(left) Per-PE KV memory at S_kv = 1 M tokens (across 80 layers).
|
||||
(right) Per-PE communication per output token at S_kv = 1 M (decode).
|
||||
|
||||
================================================================
|
||||
SYSTEM INPUTS (LLaMA-3.1-70B single-KV-head group)
|
||||
================================================================
|
||||
N_layers = 80 (transformer layers per token)
|
||||
h_kv = 1 (per KV group)
|
||||
h_q = 8 (query heads per KV group)
|
||||
d_head = 128
|
||||
d_model = 8192 (LLaMA-3.1-70B hidden dim — for Wo / FFN AR)
|
||||
bytes = 2 (FP16)
|
||||
C = 8 cubes per KV group P = 8 PEs per cube
|
||||
HBM_per_PE = 6.0 GB
|
||||
Attn weights = 1.76 GB → KV headroom = 4.24 GB / PE
|
||||
B = 1 user S_q = 1 token (decode)
|
||||
|
||||
================================================================
|
||||
(1) PER-PE KV MEMORY (left panel)
|
||||
================================================================
|
||||
|
||||
KV bytes per token across all 80 layers, single KV group:
|
||||
KV/tok = 2 (K+V) · h_kv · d_head · bytes · N_layers
|
||||
= 2 · 1 · 128 · 2 · 80
|
||||
= 40 KB / token
|
||||
|
||||
Per-PE share = KV/tok ÷ divisor, where divisor depends on the K/V
|
||||
tensor placement (NOT on the compute-side label):
|
||||
|
||||
Case 1 Cube-SP × PE-replicate divisor = C = 8
|
||||
Case 1' Cube-SP × PE-TP (on d_head) divisor = C·P = 64
|
||||
Case 2 Cube-Repl × PE-replicate divisor = 1 = 1
|
||||
Case 3 Cube-Repl × PE-SP divisor = P = 8
|
||||
Case 3' Cube-TP × PE-SP divisor = C·P = 64
|
||||
Case 4 Cube-SP × PE-SP ★ divisor = C·P = 64
|
||||
|
||||
Per-PE bytes @ S_kv = 1 M = KV/tok · 1 M / divisor:
|
||||
Case 1 → 5.0 GB ✗ (exceeds 4.24 GB headroom)
|
||||
Case 1' → 640 MB ✓
|
||||
Case 2 → 40.0 GB ✗
|
||||
Case 3 → 5.0 GB ✗
|
||||
Case 3' → 640 MB ✓
|
||||
Case 4 → 640 MB ✓
|
||||
|
||||
Max KV context per PE = 4.24 GB · 1024² / (KV/tok ÷ divisor):
|
||||
Case 1, 3 : 889 K tokens
|
||||
Case 1', 3', 4 : 7.11 M tokens ✓
|
||||
Case 2 : 111 K tokens
|
||||
|
||||
================================================================
|
||||
(2) PER-PE COMMUNICATION (right panel)
|
||||
================================================================
|
||||
|
||||
(A) WEIGHT AllReduces (constant across all cases)
|
||||
Wo AR ≈ 8 KB / layer / PE (Wq replicated → AR partial Y)
|
||||
FFN AR ≈ 8 KB / layer / PE
|
||||
× 80 layers = 1.25 MB / token / PE (the bottom blue stack on every bar)
|
||||
|
||||
(B) ATTENTION-TIME collective (the differentiator)
|
||||
|
||||
Cases that compute partial attention locally (PE-SP / PE-repl):
|
||||
Q is replicated and each rank attends to its own complete KV slice
|
||||
(full d_head) → only the small (m, ℓ, O) triple needs AllReducing.
|
||||
Payload (h_q · S_q · d_head · 2) ≈ 4 KB per AR step, hierarchical ≈ 32 KB / layer.
|
||||
CONSTANT in S_kv.
|
||||
|
||||
Case 1 : inter-cube AR on (m,ℓ,O) → ~32 KB / layer
|
||||
Case 2 : nothing → 0
|
||||
Case 3 : intra-cube AR on (m,ℓ,O) → ~32 KB / layer
|
||||
Case 4 ★ : intra + inter-cube AR on (m,ℓ,O) → ~64 KB / layer
|
||||
|
||||
d_head-TP variants (1', 3'):
|
||||
Each rank holds only d_head/divisor dims → Q·K^T produces only
|
||||
partial sums → must AR PARTIAL SCORES before softmax. The score
|
||||
tile per AR is (h_q · S_q · S_kv / slices_for_seq) · 2 bytes,
|
||||
which SCALES with S_kv.
|
||||
|
||||
Case 1' : intra-cube partial-score AR
|
||||
payload = h_q · S_q · (S_kv/C) · 2 = 2 MB / layer at 1 M
|
||||
Case 3' : inter-cube partial-score AR (UCIe, slower than NoC)
|
||||
payload = h_q · S_q · (S_kv/P) · 2 = 2 MB / layer at 1 M
|
||||
|
||||
Plus tiny (m,ℓ,O) AR for both 1' and 3' to combine partial attentions
|
||||
across the remaining axis (~32 KB / layer).
|
||||
|
||||
(C) Total per output token per PE (sum over 80 layers, S_kv = 1 M):
|
||||
Case 1 ≈ 4 MB (1.25 MB Wo+FFN + 2.5 MB inter-cube AR)
|
||||
Case 1' ≈ 166 MB (1.25 MB + 160 MB partial-score AR + ~5 MB other)
|
||||
Case 2 ≈ 1.2 MB (Wo+FFN only)
|
||||
Case 3 ≈ 4 MB (1.25 MB + 2.5 MB intra-cube AR)
|
||||
Case 3' ≈ 166 MB (1.25 MB + 160 MB partial-score AR + ~5 MB other)
|
||||
Case 4 ★ ≈ 6 MB (1.25 MB + 5 MB 2-phase AR)
|
||||
|
||||
================================================================
|
||||
KEY TAKEAWAYS
|
||||
================================================================
|
||||
- Memory winners (640 MB / PE @ 1M): Cases 1', 3', 4 (any 64-way sharding).
|
||||
- Comm winner among those three: Case 4 (6 MB). Cases 1' and 3' both
|
||||
pay ~160 MB because the d_head-sharded variant ARs a 2 MB partial-
|
||||
score tile every layer — the score tensor scales with S_kv while
|
||||
the (m, ℓ, O) triple does not.
|
||||
- Cases 1 and 3 (8-way sharding) don't fit 1M context (5 GB vs 4.24 GB
|
||||
headroom).
|
||||
- Case 2 has the cheapest comm but the most memory (40 GB / PE).
|
||||
- Case 4 ★ is the Pareto-best: fits 1M context AND lowest comm
|
||||
among memory-feasible options.
|
||||
|
||||
Output PNG:
|
||||
src/kernbench/benches/1H_milestone_output/gqa/long_ctx/
|
||||
gqa_4cases_summary.png
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import textwrap
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.patches as mpatches
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
# ── System constants (LLaMA-3.1-70B single-KV-head group, decode) ───
|
||||
_N_LAYERS = 80
|
||||
_H_KV = 1
|
||||
_H_Q = 8
|
||||
_D_HEAD = 128
|
||||
_D_MODEL = 8192
|
||||
_BYTES_PER_ELEM = 2 # FP16
|
||||
_C = 8 # cubes per KV group
|
||||
_P = 8 # PEs per cube
|
||||
_B = 1
|
||||
_S_Q = 1 # decode
|
||||
_HBM_PER_PE_GB = 6.0
|
||||
_WEIGHTS_PER_PE_GB = 1.76
|
||||
_HEADROOM_GB = _HBM_PER_PE_GB - _WEIGHTS_PER_PE_GB # 4.24 GB
|
||||
|
||||
_HEADLINE_S_KV = 1 << 20 # 1 Mi tokens
|
||||
|
||||
# Per-token KV bytes (single KV group, all 80 layers).
|
||||
_KV_PER_TOK_BYTES = (
|
||||
2 * _H_KV * _D_HEAD * _BYTES_PER_ELEM * _N_LAYERS
|
||||
) # 40 KB
|
||||
|
||||
# Per-token Wo + FFN AR (constant across cases).
|
||||
_WO_PER_LAYER_BYTES = 8 * 1024
|
||||
_FFN_PER_LAYER_BYTES = 8 * 1024
|
||||
|
||||
# ── Per-PE attention weight breakdown (single KV-head group) ────────
|
||||
#
|
||||
# LLaMA-3.1-70B single-KV-head group dimensions:
|
||||
# d_model = 8192
|
||||
# h_q per group = 8 (8 query heads attend to 1 KV head per group)
|
||||
# h_kv per group = 1
|
||||
# d_head = 128
|
||||
# FP16 (2 bytes)
|
||||
#
|
||||
# Per layer, per PE:
|
||||
# Wq shape (d_model, h_q · d_head) REPLICATED across all 64 PEs of group
|
||||
# = 8192 · 8·128 · 2 bytes = 16 MB / layer / PE
|
||||
# Wk shape (d_model, h_kv · d_head) REPLICATED
|
||||
# = 8192 · 1·128 · 2 bytes = 2 MB / layer / PE
|
||||
# Wv shape (d_model, h_kv · d_head) REPLICATED
|
||||
# = 8192 · 1·128 · 2 bytes = 2 MB / layer / PE
|
||||
# Wo shape (h_q · d_head, d_model) ROW-SPLIT across C=8 cubes,
|
||||
# replicated within cube
|
||||
# = (1024/C) · 8192 · 2 bytes = 2 MB / layer / PE
|
||||
# ────────────────────────────────────────────────────────────────
|
||||
# total attn weights / layer / PE = 22 MB
|
||||
# × 80 layers = 1.76 GB / PE (decimal GB)
|
||||
#
|
||||
# KV headroom per PE = HBM_PER_PE - weights = 6.0 - 1.76 = 4.24 GB.
|
||||
# (FFN weights are accounted for in a separate budget, not in this 4.24 GB.)
|
||||
|
||||
# Slide-17 convention: per-layer values in binary MB (MiB), totals in
|
||||
# "GB" formed by ×80 layers ÷ 1000 — giving the canonical 1.76 GB total
|
||||
# and 4.24 GB headroom that the slide-17 chart reports.
|
||||
_WQ_MB_PER_LAYER = (_D_MODEL * _H_Q * _D_HEAD * _BYTES_PER_ELEM) / (1024 ** 2) # 16
|
||||
_WK_MB_PER_LAYER = (_D_MODEL * _H_KV * _D_HEAD * _BYTES_PER_ELEM) / (1024 ** 2) # 2
|
||||
_WV_MB_PER_LAYER = (_D_MODEL * _H_KV * _D_HEAD * _BYTES_PER_ELEM) / (1024 ** 2) # 2
|
||||
_WO_MB_PER_LAYER = ((_H_Q * _D_HEAD) // _C * _D_MODEL
|
||||
* _BYTES_PER_ELEM) / (1024 ** 2) # 2
|
||||
|
||||
_WQ_GB = _WQ_MB_PER_LAYER * _N_LAYERS / 1000 # 1.28 GB
|
||||
_WK_GB = _WK_MB_PER_LAYER * _N_LAYERS / 1000 # 0.16 GB
|
||||
_WV_GB = _WV_MB_PER_LAYER * _N_LAYERS / 1000 # 0.16 GB
|
||||
_WO_GB = _WO_MB_PER_LAYER * _N_LAYERS / 1000 # 0.16 GB
|
||||
_WEIGHTS_GB = _WQ_GB + _WK_GB + _WV_GB + _WO_GB # 1.76 GB
|
||||
|
||||
# (m, ℓ, O) AR cost per layer, derived from the kernel topology.
|
||||
#
|
||||
# Per-PE payload T for one (m, ℓ, O) merge step:
|
||||
# O — h_q · S_q · d_head · 2 bytes (FP16)
|
||||
# m — h_q · S_q · 4 bytes (FP32)
|
||||
# ℓ — h_q · S_q · 4 bytes (FP32)
|
||||
# T = h_q · S_q · (d_head · 2 + 8) ≈ 2.1 KB
|
||||
#
|
||||
# Reduce algorithm: the decode kernels use hierarchical reduce-only
|
||||
# (chain / tree) rather than ring all-reduce, because the merged result
|
||||
# only needs to land on the cube that runs the downstream Wo gemm —
|
||||
# not on every PE. For a chain of N participants the total traffic is
|
||||
# (N-1)·T and the per-PE average is (N-1)/N · T.
|
||||
#
|
||||
# The previous version used a hard-coded 32 KB / layer placeholder which
|
||||
# overestimated the per-PE cost by ~6× (and ~3× even under a hypothetical
|
||||
# ring-AR assumption). See `gqa_long_ctx_6cases_measured_comm.json` for
|
||||
# the measured numbers this matches against.
|
||||
_MLO_PAYLOAD_BYTES = _H_Q * _S_Q * (_D_HEAD * _BYTES_PER_ELEM + 8)
|
||||
|
||||
|
||||
def _reduce_chain_per_pe_bytes(n_participants: int) -> int:
|
||||
"""Per-PE average bytes for a single-stage chain/tree reduce of T."""
|
||||
if n_participants <= 1:
|
||||
return 0
|
||||
return _MLO_PAYLOAD_BYTES * (n_participants - 1) // n_participants
|
||||
|
||||
|
||||
_MLO_INTRA_BYTES_PER_LAYER = _reduce_chain_per_pe_bytes(_P) # PE-axis
|
||||
_MLO_INTER_BYTES_PER_LAYER = _reduce_chain_per_pe_bytes(_C) # cube-axis
|
||||
|
||||
# Cases — renumbered in MEMORY-DESCENDING ORDER (left to right):
|
||||
# Case 1 (40 GB) : no sharding
|
||||
# Cases 2, 3 ( 5 GB) : single-axis sharding (cube OR PE only)
|
||||
# Cases 4, 5, 6 (640 MB) : two-axis sharding (cube AND PE)
|
||||
# — Case 6 ★ is the Pareto-best (lowest comm)
|
||||
_CASES = (1, 2, 3, 4, 5, 6)
|
||||
|
||||
_DIVISOR = {
|
||||
1: 1, # Cube-Repl × PE-replicate — no sharding
|
||||
2: _C, # Cube-SP × PE-replicate — cube-axis only
|
||||
3: _P, # Cube-Repl × PE-SP — PE-axis (S_kv) only
|
||||
4: _C * _P, # Cube-SP × PE-TP (d_head) — 64-way (d_head intra)
|
||||
5: _C * _P, # Cube-TP × PE-SP — 64-way (d_head inter)
|
||||
6: _C * _P, # Cube-SP × PE-SP — 64-way (S_kv both axes) ★
|
||||
}
|
||||
|
||||
_CASE_LABEL = {
|
||||
1: "Case 1\nCube-Repl\nPE-repl",
|
||||
2: "Case 2\nCube-SP\nPE-repl",
|
||||
3: "Case 3\nCube-Repl\nPE-SP",
|
||||
4: "Case 4\nCube-SP\nPE-TP",
|
||||
5: "Case 5\nCube-TP\nPE-SP",
|
||||
6: "Case 6 ★\nCube-SP\nPE-SP",
|
||||
}
|
||||
_CASE_COLOR = {
|
||||
1: "#C0504D", # red — worst memory (no sharding)
|
||||
2: "#E0834A", # orange — single-axis sharded (cube)
|
||||
3: "#EBA854", # tan — single-axis sharded (PE)
|
||||
4: "#A6C2E0", # light blue — d_head-TP 64-way (PE)
|
||||
5: "#C7D8A0", # light green — d_head-TP 64-way (cube)
|
||||
6: "#8064A2", # purple — Pareto winner ★ (S_kv 64-way)
|
||||
}
|
||||
_ATTN_DESC = {
|
||||
1: "none",
|
||||
2: "online-softmax (m,ℓ,O) — inter-cube",
|
||||
3: "online-softmax (m,ℓ,O) — intra-cube",
|
||||
4: "partial scores + (m,ℓ,O) merge (d_head-TP)",
|
||||
5: "partial scores + (m,ℓ,O) merge (d_head-TP)",
|
||||
6: "online-softmax (m,ℓ,O) — intra + inter",
|
||||
}
|
||||
|
||||
_WO_COLOR = "#9EC5E8"
|
||||
_FFN_COLOR = "#4A78B8"
|
||||
_ATTN_COLOR = "#E07A3F"
|
||||
|
||||
_OUT_DIR = (
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ "src" / "kernbench" / "benches"
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_MEASURED_JSON = _OUT_DIR / "gqa_long_ctx_6cases_measured_comm.json"
|
||||
|
||||
|
||||
def _load_measured() -> dict[int, float] | None:
|
||||
"""Load measured per-PE comm bytes (already scaled to S_kv=1M).
|
||||
|
||||
Returns {case_id: per_token_total_bytes} or None if JSON missing.
|
||||
Produced by scripts/paper/measure_gqa_decode_placement_comm.py.
|
||||
"""
|
||||
if not _MEASURED_JSON.exists():
|
||||
return None
|
||||
data = json.loads(_MEASURED_JSON.read_text())
|
||||
return {
|
||||
int(cid): info["per_pe_total_bytes_per_token_at_1M"]
|
||||
for cid, info in data["cases"].items()
|
||||
}
|
||||
|
||||
|
||||
# ── Formulae ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def kv_per_pe_bytes(case, s_kv: int) -> int:
|
||||
"""Per-PE KV bytes at the given S_kv."""
|
||||
return _KV_PER_TOK_BYTES * s_kv // _DIVISOR[case]
|
||||
|
||||
|
||||
def max_s_kv(case) -> int:
|
||||
"""Max KV context per PE given the 4.24 GB headroom."""
|
||||
headroom_bytes = int(_HEADROOM_GB * (1 << 30))
|
||||
bytes_per_tok = _KV_PER_TOK_BYTES // _DIVISOR[case]
|
||||
return headroom_bytes // bytes_per_tok
|
||||
|
||||
|
||||
def _partial_score_bytes_per_layer(s_kv: int, slices_for_seq: int) -> int:
|
||||
"""AR payload for partial scores in d_head-TP cases."""
|
||||
return _H_Q * _S_Q * (s_kv // slices_for_seq) * _BYTES_PER_ELEM
|
||||
|
||||
|
||||
def attn_comm_per_layer_bytes(case: int, s_kv: int) -> int:
|
||||
"""Per-layer per-PE attention-time comm bytes (decode, B=1, S_q=1).
|
||||
|
||||
Case numbering follows the memory-descending order defined in
|
||||
_CASES (1=no sharding, 6=Pareto winner).
|
||||
"""
|
||||
if case == 1: # Cube-Repl × PE-repl — no sharding
|
||||
return 0
|
||||
if case == 2: # Cube-SP × PE-repl — inter-cube AR on (m,ℓ,O)
|
||||
return _MLO_INTER_BYTES_PER_LAYER
|
||||
if case == 3: # Cube-Repl × PE-SP — intra-cube AR on (m,ℓ,O)
|
||||
return _MLO_INTRA_BYTES_PER_LAYER
|
||||
if case == 4: # Cube-SP × PE-TP(d_head) — partial-score AR
|
||||
return (_partial_score_bytes_per_layer(s_kv, _C)
|
||||
+ _MLO_INTRA_BYTES_PER_LAYER
|
||||
+ _MLO_INTER_BYTES_PER_LAYER)
|
||||
if case == 5: # Cube-TP(d_head) × PE-SP — partial-score AR
|
||||
return (_partial_score_bytes_per_layer(s_kv, _P)
|
||||
+ _MLO_INTRA_BYTES_PER_LAYER
|
||||
+ _MLO_INTER_BYTES_PER_LAYER)
|
||||
if case == 6: # Cube-SP × PE-SP — 2-phase (m,ℓ,O) AR ★
|
||||
return _MLO_INTRA_BYTES_PER_LAYER + _MLO_INTER_BYTES_PER_LAYER
|
||||
raise ValueError(f"unknown case {case}")
|
||||
|
||||
|
||||
def per_token_bytes(case, s_kv: int) -> tuple[int, int, int]:
|
||||
"""(Wo AR, FFN AR, Attn) bytes per output token per PE — × 80 layers."""
|
||||
wo = _WO_PER_LAYER_BYTES * _N_LAYERS
|
||||
ffn = _FFN_PER_LAYER_BYTES * _N_LAYERS
|
||||
attn = attn_comm_per_layer_bytes(case, s_kv) * _N_LAYERS
|
||||
return wo, ffn, attn
|
||||
|
||||
|
||||
# ── Formatters ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _fmt_bytes(b: float) -> str:
|
||||
if b >= (1 << 30):
|
||||
return f"{b / (1 << 30):.2f} GB"
|
||||
if b >= (1 << 20):
|
||||
return f"{b / (1 << 20):.1f} MB"
|
||||
if b >= (1 << 10):
|
||||
return f"{b / (1 << 10):.0f} KB"
|
||||
return f"{b:.0f} B"
|
||||
|
||||
|
||||
def _fmt_tokens(n: int) -> str:
|
||||
if n >= 1_000_000:
|
||||
return f"{n / 1_000_000:.2f} M"
|
||||
if n >= 1_000:
|
||||
return f"{n / 1_000:.0f} K"
|
||||
return f"{n}"
|
||||
|
||||
|
||||
# ── Panels ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _plot_budget(ax) -> None:
|
||||
"""HBM budget per PE — stacked weights + KV headroom + HBM ceiling.
|
||||
|
||||
Wq is the dominant weight slice (~1.28 GB). Wk, Wv, Wo are each
|
||||
small (~0.16 GB) so their slice labels would overlap on the bar —
|
||||
they're shown in the legend only, and only Wq + KV-headroom get
|
||||
on-bar annotations.
|
||||
"""
|
||||
components = [
|
||||
("Wq (REPL)", _WQ_GB, "#7B9CC4"),
|
||||
("Wk (REPL)", _WK_GB, "#A0BBD8"),
|
||||
("Wv (REPL)", _WV_GB, "#C5D6E8"),
|
||||
("Wo (cube-split)", _WO_GB, "#E2EAF3"),
|
||||
("KV cache headroom",
|
||||
_HBM_PER_PE_GB - _WEIGHTS_GB, "#9BBB59"),
|
||||
]
|
||||
bottom = 0.0
|
||||
for label, val, color in components:
|
||||
ax.bar(0, val, bottom=bottom, color=color, edgecolor="black",
|
||||
width=0.7, label=f"{label} · {val:.2f} GB")
|
||||
# Only annotate slices thick enough to fit text without overlap.
|
||||
if val >= 0.50:
|
||||
ax.text(0, bottom + val / 2, f"{label}\n{val:.2f} GB",
|
||||
ha="center", va="center", fontsize=9, weight="bold")
|
||||
bottom += val
|
||||
ax.axhline(_HBM_PER_PE_GB, color="red", ls="--", lw=1.4,
|
||||
label=f"HBM = {_HBM_PER_PE_GB} GB")
|
||||
ax.set_xticks([0])
|
||||
ax.set_xticklabels(["per-PE HBM"], fontsize=10)
|
||||
ax.set_ylabel("GB per PE")
|
||||
ax.set_ylim(0, _HBM_PER_PE_GB * 1.10)
|
||||
ax.set_title(
|
||||
f"Per-PE HBM budget\n"
|
||||
f"weights {_WEIGHTS_GB:.2f} GB + KV = {_HBM_PER_PE_GB} GB",
|
||||
fontsize=10,
|
||||
)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
ax.legend(loc="upper right", fontsize=7.5, framealpha=0.92)
|
||||
|
||||
|
||||
def _plot_memory(ax) -> None:
|
||||
vals_gb = [kv_per_pe_bytes(c, _HEADLINE_S_KV) / (1 << 30) for c in _CASES]
|
||||
labels = [_CASE_LABEL[c] for c in _CASES]
|
||||
colors = [_CASE_COLOR[c] for c in _CASES]
|
||||
x = list(range(len(_CASES)))
|
||||
bars = ax.bar(x, vals_gb, color=colors, width=0.65)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(labels, fontsize=9)
|
||||
ax.set_ylabel("KV bytes per PE (GB, log)")
|
||||
ax.set_yscale("log")
|
||||
ax.set_title(
|
||||
f"Per-PE KV memory at S_kv = {_HEADLINE_S_KV:,} tokens "
|
||||
f"(across {_N_LAYERS} layers, FP16)",
|
||||
fontsize=11,
|
||||
)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5, which="both")
|
||||
ax.axhline(_HEADROOM_GB, color="red", ls="--", lw=1.4,
|
||||
label=f"KV headroom = {_HEADROOM_GB} GB / PE")
|
||||
for bar, v_gb in zip(bars, vals_gb):
|
||||
v_bytes = v_gb * (1 << 30)
|
||||
fits = v_gb <= _HEADROOM_GB
|
||||
ax.text(bar.get_x() + bar.get_width() / 2,
|
||||
v_gb * 1.10,
|
||||
_fmt_bytes(v_bytes) + (" ✓" if fits else " ✗"),
|
||||
ha="center", va="bottom", fontsize=9,
|
||||
color="green" if fits else "red", weight="bold")
|
||||
ax.legend(loc="upper right", fontsize=9)
|
||||
|
||||
|
||||
def _plot_comm(ax, *, mode: str = "analytical") -> None:
|
||||
"""Per-PE comm panel.
|
||||
|
||||
mode = "analytical": single solid bars from per_token_bytes formula.
|
||||
mode = "paired" : analytical (solid) + simulator-measured
|
||||
(hatched) side-by-side per case, when the
|
||||
measurement JSON is available.
|
||||
"""
|
||||
labels = [_CASE_LABEL[c] for c in _CASES]
|
||||
|
||||
wo_mb_list: list[float] = []
|
||||
ffn_mb_list: list[float] = []
|
||||
attn_mb: list[float] = []
|
||||
for c in _CASES:
|
||||
wo, ffn, attn = per_token_bytes(c, _HEADLINE_S_KV)
|
||||
wo_mb_list.append(wo / (1 << 20))
|
||||
ffn_mb_list.append(ffn / (1 << 20))
|
||||
attn_mb.append(attn / (1 << 20))
|
||||
|
||||
measured = _load_measured() if mode == "paired" else None
|
||||
paired = measured is not None
|
||||
source_tag = "analytical (solid) vs simulator-measured (hatched)" \
|
||||
if paired else "analytical"
|
||||
|
||||
n_cases = len(_CASES)
|
||||
x = list(range(n_cases))
|
||||
bar_w = 0.36 if paired else 0.65
|
||||
x_ana = [xi - bar_w / 2 for xi in x] if paired else x
|
||||
x_meas = [xi + bar_w / 2 for xi in x] if paired else None
|
||||
|
||||
# Analytical bars (solid).
|
||||
ax.bar(x_ana, wo_mb_list, width=bar_w,
|
||||
color=_WO_COLOR, edgecolor="black")
|
||||
ax.bar(x_ana, ffn_mb_list, width=bar_w, bottom=wo_mb_list,
|
||||
color=_FFN_COLOR, edgecolor="black")
|
||||
bottoms_attn = [w + f for w, f in zip(wo_mb_list, ffn_mb_list)]
|
||||
ax.bar(x_ana, attn_mb, width=bar_w, bottom=bottoms_attn,
|
||||
color=_ATTN_COLOR, edgecolor="black")
|
||||
|
||||
# Measured bars (hatched) — same Wo+FFN base, attn from op_log.
|
||||
meas_attn_mb: list[float] = []
|
||||
if paired:
|
||||
for i, c in enumerate(_CASES):
|
||||
meas_total = measured.get(c, 0) / (1 << 20)
|
||||
meas_attn_mb.append(
|
||||
max(meas_total - wo_mb_list[i] - ffn_mb_list[i], 0.0))
|
||||
ax.bar(x_meas, wo_mb_list, width=bar_w,
|
||||
color=_WO_COLOR, edgecolor="black",
|
||||
hatch="///", alpha=0.85)
|
||||
ax.bar(x_meas, ffn_mb_list, width=bar_w, bottom=wo_mb_list,
|
||||
color=_FFN_COLOR, edgecolor="black",
|
||||
hatch="///", alpha=0.85)
|
||||
ax.bar(x_meas, meas_attn_mb, width=bar_w, bottom=bottoms_attn,
|
||||
color=_ATTN_COLOR, edgecolor="black",
|
||||
hatch="///", alpha=0.85)
|
||||
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(labels, fontsize=9)
|
||||
ax.set_ylabel("Comm per token per PE (MB, log)")
|
||||
ax.set_yscale("log")
|
||||
title = (
|
||||
f"Per-PE communication at S_kv = {_HEADLINE_S_KV:,} tokens "
|
||||
f"(decode S_q=1, B=1; {_N_LAYERS} layers) — {source_tag}"
|
||||
)
|
||||
if paired:
|
||||
title += "\n(simulator measured at S_kv = 8K; " \
|
||||
"partial-score AR scaled ×128 to S_kv = 1M)"
|
||||
ax.set_title(title, fontsize=10)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5, which="both")
|
||||
|
||||
totals_ana = [wo_mb_list[i] + ffn_mb_list[i] + attn_mb[i]
|
||||
for i in range(n_cases)]
|
||||
ymax = max(totals_ana)
|
||||
if paired:
|
||||
ymax = max(ymax, max(
|
||||
(measured.get(c, 0) / (1 << 20)) for c in _CASES))
|
||||
ax.set_ylim(top=ymax * 22)
|
||||
|
||||
for i, c in enumerate(_CASES):
|
||||
ana_bytes = totals_ana[i] * (1 << 20)
|
||||
if paired:
|
||||
meas_bytes = measured.get(c, 0)
|
||||
top_y = max(totals_ana[i], meas_bytes / (1 << 20))
|
||||
label = (f"ana: {_fmt_bytes(ana_bytes)}\n"
|
||||
f"sim: {_fmt_bytes(meas_bytes)}")
|
||||
else:
|
||||
top_y = totals_ana[i]
|
||||
label = _fmt_bytes(ana_bytes)
|
||||
ax.text(x[i], top_y * 1.5, label,
|
||||
ha="center", va="bottom",
|
||||
fontsize=8 if paired else 9,
|
||||
weight="bold", linespacing=1.05)
|
||||
|
||||
# Attention-time AR descriptor — placed at the attention-segment
|
||||
# midpoint, centered between the analytical and measured bars so
|
||||
# it visually spans both. Coloured the same as the attention bar
|
||||
# (no background box) so it lives within the case's bar zone.
|
||||
# Wrapped narrow so each line fits inside the bar-pair width.
|
||||
wrapped = textwrap.fill(_ATTN_DESC[c], width=14)
|
||||
if attn_mb[i] > 0:
|
||||
mid = bottoms_attn[i] + attn_mb[i] / 2
|
||||
ax.text(x[i], mid, wrapped,
|
||||
ha="center", va="center",
|
||||
fontsize=7 if paired else 8,
|
||||
color="black", weight="bold",
|
||||
linespacing=1.0)
|
||||
else:
|
||||
ax.text(x[i], totals_ana[i] * 0.4, wrapped,
|
||||
ha="center",
|
||||
fontsize=7 if paired else 8,
|
||||
color="grey", style="italic",
|
||||
linespacing=1.0)
|
||||
|
||||
legend_handles = [
|
||||
mpatches.Patch(facecolor=_WO_COLOR, edgecolor="black",
|
||||
label=f"Wo AR (× {_N_LAYERS} layers)"),
|
||||
mpatches.Patch(facecolor=_FFN_COLOR, edgecolor="black",
|
||||
label=f"FFN AR (× {_N_LAYERS} layers)"),
|
||||
mpatches.Patch(facecolor=_ATTN_COLOR, edgecolor="black",
|
||||
label="Attn-time collective"),
|
||||
]
|
||||
if paired:
|
||||
legend_handles.append(
|
||||
mpatches.Patch(facecolor="white", edgecolor="black",
|
||||
hatch="///", label="simulator-measured"))
|
||||
ax.legend(handles=legend_handles, loc="upper right", fontsize=8,
|
||||
framealpha=0.92)
|
||||
|
||||
|
||||
def main() -> Path:
|
||||
_OUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# (a) Per-PE HBM budget — standalone PNG.
|
||||
fig_b, ax_b = plt.subplots(figsize=(4.0, 6.0))
|
||||
_plot_budget(ax_b)
|
||||
fig_b.tight_layout()
|
||||
out_b = _OUT_DIR / "gqa_long_ctx_6cases_hbm_budget.png"
|
||||
fig_b.savefig(out_b, dpi=150)
|
||||
plt.close(fig_b)
|
||||
print(f"wrote {out_b}")
|
||||
|
||||
# (b) Combined 3-panel summary — HBM budget + KV memory + comm.
|
||||
fig = plt.figure(figsize=(22.0, 6.5))
|
||||
gs = fig.add_gridspec(1, 3, width_ratios=[0.7, 1.6, 1.6], wspace=0.22)
|
||||
ax_b2 = fig.add_subplot(gs[0, 0])
|
||||
ax_m = fig.add_subplot(gs[0, 1])
|
||||
ax_c = fig.add_subplot(gs[0, 2])
|
||||
_plot_budget(ax_b2)
|
||||
_plot_memory(ax_m)
|
||||
_plot_comm(ax_c)
|
||||
fig.tight_layout()
|
||||
out = _OUT_DIR / "gqa_long_ctx_6cases_summary.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
print(f"wrote {out}")
|
||||
|
||||
# (c) 2-panel companion (analytical only).
|
||||
fig2 = plt.figure(figsize=(18.0, 6.5))
|
||||
gs2 = fig2.add_gridspec(1, 2, width_ratios=[1.0, 1.0], wspace=0.18)
|
||||
ax_m2 = fig2.add_subplot(gs2[0, 0])
|
||||
ax_c2 = fig2.add_subplot(gs2[0, 1])
|
||||
_plot_memory(ax_m2)
|
||||
_plot_comm(ax_c2, mode="analytical")
|
||||
fig2.tight_layout()
|
||||
out2 = _OUT_DIR / "gqa_long_ctx_6cases_memory_comm_analytical.png"
|
||||
fig2.savefig(out2, dpi=150)
|
||||
plt.close(fig2)
|
||||
print(f"wrote {out2}")
|
||||
|
||||
# (d) 2-panel companion — analytical vs simulator-measured paired.
|
||||
fig3 = plt.figure(figsize=(19.0, 6.5))
|
||||
gs3 = fig3.add_gridspec(1, 2, width_ratios=[1.0, 1.0], wspace=0.18)
|
||||
ax_m3 = fig3.add_subplot(gs3[0, 0])
|
||||
ax_c3 = fig3.add_subplot(gs3[0, 1])
|
||||
_plot_memory(ax_m3)
|
||||
_plot_comm(ax_c3, mode="paired")
|
||||
fig3.tight_layout()
|
||||
out3 = _OUT_DIR / "gqa_long_ctx_6cases_memory_comm_paired.png"
|
||||
fig3.savefig(out3, dpi=150)
|
||||
plt.close(fig3)
|
||||
print(f"wrote {out3}")
|
||||
|
||||
# Paper-ready table to stdout.
|
||||
print()
|
||||
print(f" {'Case':<27} {'KV/tok·PE':>12} {'KV @ 1M':>11} "
|
||||
f"{'Max S_kv':>11} {'Comm @ 1M':>12}")
|
||||
print(" " + "-" * 80)
|
||||
for c in _CASES:
|
||||
kv_per_tok_pe = _KV_PER_TOK_BYTES / _DIVISOR[c]
|
||||
mem_1m = kv_per_pe_bytes(c, _HEADLINE_S_KV)
|
||||
max_s = max_s_kv(c)
|
||||
comm_1m = sum(per_token_bytes(c, _HEADLINE_S_KV))
|
||||
label = _CASE_LABEL[c].replace(chr(10), " ")
|
||||
print(f" {label:<27} "
|
||||
f"{kv_per_tok_pe / 1024:>9.3f} KB "
|
||||
f"{_fmt_bytes(mem_1m):>11} "
|
||||
f"{_fmt_tokens(max_s):>11} "
|
||||
f"{_fmt_bytes(comm_1m):>12}")
|
||||
return out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,131 @@
|
||||
"""Latency breakdown bar chart for the Case-6 composite-command decode study.
|
||||
|
||||
Reads sweep_decode_composite.json and writes a single-figure stacked bar
|
||||
chart comparing three variants — primitive hand-tiled (16×16×16),
|
||||
composite GEMM, composite + softmax_merge — at the S_kv = 131 072 point:
|
||||
|
||||
bottom stack: PE_CPU dispatch time (from pe_cpu_dispatch_cycles, ns at
|
||||
1 GHz — ADR-0064 Rev2 D3)
|
||||
top stack: engine time (latency_ns − dispatch cycles) — DMA / GEMM /
|
||||
MATH / IPCQ work the engine flushes on the critical path
|
||||
|
||||
The dispatch/engine split is a first-order breakdown; in reality the
|
||||
two paths overlap partially. The note on the plot calls this out.
|
||||
|
||||
Run (after the composite bench sweep):
|
||||
python scripts/paper/paper_plot_gqa_decode_composite_breakdown.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt # noqa: E402
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
_FIG_DIR = (
|
||||
_REPO_ROOT / "src" / "kernbench" / "benches"
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_SWEEP_JSON = _FIG_DIR / "sweep_decode_composite.json"
|
||||
_PAPER_FIG_DIR = (
|
||||
_REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
|
||||
)
|
||||
|
||||
# The three variants to compare (coarse `primitive` intentionally dropped).
|
||||
_ORDER = ("primitive_tiled", "composite", "composite_extended")
|
||||
_LABELS = {
|
||||
"primitive_tiled": "primitive hand-tiled\n(16×16×16)",
|
||||
"composite": "composite\nGEMM",
|
||||
"composite_extended": "composite +\nsoftmax_merge",
|
||||
}
|
||||
_S_KV_TARGET = 131_072
|
||||
|
||||
|
||||
def main() -> None:
|
||||
sweep = json.loads(_SWEEP_JSON.read_text())
|
||||
rows = {(r["variant"], r["S_kv"]): r for r in sweep["rows"]}
|
||||
|
||||
dispatch_us = []
|
||||
engine_us = []
|
||||
totals_us = []
|
||||
for v in _ORDER:
|
||||
r = rows[(v, _S_KV_TARGET)]
|
||||
disp_ns = r["pe_cpu_dispatch_cycles"]
|
||||
total_ns = r["latency_ns"]
|
||||
eng_ns = max(0.0, total_ns - disp_ns)
|
||||
dispatch_us.append(disp_ns / 1e3)
|
||||
engine_us.append(eng_ns / 1e3)
|
||||
totals_us.append(total_ns / 1e3)
|
||||
|
||||
xs = list(range(len(_ORDER)))
|
||||
labels = [_LABELS[v] for v in _ORDER]
|
||||
|
||||
fig, ax = plt.subplots(figsize=(7.5, 5.0))
|
||||
|
||||
bars_eng = ax.bar(
|
||||
xs, engine_us,
|
||||
color="#4f8a4f", edgecolor="#2a4a2a", label="engine (DMA + GEMM + MATH + IPCQ)",
|
||||
)
|
||||
bars_disp = ax.bar(
|
||||
xs, dispatch_us, bottom=engine_us,
|
||||
color="#c0504d", edgecolor="#5a2624", label="PE_CPU dispatch",
|
||||
)
|
||||
|
||||
# Segment value labels — engine at mid, dispatch at mid of its stack.
|
||||
for i, (eng, disp, tot) in enumerate(zip(engine_us, dispatch_us, totals_us)):
|
||||
ax.text(i, eng / 2, f"{eng:.1f} µs",
|
||||
ha="center", va="center", fontsize=9, color="white")
|
||||
if disp / max(totals_us) > 0.04: # only label if visible
|
||||
ax.text(i, eng + disp / 2, f"{disp:.1f} µs",
|
||||
ha="center", va="center", fontsize=9, color="white")
|
||||
else:
|
||||
ax.annotate(f"{disp:.2f} µs",
|
||||
xy=(i, tot), xytext=(0, 6),
|
||||
textcoords="offset points",
|
||||
ha="center", va="bottom", fontsize=8,
|
||||
color="#5a2624")
|
||||
offset_pts = 14 if disp / max(totals_us) > 0.04 else 20
|
||||
ax.annotate(f"total {tot:.1f}",
|
||||
xy=(i, tot), xytext=(0, offset_pts),
|
||||
textcoords="offset points",
|
||||
ha="center", va="bottom", fontsize=9,
|
||||
color="#333", fontweight="bold")
|
||||
|
||||
ax.set_xticks(xs)
|
||||
ax.set_xticklabels(labels, fontsize=10)
|
||||
ax.set_ylabel("time (µs)")
|
||||
ax.set_title(
|
||||
f"Case-6 decode latency breakdown at $S_{{kv}}=${_S_KV_TARGET // 1024}K\n"
|
||||
"(engine dominates all three — primitive hand-tiled adds "
|
||||
f"{dispatch_us[0]:.0f} µs PE_CPU dispatch overhead)",
|
||||
fontsize=11,
|
||||
)
|
||||
ax.legend(loc="upper right", fontsize=9)
|
||||
ax.grid(True, axis="y", ls=":", alpha=0.5)
|
||||
ax.set_ylim(0, max(totals_us) * 1.15)
|
||||
|
||||
fig.text(
|
||||
0.5, 0.02,
|
||||
"First-order breakdown: dispatch and engine paths overlap partially on the real "
|
||||
"critical path;\ntreat the split as an upper bound on dispatch's contribution.",
|
||||
ha="center", fontsize=8, color="#666",
|
||||
)
|
||||
fig.tight_layout(rect=(0, 0.06, 1, 1))
|
||||
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_composite_breakdown.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
print(f"wrote {out}")
|
||||
|
||||
if _PAPER_FIG_DIR.is_dir():
|
||||
dst = _PAPER_FIG_DIR / out.name
|
||||
dst.write_bytes(out.read_bytes())
|
||||
print(f"copied {dst}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,13 +1,25 @@
|
||||
"""Comparative figures for milestone-gqa-decode-long-ctx-4cases.
|
||||
|
||||
Reads sweep.json (emitted by ``kernbench run --bench
|
||||
milestone-gqa-decode-long-ctx-4cases``) and writes four PNGs into
|
||||
``docs/report/1H-codesign-paper/figures/``:
|
||||
Reads sweep_decode.json (emitted by the milestone-1h-gqa bench) and
|
||||
writes four PNGs into the same bench-output dir
|
||||
(src/kernbench/benches/1H_milestone_output/gqa/long_ctx/):
|
||||
|
||||
gqa_decode_long_ctx_4cases_latency.png end-to-end latency per case
|
||||
gqa_decode_long_ctx_4cases_traffic.png ipcq/dma op-count breakdown
|
||||
gqa_decode_long_ctx_4cases_memory.png per-PE KV bytes per case
|
||||
gqa_decode_long_ctx_4cases_parallelism.png per-PE S_local (compute work)
|
||||
gqa_decode_long_ctx_6cases_latency.png end-to-end latency per case
|
||||
gqa_decode_long_ctx_6cases_traffic.png ipcq/dma op-count breakdown
|
||||
gqa_decode_long_ctx_6cases_memory.png per-PE KV bytes per case
|
||||
gqa_decode_long_ctx_6cases_parallelism.png per-PE S_local (compute work)
|
||||
|
||||
Filename still says "4cases" for backwards compat, but the script now
|
||||
covers all SIX kv-sharding strategies from the analytical chart
|
||||
(`gqa_4cases_summary.png`) — the original 4 plus the two new
|
||||
d_head-TP variants:
|
||||
|
||||
Case 1 Cube-Repl × PE-repl (PE-TP doesn't shard KV)
|
||||
Case 2 Cube-SP × PE-repl
|
||||
Case 3 Cube-Repl × PE-SP
|
||||
Case 4 Cube-SP × PE-TP-d_head ← NEW
|
||||
Case 5 Cube-TP-d_head × PE-SP ← NEW
|
||||
Case 6 ★ Cube-SP × PE-SP (Pareto-best)
|
||||
|
||||
Run (after the bench):
|
||||
GQA_DECODE_LONG_CTX_4CASES_RUN=1 python -m kernbench.cli.main run \\
|
||||
@@ -32,16 +44,27 @@ _FIG_DIR = (
|
||||
)
|
||||
_SWEEP_JSON = _FIG_DIR / "sweep_decode.json"
|
||||
|
||||
# Panel name → (short label, case ordinal for left-to-right plot order).
|
||||
# Panel name → (short label, case ordinal, accent flag) using the
|
||||
# analytical chart's memory-descending ordering. PE-TP doesn't shard
|
||||
# KV memory, so the cube_repl_pe_tp panel maps to Case 1 (no
|
||||
# sharding, KV-wise) and cube_sp_pe_tp panel maps to Case 2.
|
||||
_NORMAL, _OVERFLOW, _PARETO = "normal", "overflow", "pareto"
|
||||
|
||||
_CASE_INFO = {
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp": (
|
||||
"Case 1\nCube-SP × PE-TP", 1),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp": (
|
||||
"Case 2\nCube-Repl × PE-TP", 2),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp": (
|
||||
"Case 3\nCube-Repl × PE-SP", 3),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp": (
|
||||
"Case 4 ★\nCube-SP × PE-SP", 4),
|
||||
# panel name label ord flag
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_repl_pe_tp": ("Case 1\nCube-Repl × PE-repl", 1, _OVERFLOW),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp": ("Case 2\nCube-SP × PE-repl", 2, _OVERFLOW),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_repl_pe_sp": ("Case 3\nCube-Repl × PE-SP", 3, _OVERFLOW),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_tp_dhead": ("Case 4\nCube-SP × PE-TP-d_head", 4, _NORMAL),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_tp_dhead_pe_sp": ("Case 5\nCube-TP-d_head × PE-SP", 5, _NORMAL),
|
||||
"single_kv_group_decode_long_ctx_gqa_cube_sp_pe_sp": ("Case 6 ★\nCube-SP × PE-SP", 6, _PARETO),
|
||||
}
|
||||
|
||||
# Bar fill colour per flag (used by every panel).
|
||||
_FLAG_COLOR = {
|
||||
_NORMAL: "#888888", # neutral grey
|
||||
_OVERFLOW: "#c0504d", # red — fails the per-PE HBM budget
|
||||
_PARETO: "#3b6ea5", # blue — Pareto-best
|
||||
}
|
||||
|
||||
|
||||
@@ -53,23 +76,26 @@ def _sorted_by_case(rows: list[dict]) -> list[dict]:
|
||||
return sorted(rows, key=lambda r: _CASE_INFO[r["panel"]][1])
|
||||
|
||||
|
||||
def _bar_colors(rows: list[dict]) -> list[str]:
|
||||
return [_FLAG_COLOR[_CASE_INFO[r["panel"]][2]] for r in rows]
|
||||
|
||||
|
||||
def _plot_latency(rows: list[dict]) -> Path:
|
||||
rows = _sorted_by_case(rows)
|
||||
labels = [_CASE_INFO[r["panel"]][0] for r in rows]
|
||||
lat_us = [r["latency_ns"] / 1e3 for r in rows]
|
||||
colors = ["#888", "#888", "#888", "#3b6ea5"] # Case 4 highlighted
|
||||
fig, ax = plt.subplots(figsize=(8.0, 4.5))
|
||||
bars = ax.bar(labels, lat_us, color=colors, width=0.6)
|
||||
fig, ax = plt.subplots(figsize=(12.0, 4.8))
|
||||
bars = ax.bar(labels, lat_us, color=_bar_colors(rows), width=0.6)
|
||||
ax.set_ylabel("end-to-end latency (µs)")
|
||||
ax.set_title(
|
||||
"Long-context decode 4-cases — end-to-end latency per case\n"
|
||||
"Long-context decode 6-cases — end-to-end latency per case\n"
|
||||
"LLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs)"
|
||||
)
|
||||
ax.bar_label(bars, fmt="%.1f", padding=3, fontsize=9)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
ax.set_ylim(0, max(lat_us) * 1.15)
|
||||
fig.tight_layout()
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_4cases_latency.png"
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_6cases_latency.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
return out
|
||||
@@ -83,18 +109,18 @@ def _plot_traffic(rows: list[dict]) -> Path:
|
||||
disp = ["IPCQ copy", "DMA read", "DMA write"]
|
||||
colors = ["#c0504d", "#9bbb59", "#8064a2"]
|
||||
w = 0.25
|
||||
fig, ax = plt.subplots(figsize=(9.0, 4.5))
|
||||
fig, ax = plt.subplots(figsize=(11.0, 4.5))
|
||||
for i, (k, d, c) in enumerate(zip(keys, disp, colors)):
|
||||
vals = [r["op_log_summary"][k] for r in rows]
|
||||
ax.bar([xi + (i - 1) * w for xi in x], vals, width=w, label=d, color=c)
|
||||
ax.set_xticks(list(x))
|
||||
ax.set_xticklabels(labels, fontsize=9)
|
||||
ax.set_ylabel("op count")
|
||||
ax.set_title("Long-context decode 4-cases — op-count breakdown per case")
|
||||
ax.set_title("Long-context decode 6-cases — op-count breakdown per case")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
fig.tight_layout()
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_4cases_traffic.png"
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_6cases_traffic.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
return out
|
||||
@@ -103,27 +129,41 @@ def _plot_traffic(rows: list[dict]) -> Path:
|
||||
def _s_local_per_pe(panel: str, *, S_kv: int, C: int, P: int) -> int:
|
||||
"""S_local (token count) each PE attends over locally.
|
||||
|
||||
Encodes the cube/pe sharding axes from the panel name:
|
||||
cube_sp_pe_tp (Case 1): S_kv / C (pe=replicate within cube)
|
||||
cube_repl_pe_tp (Case 2): S_kv (pe=replicate; only 1 PE works)
|
||||
cube_repl_pe_sp (Case 3): S_kv / P (pe=row_wise within cube)
|
||||
cube_sp_pe_sp (Case 4): S_kv / (C·P) (★ 64-way split)
|
||||
cube_repl_pe_tp (Case 1): S_kv (no sharding, KV-wise)
|
||||
cube_sp_pe_tp (Case 2): S_kv / C (cube splits S_kv, PEs replicate)
|
||||
cube_repl_pe_sp (Case 3): S_kv / P
|
||||
cube_sp_pe_tp_dhead (Case 4): S_kv / C (cube splits S_kv, PE splits d_head)
|
||||
cube_tp_dhead_pe_sp (Case 5): S_kv / P (cube splits d_head, PE splits S_kv)
|
||||
cube_sp_pe_sp (Case 6 ★): S_kv / (C·P)
|
||||
"""
|
||||
S_per_cube = S_kv if "cube_repl" in panel else S_kv // C
|
||||
return S_per_cube // P if "pe_sp" in panel else S_per_cube
|
||||
cube_splits_s = "cube_sp" in panel
|
||||
pe_splits_s = "pe_sp" in panel
|
||||
S_per_cube = S_kv // C if cube_splits_s else S_kv
|
||||
return S_per_cube // P if pe_splits_s else S_per_cube
|
||||
|
||||
|
||||
def _d_head_per_pe(panel: str, *, d_head: int, C: int, P: int) -> int:
|
||||
"""d_head dims each PE owns (Cases 4 and 5 shard d_head)."""
|
||||
if "cube_tp_dhead" in panel: # Case 5: cube shards d_head
|
||||
return d_head // C
|
||||
if "pe_tp_dhead" in panel: # Case 4: PE shards d_head
|
||||
return d_head // P
|
||||
return d_head # Cases 1, 2, 3, 6: full d_head per PE
|
||||
|
||||
|
||||
def _active_pe_count(panel: str, *, C: int, P: int) -> int:
|
||||
"""Number of PEs doing non-idle attention work.
|
||||
|
||||
cube_sp_pe_tp (Case 1): C (PE 0 of each cube; 7 PEs idle per cube)
|
||||
cube_repl_pe_tp (Case 2): 1 (only PE 0 of CUBE 0)
|
||||
cube_repl_pe_sp (Case 3): C·P (all PEs busy, but cubes are redundant)
|
||||
cube_sp_pe_sp (Case 4): C·P (all 64 PEs doing unique work)
|
||||
cube_repl_pe_tp (Case 1): 1 (PE-TP idle for B=1; only one PE works)
|
||||
cube_sp_pe_tp (Case 2): C (PE 0 of each cube; 7 PEs idle per cube)
|
||||
cube_repl_pe_sp (Case 3): C·P (all PEs busy, cube-side redundant)
|
||||
cube_sp_pe_tp_dhead (Case 4): C·P (PE shards d_head — all 64 active)
|
||||
cube_tp_dhead_pe_sp (Case 5): C·P (PE shards S_kv — all active)
|
||||
cube_sp_pe_sp (Case 6 ★): C·P (all 64 PEs doing unique work)
|
||||
"""
|
||||
if "cube_repl" in panel and "pe_tp" in panel:
|
||||
if "cube_repl" in panel and "pe_tp" in panel and "dhead" not in panel:
|
||||
return 1
|
||||
if "cube_sp" in panel and "pe_tp" in panel:
|
||||
if "cube_sp" in panel and "pe_tp" in panel and "dhead" not in panel:
|
||||
return C
|
||||
return C * P
|
||||
|
||||
@@ -132,11 +172,12 @@ def _kv_bytes_per_pe(panel: str, *, S_kv: int, h_kv: int,
|
||||
d_head: int, C: int, P: int) -> int:
|
||||
"""KV bytes a single PE references (K + V, f16, 2 B/elem)."""
|
||||
s_local = _s_local_per_pe(panel, S_kv=S_kv, C=C, P=P)
|
||||
return 2 * s_local * h_kv * d_head * 2
|
||||
d_local = _d_head_per_pe(panel, d_head=d_head, C=C, P=P)
|
||||
return 2 * s_local * h_kv * d_local * 2
|
||||
|
||||
|
||||
def _plot_memory(rows: list[dict]) -> Path:
|
||||
"""Per-PE KV bytes — Case 4 wins (64-way split)."""
|
||||
"""Per-PE KV bytes — Case 6 ★ wins (64-way split)."""
|
||||
rows = _sorted_by_case(rows)
|
||||
labels = [_CASE_INFO[r["panel"]][0] for r in rows]
|
||||
mib_per_pe = [
|
||||
@@ -146,26 +187,25 @@ def _plot_memory(rows: list[dict]) -> Path:
|
||||
) / (1024 * 1024)
|
||||
for r in rows
|
||||
]
|
||||
colors = ["#888", "#c0504d", "#888", "#3b6ea5"] # 4 highlighted, 2 marked red
|
||||
fig, ax = plt.subplots(figsize=(8.0, 4.5))
|
||||
bars = ax.bar(labels, mib_per_pe, color=colors, width=0.6)
|
||||
fig, ax = plt.subplots(figsize=(12.0, 4.8))
|
||||
bars = ax.bar(labels, mib_per_pe, color=_bar_colors(rows), width=0.6)
|
||||
ax.set_ylabel("KV bytes per PE (MiB, K + V, f16)")
|
||||
ax.set_title(
|
||||
"Long-context decode 4-cases — KV memory per PE\n"
|
||||
"Long-context decode 6-cases — KV memory per PE\n"
|
||||
"(one KV-head group; per-layer, per-token state)"
|
||||
)
|
||||
ax.bar_label(bars, fmt="%.3f", padding=3, fontsize=9)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
ax.set_ylim(0, max(mib_per_pe) * 1.15)
|
||||
fig.tight_layout()
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_4cases_memory.png"
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_6cases_memory.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
return out
|
||||
|
||||
|
||||
def _plot_parallelism(rows: list[dict]) -> Path:
|
||||
"""Total active PE-token compute load — exposes Case 3's redundancy."""
|
||||
"""Total active PE-token compute load — exposes redundant-work cases."""
|
||||
rows = _sorted_by_case(rows)
|
||||
labels = [_CASE_INFO[r["panel"]][0] for r in rows]
|
||||
total_work = [
|
||||
@@ -173,19 +213,19 @@ def _plot_parallelism(rows: list[dict]) -> Path:
|
||||
* _s_local_per_pe(r["panel"], S_kv=r["S_kv"], C=r["C"], P=r["P"])
|
||||
for r in rows
|
||||
]
|
||||
colors = ["#888", "#888", "#c0504d", "#3b6ea5"] # 4 highlighted, 3 marked red
|
||||
fig, ax = plt.subplots(figsize=(8.0, 4.5))
|
||||
bars = ax.bar(labels, total_work, color=colors, width=0.6)
|
||||
fig, ax = plt.subplots(figsize=(12.0, 4.8))
|
||||
bars = ax.bar(labels, total_work, color=_bar_colors(rows), width=0.6)
|
||||
ax.set_ylabel("active-PE × S_local (PE-tokens; lower ⇒ less wasted work)")
|
||||
ax.set_title(
|
||||
"Long-context decode 4-cases — total compute load across active PEs\n"
|
||||
"(Case 3 replicates the full K/V across 8 cubes ⇒ 8× wasted PE-tokens)"
|
||||
"Long-context decode 6-cases — total compute load across active PEs\n"
|
||||
"(Case 3 replicates KV across 8 cubes → 8× wasted PE-tokens; "
|
||||
"Case 6 ★ is fully parallel without replication)"
|
||||
)
|
||||
ax.bar_label(bars, fmt="%d", padding=3, fontsize=9)
|
||||
ax.grid(axis="y", ls=":", alpha=0.5)
|
||||
ax.set_ylim(0, max(total_work) * 1.15)
|
||||
fig.tight_layout()
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_4cases_parallelism.png"
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_6cases_parallelism.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
return out
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
"""Comparative figure for the Case-6 composite-command decode study.
|
||||
|
||||
Reads sweep_decode_composite.json (emitted by milestone-1h-gqa, sweep
|
||||
``composite``) and writes one two-panel PNG into the bench-output dir:
|
||||
|
||||
gqa_decode_long_ctx_composite.png
|
||||
Left — end-to-end decode latency (µs) vs context length, per command
|
||||
form (primitive / composite / composite_extended).
|
||||
Right — PE_CPU command count vs context length: the hand-tiled
|
||||
primitive kernel issues O(n_tiles) commands (rises with
|
||||
context), while the coarse composite forms issue O(1) and
|
||||
*saturate* — PE_SCHEDULER absorbs the per-tile fan-out.
|
||||
|
||||
The x-axis is the global context length S_kv; each PE owns
|
||||
S_local = S_kv/(C·P=64) tokens, so at the 1M production point every PE
|
||||
runs Q·Kᵀ of (G·T_q, d_head)·(d_head, 16384) and P·V of
|
||||
(G·T_q, 16384)·(16384, d_head).
|
||||
|
||||
Run (after the bench):
|
||||
GQA_1H_RUN=1 GQA_1H_SWEEPS=composite python -m kernbench.cli.main run \\
|
||||
--bench milestone-1h-gqa --topology topology.yaml
|
||||
python scripts/paper/paper_plot_gqa_decode_long_ctx_composite.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt # noqa: E402
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
_FIG_DIR = (
|
||||
_REPO_ROOT / "src" / "kernbench" / "benches"
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_SWEEP_JSON = _FIG_DIR / "sweep_decode_composite.json"
|
||||
_PAPER_FIG_DIR = (
|
||||
_REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
|
||||
)
|
||||
|
||||
_N_RANKS = 64 # C·P for the Case-6 64-way split.
|
||||
|
||||
# variant key → (display label, colour, marker)
|
||||
_VARIANT_STYLE = {
|
||||
"primitive_tiled": ("primitive hand-tiled (16×16×16)", "#8b5a2b", "D"),
|
||||
"primitive": ("primitive (tl.dot, hand-tiled)", "#c0504d", "o"),
|
||||
"composite": ("composite GEMM", "#3b6ea5", "s"),
|
||||
"composite_extended": ("composite + softmax_merge", "#4f8a4f", "^"),
|
||||
}
|
||||
_ORDER = ("primitive_tiled", "composite", "composite_extended")
|
||||
|
||||
|
||||
def _load() -> dict:
|
||||
return json.loads(_SWEEP_JSON.read_text())
|
||||
|
||||
|
||||
def _series(rows: list[dict], variant: str, key: str):
|
||||
"""Sorted (S_kv, value) series for a variant, skipping null values
|
||||
(latency is only measured over the tractable S_kv subset)."""
|
||||
pts = sorted(
|
||||
((r["S_kv"], r[key]) for r in rows
|
||||
if r["variant"] == variant and r.get(key) is not None),
|
||||
key=lambda t: t[0],
|
||||
)
|
||||
return [p[0] for p in pts], [p[1] for p in pts]
|
||||
|
||||
|
||||
def _xticklabels(s_kvs: list[int]) -> list[str]:
|
||||
out = []
|
||||
for s in s_kvs:
|
||||
if s >= 1 << 20:
|
||||
out.append(f"{s // (1 << 20)}M")
|
||||
else:
|
||||
out.append(f"{s // 1024}K")
|
||||
return out
|
||||
|
||||
|
||||
def main() -> None:
|
||||
sweep = _load()
|
||||
rows = sweep["rows"]
|
||||
s_kv_op = sweep["s_kv_opcount"]
|
||||
s_kv_lat = sweep["s_kv_latency"]
|
||||
|
||||
fig, (ax_lat, ax_cmd) = plt.subplots(1, 2, figsize=(13.0, 4.8))
|
||||
|
||||
for v in _ORDER:
|
||||
label, color, marker = _VARIANT_STYLE[v]
|
||||
xs, lat = _series(rows, v, "latency_ns")
|
||||
ax_lat.plot(xs, [y / 1e3 for y in lat], marker=marker,
|
||||
color=color, label=label, lw=2)
|
||||
xs, cmds = _series(rows, v, "pe_cpu_cmd_count")
|
||||
ax_cmd.plot(xs, cmds, marker=marker, color=color, label=label, lw=2)
|
||||
|
||||
ax_lat.set_xticks(s_kv_lat)
|
||||
ax_lat.set_xticklabels(_xticklabels(s_kv_lat))
|
||||
ax_cmd.set_xticks(s_kv_op)
|
||||
ax_cmd.set_xticklabels(_xticklabels(s_kv_op), fontsize=8)
|
||||
for ax in (ax_lat, ax_cmd):
|
||||
ax.set_xscale("log", base=2)
|
||||
ax.set_xlabel(
|
||||
r"context length $S_{kv}$ "
|
||||
r"($S_{\mathrm{local}}=S_{kv}/64$ per PE)"
|
||||
)
|
||||
ax.grid(True, ls=":", alpha=0.5)
|
||||
ax.legend(fontsize=9)
|
||||
|
||||
ax_lat.set_ylabel("end-to-end decode latency (µs)")
|
||||
ax_lat.set_title(
|
||||
"Case-6 decode latency per command form\n"
|
||||
"(memory-bound: command form does not move the critical path)"
|
||||
)
|
||||
ax_cmd.set_ylabel("PE_CPU commands issued")
|
||||
ax_cmd.set_title(
|
||||
"PE_CPU command count per command form\n"
|
||||
"(primitive O(n$_\\mathrm{tiles}$) rises; composite O(1) saturates)"
|
||||
)
|
||||
ax_cmd.axvline(1 << 20, color="#888", ls="--", lw=1, alpha=0.7)
|
||||
ax_cmd.annotate("1M production\ncontext", xy=(1 << 20, 0),
|
||||
xytext=(1 << 18, 0.6), fontsize=8,
|
||||
textcoords=("data", "axes fraction"), ha="right",
|
||||
color="#555")
|
||||
|
||||
fig.suptitle(
|
||||
"Case-6 (Cube-SP × PE-SP) long-context decode — use of composite "
|
||||
"commands\nLLaMA-3.1-70B single-KV-head group (8 cubes × 8 PEs), "
|
||||
"$T_q{=}1$",
|
||||
fontsize=11,
|
||||
)
|
||||
fig.tight_layout(rect=(0, 0, 1, 0.94))
|
||||
|
||||
out = _FIG_DIR / "gqa_decode_long_ctx_composite.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
print(f"wrote {out}")
|
||||
|
||||
# Mirror into the paper figures dir (derived artifact).
|
||||
if _PAPER_FIG_DIR.is_dir():
|
||||
dst = _PAPER_FIG_DIR / out.name
|
||||
dst.write_bytes(out.read_bytes())
|
||||
print(f"copied {dst}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,192 @@
|
||||
"""Bar plot for the multi-model GQA composite bench.
|
||||
|
||||
Topology per model: cubes per KV group C = h_q, P = 8 PEs / cube.
|
||||
|
||||
Reads ``sweep_decode_models.json`` (produced by
|
||||
``gqa_decode_long_ctx_models.py``) and writes a two-panel PNG
|
||||
comparing end-to-end latency and PE_CPU dispatch across six GQA models
|
||||
where each model uses a topology sized to match its G value
|
||||
(cubes per KV group = h_q per KV group; PEs per cube = 8):
|
||||
|
||||
Left panel — end-to-end decode latency (µs), one bar per model, sorted
|
||||
by G. Bar labels show C, N, and latency.
|
||||
Right panel — PE_CPU command count per model.
|
||||
|
||||
Run (after the bench):
|
||||
python scripts/paper/paper_plot_gqa_decode_models.py
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt # noqa: E402
|
||||
import numpy as np # noqa: E402
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
_FIG_DIR = (
|
||||
_REPO_ROOT / "src" / "kernbench" / "benches"
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_SWEEP_JSON = _FIG_DIR / "sweep_decode_models.json"
|
||||
_PAPER_FIG_DIR = (
|
||||
_REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
|
||||
)
|
||||
|
||||
# Model display labels
|
||||
_MODEL_LABELS = {
|
||||
"gemma2-27b": "Gemma 2 27B\n(G=2)",
|
||||
"llama3-8b": "LLaMA-3 8B\n(G=4)",
|
||||
"qwen2.5-7b": "Qwen 2.5 7B\n(G=7)",
|
||||
"llama3-70b": "LLaMA-3 70B\n(G=8)",
|
||||
"qwen2.5-72b": "Qwen 2.5 72B\n(G=8)",
|
||||
"command-r-plus": "Command R+\n(G=12)",
|
||||
}
|
||||
|
||||
# Bar color per model family
|
||||
_FAMILY_COLOUR = {
|
||||
"Google": "#4f8a4f",
|
||||
"Meta": "#3b6ea5",
|
||||
"Alibaba": "#c86432",
|
||||
"Cohere": "#8b5a2b",
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
sweep = json.loads(_SWEEP_JSON.read_text())
|
||||
rows_by_model = {r["model"]: r for r in sweep["rows"]}
|
||||
# Sort models by G (ascending) so the trend is left-to-right.
|
||||
ordered = sorted(_MODEL_LABELS.keys(), key=lambda m: rows_by_model[m]["G"])
|
||||
|
||||
fig, (ax_lat, ax_cmd, ax_brk) = plt.subplots(1, 3, figsize=(18.5, 5.0))
|
||||
xs = np.arange(len(ordered))
|
||||
|
||||
# ── Left: latency ─────────────────────────────────────────────
|
||||
latencies = [rows_by_model[m]["latency_ns"] / 1e3 for m in ordered]
|
||||
colours = [_FAMILY_COLOUR[rows_by_model[m]["family"]] for m in ordered]
|
||||
bars = ax_lat.bar(xs, latencies, 0.62,
|
||||
color=colours, edgecolor="#222", linewidth=0.7)
|
||||
for i, m in enumerate(ordered):
|
||||
r = rows_by_model[m]
|
||||
# Latency label above bar
|
||||
ax_lat.text(i, latencies[i] + max(latencies) * 0.015,
|
||||
f"{latencies[i]:.1f} µs",
|
||||
ha="center", va="bottom", fontsize=9,
|
||||
fontweight="bold", color="#222")
|
||||
# Topology label inside bar
|
||||
ax_lat.text(i, latencies[i] / 2,
|
||||
f"cubes = {r['C']}",
|
||||
ha="center", va="center", fontsize=9, color="white")
|
||||
|
||||
ax_lat.set_xticks(xs)
|
||||
ax_lat.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
|
||||
ax_lat.set_ylabel("end-to-end decode latency (µs)")
|
||||
ax_lat.set_title(
|
||||
f"Composite Case-6 decode latency at $S_{{kv}}=128$K\n"
|
||||
"topology per model: cubes per KV group = h_q of the KV group",
|
||||
fontsize=11,
|
||||
)
|
||||
ax_lat.grid(True, axis="y", ls=":", alpha=0.5)
|
||||
ax_lat.set_ylim(0, max(latencies) * 1.18)
|
||||
|
||||
# Legend (family)
|
||||
from matplotlib.patches import Patch
|
||||
seen_families = []
|
||||
handles = []
|
||||
for m in ordered:
|
||||
fam = rows_by_model[m]["family"]
|
||||
if fam not in seen_families:
|
||||
seen_families.append(fam)
|
||||
handles.append(Patch(facecolor=_FAMILY_COLOUR[fam],
|
||||
edgecolor="#222", label=fam))
|
||||
ax_lat.legend(handles=handles, loc="upper left", fontsize=9,
|
||||
title="family")
|
||||
|
||||
# ── Right: PE_CPU dispatch ─────────────────────────────────────
|
||||
cmds = [rows_by_model[m]["pe_cpu_cmd_count"] for m in ordered]
|
||||
bars2 = ax_cmd.bar(xs, cmds, 0.62,
|
||||
color=colours, edgecolor="#222", linewidth=0.7)
|
||||
for i, m in enumerate(ordered):
|
||||
ax_cmd.text(i, cmds[i] + max(cmds) * 0.02,
|
||||
f"{cmds[i]}",
|
||||
ha="center", va="bottom", fontsize=9,
|
||||
fontweight="bold", color="#222")
|
||||
r = rows_by_model[m]
|
||||
ax_cmd.text(i, cmds[i] / 2,
|
||||
f"cubes = {r['C']}",
|
||||
ha="center", va="center", fontsize=9, color="white")
|
||||
|
||||
ax_cmd.set_xticks(xs)
|
||||
ax_cmd.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
|
||||
ax_cmd.set_ylabel("PE_CPU commands issued")
|
||||
ax_cmd.set_title(
|
||||
"PE_CPU dispatch (composite, per KV group)\n"
|
||||
"growing sub-mesh width or height → more reduce hops",
|
||||
fontsize=11,
|
||||
)
|
||||
ax_cmd.grid(True, axis="y", ls=":", alpha=0.5)
|
||||
ax_cmd.set_ylim(0, max(cmds) * 1.18)
|
||||
|
||||
# ── Third: matmul vs comm (op-kind occupancy) ─────────────────────
|
||||
# Only the two kinds that carry useful work — matmul (GEMM + MATH
|
||||
# occupancy summed across all engines) and comm (DMA occupancy).
|
||||
# These are op-log sums across components, not critical-path
|
||||
# attribution; parallelism means they don't sum to wall-clock
|
||||
# latency. Bars are absolute µs so the reduce-cost growth with C
|
||||
# is visible.
|
||||
matmul_us = [rows_by_model[m].get("matmul_ns", 0.0) / 1e3 for m in ordered]
|
||||
comm_us = [rows_by_model[m].get("comm_ns", 0.0) / 1e3 for m in ordered]
|
||||
max_v = max(max(matmul_us), max(comm_us))
|
||||
width = 0.35
|
||||
|
||||
ax_brk.bar(xs - width / 2, matmul_us, width,
|
||||
color="#3b6ea5", edgecolor="#222", linewidth=0.6,
|
||||
label="matmul (GEMM + MATH)")
|
||||
ax_brk.bar(xs + width / 2, comm_us, width,
|
||||
color="#c0504d", edgecolor="#222", linewidth=0.6,
|
||||
label="communication (DMA)")
|
||||
for i in range(len(ordered)):
|
||||
ax_brk.text(xs[i] - width / 2, matmul_us[i] + max_v * 0.015,
|
||||
f"{matmul_us[i]:.1f}",
|
||||
ha="center", va="bottom", fontsize=8,
|
||||
color="#222", fontweight="bold")
|
||||
ax_brk.text(xs[i] + width / 2, comm_us[i] + max_v * 0.015,
|
||||
f"{comm_us[i]:.1f}",
|
||||
ha="center", va="bottom", fontsize=8,
|
||||
color="#222", fontweight="bold")
|
||||
|
||||
ax_brk.set_xticks(xs)
|
||||
ax_brk.set_xticklabels([_MODEL_LABELS[m] for m in ordered], fontsize=9)
|
||||
ax_brk.set_ylabel("op-log occupancy (µs, summed across engines)")
|
||||
ax_brk.set_title(
|
||||
"Matmul vs communication engine work per model\n"
|
||||
"sum of GEMM/MATH and DMA occupancy — not wall-clock (overlap)",
|
||||
fontsize=11,
|
||||
)
|
||||
ax_brk.grid(True, axis="y", ls=":", alpha=0.5)
|
||||
ax_brk.set_ylim(0, max_v * 1.20)
|
||||
ax_brk.legend(loc="upper left", fontsize=9)
|
||||
|
||||
fig.suptitle(
|
||||
"Multi-model attention comparison — Case-6 composite kernel"
|
||||
" ($S_{kv}=128$K, $T_q=1$, P=8 PEs / cube)",
|
||||
fontsize=12, fontweight="bold",
|
||||
)
|
||||
fig.tight_layout(rect=(0, 0, 1, 0.93))
|
||||
|
||||
out = _FIG_DIR / "gqa_decode_models.png"
|
||||
fig.savefig(out, dpi=150)
|
||||
plt.close(fig)
|
||||
print(f"wrote {out}")
|
||||
|
||||
if _PAPER_FIG_DIR.is_dir():
|
||||
dst = _PAPER_FIG_DIR / out.name
|
||||
dst.write_bytes(out.read_bytes())
|
||||
print(f"copied {dst}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||