Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 ywkang
|
a4a2683aad |
@@ -1,243 +0,0 @@
|
||||
# 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).
|
||||
@@ -1,848 +0,0 @@
|
||||
# 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.
|
||||
Binary file not shown.
Binary file not shown.
|
Before Width: | Height: | Size: 149 KiB After Width: | Height: | Size: 136 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 73 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 169 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 136 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 72 KiB |
BIN
Binary file not shown.
|
Before Width: | Height: | Size: 96 KiB |
BIN
Binary file not shown.
|
Before Width: | Height: | Size: 133 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 70 KiB |
BIN
Binary file not shown.
|
Before Width: | Height: | Size: 73 KiB |
BIN
Binary file not shown.
|
Before Width: | Height: | Size: 75 KiB |
BIN
Binary file not shown.
|
Before Width: | Height: | Size: 84 KiB |
@@ -28,7 +28,7 @@
|
||||
\date{
|
||||
\small
|
||||
AGI Computing Lab, System Technology Group\\
|
||||
2026 Q1-Q3 Report
|
||||
2026 H1 Report
|
||||
}
|
||||
|
||||
\begin{document}
|
||||
@@ -40,10 +40,8 @@ AGI Computing Lab, System Technology Group\\
|
||||
\input{sections/03-gemm}
|
||||
\input{sections/04-allreduce}
|
||||
\input{sections/05-gqa}
|
||||
\input{sections/06-agentic}
|
||||
\input{sections/07-hw-spec-search}
|
||||
\input{sections/08-discussion}
|
||||
\input{sections/09-conclusion}
|
||||
\input{sections/10-future-work}
|
||||
\input{sections/06-discussion}
|
||||
\input{sections/07-conclusion}
|
||||
\input{sections/08-future-work}
|
||||
|
||||
\end{document}
|
||||
|
||||
@@ -111,220 +111,29 @@ redundant compute each one induces.
|
||||
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
|
||||
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
|
||||
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).
|
||||
% TODO: CUBE <-> KV-head mapping diagram for the short-context regime
|
||||
% (h_kv=8 KV heads -> 8 CUBEs, 1:1; intra-CUBE PE usage).
|
||||
% Bench code: src/kernbench/benches/gqa_helpers/short_ctx/
|
||||
|
||||
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}
|
||||
% TODO: prefill performance figure (latency, stage breakdown).
|
||||
% TODO: decode performance figure (latency, stage breakdown).
|
||||
% Bench output for short_ctx to be generated.
|
||||
|
||||
\subsection{Inference with Long-Context Length}
|
||||
\label{sec:gqa-long}
|
||||
@@ -396,17 +205,22 @@ 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.
|
||||
The measurements select the right placement for this regime. The fastest
|
||||
raw latency (Case~3, \SI{17.8}{\micro\second}) comes from replicating the
|
||||
full KV cache into every CUBE---which is exactly the placement the
|
||||
long-context memory budget forbids, and which the parallelism panel shows
|
||||
wastes 8$\times$ the compute. Restricting attention to the placements that
|
||||
fit production-context memory (the 64-way splits, Cases~4--6), the choice
|
||||
is Case~6~$\star$: it is the fastest of the three
|
||||
(\SI{30.6}{\micro\second}), and the op-count panel shows why---its
|
||||
softmax-state-only reduction charges 189 IPCQ copies and one DMA write
|
||||
against the 280 copies and 8 DMA writes the $d_{\text{head}}$-split
|
||||
Cases~4--5 pay for their partial-score all-reduce. For long-context
|
||||
decode, then, the appropriate data placement is the both-axes sequence
|
||||
shard (Case~6): it is the cheapest-communicating member of the only
|
||||
memory-feasible family, 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}
|
||||
@@ -470,7 +284,9 @@ 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 for this memory-bound decode \emph{at the 64-way production scale} (a
|
||||
single-rank caveat follows, Figure~\ref{fig:gqa-decode-stream}); 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
|
||||
@@ -481,8 +297,49 @@ 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
|
||||
\paragraph{Isolating the rank: the masked streaming win.} That neutrality
|
||||
is a property of the \emph{full 64-way} critical path, not of the local
|
||||
attention: at production scale the inter-CUBE $(m,\ell,O)$ reduce tail and
|
||||
shared-HBM contention set the wall clock, so a faster local attention does
|
||||
not surface. Stripping those away---a single rank, no cross-CUBE reduce,
|
||||
swept over the per-rank context $S_{kv}$ (so $S_{kv}{=}16$\,K here is the
|
||||
per-PE load of a 1M-token, 64-way-sharded decode)---exposes the local
|
||||
attention directly (Figure~\ref{fig:gqa-decode-stream}), and the composite
|
||||
\emph{does} win, by \SI{25}{}--\SI{28}{\percent}. The reason is the
|
||||
memory-bound mirror of prefill: its scheduler-streamed concurrent per-tile
|
||||
DMAs keep the HBM pipeline full and reach \SI{233}{\giga\byte\per\second}
|
||||
---\SI{91}{\percent} of the per-rank \SI{256}{\giga\byte\per\second}
|
||||
roofline---whereas the primitive kernel's blocking \textsf{tl.dot}
|
||||
serializes one tile DMA at a time and plateaus at
|
||||
\SI{166}{\giga\byte\per\second}. So even for memory-bound decode the
|
||||
composite is not \emph{only} a CPU-issue optimization---it also extracts
|
||||
bandwidth---but that latency benefit materializes only when the local
|
||||
attention is on the critical path, which at 64-way production scale it is
|
||||
not.
|
||||
|
||||
\begin{figure}[t]
|
||||
\centering
|
||||
\includegraphics[width=\linewidth]{gqa_decode_streaming.png}
|
||||
\caption{Single-rank memory-bound decode ($T_q{=}1$, $M{=}8$), three
|
||||
command forms, swept over per-rank context. \emph{Left:} end-to-end
|
||||
latency---the composite forms run \SI{25}{}--\SI{28}{\percent} below the
|
||||
primitive, a gap that widens with context. \emph{Right:} achieved HBM
|
||||
bandwidth against the per-rank \SI{256}{\giga\byte\per\second} roofline.
|
||||
The primitive's blocking load$\rightarrow$dot serializes the KV stream and
|
||||
plateaus at \SI{166}{\giga\byte\per\second}; the composite forms pipeline
|
||||
concurrent per-tile DMAs through the scheduler and reach
|
||||
\SI{233}{\giga\byte\per\second}. This is the memory-bound mirror of the
|
||||
prefill result (Figure~\ref{fig:gqa-prefill-cb}): there the composite
|
||||
approaches the MAC roofline, here the bandwidth roofline. Capped at
|
||||
$16$\,K---the plain composite materializes the full $(M,S_{kv})$ scores in
|
||||
TCM, so beyond that only the \textsf{softmax\_merge} recipe, which tiles
|
||||
the softmax, stays within scratch.}
|
||||
\label{fig:gqa-decode-stream}
|
||||
\end{figure}
|
||||
|
||||
\paragraph{The compute-bound mirror: prefill.} Decode's \emph{production}
|
||||
verdict---command form is latency-neutral at 64-way scale---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
|
||||
@@ -530,24 +387,24 @@ 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.
|
||||
|
||||
\subsection{Summary}
|
||||
\subsection{Comprehensive Analysis}
|
||||
\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. What this implies for hardware investment
|
||||
across the report as a whole is taken up in the discussion.
|
||||
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.
|
||||
|
||||
% TODO: cross-regime DP (data parallelism) applicability:
|
||||
% - Does Case-4 long-context placement compose with batch-level DP
|
||||
|
||||
@@ -1,309 +0,0 @@
|
||||
\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.
|
||||
@@ -1,67 +0,0 @@
|
||||
\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.
|
||||
+9
-20
@@ -1,4 +1,4 @@
|
||||
\section{Future Work}
|
||||
\section{Future Work --- 2H}
|
||||
\label{sec:future}
|
||||
|
||||
The 1H work covered attention end to end. The natural next step is to
|
||||
@@ -30,22 +30,11 @@ 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: 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.
|
||||
\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.
|
||||
@@ -45,28 +45,17 @@
|
||||
6. **Fused Grouped Query Attention** (composite + PE_IPCQ) — `sections/05-gqa.tex`
|
||||
necessity · design · results · analysis. (+ GQA seq/head/user configs)
|
||||
|
||||
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`
|
||||
7. **Discussion** — `sections/06-discussion.tex`
|
||||
Which HW changes are meaningful, and under what regimes (cross-cutting).
|
||||
|
||||
10. **Conclusion** — `sections/09-conclusion.tex`
|
||||
The codesign thesis, stated plainly, supported by the measured results.
|
||||
8. **Conclusion** — `sections/07-conclusion.tex`
|
||||
The codesign thesis, stated plainly, supported by the measured results.
|
||||
|
||||
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.
|
||||
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.
|
||||
|
||||
12. **References** *(optional)* — external literature only (FlashAttention,
|
||||
10. **References** *(optional)* — external literature only (FlashAttention,
|
||||
Megatron-LM, GPT-3, Llama 3). No ADR/SPEC entries.
|
||||
|
||||
## Per-section structure (§4/§5/§6)
|
||||
|
||||
@@ -1,12 +0,0 @@
|
||||
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
|
||||
|
@@ -1,69 +0,0 @@
|
||||
"""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()
|
||||
@@ -1,231 +0,0 @@
|
||||
"""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.")
|
||||
@@ -1,17 +0,0 @@
|
||||
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
|
||||
|
@@ -1,17 +0,0 @@
|
||||
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
|
||||
|
@@ -1,17 +0,0 @@
|
||||
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
|
||||
|
@@ -1,17 +0,0 @@
|
||||
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
|
||||
|
@@ -1,17 +0,0 @@
|
||||
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
|
||||
|
@@ -1,17 +0,0 @@
|
||||
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
|
||||
|
@@ -1,17 +0,0 @@
|
||||
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
|
||||
|
@@ -1,17 +0,0 @@
|
||||
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
|
||||
|
@@ -1,131 +0,0 @@
|
||||
"""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()
|
||||
@@ -45,12 +45,11 @@ _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", "^"),
|
||||
"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")
|
||||
_ORDER = ("primitive", "composite", "composite_extended")
|
||||
|
||||
|
||||
def _load() -> dict:
|
||||
|
||||
@@ -1,192 +0,0 @@
|
||||
"""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()
|
||||
@@ -0,0 +1,124 @@
|
||||
"""Comparative figure for the memory-bound decode-streaming composite study.
|
||||
|
||||
Reads sweep_decode_streaming.json (emitted by milestone-1h-gqa, sweep
|
||||
``decode_streaming``) and writes one two-panel PNG:
|
||||
|
||||
gqa_decode_streaming.png
|
||||
Left — end-to-end single-rank decode latency (µs) vs per-rank context.
|
||||
Right — achieved HBM bandwidth (GB/s) vs context, against the
|
||||
256 GB/s per-rank roofline.
|
||||
|
||||
The memory-bound mirror of the compute-bound prefill figure. With T_q=1
|
||||
the GEMMs are skinny (M=8) and the kernel is bound by streaming the KV
|
||||
cache. Isolating a single rank (no inter-CUBE reduce) reveals what the
|
||||
64-way Case-6 decode masks: the composite command still wins, not by
|
||||
feeding the MAC array but by keeping the DMA pipeline full — its
|
||||
scheduler-streamed concurrent tile DMAs extract ~230 GB/s (near the
|
||||
256 GB/s roofline) while the primitive kernel's blocking tl.dot serializes
|
||||
one tile DMA at a time and plateaus at ~166 GB/s. That bandwidth gap is a
|
||||
~25-28 % latency win that grows nowhere near prefill's compute-bound
|
||||
margin but is decidedly not zero.
|
||||
|
||||
Run (after the bench):
|
||||
GQA_1H_RUN=1 GQA_1H_SWEEPS=decode_streaming python -m kernbench.cli.main \\
|
||||
run --bench milestone-1h-gqa --topology topology.yaml
|
||||
python scripts/paper/paper_plot_gqa_decode_streaming.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_streaming.json"
|
||||
_PAPER_FIG_DIR = (
|
||||
_REPO_ROOT / "docs" / "report" / "1H-codesign-paper" / "figures"
|
||||
)
|
||||
|
||||
# Per-rank HBM roofline: 8 pseudo-channels × 32 GB/s (topology.yaml
|
||||
# hbm_ctrl.num_pcs / pc_bw_gbs; = pe_dma_to_noc_bw_gbs).
|
||||
_PEAK_HBM_GBS = 256.0
|
||||
|
||||
_VARIANT_STYLE = {
|
||||
"primitive": ("primitive (tl.dot, hand-tiled)", "#c0504d", "o"),
|
||||
"composite": ("composite GEMM", "#3b6ea5", "s"),
|
||||
"composite_extended": ("composite + softmax_merge", "#4f8a4f", "^"),
|
||||
}
|
||||
_ORDER = ("primitive", "composite", "composite_extended")
|
||||
|
||||
|
||||
def _ctx_label(c: int) -> str:
|
||||
return f"{c // 1024}K" if c >= 1024 else str(c)
|
||||
|
||||
|
||||
def _series(rows, variant, key):
|
||||
pts = sorted(((r["s_kv"], r[key]) for r in rows
|
||||
if r["variant"] == variant), key=lambda t: t[0])
|
||||
return [p[0] for p in pts], [p[1] for p in pts]
|
||||
|
||||
|
||||
def main() -> None:
|
||||
sweep = json.loads(_SWEEP_JSON.read_text())
|
||||
rows = sweep["rows"]
|
||||
ctxs = sweep["s_kv_points"]
|
||||
|
||||
fig, (ax_lat, ax_bw) = 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, bw = _series(rows, v, "achieved_bw_gbs")
|
||||
ax_bw.plot(xs, bw, marker=marker, color=color, label=label, lw=2)
|
||||
|
||||
for ax in (ax_lat, ax_bw):
|
||||
ax.set_xscale("log", base=2)
|
||||
ax.set_xticks(ctxs)
|
||||
ax.set_xticklabels([_ctx_label(c) for c in ctxs])
|
||||
ax.set_xlabel(r"per-rank context length $S_{kv}$ ($T_q{=}1$)")
|
||||
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("Single-rank memory-bound decode latency per command form")
|
||||
|
||||
ax_bw.set_ylabel("achieved HBM bandwidth (GB/s)")
|
||||
ax_bw.set_title(
|
||||
"HBM bandwidth — composite keeps the DMA pipe full; primitive plateaus"
|
||||
)
|
||||
ax_bw.axhline(_PEAK_HBM_GBS, color="#888", ls="--", lw=1, alpha=0.7)
|
||||
ax_bw.text(ctxs[0], _PEAK_HBM_GBS - 8, "256 GB/s roofline",
|
||||
fontsize=8, color="#555", va="top")
|
||||
ax_bw.set_ylim(0, _PEAK_HBM_GBS * 1.08)
|
||||
|
||||
fig.suptitle(
|
||||
"Memory-bound decode streaming — use of composite commands\n"
|
||||
"single-rank, GQA single-KV-head group ($h_q{=}8$, $d_{\\text{head}}"
|
||||
"{=}128$); $M{=}8$ skinny, KV-streaming-bound",
|
||||
fontsize=11,
|
||||
)
|
||||
fig.tight_layout(rect=(0, 0, 1, 0.92))
|
||||
|
||||
out = _FIG_DIR / "gqa_decode_streaming.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()
|
||||
BIN
Binary file not shown.
|
Before Width: | Height: | Size: 149 KiB After Width: | Height: | Size: 136 KiB |
BIN
Binary file not shown.
|
Before Width: | Height: | Size: 73 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 169 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 136 KiB |
+22
-22
@@ -1,7 +1,7 @@
|
||||
{
|
||||
"version": 2,
|
||||
"variants": [
|
||||
"primitive_tiled",
|
||||
"primitive",
|
||||
"composite",
|
||||
"composite_extended"
|
||||
],
|
||||
@@ -21,7 +21,7 @@
|
||||
],
|
||||
"rows": [
|
||||
{
|
||||
"variant": "primitive_tiled",
|
||||
"variant": "primitive",
|
||||
"S_kv": 8192,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
@@ -29,9 +29,9 @@
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 414,
|
||||
"pe_cpu_dispatch_cycles": 3918,
|
||||
"latency_ns": 63285.01999999998
|
||||
"pe_cpu_cmd_count": 96,
|
||||
"pe_cpu_dispatch_cycles": 930,
|
||||
"latency_ns": 30646.00300000079
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
@@ -60,7 +60,7 @@
|
||||
"latency_ns": 30483.359500000362
|
||||
},
|
||||
{
|
||||
"variant": "primitive_tiled",
|
||||
"variant": "primitive",
|
||||
"S_kv": 65536,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
@@ -68,9 +68,9 @@
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 2654,
|
||||
"pe_cpu_dispatch_cycles": 24974,
|
||||
"latency_ns": 491824.02999999997
|
||||
"pe_cpu_cmd_count": 96,
|
||||
"pe_cpu_dispatch_cycles": 930,
|
||||
"latency_ns": 231579.37900000165
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
@@ -99,7 +99,7 @@
|
||||
"latency_ns": 231211.1740000015
|
||||
},
|
||||
{
|
||||
"variant": "primitive_tiled",
|
||||
"variant": "primitive",
|
||||
"S_kv": 131072,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
@@ -107,9 +107,9 @@
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 5235,
|
||||
"pe_cpu_dispatch_cycles": 49242,
|
||||
"latency_ns": 959456.0549999998
|
||||
"pe_cpu_cmd_count": 118,
|
||||
"pe_cpu_dispatch_cycles": 1145,
|
||||
"latency_ns": 461119.51900000183
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
@@ -138,7 +138,7 @@
|
||||
"latency_ns": 460552.9805000025
|
||||
},
|
||||
{
|
||||
"variant": "primitive_tiled",
|
||||
"variant": "primitive",
|
||||
"S_kv": 262144,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
@@ -146,8 +146,8 @@
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 10397,
|
||||
"pe_cpu_dispatch_cycles": 97778,
|
||||
"pe_cpu_cmd_count": 162,
|
||||
"pe_cpu_dispatch_cycles": 1575,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
@@ -177,7 +177,7 @@
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
"variant": "primitive_tiled",
|
||||
"variant": "primitive",
|
||||
"S_kv": 524288,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
@@ -185,8 +185,8 @@
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 20721,
|
||||
"pe_cpu_dispatch_cycles": 194850,
|
||||
"pe_cpu_cmd_count": 250,
|
||||
"pe_cpu_dispatch_cycles": 2435,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
@@ -216,7 +216,7 @@
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
"variant": "primitive_tiled",
|
||||
"variant": "primitive",
|
||||
"S_kv": 1048576,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
@@ -224,8 +224,8 @@
|
||||
"d_head": 128,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"pe_cpu_cmd_count": 41369,
|
||||
"pe_cpu_dispatch_cycles": 388994,
|
||||
"pe_cpu_cmd_count": 426,
|
||||
"pe_cpu_dispatch_cycles": 4155,
|
||||
"latency_ns": null
|
||||
},
|
||||
{
|
||||
|
||||
@@ -1,173 +0,0 @@
|
||||
{
|
||||
"version": 1,
|
||||
"variant": "composite",
|
||||
"S_kv": 131072,
|
||||
"P": 8,
|
||||
"note": "C = G (G-matched topology per model)",
|
||||
"models": [
|
||||
"llama3-8b",
|
||||
"llama3-70b",
|
||||
"qwen2.5-7b",
|
||||
"qwen2.5-72b",
|
||||
"gemma2-27b",
|
||||
"command-r-plus"
|
||||
],
|
||||
"rows": [
|
||||
{
|
||||
"model": "llama3-8b",
|
||||
"family": "Meta",
|
||||
"params_b": 8,
|
||||
"h_q": 4,
|
||||
"h_kv": 1,
|
||||
"G": 4,
|
||||
"d_head": 128,
|
||||
"full_h_q": 32,
|
||||
"full_h_kv": 8,
|
||||
"hidden": 4096,
|
||||
"layers": 32,
|
||||
"C": 4,
|
||||
"P": 8,
|
||||
"N": 32,
|
||||
"T_q": 1,
|
||||
"S_kv": 131072,
|
||||
"S_local": 4096,
|
||||
"variant": "composite",
|
||||
"pe_cpu_cmd_count": 44,
|
||||
"pe_cpu_dispatch_cycles": 494,
|
||||
"latency_ns": 394539.4071250012,
|
||||
"matmul_ns": 4783.8393749997485,
|
||||
"comm_ns": 1812.7050000057789,
|
||||
"other_ns": 1684807746.5613055
|
||||
},
|
||||
{
|
||||
"model": "llama3-70b",
|
||||
"family": "Meta",
|
||||
"params_b": 70,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"G": 8,
|
||||
"d_head": 128,
|
||||
"full_h_q": 64,
|
||||
"full_h_kv": 8,
|
||||
"hidden": 8192,
|
||||
"layers": 80,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"N": 64,
|
||||
"T_q": 1,
|
||||
"S_kv": 131072,
|
||||
"S_local": 2048,
|
||||
"variant": "composite",
|
||||
"pe_cpu_cmd_count": 94,
|
||||
"pe_cpu_dispatch_cycles": 978,
|
||||
"latency_ns": 460651.3170000027,
|
||||
"matmul_ns": 10069.639499999117,
|
||||
"comm_ns": 4287.750000016997,
|
||||
"other_ns": 2079259822.423066
|
||||
},
|
||||
{
|
||||
"model": "qwen2.5-7b",
|
||||
"family": "Alibaba",
|
||||
"params_b": 7,
|
||||
"h_q": 7,
|
||||
"h_kv": 1,
|
||||
"G": 7,
|
||||
"d_head": 128,
|
||||
"full_h_q": 28,
|
||||
"full_h_kv": 4,
|
||||
"hidden": 3584,
|
||||
"layers": 28,
|
||||
"C": 7,
|
||||
"P": 8,
|
||||
"N": 56,
|
||||
"T_q": 1,
|
||||
"S_kv": 131072,
|
||||
"S_local": 2340,
|
||||
"variant": "composite",
|
||||
"pe_cpu_cmd_count": 77,
|
||||
"pe_cpu_dispatch_cycles": 813,
|
||||
"latency_ns": 456482.3358750015,
|
||||
"matmul_ns": 8809.787499998813,
|
||||
"comm_ns": 3486.4500000112457,
|
||||
"other_ns": 2032890493.459298
|
||||
},
|
||||
{
|
||||
"model": "qwen2.5-72b",
|
||||
"family": "Alibaba",
|
||||
"params_b": 72,
|
||||
"h_q": 8,
|
||||
"h_kv": 1,
|
||||
"G": 8,
|
||||
"d_head": 128,
|
||||
"full_h_q": 64,
|
||||
"full_h_kv": 8,
|
||||
"hidden": 8192,
|
||||
"layers": 80,
|
||||
"C": 8,
|
||||
"P": 8,
|
||||
"N": 64,
|
||||
"T_q": 1,
|
||||
"S_kv": 131072,
|
||||
"S_local": 2048,
|
||||
"variant": "composite",
|
||||
"pe_cpu_cmd_count": 94,
|
||||
"pe_cpu_dispatch_cycles": 978,
|
||||
"latency_ns": 460651.3170000027,
|
||||
"matmul_ns": 10069.639499999117,
|
||||
"comm_ns": 4287.750000016997,
|
||||
"other_ns": 2079259822.423066
|
||||
},
|
||||
{
|
||||
"model": "gemma2-27b",
|
||||
"family": "Google",
|
||||
"params_b": 27,
|
||||
"h_q": 2,
|
||||
"h_kv": 1,
|
||||
"G": 2,
|
||||
"d_head": 128,
|
||||
"full_h_q": 32,
|
||||
"full_h_kv": 16,
|
||||
"hidden": 4608,
|
||||
"layers": 46,
|
||||
"C": 2,
|
||||
"P": 8,
|
||||
"N": 16,
|
||||
"T_q": 1,
|
||||
"S_kv": 131072,
|
||||
"S_local": 8192,
|
||||
"variant": "composite",
|
||||
"pe_cpu_cmd_count": 60,
|
||||
"pe_cpu_dispatch_cycles": 648,
|
||||
"latency_ns": 263522.86775000195,
|
||||
"matmul_ns": 2326.0945000193315,
|
||||
"comm_ns": 907.5250000030501,
|
||||
"other_ns": 1235149167.507364
|
||||
},
|
||||
{
|
||||
"model": "command-r-plus",
|
||||
"family": "Cohere",
|
||||
"params_b": 104,
|
||||
"h_q": 12,
|
||||
"h_kv": 1,
|
||||
"G": 12,
|
||||
"d_head": 128,
|
||||
"full_h_q": 96,
|
||||
"full_h_kv": 8,
|
||||
"hidden": 12288,
|
||||
"layers": 64,
|
||||
"C": 12,
|
||||
"P": 8,
|
||||
"N": 96,
|
||||
"T_q": 1,
|
||||
"S_kv": 131072,
|
||||
"S_local": 1365,
|
||||
"variant": "composite",
|
||||
"pe_cpu_cmd_count": 111,
|
||||
"pe_cpu_dispatch_cycles": 1143,
|
||||
"latency_ns": 491928.65900000196,
|
||||
"matmul_ns": 15894.392999998527,
|
||||
"comm_ns": 9904.505000027013,
|
||||
"other_ns": 2372909985.3171477
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
{
|
||||
"version": 1,
|
||||
"variants": [
|
||||
"primitive",
|
||||
"composite",
|
||||
"composite_extended"
|
||||
],
|
||||
"s_kv_points": [
|
||||
2048,
|
||||
4096,
|
||||
8192,
|
||||
16384
|
||||
],
|
||||
"rows": [
|
||||
{
|
||||
"variant": "primitive",
|
||||
"s_kv": 2048,
|
||||
"M": 8,
|
||||
"latency_ns": 6188.437999999816,
|
||||
"gemm_busy_ns": 1048.576000000001,
|
||||
"dma_busy_ns": 6322.0,
|
||||
"kv_bytes": 1048576.0,
|
||||
"achieved_bw_gbs": 169.44114169036374,
|
||||
"achieved_tflops": 1.3555291335229098,
|
||||
"mac_util": 0.16944114169036373,
|
||||
"dma_occupancy": 1.021582505957107
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"s_kv": 2048,
|
||||
"M": 8,
|
||||
"latency_ns": 4836.115999999989,
|
||||
"gemm_busy_ns": 6629.632000000123,
|
||||
"dma_busy_ns": 267480.720000013,
|
||||
"kv_bytes": 1048576.0,
|
||||
"achieved_bw_gbs": 216.82192900253062,
|
||||
"achieved_tflops": 1.734575432020245,
|
||||
"mac_util": 0.2168219290025306,
|
||||
"dma_occupancy": 55.30899589671001
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"s_kv": 2048,
|
||||
"M": 8,
|
||||
"latency_ns": 4737.751999999986,
|
||||
"gemm_busy_ns": 8481.408000000116,
|
||||
"dma_busy_ns": 251398.7400000122,
|
||||
"kv_bytes": 1048576.0,
|
||||
"achieved_bw_gbs": 221.32353065335693,
|
||||
"achieved_tflops": 1.7705882452268553,
|
||||
"mac_util": 0.22132353065335691,
|
||||
"dma_occupancy": 53.06287454472352
|
||||
},
|
||||
{
|
||||
"variant": "primitive",
|
||||
"s_kv": 4096,
|
||||
"M": 8,
|
||||
"latency_ns": 12510.006000000496,
|
||||
"gemm_busy_ns": 2097.152000000002,
|
||||
"dma_busy_ns": 12602.0,
|
||||
"kv_bytes": 2097152.0,
|
||||
"achieved_bw_gbs": 167.6379691584414,
|
||||
"achieved_tflops": 1.3411037532675312,
|
||||
"mac_util": 0.1676379691584414,
|
||||
"dma_occupancy": 1.0073536335633653
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"s_kv": 4096,
|
||||
"M": 8,
|
||||
"latency_ns": 9360.883999999554,
|
||||
"gemm_busy_ns": 19520.799999939867,
|
||||
"dma_busy_ns": 1059043.5999999335,
|
||||
"kv_bytes": 2097152.0,
|
||||
"achieved_bw_gbs": 224.03354213128802,
|
||||
"achieved_tflops": 1.792268337050304,
|
||||
"mac_util": 0.224033542131288,
|
||||
"dma_occupancy": 113.13499878857424
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"s_kv": 4096,
|
||||
"M": 8,
|
||||
"latency_ns": 9198.135999999562,
|
||||
"gemm_busy_ns": 39531.583999941715,
|
||||
"dma_busy_ns": 1026582.7399999356,
|
||||
"kv_bytes": 2097152.0,
|
||||
"achieved_bw_gbs": 227.99749862364504,
|
||||
"achieved_tflops": 1.8239799889891604,
|
||||
"mac_util": 0.22799749862364505,
|
||||
"dma_occupancy": 111.60769312390951
|
||||
},
|
||||
{
|
||||
"variant": "primitive",
|
||||
"s_kv": 8192,
|
||||
"M": 8,
|
||||
"latency_ns": 25153.141999997388,
|
||||
"gemm_busy_ns": 4194.304000000004,
|
||||
"dma_busy_ns": 25162.0,
|
||||
"kv_bytes": 4194304.0,
|
||||
"achieved_bw_gbs": 166.7506985807354,
|
||||
"achieved_tflops": 1.334005588645883,
|
||||
"mac_util": 0.16675069858073538,
|
||||
"dma_occupancy": 1.0003521627637062
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"s_kv": 8192,
|
||||
"M": 8,
|
||||
"latency_ns": 18419.187999999034,
|
||||
"gemm_busy_ns": 64177.11999976149,
|
||||
"dma_busy_ns": 4214541.840000685,
|
||||
"kv_bytes": 4194304.0,
|
||||
"achieved_bw_gbs": 227.713838416776,
|
||||
"achieved_tflops": 1.821710707334208,
|
||||
"mac_util": 0.227713838416776,
|
||||
"dma_occupancy": 228.81257523409317
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"s_kv": 8192,
|
||||
"M": 8,
|
||||
"latency_ns": 18128.439999999027,
|
||||
"gemm_busy_ns": 169650.68799976518,
|
||||
"dma_busy_ns": 4149323.2200006745,
|
||||
"kv_bytes": 4194304.0,
|
||||
"achieved_bw_gbs": 231.365964197704,
|
||||
"achieved_tflops": 1.850927713581632,
|
||||
"mac_util": 0.231365964197704,
|
||||
"dma_occupancy": 228.88473691067168
|
||||
},
|
||||
{
|
||||
"variant": "primitive",
|
||||
"s_kv": 16384,
|
||||
"M": 8,
|
||||
"latency_ns": 50439.414000008954,
|
||||
"gemm_busy_ns": 8388.607999999076,
|
||||
"dma_busy_ns": 50282.0,
|
||||
"kv_bytes": 8388608.0,
|
||||
"achieved_bw_gbs": 166.31057609032712,
|
||||
"achieved_tflops": 1.330484608722617,
|
||||
"mac_util": 0.16631057609032712,
|
||||
"dma_occupancy": 0.9968791469304357
|
||||
},
|
||||
{
|
||||
"variant": "composite",
|
||||
"s_kv": 16384,
|
||||
"M": 8,
|
||||
"latency_ns": 36535.79600000568,
|
||||
"gemm_busy_ns": 228978.46400288073,
|
||||
"dma_busy_ns": 16815028.239995122,
|
||||
"kv_bytes": 8388608.0,
|
||||
"achieved_bw_gbs": 229.59970545047648,
|
||||
"achieved_tflops": 1.8367976436038118,
|
||||
"mac_util": 0.22959970545047648,
|
||||
"dma_occupancy": 460.2343477063565
|
||||
},
|
||||
{
|
||||
"variant": "composite_extended",
|
||||
"s_kv": 16384,
|
||||
"M": 8,
|
||||
"latency_ns": 35989.048000005685,
|
||||
"gemm_busy_ns": 701990.3680028584,
|
||||
"dma_busy_ns": 16684294.09999516,
|
||||
"kv_bytes": 8388608.0,
|
||||
"achieved_bw_gbs": 233.0877993771515,
|
||||
"achieved_tflops": 1.864702395017212,
|
||||
"mac_util": 0.2330877993771515,
|
||||
"dma_occupancy": 463.5936493789034
|
||||
}
|
||||
]
|
||||
}
|
||||
+4
-4
@@ -37,9 +37,9 @@ Topology / SFR:
|
||||
from __future__ import annotations
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
|
||||
_ROOT_CUBE,
|
||||
_merge_running,
|
||||
reduce_mlo,
|
||||
root_cube_for,
|
||||
)
|
||||
|
||||
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
|
||||
@@ -108,9 +108,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel(
|
||||
tl.copy_to(O_local, O_new)
|
||||
|
||||
# ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
|
||||
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
|
||||
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
|
||||
|
||||
# ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
|
||||
if pe_id == 0 and cube_id == root_cube_for(C):
|
||||
# ── Final normalise + store (root only: PE 0 of CUBE 6) ──
|
||||
if pe_id == 0 and cube_id == _ROOT_CUBE:
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr, O_final)
|
||||
|
||||
+4
-4
@@ -18,8 +18,8 @@ under the ADR-0064 Rev2 structural model (the CPU-offload win).
|
||||
from __future__ import annotations
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
|
||||
_ROOT_CUBE,
|
||||
reduce_mlo,
|
||||
root_cube_for,
|
||||
)
|
||||
|
||||
|
||||
@@ -63,9 +63,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel(
|
||||
tl.composite(op="gemm", a=exp_scores, b=V, out=O_local) # P·V, one command
|
||||
|
||||
# ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
|
||||
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
|
||||
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
|
||||
|
||||
# ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
|
||||
if pe_id == 0 and cube_id == root_cube_for(C):
|
||||
# ── Final normalise + store (root only: PE 0 of CUBE 6) ──
|
||||
if pe_id == 0 and cube_id == _ROOT_CUBE:
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr, O_final)
|
||||
|
||||
+4
-4
@@ -27,8 +27,8 @@ numeric parity of the recipe's 8 MATH ops is a separate follow-up
|
||||
from __future__ import annotations
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
|
||||
_ROOT_CUBE,
|
||||
reduce_mlo,
|
||||
root_cube_for,
|
||||
)
|
||||
|
||||
# Small primitive seed slice that establishes the running (m, ℓ, O) the
|
||||
@@ -86,9 +86,9 @@ def gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel(
|
||||
)
|
||||
|
||||
# ── Two-level (m, ℓ, O) reduce-to-root (shared helper) ──
|
||||
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
|
||||
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
|
||||
|
||||
# ── Final normalise + store (root only: PE 0 of the C-derived root cube) ──
|
||||
if pe_id == 0 and cube_id == root_cube_for(C):
|
||||
# ── Final normalise + store (root only: PE 0 of CUBE 6) ──
|
||||
if pe_id == 0 and cube_id == _ROOT_CUBE:
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr, O_final)
|
||||
|
||||
-215
@@ -1,215 +0,0 @@
|
||||
"""GQA decode kernel — Case 6, **primitive hand-tiled 16×16×16** (per-block DMA + MAC-side accumulate).
|
||||
|
||||
Same Case-6 placement and (m, ℓ, O) reduce as the primitive baseline
|
||||
(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the difference is
|
||||
that each local-attention matmul is *hand-blocked into 16×16×16 GEMMs*
|
||||
(mac=16), **with each block re-fetching its operand slices from HBM**
|
||||
(no operand cache), and each (mi, ni) output tile updated **implicitly**
|
||||
by successive K-inner ``GemmCmd`` writes into the shared (M, N) output —
|
||||
mirroring a MAC-array accumulator register that latches across K-inner
|
||||
cycles on real accelerators.
|
||||
|
||||
Per K-inner block: two ``tl.load`` DMAs (Q and K slice for Q·Kᵀ, or one
|
||||
V slice for P·V) → one ``GemmCmd`` emitted directly via ``tl._emit`` with
|
||||
``out`` bound to the shared (M, N) handle and ``m/k/n`` overridden to
|
||||
the 16³ block dims. All K-inner iterations at a given (mi, ni) write
|
||||
into the same output tile; no compiler-emitted ``MathCmd`` accumulator
|
||||
chain.
|
||||
|
||||
This models a strict-streaming architecture (no HBM-operand cache) — the
|
||||
worst-case dispatch endpoint: every 16³ compute block pays the ADR-0064
|
||||
D8 single-op FIXED overhead on DMA (per operand slice) AND on GEMM,
|
||||
exposing the full PE_CPU dispatch pressure that composite forms
|
||||
(ADR-0065) absorb into PE_SCHEDULER.
|
||||
|
||||
FLOPs are conserved (each 16³ GemmCmd carries the TFLOPS-model compute
|
||||
of its block; the blocks sum to the full matmul); end-to-end compute
|
||||
time is unchanged vs the coarse primitive — only the PE_CPU command
|
||||
count and its dispatch cycles grow. Inputs are zero (decode bench
|
||||
convention), so the "overwrite instead of accumulate" is semantically
|
||||
equivalent to sum (0 = 0 + 0), and ``GemmCmd`` writes populate the
|
||||
shared ``out`` handle so the downstream softmax's strict ``MathCmd``
|
||||
reads succeed — the kernel runs in engine (data) mode.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from math import ceil
|
||||
|
||||
from kernbench.common.pe_commands import GemmCmd
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
|
||||
_merge_running,
|
||||
reduce_mlo,
|
||||
root_cube_for,
|
||||
)
|
||||
|
||||
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
|
||||
MAC = 16 # 16×16×16 MAC-array blocking granularity.
|
||||
|
||||
|
||||
def _blocked_dot_qk_streamed(a_ptr, b_ptr, M, K, N, *, tl):
|
||||
"""Q·Kᵀ hand-blocked, per-block DMA of both operands, MAC-side accumulate.
|
||||
|
||||
For each 16³ block (mi, ni, ki): load 16×16 A and B slices from HBM,
|
||||
emit a ``GemmCmd`` directly with ``out`` bound to the shared (M, N)
|
||||
handle and ``m/k/n`` overridden to the 16³ block dims. All K-inner
|
||||
iterations at a given (mi, ni) write into the same shared ``out`` tile
|
||||
— accumulation is implicit (MAC-side latching, mirroring how real
|
||||
accelerators feed a per-tile accumulator register instead of running
|
||||
a compiler-emitted MathCmd add chain).
|
||||
|
||||
Under the zero-input decode-bench convention, the "overwrite instead
|
||||
of accumulate" is semantically equivalent to sum (0 = 0 + 0), and
|
||||
``GemmCmd`` writes populate ``out.addr`` in ``MemoryStore`` so the
|
||||
downstream softmax's strict ``MathCmd`` reads succeed. Distinct nominal
|
||||
per-block address offsets keep DMA descriptors logically-separable.
|
||||
"""
|
||||
# Full-shape operand handles for data-mode correctness. Under the
|
||||
# simulator's DataExecutor, GemmCmd computes np.matmul over these
|
||||
# shapes and writes an (M, N) result to out.addr — populating the
|
||||
# shared handle with a valid region so the downstream softmax's
|
||||
# strict MathCmd reads succeed. In engine timing, the m/k/n
|
||||
# override on each GemmCmd charges only the 16³ block work.
|
||||
A_full = tl.load(a_ptr, shape=(M, K), dtype="f16")
|
||||
B_full = tl.load(b_ptr, shape=(K, N), dtype="f16")
|
||||
out = tl._make_compute_out(shape=(M, N), dtype="f16")
|
||||
n_m = ceil(M / MAC)
|
||||
n_k = ceil(K / MAC)
|
||||
n_n = ceil(N / MAC)
|
||||
for mi in range(n_m):
|
||||
bm = min(MAC, M - mi * MAC)
|
||||
for ni in range(n_n):
|
||||
bn = min(MAC, N - ni * MAC)
|
||||
# Recycle per (mi, ni) — the per-block slice DMAs live only
|
||||
# during this tile's processing; the shared out survives.
|
||||
with tl.scratch_scope():
|
||||
for ki in range(n_k):
|
||||
bk = min(MAC, K - ki * MAC)
|
||||
# Per-block DMAs model the streaming architecture
|
||||
# (no HBM-operand cache). Row-major byte offsets:
|
||||
# A is (M, K), B is (K, N). Nominal handles — the
|
||||
# GemmCmd below uses the full-shape handles above
|
||||
# for data-mode compute; these per-block loads pay
|
||||
# their ADR-0064 dispatch cost as descriptor work.
|
||||
_ = tl.load(
|
||||
a_ptr + mi * MAC * K * 2 + ki * MAC * 2,
|
||||
shape=(bm, bk), dtype="f16",
|
||||
)
|
||||
_ = tl.load(
|
||||
b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
|
||||
shape=(bk, bn), dtype="f16",
|
||||
)
|
||||
tl._emit(GemmCmd(
|
||||
a=A_full, b=B_full, out=out,
|
||||
m=bm, k=bk, n=bn,
|
||||
))
|
||||
return out
|
||||
|
||||
|
||||
def _blocked_dot_pv_streamed(A_handle, b_ptr, M, K, N, *, tl):
|
||||
"""P·V hand-blocked; only V streams (P is TCM-resident post-softmax).
|
||||
|
||||
Same structure as the Q·Kᵀ variant — per block: 1 DMA (V slice) + 1
|
||||
``GemmCmd`` writing to the shared (M, N) ``out`` handle with implicit
|
||||
K-inner accumulation via successive block writes.
|
||||
|
||||
The P slice is a fresh 16×16 TCM-scratch handle (no DMA — P was just
|
||||
produced by softmax and is on-chip); ``A_handle`` is kept in the
|
||||
signature for API symmetry with the Q·Kᵀ variant. Zero-input
|
||||
convention applies throughout.
|
||||
"""
|
||||
# A_handle (= exp_scores from softmax) is TCM-resident with
|
||||
# shape (M, K), already populated by tl.exp's DataExecutor. Only V
|
||||
# needs a full-shape coarse load from HBM for data-mode correctness.
|
||||
B_full = tl.load(b_ptr, shape=(K, N), dtype="f16")
|
||||
out = tl._make_compute_out(shape=(M, N), dtype="f16")
|
||||
n_m = ceil(M / MAC)
|
||||
n_k = ceil(K / MAC)
|
||||
n_n = ceil(N / MAC)
|
||||
for _mi in range(n_m):
|
||||
for ni in range(n_n):
|
||||
bn = min(MAC, N - ni * MAC)
|
||||
with tl.scratch_scope():
|
||||
for ki in range(n_k):
|
||||
bk = min(MAC, K - ki * MAC)
|
||||
# Per-block V load — streaming-architecture dispatch
|
||||
# cost. GemmCmd below uses B_full for compute.
|
||||
_ = tl.load(
|
||||
b_ptr + ki * MAC * N * 2 + ni * MAC * 2,
|
||||
shape=(bk, bn), dtype="f16",
|
||||
)
|
||||
bm = min(MAC, M - _mi * MAC)
|
||||
tl._emit(GemmCmd(
|
||||
a=A_handle, b=B_full, out=out,
|
||||
m=bm, k=bk, n=bn,
|
||||
))
|
||||
return out
|
||||
|
||||
|
||||
def gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
P: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Case-6 decode, primitive-TILED streamed (per-block DMA + 16³ GEMM)."""
|
||||
G = h_q // h_kv
|
||||
n_ranks = C * P
|
||||
S_local = S_kv // n_ranks
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
|
||||
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
|
||||
KV_ROW_BYTES = d_head * 2 # f16
|
||||
|
||||
# Bootstrap tile (tile 0). Establishes persistent (m_local, l_local,
|
||||
# O_local). Cannot fold into Tiles 1..N loop: persistent tensors must
|
||||
# live outside tl.scratch_scope or scope teardown discards them.
|
||||
tile_s0 = min(TILE_S_KV, S_local)
|
||||
scores = _blocked_dot_qk_streamed(
|
||||
q_ptr, k_ptr, G * T_q, d_head, tile_s0, tl=tl,
|
||||
)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
exp_scores = tl.exp(centered)
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
O_local = _blocked_dot_pv_streamed(
|
||||
exp_scores, v_ptr, G * T_q, tile_s0, d_head, tl=tl,
|
||||
)
|
||||
|
||||
for tile_idx in range(1, n_tiles):
|
||||
tile_start = tile_idx * TILE_S_KV
|
||||
tile_s = min(TILE_S_KV, S_local - tile_start)
|
||||
with tl.scratch_scope():
|
||||
scores_t = _blocked_dot_qk_streamed(
|
||||
q_ptr, k_ptr + tile_start * KV_ROW_BYTES,
|
||||
G * T_q, d_head, tile_s, tl=tl,
|
||||
)
|
||||
m_tile = tl.max(scores_t, axis=-1)
|
||||
centered_t = scores_t - m_tile
|
||||
exp_scores_t = tl.exp(centered_t)
|
||||
l_tile = tl.sum(exp_scores_t, axis=-1)
|
||||
O_tile = _blocked_dot_pv_streamed(
|
||||
exp_scores_t, v_ptr + tile_start * KV_ROW_BYTES,
|
||||
G * T_q, tile_s, d_head, tl=tl,
|
||||
)
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
|
||||
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, tl=tl)
|
||||
|
||||
if pe_id == 0 and cube_id == root_cube_for(C):
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr, O_final)
|
||||
+129
@@ -0,0 +1,129 @@
|
||||
"""GQA decode kernel — Case 6, **primitive-TILED** (16×16×16 MAC blocking).
|
||||
|
||||
Same Case-6 placement and (m, ℓ, O) reduce as the primitive baseline
|
||||
(``_gqa_attention_decode_long_ctx_cube_sp_pe_sp``); the only difference is
|
||||
that each local-attention matmul is *hand-blocked into 16×16×16 GemmCmds*
|
||||
(mac=16) instead of one coarse ``tl.dot`` per tile. This models a kernel
|
||||
that issues the MAC-array fan-out from PE_CPU itself: the per-block GemmCmd
|
||||
count is ``ceil(M/16)·ceil(K/16)·ceil(N/16)`` per matmul, charging the full
|
||||
ADR-0064 dispatch cost for every block — the "dispatch explosion" the
|
||||
composite form offloads to PE_SCHEDULER.
|
||||
|
||||
FLOPs are conserved (each 16³ GemmCmd carries the TFLOPS-model compute of
|
||||
its block; the blocks sum to the full matmul), so end-to-end compute time
|
||||
is unchanged vs the coarse primitive — only the PE_CPU command count and
|
||||
its dispatch cycles grow. Inputs are zero (decode bench convention), so the
|
||||
blocked accumulation is identically zero; the kernel returns a single
|
||||
zeroed ``(M, N)`` output handle that the downstream softmax consumes — no
|
||||
per-block accumulation handle needed for this zero-input study.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from math import ceil
|
||||
|
||||
from kernbench.common.pe_commands import GemmCmd
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_mlo_reduce import (
|
||||
_ROOT_CUBE,
|
||||
_merge_running,
|
||||
reduce_mlo,
|
||||
)
|
||||
|
||||
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
|
||||
MAC = 16 # 16×16×16 MAC-array blocking granularity.
|
||||
|
||||
|
||||
def _blocked_dot(A, B, *, tl):
|
||||
"""``A @ B`` issued as one GemmCmd per 16×16×16 block.
|
||||
|
||||
``A``:(M, K), ``B``:(K, N) → out:(M, N). Emits
|
||||
``ceil(M/16)·ceil(K/16)·ceil(N/16)`` GemmCmds (each charged the full
|
||||
ADR-0064 dispatch cost via ``tl.dot``'s emit path). **All blocks write
|
||||
the same ``(M, N)`` output handle** (``out``), and that handle is
|
||||
returned — so downstream softmax ops depend on it exactly like the
|
||||
coarse ``tl.dot`` path (the engine tracks the producer by output handle
|
||||
id, and ``GemmCmd`` is a blocking PE_CPU command, so the block GEMMs
|
||||
serialize on the critical path ahead of the consuming ``tl.max``/
|
||||
``tl.exp``). K is innermost so each (mi, ni) output tile accumulates
|
||||
across the K blocks into the same handle.
|
||||
"""
|
||||
M, K = A.shape[-2], A.shape[-1]
|
||||
K2, N = B.shape[-2], B.shape[-1]
|
||||
assert K == K2, f"blocked_dot shape mismatch K={K} != {K2}"
|
||||
out = tl._make_compute_out(shape=(M, N), dtype="f16")
|
||||
# One GemmCmd per 16³ block, all writing the shared `out` handle. A/B
|
||||
# reuse the full operand handles at block dims (m,k,n = 16, or the
|
||||
# ragged tail) — operands are TCM-resident (pinned), so this charges
|
||||
# dispatch + the block's TFLOPS-model compute. K innermost = accumulate
|
||||
# the K blocks into each (mi, ni) output tile of the shared handle.
|
||||
for _mi in range(ceil(M / MAC)):
|
||||
bm = min(MAC, M - _mi * MAC)
|
||||
for _ni in range(ceil(N / MAC)):
|
||||
bn = min(MAC, N - _ni * MAC)
|
||||
for _ki in range(ceil(K / MAC)):
|
||||
bk = min(MAC, K - _ki * MAC)
|
||||
tl._emit(GemmCmd(a=A, b=B, out=out, m=bm, k=bk, n=bn))
|
||||
return out
|
||||
|
||||
|
||||
def gqa_attention_decode_long_ctx_cube_sp_pe_sp_tiled_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
P: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Case-6 decode, primitive-TILED (16×16×16 GemmCmd blocking)."""
|
||||
G = h_q // h_kv
|
||||
n_ranks = C * P
|
||||
S_local = S_kv // n_ranks
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
|
||||
Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
|
||||
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
|
||||
KV_ROW_BYTES = d_head * 2 # f16
|
||||
|
||||
tile_s0 = min(TILE_S_KV, S_local)
|
||||
K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
|
||||
V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
|
||||
scores = _blocked_dot(Q, K_T, tl=tl)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
exp_scores = tl.exp(centered)
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
O_local = _blocked_dot(exp_scores, V, tl=tl)
|
||||
|
||||
for tile_idx in range(1, n_tiles):
|
||||
tile_start = tile_idx * TILE_S_KV
|
||||
tile_s = min(TILE_S_KV, S_local - tile_start)
|
||||
with tl.scratch_scope():
|
||||
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(tile_s, d_head), dtype="f16")
|
||||
scores_t = _blocked_dot(Q, K_T_t, tl=tl)
|
||||
m_tile = tl.max(scores_t, axis=-1)
|
||||
centered_t = scores_t - m_tile
|
||||
exp_scores_t = tl.exp(centered_t)
|
||||
l_tile = tl.sum(exp_scores_t, axis=-1)
|
||||
O_tile = _blocked_dot(exp_scores_t, V_t, tl=tl)
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
|
||||
reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, tl=tl)
|
||||
|
||||
if pe_id == 0 and cube_id == _ROOT_CUBE:
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr, O_final)
|
||||
@@ -1,66 +1,33 @@
|
||||
"""Shared (m, ℓ, O) reduce for the Case-6 long-context decode kernels.
|
||||
|
||||
Extracted from the Cube-SP × PE-SP decode kernel so the four command-form
|
||||
variants (primitive / primitive-tiled / composite / composite_extended)
|
||||
Extracted from the Cube-SP × PE-SP decode kernel so the three
|
||||
command-form variants (primitive / composite / composite_extended)
|
||||
share one identical reduce. The local attention differs per variant;
|
||||
the reduce does not. Byte-equal behavior at C=8 is guarded by
|
||||
``tests/attention/test_gqa_decode_long_ctx_composite.py``.
|
||||
the reduce does not. Behavior is byte-equal to the pre-extraction inline
|
||||
reduce (guarded by tests/attention/test_gqa_decode_long_ctx_composite.py).
|
||||
|
||||
The reduce is two-level:
|
||||
• intra-CUBE 8-way (row chain along intra_W + col bridge along
|
||||
intra_N) over the 2×4 PE grid;
|
||||
• inter-CUBE 2-phase lrab-adapted center-root reduce over a
|
||||
``sub_w × sub_h`` CUBE sub-mesh (ADR-0060 §4.2): Phase 1 does a
|
||||
bidirectional row reduce converging at ``root_col = sub_w // 2``;
|
||||
Phase 2 does a bidirectional col reduce on ``root_col`` converging
|
||||
at ``root_row = sub_h // 2``.
|
||||
|
||||
The submesh dimensions ``(sub_w, sub_h)`` and root are **derived from
|
||||
C** (number of cubes launched per KV group), not hardcoded, so the
|
||||
same reduce works for any ``C ≥ 1``. Peer-existence guards on every
|
||||
send/recv elide operations that would target un-launched neighbors
|
||||
(needed for non-rectangular C like 7 or 12; safe no-ops when the
|
||||
neighbor exists). For ``C = 8`` this reproduces the previous
|
||||
hardcoded ``{sub_w=4, sub_h=2, root_col=2, root_row=1, root_cube=6}``
|
||||
exactly.
|
||||
|
||||
Plain ``+`` in the tree is replaced by log-sum-exp ``_merge_running``.
|
||||
• inter-CUBE 2-phase lrab over the 4×2 CUBE sub-mesh (ADR-0060 §4.2
|
||||
lrab-adapted center-root reduce): Phase 1 row reduce converges at
|
||||
root_col=2; Phase 2 col reduce on root_col converges at root_row=1;
|
||||
root cube id = 6.
|
||||
Plain ``+`` is replaced by log-sum-exp ``_merge_running``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
# lrab geometry for the C=8 single-KV-head group (4×2 cube sub-mesh).
|
||||
_SUB_W = 4
|
||||
_SUB_H = 2
|
||||
_ROOT_COL = _SUB_W // 2 # 2
|
||||
_ROOT_ROW = _SUB_H // 2 # 1
|
||||
_ROOT_CUBE = _ROOT_ROW * _SUB_W + _ROOT_COL # 6
|
||||
|
||||
PE_GRID_COLS = 4
|
||||
|
||||
|
||||
def _submesh_for(C):
|
||||
"""Derive ``(sub_w, sub_h, root_col, root_row, root_cube)`` for a
|
||||
C-cube inter-cube reduce.
|
||||
|
||||
Submesh is at most 4 columns wide (matching the SIP's 4×4 cube
|
||||
mesh); rows fill left-to-right, top-to-bottom. For the Case-6
|
||||
baseline C=8, this returns ``(4, 2, 2, 1, 6)`` — identical to
|
||||
the previously hardcoded values.
|
||||
"""
|
||||
sub_w = min(C, 4)
|
||||
sub_h = (C + sub_w - 1) // sub_w
|
||||
root_col = sub_w // 2
|
||||
root_row = sub_h // 2
|
||||
root_cube = root_row * sub_w + root_col
|
||||
return sub_w, sub_h, root_col, root_row, root_cube
|
||||
|
||||
|
||||
def root_cube_for(C):
|
||||
"""Root cube id for a C-cube reduce. Callers use this to gate the
|
||||
final ``tl.store`` in the kernel epilogue."""
|
||||
_, _, _, _, root_cube = _submesh_for(C)
|
||||
return root_cube
|
||||
|
||||
|
||||
# Backward-compat alias: preserves the Case-6 baseline value (6) for
|
||||
# any external code that still imports the module-level constant.
|
||||
_ROOT_CUBE = root_cube_for(8)
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
|
||||
m_new = tl.maximum(m_local, m_other)
|
||||
@@ -71,23 +38,13 @@ def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, *, tl):
|
||||
"""Two-level (m, ℓ, O) reduce-to-root (PE 0 of ``root_cube_for(C)``).
|
||||
In place.
|
||||
def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, P, *, tl):
|
||||
"""Two-level (m, ℓ, O) reduce-to-root (PE 0 of CUBE 6). In place.
|
||||
|
||||
Updates ``m_local`` / ``l_local`` / ``O_local`` in place via
|
||||
``copy_to``; after this call only PE 0 of ``root_cube_for(C)``
|
||||
holds the fully-merged result.
|
||||
|
||||
Peer-existence guards on every inter-cube send/recv ensure the
|
||||
reduce completes for any launched C (rectangular or not); a
|
||||
send/recv toward a non-launched neighbor is skipped. For C = 8
|
||||
every neighbor exists → guards are all True → byte-equal to the
|
||||
previously hardcoded Case-6 reduce.
|
||||
Updates ``m_local``/``l_local``/``O_local`` in place via ``copy_to``;
|
||||
after this call only PE 0 of CUBE ``_ROOT_CUBE`` holds the full result.
|
||||
"""
|
||||
sub_w, sub_h, root_col, root_row, _root_cube = _submesh_for(C)
|
||||
|
||||
# ── Intra-CUBE reduce (unchanged — depends only on P) ─────────────
|
||||
# ── Intra-CUBE reduce: row chain (intra_W) + col bridge (intra_N) ──
|
||||
pe_col = pe_id % PE_GRID_COLS
|
||||
pe_row = pe_id // PE_GRID_COLS
|
||||
pe_cols_used = min(PE_GRID_COLS, P)
|
||||
@@ -127,126 +84,106 @@ def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, *, tl):
|
||||
tl.send(dir="intra_N", src=l_local)
|
||||
tl.send(dir="intra_N", src=O_local)
|
||||
|
||||
# ── Inter-CUBE lrab-adapted center-root reduce ────────────────────
|
||||
# ── Inter-CUBE lrab-adapted center-root reduce (ADR-0060 §4.2) ──
|
||||
# Only PE 0 of each CUBE participates. Adapts Phases 1-2 of
|
||||
# lrab_hierarchical_allreduce.py: bidirectional row reduce converges
|
||||
# at ``root_col``; bidirectional col reduce on ``root_col`` converges
|
||||
# at ``root_row``. Plain ``+`` replaced by log-sum-exp
|
||||
# ``_merge_running``. Peer-existence guards elide sends/recvs whose
|
||||
# neighbor cube was not launched.
|
||||
# at root_col; bidirectional col reduce on root_col converges at
|
||||
# root_row. Plain ``+`` replaced by log-sum-exp ``_merge_running``.
|
||||
if pe_id == 0:
|
||||
row = cube_id // sub_w
|
||||
col = cube_id % sub_w
|
||||
row = cube_id // _SUB_W
|
||||
col = cube_id % _SUB_W
|
||||
|
||||
# Peer existence within the launched set of C cubes.
|
||||
east_exists = (col + 1 < sub_w) and (cube_id + 1 < C)
|
||||
west_exists = col > 0
|
||||
south_exists = cube_id + sub_w < C
|
||||
north_exists = row > 0
|
||||
|
||||
# Phase 1: row reduce — converge at col == root_col.
|
||||
# Phase 1: row reduce — converge at col == _ROOT_COL.
|
||||
if col == 0:
|
||||
if east_exists:
|
||||
tl.send(dir="E", src=m_local)
|
||||
tl.send(dir="E", src=l_local)
|
||||
tl.send(dir="E", src=O_local)
|
||||
elif 0 < col < root_col:
|
||||
if west_exists:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if east_exists:
|
||||
tl.send(dir="E", src=m_local)
|
||||
tl.send(dir="E", src=l_local)
|
||||
tl.send(dir="E", src=O_local)
|
||||
elif col == root_col:
|
||||
if west_exists:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if east_exists:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
elif root_col < col < sub_w - 1:
|
||||
if east_exists:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if west_exists:
|
||||
tl.send(dir="W", src=m_local)
|
||||
tl.send(dir="W", src=l_local)
|
||||
tl.send(dir="W", src=O_local)
|
||||
elif col == sub_w - 1:
|
||||
if west_exists:
|
||||
tl.send(dir="W", src=m_local)
|
||||
tl.send(dir="W", src=l_local)
|
||||
tl.send(dir="W", src=O_local)
|
||||
tl.send(dir="E", src=m_local)
|
||||
tl.send(dir="E", src=l_local)
|
||||
tl.send(dir="E", src=O_local)
|
||||
elif 0 < col < _ROOT_COL:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
tl.send(dir="E", src=m_local)
|
||||
tl.send(dir="E", src=l_local)
|
||||
tl.send(dir="E", src=O_local)
|
||||
elif col == _ROOT_COL:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="W", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="W", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="W", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
elif _ROOT_COL < col < _SUB_W - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="E", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="E", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="E", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
tl.send(dir="W", src=m_local)
|
||||
tl.send(dir="W", src=l_local)
|
||||
tl.send(dir="W", src=O_local)
|
||||
elif col == _SUB_W - 1:
|
||||
tl.send(dir="W", src=m_local)
|
||||
tl.send(dir="W", src=l_local)
|
||||
tl.send(dir="W", src=O_local)
|
||||
|
||||
# Phase 2: col reduce on col == root_col — converge at row == root_row.
|
||||
if col == root_col:
|
||||
# Phase 2: col reduce on col == _ROOT_COL — converge at row == _ROOT_ROW.
|
||||
if col == _ROOT_COL:
|
||||
if row == 0:
|
||||
if south_exists:
|
||||
tl.send(dir="S", src=m_local)
|
||||
tl.send(dir="S", src=l_local)
|
||||
tl.send(dir="S", src=O_local)
|
||||
elif 0 < row < root_row:
|
||||
if north_exists:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if south_exists:
|
||||
tl.send(dir="S", src=m_local)
|
||||
tl.send(dir="S", src=l_local)
|
||||
tl.send(dir="S", src=O_local)
|
||||
elif row == root_row:
|
||||
if north_exists:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if sub_h - 1 > root_row and south_exists:
|
||||
tl.send(dir="S", src=m_local)
|
||||
tl.send(dir="S", src=l_local)
|
||||
tl.send(dir="S", src=O_local)
|
||||
elif 0 < row < _ROOT_ROW:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
tl.send(dir="S", src=m_local)
|
||||
tl.send(dir="S", src=l_local)
|
||||
tl.send(dir="S", src=O_local)
|
||||
elif row == _ROOT_ROW:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="N", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="N", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="N", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if _SUB_H - 1 > _ROOT_ROW:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
|
||||
@@ -257,24 +194,21 @@ def reduce_mlo(pe_id, cube_id, m_local, l_local, O_local, C, P, *, tl):
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
elif root_row < row < sub_h - 1:
|
||||
if south_exists:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if north_exists:
|
||||
tl.send(dir="N", src=m_local)
|
||||
tl.send(dir="N", src=l_local)
|
||||
tl.send(dir="N", src=O_local)
|
||||
elif row == sub_h - 1 and sub_h - 1 > root_row:
|
||||
if north_exists:
|
||||
tl.send(dir="N", src=m_local)
|
||||
tl.send(dir="N", src=l_local)
|
||||
tl.send(dir="N", src=O_local)
|
||||
elif _ROOT_ROW < row < _SUB_H - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="S", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="S", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="S", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
tl.send(dir="N", src=m_local)
|
||||
tl.send(dir="N", src=l_local)
|
||||
tl.send(dir="N", src=O_local)
|
||||
elif row == _SUB_H - 1 and _SUB_H - 1 > _ROOT_ROW:
|
||||
tl.send(dir="N", src=m_local)
|
||||
tl.send(dir="N", src=l_local)
|
||||
tl.send(dir="N", src=O_local)
|
||||
|
||||
@@ -24,15 +24,15 @@ from __future__ import annotations
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import ( # noqa: E501
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite import ( # noqa: E501
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext import ( # noqa: E501
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16 import ( # noqa: E501
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_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
|
||||
@@ -50,12 +50,12 @@ _SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_composite.json"
|
||||
# Each kernel implements the same Case-6 placement; only the per-tile
|
||||
# command form differs (see module docstring).
|
||||
_VARIANT_KERNELS = {
|
||||
"primitive_tiled": gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel,
|
||||
"primitive": gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel,
|
||||
"composite": gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
|
||||
"composite_extended":
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_ext_kernel,
|
||||
}
|
||||
_VARIANTS = ("primitive_tiled", "composite", "composite_extended")
|
||||
_VARIANTS = ("primitive", "composite", "composite_extended")
|
||||
|
||||
# Op-count (PE_CPU dispatch) is computed at emit time (exact, instant), so
|
||||
# it spans up to the 1M production point (S_local = 1M/64 = 16384, 16
|
||||
|
||||
@@ -1,262 +0,0 @@
|
||||
"""GQA composite kernel across models — per-KV-group multi-model comparison.
|
||||
|
||||
The topology per model uses C = h_q (cubes per KV group = query heads
|
||||
per KV group in the model). P = 8 PEs / cube is fixed to the physical
|
||||
layout. So Gemma-2 27B (G=2) runs on a 2-cube slice; Command R+ (G=12)
|
||||
runs on 12 cubes; LLaMA-3-70B (G=8) matches the Case-6 baseline
|
||||
exactly.
|
||||
|
||||
Companion to ``gqa_decode_long_ctx_models.py``. That bench sweeps a
|
||||
fixed set of topology sizes (N ∈ {16, 32, 64}) across six GQA models
|
||||
and showed that latency is model-agnostic when only G varies.
|
||||
|
||||
This bench takes a different slice: it fixes the topology to match
|
||||
each model's G ratio — one cube per Q head in the KV group:
|
||||
|
||||
C = G (cubes per KV group = h_q of the KV group)
|
||||
P = 8 (fixed, physical PEs per cube)
|
||||
N = C · P = 8·G
|
||||
|
||||
Six models, one topology per model, one context length:
|
||||
|
||||
Gemma 2 27B G=2 → C=2 → N=16
|
||||
LLaMA-3 8B G=4 → C=4 → N=32
|
||||
Qwen 2.5 7B G=7 → C=7 → N=56 (2 KV groups per SIP, 2 cubes idle)
|
||||
LLaMA-3 70B G=8 → C=8 → N=64
|
||||
Qwen 2.5 72B G=8 → C=8 → N=64
|
||||
Command R+ G=12 → C=12 → N=96 (1 KV group per SIP, 4 cubes idle)
|
||||
|
||||
Physical SIP is 4×4 = 16 cubes (topology.yaml sip.cube_mesh: {w:4, h:4}),
|
||||
so all six configurations fit in one SIP.
|
||||
|
||||
Total: 6 engine runs at S_kv = 131 072 (~40 min wall-clock).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite import ( # noqa: E501
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_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
|
||||
|
||||
_OUTPUT_DIR = (
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_models.json"
|
||||
|
||||
|
||||
# ── Model registry (per-KV-group shapes) ─────────────────────────────
|
||||
|
||||
|
||||
# Kernel operates on ONE KV group at a time. h_q / h_kv here are the
|
||||
# per-group counts (not the model's totals); G = h_q / h_kv.
|
||||
# Full-model context (full_h_q, full_h_kv, hidden, layers) is metadata
|
||||
# for annotation only — it does not enter the kernel launch.
|
||||
_MODELS = {
|
||||
"llama3-8b": dict(
|
||||
family="Meta", params_b=8,
|
||||
h_q=4, h_kv=1, d_head=128,
|
||||
full_h_q=32, full_h_kv=8, hidden=4096, layers=32,
|
||||
),
|
||||
"llama3-70b": dict(
|
||||
family="Meta", params_b=70,
|
||||
h_q=8, h_kv=1, d_head=128,
|
||||
full_h_q=64, full_h_kv=8, hidden=8192, layers=80,
|
||||
),
|
||||
"qwen2.5-7b": dict(
|
||||
family="Alibaba", params_b=7,
|
||||
h_q=7, h_kv=1, d_head=128,
|
||||
full_h_q=28, full_h_kv=4, hidden=3584, layers=28,
|
||||
),
|
||||
"qwen2.5-72b": dict(
|
||||
family="Alibaba", params_b=72,
|
||||
h_q=8, h_kv=1, d_head=128,
|
||||
full_h_q=64, full_h_kv=8, hidden=8192, layers=80,
|
||||
),
|
||||
"gemma2-27b": dict(
|
||||
family="Google", params_b=27,
|
||||
h_q=2, h_kv=1, d_head=128,
|
||||
full_h_q=32, full_h_kv=16, hidden=4608, layers=46,
|
||||
),
|
||||
"command-r-plus": dict(
|
||||
family="Cohere", params_b=104,
|
||||
h_q=12, h_kv=1, d_head=128,
|
||||
full_h_q=96, full_h_kv=8, hidden=12288, layers=64,
|
||||
),
|
||||
}
|
||||
_MODEL_ORDER = tuple(_MODELS.keys())
|
||||
|
||||
_S_KV = 131_072 # 128 K context
|
||||
_T_Q = 1
|
||||
_P = 8 # fixed physical PEs per cube
|
||||
|
||||
|
||||
def _c_for(model_key: str) -> int:
|
||||
m = _MODELS[model_key]
|
||||
return m["h_q"] // m["h_kv"]
|
||||
|
||||
|
||||
def _end_to_end_ns(op_log) -> float:
|
||||
if not op_log:
|
||||
return 0.0
|
||||
return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
|
||||
|
||||
|
||||
def _op_breakdown_ns(op_log) -> dict[str, float]:
|
||||
"""Sum occupancy per op_kind across all components (not
|
||||
critical-path — ops overlap on independent components). First-order
|
||||
view of where time is spent inside the engine."""
|
||||
totals = {"matmul": 0.0, "comm": 0.0, "other": 0.0}
|
||||
for r in op_log:
|
||||
dur = r.t_end - r.t_start
|
||||
if r.op_kind in ("gemm", "math"):
|
||||
totals["matmul"] += dur
|
||||
elif r.op_kind == "memory":
|
||||
totals["comm"] += dur
|
||||
else:
|
||||
totals["other"] += dur
|
||||
return totals
|
||||
|
||||
|
||||
def _run_panel_fn(model_key: str, C: int, P: int, S_kv: int):
|
||||
m = _MODELS[model_key]
|
||||
panel = f"decode_models_{model_key}_c{C}p{P}_s{S_kv}"
|
||||
|
||||
def _bench_fn(ctx):
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_full = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
|
||||
num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((_T_Q, m["h_q"] * m["d_head"]),
|
||||
dtype="f16", dp=dp_full, name=f"{panel}_q")
|
||||
k = ctx.zeros((S_kv, m["h_kv"] * m["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name=f"{panel}_k")
|
||||
v = ctx.zeros((S_kv, m["h_kv"] * m["d_head"]),
|
||||
dtype="f16", dp=dp_kv, name=f"{panel}_v")
|
||||
o = ctx.empty((_T_Q, m["h_q"] * m["d_head"]),
|
||||
dtype="f16", dp=dp_full, name=f"{panel}_o")
|
||||
ctx.launch(
|
||||
panel,
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel,
|
||||
q, k, v, o,
|
||||
_T_Q, S_kv, m["h_q"], m["h_kv"], m["d_head"], C, P,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
return _bench_fn
|
||||
|
||||
|
||||
def _emit_dispatch(model_key: str, C: int, P: int,
|
||||
S_kv: int) -> tuple[int, float]:
|
||||
from kernbench.common.pe_commands import PeCpuOverheadCmd
|
||||
from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
|
||||
from kernbench.triton_emu.tl_context import TLContext, run_kernel
|
||||
|
||||
m = _MODELS[model_key]
|
||||
cube_id = min(6, C - 1)
|
||||
tl = TLContext(
|
||||
pe_id=0, num_programs=P, cost_model=DEFAULT_PE_COST_MODEL,
|
||||
cube_id=cube_id, num_cubes=C, scratch_base=1 << 61,
|
||||
scratch_size=1 << 20,
|
||||
)
|
||||
run_kernel(
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_composite_kernel, tl,
|
||||
0x1000, 0x2000, 0x3000, 0x4000,
|
||||
_T_Q, S_kv, m["h_q"], m["h_kv"], m["d_head"], C, P,
|
||||
)
|
||||
cmds = [c for c in tl.commands if isinstance(c, PeCpuOverheadCmd)]
|
||||
return len(cmds), sum(c.cycles for c in cmds)
|
||||
|
||||
|
||||
def _engine_run(model_key: str, C: int, P: int,
|
||||
S_kv: int, topology: str):
|
||||
"""Run the engine sim; return (latency_ns, op_kind_breakdown)."""
|
||||
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
|
||||
|
||||
topo = resolve_topology(topology)
|
||||
result = run_bench(
|
||||
topology=topo, bench_fn=_run_panel_fn(model_key, C, P, 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"gqa-decode-models {model_key} C={C} P={P} S={S_kv} "
|
||||
f"failed: {result.completion}"
|
||||
)
|
||||
op_log = result.engine.op_log
|
||||
latency = _end_to_end_ns(op_log)
|
||||
breakdown = _op_breakdown_ns(op_log)
|
||||
return latency, breakdown
|
||||
|
||||
|
||||
def run_sweep(topology: str = "topology.yaml") -> int:
|
||||
"""One row per model at its per-KV-group topology (C = h_q). Writes
|
||||
sweep_decode_models.json."""
|
||||
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
rows = []
|
||||
for model_key in _MODEL_ORDER:
|
||||
m = _MODELS[model_key]
|
||||
C = _c_for(model_key)
|
||||
N = C * _P
|
||||
n_cmds, cycles = _emit_dispatch(model_key, C, _P, _S_KV)
|
||||
latency, breakdown = _engine_run(model_key, C, _P, _S_KV, topology)
|
||||
rows.append({
|
||||
"model": model_key,
|
||||
"family": m["family"],
|
||||
"params_b": m["params_b"],
|
||||
"h_q": m["h_q"],
|
||||
"h_kv": m["h_kv"],
|
||||
"G": C,
|
||||
"d_head": m["d_head"],
|
||||
"full_h_q": m["full_h_q"],
|
||||
"full_h_kv": m["full_h_kv"],
|
||||
"hidden": m["hidden"],
|
||||
"layers": m["layers"],
|
||||
"C": C, "P": _P, "N": N,
|
||||
"T_q": _T_Q,
|
||||
"S_kv": _S_KV,
|
||||
"S_local": _S_KV // N,
|
||||
"variant": "composite",
|
||||
"pe_cpu_cmd_count": n_cmds,
|
||||
"pe_cpu_dispatch_cycles": cycles,
|
||||
"latency_ns": latency,
|
||||
# Op-kind occupancy (summed across all components — not
|
||||
# critical-path; first-order view of where engine time goes)
|
||||
"matmul_ns": breakdown["matmul"],
|
||||
"comm_ns": breakdown["comm"],
|
||||
"other_ns": breakdown["other"],
|
||||
})
|
||||
print(
|
||||
f" {model_key:<18} G={C:>2} C={C:>2} P={_P} N={N:>3} "
|
||||
f"latency={latency/1e3:>7.1f} µs cmds={n_cmds:>4} "
|
||||
f"matmul={breakdown['matmul']/1e3:>7.1f} "
|
||||
f"comm={breakdown['comm']/1e3:>7.1f}"
|
||||
)
|
||||
sweep = {
|
||||
"version": 1,
|
||||
"variant": "composite",
|
||||
"S_kv": _S_KV,
|
||||
"P": _P,
|
||||
"note": "C = h_q per model (cubes per KV group = query heads per KV group)",
|
||||
"models": list(_MODEL_ORDER),
|
||||
"rows": rows,
|
||||
}
|
||||
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
|
||||
print(f" gqa-decode-models: {len(rows)} rows -> {_SWEEP_JSON}")
|
||||
return len(rows)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_sweep()
|
||||
@@ -0,0 +1,168 @@
|
||||
"""milestone-1h-gqa: memory-bound decode streaming composite-command study.
|
||||
|
||||
Single-rank companion to the compute-bound prefill study
|
||||
(``gqa_prefill_compute_bound``). Same three command-form kernels
|
||||
(``_gqa_prefill_compute_bound``), same single-rank (C=P=1) harness, but
|
||||
driven with **T_q=1** (decode) and swept over a large S_kv. With T_q=1 the
|
||||
score / context GEMMs are skinny (M = G·T_q = 8), so the MAC array is
|
||||
barely fed and the kernel is **memory-bound** — streaming the KV cache out
|
||||
of HBM dominates. This is the regime mirror of prefill: command form is
|
||||
*latency-neutral* here, because the bottleneck is data movement, not
|
||||
issue.
|
||||
|
||||
Running single-rank (no cross-CUBE (m,ℓ,O) reduce) isolates the
|
||||
local-attention streaming vs. issue trade-off cleanly — unlike the 64-way
|
||||
Case-6 decode sweep, whose small-S_kv end-to-end latency is masked by the
|
||||
inter-CUBE reduce tail. S_kv here is the *per-rank* context, so S_kv=64K
|
||||
single-rank is the per-PE load of a 4M-token, 64-way-sharded decode.
|
||||
|
||||
Records per (variant, S_kv) the end-to-end latency, GEMM/DMA engine busy
|
||||
time, MAC utilization (achieved ÷ peak), and DMA occupancy (dma_busy ÷
|
||||
e2e), so the comparative plot can show all three command forms landing on
|
||||
the same latency curve while the MAC array stays floored and the DMA
|
||||
channel stays saturated.
|
||||
|
||||
Runs in data mode (engine latency). Gated via the umbrella
|
||||
``GQA_1H_SWEEPS=decode_streaming``.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
|
||||
gqa_prefill_composite_ext_kernel,
|
||||
gqa_prefill_composite_kernel,
|
||||
gqa_prefill_primitive_kernel,
|
||||
)
|
||||
from kernbench.policy.placement.dp import DPPolicy
|
||||
|
||||
_OUTPUT_DIR = (
|
||||
Path(__file__).resolve().parents[2]
|
||||
/ "1H_milestone_output" / "gqa" / "long_ctx"
|
||||
)
|
||||
_SWEEP_JSON = _OUTPUT_DIR / "sweep_decode_streaming.json"
|
||||
|
||||
# Same kernels as the prefill study — they are general attention kernels;
|
||||
# the regime is set by T_q (=1 here → memory-bound).
|
||||
_VARIANT_KERNELS = {
|
||||
"primitive": gqa_prefill_primitive_kernel,
|
||||
"composite": gqa_prefill_composite_kernel,
|
||||
"composite_extended": gqa_prefill_composite_ext_kernel,
|
||||
}
|
||||
_VARIANTS = ("primitive", "composite", "composite_extended")
|
||||
|
||||
# Per-rank context length. T_q=1, so M = G·T_q = 8 (skinny, memory-bound)
|
||||
# at every point. Capped at 16K: the *plain* composite materializes the
|
||||
# full (M, S_kv) scores in TCM scratch (~48·S_kv B incl. the exp transients),
|
||||
# which overflows the 1 MB kernel scratch beyond ~21K — itself a sign that
|
||||
# the recipe (composite_extended), which tiles the softmax, is what long
|
||||
# context actually needs. S_kv is per-rank, so 16K single-rank already
|
||||
# equals the per-PE load of a 1M-token, 64-way-sharded decode.
|
||||
_S_KV_POINTS = (2048, 4096, 8192, 16384)
|
||||
|
||||
_T_Q = 1
|
||||
_H_Q, _H_KV, _D_HEAD = 8, 1, 128
|
||||
_PEAK_TFLOPS = 8.0 # per-PE f16 GEMM peak (topology.yaml pe_gemm.peak_tflops_f16)
|
||||
|
||||
|
||||
def _run_panel_fn(variant: str, s_kv: int):
|
||||
kernel = _VARIANT_KERNELS[variant]
|
||||
panel = f"decode_stream_{variant}_s{s_kv}"
|
||||
|
||||
def _bench_fn(ctx):
|
||||
dp = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=1, num_pes=1)
|
||||
q = ctx.zeros((_T_Q, _H_Q * _D_HEAD),
|
||||
dtype="f16", dp=dp, name=f"{panel}_q")
|
||||
k = ctx.zeros((s_kv, _H_KV * _D_HEAD),
|
||||
dtype="f16", dp=dp, name=f"{panel}_k")
|
||||
v = ctx.zeros((s_kv, _H_KV * _D_HEAD),
|
||||
dtype="f16", dp=dp, name=f"{panel}_v")
|
||||
o = ctx.empty((_T_Q, _H_Q * _D_HEAD),
|
||||
dtype="f16", dp=dp, name=f"{panel}_o")
|
||||
ctx.launch(panel, kernel, q, k, v, o,
|
||||
_T_Q, s_kv, _H_Q, _H_KV, _D_HEAD, 1, 1,
|
||||
_auto_dim_remap=False)
|
||||
|
||||
return _bench_fn
|
||||
|
||||
|
||||
def _end_to_end_ns(op_log) -> float:
|
||||
if not op_log:
|
||||
return 0.0
|
||||
return max(r.t_end for r in op_log) - min(r.t_start for r in op_log)
|
||||
|
||||
|
||||
def _engine_busy_ns(op_log, suffix: str) -> float:
|
||||
return sum(r.t_end - r.t_start
|
||||
for r in op_log if r.component_id.endswith("." + suffix))
|
||||
|
||||
|
||||
def _run_panel(variant: str, s_kv: int, topology: str) -> dict:
|
||||
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
|
||||
|
||||
topo = resolve_topology(topology)
|
||||
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"gqa-decode-streaming {variant}@{s_kv} failed: {result.completion}"
|
||||
)
|
||||
op_log = result.engine.op_log
|
||||
e2e = _end_to_end_ns(op_log)
|
||||
gemm = _engine_busy_ns(op_log, "pe_gemm")
|
||||
dma = _engine_busy_ns(op_log, "pe_dma")
|
||||
G = _H_Q // _H_KV
|
||||
M = G * _T_Q
|
||||
# Useful attention flops (Q·Kᵀ + P·V), single rank.
|
||||
useful_flops = 4.0 * M * _D_HEAD * s_kv
|
||||
# KV-cache bytes streamed from HBM (K + V, f16) — the memory-bound
|
||||
# denominator. achieved_bw = bytes / e2e (GB/s) measures how close the
|
||||
# command form gets to the HBM roofline (the memory-bound mirror of
|
||||
# MAC utilization in the compute-bound prefill study).
|
||||
kv_bytes = 2.0 * s_kv * _D_HEAD * 2.0
|
||||
return {
|
||||
"variant": variant,
|
||||
"s_kv": s_kv,
|
||||
"M": M,
|
||||
"latency_ns": e2e,
|
||||
"gemm_busy_ns": gemm,
|
||||
"dma_busy_ns": dma,
|
||||
"kv_bytes": kv_bytes,
|
||||
"achieved_bw_gbs": (kv_bytes / e2e) if e2e > 0 else 0.0,
|
||||
"achieved_tflops": (useful_flops / e2e / 1e3) if e2e > 0 else 0.0,
|
||||
"mac_util": (useful_flops / e2e / 1e3 / _PEAK_TFLOPS) if e2e > 0 else 0.0,
|
||||
"dma_occupancy": (dma / e2e) if e2e > 0 else 0.0,
|
||||
}
|
||||
|
||||
|
||||
def run_sweep(topology: str = "topology.yaml") -> int:
|
||||
"""Drive all (variant, S_kv) decode-streaming panels; write sweep.json."""
|
||||
_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
rows = [
|
||||
_run_panel(variant, s_kv, topology)
|
||||
for s_kv in _S_KV_POINTS
|
||||
for variant in _VARIANTS
|
||||
]
|
||||
sweep = {
|
||||
"version": 1,
|
||||
"variants": list(_VARIANTS),
|
||||
"s_kv_points": list(_S_KV_POINTS),
|
||||
"rows": rows,
|
||||
}
|
||||
_SWEEP_JSON.write_text(json.dumps(sweep, indent=2))
|
||||
print(f" gqa-decode-streaming: {len(rows)} rows -> {_SWEEP_JSON}")
|
||||
return len(rows)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_sweep()
|
||||
@@ -1,74 +1,37 @@
|
||||
"""GQA decode kernel: short context, attention only, multi-tile per PE (1).
|
||||
"""GQA fused-attention decode kernel — short context (ADR-0060 §B.split.2).
|
||||
|
||||
Unified A1/A2/A4/B decode mapping per ADR-0070 (phase mirror of
|
||||
``_gqa_attention_prefill_short.py``). Mode selected at launch via
|
||||
``kv_per_cube ∈ {1, 2, 4, 8}``:
|
||||
Short context (``S_kv < 256K``): each CUBE owns ``kv_per_cube`` whole
|
||||
KV heads, with no S_kv sharding across CUBEs and no inter-CUBE reduce.
|
||||
PE-SP within each CUBE: the ``P`` PEs split into ``kv_per_cube`` groups
|
||||
of ``P/kv_per_cube`` PEs each; each group does PE-SP across the group
|
||||
for one owned head, then the group's root PE stores its head's output.
|
||||
|
||||
Mode kv_per_cube C group_size Reduce topology
|
||||
---- ----------- ------- ---------- -----------------------------
|
||||
A1 1 h_kv P (=8) row 0 + col bridge + row 1
|
||||
A2 2 h_kv/2 P/2 (=4) row chain only
|
||||
A4 4 h_kv/4 P/4 (=2) single intra_W hop
|
||||
B 8 1 1 (no reduce, single PE)
|
||||
The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so
|
||||
per-rank scratch is bounded by ``TILE_S_KV``.
|
||||
|
||||
Cube ``h`` owns ``kv_per_cube`` whole KV heads. Within each cube the
|
||||
8 PEs split into ``kv_per_cube`` groups of ``group_size = P/kv_per_cube``
|
||||
PEs; each PE owns ``S_local = S_kv/group_size`` KV tokens (sequence
|
||||
shard). FA2 fuses ``G = h_q/h_kv`` Q heads into the M dim of one GEMM
|
||||
per tile. After ``n_tiles_per_pe = S_local/TILE_S_KV`` tiles, partials
|
||||
``(m, ℓ, O)`` chain-reduce up to group root (PE 0), which normalizes
|
||||
and stores the cube's O slab.
|
||||
Group layout on the 2×4 PE grid:
|
||||
kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
|
||||
kv_per_cube=2, group=4 PEs (one row): row chain only.
|
||||
kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
|
||||
kv_per_cube=8, group=1 PE: no chain — direct write.
|
||||
|
||||
Per ADR-0011 D-VA1 the kernel computes its own shard base offset from
|
||||
``program_id(axis=0)`` (PE id) and ``program_id(axis=1)`` (cube id);
|
||||
deploy places shards in HBM but the kernel must address them.
|
||||
|
||||
Tensor layouts (host-side, mode-invariant byte totals):
|
||||
- Q: ``(kv_per_cube·T_q, h_q·d_head/kv_per_cube)``, T_q=1.
|
||||
dp=(cube=column_wise, pe=replicate). Caller pre-scales by
|
||||
``1/sqrt(d_head)``.
|
||||
- K: ``(h_kv·S_kv·d_head/TILE_S_KV, TILE_S_KV)`` tile-major.
|
||||
dp=(cube=row_wise, pe=row_wise).
|
||||
- V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
|
||||
- O: same dp as Q; only group root (pe_in_group==0) stores.
|
||||
|
||||
Configuration constraints:
|
||||
kv_per_cube ∈ {1, 2, 4, 8}, T_q == 1, P == 8,
|
||||
C == h_kv/kv_per_cube, h_q % h_kv == 0,
|
||||
S_kv % (group_size·TILE_S_KV) == 0.
|
||||
|
||||
Out of scope: causal mask (decode is auto-causal), f32 accumulator.
|
||||
Layout caveats:
|
||||
- K, V: ``(h_kv·S_kv, d_head)`` head-stacked, deployed with
|
||||
``dp = (cube=row_wise, pe=row_wise)`` so each PE's chunk is
|
||||
contiguously ``(S_local, d_head)`` at its own shard. K loaded as
|
||||
``(d_head, S_local)`` via byte-conserving reshape (ADR-0060 §3).
|
||||
- Q: replicated ``(T_q, h_q·d_head)``, reshaped byte-conservingly to
|
||||
``(h_q·T_q, d_head)``. Kernel computes attention for ALL Q rows
|
||||
against the group's owned K head; only the group's owned head rows
|
||||
are semantically meaningful (correct for zero / symmetric inputs).
|
||||
- O: replicated; each group root writes its head's
|
||||
``(h_q·T_q, d_head)`` result at disjoint PE-local addresses.
|
||||
- GEMMs use ``tl.dot`` (no composite epilogue / ``softmax_scale``).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024
|
||||
|
||||
|
||||
def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
|
||||
h_q: int, h_kv: int, S_kv: int) -> None:
|
||||
"""Validate decode kernel configuration before the run.
|
||||
"""
|
||||
if kv_per_cube not in (1, 2, 4, 8):
|
||||
raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
|
||||
if T_q != 1:
|
||||
raise ValueError(f"decode requires T_q == 1; got {T_q}")
|
||||
if P != 8:
|
||||
raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
|
||||
if h_q % h_kv != 0:
|
||||
raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
|
||||
if C != h_kv // kv_per_cube:
|
||||
raise ValueError(
|
||||
f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
|
||||
f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
|
||||
)
|
||||
group_size = P // kv_per_cube
|
||||
if S_kv % (group_size * TILE_S_KV) != 0:
|
||||
raise ValueError(
|
||||
f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
|
||||
f"({group_size * TILE_S_KV}); each PE's sequence shard must "
|
||||
f"be a whole number of tiles"
|
||||
)
|
||||
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
@@ -94,47 +57,35 @@ def gqa_attention_decode_short_kernel(
|
||||
C: int,
|
||||
P: int,
|
||||
kv_per_cube: int,
|
||||
cube_base: int = 0,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Unified decode: sequence-shard + chain reduce + FA2 (ADR-0070).
|
||||
"""
|
||||
"""Short-context GQA decode with PE-parallel heads + intra-group PE-SP."""
|
||||
group_size = P // kv_per_cube
|
||||
S_local = S_kv // group_size
|
||||
n_tiles_per_pe = S_local // TILE_S_KV
|
||||
|
||||
G = h_q // h_kv
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
# User-local cube index: a batched user placed at cube_start=cube_base
|
||||
# addresses its own shards from a 0-based cube index (default 0 = single user).
|
||||
cube_local = cube_id - cube_base
|
||||
pe_in_group = pe_id % group_size
|
||||
group_id_in_cube = pe_id // group_size
|
||||
S_local = S_kv // group_size
|
||||
|
||||
KV_ROW_BYTES = d_head * 2
|
||||
K_TILE_BYTES = d_head * TILE_S_KV * 2
|
||||
Q_ROW_BYTES = G * d_head * 2
|
||||
K_HEAD_BYTES = S_kv * d_head * 2
|
||||
V_HEAD_BYTES = S_kv * KV_ROW_BYTES
|
||||
# ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
|
||||
Q = tl.load(q_ptr, shape=(h_q * T_q, d_head), dtype="f16")
|
||||
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
|
||||
KV_ROW_BYTES = d_head * 2 # f16
|
||||
|
||||
# Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
|
||||
q_base = q_ptr + cube_local * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
|
||||
k_shard_base = (k_ptr
|
||||
+ cube_local * kv_per_cube * K_HEAD_BYTES
|
||||
+ group_id_in_cube * K_HEAD_BYTES
|
||||
+ pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
|
||||
v_shard_base = (v_ptr
|
||||
+ cube_local * kv_per_cube * V_HEAD_BYTES
|
||||
+ group_id_in_cube * V_HEAD_BYTES
|
||||
+ pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)
|
||||
|
||||
Q = tl.load(q_base, shape=(G, d_head), dtype="f16")
|
||||
|
||||
# ── Tile 0: establish (m, ℓ, O) ──────────────────────────────────
|
||||
K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.load(v_shard_base, shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
# Tile 0: establishes persistent (m_local, l_local, O_local).
|
||||
#
|
||||
# Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
|
||||
# - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
|
||||
# otherwise scope teardown discards them before the next tile's
|
||||
# merge can read them;
|
||||
# - kernbench has no scratch-backed initializer — ``tl.zeros`` /
|
||||
# ``tl.full`` return addr=0 handles with no backing storage, so
|
||||
# they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
|
||||
# So Tile 0 computes the initial running state directly; Tiles 1..N
|
||||
# fold into it. Triton port: limitation does not apply (SSA tensors
|
||||
# stay live across iterations) — a single unified loop suffices.
|
||||
tile_s0 = min(TILE_S_KV, S_local)
|
||||
K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
|
||||
V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
|
||||
scores = tl.dot(Q, K_T)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
@@ -142,13 +93,17 @@ def gqa_attention_decode_short_kernel(
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
O_local = tl.dot(exp_scores, V)
|
||||
|
||||
# ── Tiles 1..n_tiles_per_pe-1: sweep this PE's sequence shard ──
|
||||
for tile_idx in range(1, n_tiles_per_pe):
|
||||
# Tiles 1..n_tiles-1: fold into running state via online-softmax merge.
|
||||
# Triton port: drop the ``with tl.scratch_scope():`` line and replace
|
||||
# each ``copy_to`` with a Python rebind.
|
||||
for tile_idx in range(1, n_tiles):
|
||||
tile_start = tile_idx * TILE_S_KV
|
||||
tile_s = min(TILE_S_KV, S_local - tile_start)
|
||||
with tl.scratch_scope():
|
||||
K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(tile_s, d_head), dtype="f16")
|
||||
scores_t = tl.dot(Q, K_T_t)
|
||||
m_tile = tl.max(scores_t, axis=-1)
|
||||
centered_t = scores_t - m_tile
|
||||
@@ -162,13 +117,13 @@ def gqa_attention_decode_short_kernel(
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
|
||||
# ── 2x4 mesh chain reduce geometry within the group ──
|
||||
# ── Communication: within-group chain reduce-to-root (Level-2 only) ──
|
||||
group_cols = min(4, group_size)
|
||||
group_rows = (group_size + group_cols - 1) // group_cols
|
||||
pe_col_in_group = pe_in_group % group_cols
|
||||
pe_row_in_group = pe_in_group // group_cols
|
||||
|
||||
# Row chain (intra_W, leftward) — every group with group_cols > 1.
|
||||
# Row chain (within group's row, along intra_W, leftward).
|
||||
if group_cols > 1:
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
with tl.scratch_scope():
|
||||
@@ -186,7 +141,7 @@ def gqa_attention_decode_short_kernel(
|
||||
tl.send(dir="intra_W", src=l_local)
|
||||
tl.send(dir="intra_W", src=O_local)
|
||||
|
||||
# Col bridge (intra_N, row-1 col-0 → row-0 col-0) — only A1 (group_rows > 1).
|
||||
# Col bridge (within group, along intra_N, row-1 → row-0).
|
||||
if pe_col_in_group == 0 and group_rows > 1:
|
||||
if pe_row_in_group < group_rows - 1:
|
||||
with tl.scratch_scope():
|
||||
@@ -204,10 +159,7 @@ def gqa_attention_decode_short_kernel(
|
||||
tl.send(dir="intra_N", src=l_local)
|
||||
tl.send(dir="intra_N", src=O_local)
|
||||
|
||||
# ── Final normalize + store (group root, pe_in_group == 0) ──
|
||||
# ── Final normalise + store (group root only) ──
|
||||
if pe_in_group == 0:
|
||||
O_final = O_local / l_local
|
||||
o_base = (o_ptr
|
||||
+ cube_local * kv_per_cube * Q_ROW_BYTES
|
||||
+ group_id_in_cube * Q_ROW_BYTES)
|
||||
tl.store(o_base, O_final)
|
||||
tl.store(o_ptr, O_final)
|
||||
|
||||
@@ -1,193 +0,0 @@
|
||||
"""GQA decode kernel: composite GEMM-only variant (2).
|
||||
|
||||
Identical mapping to ``_gqa_attention_decode_short.py`` (sequence-shard
|
||||
+ chain reduce + FA2 head fusion, unified A1/A2/A4/B). Difference vs
|
||||
first-level: the Q·Kᵀ GEMM is issued via ``tl.composite(op="gemm", ...)``
|
||||
instead of ``tl.dot``.
|
||||
|
||||
This is the GEMM-only tier: Q·Kᵀ is a composite (operands ``a=Q`` /
|
||||
``b=K_T`` are pinned ``tl.load`` results), while P·V stays a plain
|
||||
``tl.dot`` and softmax stays a primitive MATH chain — no
|
||||
``softmax_merge`` fusion. Folding the per-tile softmax into a P·V
|
||||
composite (the ``softmax_merge`` prologue making ``P`` a pinned
|
||||
primary-out bound to the head GEMM) is variant (3) in
|
||||
``_gqa_attention_decode_short_composite_fused.py``.
|
||||
|
||||
Three-variant comparison:
|
||||
(1) without composite : ``_gqa_attention_decode_short.py``
|
||||
(2) with composite (GEMM-only, no fuse) : this file
|
||||
(3) with composite + softmax_merge fuse : ``…_composite_fused.py``
|
||||
|
||||
Shard addressing, layouts, and caller contract are identical to the
|
||||
first-level decode kernel (ADR-0011 D-VA1).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024
|
||||
|
||||
|
||||
def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
|
||||
h_q: int, h_kv: int, S_kv: int) -> None:
|
||||
"""Validate composite-decode config — caller-side, sim-cost 0.
|
||||
|
||||
Mirrors first-level decode ``_validate_config``.
|
||||
"""
|
||||
if kv_per_cube not in (1, 2, 4, 8):
|
||||
raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
|
||||
if T_q != 1:
|
||||
raise ValueError(f"decode requires T_q == 1; got {T_q}")
|
||||
if P != 8:
|
||||
raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
|
||||
if h_q % h_kv != 0:
|
||||
raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
|
||||
if C != h_kv // kv_per_cube:
|
||||
raise ValueError(
|
||||
f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
|
||||
f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
|
||||
)
|
||||
group_size = P // kv_per_cube
|
||||
if S_kv % (group_size * TILE_S_KV) != 0:
|
||||
raise ValueError(
|
||||
f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
|
||||
f"({group_size * TILE_S_KV}); each PE's sequence shard must "
|
||||
f"be a whole number of tiles"
|
||||
)
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
"""Used only by the intra-group chain reduce below (recipe handles
|
||||
per-tile fold internally)."""
|
||||
m_new = tl.maximum(m_local, m_other)
|
||||
scale_old = tl.exp(m_local - m_new)
|
||||
scale_new = tl.exp(m_other - m_new)
|
||||
l_new = l_local * scale_old + l_other * scale_new
|
||||
O_new = O_local * scale_old + O_other * scale_new
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def gqa_attention_decode_short_composite_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
P: int,
|
||||
kv_per_cube: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Composite-GEMM decode — same mapping as first-level + tl.composite.
|
||||
|
||||
Caller must invoke ``_validate_config(...)`` first.
|
||||
"""
|
||||
group_size = P // kv_per_cube
|
||||
S_local = S_kv // group_size
|
||||
n_tiles_per_pe = S_local // TILE_S_KV
|
||||
|
||||
G = h_q // h_kv
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
pe_in_group = pe_id % group_size
|
||||
group_id_in_cube = pe_id // group_size
|
||||
|
||||
KV_ROW_BYTES = d_head * 2
|
||||
K_TILE_BYTES = d_head * TILE_S_KV * 2
|
||||
Q_ROW_BYTES = G * d_head * 2
|
||||
K_HEAD_BYTES = S_kv * d_head * 2
|
||||
V_HEAD_BYTES = S_kv * KV_ROW_BYTES
|
||||
|
||||
# Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
|
||||
q_base = q_ptr + cube_id * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
|
||||
k_shard_base = (k_ptr
|
||||
+ cube_id * kv_per_cube * K_HEAD_BYTES
|
||||
+ group_id_in_cube * K_HEAD_BYTES
|
||||
+ pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
|
||||
v_shard_base = (v_ptr
|
||||
+ cube_id * kv_per_cube * V_HEAD_BYTES
|
||||
+ group_id_in_cube * V_HEAD_BYTES
|
||||
+ pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)
|
||||
|
||||
Q = tl.load(q_base, shape=(G, d_head), dtype="f16")
|
||||
|
||||
# ── Tile 0: establish (m, ℓ, O) — Q·Kᵀ composite, softmax + P·V primitives ──
|
||||
K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.load(v_shard_base, shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
scores = tl.composite(op="gemm", a=Q, b=K_T)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
exp_scores = tl.exp(centered)
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
O_local = tl.dot(exp_scores, V) # P·V stays primitive in the GEMM-only tier
|
||||
|
||||
# ── Tiles 1..n_tiles_per_pe-1: Q·Kᵀ composite + softmax + P·V tl.dot ──
|
||||
for tile_idx in range(1, n_tiles_per_pe):
|
||||
with tl.scratch_scope():
|
||||
K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
|
||||
m_tile = tl.max(scores_t, axis=-1)
|
||||
centered_t = scores_t - m_tile
|
||||
exp_scores_t = tl.exp(centered_t)
|
||||
l_tile = tl.sum(exp_scores_t, axis=-1)
|
||||
O_tile = tl.dot(exp_scores_t, V_t)
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
|
||||
# ── Chain reduce ──
|
||||
group_cols = min(4, group_size)
|
||||
group_rows = (group_size + group_cols - 1) // group_cols
|
||||
pe_col_in_group = pe_in_group % group_cols
|
||||
pe_row_in_group = pe_in_group // group_cols
|
||||
|
||||
if group_cols > 1:
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if pe_col_in_group > 0:
|
||||
tl.send(dir="intra_W", src=m_local)
|
||||
tl.send(dir="intra_W", src=l_local)
|
||||
tl.send(dir="intra_W", src=O_local)
|
||||
|
||||
if pe_col_in_group == 0 and group_rows > 1:
|
||||
if pe_row_in_group < group_rows - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if pe_row_in_group > 0:
|
||||
tl.send(dir="intra_N", src=m_local)
|
||||
tl.send(dir="intra_N", src=l_local)
|
||||
tl.send(dir="intra_N", src=O_local)
|
||||
|
||||
if pe_in_group == 0:
|
||||
O_final = O_local / l_local
|
||||
o_base = (o_ptr
|
||||
+ cube_id * kv_per_cube * Q_ROW_BYTES
|
||||
+ group_id_in_cube * Q_ROW_BYTES)
|
||||
tl.store(o_base, O_final)
|
||||
-195
@@ -1,195 +0,0 @@
|
||||
"""GQA decode kernel: composite + softmax_merge fused variant (3).
|
||||
|
||||
Identical mapping to ``_gqa_attention_decode_short.py``. Difference vs
|
||||
the GEMM-only composite: per-tile softmax is folded into
|
||||
the P·V composite via the ``softmax_merge`` prologue recipe, eliminating
|
||||
the GEMM/MATH engine bubble that a flat primitives chain incurs.
|
||||
|
||||
Per-tile structure (ADR-0065 two-composite pattern, cf. decode_opt2):
|
||||
tile 0 primitives establish (m, ℓ, O)
|
||||
tile k>0 #1 Q·Kᵀ composite → scores (pinned primary-out)
|
||||
#2 softmax_merge prologue → P (pinned, auto-binds to GEMM a)
|
||||
+ P·V composite, ``out=O_local``
|
||||
+ add epilogue folding the result into O_local
|
||||
|
||||
The intra-group chain reduce after the tile loop is unchanged.
|
||||
|
||||
Three-variant comparison:
|
||||
(1) without composite : ``_gqa_attention_decode_short.py``
|
||||
(2) with composite (GEMM-only, no fuse) : ``…_composite.py``
|
||||
(3) with composite + softmax_merge fuse : this file
|
||||
|
||||
Shard addressing, layouts, and caller contract are identical to the
|
||||
first-level decode kernel (ADR-0011 D-VA1).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024
|
||||
|
||||
|
||||
def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
|
||||
h_q: int, h_kv: int, S_kv: int) -> None:
|
||||
"""Validate composite-decode config — caller-side, sim-cost 0.
|
||||
|
||||
Mirrors first-level decode ``_validate_config``.
|
||||
"""
|
||||
if kv_per_cube not in (1, 2, 4, 8):
|
||||
raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
|
||||
if T_q != 1:
|
||||
raise ValueError(f"decode requires T_q == 1; got {T_q}")
|
||||
if P != 8:
|
||||
raise ValueError(f"2x4 mesh chain-reduce geometry requires P == 8; got {P}")
|
||||
if h_q % h_kv != 0:
|
||||
raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
|
||||
if C != h_kv // kv_per_cube:
|
||||
raise ValueError(
|
||||
f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
|
||||
f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
|
||||
)
|
||||
group_size = P // kv_per_cube
|
||||
if S_kv % (group_size * TILE_S_KV) != 0:
|
||||
raise ValueError(
|
||||
f"S_kv ({S_kv}) must be a multiple of group_size·TILE_S_KV "
|
||||
f"({group_size * TILE_S_KV}); each PE's sequence shard must "
|
||||
f"be a whole number of tiles"
|
||||
)
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
"""Used only by the intra-group chain reduce below (recipe handles
|
||||
per-tile fold internally)."""
|
||||
m_new = tl.maximum(m_local, m_other)
|
||||
scale_old = tl.exp(m_local - m_new)
|
||||
scale_new = tl.exp(m_other - m_new)
|
||||
l_new = l_local * scale_old + l_other * scale_new
|
||||
O_new = O_local * scale_old + O_other * scale_new
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def gqa_attention_decode_short_composite_fused_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
P: int,
|
||||
kv_per_cube: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Composite-GEMM decode — same mapping as first-level + tl.composite.
|
||||
|
||||
Caller must invoke ``_validate_config(...)`` first.
|
||||
"""
|
||||
group_size = P // kv_per_cube
|
||||
S_local = S_kv // group_size
|
||||
n_tiles_per_pe = S_local // TILE_S_KV
|
||||
|
||||
G = h_q // h_kv
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
pe_in_group = pe_id % group_size
|
||||
group_id_in_cube = pe_id // group_size
|
||||
|
||||
KV_ROW_BYTES = d_head * 2
|
||||
K_TILE_BYTES = d_head * TILE_S_KV * 2
|
||||
Q_ROW_BYTES = G * d_head * 2
|
||||
K_HEAD_BYTES = S_kv * d_head * 2
|
||||
V_HEAD_BYTES = S_kv * KV_ROW_BYTES
|
||||
|
||||
# Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
|
||||
q_base = q_ptr + cube_id * kv_per_cube * Q_ROW_BYTES + group_id_in_cube * Q_ROW_BYTES
|
||||
k_shard_base = (k_ptr
|
||||
+ cube_id * kv_per_cube * K_HEAD_BYTES
|
||||
+ group_id_in_cube * K_HEAD_BYTES
|
||||
+ pe_in_group * n_tiles_per_pe * K_TILE_BYTES)
|
||||
v_shard_base = (v_ptr
|
||||
+ cube_id * kv_per_cube * V_HEAD_BYTES
|
||||
+ group_id_in_cube * V_HEAD_BYTES
|
||||
+ pe_in_group * n_tiles_per_pe * TILE_S_KV * KV_ROW_BYTES)
|
||||
|
||||
Q = tl.load(q_base, shape=(G, d_head), dtype="f16")
|
||||
|
||||
# ── Tile 0: establish running (m, ℓ, O) with primitives ──────────
|
||||
# (Reference: decode_opt2 — running state is set up with tl.dot/MATH
|
||||
# primitives, not composite. Recipe-driven composite enters in tile 1+.)
|
||||
K_T = tl.load(k_shard_base, shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.load(v_shard_base, shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
scores = tl.dot(Q, K_T)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
exp_scores = tl.exp(centered)
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
O_local = tl.dot(exp_scores, V)
|
||||
|
||||
# ── Tiles 1..n_tiles_per_pe-1: two composites per tile ──────────
|
||||
# #1 Q·Kᵀ composite → scores (pinned primary-out, fed into #2).
|
||||
# #2 softmax_merge prologue + P·V GEMM + add epilogue, all in one
|
||||
# composite: updates (m, ℓ) in place, computes P, runs P·V with
|
||||
# pinned auto-bind, and folds the result into O_local.
|
||||
for tile_idx in range(1, n_tiles_per_pe):
|
||||
with tl.scratch_scope():
|
||||
K_T_t = tl.load(k_shard_base + tile_idx * K_TILE_BYTES,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.load(v_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
|
||||
tl.composite(
|
||||
prologue=[{"op": "softmax_merge", "s": scores_t,
|
||||
"m": m_local, "l": l_local, "O": O_local}],
|
||||
op="gemm", b=V_t, out=O_local,
|
||||
epilogue=[{"op": "add", "other": O_local}],
|
||||
)
|
||||
|
||||
# ── Chain reduce ──
|
||||
group_cols = min(4, group_size)
|
||||
group_rows = (group_size + group_cols - 1) // group_cols
|
||||
pe_col_in_group = pe_in_group % group_cols
|
||||
pe_row_in_group = pe_in_group // group_cols
|
||||
|
||||
if group_cols > 1:
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if pe_col_in_group > 0:
|
||||
tl.send(dir="intra_W", src=m_local)
|
||||
tl.send(dir="intra_W", src=l_local)
|
||||
tl.send(dir="intra_W", src=O_local)
|
||||
|
||||
if pe_col_in_group == 0 and group_rows > 1:
|
||||
if pe_row_in_group < group_rows - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if pe_row_in_group > 0:
|
||||
tl.send(dir="intra_N", src=m_local)
|
||||
tl.send(dir="intra_N", src=l_local)
|
||||
tl.send(dir="intra_N", src=O_local)
|
||||
|
||||
if pe_in_group == 0:
|
||||
O_final = O_local / l_local
|
||||
o_base = (o_ptr
|
||||
+ cube_id * kv_per_cube * Q_ROW_BYTES
|
||||
+ group_id_in_cube * Q_ROW_BYTES)
|
||||
tl.store(o_base, O_final)
|
||||
@@ -1,45 +1,15 @@
|
||||
"""GQA prefill kernel: short context, attention only, multi-tile (1).
|
||||
"""GQA fused-attention prefill kernel — short context (ADR-0060 §B.split.2).
|
||||
|
||||
Unified A1/A2/A4/B prefill mapping per ADR-0070 (supersedes
|
||||
ADR-0060 §B.split.2 prefill-short clause). Mode selected at launch
|
||||
via ``kv_per_cube ∈ {1, 2, 4, 8}``:
|
||||
Prefill analogue of ``_gqa_decode_short.py`` — same CUBE/PE layout
|
||||
(``kv_per_cube`` heads per CUBE, group-PE-SP, within-group chain reduce,
|
||||
no inter-CUBE reduce). The only structural difference from short decode:
|
||||
``T_q`` may be > 1 (prefill processes multiple query tokens) and Q is
|
||||
shaped ``(T_q, h_kv·d_head)`` — one Q head per KV head, no GQA M-fold.
|
||||
|
||||
Mode kv_per_cube C group_size Broadcast topology
|
||||
---- ----------- ------- ---------- ---------------------------
|
||||
A1 1 h_kv P (=8) row 0 + col bridge + row 1
|
||||
A2 2 h_kv/2 P/2 (=4) row chain only
|
||||
A4 4 h_kv/4 P/4 (=2) single intra_E hop
|
||||
B 8 1 1 (no broadcast, single PE)
|
||||
The local attention uses an S_kv-axis tile sweep (ADR-0063 §A.2) so
|
||||
per-rank scratch is bounded by ``TILE_S_KV``.
|
||||
|
||||
Cube ``h`` owns ``kv_per_cube`` whole KV heads. Within each cube the
|
||||
8 PEs split into ``kv_per_cube`` groups of ``group_size = P/kv_per_cube``
|
||||
PEs; the group's T_q rows are floor-balanced across its PEs (Q-tile
|
||||
split). FA2 fuses ``G = h_q/h_kv`` Q heads into one batched GEMM per
|
||||
PE per tile. 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 — every PE writes its own
|
||||
``(T_q_pe·G, d_head)`` slab to disjoint rows of O.
|
||||
|
||||
Per ADR-0011 D-VA1 the kernel computes its own shard base offset from
|
||||
``program_id(axis=0)`` (PE id) and ``program_id(axis=1)`` (cube id);
|
||||
deploy places shards in HBM but the kernel must address them.
|
||||
|
||||
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). Caller pre-scales by
|
||||
``1/sqrt(d_head)``.
|
||||
- K: ``(h_kv·n_tiles·d_head, TILE_S_KV)`` tile-major.
|
||||
dp=(cube=row_wise, pe=replicate).
|
||||
- V: ``(h_kv·S_kv, d_head)`` native. Same dp as K.
|
||||
- O: same dp as Q; every PE stores.
|
||||
|
||||
Configuration constraints:
|
||||
kv_per_cube ∈ {1, 2, 4, 8}, P == 8,
|
||||
C == h_kv/kv_per_cube, h_q % h_kv == 0,
|
||||
T_q >= group_size, S_kv % TILE_S_KV == 0.
|
||||
|
||||
Out of scope: causal mask, f32 accumulator. See ADR-0070 §Known
|
||||
Limitations.
|
||||
No Ring KV here — each owned head is fully resident at its CUBE.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -47,33 +17,6 @@ from __future__ import annotations
|
||||
TILE_S_KV = 1024 # ADR-0063 §A.2 S_kv-axis tile sweep (per-tile width).
|
||||
|
||||
|
||||
def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
|
||||
h_q: int, h_kv: int, S_kv: int) -> None:
|
||||
"""Validate prefill kernel configuration before the run.
|
||||
"""
|
||||
if kv_per_cube not in (1, 2, 4, 8):
|
||||
raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
|
||||
if P != 8:
|
||||
raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}")
|
||||
if h_q % h_kv != 0:
|
||||
raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
|
||||
if C != h_kv // kv_per_cube:
|
||||
raise ValueError(
|
||||
f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
|
||||
f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
|
||||
)
|
||||
group_size = P // kv_per_cube
|
||||
if T_q < group_size:
|
||||
raise ValueError(
|
||||
f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
|
||||
f"has q-tile work for the broadcast chain"
|
||||
)
|
||||
if S_kv % TILE_S_KV != 0:
|
||||
raise ValueError(
|
||||
f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
|
||||
)
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
"""Online-softmax merge of two partial ``(m, ℓ, O)`` triples."""
|
||||
m_new = tl.maximum(m_local, m_other)
|
||||
@@ -91,7 +34,6 @@ def gqa_attention_prefill_short_kernel(
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
@@ -100,79 +42,32 @@ def gqa_attention_prefill_short_kernel(
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Unified prefill: Q-tile split + FA2 + IPCQ KV broadcast (ADR-0070)."""
|
||||
"""Short-context prefill with PE-parallel heads + intra-group PE-SP."""
|
||||
group_size = P // kv_per_cube
|
||||
|
||||
G = h_q // h_kv
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
pe_in_group = pe_id % group_size
|
||||
group_id_in_cube = pe_id // group_size
|
||||
S_local = S_kv // group_size
|
||||
|
||||
# Floor-balanced Q-tile partition.
|
||||
q_start = (T_q * pe_in_group) // group_size
|
||||
q_end = (T_q * (pe_in_group + 1)) // group_size
|
||||
T_q_pe = q_end - q_start
|
||||
|
||||
n_tiles = S_kv // TILE_S_KV
|
||||
|
||||
Q_ROW_BYTES = G * d_head * 2
|
||||
KV_ROW_BYTES = d_head * 2
|
||||
K_TILE_BYTES = d_head * TILE_S_KV * 2
|
||||
K_HEAD_BYTES = n_tiles * K_TILE_BYTES
|
||||
V_HEAD_BYTES = S_kv * KV_ROW_BYTES
|
||||
|
||||
# Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
|
||||
q_base = (q_ptr
|
||||
+ cube_id * kv_per_cube * T_q * Q_ROW_BYTES
|
||||
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
|
||||
k_head_shard_base = (k_ptr
|
||||
+ cube_id * kv_per_cube * K_HEAD_BYTES
|
||||
+ group_id_in_cube * K_HEAD_BYTES)
|
||||
v_head_shard_base = (v_ptr
|
||||
+ cube_id * kv_per_cube * V_HEAD_BYTES
|
||||
+ group_id_in_cube * V_HEAD_BYTES)
|
||||
|
||||
Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16")
|
||||
|
||||
# 2x4 mesh broadcast geometry within the group.
|
||||
group_cols = min(4, group_size)
|
||||
group_rows = (group_size + group_cols - 1) // group_cols
|
||||
pe_col_in_group = pe_in_group % group_cols
|
||||
pe_row_in_group = pe_in_group // group_cols
|
||||
|
||||
# ──────────────────────────────────────────────────────────
|
||||
# Tile 0: KV broadcast + establish persistent (m, ℓ, O).
|
||||
# Persistent state lives OUTSIDE scratch_scope (ADR-0063 §A.3).
|
||||
# ──────────────────────────────────────────────────────────
|
||||
if pe_in_group == 0:
|
||||
K_T = tl.load(k_head_shard_base,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.load(v_head_shard_base,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
if group_rows > 1:
|
||||
tl.send(dir="intra_S", src=K_T)
|
||||
tl.send(dir="intra_S", src=V)
|
||||
elif pe_col_in_group == 0 and pe_row_in_group > 0:
|
||||
K_T = tl.recv(dir="intra_N",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.recv(dir="intra_N",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
else:
|
||||
K_T = tl.recv(dir="intra_W",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.recv(dir="intra_W",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
# ── Local attention (S_kv-axis tile sweep, ADR-0063 §A.2) ──
|
||||
Q = tl.load(q_ptr, shape=(h_kv * T_q, d_head), dtype="f16")
|
||||
n_tiles = (S_local + TILE_S_KV - 1) // TILE_S_KV
|
||||
KV_ROW_BYTES = d_head * 2 # f16
|
||||
|
||||
# Tile 0: establishes persistent (m_local, l_local, O_local).
|
||||
#
|
||||
# Cannot be folded into the Tiles 1..N loop (kernbench-only limitation):
|
||||
# - persistent (m, ℓ, O) must live OUTSIDE ``tl.scratch_scope``,
|
||||
# otherwise scope teardown discards them before the next tile's
|
||||
# merge can read them;
|
||||
# - kernbench has no scratch-backed initializer — ``tl.zeros`` /
|
||||
# ``tl.full`` return addr=0 handles with no backing storage, so
|
||||
# they cannot be overwritten via ``tl.copy_to`` to seed (-inf, 0, 0).
|
||||
# So Tile 0 computes the initial running state directly; Tiles 1..N
|
||||
# fold into it. Triton port: limitation does not apply (SSA tensors
|
||||
# stay live across iterations) — a single unified loop suffices.
|
||||
tile_s0 = min(TILE_S_KV, S_local)
|
||||
K_T = tl.load(k_ptr, shape=(d_head, tile_s0), dtype="f16")
|
||||
V = tl.load(v_ptr, shape=(tile_s0, d_head), dtype="f16")
|
||||
scores = tl.dot(Q, K_T)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
@@ -180,38 +75,17 @@ def gqa_attention_prefill_short_kernel(
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
O_local = tl.dot(exp_scores, V)
|
||||
|
||||
# ──────────────────────────────────────────────────────────
|
||||
# Tiles 1..n_tiles-1: broadcast + fold into running state.
|
||||
# ──────────────────────────────────────────────────────────
|
||||
# Tiles 1..n_tiles-1: fold into running state via online-softmax merge.
|
||||
# Triton port: drop the ``with tl.scratch_scope():`` line and replace
|
||||
# each ``copy_to`` with a Python rebind.
|
||||
for tile_idx in range(1, n_tiles):
|
||||
tile_start = tile_idx * TILE_S_KV
|
||||
tile_s = min(TILE_S_KV, S_local - tile_start)
|
||||
with tl.scratch_scope():
|
||||
if pe_in_group == 0:
|
||||
K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
if group_rows > 1:
|
||||
tl.send(dir="intra_S", src=K_T_t)
|
||||
tl.send(dir="intra_S", src=V_t)
|
||||
elif pe_col_in_group == 0 and pe_row_in_group > 0:
|
||||
K_T_t = tl.recv(dir="intra_N",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.recv(dir="intra_N",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
else:
|
||||
K_T_t = tl.recv(dir="intra_W",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.recv(dir="intra_W",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
K_T_t = tl.load(k_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(d_head, tile_s), dtype="f16")
|
||||
V_t = tl.load(v_ptr + tile_start * KV_ROW_BYTES,
|
||||
shape=(tile_s, d_head), dtype="f16")
|
||||
scores_t = tl.dot(Q, K_T_t)
|
||||
m_tile = tl.max(scores_t, axis=-1)
|
||||
centered_t = scores_t - m_tile
|
||||
@@ -225,10 +99,49 @@ def gqa_attention_prefill_short_kernel(
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
|
||||
# No intra-group reduce — q-tiles independent. Each PE writes its own
|
||||
# slab to disjoint rows of the cube's column-wise O slab.
|
||||
O_final = O_local / l_local
|
||||
o_base = (o_ptr
|
||||
+ cube_id * kv_per_cube * T_q * Q_ROW_BYTES
|
||||
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
|
||||
tl.store(o_base, O_final)
|
||||
# ── Communication: within-group chain reduce-to-root (Level-2 only) ──
|
||||
group_cols = min(4, group_size)
|
||||
group_rows = (group_size + group_cols - 1) // group_cols
|
||||
pe_col_in_group = pe_in_group % group_cols
|
||||
pe_row_in_group = pe_in_group // group_cols
|
||||
|
||||
# Row chain (within group's row, along intra_W, leftward).
|
||||
if group_cols > 1:
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="intra_E", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="intra_E", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="intra_E", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if pe_col_in_group > 0:
|
||||
tl.send(dir="intra_W", src=m_local)
|
||||
tl.send(dir="intra_W", src=l_local)
|
||||
tl.send(dir="intra_W", src=O_local)
|
||||
|
||||
# Col bridge (within group, along intra_N, row-1 → row-0).
|
||||
if pe_col_in_group == 0 and group_rows > 1:
|
||||
if pe_row_in_group < group_rows - 1:
|
||||
with tl.scratch_scope():
|
||||
m_other = tl.recv(dir="intra_S", shape=m_local.shape, dtype="f16")
|
||||
l_other = tl.recv(dir="intra_S", shape=l_local.shape, dtype="f16")
|
||||
O_other = tl.recv(dir="intra_S", shape=O_local.shape, dtype="f16")
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_other, l_other, O_other, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
if pe_row_in_group > 0:
|
||||
tl.send(dir="intra_N", src=m_local)
|
||||
tl.send(dir="intra_N", src=l_local)
|
||||
tl.send(dir="intra_N", src=O_local)
|
||||
|
||||
# ── Final normalise + store (group root only) ──
|
||||
if pe_in_group == 0:
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr, O_final)
|
||||
|
||||
-209
@@ -1,209 +0,0 @@
|
||||
"""GQA prefill kernel: composite GEMM-only variant (2).
|
||||
|
||||
Identical mapping to ``_gqa_attention_prefill_short.py`` (Q-tile split +
|
||||
FA2 head fusion + IPCQ KV broadcast, unified A1/A2/A4/B). Difference vs
|
||||
first-level: Q·Kᵀ uses ``tl.composite(op="gemm")`` instead of ``tl.dot``.
|
||||
``P·V`` stays a plain ``tl.dot``: this is the GEMM-only tier, where
|
||||
softmax remains a primitive MATH chain with no ``softmax_merge`` fusion.
|
||||
|
||||
Full second-level fusion (with the ``softmax_merge`` prologue making
|
||||
P a pinned primary-out bound to a P·V composite) is variant (3) in
|
||||
``_gqa_attention_prefill_short_composite_fused.py``.
|
||||
|
||||
Three-variant comparison:
|
||||
(1) without composite : ``_gqa_attention_prefill_short.py``
|
||||
(2) with composite (GEMM-only, no fuse) : this file
|
||||
(3) with composite + softmax_merge fuse : ``…_composite_fused.py``
|
||||
|
||||
Shard addressing, layouts, and caller contract are identical to the
|
||||
first-level kernel (ADR-0011 D-VA1).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024
|
||||
|
||||
|
||||
def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
|
||||
h_q: int, h_kv: int, S_kv: int) -> None:
|
||||
"""Validate composite-prefill config — caller-side, sim-cost 0.
|
||||
|
||||
Mirrors first-level prefill ``_validate_config``.
|
||||
"""
|
||||
if kv_per_cube not in (1, 2, 4, 8):
|
||||
raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
|
||||
if P != 8:
|
||||
raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}")
|
||||
if h_q % h_kv != 0:
|
||||
raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
|
||||
if C != h_kv // kv_per_cube:
|
||||
raise ValueError(
|
||||
f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
|
||||
f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
|
||||
)
|
||||
group_size = P // kv_per_cube
|
||||
if T_q < group_size:
|
||||
raise ValueError(
|
||||
f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
|
||||
f"has q-tile work for the broadcast chain"
|
||||
)
|
||||
if S_kv % TILE_S_KV != 0:
|
||||
raise ValueError(
|
||||
f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
|
||||
)
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
m_new = tl.maximum(m_local, m_other)
|
||||
scale_old = tl.exp(m_local - m_new)
|
||||
scale_new = tl.exp(m_other - m_new)
|
||||
l_new = l_local * scale_old + l_other * scale_new
|
||||
O_new = O_local * scale_old + O_other * scale_new
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def gqa_attention_prefill_short_composite_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
P: int,
|
||||
kv_per_cube: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Composite-GEMM prefill — same mapping as first-level + tl.composite.
|
||||
|
||||
Caller must invoke ``_validate_config(...)`` first.
|
||||
"""
|
||||
group_size = P // kv_per_cube
|
||||
|
||||
G = h_q // h_kv
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
pe_in_group = pe_id % group_size
|
||||
group_id_in_cube = pe_id // group_size
|
||||
|
||||
q_start = (T_q * pe_in_group) // group_size
|
||||
q_end = (T_q * (pe_in_group + 1)) // group_size
|
||||
T_q_pe = q_end - q_start
|
||||
|
||||
n_tiles = S_kv // TILE_S_KV
|
||||
|
||||
Q_ROW_BYTES = G * d_head * 2
|
||||
KV_ROW_BYTES = d_head * 2
|
||||
K_TILE_BYTES = d_head * TILE_S_KV * 2
|
||||
K_HEAD_BYTES = n_tiles * K_TILE_BYTES
|
||||
V_HEAD_BYTES = S_kv * KV_ROW_BYTES
|
||||
|
||||
# Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
|
||||
q_base = (q_ptr
|
||||
+ cube_id * kv_per_cube * T_q * Q_ROW_BYTES
|
||||
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
|
||||
k_head_shard_base = (k_ptr
|
||||
+ cube_id * kv_per_cube * K_HEAD_BYTES
|
||||
+ group_id_in_cube * K_HEAD_BYTES)
|
||||
v_head_shard_base = (v_ptr
|
||||
+ cube_id * kv_per_cube * V_HEAD_BYTES
|
||||
+ group_id_in_cube * V_HEAD_BYTES)
|
||||
|
||||
Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16")
|
||||
|
||||
group_cols = min(4, group_size)
|
||||
group_rows = (group_size + group_cols - 1) // group_cols
|
||||
pe_col_in_group = pe_in_group % group_cols
|
||||
pe_row_in_group = pe_in_group // group_cols
|
||||
|
||||
# ── Tile 0 — broadcast + persistent (m, ℓ, O) ──
|
||||
if pe_in_group == 0:
|
||||
K_T = tl.load(k_head_shard_base,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.load(v_head_shard_base,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
if group_rows > 1:
|
||||
tl.send(dir="intra_S", src=K_T)
|
||||
tl.send(dir="intra_S", src=V)
|
||||
elif pe_col_in_group == 0 and pe_row_in_group > 0:
|
||||
K_T = tl.recv(dir="intra_N",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.recv(dir="intra_N",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
else:
|
||||
K_T = tl.recv(dir="intra_W",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.recv(dir="intra_W",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
|
||||
# Q·Kᵀ is a composite GEMM; P·V stays a primitive tl.dot in this
|
||||
# GEMM-only tier (no softmax_merge fusion — that is variant 3).
|
||||
scores = tl.composite(op="gemm", a=Q, b=K_T)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
exp_scores = tl.exp(centered)
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
O_local = tl.dot(exp_scores, V)
|
||||
|
||||
# ── Tiles 1..n_tiles-1 ──
|
||||
for tile_idx in range(1, n_tiles):
|
||||
with tl.scratch_scope():
|
||||
if pe_in_group == 0:
|
||||
K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
if group_rows > 1:
|
||||
tl.send(dir="intra_S", src=K_T_t)
|
||||
tl.send(dir="intra_S", src=V_t)
|
||||
elif pe_col_in_group == 0 and pe_row_in_group > 0:
|
||||
K_T_t = tl.recv(dir="intra_N",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.recv(dir="intra_N",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
else:
|
||||
K_T_t = tl.recv(dir="intra_W",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.recv(dir="intra_W",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
# Q·Kᵀ composite; P·V stays primitive (see comment above).
|
||||
scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
|
||||
m_tile = tl.max(scores_t, axis=-1)
|
||||
centered_t = scores_t - m_tile
|
||||
exp_scores_t = tl.exp(centered_t)
|
||||
l_tile = tl.sum(exp_scores_t, axis=-1)
|
||||
O_tile = tl.dot(exp_scores_t, V_t)
|
||||
m_new, l_new, O_new = _merge_running(
|
||||
m_local, l_local, O_local, m_tile, l_tile, O_tile, tl=tl,
|
||||
)
|
||||
tl.copy_to(m_local, m_new)
|
||||
tl.copy_to(l_local, l_new)
|
||||
tl.copy_to(O_local, O_new)
|
||||
|
||||
O_final = O_local / l_local
|
||||
o_base = (o_ptr
|
||||
+ cube_id * kv_per_cube * T_q * Q_ROW_BYTES
|
||||
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
|
||||
tl.store(o_base, O_final)
|
||||
-204
@@ -1,204 +0,0 @@
|
||||
"""GQA prefill kernel: composite + softmax_merge fused variant (3).
|
||||
|
||||
Identical mapping to ``_gqa_attention_prefill_short.py``. Difference vs
|
||||
the GEMM-only composite: per-tile softmax is folded into
|
||||
the P·V composite via the ``softmax_merge`` prologue recipe, eliminating
|
||||
the GEMM/MATH engine bubble.
|
||||
|
||||
Multi-cube (A1/A2/A4): non-root PEs receive K/V via IPCQ ``tl.recv``.
|
||||
Recv'd slots in on-chip memory are pinned (read in place as a composite
|
||||
operand, not DMA-streamed from HBM).
|
||||
|
||||
Three-variant comparison:
|
||||
(1) without composite : ``_gqa_attention_prefill_short.py``
|
||||
(2) with composite (GEMM-only, no fuse) : ``…_composite.py``
|
||||
(3) with composite + softmax_merge fuse : this file
|
||||
|
||||
Shard addressing, layouts, and caller contract are identical to the
|
||||
first-level kernel (ADR-0011 D-VA1).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
TILE_S_KV = 1024
|
||||
|
||||
|
||||
def _validate_config(*, kv_per_cube: int, T_q: int, P: int, C: int,
|
||||
h_q: int, h_kv: int, S_kv: int) -> None:
|
||||
"""Validate composite-prefill config — caller-side, sim-cost 0.
|
||||
|
||||
Mirrors first-level prefill ``_validate_config``.
|
||||
"""
|
||||
if kv_per_cube not in (1, 2, 4, 8):
|
||||
raise ValueError(f"kv_per_cube must be in {{1,2,4,8}}; got {kv_per_cube}")
|
||||
if P != 8:
|
||||
raise ValueError(f"2x4 mesh broadcast geometry requires P == 8; got {P}")
|
||||
if h_q % h_kv != 0:
|
||||
raise ValueError(f"GQA group G = h_q/h_kv must be integer; got h_q={h_q}, h_kv={h_kv}")
|
||||
if C != h_kv // kv_per_cube:
|
||||
raise ValueError(
|
||||
f"C must equal h_kv/kv_per_cube = {h_kv // kv_per_cube}; "
|
||||
f"got C={C} (h_kv={h_kv}, kv_per_cube={kv_per_cube})"
|
||||
)
|
||||
group_size = P // kv_per_cube
|
||||
if T_q < group_size:
|
||||
raise ValueError(
|
||||
f"T_q ({T_q}) must be >= group_size ({group_size}) so every PE "
|
||||
f"has q-tile work for the broadcast chain"
|
||||
)
|
||||
if S_kv % TILE_S_KV != 0:
|
||||
raise ValueError(
|
||||
f"S_kv ({S_kv}) must be a multiple of TILE_S_KV ({TILE_S_KV})"
|
||||
)
|
||||
|
||||
|
||||
def _merge_running(m_local, l_local, O_local, m_other, l_other, O_other, *, tl):
|
||||
m_new = tl.maximum(m_local, m_other)
|
||||
scale_old = tl.exp(m_local - m_new)
|
||||
scale_new = tl.exp(m_other - m_new)
|
||||
l_new = l_local * scale_old + l_other * scale_new
|
||||
O_new = O_local * scale_old + O_other * scale_new
|
||||
return m_new, l_new, O_new
|
||||
|
||||
|
||||
def gqa_attention_prefill_short_composite_fused_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
P: int,
|
||||
kv_per_cube: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Composite-GEMM prefill — same mapping as first-level + tl.composite.
|
||||
|
||||
Caller must invoke ``_validate_config(...)`` first.
|
||||
"""
|
||||
group_size = P // kv_per_cube
|
||||
|
||||
G = h_q // h_kv
|
||||
pe_id = tl.program_id(axis=0)
|
||||
cube_id = tl.program_id(axis=1)
|
||||
pe_in_group = pe_id % group_size
|
||||
group_id_in_cube = pe_id // group_size
|
||||
|
||||
q_start = (T_q * pe_in_group) // group_size
|
||||
q_end = (T_q * (pe_in_group + 1)) // group_size
|
||||
T_q_pe = q_end - q_start
|
||||
|
||||
n_tiles = S_kv // TILE_S_KV
|
||||
|
||||
Q_ROW_BYTES = G * d_head * 2
|
||||
KV_ROW_BYTES = d_head * 2
|
||||
K_TILE_BYTES = d_head * TILE_S_KV * 2
|
||||
K_HEAD_BYTES = n_tiles * K_TILE_BYTES
|
||||
V_HEAD_BYTES = S_kv * KV_ROW_BYTES
|
||||
|
||||
# Global VA per ADR-0011 D-VA1: kernel computes its own shard base.
|
||||
q_base = (q_ptr
|
||||
+ cube_id * kv_per_cube * T_q * Q_ROW_BYTES
|
||||
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
|
||||
k_head_shard_base = (k_ptr
|
||||
+ cube_id * kv_per_cube * K_HEAD_BYTES
|
||||
+ group_id_in_cube * K_HEAD_BYTES)
|
||||
v_head_shard_base = (v_ptr
|
||||
+ cube_id * kv_per_cube * V_HEAD_BYTES
|
||||
+ group_id_in_cube * V_HEAD_BYTES)
|
||||
|
||||
Q = tl.load(q_base, shape=(T_q_pe * G, d_head), dtype="f16")
|
||||
|
||||
group_cols = min(4, group_size)
|
||||
group_rows = (group_size + group_cols - 1) // group_cols
|
||||
pe_col_in_group = pe_in_group % group_cols
|
||||
pe_row_in_group = pe_in_group // group_cols
|
||||
|
||||
# ── Tile 0 — broadcast + persistent (m, ℓ, O) ──
|
||||
if pe_in_group == 0:
|
||||
K_T = tl.load(k_head_shard_base,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.load(v_head_shard_base,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
if group_rows > 1:
|
||||
tl.send(dir="intra_S", src=K_T)
|
||||
tl.send(dir="intra_S", src=V)
|
||||
elif pe_col_in_group == 0 and pe_row_in_group > 0:
|
||||
K_T = tl.recv(dir="intra_N",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.recv(dir="intra_N",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
else:
|
||||
K_T = tl.recv(dir="intra_W",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V = tl.recv(dir="intra_W",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
tl.send(dir="intra_E", src=K_T)
|
||||
tl.send(dir="intra_E", src=V)
|
||||
|
||||
# Tile 0 — primitives establish (m, ℓ, O); recipe fusion enters tile 1+.
|
||||
scores = tl.dot(Q, K_T)
|
||||
m_local = tl.max(scores, axis=-1)
|
||||
centered = scores - m_local
|
||||
exp_scores = tl.exp(centered)
|
||||
l_local = tl.sum(exp_scores, axis=-1)
|
||||
O_local = tl.dot(exp_scores, V)
|
||||
|
||||
# ── Tiles 1..n_tiles-1 ──
|
||||
for tile_idx in range(1, n_tiles):
|
||||
with tl.scratch_scope():
|
||||
if pe_in_group == 0:
|
||||
K_T_t = tl.load(k_head_shard_base + tile_idx * K_TILE_BYTES,
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.load(v_head_shard_base + tile_idx * TILE_S_KV * KV_ROW_BYTES,
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
if group_rows > 1:
|
||||
tl.send(dir="intra_S", src=K_T_t)
|
||||
tl.send(dir="intra_S", src=V_t)
|
||||
elif pe_col_in_group == 0 and pe_row_in_group > 0:
|
||||
K_T_t = tl.recv(dir="intra_N",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.recv(dir="intra_N",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if group_cols > 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
else:
|
||||
K_T_t = tl.recv(dir="intra_W",
|
||||
shape=(d_head, TILE_S_KV), dtype="f16")
|
||||
V_t = tl.recv(dir="intra_W",
|
||||
shape=(TILE_S_KV, d_head), dtype="f16")
|
||||
if pe_col_in_group < group_cols - 1:
|
||||
tl.send(dir="intra_E", src=K_T_t)
|
||||
tl.send(dir="intra_E", src=V_t)
|
||||
# Two-composite fusion: Q·Kᵀ composite → softmax_merge prologue
|
||||
# binds P (pinned primary-out) to the P·V composite, folding
|
||||
# the new tile's contribution into O_local.
|
||||
scores_t = tl.composite(op="gemm", a=Q, b=K_T_t)
|
||||
tl.composite(
|
||||
prologue=[{"op": "softmax_merge", "s": scores_t,
|
||||
"m": m_local, "l": l_local, "O": O_local}],
|
||||
op="gemm", b=V_t, out=O_local,
|
||||
epilogue=[{"op": "add", "other": O_local}],
|
||||
)
|
||||
|
||||
O_final = O_local / l_local
|
||||
o_base = (o_ptr
|
||||
+ cube_id * kv_per_cube * T_q * Q_ROW_BYTES
|
||||
+ (group_id_in_cube * T_q + q_start) * Q_ROW_BYTES)
|
||||
tl.store(o_base, O_final)
|
||||
@@ -28,6 +28,9 @@ from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (
|
||||
run_sweep as _run_composite_sweep,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_streaming import (
|
||||
run_sweep as _run_decode_streaming_sweep,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
|
||||
run_sweep as _run_prefill_cb_sweep,
|
||||
)
|
||||
@@ -66,6 +69,7 @@ def run(torch) -> None:
|
||||
"decode": _run_decode_sweep,
|
||||
"composite": _run_composite_sweep,
|
||||
"prefill_cb": _run_prefill_cb_sweep,
|
||||
"decode_streaming": _run_decode_streaming_sweep,
|
||||
}
|
||||
unknown = [s for s in sweeps if s not in runners]
|
||||
if unknown:
|
||||
|
||||
@@ -117,13 +117,17 @@ class PeCpuComponent(ComponentBase):
|
||||
pe_exec_start = env.now
|
||||
scheduler_id = f"{self._pe_prefix}.pe_scheduler"
|
||||
|
||||
# ADR-0020 greenlet execution — always used so IPCQ / ring credits /
|
||||
# fabric-transfer sim events fire regardless of data-mode. KernelRunner
|
||||
# guards its store reads on ``self._store is not None``.
|
||||
# Choose execution mode: greenlet (ADR-0020) or legacy command-list
|
||||
store = getattr(self.ctx, "memory_store", None) if self.ctx else None
|
||||
composite_results = yield from self._execute_greenlet(
|
||||
env, kernel_fn, kernel_args, num_programs, scheduler_id, store,
|
||||
)
|
||||
|
||||
if store is not None:
|
||||
composite_results = yield from self._execute_greenlet(
|
||||
env, kernel_fn, kernel_args, num_programs, scheduler_id, store,
|
||||
)
|
||||
else:
|
||||
composite_results = yield from self._execute_legacy(
|
||||
env, kernel_fn, kernel_args, num_programs, scheduler_id,
|
||||
)
|
||||
|
||||
# Record PE-internal execution time
|
||||
txn.result_data["pe_exec_ns"] = env.now - pe_exec_start
|
||||
|
||||
@@ -111,13 +111,6 @@ class PathRouter:
|
||||
self._adj_all: dict[str, list[tuple[str, float]]] = defaultdict(list)
|
||||
self._adj_mcpu_dma: dict[str, list[tuple[str, float]]] = defaultdict(list)
|
||||
self._adj_local: dict[str, list[tuple[str, float]]] = defaultdict(list)
|
||||
# Memoize path lookups: adj dicts are built once here and never
|
||||
# mutated (topology is static per ADR-0006 / SPEC §0.1), so
|
||||
# (id(adj), start, goal) is a stable cache key for the router's
|
||||
# lifetime. Callers use the returned path list read-only.
|
||||
self._path_cache: dict[
|
||||
tuple[int, str, str], tuple[list[str], float]
|
||||
] = {}
|
||||
for e in graph.edges:
|
||||
w = e.routing_weight_mm if e.routing_weight_mm is not None else e.distance_mm
|
||||
self._adj_all[e.src].append((e.dst, w))
|
||||
@@ -192,14 +185,8 @@ class PathRouter:
|
||||
start: str,
|
||||
goal: str,
|
||||
) -> tuple[list[str], float]:
|
||||
cache_key = (id(adj), start, goal)
|
||||
cached = self._path_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
if start == goal:
|
||||
result = ([start], 0.0)
|
||||
self._path_cache[cache_key] = result
|
||||
return result
|
||||
return [start], 0.0
|
||||
best: dict[str, float] = {start: 0.0}
|
||||
prev: dict[str, str] = {}
|
||||
heap: list[tuple[float, str]] = [(0.0, start)]
|
||||
@@ -213,9 +200,7 @@ class PathRouter:
|
||||
cur = prev[cur]
|
||||
path.append(start)
|
||||
path.reverse()
|
||||
result = (path, d)
|
||||
self._path_cache[cache_key] = result
|
||||
return result
|
||||
return path, d
|
||||
if d > best.get(node, float("inf")):
|
||||
continue
|
||||
for neighbor, edge_dist in adj[node]:
|
||||
|
||||
@@ -570,14 +570,8 @@ class RuntimeContext:
|
||||
h = self.submit(msg)
|
||||
self.wait(h)
|
||||
|
||||
# Submit MemoryWriteMsg per shard (deploy data to device). Gated on
|
||||
# memory_store presence: under enable_data=False there is no
|
||||
# MemoryStore to populate and no Phase 2 DataExecutor replay, so the
|
||||
# per-shard sim events are pure wall-clock overhead with no effect
|
||||
# on reported kernel latency (Yangwook's max(t_end) - min(t_start)
|
||||
# formula excludes ops before the kernel starts).
|
||||
store = getattr(self.engine, "_memory_store", None)
|
||||
if pattern is not None and store is not None:
|
||||
# Submit MemoryWriteMsg per shard (deploy data to device)
|
||||
if pattern is not None:
|
||||
for shard in handle.shards:
|
||||
h = self.submit(MemoryWriteMsg(
|
||||
correlation_id=self.correlation_id,
|
||||
@@ -597,7 +591,8 @@ class RuntimeContext:
|
||||
# VA; Phase 2 DataExecutor reads via the addresses captured in
|
||||
# op_log (VA for tl.load). Without this, zero-init tensors are
|
||||
# invisible to kernels in Phase 2.
|
||||
if pattern == "zero" and handle.va_base:
|
||||
store = getattr(self.engine, "_memory_store", None)
|
||||
if store is not None and pattern == "zero" and handle.va_base:
|
||||
import numpy as np
|
||||
from kernbench.runtime_api.tensor import _numpy_dtype
|
||||
np_dtype = _numpy_dtype(dtype)
|
||||
|
||||
@@ -53,16 +53,14 @@ class GraphEngine:
|
||||
self._events: dict[str, simpy.Event] = {}
|
||||
self._counter = 0
|
||||
overrides = component_overrides or {}
|
||||
# ADR-0020: optional data execution support. OpLogger is always
|
||||
# created so op_log-based latency (max(t_end) - min(t_start)) is
|
||||
# available in both modes; MemoryStore is created only when
|
||||
# enable_data=True so Phase 2 DataExecutor replay is gated on it.
|
||||
# ADR-0020: optional data execution support
|
||||
self._op_logger = None
|
||||
self._memory_store = None
|
||||
if enable_data:
|
||||
from kernbench.sim_engine.memory_store import MemoryStore
|
||||
from kernbench.sim_engine.op_log import OpLogger
|
||||
self._memory_store = MemoryStore()
|
||||
from kernbench.sim_engine.op_log import OpLogger
|
||||
self._op_logger = OpLogger(memory_store=self._memory_store)
|
||||
self._op_logger = OpLogger(memory_store=self._memory_store)
|
||||
# Cursor for incremental Phase 2 replay (ADR-0020 D6).
|
||||
# SimPy env.now is monotonic so newly logged records always sort
|
||||
# to the tail; the cursor remains valid across waits.
|
||||
|
||||
@@ -697,10 +697,6 @@ class TLContext:
|
||||
nbytes=self._nbytes(shape, dtype),
|
||||
data=data,
|
||||
space=slot_space,
|
||||
# On-chip (tcm/sram) IPCQ slot: read in place as a composite
|
||||
# operand, not DMA_READ of its non-HBM slot address. An HBM
|
||||
# slot is a valid PA, so it stays unpinned (DMA streams it).
|
||||
pinned=slot_space != "hbm",
|
||||
)
|
||||
return self._make_handle(addr=0, shape=shape, dtype=dtype)
|
||||
|
||||
@@ -742,7 +738,6 @@ class TLContext:
|
||||
nbytes=self._nbytes(shape, dtype),
|
||||
data=None,
|
||||
space=slot_space,
|
||||
pinned=slot_space != "hbm",
|
||||
)
|
||||
return self._make_handle(addr=0, shape=shape, dtype=dtype)
|
||||
|
||||
@@ -1110,7 +1105,6 @@ class TLContext:
|
||||
nbytes=self._nbytes(handle.cmd.shape, handle.cmd.dtype),
|
||||
data=data,
|
||||
space=slot_space,
|
||||
pinned=slot_space != "hbm",
|
||||
)
|
||||
handle.resolved = True
|
||||
handle.result = th
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,459 +0,0 @@
|
||||
"""Auto-explore parallelism configuration space and rank by Pareto frontier.
|
||||
|
||||
Extends the existing ``autosuggest`` (memory-only, single winner) to the
|
||||
full 9-knob search (CP, TP, PP, DP, kv_shard_mode, ffn_shard_scope,
|
||||
tp_placement, cp_placement, cp_ring_variant) with four objectives:
|
||||
|
||||
- min total_latency_ns (single-request decode/prefill step)
|
||||
- max throughput_tok_s (naive tokens/sec = 1 / latency for single-req)
|
||||
- max efficiency_score (geo-mean of compute + BW utilization)
|
||||
- min pes_used
|
||||
|
||||
Reuses ``stage_latencies`` and ``memory_layout`` as the physics; adds
|
||||
enumeration + 4D Pareto sort on top.
|
||||
|
||||
Analytical model runs in ~microseconds per config, so the full sweep
|
||||
(~5-10k feasible configs after pruning) completes in a few seconds.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass, field, replace
|
||||
|
||||
from .memory_layout import compute_memory
|
||||
from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig
|
||||
from .stage_latencies import all_ffn_stages, all_stages
|
||||
|
||||
# ── Parallelism sensitivity sweep values (see compute_parallelism_sensitivity)
|
||||
# Multiples of 2 (finer grid than powers of 2 alone), capped at values that
|
||||
# are physically meaningful for the modelled topology.
|
||||
_PARALLELISM_SWEEP_VALUES = {
|
||||
"cp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64, 96, 128, 192, 256),
|
||||
"tp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 24, 32, 48, 64),
|
||||
"pp": (1, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 32),
|
||||
"dp": (1, 2, 4, 6, 8, 12, 16),
|
||||
"ep": (1, 2, 4, 6, 8, 12, 16, 32),
|
||||
}
|
||||
|
||||
|
||||
# ── Search space ─────────────────────────────────────────────────────
|
||||
|
||||
_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
|
||||
_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
|
||||
_PP_OPTIONS = (1, 2, 4, 8, 16)
|
||||
_DP_OPTIONS = (1, 2, 4)
|
||||
_KV_SHARD_MODES = ("split", "replicate")
|
||||
_FFN_SHARD_SCOPES = ("TP", "TP+CP", "TP+CP+DP")
|
||||
_TP_PLACEMENTS = ("pe", "cube")
|
||||
_CP_PLACEMENTS = ("cube", "pe")
|
||||
_CP_RING_VARIANTS = ("kv", "qoml")
|
||||
|
||||
|
||||
# ── Result types ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConfigScore:
|
||||
"""A single (config, computed-metrics) tuple. All floats in SI units
|
||||
(seconds, tokens/sec, dimensionless) unless suffixed otherwise."""
|
||||
|
||||
# ── Config (the 9 knobs) ─────
|
||||
cp: int
|
||||
tp: int
|
||||
pp: int
|
||||
dp: int
|
||||
kv_shard_mode: str
|
||||
ffn_shard_scope: str
|
||||
tp_placement: str
|
||||
cp_placement: str
|
||||
cp_ring_variant: str
|
||||
|
||||
# ── Objectives ─────
|
||||
total_latency_ns: float # ↓ minimize
|
||||
throughput_tok_s: float # ↑ maximize
|
||||
efficiency_score: float # ↑ maximize (0..1)
|
||||
pes_used: int # ↓ minimize
|
||||
|
||||
# ── Info-only (not ranked on) ─────
|
||||
hbm_utilization: float # bytes_used / budget (0..1+; may exceed 1 if over-budget)
|
||||
weights_gb: float
|
||||
kv_gb: float
|
||||
transient_gb: float
|
||||
sips_used: int
|
||||
fits_memory: bool
|
||||
placement_valid: bool
|
||||
reason: str = ""
|
||||
|
||||
@property
|
||||
def latency_us(self) -> float:
|
||||
return self.total_latency_ns / 1e3
|
||||
|
||||
@property
|
||||
def latency_ms(self) -> float:
|
||||
return self.total_latency_ns / 1e6
|
||||
|
||||
def as_topology(self, s_kv: int, mode: str) -> TopologyConfig:
|
||||
"""Reconstruct the TopologyConfig this score was computed for."""
|
||||
return TopologyConfig(
|
||||
cp=self.cp, tp=self.tp, pp=self.pp, dp=self.dp,
|
||||
s_kv=s_kv, mode=mode,
|
||||
kv_shard_mode=self.kv_shard_mode,
|
||||
ffn_shard_scope=self.ffn_shard_scope,
|
||||
tp_placement=self.tp_placement,
|
||||
cp_placement=self.cp_placement,
|
||||
cp_ring_variant=self.cp_ring_variant,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AutoExploreResult:
|
||||
model_name: str
|
||||
s_kv: int
|
||||
mode: str
|
||||
total_enumerated: int # after basic domain pruning
|
||||
total_feasible: int # after memory + placement checks
|
||||
all_scores: list[ConfigScore] = field(default_factory=list) # every feasible config
|
||||
pareto_scores: list[ConfigScore] = field(default_factory=list) # non-dominated set
|
||||
|
||||
|
||||
# ── Enumeration + pruning ────────────────────────────────────────────
|
||||
|
||||
|
||||
def enumerate_configs(
|
||||
model: ModelConfig,
|
||||
s_kv: int,
|
||||
mode: str,
|
||||
) -> Iterator[TopologyConfig]:
|
||||
"""Yield every domain-valid TopologyConfig.
|
||||
|
||||
Domain rules (fast pruning; no memory check yet):
|
||||
- PP ≤ model.layers (can't have more stages than layers)
|
||||
- TP ≤ 4 × model.h_q (unrealistic head-dim splits above this)
|
||||
- Skip ffn_shard_scope containing 'DP' when DP=1 (redundant with plain 'TP+CP')
|
||||
- Skip cp_ring_variant='qoml' when CP=1 (no ring, variant is a no-op)
|
||||
"""
|
||||
for cp in _CP_OPTIONS:
|
||||
for tp in _TP_OPTIONS:
|
||||
if tp > 4 * model.h_q:
|
||||
continue
|
||||
for pp in _PP_OPTIONS:
|
||||
if pp > model.layers:
|
||||
continue
|
||||
for dp in _DP_OPTIONS:
|
||||
for kv_mode in _KV_SHARD_MODES:
|
||||
for ffn_scope in _FFN_SHARD_SCOPES:
|
||||
if "DP" in ffn_scope and dp == 1:
|
||||
continue
|
||||
for tp_place in _TP_PLACEMENTS:
|
||||
for cp_place in _CP_PLACEMENTS:
|
||||
for cp_ring in _CP_RING_VARIANTS:
|
||||
if cp == 1 and cp_ring == "qoml":
|
||||
continue
|
||||
yield TopologyConfig(
|
||||
cp=cp, tp=tp, pp=pp, dp=dp,
|
||||
s_kv=s_kv, mode=mode,
|
||||
kv_shard_mode=kv_mode,
|
||||
ffn_shard_scope=ffn_scope,
|
||||
tp_placement=tp_place,
|
||||
cp_placement=cp_place,
|
||||
cp_ring_variant=cp_ring,
|
||||
)
|
||||
|
||||
|
||||
# ── Scoring ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _sum_visible_latency(
|
||||
cfg: FullConfig,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True,
|
||||
) -> float:
|
||||
"""Total single-request latency (seconds) across all model layers.
|
||||
|
||||
A single request traverses every layer sequentially, whether the layers
|
||||
sit on one PP stage or are spread across many:
|
||||
- PP=1 : one rank holds all L layers, latency = L × per_layer
|
||||
- PP=K : K ranks each hold L/K layers, but request crosses all K
|
||||
stages sequentially → still L × per_layer
|
||||
|
||||
PP therefore does NOT reduce single-request latency; it only improves
|
||||
throughput under batching. This function is the single-request cost, so
|
||||
we multiply by full model.layers regardless of PP.
|
||||
|
||||
Scope selection (both default True — full per-token cost):
|
||||
- ``include_attention=True, include_ffn=True`` → full transformer
|
||||
- ``include_attention=True, include_ffn=False`` → attention only
|
||||
- ``include_attention=False, include_ffn=True`` → FFN / MoE only
|
||||
- ``include_attention=False, include_ffn=False`` → zero (rejected caller-side)
|
||||
"""
|
||||
attn = sum(s.visible_s for s in all_stages(cfg)) if include_attention else 0.0
|
||||
ffn = sum(s.visible_s for s in all_ffn_stages(cfg)) if include_ffn else 0.0
|
||||
per_layer = attn + ffn
|
||||
return per_layer * cfg.model.layers
|
||||
|
||||
|
||||
def _efficiency(cfg: FullConfig, latency_s: float,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True) -> float:
|
||||
"""Geo-mean of compute-util and BW-util. Range ~ (0, 1].
|
||||
|
||||
- compute_util = achieved_flops / (peak_flops × pes × latency)
|
||||
- bw_util = achieved_bytes / (peak_bw × pes × latency)
|
||||
|
||||
Scope flags mirror :func:`_sum_visible_latency`.
|
||||
"""
|
||||
if latency_s <= 0:
|
||||
return 0.0
|
||||
attn = all_stages(cfg) if include_attention else []
|
||||
ffn = all_ffn_stages(cfg) if include_ffn else []
|
||||
layers = math.ceil(cfg.model.layers / cfg.topo.pp)
|
||||
total_flops = layers * sum(s.flops for s in attn + ffn)
|
||||
total_bytes = layers * sum(s.mem_bytes for s in attn + ffn)
|
||||
|
||||
pes = cfg.topo.total_pes
|
||||
peak_flops = cfg.machine.peak_flops * pes
|
||||
peak_bw = cfg.machine.bw_hbm * pes
|
||||
|
||||
compute_util = total_flops / (peak_flops * latency_s) if peak_flops > 0 else 0.0
|
||||
bw_util = total_bytes / (peak_bw * latency_s) if peak_bw > 0 else 0.0
|
||||
compute_util = min(1.0, max(0.0, compute_util))
|
||||
bw_util = min(1.0, max(0.0, bw_util))
|
||||
# Geo-mean; if either is 0 the score is 0 (avoids overrewarding lopsided configs).
|
||||
return math.sqrt(compute_util * bw_util)
|
||||
|
||||
|
||||
def score_config(cfg: FullConfig,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True) -> ConfigScore:
|
||||
"""Compute all 4 objectives + info fields for one config.
|
||||
|
||||
Feasibility (memory + placement) is stored but does NOT gate scoring —
|
||||
infeasible configs get returned with fits_memory=False so callers can
|
||||
filter or display them.
|
||||
|
||||
Scope flags (default: full transformer) restrict *latency + efficiency*
|
||||
to attention only, FFN only, or both. Memory feasibility is unchanged —
|
||||
still checks weights+KV+transient fit in per-PE HBM since the model
|
||||
still exists physically regardless of what the caller is scoring.
|
||||
"""
|
||||
mem = compute_memory(cfg)
|
||||
placement_ok = cfg.topo.placement_valid
|
||||
|
||||
latency_s = _sum_visible_latency(
|
||||
cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
throughput = 1.0 / latency_s if latency_s > 0 else 0.0
|
||||
efficiency = (
|
||||
_efficiency(cfg, latency_s,
|
||||
include_attention=include_attention, include_ffn=include_ffn)
|
||||
if latency_s > 0 else 0.0
|
||||
)
|
||||
fits = not mem.over_budget
|
||||
|
||||
reason = ""
|
||||
if not fits:
|
||||
reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
|
||||
f"exceeds per-PE budget ({mem.budget_bytes/1e9:.2f} GB)")
|
||||
elif not placement_ok:
|
||||
reason = (f"intra-cube demand ({cfg.topo.intra_cube_dims}) "
|
||||
f"exceeds PEs/cube ({cfg.topo.pes_per_cube_hw})")
|
||||
|
||||
return ConfigScore(
|
||||
cp=cfg.topo.cp, tp=cfg.topo.tp, pp=cfg.topo.pp, dp=cfg.topo.dp,
|
||||
kv_shard_mode=cfg.topo.kv_shard_mode,
|
||||
ffn_shard_scope=cfg.topo.ffn_shard_scope,
|
||||
tp_placement=cfg.topo.tp_placement,
|
||||
cp_placement=cfg.topo.cp_placement,
|
||||
cp_ring_variant=cfg.topo.cp_ring_variant,
|
||||
total_latency_ns=latency_s * 1e9,
|
||||
throughput_tok_s=throughput,
|
||||
efficiency_score=efficiency,
|
||||
pes_used=cfg.topo.total_pes,
|
||||
hbm_utilization=mem.used_bytes / mem.budget_bytes if mem.budget_bytes > 0 else 0.0,
|
||||
weights_gb=mem.weights_bytes / 1e9,
|
||||
kv_gb=mem.kv_cache_bytes / 1e9,
|
||||
transient_gb=mem.transient_bytes / 1e9,
|
||||
sips_used=cfg.topo.sips_used,
|
||||
fits_memory=fits,
|
||||
placement_valid=placement_ok,
|
||||
reason=reason,
|
||||
)
|
||||
|
||||
|
||||
# ── Pareto sort ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _dominates(a: ConfigScore, b: ConfigScore) -> bool:
|
||||
"""True iff a is no worse than b on every axis AND strictly better on ≥ 1.
|
||||
|
||||
Axes (with direction) — 3D Pareto:
|
||||
total_latency_ns ↓ (a ≤ b)
|
||||
pes_used ↓ (a ≤ b) — proxy for cost / deployment size
|
||||
efficiency_score ↑ (a ≥ b) — geo-mean of compute+BW utilization
|
||||
|
||||
Note: ``throughput_tok_s`` is deliberately NOT a Pareto axis. For a
|
||||
single-request analysis, throughput = 1 / latency, so it's collinear
|
||||
with latency and would collapse the frontier. It stays on ConfigScore
|
||||
for display but doesn't participate in domination.
|
||||
"""
|
||||
no_worse = (
|
||||
a.total_latency_ns <= b.total_latency_ns
|
||||
and a.pes_used <= b.pes_used
|
||||
and a.efficiency_score >= b.efficiency_score
|
||||
)
|
||||
if not no_worse:
|
||||
return False
|
||||
strictly_better = (
|
||||
a.total_latency_ns < b.total_latency_ns
|
||||
or a.pes_used < b.pes_used
|
||||
or a.efficiency_score > b.efficiency_score
|
||||
)
|
||||
return strictly_better
|
||||
|
||||
|
||||
def pareto_frontier(scores: list[ConfigScore]) -> list[ConfigScore]:
|
||||
"""Extract non-dominated set on (latency↓, pe↓, throughput↑, efficiency↑).
|
||||
|
||||
Only feasible configs (fits_memory AND placement_valid) participate; the
|
||||
non-feasible are excluded from the frontier (they can still be returned in
|
||||
``all_scores`` for informational display).
|
||||
|
||||
O(N²) — fine for N up to ~50k with early exits.
|
||||
"""
|
||||
feasible = [s for s in scores if s.fits_memory and s.placement_valid]
|
||||
frontier: list[ConfigScore] = []
|
||||
for i, a in enumerate(feasible):
|
||||
dominated = False
|
||||
for j, b in enumerate(feasible):
|
||||
if i == j:
|
||||
continue
|
||||
if _dominates(b, a):
|
||||
dominated = True
|
||||
break
|
||||
if not dominated:
|
||||
frontier.append(a)
|
||||
return frontier
|
||||
|
||||
|
||||
# ── Parallelism sensitivity ─────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParallelismSensitivityRow:
|
||||
"""One knob's sweep result: latency + fit at every tested value.
|
||||
|
||||
All *other* parallelism knobs (including HW) are held at their baseline
|
||||
values, so this isolates the effect of just this one knob.
|
||||
"""
|
||||
knob: str # "cp" / "tp" / "pp" / "dp" / "ep"
|
||||
values: list[int] # tested values (from _PARALLELISM_SWEEP_VALUES)
|
||||
latencies_ns: list[float] # one per value; NaN when infeasible
|
||||
fits_flags: list[bool] # per-value memory+placement feasibility
|
||||
baseline_value: int # what the baseline config uses for this knob
|
||||
baseline_latency_ns: float
|
||||
|
||||
|
||||
def compute_parallelism_sensitivity(
|
||||
baseline: ConfigScore,
|
||||
model: ModelConfig,
|
||||
machine: MachineParams,
|
||||
s_kv: int,
|
||||
mode: str,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True,
|
||||
) -> list[ParallelismSensitivityRow]:
|
||||
"""For each parallelism knob (CP, TP, PP, DP, EP), sweep its values
|
||||
holding the OTHER knobs fixed at the baseline. Reports latency + memory
|
||||
fit per value.
|
||||
|
||||
Answers: 'if I only change this one knob, how does latency and memory
|
||||
fit change?' — the complement to auto_hardware.compute_sensitivity
|
||||
which does the same for hardware knobs.
|
||||
"""
|
||||
baseline_topo = baseline.as_topology(s_kv, mode)
|
||||
|
||||
rows: list[ParallelismSensitivityRow] = []
|
||||
for knob, values in _PARALLELISM_SWEEP_VALUES.items():
|
||||
baseline_val = getattr(baseline_topo, knob, 1)
|
||||
# EP isn't stored on ConfigScore's as_topology output today (dp
|
||||
# attribute exists but ep does not always). Fall back to 1.
|
||||
latencies: list[float] = []
|
||||
fits: list[bool] = []
|
||||
for v in values:
|
||||
# Skip nonsensical values that would violate hard bounds.
|
||||
if knob == "pp" and v > model.layers:
|
||||
latencies.append(float("nan"))
|
||||
fits.append(False)
|
||||
continue
|
||||
if knob == "tp" and v > 4 * model.h_q:
|
||||
latencies.append(float("nan"))
|
||||
fits.append(False)
|
||||
continue
|
||||
# Build a variant TopologyConfig with just this one knob changed.
|
||||
trial = replace(baseline_topo, **{knob: v})
|
||||
cfg = FullConfig(model=model, topo=trial, machine=machine)
|
||||
score = score_config(
|
||||
cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
latencies.append(score.total_latency_ns)
|
||||
fits.append(score.fits_memory and score.placement_valid)
|
||||
rows.append(ParallelismSensitivityRow(
|
||||
knob=knob,
|
||||
values=list(values),
|
||||
latencies_ns=latencies,
|
||||
fits_flags=fits,
|
||||
baseline_value=int(baseline_val),
|
||||
baseline_latency_ns=baseline.total_latency_ns,
|
||||
))
|
||||
return rows
|
||||
|
||||
|
||||
# ── Top-level driver ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def run_auto_explore(
|
||||
model: ModelConfig,
|
||||
machine: MachineParams,
|
||||
s_kv: int,
|
||||
mode: str = "decode",
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True,
|
||||
b: int = 1,
|
||||
) -> AutoExploreResult:
|
||||
"""Enumerate all configs, score each, extract Pareto frontier.
|
||||
|
||||
Returns both the full ``all_scores`` list (for the table view) and the
|
||||
``pareto_scores`` subset (for the scatter/highlight view). Both are sorted
|
||||
by total_latency_ns ascending.
|
||||
|
||||
Scope: at least one of ``include_attention`` / ``include_ffn`` must be
|
||||
True. Setting both False would give zero latency for every config —
|
||||
the caller should not do that.
|
||||
"""
|
||||
all_scores: list[ConfigScore] = []
|
||||
total_enumerated = 0
|
||||
for topo in enumerate_configs(model, s_kv, mode):
|
||||
# Inject the caller's batch B into every candidate.
|
||||
topo.b = b
|
||||
total_enumerated += 1
|
||||
cfg = FullConfig(model=model, topo=topo, machine=machine)
|
||||
all_scores.append(score_config(
|
||||
cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
))
|
||||
|
||||
all_scores.sort(key=lambda s: s.total_latency_ns)
|
||||
pareto = pareto_frontier(all_scores)
|
||||
pareto.sort(key=lambda s: s.total_latency_ns)
|
||||
|
||||
feasible_count = sum(1 for s in all_scores if s.fits_memory and s.placement_valid)
|
||||
|
||||
return AutoExploreResult(
|
||||
model_name=model.name,
|
||||
s_kv=s_kv,
|
||||
mode=mode,
|
||||
total_enumerated=total_enumerated,
|
||||
total_feasible=feasible_count,
|
||||
all_scores=all_scores,
|
||||
pareto_scores=pareto,
|
||||
)
|
||||
@@ -1,429 +0,0 @@
|
||||
"""Joint hardware × parallelism exploration.
|
||||
|
||||
For a fixed model + workload, sweep both:
|
||||
- hardware parameters (MachineParams: pe_hbm_gb, bw_hbm_gbs, peak_tflops_f16,
|
||||
bw_intra_gbs, bw_inter_gbs, bw_intersip_gbs), and
|
||||
- parallelism knobs (a reduced subset of auto_explore's 9 knobs)
|
||||
|
||||
Then rank on:
|
||||
- latency ↓
|
||||
- hardware cost proxy ↓ (normalized sum of knob-over-default)
|
||||
|
||||
Also computes per-knob sensitivity — how much latency drops when each
|
||||
hardware knob is doubled from its baseline value. Useful for HW co-design
|
||||
("which knob to invest in first?").
|
||||
|
||||
The default 'balanced' depth: 64 HW candidates × ~2k parallelism configs =
|
||||
~130k joint evaluations, ~1-5 s at analytical speeds. Coarse and
|
||||
two-stage variants trade coverage for speed.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass, field, replace
|
||||
|
||||
from .auto_explore import ConfigScore, score_config
|
||||
from .autosuggest import auto_suggest
|
||||
from .memory_layout import compute_memory
|
||||
from .model_config import FullConfig, MachineParams, ModelConfig, TopologyConfig
|
||||
|
||||
|
||||
# ── Hardware search space ────────────────────────────────────────────
|
||||
#
|
||||
# For each depth level, values per knob. Defaults are always included so the
|
||||
# baseline configuration always appears in the sweep.
|
||||
|
||||
_HW_KNOB_DEFAULTS = {
|
||||
"pe_hbm_gb": 6.0,
|
||||
"bw_hbm_gbs": 256.0,
|
||||
"peak_tflops_f16": 8.0,
|
||||
"bw_intra_gbs": 512.0,
|
||||
"bw_inter_gbs": 128.0,
|
||||
"bw_intersip_gbs": 50.0,
|
||||
}
|
||||
|
||||
_HW_VALUES_BALANCED = {
|
||||
"pe_hbm_gb": [6.0, 12.0],
|
||||
"bw_hbm_gbs": [256.0, 512.0],
|
||||
"peak_tflops_f16": [8.0, 16.0],
|
||||
"bw_intra_gbs": [512.0, 1024.0],
|
||||
"bw_inter_gbs": [128.0, 256.0],
|
||||
"bw_intersip_gbs": [50.0, 100.0],
|
||||
}
|
||||
|
||||
_HW_VALUES_COARSE = {
|
||||
"pe_hbm_gb": [6.0, 12.0, 24.0],
|
||||
"bw_hbm_gbs": [256.0, 512.0, 1024.0],
|
||||
"peak_tflops_f16": [8.0, 16.0, 32.0],
|
||||
"bw_intra_gbs": [512.0, 1024.0, 2048.0],
|
||||
"bw_inter_gbs": [128.0, 256.0, 512.0],
|
||||
"bw_intersip_gbs": [50.0, 100.0, 200.0],
|
||||
}
|
||||
|
||||
# For sensitivity: "double each knob independently from baseline."
|
||||
_SENSITIVITY_KNOBS = list(_HW_KNOB_DEFAULTS.keys())
|
||||
|
||||
|
||||
# ── Reduced parallelism search space (for per-HW inner loop) ─────────
|
||||
#
|
||||
# The full auto_explore has 28,800 configs; using it inside a HW loop is
|
||||
# too slow. Reduce to CP × TP × PP × DP × kv_shard_mode with common
|
||||
# defaults for the remaining 4 knobs (~2k configs).
|
||||
|
||||
_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
|
||||
_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
|
||||
_PP_OPTIONS = (1, 2, 4, 8, 16)
|
||||
_DP_OPTIONS = (1, 2, 4)
|
||||
_KV_SHARD_MODES = ("split", "replicate")
|
||||
|
||||
|
||||
# ── Result types ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class HardwareCandidate:
|
||||
pe_hbm_gb: float
|
||||
bw_hbm_gbs: float
|
||||
peak_tflops_f16: float
|
||||
bw_intra_gbs: float
|
||||
bw_inter_gbs: float
|
||||
bw_intersip_gbs: float
|
||||
|
||||
def as_machine(self) -> MachineParams:
|
||||
"""Build a MachineParams from this HW candidate, holding alphas +
|
||||
compute_util at their MachineParams defaults."""
|
||||
return MachineParams(
|
||||
pe_hbm_gb=self.pe_hbm_gb,
|
||||
bw_hbm_gbs=self.bw_hbm_gbs,
|
||||
peak_tflops_f16=self.peak_tflops_f16,
|
||||
bw_intra_gbs=self.bw_intra_gbs,
|
||||
bw_inter_gbs=self.bw_inter_gbs,
|
||||
bw_intersip_gbs=self.bw_intersip_gbs,
|
||||
)
|
||||
|
||||
@property
|
||||
def cost_score(self) -> float:
|
||||
"""Normalized cost proxy: sum of (knob / default). 6.0 at defaults.
|
||||
|
||||
This is NOT dollars — it's a rough silicon-area / capability proxy. A
|
||||
higher cost_score means "more capable hardware" (more HBM, faster BW,
|
||||
etc.). Under identical latency, lower cost_score wins.
|
||||
"""
|
||||
return sum(
|
||||
getattr(self, k) / _HW_KNOB_DEFAULTS[k]
|
||||
for k in _HW_KNOB_DEFAULTS
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class JointScore:
|
||||
"""One (hardware, parallelism) pair with its computed latency + cost."""
|
||||
hardware: HardwareCandidate
|
||||
parallelism: ConfigScore
|
||||
total_latency_ns: float
|
||||
cost_score: float
|
||||
|
||||
@property
|
||||
def latency_ms(self) -> float:
|
||||
return self.total_latency_ns / 1e6
|
||||
|
||||
|
||||
@dataclass
|
||||
class SensitivityRow:
|
||||
knob: str
|
||||
baseline_value: float
|
||||
doubled_value: float
|
||||
baseline_latency_ns: float
|
||||
doubled_latency_ns: float
|
||||
|
||||
@property
|
||||
def rel_speedup(self) -> float:
|
||||
"""1 - (doubled/baseline). Positive = doubling this knob speeds things up."""
|
||||
if self.baseline_latency_ns <= 0:
|
||||
return 0.0
|
||||
return 1.0 - (self.doubled_latency_ns / self.baseline_latency_ns)
|
||||
|
||||
|
||||
@dataclass
|
||||
class JointExploreResult:
|
||||
model_name: str
|
||||
s_kv: int
|
||||
mode: str
|
||||
depth: str
|
||||
total_hw: int
|
||||
total_joint: int
|
||||
all_scores: list[JointScore] = field(default_factory=list)
|
||||
pareto_scores: list[JointScore] = field(default_factory=list)
|
||||
sensitivity: list[SensitivityRow] = field(default_factory=list)
|
||||
|
||||
|
||||
# ── Hardware enumeration ─────────────────────────────────────────────
|
||||
|
||||
|
||||
def enumerate_hardware(depth: str = "balanced") -> Iterator[HardwareCandidate]:
|
||||
"""Yield HardwareCandidates for the given sweep depth.
|
||||
|
||||
depth ∈ {"two_stage", "balanced", "coarse"}:
|
||||
- two_stage: only the default HW (1 candidate; parallelism sweep dominates)
|
||||
- balanced: 2 values per knob → 2^6 = 64 candidates
|
||||
- coarse: 3 values per knob → 3^6 = 729 candidates
|
||||
"""
|
||||
if depth == "two_stage":
|
||||
vals = {k: [v] for k, v in _HW_KNOB_DEFAULTS.items()}
|
||||
elif depth == "coarse":
|
||||
vals = _HW_VALUES_COARSE
|
||||
else: # balanced (default)
|
||||
vals = _HW_VALUES_BALANCED
|
||||
|
||||
for pe_hbm in vals["pe_hbm_gb"]:
|
||||
for bw_hbm in vals["bw_hbm_gbs"]:
|
||||
for tflops in vals["peak_tflops_f16"]:
|
||||
for bw_intra in vals["bw_intra_gbs"]:
|
||||
for bw_inter in vals["bw_inter_gbs"]:
|
||||
for bw_intersip in vals["bw_intersip_gbs"]:
|
||||
yield HardwareCandidate(
|
||||
pe_hbm_gb=pe_hbm,
|
||||
bw_hbm_gbs=bw_hbm,
|
||||
peak_tflops_f16=tflops,
|
||||
bw_intra_gbs=bw_intra,
|
||||
bw_inter_gbs=bw_inter,
|
||||
bw_intersip_gbs=bw_intersip,
|
||||
)
|
||||
|
||||
|
||||
# ── Reduced parallelism search per HW ────────────────────────────────
|
||||
|
||||
|
||||
def _iter_reduced_parallelism(
|
||||
model: ModelConfig, s_kv: int, mode: str,
|
||||
) -> Iterator[TopologyConfig]:
|
||||
"""Reduced sweep: CP × TP × PP × DP × kv_shard_mode.
|
||||
|
||||
Other 4 knobs held at defaults chosen to be latency-friendly for the
|
||||
decode/prefill common case:
|
||||
- ffn_shard_scope = "TP+CP" (common good choice; trades a small AR
|
||||
for lower per-PE weight bytes)
|
||||
- tp_placement = "cube" (packs TP across cubes; UCIe D2D)
|
||||
- cp_placement = "pe" (packs CP inside cubes; intra NoC)
|
||||
- cp_ring_variant = "qoml" for decode, "kv" for prefill
|
||||
|
||||
Total: 8 × 6 × 5 × 4 × 2 = 1,920 configs per HW candidate.
|
||||
"""
|
||||
cp_ring = "qoml" if mode == "decode" else "kv"
|
||||
for cp in _CP_OPTIONS:
|
||||
for tp in _TP_OPTIONS:
|
||||
if tp > 4 * model.h_q:
|
||||
continue
|
||||
for pp in _PP_OPTIONS:
|
||||
if pp > model.layers:
|
||||
continue
|
||||
for dp in _DP_OPTIONS:
|
||||
for kv_mode in _KV_SHARD_MODES:
|
||||
# cp_ring=qoml with cp=1 is a no-op; fall back to kv.
|
||||
ring = "kv" if cp == 1 else cp_ring
|
||||
yield TopologyConfig(
|
||||
cp=cp, tp=tp, pp=pp, dp=dp,
|
||||
s_kv=s_kv, mode=mode,
|
||||
kv_shard_mode=kv_mode,
|
||||
ffn_shard_scope="TP+CP",
|
||||
tp_placement="cube",
|
||||
cp_placement="pe",
|
||||
cp_ring_variant=ring,
|
||||
)
|
||||
|
||||
|
||||
def _best_parallelism_for_hw(
|
||||
model: ModelConfig, machine: MachineParams,
|
||||
s_kv: int, mode: str,
|
||||
include_attention: bool = True, include_ffn: bool = True,
|
||||
b: int = 1,
|
||||
) -> ConfigScore | None:
|
||||
"""Return the latency-minimum feasible parallelism for this HW.
|
||||
|
||||
Feasibility: fits memory + placement_valid. Returns None if nothing fits.
|
||||
Scope flags are forwarded to score_config so the "latency" ranked on
|
||||
matches the caller's attention/FFN/full choice.
|
||||
"""
|
||||
best: ConfigScore | None = None
|
||||
best_key: tuple | None = None
|
||||
for topo in _iter_reduced_parallelism(model, s_kv, mode):
|
||||
topo.b = b
|
||||
cfg = FullConfig(model=model, topo=topo, machine=machine)
|
||||
s = score_config(cfg, include_attention=include_attention,
|
||||
include_ffn=include_ffn)
|
||||
if not (s.fits_memory and s.placement_valid):
|
||||
continue
|
||||
# Compound tiebreaker: prefer fewer PEs and lower HBM % when latency
|
||||
# ties across variant knobs (cp_ring, placement, kv_shard_mode, etc.).
|
||||
_k = (s.total_latency_ns, s.pes_used, s.hbm_utilization)
|
||||
if best is None or _k < best_key:
|
||||
best = s
|
||||
best_key = _k
|
||||
return best
|
||||
|
||||
|
||||
def _best_parallelism_two_stage(
|
||||
model: ModelConfig, machine: MachineParams,
|
||||
s_kv: int, mode: str,
|
||||
include_attention: bool = True, include_ffn: bool = True,
|
||||
b: int = 1,
|
||||
) -> ConfigScore | None:
|
||||
"""Fast fallback: use autosuggest's memory-min (CP,TP,PP) then score it."""
|
||||
sug = auto_suggest(model, machine, s_kv, mode)
|
||||
if not sug.fits:
|
||||
return None
|
||||
topo = TopologyConfig(
|
||||
cp=sug.cp, tp=sug.tp, pp=sug.pp, dp=1, b=b,
|
||||
s_kv=s_kv, mode=mode,
|
||||
kv_shard_mode="split",
|
||||
ffn_shard_scope="TP+CP",
|
||||
tp_placement="cube",
|
||||
cp_placement="pe",
|
||||
cp_ring_variant="qoml" if mode == "decode" and sug.cp > 1 else "kv",
|
||||
)
|
||||
cfg = FullConfig(model=model, topo=topo, machine=machine)
|
||||
s = score_config(cfg, include_attention=include_attention,
|
||||
include_ffn=include_ffn)
|
||||
return s if s.fits_memory and s.placement_valid else None
|
||||
|
||||
|
||||
# ── Pareto over joint (latency, cost) ────────────────────────────────
|
||||
|
||||
|
||||
def _pareto_2d(scores: list[JointScore]) -> list[JointScore]:
|
||||
"""Non-dominated set on (total_latency_ns ↓, cost_score ↓). O(N²)."""
|
||||
frontier: list[JointScore] = []
|
||||
for i, a in enumerate(scores):
|
||||
dominated = False
|
||||
for j, b in enumerate(scores):
|
||||
if i == j:
|
||||
continue
|
||||
no_worse = (
|
||||
b.total_latency_ns <= a.total_latency_ns
|
||||
and b.cost_score <= a.cost_score
|
||||
)
|
||||
strictly_better = (
|
||||
b.total_latency_ns < a.total_latency_ns
|
||||
or b.cost_score < a.cost_score
|
||||
)
|
||||
if no_worse and strictly_better:
|
||||
dominated = True
|
||||
break
|
||||
if not dominated:
|
||||
frontier.append(a)
|
||||
return frontier
|
||||
|
||||
|
||||
# ── Per-knob sensitivity ─────────────────────────────────────────────
|
||||
|
||||
|
||||
def compute_sensitivity(
|
||||
baseline_hw: HardwareCandidate,
|
||||
parallelism: ConfigScore,
|
||||
model: ModelConfig, s_kv: int, mode: str,
|
||||
include_attention: bool = True, include_ffn: bool = True,
|
||||
) -> list[SensitivityRow]:
|
||||
"""For each HW knob, double it (holding others at baseline) and measure
|
||||
the latency change. Same parallelism used throughout so we isolate the
|
||||
HW knob's effect."""
|
||||
rows: list[SensitivityRow] = []
|
||||
baseline_machine = baseline_hw.as_machine()
|
||||
baseline_topo = parallelism.as_topology(s_kv, mode)
|
||||
baseline_cfg = FullConfig(
|
||||
model=model, topo=baseline_topo, machine=baseline_machine,
|
||||
)
|
||||
baseline_score = score_config(
|
||||
baseline_cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
baseline_latency = baseline_score.total_latency_ns
|
||||
|
||||
for knob in _SENSITIVITY_KNOBS:
|
||||
doubled_val = 2.0 * getattr(baseline_hw, knob)
|
||||
doubled_hw = replace(baseline_hw, **{knob: doubled_val})
|
||||
doubled_cfg = FullConfig(
|
||||
model=model, topo=baseline_topo, machine=doubled_hw.as_machine(),
|
||||
)
|
||||
doubled_score = score_config(
|
||||
doubled_cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
rows.append(SensitivityRow(
|
||||
knob=knob,
|
||||
baseline_value=getattr(baseline_hw, knob),
|
||||
doubled_value=doubled_val,
|
||||
baseline_latency_ns=baseline_latency,
|
||||
doubled_latency_ns=doubled_score.total_latency_ns,
|
||||
))
|
||||
# Sort by biggest speedup first (most sensitive knob at the top).
|
||||
rows.sort(key=lambda r: -r.rel_speedup)
|
||||
return rows
|
||||
|
||||
|
||||
# ── Top-level driver ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def joint_explore(
|
||||
model: ModelConfig,
|
||||
s_kv: int,
|
||||
mode: str,
|
||||
depth: str = "balanced",
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True,
|
||||
b: int = 1,
|
||||
) -> JointExploreResult:
|
||||
"""Sweep HW candidates × parallelism, return joint Pareto + sensitivity.
|
||||
|
||||
Sensitivity is computed around the LATENCY-MINIMUM joint point (the
|
||||
"best fast" config), doubling each HW knob one at a time.
|
||||
|
||||
Scope flags select which stages contribute to the summed latency —
|
||||
both the per-HW parallelism search and the sensitivity ranking use
|
||||
the same restriction.
|
||||
"""
|
||||
all_scores: list[JointScore] = []
|
||||
for hw in enumerate_hardware(depth):
|
||||
machine = hw.as_machine()
|
||||
if depth == "two_stage":
|
||||
par = _best_parallelism_two_stage(
|
||||
model, machine, s_kv, mode,
|
||||
include_attention=include_attention, include_ffn=include_ffn,
|
||||
b=b,
|
||||
)
|
||||
else:
|
||||
par = _best_parallelism_for_hw(
|
||||
model, machine, s_kv, mode,
|
||||
include_attention=include_attention, include_ffn=include_ffn,
|
||||
b=b,
|
||||
)
|
||||
if par is None:
|
||||
continue
|
||||
all_scores.append(JointScore(
|
||||
hardware=hw,
|
||||
parallelism=par,
|
||||
total_latency_ns=par.total_latency_ns,
|
||||
cost_score=hw.cost_score,
|
||||
))
|
||||
|
||||
all_scores.sort(key=lambda s: s.total_latency_ns)
|
||||
pareto = _pareto_2d(all_scores)
|
||||
pareto.sort(key=lambda s: s.total_latency_ns)
|
||||
|
||||
sensitivity: list[SensitivityRow] = []
|
||||
if all_scores:
|
||||
best = all_scores[0]
|
||||
sensitivity = compute_sensitivity(
|
||||
best.hardware, best.parallelism, model, s_kv, mode,
|
||||
include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
|
||||
return JointExploreResult(
|
||||
model_name=model.name,
|
||||
s_kv=s_kv,
|
||||
mode=mode,
|
||||
depth=depth,
|
||||
total_hw=sum(1 for _ in enumerate_hardware(depth)),
|
||||
total_joint=len(all_scores),
|
||||
all_scores=all_scores,
|
||||
pareto_scores=pareto,
|
||||
sensitivity=sensitivity,
|
||||
)
|
||||
@@ -1,146 +0,0 @@
|
||||
"""Auto-suggest (CP, TP, PP) that fits the model in the per-PE budget.
|
||||
|
||||
Strategy: iterate over candidate (CP, TP, PP) triples in order of
|
||||
increasing total PE count and return the first one that satisfies
|
||||
memory + kernel-support constraints with >10% slack. If none fits, return
|
||||
the best-effort configuration and flag over-budget.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Iterable
|
||||
|
||||
from .model_config import FullConfig, ModelConfig, TopologyConfig, MachineParams
|
||||
from .memory_layout import compute_memory
|
||||
|
||||
|
||||
# Candidate parallelism dimensions (powers of 2 mostly, up to sensible limits).
|
||||
_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
|
||||
_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
|
||||
_PP_OPTIONS = (1, 2, 4, 8, 16, 32)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Suggestion:
|
||||
cp: int
|
||||
tp: int
|
||||
pp: int
|
||||
weights_gb: float
|
||||
kv_gb: float
|
||||
transient_gb: float
|
||||
slack_gb: float
|
||||
fits: bool
|
||||
pes_used: int
|
||||
sips_used: int
|
||||
cubes_used: int = 0 # picked to minimize this (see _score_candidate)
|
||||
cp_placement: str = "cube" # "pe" packs CP into intra-cube PEs when it fits
|
||||
reason: str = ""
|
||||
|
||||
|
||||
def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]:
|
||||
"""Yield (CP, TP, PP) in order of increasing PE count."""
|
||||
triples: list[tuple[int, int, int]] = []
|
||||
for tp in _TP_OPTIONS:
|
||||
for cp in _CP_OPTIONS:
|
||||
for pp in _PP_OPTIONS:
|
||||
# PP must not exceed layer count.
|
||||
if pp > model.layers:
|
||||
continue
|
||||
# Skip TP > H_q * some factor (unrealistic).
|
||||
if tp > model.h_q * 4:
|
||||
continue
|
||||
triples.append((cp, tp, pp))
|
||||
# Sort by pe_count then by (pp, tp, cp) — prefer smaller PP first
|
||||
# (avoids pipeline bubbles), then smaller TP (avoids head-dim split).
|
||||
triples.sort(key=lambda t: (t[0] * t[1] * t[2], t[2], t[1], t[0]))
|
||||
for t in triples:
|
||||
yield t
|
||||
|
||||
|
||||
def _score_candidate(cp: int, tp: int, pp: int,
|
||||
model: ModelConfig, machine: MachineParams,
|
||||
s_kv: int, mode: str,
|
||||
slack_frac: float = 0.10,
|
||||
b: int = 1,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True) -> Suggestion:
|
||||
# For each triple, try both cp_placement options and keep the one
|
||||
# with fewer cubes (breaks the historical "CP always across cubes"
|
||||
# default when a smaller pack is possible). The pe placement is only
|
||||
# valid when CP·TP fits within a single cube's PE count.
|
||||
best_topo: TopologyConfig | None = None
|
||||
best_placement = "cube"
|
||||
_pe_per_cube_hw = TopologyConfig().pes_per_cube_hw # instance-invariant HW const
|
||||
_placements_to_try = ["cube"]
|
||||
if cp * tp <= _pe_per_cube_hw:
|
||||
_placements_to_try.append("pe")
|
||||
for _place in _placements_to_try:
|
||||
_t = TopologyConfig(
|
||||
cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode, b=b,
|
||||
cp_placement=_place,
|
||||
)
|
||||
if best_topo is None or _t.cubes_used < best_topo.cubes_used:
|
||||
best_topo = _t
|
||||
best_placement = _place
|
||||
topo = best_topo
|
||||
cfg = FullConfig(model=model, topo=topo, machine=machine)
|
||||
mem = compute_memory(
|
||||
cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
)
|
||||
|
||||
fits = (not mem.over_budget
|
||||
and mem.slack_bytes >= slack_frac * mem.budget_bytes)
|
||||
reason = ""
|
||||
if mem.over_budget:
|
||||
reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
|
||||
f"exceeds budget ({mem.budget_bytes/1e9:.2f} GB)")
|
||||
elif not fits:
|
||||
reason = f"slack ({mem.slack_bytes/1e9:.2f} GB) below 10% of budget"
|
||||
return Suggestion(
|
||||
cp=cp, tp=tp, pp=pp,
|
||||
weights_gb=mem.weights_bytes / 1e9,
|
||||
kv_gb=mem.kv_cache_bytes / 1e9,
|
||||
transient_gb=mem.transient_bytes / 1e9,
|
||||
slack_gb=mem.slack_bytes / 1e9,
|
||||
fits=fits,
|
||||
pes_used=topo.total_pes,
|
||||
sips_used=topo.sips_used,
|
||||
cubes_used=topo.cubes_used,
|
||||
cp_placement=best_placement,
|
||||
reason=reason,
|
||||
)
|
||||
|
||||
|
||||
def auto_suggest(model: ModelConfig, machine: MachineParams,
|
||||
s_kv: int, mode: str = "decode",
|
||||
slack_frac: float = 0.10,
|
||||
b: int = 1,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True) -> Suggestion:
|
||||
"""Return the smallest deployment (fewer cubes, then fewer PEs)
|
||||
that fits.
|
||||
|
||||
Sort key: (cubes_used ↑, pes_used ↑, pp ↑, tp ↑, cp ↑).
|
||||
Fewer cubes wins first because a cube is the physical die-level
|
||||
hardware unit; PE count is the tiebreaker. Preferring fewer PP then
|
||||
TP then CP keeps the earlier historical bias (avoid pipeline
|
||||
bubbles / head-dim splits) among equal-cube-and-PE ties.
|
||||
|
||||
If no candidate fits, returns the best-effort one (highest slack,
|
||||
even if negative) with fits=False.
|
||||
"""
|
||||
scored: list[Suggestion] = []
|
||||
for cp, tp, pp in _iter_candidates(model):
|
||||
scored.append(
|
||||
_score_candidate(cp, tp, pp, model, machine, s_kv, mode,
|
||||
slack_frac, b=b,
|
||||
include_attention=include_attention,
|
||||
include_ffn=include_ffn)
|
||||
)
|
||||
scored.sort(key=lambda s: (s.cubes_used, s.pes_used, s.pp, s.tp, s.cp))
|
||||
|
||||
for s in scored:
|
||||
if s.fits:
|
||||
return s
|
||||
# No fit — return the best-effort (largest slack, closest to fitting).
|
||||
return max(scored, key=lambda s: s.slack_gb)
|
||||
@@ -1,518 +0,0 @@
|
||||
"""Chip-level roofline math: AI, B*, L*, per-token latency curves.
|
||||
|
||||
Surfaces the arithmetic-intensity story from LLM-serving practice:
|
||||
|
||||
- **AI** = C / W (peak FLOPs per byte of HBM bandwidth).
|
||||
- **B\*** = C * b / (2 * W) * sparsity — the critical batch size at
|
||||
which weight-fetch time equals compute time for one decode step.
|
||||
Sparsity = N_total / N_active (MoE factor; 1 for dense).
|
||||
- **L\*** = 2 * N_active / (AI * kv_bytes_per_token) — the balance
|
||||
context length at which KV-read time equals compute time.
|
||||
- **B_knee(S_kv)** = B* / (1 - S_kv/L*) — the batch size where the
|
||||
cost curve bends (weight-fetch drops below the compute+KV floor).
|
||||
Diverges at S_kv = L* and no knee exists past it.
|
||||
|
||||
Per-token decode-step latency, per PE (dense-approx, no comm):
|
||||
|
||||
t(B, S_kv) = N_active * b / (W * B) <-- weight fetch, 1/B
|
||||
+ 2 * N_active / C <-- compute (peak), flat
|
||||
+ S_kv * kv_bpt / W <-- KV read, flat
|
||||
|
||||
All numbers per PE / per one forward pass. **Peak roofline — no
|
||||
utilization factor.** Comm cost and TP/CP sharding are intentionally
|
||||
NOT in the roofline — this is the back-of-envelope chip-vs-model view
|
||||
the transcript talks about, not the full latency model that
|
||||
stage_latencies.py builds. With this convention weight_s == compute_s
|
||||
exactly at B*.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, replace
|
||||
|
||||
from .model_config import FullConfig, MachineParams, ModelConfig
|
||||
|
||||
|
||||
# Per-token bytes of BF16 MAC arithmetic: one multiply + one add = 2 FLOPs.
|
||||
_FLOPS_PER_PARAM_PER_TOKEN = 2
|
||||
|
||||
|
||||
# ── Chip / model derived quantities ────────────────────────────────
|
||||
|
||||
|
||||
def arithmetic_intensity(machine: MachineParams) -> float:
|
||||
"""FLOPs per byte of HBM bandwidth. Peak roofline; utilization
|
||||
is applied only in the compute-time formula, not here."""
|
||||
return machine.peak_flops / machine.bw_hbm
|
||||
|
||||
|
||||
def total_active_params(model: ModelConfig) -> int:
|
||||
"""Full-model parameter count (attention + FFN, all layers).
|
||||
|
||||
Attention: 4 projections × hidden × H_q * d_head effective per layer.
|
||||
(W_Q hidden×H_q*d_h, W_O H_q*d_h×hidden, W_K/W_V hidden×H_kv*d_h.)
|
||||
FFN: 3 × hidden × ffn_dim per layer (gate, up, down).
|
||||
"""
|
||||
m = model
|
||||
attn = (
|
||||
m.hidden * m.h_q * m.d_head # W_Q
|
||||
+ m.hidden * m.h_kv * m.d_head * 2 # W_K + W_V
|
||||
+ m.h_q * m.d_head * m.hidden # W_O
|
||||
)
|
||||
ffn = 3 * m.hidden * m.ffn_dim
|
||||
return (attn + ffn) * m.layers
|
||||
|
||||
|
||||
def kv_bytes_per_token(model: ModelConfig) -> int:
|
||||
"""Bytes of KV cache one new token adds across ALL layers, per one
|
||||
sequence, un-sharded (K + V, H_kv heads * d_h * bytes)."""
|
||||
m = model
|
||||
return 2 * m.h_kv * m.d_head * m.bytes_per_elem * m.layers
|
||||
|
||||
|
||||
def critical_batch(machine: MachineParams, model: ModelConfig,
|
||||
sparsity: float = 1.0) -> float:
|
||||
"""B* = C * b / (2 * W) * sparsity.
|
||||
|
||||
Sparsity = N_total / N_active (>= 1). Dense = 1. MoE 8-of-256 = 8.
|
||||
"""
|
||||
b = model.bytes_per_elem
|
||||
ai = arithmetic_intensity(machine)
|
||||
return ai * b / _FLOPS_PER_PARAM_PER_TOKEN * sparsity
|
||||
|
||||
|
||||
def balance_context(machine: MachineParams, model: ModelConfig) -> float:
|
||||
"""L* = 2 * N_active / (AI * kv_bpt).
|
||||
|
||||
Context length (in tokens) at which per-step KV read matches the
|
||||
per-step compute cost. Beyond L*, the KV term is dominant and no
|
||||
batch size gets you compute-bound.
|
||||
"""
|
||||
n_active = total_active_params(model)
|
||||
ai = arithmetic_intensity(machine)
|
||||
kv_bpt = kv_bytes_per_token(model)
|
||||
return _FLOPS_PER_PARAM_PER_TOKEN * n_active / (ai * kv_bpt)
|
||||
|
||||
|
||||
def knee_batch(machine: MachineParams, model: ModelConfig,
|
||||
s_kv: int) -> float | None:
|
||||
"""B_knee(S_kv) = B* / (1 - S_kv/L*).
|
||||
|
||||
Returns None when S_kv >= L* (no knee exists — the total-cost
|
||||
curve never touches the compute floor).
|
||||
"""
|
||||
b_star = critical_batch(machine, model)
|
||||
l_star = balance_context(machine, model)
|
||||
r = s_kv / l_star
|
||||
if r >= 1.0:
|
||||
return None
|
||||
return b_star / (1.0 - r)
|
||||
|
||||
|
||||
# ── Per-token latency curves ───────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class RooflinePoint:
|
||||
batch: int
|
||||
weight_s: float # weight fetch time, 1/B
|
||||
compute_s: float # compute time, flat
|
||||
kv_s: float # KV read time, flat
|
||||
total_s: float
|
||||
|
||||
|
||||
def per_token_latency_curve(machine: MachineParams, model: ModelConfig,
|
||||
batch_range: list[int],
|
||||
s_kv: int) -> list[RooflinePoint]:
|
||||
"""Per-token decode-step latency curve across a range of batch sizes.
|
||||
|
||||
Returns one point per batch. All times per PE, dense-approx,
|
||||
utilization from machine.compute_util. Comm and TP/CP sharding are
|
||||
excluded — this is the roofline model.
|
||||
"""
|
||||
n_active = total_active_params(model)
|
||||
b = model.bytes_per_elem
|
||||
weight_bytes = n_active * b
|
||||
compute_flops = _FLOPS_PER_PARAM_PER_TOKEN * n_active
|
||||
kv_read_bytes = s_kv * kv_bytes_per_token(model)
|
||||
|
||||
compute_s = compute_flops / machine.peak_flops
|
||||
kv_s = kv_read_bytes / machine.bw_hbm
|
||||
|
||||
points: list[RooflinePoint] = []
|
||||
for bs in batch_range:
|
||||
weight_s = weight_bytes / machine.bw_hbm / max(1, bs)
|
||||
points.append(RooflinePoint(
|
||||
batch=bs,
|
||||
weight_s=weight_s,
|
||||
compute_s=compute_s,
|
||||
kv_s=kv_s,
|
||||
total_s=weight_s + compute_s + kv_s,
|
||||
))
|
||||
return points
|
||||
|
||||
|
||||
def bound_regime(machine: MachineParams, model: ModelConfig,
|
||||
batch: int, s_kv: int) -> str:
|
||||
"""Which term dominates at the current (batch, S_kv) point.
|
||||
|
||||
Returns 'memory-bound' if weight_fetch is the largest term,
|
||||
'kv-bound' if KV read is largest, 'compute-bound' if compute.
|
||||
"""
|
||||
pts = per_token_latency_curve(machine, model, [batch], s_kv)
|
||||
p = pts[0]
|
||||
parts = {"memory-bound": p.weight_s,
|
||||
"kv-bound": p.kv_s,
|
||||
"compute-bound": p.compute_s}
|
||||
return max(parts, key=parts.get)
|
||||
|
||||
|
||||
# ── Regime-dependent cost terms ────────────────────────────────────
|
||||
|
||||
|
||||
def t_mem_short(machine: MachineParams, model: ModelConfig,
|
||||
batch: int) -> float:
|
||||
"""Per-token weight-fetch time (short-context regime term).
|
||||
|
||||
N_active · b / (W · B). Shrinks as B grows — this is what
|
||||
batching amortizes.
|
||||
"""
|
||||
return (total_active_params(model) * model.bytes_per_elem
|
||||
/ (machine.bw_hbm * max(1, batch)))
|
||||
|
||||
|
||||
def t_mem_long(machine: MachineParams, model: ModelConfig,
|
||||
s_kv: int) -> float:
|
||||
"""Per-token KV-read time (long-context regime term).
|
||||
|
||||
S_kv · kv_bpt / W. Independent of B — each sequence reads its
|
||||
own KV cache; batching doesn't help.
|
||||
"""
|
||||
return s_kv * kv_bytes_per_token(model) / machine.bw_hbm
|
||||
|
||||
|
||||
def t_com(machine: MachineParams, model: ModelConfig) -> float:
|
||||
"""Per-token compute time. Same in both regimes: 2·N/C, peak."""
|
||||
return _FLOPS_PER_PARAM_PER_TOKEN * total_active_params(model) / machine.peak_flops
|
||||
|
||||
|
||||
# ── "Good" batch / context recommendations ─────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class BatchRecommendation:
|
||||
target: float # Pope's rule: 2 × B*
|
||||
b_star: float # B* itself
|
||||
effective: float # what we recommend using
|
||||
reason: str # short explanation
|
||||
|
||||
|
||||
@dataclass
|
||||
class ContextRecommendation:
|
||||
l_star: float # balance context length
|
||||
max_efficient: float # same as l_star (compute-friendly ceiling)
|
||||
utilization_at: float # utilization at current s_kv
|
||||
reason: str
|
||||
|
||||
|
||||
def good_batch(machine: MachineParams, model: ModelConfig,
|
||||
sparsity: float = 1.0) -> BatchRecommendation:
|
||||
"""Recommended batch size: 2 × B* (Pope's rule of thumb).
|
||||
|
||||
Below B*: memory-bound, doubling B halves cost/token.
|
||||
At 2×B*: 50% excess over compute floor — the sweet spot.
|
||||
Beyond 3×B*: diminishing returns; latency keeps growing linearly.
|
||||
"""
|
||||
b_star = critical_batch(machine, model, sparsity)
|
||||
target = 2 * b_star
|
||||
return BatchRecommendation(
|
||||
target=target, b_star=b_star, effective=target,
|
||||
reason=(f"2·B* = 2 · {b_star:.0f} = {target:.0f}. "
|
||||
"Below B*: memory-bound (doubling B halves cost/token). "
|
||||
"Beyond 3·B*: diminishing returns."),
|
||||
)
|
||||
|
||||
|
||||
def good_context(machine: MachineParams, model: ModelConfig,
|
||||
s_kv: int) -> ContextRecommendation:
|
||||
"""Recommended max context: L* — the compute-friendly ceiling.
|
||||
|
||||
Below L*: KV read is cheap relative to compute → good utilization.
|
||||
Above L*: KV bandwidth wall → utilization = 1/(1 + S_kv/L*).
|
||||
"""
|
||||
l_star = balance_context(machine, model)
|
||||
util = utilization_at(s_kv, l_star)
|
||||
return ContextRecommendation(
|
||||
l_star=l_star, max_efficient=l_star, utilization_at=util,
|
||||
reason=(f"L* = {l_star:,.0f} tokens. Below L*: compute-bound "
|
||||
f"(good util). At {s_kv:,} tokens: peak utilization ≈ "
|
||||
f"{util*100:.1f}% (1 / (1 + S_kv/L*))."),
|
||||
)
|
||||
|
||||
|
||||
def utilization_at(s_kv: int, l_star: float) -> float:
|
||||
"""Peak compute utilization at context length s_kv, given L*.
|
||||
|
||||
util = compute / (compute + KV_read) = 1 / (1 + S_kv/L*).
|
||||
At S_kv=0: 100%. At S_kv=L*: 50%. At 2·L*: 33.3%. At 5·L*: 16.7%.
|
||||
"""
|
||||
return 1.0 / (1.0 + s_kv / l_star)
|
||||
|
||||
|
||||
# ── Per-step latency (undivided by B) ─────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class StepLatencyPoint:
|
||||
batch: int
|
||||
weight_s: float # N·b / W — flat in B (loaded once per step)
|
||||
compute_s: float # 2·N·B / C — linear in B
|
||||
kv_s: float # B · S_kv · kv_bpt / W — linear in B
|
||||
total_s: float
|
||||
|
||||
|
||||
def step_latency_curve(machine: MachineParams, model: ModelConfig,
|
||||
batch_range: list[int],
|
||||
s_kv: int) -> list[StepLatencyPoint]:
|
||||
"""Total time of one decode step (one forward pass), across a
|
||||
range of batch sizes. NOT divided by B — this is the SLO view.
|
||||
|
||||
step_weight = N·b / W (batch-invariant)
|
||||
step_compute = 2·N·B / C (linear in B)
|
||||
step_kv = B · S_kv · kv_bpt / W (linear in B)
|
||||
|
||||
Per-token cost = step_total / B — the two views trade off:
|
||||
bigger B lowers cost/token but raises step latency.
|
||||
"""
|
||||
n_active = total_active_params(model)
|
||||
b = model.bytes_per_elem
|
||||
weight_bytes = n_active * b
|
||||
kv_bpt = kv_bytes_per_token(model)
|
||||
|
||||
step_weight = weight_bytes / machine.bw_hbm
|
||||
|
||||
points: list[StepLatencyPoint] = []
|
||||
for bs in batch_range:
|
||||
B = max(1, bs)
|
||||
step_compute = _FLOPS_PER_PARAM_PER_TOKEN * n_active * B / machine.peak_flops
|
||||
step_kv = B * s_kv * kv_bpt / machine.bw_hbm
|
||||
total = step_weight + step_compute + step_kv
|
||||
points.append(StepLatencyPoint(
|
||||
batch=bs,
|
||||
weight_s=step_weight,
|
||||
compute_s=step_compute,
|
||||
kv_s=step_kv,
|
||||
total_s=total,
|
||||
))
|
||||
return points
|
||||
|
||||
|
||||
# ── PE memory budget curves ───────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class MemoryBudgetPoint:
|
||||
axis_val: int # S_kv or B being swept
|
||||
weights_gb: float
|
||||
kv_gb: float
|
||||
transient_gb: float
|
||||
used_gb: float
|
||||
free_gb: float # max(0, hbm_gb - used_gb)
|
||||
over_budget: bool
|
||||
|
||||
|
||||
def _budget_point(weights_bytes: int, kv_bytes: int, transient_bytes: int,
|
||||
hbm_bytes: int, axis_val: int) -> MemoryBudgetPoint:
|
||||
used = weights_bytes + kv_bytes + transient_bytes
|
||||
return MemoryBudgetPoint(
|
||||
axis_val=axis_val,
|
||||
weights_gb=weights_bytes / 1e9,
|
||||
kv_gb=kv_bytes / 1e9,
|
||||
transient_gb=transient_bytes / 1e9,
|
||||
used_gb=used / 1e9,
|
||||
free_gb=max(0, hbm_bytes - used) / 1e9,
|
||||
over_budget=(used > hbm_bytes),
|
||||
)
|
||||
|
||||
|
||||
def memory_budget_curve_vs_skv(cfg: FullConfig,
|
||||
s_kv_range: list[int],
|
||||
batch: int) -> list[MemoryBudgetPoint]:
|
||||
"""Per-PE memory as S_kv sweeps. Uses cfg's current sharding
|
||||
(CP, TP, PP). Sets topo.b = batch. Returns one point per S_kv."""
|
||||
from .memory_layout import (
|
||||
per_pe_kv_cache_bytes, per_pe_transient_bytes, per_pe_weight_bytes,
|
||||
)
|
||||
hbm_bytes = int(cfg.machine.pe_budget_bytes)
|
||||
weights_bytes = per_pe_weight_bytes(cfg)
|
||||
transient_bytes = per_pe_transient_bytes(cfg)
|
||||
points: list[MemoryBudgetPoint] = []
|
||||
for skv in s_kv_range:
|
||||
swept_topo = replace(cfg.topo, s_kv=int(skv), b=max(1, int(batch)))
|
||||
swept_cfg = FullConfig(model=cfg.model, topo=swept_topo,
|
||||
machine=cfg.machine)
|
||||
kv_bytes = per_pe_kv_cache_bytes(swept_cfg)
|
||||
points.append(_budget_point(weights_bytes, kv_bytes,
|
||||
transient_bytes, hbm_bytes, int(skv)))
|
||||
return points
|
||||
|
||||
|
||||
def memory_budget_curve_vs_batch(cfg: FullConfig,
|
||||
b_range: list[int],
|
||||
s_kv: int) -> list[MemoryBudgetPoint]:
|
||||
"""Per-PE memory as B sweeps. Sets topo.s_kv = s_kv."""
|
||||
from .memory_layout import (
|
||||
per_pe_kv_cache_bytes, per_pe_transient_bytes, per_pe_weight_bytes,
|
||||
)
|
||||
hbm_bytes = int(cfg.machine.pe_budget_bytes)
|
||||
weights_bytes = per_pe_weight_bytes(cfg)
|
||||
transient_bytes = per_pe_transient_bytes(cfg)
|
||||
points: list[MemoryBudgetPoint] = []
|
||||
for bs in b_range:
|
||||
swept_topo = replace(cfg.topo, b=max(1, int(bs)), s_kv=int(s_kv))
|
||||
swept_cfg = FullConfig(model=cfg.model, topo=swept_topo,
|
||||
machine=cfg.machine)
|
||||
kv_bytes = per_pe_kv_cache_bytes(swept_cfg)
|
||||
points.append(_budget_point(weights_bytes, kv_bytes,
|
||||
transient_bytes, hbm_bytes, int(bs)))
|
||||
return points
|
||||
|
||||
|
||||
# ── AI / B* sensitivity to hardware knobs ──────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class AISensitivityPoint:
|
||||
multiplier: float # scale factor applied to the base machine
|
||||
peak_tflops: float
|
||||
bw_gbs: float
|
||||
ai: float # C / W
|
||||
b_star: float # C·b/(2·W)
|
||||
|
||||
|
||||
# ── GPU count sizing (three-axis) ─────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class GpuSizingResult:
|
||||
# Inputs echoed back for display / debugging
|
||||
n_users: int
|
||||
avg_ctx_tokens: int
|
||||
tpot_slo_s: float
|
||||
# Per-replica values
|
||||
pes_axis_a_capacity: int # PEs to hold weights alone (bare floor)
|
||||
pes_axis_b_kv: int # PEs to hold weights + KV of this replica's users
|
||||
pes_per_replica: int # max(A, B)
|
||||
# Throughput axis
|
||||
b_at_slo: int # largest per-replica B satisfying SLO
|
||||
users_per_replica: int # min(n_users, b_at_slo), at least 1
|
||||
n_replicas: int # ceil(n_users / users_per_replica)
|
||||
# Grand total
|
||||
total_pes: int
|
||||
binding_axis: str # 'capacity' | 'kv' | 'throughput'
|
||||
|
||||
|
||||
def max_batch_within_slo(machine: MachineParams, model: ModelConfig,
|
||||
s_kv: int, tpot_slo_s: float) -> int:
|
||||
"""Largest per-replica B such that decode step latency ≤ SLO.
|
||||
|
||||
step_latency(B) = N·b/W + 2·N·B/C + B·S_kv·kv_bpt/W
|
||||
= weight_fetch + B · (compute + kv_read)
|
||||
|
||||
Returns 0 if even B=1 exceeds SLO (weight fetch alone too large,
|
||||
or per-sequence term already blows the budget).
|
||||
"""
|
||||
n = total_active_params(model)
|
||||
b_elem = model.bytes_per_elem
|
||||
weight_time = n * b_elem / machine.bw_hbm
|
||||
per_seq_time = (_FLOPS_PER_PARAM_PER_TOKEN * n / machine.peak_flops
|
||||
+ s_kv * kv_bytes_per_token(model) / machine.bw_hbm)
|
||||
if weight_time + per_seq_time > tpot_slo_s:
|
||||
return 0
|
||||
remaining = tpot_slo_s - weight_time
|
||||
b_max = int(remaining / per_seq_time)
|
||||
return max(1, b_max)
|
||||
|
||||
|
||||
def size_deployment(machine: MachineParams, model: ModelConfig,
|
||||
n_users: int, avg_ctx: int,
|
||||
tpot_slo_s: float) -> GpuSizingResult:
|
||||
"""Three-axis GPU count for a target workload.
|
||||
|
||||
Axis A — capacity floor: PEs to hold one replica's weights.
|
||||
Axis B — KV headroom: PEs to hold weights + all KV of the
|
||||
users assigned to this replica.
|
||||
Axis C — throughput SLO: replicas needed so per-replica B ≤ b_at_slo.
|
||||
|
||||
Result: pes_per_replica = max(A, B); total = pes_per_replica × replicas.
|
||||
Binding axis is whichever grew the count the most.
|
||||
"""
|
||||
n = total_active_params(model)
|
||||
b_elem = model.bytes_per_elem
|
||||
weight_bytes = n * b_elem
|
||||
hbm_pe = int(machine.pe_hbm_gb * 1e9)
|
||||
kv_bpt = kv_bytes_per_token(model)
|
||||
|
||||
# Axis C: throughput
|
||||
b_at_slo = max_batch_within_slo(machine, model, avg_ctx, tpot_slo_s)
|
||||
users_per_replica = max(1, min(n_users, b_at_slo)) if b_at_slo > 0 else 1
|
||||
n_replicas = (n_users + users_per_replica - 1) // users_per_replica
|
||||
|
||||
# Axis A: bare weights
|
||||
pes_a = (weight_bytes + hbm_pe - 1) // hbm_pe
|
||||
|
||||
# Axis B: weights + KV load for this replica's users
|
||||
kv_per_replica = users_per_replica * avg_ctx * kv_bpt
|
||||
pes_b = (weight_bytes + kv_per_replica + hbm_pe - 1) // hbm_pe
|
||||
|
||||
pes_per_replica = max(pes_a, pes_b)
|
||||
total = pes_per_replica * n_replicas
|
||||
|
||||
if n_replicas > 1:
|
||||
binding = "throughput"
|
||||
elif pes_b > pes_a:
|
||||
binding = "kv"
|
||||
else:
|
||||
binding = "capacity"
|
||||
|
||||
return GpuSizingResult(
|
||||
n_users=n_users,
|
||||
avg_ctx_tokens=avg_ctx,
|
||||
tpot_slo_s=tpot_slo_s,
|
||||
pes_axis_a_capacity=int(pes_a),
|
||||
pes_axis_b_kv=int(pes_b),
|
||||
pes_per_replica=int(pes_per_replica),
|
||||
b_at_slo=b_at_slo,
|
||||
users_per_replica=int(users_per_replica),
|
||||
n_replicas=int(n_replicas),
|
||||
total_pes=int(total),
|
||||
binding_axis=binding,
|
||||
)
|
||||
|
||||
|
||||
def ai_sensitivity_curve(machine: MachineParams, model: ModelConfig,
|
||||
multipliers: list[float],
|
||||
axis: str = "flops") -> list[AISensitivityPoint]:
|
||||
"""Sweep FLOPs OR BW while holding the other fixed. Returns AI + B*
|
||||
at each multiplier.
|
||||
|
||||
axis='flops' → scales peak_tflops_f16 (AI grows linearly).
|
||||
axis='bw' → scales bw_hbm_gbs (AI shrinks inversely).
|
||||
"""
|
||||
if axis not in ("flops", "bw"):
|
||||
raise ValueError(f"axis must be 'flops' or 'bw', got {axis!r}")
|
||||
base_flops = machine.peak_tflops_f16
|
||||
base_bw = machine.bw_hbm_gbs
|
||||
points: list[AISensitivityPoint] = []
|
||||
for k in multipliers:
|
||||
if axis == "flops":
|
||||
m = replace(machine, peak_tflops_f16=base_flops * k)
|
||||
else:
|
||||
m = replace(machine, bw_hbm_gbs=base_bw * k)
|
||||
points.append(AISensitivityPoint(
|
||||
multiplier=k,
|
||||
peak_tflops=m.peak_tflops_f16,
|
||||
bw_gbs=m.bw_hbm_gbs,
|
||||
ai=arithmetic_intensity(m),
|
||||
b_star=critical_batch(m, model),
|
||||
))
|
||||
return points
|
||||
@@ -1,248 +0,0 @@
|
||||
"""Per-PE memory footprint: weights + KV cache + activations + slack.
|
||||
|
||||
All formulas are per-PE, per-PP-stage. If PP > 1, each PE holds only
|
||||
its stage's layers (layers/PP).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .model_config import FullConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class MemoryBreakdown:
|
||||
weights_bytes: int
|
||||
kv_cache_bytes: int
|
||||
transient_bytes: int
|
||||
budget_bytes: int
|
||||
|
||||
@property
|
||||
def used_bytes(self) -> int:
|
||||
return self.weights_bytes + self.kv_cache_bytes + self.transient_bytes
|
||||
|
||||
@property
|
||||
def slack_bytes(self) -> int:
|
||||
return max(0, self.budget_bytes - self.used_bytes)
|
||||
|
||||
@property
|
||||
def over_budget(self) -> bool:
|
||||
return self.used_bytes > self.budget_bytes
|
||||
|
||||
|
||||
def per_pe_weight_bytes(cfg: FullConfig,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True) -> int:
|
||||
"""Attention + FFN weights per PE, bf16. Divided by TP, PP, and EP.
|
||||
|
||||
EP divides the FFN (experts) across ranks; attention weights are
|
||||
unaffected. When kv_shard_mode='replicate' and TP > H_kv, each KV
|
||||
head is replicated (per-PE W_K/W_V size doesn't drop below one head).
|
||||
|
||||
Scope flags let callers count only attention or only FFN weights —
|
||||
useful for "smallest deployment for just this block" sizing. Both
|
||||
True (default) matches the traditional full-model per-PE weight
|
||||
footprint.
|
||||
"""
|
||||
m = cfg.model
|
||||
tp = cfg.topo.tp
|
||||
pp = cfg.topo.pp
|
||||
ep = max(1, cfg.topo.ep)
|
||||
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
if cfg.topo.kv_shard_mode == "replicate":
|
||||
# Each KV head held fully; replicated across ranks if TP > H_kv.
|
||||
hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
|
||||
else: # "split"
|
||||
# Head-dim split: fractional head allowed (< 1 head bytes when TP > H_kv).
|
||||
hkv_per_pe_bytes = m.h_kv / tp
|
||||
|
||||
per_layer_attn = 0
|
||||
if include_attention:
|
||||
per_layer_attn = (
|
||||
m.hidden * hq_per_pe * m.d_head + # W_Q
|
||||
m.hidden * hkv_per_pe_bytes * m.d_head + # W_K
|
||||
m.hidden * hkv_per_pe_bytes * m.d_head + # W_V
|
||||
hq_per_pe * m.d_head * m.hidden # W_O
|
||||
)
|
||||
|
||||
per_layer_ffn = 0
|
||||
if include_ffn:
|
||||
# FFN divisor: TP (default) or TP*CP or TP*CP*DP if user opts in.
|
||||
# EP further divides (MoE experts).
|
||||
ffn_div = cfg.ffn_shard_divisor * ep
|
||||
per_layer_ffn = 3 * m.hidden * (m.ffn_dim // ffn_div)
|
||||
|
||||
per_layer = (per_layer_attn + per_layer_ffn) * m.bytes_per_elem
|
||||
layers_per_stage = (m.layers + pp - 1) // pp
|
||||
return int(per_layer * layers_per_stage)
|
||||
|
||||
|
||||
def per_pe_kv_cache_bytes(cfg: FullConfig) -> int:
|
||||
"""K + V per PE, across all layers this stage holds.
|
||||
|
||||
- CP shards the sequence dim -> S_local = S_kv/CP tokens per PE.
|
||||
- TP splits KV heads across ranks. If kv_shard_mode='replicate' and
|
||||
TP > H_kv, each KV head is duplicated across TP/H_kv ranks so
|
||||
per-PE storage doesn't shrink below 1 head.
|
||||
- PP: only layers_per_stage layers stored per PE.
|
||||
- Batch (B): each concurrent request keeps its own KV cache slice.
|
||||
"""
|
||||
m = cfg.model
|
||||
pp = cfg.topo.pp
|
||||
tp = cfg.topo.tp
|
||||
B = max(1, cfg.topo.b)
|
||||
if cfg.topo.kv_shard_mode == "replicate":
|
||||
hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
|
||||
else:
|
||||
hkv_per_pe_bytes = m.h_kv / tp
|
||||
layers_per_stage = (m.layers + pp - 1) // pp
|
||||
per_layer = 2 * cfg.topo.s_local * hkv_per_pe_bytes * m.d_head * m.bytes_per_elem
|
||||
return int(per_layer * layers_per_stage * B)
|
||||
|
||||
|
||||
def per_pe_transient_bytes(cfg: FullConfig) -> int:
|
||||
"""Rough peak transient (activations, GEMM outputs) per PE."""
|
||||
m = cfg.model
|
||||
T_q = cfg.topo.T_q
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
|
||||
if cfg.topo.mode == "decode":
|
||||
return 4 * (m.hidden + m.head_dim_total_q // cfg.topo.tp) * m.bytes_per_elem
|
||||
else:
|
||||
TILE = 1024
|
||||
tile_score = hq_per_pe * T_q * TILE * m.bytes_per_elem
|
||||
return int(2 * tile_score + T_q * m.hidden * m.bytes_per_elem)
|
||||
|
||||
|
||||
def compute_memory(cfg: FullConfig,
|
||||
include_attention: bool = True,
|
||||
include_ffn: bool = True) -> MemoryBreakdown:
|
||||
"""Per-PE memory breakdown. Scope flags let callers count only the
|
||||
attention or only the FFN block (KV cache goes with the attention
|
||||
block; transient activation buffer is small either way and left in).
|
||||
"""
|
||||
kv_bytes = per_pe_kv_cache_bytes(cfg) if include_attention else 0
|
||||
return MemoryBreakdown(
|
||||
weights_bytes=per_pe_weight_bytes(
|
||||
cfg, include_attention=include_attention, include_ffn=include_ffn,
|
||||
),
|
||||
kv_cache_bytes=kv_bytes,
|
||||
transient_bytes=per_pe_transient_bytes(cfg),
|
||||
budget_bytes=cfg.machine.pe_budget_bytes,
|
||||
)
|
||||
|
||||
|
||||
def _one_layer_row(name: str,
|
||||
global_shape: tuple[int, int],
|
||||
per_pe_shape: tuple[int, int],
|
||||
bytes_per_elem: int) -> dict:
|
||||
"""Per-tensor row for ONE layer (both global and per-PE shard)."""
|
||||
p_global = global_shape[0] * global_shape[1]
|
||||
p_per_pe = per_pe_shape[0] * per_pe_shape[1]
|
||||
return {
|
||||
"Tensor": name,
|
||||
"Global shape": f"({global_shape[0]}, {global_shape[1]})",
|
||||
"Params/layer": f"{p_global/1e6:.2f} M",
|
||||
"Bytes/layer (global)": f"{p_global * bytes_per_elem / 1e6:.2f} MB",
|
||||
"Per-PE shape": f"({per_pe_shape[0]}, {per_pe_shape[1]})",
|
||||
"Bytes/layer (per PE)": f"{p_per_pe * bytes_per_elem / 1e6:.2f} MB",
|
||||
"_p_global": p_global,
|
||||
"_p_per_pe": p_per_pe,
|
||||
}
|
||||
|
||||
|
||||
def attention_weight_rows(cfg: FullConfig) -> list[dict]:
|
||||
"""Per-tensor rows for attention weights (one layer)."""
|
||||
m = cfg.model
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
hkv_per_pe = max(1, m.h_kv // cfg.topo.tp)
|
||||
return [
|
||||
_one_layer_row("W_Q", (m.hidden, m.h_q * m.d_head),
|
||||
(m.hidden, hq_per_pe * m.d_head),
|
||||
m.bytes_per_elem),
|
||||
_one_layer_row("W_K", (m.hidden, m.h_kv * m.d_head),
|
||||
(m.hidden, hkv_per_pe * m.d_head),
|
||||
m.bytes_per_elem),
|
||||
_one_layer_row("W_V", (m.hidden, m.h_kv * m.d_head),
|
||||
(m.hidden, hkv_per_pe * m.d_head),
|
||||
m.bytes_per_elem),
|
||||
_one_layer_row("W_O", (m.h_q * m.d_head, m.hidden),
|
||||
(hq_per_pe * m.d_head, m.hidden),
|
||||
m.bytes_per_elem),
|
||||
]
|
||||
|
||||
|
||||
def ffn_weight_rows(cfg: FullConfig) -> list[dict]:
|
||||
"""Per-tensor rows for FFN weights (one layer, activated for MoE)."""
|
||||
m = cfg.model
|
||||
ffn_per_pe = m.ffn_dim // cfg.topo.tp
|
||||
return [
|
||||
_one_layer_row("W_gate", (m.hidden, m.ffn_dim),
|
||||
(m.hidden, ffn_per_pe), m.bytes_per_elem),
|
||||
_one_layer_row("W_up", (m.hidden, m.ffn_dim),
|
||||
(m.hidden, ffn_per_pe), m.bytes_per_elem),
|
||||
_one_layer_row("W_down", (m.ffn_dim, m.hidden),
|
||||
(ffn_per_pe, m.hidden), m.bytes_per_elem),
|
||||
]
|
||||
|
||||
|
||||
def kv_cache_rows(cfg: FullConfig) -> list[dict]:
|
||||
"""Per-tensor row for K and V cache (one layer)."""
|
||||
m = cfg.model
|
||||
hkv_per_pe = max(1, m.h_kv // cfg.topo.tp)
|
||||
b = m.bytes_per_elem
|
||||
global_k = cfg.topo.s_kv * m.h_kv * m.d_head
|
||||
per_pe_k = cfg.topo.s_local * hkv_per_pe * m.d_head
|
||||
return [
|
||||
{
|
||||
"Tensor": "K cache",
|
||||
"Global shape": f"({cfg.topo.s_kv:,}, {m.h_kv * m.d_head})",
|
||||
"Params/layer": f"{global_k/1e6:.2f} M",
|
||||
"Bytes/layer (global)": f"{global_k * b / 1e6:.2f} MB",
|
||||
"Per-PE shape": f"({cfg.topo.s_local:,}, {hkv_per_pe * m.d_head})",
|
||||
"Bytes/layer (per PE)": f"{per_pe_k * b / 1e6:.2f} MB",
|
||||
"_p_global": global_k,
|
||||
"_p_per_pe": per_pe_k,
|
||||
},
|
||||
{
|
||||
"Tensor": "V cache",
|
||||
"Global shape": f"({cfg.topo.s_kv:,}, {m.h_kv * m.d_head})",
|
||||
"Params/layer": f"{global_k/1e6:.2f} M",
|
||||
"Bytes/layer (global)": f"{global_k * b / 1e6:.2f} MB",
|
||||
"Per-PE shape": f"({cfg.topo.s_local:,}, {hkv_per_pe * m.d_head})",
|
||||
"Bytes/layer (per PE)": f"{per_pe_k * b / 1e6:.2f} MB",
|
||||
"_p_global": global_k,
|
||||
"_p_per_pe": per_pe_k,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def sum_bytes_all_layers(rows: list[dict], cfg: FullConfig,
|
||||
per_pe: bool = False) -> int:
|
||||
"""Sum bytes across all rows × N layers (or layers/PP for per_pe)."""
|
||||
b = cfg.model.bytes_per_elem
|
||||
layers = cfg.model.layers
|
||||
if per_pe:
|
||||
layers = (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
|
||||
key = "_p_per_pe" if per_pe else "_p_global"
|
||||
return sum(r[key] for r in rows) * b * layers
|
||||
|
||||
|
||||
def total_weight_bytes_full_model(cfg: FullConfig) -> int:
|
||||
"""Total bytes of ALL weights (attention + FFN) for the full unsharded model."""
|
||||
m = cfg.model
|
||||
attn_per_layer = (
|
||||
m.hidden * m.h_q * m.d_head + # W_Q
|
||||
m.hidden * m.h_kv * m.d_head * 2 + # W_K, W_V
|
||||
m.h_q * m.d_head * m.hidden # W_O
|
||||
)
|
||||
ffn_per_layer = 3 * m.hidden * m.ffn_dim
|
||||
return (attn_per_layer + ffn_per_layer) * m.bytes_per_elem * m.layers
|
||||
|
||||
|
||||
def total_kv_bytes_full_model(cfg: FullConfig) -> int:
|
||||
"""Total bytes of KV cache for full unsharded model at cfg.topo.s_kv."""
|
||||
m = cfg.model
|
||||
per_layer = 2 * cfg.topo.s_kv * m.h_kv * m.d_head * m.bytes_per_elem
|
||||
return per_layer * m.layers
|
||||
@@ -1,311 +0,0 @@
|
||||
"""Model + machine parameters for the analytical visualization tool.
|
||||
|
||||
All numbers are per-rank / per-PE unless noted. bytes are counted in
|
||||
bytes (b), FLOPs in raw ops, bandwidth in bytes/second, latency in
|
||||
seconds. Convert to μs / GB / GFLOP at display time only.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelConfig:
|
||||
"""Transformer model dimensions."""
|
||||
name: str = "Qwen 3 8B"
|
||||
hidden: int = 4096
|
||||
ffn_dim: int = 12288
|
||||
h_q: int = 32 # Q heads
|
||||
h_kv: int = 8 # KV heads
|
||||
d_head: int = 128
|
||||
layers: int = 36
|
||||
bytes_per_elem: int = 2 # bf16
|
||||
|
||||
@property
|
||||
def group_size(self) -> int:
|
||||
"""GQA group size = H_q / H_kv."""
|
||||
return self.h_q // self.h_kv
|
||||
|
||||
@property
|
||||
def head_dim_total_q(self) -> int:
|
||||
return self.h_q * self.d_head
|
||||
|
||||
@property
|
||||
def head_dim_total_kv(self) -> int:
|
||||
return self.h_kv * self.d_head
|
||||
|
||||
|
||||
@dataclass
|
||||
class TopologyConfig:
|
||||
"""Physical mapping — DP × PP × CP × TP (× EP for MoE)."""
|
||||
cp: int = 4 # inter-cube ring size (sequence shard)
|
||||
tp: int = 8 # PEs per cube (head/hidden shard)
|
||||
pp: int = 1 # pipeline stages (layer shard)
|
||||
dp: int = 1 # data-parallel replicas
|
||||
ep: int = 1 # expert-parallel degree (MoE only)
|
||||
b: int = 1 # batch size (concurrent requests per PP stage)
|
||||
s_kv: int = 4096
|
||||
mode: str = "decode" # "decode" or "prefill"
|
||||
kv_shard_mode: str = "split" # "split" or "replicate" (used when TP > H_kv)
|
||||
# FFN sharding scope: default TP+CP (across cube group) trades a small
|
||||
# extra AllReduce for lower per-PE weight bytes.
|
||||
ffn_shard_scope: str = "TP+CP" # "TP" | "TP+CP" | "TP+CP+DP"
|
||||
# Inter-SIP interconnect: how physical SIPs are wired together.
|
||||
# "ring" (1D wrap), "mesh2d" (grid, no wrap), "torus2d" (grid + wrap).
|
||||
sip_topology: str = "ring"
|
||||
# Placement of parallelism dims on the hierarchy.
|
||||
# "pe" -> dim lives on PEs within a single cube (intra-cube).
|
||||
# "cube" -> dim lives across cubes (inter-cube).
|
||||
tp_placement: str = "pe" # default: TP fills PEs of a cube
|
||||
cp_placement: str = "cube" # default: CP fills cubes of a SIP
|
||||
# CP ring variant: what gets passed around the CP ring each hop.
|
||||
# "kv" -> K and V shards rotate (Q,O,m,l stay local). Bytes/hop scales
|
||||
# with S_local * H_kv * d_h. Good for prefill (large T_q).
|
||||
# "qoml" -> Q and running (O, m, l) rotate; K,V stay local. Bytes/hop
|
||||
# scales with T_q * H_q * d_h. Cheap for decode (T_q=1).
|
||||
cp_ring_variant: str = "kv"
|
||||
pes_per_cube_hw: int = 8 # SIP hardware constant
|
||||
cubes_per_sip_hw: int = 16
|
||||
|
||||
@property
|
||||
def pes_per_stage(self) -> int:
|
||||
"""PEs per pipeline stage (all stages have the same size)."""
|
||||
return self.cp * self.tp
|
||||
|
||||
@property
|
||||
def pes_per_replica(self) -> int:
|
||||
"""PEs per one full model replica (CP*TP*PP)."""
|
||||
return self.cp * self.tp * self.pp
|
||||
|
||||
@property
|
||||
def total_pes(self) -> int:
|
||||
return self.pes_per_replica * self.dp
|
||||
|
||||
# ── Placement-aware layout ────────────────────────────────
|
||||
@property
|
||||
def intra_cube_dims(self) -> int:
|
||||
"""Product of dims placed on PE level (must fit in one cube, else spills)."""
|
||||
n = 1
|
||||
if self.tp_placement == "pe":
|
||||
n *= self.tp
|
||||
if self.cp_placement == "pe":
|
||||
n *= self.cp
|
||||
return n
|
||||
|
||||
@property
|
||||
def inter_cube_dims(self) -> int:
|
||||
"""Product of dims placed at cube level = distinct cube groups per stage."""
|
||||
n = 1
|
||||
if self.tp_placement == "cube":
|
||||
n *= self.tp
|
||||
if self.cp_placement == "cube":
|
||||
n *= self.cp
|
||||
return n
|
||||
|
||||
@property
|
||||
def pes_per_cube_used(self) -> int:
|
||||
"""How many PEs are live in each cube."""
|
||||
return min(self.intra_cube_dims, self.pes_per_cube_hw)
|
||||
|
||||
@property
|
||||
def cubes_per_stage(self) -> int:
|
||||
"""Cubes needed per PP stage under current placement.
|
||||
If intra-cube demand exceeds PEs/cube, the intra dim spills to cubes."""
|
||||
spill = max(1, (self.intra_cube_dims + self.pes_per_cube_hw - 1)
|
||||
// self.pes_per_cube_hw)
|
||||
return self.inter_cube_dims * spill
|
||||
|
||||
@property
|
||||
def cubes_used(self) -> int:
|
||||
"""Physical cubes used across all PP stages and DP replicas."""
|
||||
return self.cubes_per_stage * self.pp * self.dp
|
||||
|
||||
@property
|
||||
def sips_used(self) -> int:
|
||||
"""Number of SIPs used (each SIP = pes_per_cube_hw × cubes_per_sip_hw)."""
|
||||
cubes_per_sip = self.cubes_per_sip_hw
|
||||
return max(1, (self.cubes_used + cubes_per_sip - 1) // cubes_per_sip)
|
||||
|
||||
@property
|
||||
def placement_valid(self) -> bool:
|
||||
"""False when intra-cube demand exceeds one cube (dim spills to cubes)."""
|
||||
return self.intra_cube_dims <= self.pes_per_cube_hw
|
||||
|
||||
@property
|
||||
def tp_spans_cubes(self) -> int:
|
||||
"""How many cubes one TP group spans (>=1). Depends on placement:
|
||||
- tp_placement=pe: 1 if TP fits in a cube, else spill.
|
||||
- tp_placement=cube: TP itself is the number of cubes per TP group.
|
||||
"""
|
||||
if self.tp_placement == "cube":
|
||||
return max(1, self.tp)
|
||||
return max(1, (self.tp + self.pes_per_cube_hw - 1) // self.pes_per_cube_hw)
|
||||
|
||||
@property
|
||||
def cp_intra_sip_hops(self) -> int:
|
||||
"""CP ring hops that stay within one SIP."""
|
||||
if self.cp <= 1:
|
||||
return 0
|
||||
sips_in_ring = self.sips_used # one PP stage's CP ring may span SIPs
|
||||
return max(0, (self.cp - 1) - (sips_in_ring - 1))
|
||||
|
||||
@property
|
||||
def cp_inter_sip_hops(self) -> int:
|
||||
"""CP ring hops that cross SIP boundaries."""
|
||||
if self.cp <= 1:
|
||||
return 0
|
||||
# Approximation: one inter-SIP hop per SIP boundary in the ring.
|
||||
# For CP=32 across 2 SIPs → 1 inter-SIP + 30 intra-SIP hops.
|
||||
return max(0, self.sips_used - 1)
|
||||
|
||||
@property
|
||||
def T_q(self) -> int:
|
||||
return 1 if self.mode == "decode" else self.s_kv // self.cp
|
||||
|
||||
@property
|
||||
def s_local(self) -> int:
|
||||
return self.s_kv // self.cp
|
||||
|
||||
@property
|
||||
def layers_per_stage(self) -> str:
|
||||
"""Layers per PP stage. Requires model.layers to render — leave as fn."""
|
||||
return "N_layers / PP"
|
||||
|
||||
def tp_link_tier(self) -> str:
|
||||
"""Return which BW tier a TP AllReduce uses: 'intra' | 'inter' | 'intersip'.
|
||||
- tp_placement=pe and TP fits in one cube -> 'intra'
|
||||
- tp_placement=cube and TP fits in one SIP -> 'inter'
|
||||
- otherwise -> 'intersip'
|
||||
"""
|
||||
if self.tp_placement == "pe":
|
||||
if self.tp <= self.pes_per_cube_hw:
|
||||
return "intra"
|
||||
# spilled to multiple cubes: cross-cube
|
||||
if self.tp_spans_cubes <= self.cubes_per_sip_hw:
|
||||
return "inter"
|
||||
return "intersip"
|
||||
else: # tp_placement == "cube"
|
||||
if self.tp <= self.cubes_per_sip_hw:
|
||||
return "inter"
|
||||
return "intersip"
|
||||
|
||||
def cp_link_tier(self) -> str:
|
||||
"""Return BW tier for the CP ring: 'intra' | 'inter' | 'intersip'."""
|
||||
if self.cp_placement == "pe":
|
||||
return "intra"
|
||||
# cp_placement == "cube"
|
||||
if self.cp <= self.cubes_per_sip_hw:
|
||||
return "inter"
|
||||
return "intersip"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MachineParams:
|
||||
"""SIP hardware parameters — all user-adjustable via sliders.
|
||||
|
||||
Bandwidth tiers, from fastest to slowest:
|
||||
HBM (per-PE local memory)
|
||||
> intra-cube (PE↔PE on same die)
|
||||
> inter-cube (UCIe D2D between dies on same SIP)
|
||||
> inter-SIP (C2C / RDMA-class, across SIPs / servers)
|
||||
"""
|
||||
# Per-PE compute peak (TFLOPs f16).
|
||||
peak_tflops_f16: float = 8.0
|
||||
# Per-PE HBM budget (GB) — the memory ceiling per PE.
|
||||
pe_hbm_gb: float = 6.0
|
||||
# HBM per-PE bandwidth (GB/s).
|
||||
bw_hbm_gbs: float = 256.0
|
||||
# Intra-cube PE↔PE link BW (GB/s).
|
||||
bw_intra_gbs: float = 512.0
|
||||
# Inter-cube (UCIe D2D) BW (GB/s).
|
||||
bw_inter_gbs: float = 128.0
|
||||
# Inter-SIP (C2C / RDMA) BW (GB/s).
|
||||
bw_intersip_gbs: float = 50.0
|
||||
# Per-hop latencies (ns).
|
||||
alpha_intra_ns: float = 20.0
|
||||
alpha_inter_ns: float = 100.0
|
||||
alpha_intersip_ns: float = 1000.0
|
||||
# Achievable utilization for compute-bound stages (0-1).
|
||||
compute_util: float = 0.8
|
||||
|
||||
@property
|
||||
def peak_flops(self) -> float:
|
||||
return self.peak_tflops_f16 * 1e12
|
||||
|
||||
@property
|
||||
def bw_hbm(self) -> float:
|
||||
return self.bw_hbm_gbs * 1e9
|
||||
|
||||
@property
|
||||
def bw_intra(self) -> float:
|
||||
return self.bw_intra_gbs * 1e9
|
||||
|
||||
@property
|
||||
def bw_inter(self) -> float:
|
||||
return self.bw_inter_gbs * 1e9
|
||||
|
||||
@property
|
||||
def bw_intersip(self) -> float:
|
||||
return self.bw_intersip_gbs * 1e9
|
||||
|
||||
@property
|
||||
def alpha_intra(self) -> float:
|
||||
return self.alpha_intra_ns * 1e-9
|
||||
|
||||
@property
|
||||
def alpha_inter(self) -> float:
|
||||
return self.alpha_inter_ns * 1e-9
|
||||
|
||||
@property
|
||||
def alpha_intersip(self) -> float:
|
||||
return self.alpha_intersip_ns * 1e-9
|
||||
|
||||
@property
|
||||
def pe_budget_bytes(self) -> int:
|
||||
return int(self.pe_hbm_gb * 1e9)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FullConfig:
|
||||
"""Bundled config for one analysis run."""
|
||||
model: ModelConfig = field(default_factory=ModelConfig)
|
||||
topo: TopologyConfig = field(default_factory=TopologyConfig)
|
||||
machine: MachineParams = field(default_factory=MachineParams)
|
||||
|
||||
@property
|
||||
def h_q_per_pe(self) -> int:
|
||||
return max(1, self.model.h_q // self.topo.tp)
|
||||
|
||||
@property
|
||||
def h_kv_per_pe(self) -> float:
|
||||
"""Fractional if TP > H_kv (head-dim split needed)."""
|
||||
return self.model.h_kv / self.topo.tp
|
||||
|
||||
@property
|
||||
def d_per_pe(self) -> int:
|
||||
return self.model.hidden // self.topo.tp
|
||||
|
||||
@property
|
||||
def ffn_per_pe(self) -> int:
|
||||
return self.model.ffn_dim // self.topo.tp
|
||||
|
||||
@property
|
||||
def kv_replication_needed(self) -> bool:
|
||||
"""True if TP > H_kv (each KV head shared across TP/H_kv ranks)."""
|
||||
return self.topo.tp > self.model.h_kv
|
||||
|
||||
@property
|
||||
def head_dim_split_factor(self) -> int:
|
||||
"""Number of TP ranks sharing one head (via d_head split)."""
|
||||
return max(1, self.topo.tp // self.model.h_kv)
|
||||
|
||||
@property
|
||||
def ffn_shard_divisor(self) -> int:
|
||||
"""How much the FFN dim is sharded across ranks."""
|
||||
scope = self.topo.ffn_shard_scope
|
||||
d = self.topo.tp
|
||||
if "CP" in scope:
|
||||
d *= self.topo.cp
|
||||
if "DP" in scope:
|
||||
d *= self.topo.dp
|
||||
return max(1, d)
|
||||
@@ -1,126 +0,0 @@
|
||||
"""Model preset library for the analytical visualization tool.
|
||||
|
||||
Comprehensive coverage of popular MHA, GQA, MQA, and MLA models across
|
||||
sizes. MoE models are approximated as dense with (ffn_dim = activated
|
||||
experts × per_expert_ffn) for compute-side estimates; the note field
|
||||
flags this. MLA models are approximated as GQA with matching H_kv until
|
||||
we add proper MLA support.
|
||||
|
||||
Source of numeric params: model cards / config.json from HuggingFace.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .model_config import ModelConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class Preset:
|
||||
label: str
|
||||
model: ModelConfig
|
||||
family: str = "" # e.g. "Llama 3", "Qwen 3"
|
||||
attn_type: str = "" # "MHA" | "GQA" | "MQA" | "MLA"
|
||||
note: str = ""
|
||||
|
||||
|
||||
def _mk(name: str, hidden: int, ffn: int, hq: int, hkv: int,
|
||||
dh: int, layers: int, *, family: str = "", attn: str = "GQA",
|
||||
note: str = "") -> Preset:
|
||||
return Preset(
|
||||
label=name,
|
||||
model=ModelConfig(
|
||||
name=name, hidden=hidden, ffn_dim=ffn,
|
||||
h_q=hq, h_kv=hkv, d_head=dh, layers=layers,
|
||||
),
|
||||
family=family, attn_type=attn, note=note,
|
||||
)
|
||||
|
||||
|
||||
PRESETS: dict[str, Preset] = {
|
||||
# ── MHA (H_q == H_kv) ─────────────────────────────────────────
|
||||
"GPT-2 (117M)": _mk("GPT-2 117M", 768, 3072, 12, 12, 64, 12, family="GPT-2", attn="MHA"),
|
||||
"GPT-2 XL (1.5B)": _mk("GPT-2 XL 1.5B", 1600, 6400, 25, 25, 64, 48, family="GPT-2", attn="MHA"),
|
||||
"Llama 2 7B": _mk("Llama 2 7B", 4096, 11008, 32, 32, 128, 32, family="Llama 2", attn="MHA"),
|
||||
"Llama 2 13B": _mk("Llama 2 13B", 5120, 13824, 40, 40, 128, 40, family="Llama 2", attn="MHA"),
|
||||
"Yi 6B": _mk("Yi 6B", 4096, 11008, 32, 4, 128, 32, family="Yi", attn="GQA"),
|
||||
"Yi 9B": _mk("Yi 9B", 4096, 11008, 32, 4, 128, 48, family="Yi", attn="GQA"),
|
||||
"OPT 30B": _mk("OPT 30B", 7168, 28672, 56, 56, 128, 48, family="OPT", attn="MHA"),
|
||||
|
||||
# ── GQA (H_q > H_kv > 1) ──────────────────────────────────────
|
||||
"Qwen 2 0.5B": _mk("Qwen 2 0.5B", 896, 4864, 14, 2, 64, 24, family="Qwen 2"),
|
||||
"Qwen 2 1.5B": _mk("Qwen 2 1.5B", 1536, 8960, 12, 2, 128, 28, family="Qwen 2"),
|
||||
"Qwen 2 7B": _mk("Qwen 2 7B", 3584, 18944, 28, 4, 128, 28, family="Qwen 2"),
|
||||
"Qwen 2 72B": _mk("Qwen 2 72B", 8192, 29568, 64, 8, 128, 80, family="Qwen 2"),
|
||||
|
||||
"Qwen 3 0.6B": _mk("Qwen 3 0.6B", 1024, 3072, 16, 8, 128, 28, family="Qwen 3"),
|
||||
"Qwen 3 1.7B": _mk("Qwen 3 1.7B", 2048, 6144, 16, 8, 128, 28, family="Qwen 3"),
|
||||
"Qwen 3 4B": _mk("Qwen 3 4B", 2560, 9728, 32, 8, 128, 36, family="Qwen 3"),
|
||||
"Qwen 3 8B": _mk("Qwen 3 8B", 4096, 12288, 32, 8, 128, 36, family="Qwen 3"),
|
||||
"Qwen 3 14B": _mk("Qwen 3 14B", 5120, 17408, 40, 8, 128, 40, family="Qwen 3"),
|
||||
"Qwen 3 32B": _mk("Qwen 3 32B", 5120, 25600, 64, 8, 128, 64, family="Qwen 3"),
|
||||
|
||||
"Llama 3 8B": _mk("Llama 3 8B", 4096, 14336, 32, 8, 128, 32, family="Llama 3"),
|
||||
"Llama 3 70B": _mk("Llama 3 70B", 8192, 28672, 64, 8, 128, 80, family="Llama 3"),
|
||||
"Llama 3.1 8B": _mk("Llama 3.1 8B", 4096, 14336, 32, 8, 128, 32, family="Llama 3.1"),
|
||||
"Llama 3.1 70B": _mk("Llama 3.1 70B", 8192, 28672, 64, 8, 128, 80, family="Llama 3.1"),
|
||||
"Llama 3.1 405B": _mk("Llama 3.1 405B", 16384,53248, 128,8, 128, 126,family="Llama 3.1"),
|
||||
"Llama 3.2 1B": _mk("Llama 3.2 1B", 2048, 8192, 32, 8, 64, 16, family="Llama 3.2"),
|
||||
"Llama 3.2 3B": _mk("Llama 3.2 3B", 3072, 8192, 24, 8, 128, 28, family="Llama 3.2"),
|
||||
|
||||
"Mistral 7B": _mk("Mistral 7B v0.3", 4096, 14336, 32, 8, 128, 32, family="Mistral"),
|
||||
"Mistral Nemo 12B": _mk("Mistral Nemo 12B", 5120, 14336, 32, 8, 128, 40, family="Mistral"),
|
||||
"Mistral Small 22B": _mk("Mistral Small 22B",6144, 16384, 48, 8, 128, 56, family="Mistral"),
|
||||
"Mistral Large 123B": _mk("Mistral Large 123B",12288,28672,96, 8, 128, 88, family="Mistral"),
|
||||
|
||||
"Gemma 2 2B": _mk("Gemma 2 2B", 2304, 9216, 8, 4, 256, 26, family="Gemma 2"),
|
||||
"Gemma 2 9B": _mk("Gemma 2 9B", 3584, 14336, 16, 8, 256, 42, family="Gemma 2"),
|
||||
"Gemma 2 27B": _mk("Gemma 2 27B", 4608, 36864, 32, 16, 128, 46, family="Gemma 2"),
|
||||
|
||||
"Phi 3 mini (3.8B)": _mk("Phi 3 mini 3.8B", 3072, 8192, 32, 32, 96, 32, family="Phi 3", attn="MHA"),
|
||||
"Phi 3 small (7B)": _mk("Phi 3 small 7B", 4096, 14336, 32, 8, 128, 32, family="Phi 3"),
|
||||
"Phi 3 medium (14B)": _mk("Phi 3 medium 14B", 5120, 17920, 40, 10, 128, 40, family="Phi 3"),
|
||||
|
||||
"Yi 34B": _mk("Yi 34B", 7168, 20480, 56, 8, 128, 60, family="Yi"),
|
||||
"Command R+ 104B": _mk("Command R+ 104B", 12288,33792, 96, 8, 128, 64, family="Cohere"),
|
||||
|
||||
# ── MoE (activated-only approximation) ────────────────────────
|
||||
"Mixtral 8x7B (MoE)": _mk("Mixtral 8x7B", 4096, 14336*2, 32, 8, 128, 32,
|
||||
family="Mistral MoE",
|
||||
note="MoE: 8 experts x 2 activated. FFN ~ 2 x per_expert."),
|
||||
"Mixtral 8x22B (MoE)":_mk("Mixtral 8x22B", 6144, 16384*2, 48, 8, 128, 56,
|
||||
family="Mistral MoE",
|
||||
note="MoE: 8 experts x 2 activated."),
|
||||
"Qwen 3 30B (MoE)": _mk("Qwen 3 30B (MoE)", 2048, 768*8, 32, 4, 128, 48,
|
||||
family="Qwen 3 MoE",
|
||||
note="128 experts x 8 activated per token."),
|
||||
"Qwen 3 235B (MoE)": _mk("Qwen 3 235B (MoE)",4096, 1536*8, 64, 4, 128, 94,
|
||||
family="Qwen 3 MoE",
|
||||
note="128 experts x 8 activated. Dense-approx overstates weight."),
|
||||
"DeepSeek V2 (dense-approx)":
|
||||
_mk("DeepSeek V2", 5120, 12288, 128, 128, 128, 60,
|
||||
family="DeepSeek", attn="MLA",
|
||||
note="MLA compresses KV to ~576 B/token; d_c=512 latent. Dense-approx."),
|
||||
"DeepSeek V3 (dense-approx)":
|
||||
_mk("DeepSeek V3", 7168, 16384, 128, 128, 128, 61,
|
||||
family="DeepSeek", attn="MLA",
|
||||
note="MoE 256 x 8 activated + MLA. Both approximations."),
|
||||
"Grok-1 314B (MoE)": _mk("Grok-1 314B", 6144, 32768*2, 48, 8, 128, 64,
|
||||
family="xAI",
|
||||
note="MoE 8 experts x 2 activated."),
|
||||
|
||||
# ── MQA (H_kv == 1) ───────────────────────────────────────────
|
||||
"Falcon 7B (MQA)": _mk("Falcon 7B", 4544, 18176, 71, 1, 64, 32, family="Falcon", attn="MQA"),
|
||||
"Falcon 40B (MQA)": _mk("Falcon 40B", 8192, 32768, 128,1, 64, 60, family="Falcon", attn="MQA"),
|
||||
|
||||
# ── Custom slot ──────────────────────────────────────────────
|
||||
"Custom": Preset(label="Custom", model=ModelConfig(name="Custom")),
|
||||
}
|
||||
|
||||
|
||||
def preset_names_by_family() -> dict[str, list[str]]:
|
||||
"""Return preset names grouped by family (for a nested dropdown)."""
|
||||
grouped: dict[str, list[str]] = {}
|
||||
for name, p in PRESETS.items():
|
||||
grouped.setdefault(p.family or "Other", []).append(name)
|
||||
return grouped
|
||||
@@ -1,235 +0,0 @@
|
||||
"""Replication (waste) accounting + optimization opportunities.
|
||||
|
||||
Two things this module produces:
|
||||
|
||||
1. `replication_report(cfg)` — where memory is being *duplicated* across
|
||||
the deployment, per category (attention weights, FFN weights, KV cache
|
||||
under replicate mode, DP replicas). Answers "if I lifted this
|
||||
duplication, how much would I save?"
|
||||
|
||||
2. `optimization_hints(cfg)` — a list of actionable hints. Each hint has
|
||||
a `category` (space | comm | compute), a `severity` (info | warn |
|
||||
good), and a short human message. UI code renders these as bullets.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .model_config import FullConfig
|
||||
from .memory_layout import per_pe_weight_bytes, per_pe_kv_cache_bytes
|
||||
|
||||
|
||||
# ── Replication accounting ────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class ReplicationEntry:
|
||||
tensor: str
|
||||
per_pe_bytes: int
|
||||
replicated_across: str # e.g. "CP (x4 copies)"
|
||||
copies: int # total # copies globally
|
||||
wasted_bytes: int # (copies - 1) * per_pe_bytes (per PE view)
|
||||
|
||||
|
||||
def _attn_weight_bytes_per_pe(cfg: FullConfig) -> int:
|
||||
"""W_Q + W_K + W_V + W_O per PE, one layer."""
|
||||
m = cfg.model
|
||||
tp = cfg.topo.tp
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
if cfg.topo.kv_shard_mode == "replicate":
|
||||
hkv_per_pe = max(1.0, m.h_kv / tp)
|
||||
else:
|
||||
hkv_per_pe = m.h_kv / tp
|
||||
per_layer = (
|
||||
m.hidden * hq_per_pe * m.d_head +
|
||||
m.hidden * hkv_per_pe * m.d_head +
|
||||
m.hidden * hkv_per_pe * m.d_head +
|
||||
hq_per_pe * m.d_head * m.hidden
|
||||
) * m.bytes_per_elem
|
||||
layers_per_stage = (m.layers + cfg.topo.pp - 1) // cfg.topo.pp
|
||||
return int(per_layer * layers_per_stage)
|
||||
|
||||
|
||||
def _ffn_weight_bytes_per_pe(cfg: FullConfig) -> int:
|
||||
"""W_gate + W_up + W_down per PE, all local layers."""
|
||||
m = cfg.model
|
||||
ep = max(1, cfg.topo.ep)
|
||||
ffn_div = cfg.ffn_shard_divisor * ep
|
||||
per_layer = 3 * m.hidden * (m.ffn_dim // max(1, ffn_div)) * m.bytes_per_elem
|
||||
layers_per_stage = (m.layers + cfg.topo.pp - 1) // cfg.topo.pp
|
||||
return int(per_layer * layers_per_stage)
|
||||
|
||||
|
||||
def _kv_bytes_per_pe(cfg: FullConfig) -> int:
|
||||
return per_pe_kv_cache_bytes(cfg)
|
||||
|
||||
|
||||
def replication_report(cfg: FullConfig) -> list[ReplicationEntry]:
|
||||
"""Enumerate replicated (duplicated) tensors and their waste."""
|
||||
entries: list[ReplicationEntry] = []
|
||||
topo = cfg.topo
|
||||
|
||||
# Attention weights are NOT sharded by CP - so each CP rank holds a copy.
|
||||
attn_pe = _attn_weight_bytes_per_pe(cfg)
|
||||
cp_copies = topo.cp
|
||||
if cp_copies > 1:
|
||||
entries.append(ReplicationEntry(
|
||||
tensor="Attn weights (W_Q/W_K/W_V/W_O)",
|
||||
per_pe_bytes=attn_pe,
|
||||
replicated_across=f"CP (x{cp_copies} identical group copies)",
|
||||
copies=cp_copies,
|
||||
wasted_bytes=attn_pe * (cp_copies - 1),
|
||||
))
|
||||
|
||||
# FFN weights: replicated across scope levels NOT included.
|
||||
# Scope=TP -> replicated across CP AND DP.
|
||||
# Scope=TP+CP -> replicated across DP only.
|
||||
# Scope=TP+CP+DP -> no replication (perfect share).
|
||||
ffn_pe = _ffn_weight_bytes_per_pe(cfg)
|
||||
scope = topo.ffn_shard_scope
|
||||
ffn_cp_copies = 1 if "CP" in scope else topo.cp
|
||||
ffn_dp_copies = 1 if "DP" in scope else topo.dp
|
||||
ffn_total_copies = ffn_cp_copies * ffn_dp_copies
|
||||
if ffn_total_copies > 1:
|
||||
reasons = []
|
||||
if ffn_cp_copies > 1:
|
||||
reasons.append(f"CP (x{ffn_cp_copies})")
|
||||
if ffn_dp_copies > 1:
|
||||
reasons.append(f"DP (x{ffn_dp_copies})")
|
||||
entries.append(ReplicationEntry(
|
||||
tensor=f"FFN weights (scope={scope})",
|
||||
per_pe_bytes=ffn_pe,
|
||||
replicated_across=" & ".join(reasons),
|
||||
copies=ffn_total_copies,
|
||||
wasted_bytes=ffn_pe * (ffn_total_copies - 1),
|
||||
))
|
||||
|
||||
# KV cache: if kv_shard_mode='replicate' and TP > H_kv, each KV head
|
||||
# is duplicated tp // h_kv times across TP ranks.
|
||||
if (cfg.kv_replication_needed
|
||||
and topo.kv_shard_mode == "replicate"):
|
||||
kv_pe = _kv_bytes_per_pe(cfg)
|
||||
rep = cfg.head_dim_split_factor # tp // h_kv
|
||||
entries.append(ReplicationEntry(
|
||||
tensor="KV cache (replicate mode)",
|
||||
per_pe_bytes=kv_pe,
|
||||
replicated_across=f"TP (x{rep} copies of each KV head)",
|
||||
copies=rep,
|
||||
wasted_bytes=kv_pe * (rep - 1),
|
||||
))
|
||||
|
||||
# DP replicas duplicate the entire per-PE footprint (attn + ffn + KV).
|
||||
if topo.dp > 1:
|
||||
total_pe = attn_pe + ffn_pe + _kv_bytes_per_pe(cfg)
|
||||
entries.append(ReplicationEntry(
|
||||
tensor="Full per-PE footprint (attn + FFN + KV)",
|
||||
per_pe_bytes=total_pe,
|
||||
replicated_across=f"DP (x{topo.dp} model replicas)",
|
||||
copies=topo.dp,
|
||||
wasted_bytes=total_pe * (topo.dp - 1),
|
||||
))
|
||||
|
||||
return entries
|
||||
|
||||
|
||||
# ── Optimization hints ────────────────────────────────────────────
|
||||
|
||||
@dataclass
|
||||
class Hint:
|
||||
category: str # "space" | "comm" | "compute" | "layout"
|
||||
severity: str # "info" | "warn" | "good"
|
||||
message: str
|
||||
|
||||
|
||||
def optimization_hints(cfg: FullConfig) -> list[Hint]:
|
||||
hints: list[Hint] = []
|
||||
m = cfg.model
|
||||
topo = cfg.topo
|
||||
|
||||
# SPACE hints -------------------------------------------------
|
||||
scope = topo.ffn_shard_scope
|
||||
ffn_pe = _ffn_weight_bytes_per_pe(cfg)
|
||||
if scope == "TP" and topo.cp > 1:
|
||||
current = ffn_pe
|
||||
new = current // topo.cp
|
||||
hints.append(Hint(
|
||||
"space", "warn",
|
||||
f"FFN scope=TP replicates FFN weights across all {topo.cp} CP "
|
||||
f"groups. Switching to **TP+CP** would drop per-PE FFN weight "
|
||||
f"from {current/1e9:.2f} GB to {new/1e9:.2f} GB "
|
||||
f"(saves {(current-new)/1e9:.2f} GB/PE) at the cost of one "
|
||||
f"extra AllReduce per layer."
|
||||
))
|
||||
if scope in ("TP", "TP+CP") and topo.dp > 1:
|
||||
base = ffn_pe if scope == "TP+CP" else ffn_pe // topo.cp
|
||||
new = base // topo.dp
|
||||
hints.append(Hint(
|
||||
"space", "info",
|
||||
f"With DP={topo.dp}, scope=**TP+CP+DP** would shard FFN across "
|
||||
f"replicas too, dropping per-PE FFN from {base/1e9:.2f} GB to "
|
||||
f"{new/1e9:.2f} GB (comes with an inter-replica AllReduce)."
|
||||
))
|
||||
if (cfg.kv_replication_needed
|
||||
and topo.kv_shard_mode == "replicate"):
|
||||
rep = cfg.head_dim_split_factor
|
||||
kv_pe = _kv_bytes_per_pe(cfg)
|
||||
hints.append(Hint(
|
||||
"space", "warn",
|
||||
f"KV mode=replicate holds each KV head {rep}x. Switching to "
|
||||
f"**split** saves ~{kv_pe*(rep-1)/rep/1e9:.2f} GB KV per PE "
|
||||
f"but adds a Score AllReduce over {rep} ranks per hop."
|
||||
))
|
||||
if topo.cp == 1 and topo.tp <= m.h_kv:
|
||||
# Suggest CP if s_kv is large and KV dominates.
|
||||
kv_pe = _kv_bytes_per_pe(cfg)
|
||||
attn_pe = _attn_weight_bytes_per_pe(cfg)
|
||||
if kv_pe > attn_pe:
|
||||
hints.append(Hint(
|
||||
"space", "info",
|
||||
f"KV cache ({kv_pe/1e9:.2f} GB) dominates weights "
|
||||
f"({attn_pe/1e9:.2f} GB) at S_kv={topo.s_kv:,}. Adding CP "
|
||||
f"shards the sequence axis (S_local = S_kv/CP)."
|
||||
))
|
||||
|
||||
# COMM hints --------------------------------------------------
|
||||
if topo.cp_inter_sip_hops > 0:
|
||||
hints.append(Hint(
|
||||
"comm", "warn",
|
||||
f"CP ring crosses {topo.cp_inter_sip_hops} SIP boundary(ies) "
|
||||
f"at {cfg.machine.bw_intersip_gbs:.0f} GB/s vs "
|
||||
f"{cfg.machine.bw_inter_gbs:.0f} GB/s intra-SIP - "
|
||||
f"consider a smaller CP or a topology that keeps CP inside one SIP."
|
||||
))
|
||||
if topo.tp_spans_cubes > 1:
|
||||
hints.append(Hint(
|
||||
"comm", "warn",
|
||||
f"TP={topo.tp} spans {topo.tp_spans_cubes} cubes, so the W_O "
|
||||
f"AllReduce runs cross-cube "
|
||||
f"({cfg.machine.bw_intra_gbs:.0f} -> "
|
||||
f"{cfg.machine.bw_inter_gbs:.0f} GB/s). A TP that fits in one "
|
||||
f"cube (TP<={topo.pes_per_cube_hw}) keeps it intra-cube."
|
||||
))
|
||||
if topo.sips_used > 1 and topo.sip_topology == "ring":
|
||||
hints.append(Hint(
|
||||
"layout", "info",
|
||||
f"With {topo.sips_used} SIPs on a ring, worst-case inter-SIP "
|
||||
f"distance is {topo.sips_used - 1} hops. Switching to **torus2d** "
|
||||
f"cuts worst-case hops to ~sqrt({topo.sips_used})."
|
||||
))
|
||||
if topo.tp <= m.h_kv:
|
||||
hints.append(Hint(
|
||||
"comm", "good",
|
||||
f"TP ({topo.tp}) <= H_kv ({m.h_kv}): each KV head lives on "
|
||||
f"exactly one TP rank - no head-split AllReduce needed."
|
||||
))
|
||||
|
||||
# COMPUTE hints -----------------------------------------------
|
||||
if topo.mode == "decode" and topo.T_q == 1:
|
||||
# Decode is memory-bound almost everywhere; note it.
|
||||
hints.append(Hint(
|
||||
"compute", "info",
|
||||
"Decode (T_q=1): most GEMMs are memory-bound. FLOPs matter far "
|
||||
"less than HBM BW here - budget attention around "
|
||||
"weight-read cost, not compute peak."
|
||||
))
|
||||
|
||||
return hints
|
||||
@@ -1,279 +0,0 @@
|
||||
"""Draw a PE-level view of one CP group (one TP group of cubes).
|
||||
|
||||
Shows how weights + KV cache are distributed across the PEs of ONE CP
|
||||
rank. Attention weights are sharded by head (H_q, H_kv split across TP);
|
||||
FFN is sharded by dim (ffn_dim / TP / EP); KV cache is sharded by CP
|
||||
(sequence) × TP (head).
|
||||
|
||||
If TP > 8 the group spans multiple cubes — layout draws all of them.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.patches as patches
|
||||
|
||||
from .model_config import FullConfig
|
||||
|
||||
|
||||
PE_COLS = 4
|
||||
PE_ROWS = 2
|
||||
PES_PER_CUBE = PE_COLS * PE_ROWS
|
||||
|
||||
|
||||
def layers_per_stage_val(cfg: FullConfig) -> int:
|
||||
return (cfg.model.layers + cfg.topo.pp - 1) // cfg.topo.pp
|
||||
|
||||
|
||||
def _q_heads_for_pe(pe_id_in_group: int, tp: int, h_q: int) -> list[int]:
|
||||
"""Return the list of Q-head indices assigned to this PE."""
|
||||
heads_per_pe = h_q / tp
|
||||
start = int(round(pe_id_in_group * heads_per_pe))
|
||||
end = int(round((pe_id_in_group + 1) * heads_per_pe))
|
||||
return list(range(start, end))
|
||||
|
||||
|
||||
def _kv_heads_for_pe(pe_id_in_group: int, tp: int, h_kv: int,
|
||||
kv_shard_mode: str) -> tuple[list[int], str]:
|
||||
"""KV heads assigned to this PE and a note about replication/split."""
|
||||
if h_kv >= tp:
|
||||
heads_per_pe = h_kv // tp
|
||||
start = pe_id_in_group * heads_per_pe
|
||||
end = start + heads_per_pe
|
||||
return list(range(start, end)), ""
|
||||
# TP > H_kv
|
||||
if kv_shard_mode == "replicate":
|
||||
rep_factor = tp // h_kv
|
||||
head_idx = pe_id_in_group // rep_factor
|
||||
return [head_idx], f"replicated x{rep_factor}"
|
||||
# split mode
|
||||
split_factor = tp // h_kv
|
||||
head_idx = pe_id_in_group // split_factor
|
||||
part = pe_id_in_group % split_factor
|
||||
return [head_idx], f"head-split ({part+1}/{split_factor})"
|
||||
|
||||
|
||||
def _per_pe_bytes(cfg: FullConfig, pe_id_in_group: int) -> dict:
|
||||
"""Return per-tensor bytes for one PE within a TP group."""
|
||||
m = cfg.model
|
||||
tp = cfg.topo.tp
|
||||
pp = cfg.topo.pp
|
||||
ep = max(1, cfg.topo.ep)
|
||||
b = m.bytes_per_elem
|
||||
layers_per_stage = (m.layers + pp - 1) // pp
|
||||
|
||||
hq_per_pe = m.h_q / tp
|
||||
if cfg.topo.kv_shard_mode == "replicate":
|
||||
hkv_per_pe = max(1.0, m.h_kv / tp)
|
||||
else:
|
||||
hkv_per_pe = m.h_kv / tp
|
||||
ffn_div = cfg.ffn_shard_divisor * ep
|
||||
ffn_per_pe = m.ffn_dim // ffn_div
|
||||
|
||||
per_layer = {
|
||||
"W_Q": int(m.hidden * hq_per_pe * m.d_head * b),
|
||||
"W_K": int(m.hidden * hkv_per_pe * m.d_head * b),
|
||||
"W_V": int(m.hidden * hkv_per_pe * m.d_head * b),
|
||||
"W_O": int(hq_per_pe * m.d_head * m.hidden * b),
|
||||
"W_gate": int(m.hidden * ffn_per_pe * b),
|
||||
"W_up": int(m.hidden * ffn_per_pe * b),
|
||||
"W_down": int(ffn_per_pe * m.hidden * b),
|
||||
}
|
||||
kv_per_layer = int(2 * cfg.topo.s_local * hkv_per_pe * m.d_head * b)
|
||||
|
||||
all_layers = {k: v * layers_per_stage for k, v in per_layer.items()}
|
||||
all_layers["KV cache"] = kv_per_layer * layers_per_stage
|
||||
# transient estimate
|
||||
if cfg.topo.mode == "decode":
|
||||
all_layers["Transient"] = int(4 * (m.hidden + hq_per_pe * m.d_head) * b)
|
||||
else:
|
||||
TILE = 1024
|
||||
all_layers["Transient"] = int(2 * hq_per_pe * cfg.topo.T_q * TILE * b
|
||||
+ cfg.topo.T_q * m.hidden * b)
|
||||
return all_layers
|
||||
|
||||
|
||||
def draw_pe_layout(cfg: FullConfig, ax=None):
|
||||
"""Draw one CP group's PEs with per-tensor weight + KV breakdown.
|
||||
|
||||
Each PE shows: PE id, Q heads, KV heads, W_Q / W_K / W_V / W_O in
|
||||
MB, FFN (gate+up+down) in MB, KV cache in GB, total in GB.
|
||||
|
||||
TP=8 -> 1 cube. TP=16 -> 2 cubes. TP=32 -> 4 cubes.
|
||||
Under placement, n_cubes reflects cubes-per-group (may include CP spill).
|
||||
"""
|
||||
tp = cfg.topo.tp
|
||||
# cubes per one cube-level group (= intra-cube dims / PEs-per-cube, rounded up).
|
||||
n_cubes = max(1, (cfg.topo.intra_cube_dims + PES_PER_CUBE - 1) // PES_PER_CUBE)
|
||||
|
||||
if ax is None:
|
||||
fig, ax = plt.subplots(figsize=(7.5 * n_cubes, 8.5))
|
||||
else:
|
||||
fig = ax.figure
|
||||
|
||||
cube_w = 7.0
|
||||
cube_h = 6.0
|
||||
cube_gap = 0.6
|
||||
|
||||
# When TP fills the cube (TP>=8), cp_placement is "cube" and CP
|
||||
# ranks span cubes — the intra-cube CP-color branch below is
|
||||
# skipped. Surface which CP rank this figure is showing so users
|
||||
# aren't guessing which of `CP` groups is drawn.
|
||||
_cube_place_tag = ""
|
||||
if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1:
|
||||
_cube_place_tag = f" [CP rank 0 of {cfg.topo.cp}]"
|
||||
|
||||
for cube_idx in range(n_cubes):
|
||||
x0 = cube_idx * (cube_w + cube_gap)
|
||||
y0 = 0
|
||||
|
||||
rect = patches.FancyBboxPatch(
|
||||
(x0, y0), cube_w, cube_h,
|
||||
boxstyle="round,pad=0.05",
|
||||
facecolor="#f8f9fa", edgecolor="#212529", linewidth=1.5,
|
||||
)
|
||||
ax.add_patch(rect)
|
||||
ax.text(x0 + cube_w / 2, y0 + cube_h + 0.1,
|
||||
f"Cube {cube_idx} (PEs {cube_idx*8}-{cube_idx*8+7})"
|
||||
+ _cube_place_tag,
|
||||
ha="center", va="bottom", fontsize=11, fontweight="bold")
|
||||
|
||||
pe_pad_x = 0.15
|
||||
pe_pad_y = 0.2
|
||||
pe_gap = 0.10
|
||||
inner_w = cube_w - 2 * pe_pad_x
|
||||
inner_h = cube_h - 2 * pe_pad_y
|
||||
pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
|
||||
pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
|
||||
|
||||
# When cp_placement=pe, CP ranks are packed intra-cube along with
|
||||
# TP ranks. Color-code by CP rank so it's obvious which PEs belong
|
||||
# to which sequence-shard group. Palette wraps beyond 8 CP ranks.
|
||||
_cp_place = cfg.topo.cp_placement
|
||||
_cp_val = cfg.topo.cp
|
||||
_cp_palette = [
|
||||
"#f0f7ff", # baseline blue (also the fallback for cp_placement=cube)
|
||||
"#e8f5e9", # light green
|
||||
"#fff3e0", # light orange
|
||||
"#f3e5f5", # light purple
|
||||
"#ffebee", # light red
|
||||
"#e0f7fa", # light cyan
|
||||
"#fce4ec", # light pink
|
||||
"#f9fbe7", # light lime
|
||||
]
|
||||
_cp_edge_palette = [
|
||||
"#3a86ff", "#2e7d32", "#ef6c00", "#7b1fa2",
|
||||
"#c62828", "#0097a7", "#c2185b", "#9e9d24",
|
||||
]
|
||||
|
||||
for pr in range(PE_ROWS):
|
||||
for pc in range(PE_COLS):
|
||||
pe_local = pr * PE_COLS + pc
|
||||
pe_id_in_group = cube_idx * PES_PER_CUBE + pe_local
|
||||
if pe_id_in_group >= cfg.topo.intra_cube_dims:
|
||||
continue
|
||||
px = x0 + pe_pad_x + pc * (pe_w + pe_gap)
|
||||
py = y0 + pe_pad_y + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
|
||||
|
||||
# Which (cp_rank, tp_rank) does this PE hold?
|
||||
# When cp_placement=pe: cp_rank = pe_id // tp, tp_rank = pe_id % tp.
|
||||
# When cp_placement=cube: every PE in this cube is one CP rank
|
||||
# (drawn on its own tile), tp_rank = pe_id % tp.
|
||||
if _cp_place == "pe" and _cp_val > 1:
|
||||
_cp_rank = pe_id_in_group // tp
|
||||
_tp_rank = pe_id_in_group % tp
|
||||
_fc = _cp_palette[_cp_rank % len(_cp_palette)]
|
||||
_ec = _cp_edge_palette[_cp_rank % len(_cp_edge_palette)]
|
||||
_cp_label = f" | CP={_cp_rank}"
|
||||
else:
|
||||
_cp_rank = None
|
||||
_tp_rank = pe_id_in_group % tp
|
||||
_fc = "#f0f7ff"
|
||||
_ec = "#3a86ff"
|
||||
# cp_placement=="cube" with CP>1: the whole cube is
|
||||
# one CP rank (rank 0 by convention — see cube title
|
||||
# tag above). Add the locator so users know what
|
||||
# rank this per-PE data belongs to.
|
||||
if _cp_place == "cube" and _cp_val > 1:
|
||||
_cp_label = " | CP=0"
|
||||
else:
|
||||
_cp_label = ""
|
||||
|
||||
pe_rect = patches.Rectangle(
|
||||
(px, py), pe_w, pe_h,
|
||||
facecolor=_fc, edgecolor=_ec, linewidth=1.2,
|
||||
)
|
||||
ax.add_patch(pe_rect)
|
||||
|
||||
# Head assignment uses TP rank (not raw pe_id) so that under
|
||||
# cp_placement=pe, all CP ranks share the same head split.
|
||||
q_heads = _q_heads_for_pe(_tp_rank, tp, cfg.model.h_q)
|
||||
kv_heads, kv_note = _kv_heads_for_pe(
|
||||
_tp_rank, tp, cfg.model.h_kv, cfg.topo.kv_shard_mode,
|
||||
)
|
||||
bytes_ = _per_pe_bytes(cfg, _tp_rank)
|
||||
weights_gb = sum(v for k, v in bytes_.items()
|
||||
if k not in ("KV cache", "Transient")) / 1e9
|
||||
kv_gb = bytes_["KV cache"] / 1e9
|
||||
trans_mb = bytes_["Transient"] / 1e6
|
||||
total_gb = sum(bytes_.values()) / 1e9
|
||||
|
||||
q_str = (f"{q_heads[0]}-{q_heads[-1]}"
|
||||
if len(q_heads) > 1 else str(q_heads[0])
|
||||
if q_heads else "-")
|
||||
kv_str = str(kv_heads[0]) if kv_heads else "-"
|
||||
if kv_note:
|
||||
kv_str += f"({kv_note})"
|
||||
|
||||
ffn_mb = (bytes_["W_gate"] + bytes_["W_up"]
|
||||
+ bytes_["W_down"]) / 1e6
|
||||
# Header (PE id and heads) bigger, then per-tensor breakdown
|
||||
header = (f"PE {pe_id_in_group} TP={_tp_rank}{_cp_label}\n"
|
||||
f"Q heads: {q_str}\n"
|
||||
f"KV head: {kv_str}")
|
||||
ax.text(px + 0.08, py + pe_h - 0.05, header,
|
||||
ha="left", va="top",
|
||||
fontsize=7.0, fontweight="bold",
|
||||
family="monospace", color="#0d47a1")
|
||||
|
||||
# Per-tensor bytes with explicit division divisors.
|
||||
# Divisor is always the PRODUCT of the parallel dims listed
|
||||
# (TP*CP means TP times CP, i.e. multiplication).
|
||||
tp_d = cfg.topo.tp
|
||||
ffn_div = cfg.ffn_shard_divisor
|
||||
ffn_scope = cfg.topo.ffn_shard_scope.replace("+", "*")
|
||||
kv_div = cfg.topo.cp * cfg.topo.tp
|
||||
L = layers_per_stage_val(cfg)
|
||||
detail = (
|
||||
f"weights (all {L} layers):\n"
|
||||
f" W_Q /TP={tp_d} : {bytes_['W_Q']/1e6:7.1f}MB\n"
|
||||
f" W_K /TP={tp_d} : {bytes_['W_K']/1e6:7.1f}MB\n"
|
||||
f" W_V /TP={tp_d} : {bytes_['W_V']/1e6:7.1f}MB\n"
|
||||
f" W_O /TP={tp_d} : {bytes_['W_O']/1e6:7.1f}MB\n"
|
||||
f" FFN /({ffn_scope})={ffn_div}: {ffn_mb:7.1f}MB\n"
|
||||
f"runtime:\n"
|
||||
f" KV /(CP*TP)={kv_div}: {kv_gb*1000:7.1f}MB\n"
|
||||
f" trans : {trans_mb:7.1f}MB\n"
|
||||
f"W total: {weights_gb*1000:7.1f}MB\n"
|
||||
f"TOTAL : {total_gb*1000:7.1f}MB"
|
||||
)
|
||||
ax.text(px + pe_w / 2, py + 0.06, detail,
|
||||
ha="center", va="bottom",
|
||||
fontsize=5.8, family="monospace", color="#1a1a1a")
|
||||
|
||||
total_w = n_cubes * cube_w + (n_cubes - 1) * cube_gap
|
||||
ax.set_xlim(-0.2, total_w + 0.2)
|
||||
ax.set_ylim(-0.3, cube_h + 0.6)
|
||||
ax.set_aspect("equal")
|
||||
ax.axis("off")
|
||||
|
||||
_title_cp = f", CP={cfg.topo.cp}"
|
||||
if cfg.topo.cp_placement == "cube" and cfg.topo.cp > 1:
|
||||
_title_cp += " (showing 1 of CP groups; others identical)"
|
||||
title = (
|
||||
f"Per-PE layout of one CP group "
|
||||
f"(TP={tp}{_title_cp}, H_q={cfg.model.h_q}, H_kv={cfg.model.h_kv})"
|
||||
f" | KV mode: {cfg.topo.kv_shard_mode}"
|
||||
)
|
||||
ax.set_title(title, fontsize=11, fontweight="bold")
|
||||
|
||||
return fig
|
||||
@@ -1,300 +0,0 @@
|
||||
"""Pipeline diagram: attention + FFN steps, colored by bound type.
|
||||
|
||||
Each stage is drawn as a rounded box in a horizontal flow:
|
||||
Attention: S1 -> S2 -> S3 -> ... -> S10 -> C2 (AllReduce)
|
||||
FFN: F1 -> F2 -> F3 -> F4 -> F5 -> CF1 (AllReduce)
|
||||
|
||||
CP ring communication (C1) and score AllReduce (C3) are drawn as
|
||||
"overhead" boxes attached above the QKt/PV region (they run
|
||||
concurrently with per-hop compute in the ring).
|
||||
|
||||
Box fill color = the stage's bound classification:
|
||||
compute-bound blue
|
||||
memory-bound orange
|
||||
comm-bound red
|
||||
trivial grey
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.patches as patches
|
||||
|
||||
from .model_config import FullConfig
|
||||
from .stage_latencies import (
|
||||
all_stages, all_ffn_stages, StageCost,
|
||||
)
|
||||
|
||||
|
||||
BOUND_COLOR = {
|
||||
"compute": "#4682b4", # blue
|
||||
"memory": "#e37400", # orange
|
||||
"comm": "#c62828", # red
|
||||
"trivial": "#a9a9a9", # grey
|
||||
}
|
||||
BOUND_LABEL = {
|
||||
"compute": "compute-bound",
|
||||
"memory": "memory-bound",
|
||||
"comm": "communication",
|
||||
"trivial": "trivial / negligible",
|
||||
}
|
||||
|
||||
|
||||
def _fmt_us(sec: float) -> str:
|
||||
us = sec * 1e6
|
||||
if us < 1:
|
||||
return f"{us*1e3:.1f}ns"
|
||||
if us < 100:
|
||||
return f"{us:.2f}us"
|
||||
return f"{us:.0f}us"
|
||||
|
||||
|
||||
def _draw_block(ax, x, y, w, h, label, cost: StageCost, edge="#333"):
|
||||
color = BOUND_COLOR.get(cost.bound, "#a9a9a9")
|
||||
box = patches.FancyBboxPatch(
|
||||
(x, y), w, h,
|
||||
boxstyle="round,pad=0.05",
|
||||
facecolor=color, edgecolor=edge, linewidth=1.2, alpha=0.85,
|
||||
)
|
||||
ax.add_patch(box)
|
||||
# Two lines: stage name (bold) + visible latency (small)
|
||||
ax.text(x + w / 2, y + h * 0.62, label,
|
||||
ha="center", va="center",
|
||||
fontsize=8.0, fontweight="bold", color="white")
|
||||
ax.text(x + w / 2, y + h * 0.25, _fmt_us(cost.visible_s),
|
||||
ha="center", va="center",
|
||||
fontsize=7.0, color="white")
|
||||
|
||||
|
||||
_SHORT_NAMES = {
|
||||
"S1": "RMSNorm",
|
||||
"S2": "W_Q",
|
||||
"S3": "W_K+W_V",
|
||||
"S4": "KV append",
|
||||
"S5": "Q.K^T",
|
||||
"S6": "softmax",
|
||||
"S7": "P.V",
|
||||
"S8": "merge",
|
||||
"S9": "norm O/l",
|
||||
"S10": "W_O",
|
||||
"C1": "CP ring",
|
||||
"C2": "TP AR",
|
||||
"C3": "Score AR",
|
||||
"F1": "RMSNorm",
|
||||
"F2": "W_gate",
|
||||
"F3": "W_up",
|
||||
"F4": "SwiGLU",
|
||||
"F5": "W_down",
|
||||
"CF1": "FFN AR",
|
||||
}
|
||||
|
||||
|
||||
def _short_name(name: str) -> str:
|
||||
"""Map 'S1 RMSNorm' -> 'RMSNorm'; fallback to first token."""
|
||||
tok = name.split()[0]
|
||||
return _SHORT_NAMES.get(tok, tok)
|
||||
|
||||
|
||||
def _draw_arrow(ax, x0, x1, y):
|
||||
ax.annotate("", xy=(x1, y), xytext=(x0, y),
|
||||
arrowprops=dict(arrowstyle="->", color="#333", lw=1.1))
|
||||
|
||||
|
||||
def _draw_ring_loop(ax, x_left, x_right, y_top, label: str, n_iters: int,
|
||||
arc_lift: float = 0.9):
|
||||
"""Draw a curved back-arrow above [x_left..x_right] with a repeat label.
|
||||
|
||||
Visualizes 'repeat this block n_iters times' (like a loop in a diagram).
|
||||
arc_lift: how high above y_top the arc sits.
|
||||
"""
|
||||
color = "#5a189a"
|
||||
lw = 1.6
|
||||
arc_y = y_top + arc_lift
|
||||
# Vertical tick up from the right edge
|
||||
ax.plot([x_right, x_right], [y_top + 0.02, arc_y],
|
||||
color=color, linewidth=lw)
|
||||
# Horizontal top of the loop
|
||||
ax.plot([x_right, x_left], [arc_y, arc_y],
|
||||
color=color, linewidth=lw)
|
||||
# Down-arrow into the left edge (a back-edge indicating the repeat)
|
||||
ax.annotate("", xy=(x_left, y_top + 0.02), xytext=(x_left, arc_y),
|
||||
arrowprops=dict(arrowstyle="->", color=color, lw=lw))
|
||||
# Loop label centered above the arc
|
||||
ax.text((x_left + x_right) / 2, arc_y + 0.08,
|
||||
label,
|
||||
ha="center", va="bottom",
|
||||
fontsize=9, fontweight="bold", color=color)
|
||||
# Iteration count as a badge on the right
|
||||
ax.text(x_right + 0.08, arc_y - 0.03,
|
||||
f"x {n_iters}",
|
||||
ha="left", va="top",
|
||||
fontsize=10, fontweight="bold", color=color,
|
||||
bbox=dict(boxstyle="round,pad=0.15", facecolor="white",
|
||||
edgecolor=color, linewidth=1.2, alpha=0.95))
|
||||
|
||||
|
||||
def _draw_row(ax, y_base, row_title, stages: list[StageCost],
|
||||
x_start: float, box_w: float, box_h: float, gap: float):
|
||||
"""Draw one row of stage boxes with arrows between them."""
|
||||
ax.text(x_start - 0.4, y_base + box_h / 2, row_title,
|
||||
ha="right", va="center",
|
||||
fontsize=11, fontweight="bold", color="#212529")
|
||||
x = x_start
|
||||
for i, st in enumerate(stages):
|
||||
_draw_block(ax, x, y_base, box_w, box_h,
|
||||
_short_name(st.name), st)
|
||||
if i < len(stages) - 1:
|
||||
_draw_arrow(ax, x + box_w + 0.02, x + box_w + gap - 0.02,
|
||||
y_base + box_h / 2)
|
||||
x += box_w + gap
|
||||
return x - gap
|
||||
|
||||
|
||||
def _draw_overlap_boxes(ax, hidden_stages, x_left, x_right, y_top,
|
||||
gap_between: float = 0.15):
|
||||
"""Draw dashed 'overlapped comm' boxes spanning x_left..x_right above y_top.
|
||||
|
||||
Boxes are widened to cover the horizontal range of the stages they
|
||||
overlap with (e.g. C1 CP ring sits above S5-S7).
|
||||
"""
|
||||
if not hidden_stages:
|
||||
return
|
||||
hy = y_top + 0.10
|
||||
n = len(hidden_stages)
|
||||
total_span = x_right - x_left
|
||||
box_span = (total_span - (n - 1) * gap_between) / max(1, n)
|
||||
box_h_local = 0.42
|
||||
hx = x_left
|
||||
for st in hidden_stages:
|
||||
box = patches.FancyBboxPatch(
|
||||
(hx, hy), box_span, box_h_local,
|
||||
boxstyle="round,pad=0.03",
|
||||
facecolor=BOUND_COLOR["comm"], edgecolor="#333",
|
||||
linewidth=1.0, alpha=0.35, linestyle="--",
|
||||
)
|
||||
ax.add_patch(box)
|
||||
ax.text(hx + box_span / 2, hy + box_h_local * 0.65,
|
||||
f"{_short_name(st.name)} (concurrent)",
|
||||
ha="center", va="center",
|
||||
fontsize=7, color="#7a0e0e", fontweight="bold")
|
||||
ax.text(hx + box_span / 2, hy + box_h_local * 0.2,
|
||||
_fmt_us(st.visible_s),
|
||||
ha="center", va="center",
|
||||
fontsize=6.5, color="#7a0e0e")
|
||||
hx += box_span + gap_between
|
||||
|
||||
|
||||
def draw_pipeline(cfg: FullConfig, ax=None):
|
||||
"""Draw attention + FFN pipeline. Boxes colored by bound type."""
|
||||
attn = all_stages(cfg)
|
||||
ffn = all_ffn_stages(cfg)
|
||||
|
||||
# Split attention: main stages, then C2 (TP AllReduce). C1/C3 = overlapped.
|
||||
attn_main = [s for s in attn if not s.name.startswith("C")]
|
||||
attn_c2 = next((s for s in attn if s.name.startswith("C2")), None)
|
||||
attn_hidden = [s for s in attn
|
||||
if s.name.startswith("C1") or s.name.startswith("C3")]
|
||||
attn_hidden = [s for s in attn_hidden if s.visible_s > 0]
|
||||
attn_row = attn_main + ([attn_c2] if attn_c2 and attn_c2.visible_s > 0 else [])
|
||||
|
||||
# FFN row: all F* + CF1 if non-trivial
|
||||
ffn_main = [s for s in ffn if not s.name.startswith("CF")]
|
||||
ffn_cf = next((s for s in ffn if s.name.startswith("CF1")), None)
|
||||
ffn_row = ffn_main + ([ffn_cf] if ffn_cf and ffn_cf.visible_s > 0 else [])
|
||||
|
||||
n_max = max(len(attn_row), len(ffn_row))
|
||||
box_w = 1.35
|
||||
box_h = 0.78
|
||||
gap = 0.18
|
||||
x_start = 1.1
|
||||
|
||||
total_w = x_start + n_max * box_w + (n_max - 1) * gap + 0.5
|
||||
total_h = 4.2
|
||||
|
||||
if ax is None:
|
||||
fig, ax = plt.subplots(figsize=(max(11, total_w * 1.1),
|
||||
max(4.5, total_h * 1.0)))
|
||||
else:
|
||||
fig = ax.figure
|
||||
|
||||
_draw_row(ax, y_base=2.6, row_title="Attention",
|
||||
stages=attn_row, x_start=x_start,
|
||||
box_w=box_w, box_h=box_h, gap=gap)
|
||||
|
||||
_draw_row(ax, y_base=0.6, row_title="FFN",
|
||||
stages=ffn_row, x_start=x_start,
|
||||
box_w=box_w, box_h=box_h, gap=gap)
|
||||
|
||||
# Ring loop indicator + concurrent-comm boxes anchored over S5..S8.
|
||||
_idx_by_prefix = {s.name.split()[0]: i for i, s in enumerate(attn_row)}
|
||||
_s5_idx = _idx_by_prefix.get("S5")
|
||||
_s7_idx = _idx_by_prefix.get("S7")
|
||||
_s8_idx = _idx_by_prefix.get("S8")
|
||||
|
||||
if cfg.topo.cp > 1 and cfg.topo.mode == "prefill" and _s5_idx is not None:
|
||||
_right_idx = _s8_idx if _s8_idx is not None else _s7_idx
|
||||
_lx = x_start + _s5_idx * (box_w + gap)
|
||||
_rx = x_start + _right_idx * (box_w + gap) + box_w
|
||||
_y_top = 2.6 + box_h
|
||||
# Concurrent-comm boxes (C1 CP ring, C3 score AR) directly above S5-S7
|
||||
_draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
|
||||
# Ring loop arc above those
|
||||
_draw_ring_loop(
|
||||
ax, _lx, _rx, _y_top,
|
||||
label="ring attention loop (K/V ring per hop; S8 merges partial O,m,l)",
|
||||
n_iters=cfg.topo.cp,
|
||||
arc_lift=1.05,
|
||||
)
|
||||
elif cfg.topo.cp > 1 and cfg.topo.mode == "decode" and _s5_idx is not None:
|
||||
_right_idx = _s7_idx if _s7_idx is not None else _s5_idx
|
||||
_lx = x_start + _s5_idx * (box_w + gap)
|
||||
_rx = x_start + _right_idx * (box_w + gap) + box_w
|
||||
_y_top = 2.6 + box_h
|
||||
# Any straggling concurrent comm (e.g. C3 if head-split active) above S5-S7
|
||||
_draw_overlap_boxes(ax, attn_hidden, _lx, _rx, _y_top)
|
||||
_draw_ring_loop(
|
||||
ax, _lx, _rx, _y_top,
|
||||
label="decode: local pass only; S8 fires O/m/l all-reduce",
|
||||
n_iters=1,
|
||||
arc_lift=1.05 if attn_hidden else 0.85,
|
||||
)
|
||||
elif attn_hidden:
|
||||
# Fallback: CP=1 but still have C3 or similar; anchor above S5-S7 if present
|
||||
if _s5_idx is not None and _s7_idx is not None:
|
||||
_lx = x_start + _s5_idx * (box_w + gap)
|
||||
_rx = x_start + _s7_idx * (box_w + gap) + box_w
|
||||
_draw_overlap_boxes(ax, attn_hidden, _lx, _rx, 2.6 + box_h)
|
||||
|
||||
# Totals
|
||||
attn_total = sum(s.visible_s for s in attn_row) \
|
||||
+ sum(s.visible_s for s in attn_hidden)
|
||||
ffn_total = sum(s.visible_s for s in ffn_row)
|
||||
ax.text(total_w - 0.3, 2.6 + box_h / 2 + 1.15,
|
||||
f"Attn total: {_fmt_us(attn_total)}",
|
||||
ha="right", va="bottom", fontsize=9,
|
||||
fontweight="bold", color="#212529")
|
||||
ax.text(total_w - 0.3, 0.6 + box_h + 0.15,
|
||||
f"FFN total: {_fmt_us(ffn_total)}",
|
||||
ha="right", va="bottom", fontsize=9,
|
||||
fontweight="bold", color="#212529")
|
||||
|
||||
# Legend
|
||||
handles = []
|
||||
for k, lbl in BOUND_LABEL.items():
|
||||
handles.append(patches.Patch(
|
||||
facecolor=BOUND_COLOR[k], edgecolor="#333",
|
||||
alpha=0.85, label=lbl,
|
||||
))
|
||||
ax.legend(handles=handles, loc="upper center",
|
||||
bbox_to_anchor=(0.5, -0.02),
|
||||
ncol=4, fontsize=8, frameon=False)
|
||||
|
||||
ax.set_xlim(0, total_w + 0.6)
|
||||
ax.set_ylim(0, 4.9)
|
||||
ax.set_aspect("auto")
|
||||
ax.axis("off")
|
||||
ax.set_title(
|
||||
f"Per-layer pipeline (one PE, {cfg.topo.mode}, T_q={cfg.topo.T_q}, "
|
||||
f"S_local={cfg.topo.s_local}) - color = dominant bound",
|
||||
fontsize=10, fontweight="bold",
|
||||
)
|
||||
return fig
|
||||
@@ -1,830 +0,0 @@
|
||||
"""Per-stage cost formulas for attention forward pass.
|
||||
|
||||
Each stage returns:
|
||||
name, formula (str with the substituted numeric form),
|
||||
compute_time, memory_time, comm_time, bound, visible_time.
|
||||
|
||||
Bound is the max of the three; "hidden" comm can appear as 0 when
|
||||
overlapped with compute.
|
||||
|
||||
All times are in seconds; convert to μs at display.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Literal
|
||||
|
||||
from .model_config import FullConfig
|
||||
|
||||
|
||||
Bound = Literal["compute", "memory", "comm", "trivial"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class StageCost:
|
||||
name: str
|
||||
formula: str
|
||||
compute_s: float
|
||||
memory_s: float
|
||||
comm_s: float
|
||||
bound: Bound
|
||||
visible_s: float
|
||||
hop_multiplier: int = 1 # some stages scale with N_cp hops
|
||||
# Detailed breakdown for the per-stage table.
|
||||
flops: int = 0 # total FLOPs (numeric)
|
||||
mem_bytes: int = 0 # HBM bytes moved (numeric)
|
||||
comm_bytes: int = 0 # comm bytes over the wire (numeric)
|
||||
flops_formula: str = "" # human-readable FLOPs formula
|
||||
mem_formula: str = "" # human-readable memory formula
|
||||
comm_formula: str = "" # human-readable comm formula
|
||||
|
||||
|
||||
def _visible(compute: float, memory: float, comm: float) -> tuple[float, Bound]:
|
||||
"""Pick the dominant time as the visible stage cost."""
|
||||
parts = {"compute": compute, "memory": memory, "comm": comm}
|
||||
dominant = max(parts, key=parts.get) # type: ignore
|
||||
return parts[dominant], dominant # type: ignore
|
||||
|
||||
|
||||
def stage_rmsnorm(cfg: FullConfig) -> StageCost:
|
||||
T_q = cfg.topo.T_q
|
||||
d = cfg.model.hidden
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
# Activation memory (x) scales with B; weight (once) is fixed.
|
||||
bytes_ = B * T_q * d * b + T_q * d * b
|
||||
flops = 4 * B * T_q * d
|
||||
mem_s = bytes_ / cfg.machine.bw_hbm
|
||||
cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
|
||||
vis, bnd = _visible(cmp_s, mem_s, 0)
|
||||
return StageCost(
|
||||
name="S1 RMSNorm",
|
||||
formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM",
|
||||
compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
|
||||
bound=bnd, visible_s=vis,
|
||||
flops=flops, mem_bytes=bytes_,
|
||||
flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}",
|
||||
mem_formula=f"B*T_q*d*b + T_q*d*b (weight) "
|
||||
f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B",
|
||||
)
|
||||
|
||||
|
||||
def _gemm_time(flops: int, weight_bytes: int, cfg: FullConfig) -> tuple[float, float]:
|
||||
cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
|
||||
mem_s = weight_bytes / cfg.machine.bw_hbm
|
||||
return cmp_s, mem_s
|
||||
|
||||
|
||||
def stage_wq(cfg: FullConfig) -> StageCost:
|
||||
T_q = cfg.topo.T_q
|
||||
d = cfg.model.hidden
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
dh = cfg.model.d_head
|
||||
# FLOPs scale with batch; weight bytes fixed (shared across batch).
|
||||
flops = 2 * B * T_q * d * (hq_per_pe * dh)
|
||||
weight_B = d * (hq_per_pe * dh) * b
|
||||
cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
|
||||
vis, bnd = _visible(cmp_s, mem_s, 0)
|
||||
return StageCost(
|
||||
name="S2 W_Q GEMM",
|
||||
formula=f"FLOPs = 2*B*T_q*d*(H_q/TP*d_h) "
|
||||
f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; "
|
||||
f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
|
||||
compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
|
||||
bound=bnd, visible_s=vis,
|
||||
flops=int(flops), mem_bytes=int(weight_B),
|
||||
flops_formula=f"2*B*T_q*d*(H_q/TP*d_h) "
|
||||
f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}",
|
||||
mem_formula=f"d*(H_q/TP*d_h)*b (weight, B-invariant) "
|
||||
f"= {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB",
|
||||
)
|
||||
|
||||
|
||||
def stage_wkv(cfg: FullConfig) -> StageCost:
|
||||
T_q = cfg.topo.T_q
|
||||
d = cfg.model.hidden
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
|
||||
dh = cfg.model.d_head
|
||||
flops_one = 2 * B * T_q * d * (hkv_per_pe * dh)
|
||||
weight_B_one = d * (hkv_per_pe * dh) * b
|
||||
flops = 2 * flops_one
|
||||
weight_B = 2 * weight_B_one
|
||||
cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
|
||||
vis, bnd = _visible(cmp_s, mem_s, 0)
|
||||
return StageCost(
|
||||
name="S3 W_K + W_V GEMM",
|
||||
formula=f"FLOPs per proj = 2*B*T_q*d*(H_kv/TP*d_h) "
|
||||
f"= 2*{B}*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2",
|
||||
compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
|
||||
bound=bnd, visible_s=vis,
|
||||
flops=int(flops), mem_bytes=int(weight_B),
|
||||
flops_formula=f"2*(2*B*T_q*d*(H_kv/TP*d_h)) "
|
||||
f"= 2*(2*{B}*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}",
|
||||
mem_formula=f"2*(d*(H_kv/TP*d_h)*b) (weights, B-invariant) "
|
||||
f"= 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB",
|
||||
)
|
||||
|
||||
|
||||
def stage_kv_append(cfg: FullConfig) -> StageCost:
|
||||
T_q = cfg.topo.T_q
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
|
||||
dh = cfg.model.d_head
|
||||
bytes_ = 2 * B * T_q * hkv_per_pe * dh * b
|
||||
mem_s = bytes_ / cfg.machine.bw_hbm
|
||||
return StageCost(
|
||||
name="S4 KV cache append",
|
||||
formula=f"bytes = 2*B*T_q*(H_kv/TP)*d_h*b "
|
||||
f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
|
||||
compute_s=0, memory_s=mem_s, comm_s=0,
|
||||
bound="memory", visible_s=mem_s,
|
||||
flops=0, mem_bytes=bytes_,
|
||||
flops_formula="0 (no matmul, just cache write)",
|
||||
mem_formula=f"2*B*T_q*(H_kv/TP)*d_h*b "
|
||||
f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
|
||||
)
|
||||
|
||||
|
||||
def _per_hop_qkT_pv(cfg: FullConfig) -> tuple[float, str]:
|
||||
"""Q·Kᵀ (and P·V has same FLOPs) per hop, per PE."""
|
||||
T_q = cfg.topo.T_q
|
||||
S_local = cfg.topo.s_local
|
||||
dh = cfg.model.d_head
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
B = max(1, cfg.topo.b)
|
||||
flops = 2 * B * T_q * S_local * dh * hq_per_pe
|
||||
cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
|
||||
formula = f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops:.2g} FLOPs/hop"
|
||||
return cmp_s, formula
|
||||
|
||||
|
||||
def _cp_compute_passes(cfg: FullConfig) -> int:
|
||||
"""How many local S5/S6/S7 passes happen per token step under CP.
|
||||
- decode: 1 (each PE computes once against its local K,V, then all-reduce O/m/l)
|
||||
- prefill: CP (K/V or Q rotates through the ring; one local pass per hop)
|
||||
"""
|
||||
if cfg.topo.cp <= 1:
|
||||
return 1
|
||||
return 1 if cfg.topo.mode == "decode" else cfg.topo.cp
|
||||
|
||||
|
||||
def stage_qkT(cfg: FullConfig) -> StageCost:
|
||||
cmp_hop, _ = _per_hop_qkT_pv(cfg)
|
||||
passes = _cp_compute_passes(cfg)
|
||||
total = cmp_hop * passes
|
||||
T_q = cfg.topo.T_q
|
||||
S_local = cfg.topo.s_local
|
||||
dh = cfg.model.d_head
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
B = max(1, cfg.topo.b)
|
||||
flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe
|
||||
total_flops = flops_per_hop * passes
|
||||
_hop_word = "pass" if passes == 1 else "hops"
|
||||
return StageCost(
|
||||
name=f"S5 Q.K^T (x{passes} {_hop_word})",
|
||||
formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = "
|
||||
f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = "
|
||||
f"{flops_per_hop:.2g} FLOPs/hop",
|
||||
compute_s=total, memory_s=0, comm_s=0,
|
||||
bound="compute", visible_s=total,
|
||||
hop_multiplier=passes,
|
||||
flops=int(total_flops), mem_bytes=0,
|
||||
flops_formula=(
|
||||
f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = "
|
||||
f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = "
|
||||
f"{total_flops:.3g}"
|
||||
),
|
||||
mem_formula="0 (scores accumulated on-chip)",
|
||||
)
|
||||
|
||||
|
||||
def stage_softmax(cfg: FullConfig) -> StageCost:
|
||||
T_q = cfg.topo.T_q
|
||||
S_local = cfg.topo.s_local
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
elems = B * hq_per_pe * T_q * S_local
|
||||
bytes_ = elems * b * 2
|
||||
mem_s_per_hop = bytes_ / cfg.machine.bw_hbm
|
||||
passes = _cp_compute_passes(cfg)
|
||||
total = mem_s_per_hop * passes
|
||||
total_bytes = bytes_ * passes
|
||||
_hop_word = "pass" if passes == 1 else "hops"
|
||||
return StageCost(
|
||||
name=f"S6 softmax (x{passes} {_hop_word})",
|
||||
formula=f"elts/hop = B*(H_q/TP)*T_q*S_local "
|
||||
f"= {B}*{hq_per_pe}*{T_q}*{S_local} = {elems:.2g}",
|
||||
compute_s=0, memory_s=total, comm_s=0,
|
||||
bound="memory", visible_s=total,
|
||||
hop_multiplier=passes,
|
||||
flops=0, mem_bytes=int(total_bytes),
|
||||
flops_formula="~O(elts) (negligible)",
|
||||
mem_formula=(
|
||||
f"{passes}*2*b*B*(H_q/TP)*T_q*S_local = "
|
||||
f"{passes}*2*{b}*{B}*{hq_per_pe}*{T_q}*{S_local} = "
|
||||
f"{total_bytes/1e6:.2f} MB"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def stage_pv(cfg: FullConfig) -> StageCost:
|
||||
cmp_hop, _ = _per_hop_qkT_pv(cfg)
|
||||
passes = _cp_compute_passes(cfg)
|
||||
total = cmp_hop * passes
|
||||
T_q = cfg.topo.T_q
|
||||
S_local = cfg.topo.s_local
|
||||
dh = cfg.model.d_head
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
B = max(1, cfg.topo.b)
|
||||
flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe
|
||||
total_flops = flops_per_hop * passes
|
||||
_hop_word = "pass" if passes == 1 else "hops"
|
||||
return StageCost(
|
||||
name=f"S7 P.V (x{passes} {_hop_word})",
|
||||
formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = "
|
||||
f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = "
|
||||
f"{flops_per_hop:.2g} FLOPs/hop",
|
||||
compute_s=total, memory_s=0, comm_s=0,
|
||||
bound="compute", visible_s=total,
|
||||
hop_multiplier=passes,
|
||||
flops=int(total_flops), mem_bytes=0,
|
||||
flops_formula=(
|
||||
f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = "
|
||||
f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = "
|
||||
f"{total_flops:.3g}"
|
||||
),
|
||||
mem_formula="0 (accumulated on-chip)",
|
||||
)
|
||||
|
||||
|
||||
def stage_merge(cfg: FullConfig) -> StageCost:
|
||||
"""S8: online-softmax merge of partial (o, m, l).
|
||||
- prefill (K/V ring): merge happens in-place across CP hops; no comm here.
|
||||
- decode: merge is the moment we all-reduce partial (o, m, l) across CP
|
||||
ranks; comm is folded into this stage (no separate C1).
|
||||
"""
|
||||
T_q = cfg.topo.T_q
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
dh = cfg.model.d_head
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
flops = 6 * B * T_q * hq_per_pe * dh * max(0, cfg.topo.cp - 1)
|
||||
cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
|
||||
|
||||
comm_s = 0.0
|
||||
comm_bytes = 0
|
||||
comm_formula = ""
|
||||
name_suffix = ""
|
||||
|
||||
if cfg.topo.mode == "decode" and cfg.topo.cp > 1:
|
||||
cp = cfg.topo.cp
|
||||
# (O + m + l) bytes per rank — scales with B (per-request partials).
|
||||
M = B * ((T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
|
||||
if cfg.topo.cp_placement == "pe":
|
||||
bw, alpha, tier = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
|
||||
elif cfg.topo.sips_used > 1:
|
||||
bw, alpha, tier = cfg.machine.bw_intersip, cfg.machine.alpha_intersip, "inter-SIP"
|
||||
else:
|
||||
bw, alpha, tier = cfg.machine.bw_inter, cfg.machine.alpha_inter, "inter-cube"
|
||||
ar_bytes = 2 * (cp - 1) / cp * M
|
||||
comm_s = ar_bytes / bw + 2 * (cp - 1) * alpha
|
||||
comm_bytes = int(ar_bytes)
|
||||
comm_formula = (
|
||||
f"PURPOSE (folded into S8 in decode): each CP rank computed\n"
|
||||
f"attention against its LOCAL slice of KV and holds a partial\n"
|
||||
f"(O, m, l). This all-reduce merges those partials across all\n"
|
||||
f"{cp} CP ranks (online-softmax combine). No separate C1 row\n"
|
||||
f"in decode because this is the only CP comm and it happens\n"
|
||||
f"exactly here, at the end of the local attention body.\n"
|
||||
f"---\n"
|
||||
f"AR of (O + m + l): 2*(CP-1)/CP * M "
|
||||
f"= 2*({cp}-1)/{cp} * {M} B "
|
||||
f"= {ar_bytes:.0f} B over {tier} ({bw/1e9:.0f} GB/s) "
|
||||
f"+ 2*({cp}-1)*alpha"
|
||||
)
|
||||
name_suffix = f" + O/m/l AR (CP={cp} ranks, {tier})"
|
||||
|
||||
visible_s = max(cmp_s, comm_s)
|
||||
if comm_s > cmp_s:
|
||||
bound = "comm"
|
||||
elif cmp_s > 0:
|
||||
bound = "compute"
|
||||
else:
|
||||
bound = "trivial"
|
||||
|
||||
return StageCost(
|
||||
name=f"S8 online-softmax merge (x{max(0, cfg.topo.cp-1)}){name_suffix}",
|
||||
formula=f"~6*B*T_q*(H_q/TP)*d_h*(C-1) "
|
||||
f"= 6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} "
|
||||
f"= {flops:.2g} FLOPs",
|
||||
compute_s=cmp_s, memory_s=0, comm_s=comm_s,
|
||||
bound=bound, visible_s=visible_s,
|
||||
flops=int(flops), mem_bytes=0, comm_bytes=comm_bytes,
|
||||
flops_formula=(
|
||||
f"6*B*T_q*(H_q/TP)*d_h*(CP-1) = "
|
||||
f"6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}"
|
||||
),
|
||||
mem_formula="0 (in-register)",
|
||||
comm_formula=comm_formula or "0",
|
||||
)
|
||||
|
||||
|
||||
def stage_normalize(cfg: FullConfig) -> StageCost:
|
||||
T_q = cfg.topo.T_q
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
dh = cfg.model.d_head
|
||||
B = max(1, cfg.topo.b)
|
||||
flops = B * T_q * hq_per_pe * dh
|
||||
cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
|
||||
return StageCost(
|
||||
name="S9 normalize O/l",
|
||||
formula=f"B*T_q*(H_q/TP)*d_h = {B}*{T_q}*{hq_per_pe}*{dh} "
|
||||
f"= {flops} divisions",
|
||||
compute_s=cmp_s, memory_s=0, comm_s=0,
|
||||
bound="trivial", visible_s=cmp_s,
|
||||
flops=int(flops), mem_bytes=0,
|
||||
flops_formula=f"B*T_q*(H_q/TP)*d_h "
|
||||
f"= {B}*{T_q}*{hq_per_pe}*{dh} = {flops}",
|
||||
mem_formula="0",
|
||||
)
|
||||
|
||||
|
||||
def stage_wo(cfg: FullConfig) -> StageCost:
|
||||
T_q = cfg.topo.T_q
|
||||
d = cfg.model.hidden
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
dh = cfg.model.d_head
|
||||
flops = 2 * B * T_q * (hq_per_pe * dh) * d
|
||||
weight_B = (hq_per_pe * dh) * d * b
|
||||
cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
|
||||
vis, bnd = _visible(cmp_s, mem_s, 0)
|
||||
return StageCost(
|
||||
name="S10 W_O GEMM",
|
||||
formula=f"FLOPs = 2*B*T_q*(H_q/TP*d_h)*d "
|
||||
f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; "
|
||||
f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
|
||||
compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
|
||||
bound=bnd, visible_s=vis,
|
||||
flops=int(flops), mem_bytes=int(weight_B),
|
||||
flops_formula=f"2*B*T_q*(H_q/TP*d_h)*d "
|
||||
f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}",
|
||||
mem_formula=f"(H_q/TP*d_h)*d*b (weight, B-invariant) "
|
||||
f"= {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB",
|
||||
)
|
||||
|
||||
|
||||
def comm_cp_ring(cfg: FullConfig) -> StageCost:
|
||||
"""CP comm — depends on mode:
|
||||
|
||||
- decode: single O/m/l all-reduce over CP ranks AFTER local S5-S8 finish.
|
||||
No per-hop ring; each PE computes attention against its local K,V once
|
||||
and then contributes (o, m, l) to the reduce.
|
||||
- prefill: per-hop ring during S5/S7 compute (either K/V or Q+O/m/l per
|
||||
cp_ring_variant).
|
||||
|
||||
BW tier depends on cp_placement:
|
||||
- cp_placement=pe: intra-cube BW for all hops
|
||||
- cp_placement=cube: inter-cube BW; long rings cross SIP boundaries
|
||||
"""
|
||||
if cfg.topo.cp <= 1:
|
||||
return StageCost(
|
||||
name="C1 CP comm", formula="C=1 -> no comm",
|
||||
compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
|
||||
)
|
||||
S_local = cfg.topo.s_local
|
||||
hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
T_q = cfg.topo.T_q
|
||||
dh = cfg.model.d_head
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
|
||||
# ── Decode: single O/m/l all-reduce after local compute ──
|
||||
if cfg.topo.mode == "decode":
|
||||
cp = cfg.topo.cp
|
||||
# per-rank O + m + l bytes — scales with B (per-request partials).
|
||||
M = B * ((T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
|
||||
if cfg.topo.cp_placement == "pe":
|
||||
bw = cfg.machine.bw_intra
|
||||
alpha = cfg.machine.alpha_intra
|
||||
tier = "intra-cube"
|
||||
elif cfg.topo.sips_used > 1:
|
||||
bw = cfg.machine.bw_intersip
|
||||
alpha = cfg.machine.alpha_intersip
|
||||
tier = "inter-SIP"
|
||||
else:
|
||||
bw = cfg.machine.bw_inter
|
||||
alpha = cfg.machine.alpha_inter
|
||||
tier = "inter-cube"
|
||||
ar_bytes = 2 * (cp - 1) / cp * M
|
||||
ar_time = ar_bytes / bw + 2 * (cp - 1) * alpha
|
||||
return StageCost(
|
||||
name=f"C1 O/m/l all-reduce ({tier}, x{cp} ranks)",
|
||||
formula=(f"decode: gather partial (O,m,l) once at end; "
|
||||
f"M = B*T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; "
|
||||
f"AR: 2*(CP-1)/CP*M / BW + 2*(CP-1)*alpha"),
|
||||
compute_s=0, memory_s=0, comm_s=ar_time,
|
||||
bound="comm", visible_s=ar_time,
|
||||
flops=0, mem_bytes=0, comm_bytes=int(ar_bytes),
|
||||
flops_formula="0",
|
||||
mem_formula="0",
|
||||
comm_formula=(
|
||||
f"PURPOSE: In decode, each CP rank computed attention against\n"
|
||||
f"its OWN slice of the KV cache. Each rank now holds a "
|
||||
f"partial\n"
|
||||
f"(O, m, l). This single all-reduce merges those partials "
|
||||
f"across\n"
|
||||
f"all {cp} CP ranks (using online-softmax math) to get the "
|
||||
f"final O.\n"
|
||||
f"---\n"
|
||||
f"AR of (O + m + l): 2*(CP-1)/CP * M "
|
||||
f"= 2*({cp}-1)/{cp} * {M} "
|
||||
f"= {ar_bytes:.0f} B over {tier} at {bw/1e9:.0f} GB/s "
|
||||
f"+ 2*({cp}-1)*alpha"
|
||||
),
|
||||
)
|
||||
|
||||
# ── Prefill: per-hop ring (K/V or Q+O/m/l) concurrent with S5-S7 ──
|
||||
if cfg.topo.cp_ring_variant == "qoml":
|
||||
M_KV = B * ((2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
|
||||
variant_desc = "Q+O/m/l ring"
|
||||
formula_bytes = (
|
||||
f"B*(2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b) "
|
||||
f"= {B}*(2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b}) "
|
||||
f"= {M_KV} B/hop"
|
||||
)
|
||||
else:
|
||||
M_KV = 2 * B * S_local * hkv_per_pe * dh * b
|
||||
variant_desc = "K/V ring"
|
||||
formula_bytes = (
|
||||
f"2*B*S_local*(H_kv/TP)*d_h*b "
|
||||
f"= 2*{B}*{S_local}*{hkv_per_pe}*{dh}*{b} "
|
||||
f"= {M_KV/1e6:.3f} MB/hop"
|
||||
)
|
||||
|
||||
if cfg.topo.cp_placement == "pe":
|
||||
# All hops intra-cube (fastest).
|
||||
intra_hops = cfg.topo.cp - 1
|
||||
inter_hops = 0
|
||||
intra_time = (intra_hops * M_KV / cfg.machine.bw_intra
|
||||
+ intra_hops * cfg.machine.alpha_intra)
|
||||
inter_time = 0.0
|
||||
else: # cp_placement == "cube"
|
||||
intra_hops = cfg.topo.cp_intra_sip_hops
|
||||
inter_hops = cfg.topo.cp_inter_sip_hops
|
||||
intra_time = (intra_hops * M_KV / cfg.machine.bw_inter
|
||||
+ intra_hops * cfg.machine.alpha_inter)
|
||||
inter_time = (inter_hops * M_KV / cfg.machine.bw_intersip
|
||||
+ inter_hops * cfg.machine.alpha_intersip)
|
||||
comm_time = intra_time + inter_time
|
||||
|
||||
# Overlap check: per-hop compute (S5+S6+S7) hides intra-SIP hops well;
|
||||
# inter-SIP hops usually dominate.
|
||||
per_hop_cmp, _ = _per_hop_qkT_pv(cfg)
|
||||
per_hop_mem_softmax = (cfg.h_q_per_pe * cfg.topo.T_q
|
||||
* S_local * cfg.model.bytes_per_elem * 2
|
||||
/ cfg.machine.bw_hbm)
|
||||
per_hop_compute = 2 * per_hop_cmp + per_hop_mem_softmax
|
||||
per_intra_ring = M_KV / cfg.machine.bw_inter + cfg.machine.alpha_inter
|
||||
per_inter_ring = M_KV / cfg.machine.bw_intersip + cfg.machine.alpha_intersip
|
||||
visible_intra = intra_hops * max(0.0, per_intra_ring - per_hop_compute)
|
||||
visible_inter = inter_hops * max(0.0, per_inter_ring - per_hop_compute)
|
||||
visible_total = visible_intra + visible_inter
|
||||
|
||||
tier_desc = f"{intra_hops} intra-SIP + {inter_hops} inter-SIP hops"
|
||||
total_ring_bytes = M_KV * (intra_hops + inter_hops)
|
||||
_var_purpose = (
|
||||
"K, V shards rotate between CP ranks each hop"
|
||||
if cfg.topo.cp_ring_variant == "kv"
|
||||
else "Q + running (O, m, l) rotate between CP ranks each hop"
|
||||
)
|
||||
return StageCost(
|
||||
name=f"C1 CP {variant_desc} ({tier_desc})",
|
||||
formula=f"{formula_bytes}; "
|
||||
f"intra: {intra_hops}*M/BW + inter: {inter_hops}*M/BW",
|
||||
compute_s=0, memory_s=0, comm_s=comm_time,
|
||||
bound="comm", visible_s=visible_total,
|
||||
flops=0, mem_bytes=0, comm_bytes=int(total_ring_bytes),
|
||||
flops_formula="0",
|
||||
mem_formula="0",
|
||||
comm_formula=(
|
||||
f"PURPOSE: CP shards the sequence axis. Each CP rank holds only\n"
|
||||
f"1/{cfg.topo.cp} of the KV cache, so to compute full attention\n"
|
||||
f"we must move data between CP ranks. Variant: {_var_purpose}.\n"
|
||||
f"Runs concurrently with S5/S6/S7 - each hop, compute of the\n"
|
||||
f"just-arrived shard overlaps with comm of the next shard.\n"
|
||||
f"---\n"
|
||||
f"M*(CP-1) with M = {formula_bytes}; "
|
||||
f"total = {total_ring_bytes/1e6:.3f} MB over the ring"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
|
||||
"""TP AllReduce on W_O output."""
|
||||
if cfg.topo.tp <= 1:
|
||||
return StageCost(
|
||||
name="C2 TP AllReduce W_O", formula="TP=1 -> no AR",
|
||||
compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
|
||||
)
|
||||
T_q = cfg.topo.T_q
|
||||
d = cfg.model.hidden
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
bytes_ = B * T_q * d * b
|
||||
tp = cfg.topo.tp
|
||||
tier = cfg.topo.tp_link_tier() # "intra" | "inter" | "intersip"
|
||||
if tier == "intra":
|
||||
bw, alpha, scope = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
|
||||
elif tier == "inter":
|
||||
bw, alpha, scope = cfg.machine.bw_inter, cfg.machine.alpha_inter, "cross-cube"
|
||||
else:
|
||||
bw, alpha, scope = cfg.machine.bw_intersip, cfg.machine.alpha_intersip, "cross-SIP"
|
||||
comm_time = 2 * (tp - 1) / tp * bytes_ / bw + 2 * (tp - 1) * alpha
|
||||
total_comm_bytes = int(2 * (tp - 1) / tp * bytes_)
|
||||
return StageCost(
|
||||
name=f"C2 TP AllReduce W_O (TP={tp} ranks, {scope})",
|
||||
formula=f"2*(TP-1)/TP * B*T_q*d*b B / BW + 2(TP-1)*alpha "
|
||||
f"[B={B}, {scope}]",
|
||||
compute_s=0, memory_s=0, comm_s=comm_time,
|
||||
bound="comm", visible_s=comm_time,
|
||||
flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
|
||||
flops_formula="0",
|
||||
mem_formula="0",
|
||||
comm_formula=(
|
||||
f"PURPOSE: W_O is row-parallel across TP ranks (each rank holds\n"
|
||||
f"1/{tp} of W_O's input rows). After the local W_O GEMM, each of\n"
|
||||
f"the {tp} TP ranks holds a PARTIAL hidden vector (sum over its\n"
|
||||
f"own Q-head slice only). This all-reduce sums those partials so\n"
|
||||
f"every rank ends up with the full hidden vector for the next\n"
|
||||
f"stage (RMSNorm -> FFN). Fires ONCE per layer over {scope} links.\n"
|
||||
f"---\n"
|
||||
f"2*(TP-1)/TP * B*T_q*d*b "
|
||||
f"= 2*({tp}-1)/{tp} * {B}*{T_q}*{d}*{b} "
|
||||
f"= {total_comm_bytes/1e6:.2f} MB over {scope} at "
|
||||
f"{bw/1e9:.0f} GB/s"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
|
||||
"""Extra AllReduce on attention scores when TP > H_kv with head-dim split.
|
||||
|
||||
Ranks sharing a KV head have partial scores; must AllReduce across the
|
||||
split_factor group to get final scores per hop.
|
||||
"""
|
||||
if not cfg.kv_replication_needed or cfg.topo.kv_shard_mode != "split":
|
||||
return StageCost(
|
||||
name="C3 Score AllReduce (head-split)",
|
||||
formula="TP <= H_kv or replicate mode: not needed",
|
||||
compute_s=0, memory_s=0, comm_s=0, bound="trivial", visible_s=0,
|
||||
)
|
||||
split = cfg.head_dim_split_factor # ranks sharing one head
|
||||
T_q = cfg.topo.T_q
|
||||
S_local = cfg.topo.s_local
|
||||
hq_per_pe = cfg.h_q_per_pe
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
# score bytes per rank per hop — scales with B (per-request scores).
|
||||
bytes_per_hop = B * hq_per_pe * T_q * S_local * b
|
||||
# split-group AllReduce (intra-cube assumed if group fits)
|
||||
bw = cfg.machine.bw_intra
|
||||
alpha = cfg.machine.alpha_intra
|
||||
per_hop = 2 * (split - 1) / split * bytes_per_hop / bw + 2 * (split - 1) * alpha
|
||||
total = per_hop * cfg.topo.cp
|
||||
total_comm_bytes = int(2 * (split - 1) / split * bytes_per_hop * cfg.topo.cp)
|
||||
return StageCost(
|
||||
name=f"C3 Score AllReduce ({split}-way, xCP hops)",
|
||||
formula=f"2*({split}-1)/{split} * B*(H_q/TP)*T_q*S_local*b / BW "
|
||||
f"= 2*({split}-1)/{split} * {B}*{hq_per_pe}*{T_q}*{S_local}*{b} "
|
||||
f"/ BW + latency",
|
||||
compute_s=0, memory_s=0, comm_s=total,
|
||||
bound="comm", visible_s=total,
|
||||
flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
|
||||
flops_formula="0",
|
||||
mem_formula="0",
|
||||
comm_formula=(
|
||||
f"PURPOSE: TP={cfg.topo.tp} > H_kv={cfg.model.h_kv}, so each KV\n"
|
||||
f"head is split across {split} TP ranks along the head-dim (d_h\n"
|
||||
f"/{split} per rank). Each rank's Q.K^T is therefore only a\n"
|
||||
f"PARTIAL dot product; the {split} ranks sharing one KV head\n"
|
||||
f"must AllReduce their partial scores to get the true score.\n"
|
||||
f"This fires per hop of the ring, so {cfg.topo.cp}x per layer.\n"
|
||||
f"To eliminate C3: set KV mode = 'replicate' (costs {split}x KV\n"
|
||||
f"memory but no per-hop score AR), or lower TP so TP <= H_kv.\n"
|
||||
f"---\n"
|
||||
f"CP * 2*(split-1)/split * (H_q/TP)*T_q*S_local*b "
|
||||
f"= {cfg.topo.cp} * 2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} "
|
||||
f"= {total_comm_bytes/1e6:.2f} MB total"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def all_stages(cfg: FullConfig) -> list[StageCost]:
|
||||
"""Ordered per-layer attention stages.
|
||||
|
||||
In prefill with CP > 1, C1 (the CP ring) is inserted between S7 and S8
|
||||
because it runs concurrently with S5-S7 per hop. Placing it there in the
|
||||
stage list makes the per-stage table read left-to-right in the order
|
||||
things actually happen. C3 (score AR) sits with C1 for the same reason.
|
||||
C2 (TP AllReduce on W_O) sits right after S10.
|
||||
|
||||
In decode, C1's comm is folded into S8, so no separate C1 row.
|
||||
"""
|
||||
stages = [
|
||||
stage_rmsnorm(cfg),
|
||||
stage_wq(cfg),
|
||||
stage_wkv(cfg),
|
||||
stage_kv_append(cfg),
|
||||
stage_qkT(cfg),
|
||||
stage_softmax(cfg),
|
||||
stage_pv(cfg),
|
||||
]
|
||||
# Concurrent comm (with S5-S7 in prefill) goes here, right after S7.
|
||||
if cfg.topo.mode == "prefill" and cfg.topo.cp > 1:
|
||||
stages.append(comm_cp_ring(cfg))
|
||||
_c3 = comm_kv_split_allreduce(cfg)
|
||||
if _c3.visible_s > 0 or _c3.comm_s > 0:
|
||||
stages.append(_c3)
|
||||
stages.extend([
|
||||
stage_merge(cfg),
|
||||
stage_normalize(cfg),
|
||||
stage_wo(cfg),
|
||||
])
|
||||
# C2 fires right after S10 (W_O produces per-TP-rank output; AR combines).
|
||||
stages.append(comm_tp_allreduce(cfg))
|
||||
return stages
|
||||
|
||||
|
||||
# ── FFN block stages (lightweight; per PE, per layer) ────────────
|
||||
def stage_ffn_rmsnorm(cfg: FullConfig) -> StageCost:
|
||||
T_q = cfg.topo.T_q
|
||||
d = cfg.model.hidden
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
# Activation memory scales with B; weight (once) is fixed.
|
||||
bytes_ = B * T_q * d * b + T_q * d * b
|
||||
flops = 4 * B * T_q * d
|
||||
mem_s = bytes_ / cfg.machine.bw_hbm
|
||||
return StageCost(
|
||||
name="F1 RMSNorm (pre-FFN)",
|
||||
formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM",
|
||||
compute_s=0, memory_s=mem_s, comm_s=0,
|
||||
bound="memory", visible_s=mem_s,
|
||||
flops=flops, mem_bytes=bytes_,
|
||||
flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}",
|
||||
mem_formula=f"B*T_q*d*b + T_q*d*b (weight) "
|
||||
f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B",
|
||||
)
|
||||
|
||||
|
||||
def _ffn_gemm(cfg: FullConfig, name: str, ffn_per_pe: int) -> StageCost:
|
||||
T_q = cfg.topo.T_q
|
||||
d = cfg.model.hidden
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
# FLOPs scale with batch; weight (shared) is fixed.
|
||||
flops = 2 * B * T_q * d * ffn_per_pe
|
||||
weight_B = d * ffn_per_pe * b
|
||||
cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
|
||||
vis, bnd = _visible(cmp_s, mem_s, 0)
|
||||
return StageCost(
|
||||
name=name,
|
||||
formula=f"FLOPs = 2*B*T_q*d*(ffn/div) "
|
||||
f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; "
|
||||
f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
|
||||
compute_s=cmp_s, memory_s=mem_s, comm_s=0,
|
||||
bound=bnd, visible_s=vis,
|
||||
flops=int(flops), mem_bytes=int(weight_B),
|
||||
flops_formula=f"2*B*T_q*d*(ffn/div) "
|
||||
f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}",
|
||||
mem_formula=f"d*(ffn/div)*b (weight, B-invariant) "
|
||||
f"= {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB",
|
||||
)
|
||||
|
||||
|
||||
def stage_ffn_gate(cfg: FullConfig) -> StageCost:
|
||||
ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
|
||||
return _ffn_gemm(cfg, "F2 W_gate GEMM", ffn_per_pe)
|
||||
|
||||
|
||||
def stage_ffn_up(cfg: FullConfig) -> StageCost:
|
||||
ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
|
||||
return _ffn_gemm(cfg, "F3 W_up GEMM", ffn_per_pe)
|
||||
|
||||
|
||||
def stage_ffn_swiglu(cfg: FullConfig) -> StageCost:
|
||||
"""SwiGLU element-wise activation on the intermediate FFN tensor."""
|
||||
T_q = cfg.topo.T_q
|
||||
ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
bytes_ = 3 * B * T_q * ffn_per_pe * b
|
||||
flops = 3 * B * T_q * ffn_per_pe # gate * silu(up) approximated as 3 flops/elt
|
||||
mem_s = bytes_ / cfg.machine.bw_hbm
|
||||
return StageCost(
|
||||
name="F4 SwiGLU act",
|
||||
formula=f"3*B*T_q*(ffn/div)*b "
|
||||
f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM",
|
||||
compute_s=0, memory_s=mem_s, comm_s=0,
|
||||
bound="memory", visible_s=mem_s,
|
||||
flops=flops, mem_bytes=bytes_,
|
||||
flops_formula=f"~3*B*T_q*(ffn/div) = 3*{B}*{T_q}*{ffn_per_pe} = {flops}",
|
||||
mem_formula=f"3*B*T_q*(ffn/div)*b "
|
||||
f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B",
|
||||
)
|
||||
|
||||
|
||||
def stage_ffn_down(cfg: FullConfig) -> StageCost:
|
||||
"""W_down: (ffn_per_pe, hidden). Row-parallel over FFN scope."""
|
||||
T_q = cfg.topo.T_q
|
||||
d = cfg.model.hidden
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
|
||||
flops = 2 * B * T_q * ffn_per_pe * d
|
||||
weight_B = ffn_per_pe * d * b
|
||||
cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
|
||||
vis, bnd = _visible(cmp_s, mem_s, 0)
|
||||
return StageCost(
|
||||
name="F5 W_down GEMM",
|
||||
formula=f"FLOPs = 2*B*T_q*(ffn/div)*d "
|
||||
f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; "
|
||||
f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
|
||||
compute_s=cmp_s, memory_s=mem_s, comm_s=0,
|
||||
bound=bnd, visible_s=vis,
|
||||
flops=int(flops), mem_bytes=int(weight_B),
|
||||
flops_formula=f"2*B*T_q*(ffn/div)*d "
|
||||
f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}",
|
||||
mem_formula=f"(ffn/div)*d*b (weight, B-invariant) "
|
||||
f"= {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB",
|
||||
)
|
||||
|
||||
|
||||
def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
|
||||
"""AllReduce on FFN output across the FFN sharding scope."""
|
||||
scope = cfg.topo.ffn_shard_scope
|
||||
divisor = cfg.ffn_shard_divisor # TP or TP*CP or TP*CP*DP
|
||||
if divisor <= 1:
|
||||
return StageCost(
|
||||
name="CF1 FFN AllReduce",
|
||||
formula="FFN scope = 1: not needed",
|
||||
compute_s=0, memory_s=0, comm_s=0,
|
||||
bound="trivial", visible_s=0,
|
||||
)
|
||||
T_q = cfg.topo.T_q
|
||||
d = cfg.model.hidden
|
||||
b = cfg.model.bytes_per_elem
|
||||
B = max(1, cfg.topo.b)
|
||||
# AR reduces the batched FFN output — bytes scale with B.
|
||||
bytes_ = B * T_q * d * b
|
||||
# Choose BW/alpha tier based on scope (rough): TP=intra-cube (or inter-cube),
|
||||
# +CP=inter-SIP possible, +DP=inter-SIP always.
|
||||
if "DP" in scope or "CP" in scope:
|
||||
bw = cfg.machine.bw_intersip if cfg.topo.sips_used > 1 else cfg.machine.bw_inter
|
||||
alpha = cfg.machine.alpha_intersip if cfg.topo.sips_used > 1 else cfg.machine.alpha_inter
|
||||
else: # TP only
|
||||
bw = cfg.machine.bw_intra if cfg.topo.tp_spans_cubes == 1 else cfg.machine.bw_inter
|
||||
alpha = cfg.machine.alpha_intra if cfg.topo.tp_spans_cubes == 1 else cfg.machine.alpha_inter
|
||||
comm_time = 2 * (divisor - 1) / divisor * bytes_ / bw + 2 * (divisor - 1) * alpha
|
||||
total_comm_bytes = int(2 * (divisor - 1) / divisor * bytes_)
|
||||
return StageCost(
|
||||
name=f"CF1 FFN AllReduce (scope={scope}, x{divisor})",
|
||||
formula=f"2*({divisor}-1)/{divisor} * B*T_q*d*b B / BW + latency "
|
||||
f"[B={B}]",
|
||||
compute_s=0, memory_s=0, comm_s=comm_time,
|
||||
bound="comm", visible_s=comm_time,
|
||||
flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
|
||||
flops_formula="0",
|
||||
mem_formula="0",
|
||||
comm_formula=(
|
||||
f"PURPOSE: W_down is row-parallel across the {divisor} ranks in\n"
|
||||
f"the FFN scope (scope={scope}). Each rank produced a PARTIAL\n"
|
||||
f"hidden vector after W_down; this all-reduce sums them so every\n"
|
||||
f"rank has the full FFN output for the next layer. Larger scope\n"
|
||||
f"= less FFN weight memory per PE but bigger AR bill.\n"
|
||||
f"---\n"
|
||||
f"2*(div-1)/div * B*T_q*d*b = 2*({divisor}-1)/{divisor} * "
|
||||
f"{B}*{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at "
|
||||
f"{bw/1e9:.0f} GB/s (scope={scope})"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def all_ffn_stages(cfg: FullConfig) -> list[StageCost]:
|
||||
return [
|
||||
stage_ffn_rmsnorm(cfg),
|
||||
stage_ffn_gate(cfg),
|
||||
stage_ffn_up(cfg),
|
||||
stage_ffn_swiglu(cfg),
|
||||
stage_ffn_down(cfg),
|
||||
comm_ffn_allreduce(cfg),
|
||||
]
|
||||
@@ -1,195 +0,0 @@
|
||||
"""Per-stage tensor shape rows for the analytical visualization.
|
||||
|
||||
Complements stage_latencies.py: while that module reports FLOPs/bytes/
|
||||
time per stage, this one reports the INPUT / WEIGHT / OUTPUT tensor
|
||||
shapes each stage operates on, per PE. Used for the 'per-stage shape'
|
||||
tables in the Streamlit app.
|
||||
|
||||
Shape strings substitute the current cfg's numeric values so the table
|
||||
reads like an inspection of the actual deployment (no symbolic-only
|
||||
form).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from .model_config import FullConfig
|
||||
|
||||
|
||||
def _s(shape: tuple) -> str:
|
||||
"""Render a tuple of dim values as '(a, b, c)'."""
|
||||
return "(" + ", ".join(str(x) for x in shape) + ")"
|
||||
|
||||
|
||||
def attn_stage_shape_rows(cfg: FullConfig) -> list[dict]:
|
||||
"""Per-PE attention shapes, one row per stage.
|
||||
|
||||
Rows follow the same S1..S10, C1..C3 order the latency table uses,
|
||||
with the same conditional insertions (C1 only in prefill+CP>1, C3
|
||||
only when kv_shard_mode='split' and TP > H_kv, C2 only when TP>1,
|
||||
S8 merge only when CP>1).
|
||||
"""
|
||||
m = cfg.model
|
||||
t = cfg.topo
|
||||
B = max(1, t.b)
|
||||
T_q = t.T_q
|
||||
S_local = t.s_local
|
||||
d = m.hidden
|
||||
dh = m.d_head
|
||||
hq_pe = cfg.h_q_per_pe
|
||||
hkv_pe = max(1, m.h_kv // t.tp)
|
||||
|
||||
rows: list[dict] = []
|
||||
|
||||
rows.append({
|
||||
"Stage": "S1", "Op": "RMSNorm",
|
||||
"Input (per PE)": _s((B, T_q, d)),
|
||||
"Weight (per PE)": _s((d,)),
|
||||
"Output (per PE)": _s((B, T_q, d)),
|
||||
})
|
||||
rows.append({
|
||||
"Stage": "S2", "Op": "W_Q GEMM",
|
||||
"Input (per PE)": _s((B, T_q, d)),
|
||||
"Weight (per PE)": _s((d, hq_pe * dh)),
|
||||
"Output (per PE)": _s((B, T_q, hq_pe * dh)),
|
||||
})
|
||||
rows.append({
|
||||
"Stage": "S3", "Op": "W_K + W_V GEMM",
|
||||
"Input (per PE)": _s((B, T_q, d)),
|
||||
"Weight (per PE)": f"{_s((d, hkv_pe * dh))} x2 (K, V)",
|
||||
"Output (per PE)": f"{_s((B, T_q, hkv_pe * dh))} x2 (K, V)",
|
||||
})
|
||||
rows.append({
|
||||
"Stage": "S4", "Op": "KV cache append",
|
||||
"Input (per PE)": f"{_s((B, T_q, hkv_pe * dh))} x2 (K, V)",
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": (
|
||||
f"KV cache: {_s((B, S_local, hkv_pe * dh))} x2 "
|
||||
f"(extends by T_q={T_q} tokens)"
|
||||
),
|
||||
})
|
||||
rows.append({
|
||||
"Stage": "S5", "Op": "Q · K^T",
|
||||
"Input (per PE)": (
|
||||
f"Q={_s((B, hq_pe, T_q, dh))}, "
|
||||
f"K^T={_s((B, hkv_pe, dh, S_local))}"
|
||||
),
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": _s((B, hq_pe, T_q, S_local)),
|
||||
})
|
||||
rows.append({
|
||||
"Stage": "S6", "Op": "softmax",
|
||||
"Input (per PE)": _s((B, hq_pe, T_q, S_local)),
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": _s((B, hq_pe, T_q, S_local)),
|
||||
})
|
||||
rows.append({
|
||||
"Stage": "S7", "Op": "P · V",
|
||||
"Input (per PE)": (
|
||||
f"P={_s((B, hq_pe, T_q, S_local))}, "
|
||||
f"V={_s((B, hkv_pe, S_local, dh))}"
|
||||
),
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": _s((B, hq_pe, T_q, dh)),
|
||||
})
|
||||
|
||||
if t.mode == "prefill" and t.cp > 1:
|
||||
variant = t.cp_ring_variant
|
||||
if variant == "kv":
|
||||
payload = (
|
||||
f"K,V shards rotate: {_s((B, S_local, hkv_pe, dh))} x2 per hop"
|
||||
)
|
||||
else:
|
||||
payload = (
|
||||
f"Q + running (O,m,l) rotate: {_s((B, hq_pe, T_q, dh))} per hop"
|
||||
)
|
||||
rows.append({
|
||||
"Stage": "C1", "Op": f"CP ring ({variant})",
|
||||
"Input (per PE)": payload,
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": "(circulating; final owner reduces)",
|
||||
})
|
||||
|
||||
if t.kv_shard_mode == "split" and t.tp > m.h_kv:
|
||||
rows.append({
|
||||
"Stage": "C3", "Op": "Score AllReduce (head-split)",
|
||||
"Input (per PE)": _s((B, hq_pe, T_q, S_local)),
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": _s((B, hq_pe, T_q, S_local)),
|
||||
})
|
||||
|
||||
if t.cp > 1:
|
||||
rows.append({
|
||||
"Stage": "S8", "Op": "online-softmax merge",
|
||||
"Input (per PE)": (
|
||||
f"O={_s((B, hq_pe, T_q, dh))}, "
|
||||
f"m,l={_s((B, hq_pe, T_q))} each"
|
||||
),
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": _s((B, hq_pe, T_q, dh)),
|
||||
})
|
||||
|
||||
rows.append({
|
||||
"Stage": "S9", "Op": "normalize O / l",
|
||||
"Input (per PE)": _s((B, hq_pe, T_q, dh)),
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": _s((B, hq_pe, T_q, dh)),
|
||||
})
|
||||
rows.append({
|
||||
"Stage": "S10", "Op": "W_O GEMM",
|
||||
"Input (per PE)": _s((B, T_q, hq_pe * dh)),
|
||||
"Weight (per PE)": _s((hq_pe * dh, d)),
|
||||
"Output (per PE)": _s((B, T_q, d)),
|
||||
})
|
||||
|
||||
if t.tp > 1:
|
||||
rows.append({
|
||||
"Stage": "C2", "Op": f"TP AllReduce (W_O, TP={t.tp})",
|
||||
"Input (per PE)": _s((B, T_q, d)),
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": _s((B, T_q, d)),
|
||||
})
|
||||
|
||||
return rows
|
||||
|
||||
|
||||
def ffn_stage_shape_rows(cfg: FullConfig) -> list[dict]:
|
||||
"""Per-PE FFN shapes, one row per stage."""
|
||||
m = cfg.model
|
||||
t = cfg.topo
|
||||
B = max(1, t.b)
|
||||
T_q = t.T_q
|
||||
d = m.hidden
|
||||
divisor = cfg.ffn_shard_divisor * max(1, t.ep)
|
||||
ffn_pe = max(1, m.ffn_dim // divisor)
|
||||
scope = t.ffn_shard_scope
|
||||
|
||||
rows: list[dict] = [
|
||||
{"Stage": "F1", "Op": "RMSNorm (pre-FFN)",
|
||||
"Input (per PE)": _s((B, T_q, d)),
|
||||
"Weight (per PE)": _s((d,)),
|
||||
"Output (per PE)": _s((B, T_q, d))},
|
||||
{"Stage": "F2", "Op": "W_gate GEMM",
|
||||
"Input (per PE)": _s((B, T_q, d)),
|
||||
"Weight (per PE)": _s((d, ffn_pe)),
|
||||
"Output (per PE)": _s((B, T_q, ffn_pe))},
|
||||
{"Stage": "F3", "Op": "W_up GEMM",
|
||||
"Input (per PE)": _s((B, T_q, d)),
|
||||
"Weight (per PE)": _s((d, ffn_pe)),
|
||||
"Output (per PE)": _s((B, T_q, ffn_pe))},
|
||||
{"Stage": "F4", "Op": "SwiGLU (gate * silu(up))",
|
||||
"Input (per PE)": f"gate, up = {_s((B, T_q, ffn_pe))} each",
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": _s((B, T_q, ffn_pe))},
|
||||
{"Stage": "F5", "Op": "W_down GEMM",
|
||||
"Input (per PE)": _s((B, T_q, ffn_pe)),
|
||||
"Weight (per PE)": _s((ffn_pe, d)),
|
||||
"Output (per PE)": _s((B, T_q, d))},
|
||||
]
|
||||
if cfg.ffn_shard_divisor > 1:
|
||||
rows.append({
|
||||
"Stage": "CF1",
|
||||
"Op": f"FFN AllReduce (scope={scope}, x{cfg.ffn_shard_divisor})",
|
||||
"Input (per PE)": _s((B, T_q, d)),
|
||||
"Weight (per PE)": "-",
|
||||
"Output (per PE)": _s((B, T_q, d)),
|
||||
})
|
||||
return rows
|
||||
@@ -1,360 +0,0 @@
|
||||
"""Pictorial view of how each weight/KV tensor is sharded across ranks.
|
||||
|
||||
Every tensor is drawn as a rectangle labelled with its (rows, cols)
|
||||
shape. Grid lines show shard boundaries; one shard (this PE's share)
|
||||
is highlighted so it's obvious which dim the split is along.
|
||||
|
||||
Sharding pattern:
|
||||
- W_Q (hidden, H_q*d_head) : split on **output cols** across TP
|
||||
- W_K (hidden, H_kv*d_head): split on **output cols** across TP
|
||||
- W_V (hidden, H_kv*d_head): split on **output cols** across TP
|
||||
- W_O (H_q*d_head, hidden) : split on **input rows** across TP
|
||||
- W_gate/up (hidden, ffn_dim): split on **output cols** across FFN scope
|
||||
- W_down (ffn_dim, hidden) : split on **input rows** across FFN scope
|
||||
- K/V cache (S_kv, H_kv*d_head): split on **rows (S axis)** by CP AND
|
||||
**cols (head axis)** by TP -> 2D shard grid
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.patches as patches
|
||||
|
||||
from .model_config import FullConfig
|
||||
|
||||
|
||||
_UNSHARDED_FACE = "#e9ecef"
|
||||
_UNSHARDED_EDGE = "#adb5bd"
|
||||
_MY_SHARD_FACE = "#3a86ff"
|
||||
_MY_SHARD_EDGE = "#0057b3"
|
||||
_TP_LINE_COLOR = "#d90429" # TP split lines
|
||||
_CP_LINE_COLOR = "#0f9d58" # CP split lines
|
||||
_FFN_LINE_COLOR = "#e37400" # FFN-scope split lines
|
||||
|
||||
|
||||
def _draw_tensor(ax, x, y, w, h, name, shape_str,
|
||||
row_splits: int = 1, col_splits: int = 1,
|
||||
my_row_shard: int = 0, my_col_shard: int = 0,
|
||||
row_line_color: str = _TP_LINE_COLOR,
|
||||
col_line_color: str = _TP_LINE_COLOR,
|
||||
row_label: str = "", col_label: str = "",
|
||||
note: str = "",
|
||||
cell_annot=None):
|
||||
"""Draw a tensor with sharding grid.
|
||||
|
||||
cell_annot: optional callable (row, col) -> str returning the physical
|
||||
label to render inside shard cell (r, c). If it returns "", cell is not
|
||||
annotated.
|
||||
"""
|
||||
"""Draw a tensor as a rectangle with grid lines showing shards."""
|
||||
# Base rectangle (unsharded background)
|
||||
rect = patches.Rectangle(
|
||||
(x, y), w, h,
|
||||
facecolor=_UNSHARDED_FACE, edgecolor=_UNSHARDED_EDGE, linewidth=1.0,
|
||||
)
|
||||
ax.add_patch(rect)
|
||||
|
||||
# Highlight this PE's shard
|
||||
shard_w = w / max(1, col_splits)
|
||||
shard_h = h / max(1, row_splits)
|
||||
sx = x + my_col_shard * shard_w
|
||||
sy = y + (row_splits - 1 - my_row_shard) * shard_h
|
||||
highlight = patches.Rectangle(
|
||||
(sx, sy), shard_w, shard_h,
|
||||
facecolor=_MY_SHARD_FACE, edgecolor=_MY_SHARD_EDGE, linewidth=1.5,
|
||||
alpha=0.8,
|
||||
)
|
||||
ax.add_patch(highlight)
|
||||
|
||||
# Horizontal grid lines (row splits)
|
||||
for r in range(1, row_splits):
|
||||
yy = y + r * shard_h
|
||||
ax.plot([x, x + w], [yy, yy],
|
||||
color=row_line_color, linewidth=1.2, linestyle="--")
|
||||
# Vertical grid lines (col splits)
|
||||
for c in range(1, col_splits):
|
||||
xx = x + c * shard_w
|
||||
ax.plot([xx, xx], [y, y + h],
|
||||
color=col_line_color, linewidth=1.2, linestyle="--")
|
||||
|
||||
# Per-cell physical annotations (e.g. "cube3 PE5") if provided.
|
||||
if cell_annot is not None:
|
||||
for r in range(row_splits):
|
||||
for c in range(col_splits):
|
||||
text = cell_annot(r, c)
|
||||
if not text:
|
||||
continue
|
||||
cx = x + c * shard_w + shard_w / 2
|
||||
cy = y + (row_splits - 1 - r) * shard_h + shard_h / 2
|
||||
ax.text(cx, cy, text,
|
||||
ha="center", va="center",
|
||||
fontsize=6.0, color="#0d1b2a",
|
||||
alpha=0.85)
|
||||
|
||||
# Stacked labels ABOVE rectangle: note (italic, top), shape, name (bold, closest)
|
||||
if note:
|
||||
ax.text(x + w / 2, y + h + 0.55, note,
|
||||
ha="center", va="bottom", fontsize=6.5,
|
||||
color="#555", style="italic")
|
||||
ax.text(x + w / 2, y + h + 0.30, shape_str,
|
||||
ha="center", va="bottom", fontsize=7, color="#666")
|
||||
ax.text(x + w / 2, y + h + 0.08, name,
|
||||
ha="center", va="bottom",
|
||||
fontsize=10, fontweight="bold")
|
||||
|
||||
# Split labels BELOW rectangle (stacked if both present)
|
||||
if row_label and col_label:
|
||||
ax.text(x + w / 2, y - 0.10, row_label,
|
||||
ha="center", va="top", fontsize=6.5,
|
||||
color=row_line_color, fontweight="bold")
|
||||
ax.text(x + w / 2, y - 0.35, col_label,
|
||||
ha="center", va="top", fontsize=6.5,
|
||||
color=col_line_color, fontweight="bold")
|
||||
elif row_label:
|
||||
ax.text(x + w / 2, y - 0.10, row_label,
|
||||
ha="center", va="top", fontsize=6.5,
|
||||
color=row_line_color, fontweight="bold")
|
||||
elif col_label:
|
||||
ax.text(x + w / 2, y - 0.10, col_label,
|
||||
ha="center", va="top", fontsize=6.5,
|
||||
color=col_line_color, fontweight="bold")
|
||||
|
||||
|
||||
def _physical_label(cfg: FullConfig, dim: str, rank: int) -> str:
|
||||
"""Where does rank `rank` of dim (`tp`|`cp`|`ffn`) live physically?
|
||||
Returns a short "cubeX PEy" style label based on the placement config.
|
||||
"""
|
||||
topo = cfg.topo
|
||||
if dim == "tp":
|
||||
if topo.tp_placement == "pe":
|
||||
# TP fills PEs within one cube (spilling if tp > 8)
|
||||
pe_in_cube = rank % topo.pes_per_cube_hw
|
||||
cube = rank // topo.pes_per_cube_hw
|
||||
return f"c{cube}p{pe_in_cube}"
|
||||
else: # tp on cube
|
||||
return f"cube{rank}"
|
||||
if dim == "cp":
|
||||
if topo.cp_placement == "pe":
|
||||
pe_in_cube = rank % topo.pes_per_cube_hw
|
||||
cube = rank // topo.pes_per_cube_hw
|
||||
return f"c{cube}p{pe_in_cube}"
|
||||
else: # cp on cube
|
||||
return f"cube{rank}"
|
||||
return ""
|
||||
|
||||
|
||||
def _kv_cell_annot(cfg: FullConfig):
|
||||
"""Return an annotation function for KV cache cells: (cp_r, tp_r) -> label."""
|
||||
def _ann(r, c):
|
||||
cp_lbl = _physical_label(cfg, "cp", r)
|
||||
tp_lbl = _physical_label(cfg, "tp", c)
|
||||
# Two lines if both are non-trivial; otherwise a single line
|
||||
if cp_lbl and tp_lbl:
|
||||
return f"{cp_lbl}\n{tp_lbl}"
|
||||
return cp_lbl or tp_lbl
|
||||
return _ann
|
||||
|
||||
|
||||
def _tp_cell_annot(cfg: FullConfig):
|
||||
"""Col-parallel tensors (W_Q/W_K/W_V/W_gate/W_up): each col = one TP rank."""
|
||||
def _ann(r, c):
|
||||
return _physical_label(cfg, "tp", c)
|
||||
return _ann
|
||||
|
||||
|
||||
def _tp_row_cell_annot(cfg: FullConfig):
|
||||
"""Row-parallel tensors split by TP (W_O): each row = one TP rank."""
|
||||
def _ann(r, c):
|
||||
return _physical_label(cfg, "tp", r)
|
||||
return _ann
|
||||
|
||||
|
||||
def draw_tensor_sharding(cfg: FullConfig, ax=None, my_pe: int = 0,
|
||||
show_physical: bool = False):
|
||||
"""Draw all tensors with their sharding overlays for PE `my_pe`.
|
||||
|
||||
show_physical: annotate each shard cell with the owning cube/PE.
|
||||
"""
|
||||
m = cfg.model
|
||||
tp = cfg.topo.tp
|
||||
cp = cfg.topo.cp
|
||||
|
||||
# PE's assigned shard index for each dim
|
||||
q_shard = my_pe % tp # W_Q col shard
|
||||
kv_shard_col = min(my_pe % tp, m.h_kv - 1) if m.h_kv >= tp else 0
|
||||
ffn_shard = my_pe % cfg.ffn_shard_divisor
|
||||
kv_row_shard = 0 # per-CP; showing CP=0 view for clarity
|
||||
|
||||
# Enlarged mode gives more room for physical annotations.
|
||||
_fig_w, _fig_h = (22, 13) if show_physical else (15, 9)
|
||||
_cell_w = 3.4 if show_physical else 2.6
|
||||
_cell_h = 3.0 if show_physical else 2.2
|
||||
_gap_x = 0.8 if show_physical else 0.55
|
||||
_gap_y = 2.0 if show_physical else 1.7
|
||||
|
||||
if ax is None:
|
||||
fig, ax = plt.subplots(figsize=(_fig_w, _fig_h))
|
||||
else:
|
||||
fig = ax.figure
|
||||
|
||||
# Grid layout: 2 rows x 4 cols of tensor rects
|
||||
# Attention row: W_Q, W_K, W_V, W_O
|
||||
# FFN + KV row: W_gate, W_up, W_down, KV cache
|
||||
cell_w = _cell_w
|
||||
cell_h = _cell_h
|
||||
gap_x = _gap_x
|
||||
gap_y = _gap_y
|
||||
|
||||
_kv_ann = _kv_cell_annot(cfg) if show_physical else None
|
||||
_tp_col_ann = _tp_cell_annot(cfg) if show_physical else None
|
||||
_tp_row_ann = _tp_row_cell_annot(cfg) if show_physical else None
|
||||
|
||||
positions = [
|
||||
(0, 1), # W_Q at row 0
|
||||
(1, 1), # W_K
|
||||
(2, 1), # W_V
|
||||
(3, 1), # W_O
|
||||
(0, 0), # W_gate at row 1
|
||||
(1, 0), # W_up
|
||||
(2, 0), # W_down
|
||||
(3, 0), # KV cache
|
||||
]
|
||||
|
||||
# ── attention tensors ────────────────────────
|
||||
tp_txt = f"TP={tp}"
|
||||
ffn_txt = f"{cfg.topo.ffn_shard_scope}={cfg.ffn_shard_divisor}"
|
||||
cp_txt = f"CP={cp}"
|
||||
|
||||
def px(col): return col * (cell_w + gap_x)
|
||||
def py(row): return row * (cell_h + gap_y)
|
||||
|
||||
# W_Q: (hidden, H_q*d_head), split on COL by TP
|
||||
_draw_tensor(ax, px(0), py(1), cell_w, cell_h,
|
||||
"W_Q", f"({m.hidden}, {m.h_q * m.d_head})",
|
||||
row_splits=1, col_splits=tp,
|
||||
my_col_shard=q_shard,
|
||||
col_line_color=_TP_LINE_COLOR,
|
||||
col_label=f"cols split by {tp_txt}",
|
||||
note="col-parallel (output dim)",
|
||||
cell_annot=_tp_col_ann)
|
||||
|
||||
# W_K, W_V: (hidden, H_kv*d_head)
|
||||
for (name, col_idx) in [("W_K", 1), ("W_V", 2)]:
|
||||
splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
|
||||
_draw_tensor(ax, px(col_idx), py(1), cell_w, cell_h,
|
||||
name, f"({m.hidden}, {m.h_kv * m.d_head})",
|
||||
row_splits=1, col_splits=splits,
|
||||
my_col_shard=min(q_shard, splits - 1),
|
||||
col_line_color=_TP_LINE_COLOR,
|
||||
col_label=f"cols split by {tp_txt}",
|
||||
note="col-parallel (heads)",
|
||||
cell_annot=_tp_col_ann)
|
||||
|
||||
# W_O: (H_q*d_head, hidden), split on ROW by TP
|
||||
_draw_tensor(ax, px(3), py(1), cell_w, cell_h,
|
||||
"W_O", f"({m.h_q * m.d_head}, {m.hidden})",
|
||||
row_splits=tp, col_splits=1,
|
||||
my_row_shard=q_shard,
|
||||
row_line_color=_TP_LINE_COLOR,
|
||||
row_label=f"rows split by {tp_txt}",
|
||||
note="row-parallel (input dim)",
|
||||
cell_annot=_tp_row_ann)
|
||||
|
||||
# ── FFN tensors ──────────────────────────────
|
||||
# W_gate, W_up: (hidden, ffn_dim), split on COL by FFN scope
|
||||
for (name, col_idx) in [("W_gate", 0), ("W_up", 1)]:
|
||||
_draw_tensor(ax, px(col_idx), py(0), cell_w, cell_h,
|
||||
name, f"({m.hidden}, {m.ffn_dim})",
|
||||
row_splits=1, col_splits=cfg.ffn_shard_divisor,
|
||||
my_col_shard=ffn_shard,
|
||||
col_line_color=_FFN_LINE_COLOR,
|
||||
col_label=f"cols split by {ffn_txt}",
|
||||
note="col-parallel")
|
||||
|
||||
# W_down: (ffn_dim, hidden), split on ROW by FFN scope
|
||||
_draw_tensor(ax, px(2), py(0), cell_w, cell_h,
|
||||
"W_down", f"({m.ffn_dim}, {m.hidden})",
|
||||
row_splits=cfg.ffn_shard_divisor, col_splits=1,
|
||||
my_row_shard=ffn_shard,
|
||||
row_line_color=_FFN_LINE_COLOR,
|
||||
row_label=f"rows split by {ffn_txt}",
|
||||
note="row-parallel")
|
||||
|
||||
# ── KV cache: (B, S_kv, H_kv*d_head), 2D shard on S_kv x heads ─
|
||||
# Batch (B) is a stacking dimension — each concurrent request holds
|
||||
# its own KV slice per PE, so total KV per PE = B * (per-slice size).
|
||||
kv_row_splits = cp
|
||||
kv_col_splits = min(tp, m.h_kv) if m.h_kv >= tp else tp
|
||||
_B = max(1, cfg.topo.b)
|
||||
_batch_suffix = f", B={_B}" if _B > 1 else ""
|
||||
_batch_note = (f"; batch B={_B} => {_B}x KV bytes per PE"
|
||||
if _B > 1 else "")
|
||||
# Draw B-1 offset "shadow" rectangles behind the KV cache to give
|
||||
# a visual hint of the batch stacking dimension. Cap at 5 shadows
|
||||
# so a large B doesn't overwhelm the diagram.
|
||||
if _B > 1:
|
||||
_kv_shadow_x = px(3)
|
||||
_kv_shadow_y = py(0)
|
||||
_n_shadows = min(_B - 1, 5)
|
||||
_offset_step = 0.15
|
||||
for _si in range(_n_shadows, 0, -1):
|
||||
_dx = _si * _offset_step
|
||||
_dy = _si * _offset_step
|
||||
_shadow = patches.Rectangle(
|
||||
(_kv_shadow_x + _dx, _kv_shadow_y + _dy),
|
||||
cell_w, cell_h,
|
||||
facecolor="#f5f5f5",
|
||||
edgecolor="#adb5bd",
|
||||
linewidth=0.8, linestyle="--", alpha=0.55,
|
||||
)
|
||||
ax.add_patch(_shadow)
|
||||
_draw_tensor(ax, px(3), py(0), cell_w, cell_h,
|
||||
"KV cache (K and V)",
|
||||
f"({m.h_kv * m.d_head}, S_kv={cfg.topo.s_kv:,}"
|
||||
f"{_batch_suffix})",
|
||||
row_splits=kv_row_splits,
|
||||
col_splits=kv_col_splits if show_physical else 1,
|
||||
my_row_shard=kv_row_shard,
|
||||
row_line_color=_CP_LINE_COLOR,
|
||||
row_label=f"rows: S split by {cp_txt}",
|
||||
col_label=f"cols: heads split by {tp_txt}",
|
||||
col_line_color=_TP_LINE_COLOR,
|
||||
note=f"2D shard: seq axis (CP) x head axis (TP)"
|
||||
f"{_batch_note}",
|
||||
cell_annot=_kv_ann)
|
||||
# In non-physical mode, KV was drawn with col_splits=1; overlay TP col
|
||||
# splits as dotted lines to hint at the 2D nature. In physical mode
|
||||
# the _draw_tensor call already drew proper vertical dashes.
|
||||
if not show_physical:
|
||||
kv_x = px(3)
|
||||
kv_y = py(0)
|
||||
shard_col_w = cell_w / max(1, kv_col_splits)
|
||||
for c in range(1, kv_col_splits):
|
||||
xx = kv_x + c * shard_col_w
|
||||
ax.plot([xx, xx], [kv_y, kv_y + cell_h],
|
||||
color=_TP_LINE_COLOR, linewidth=1.2, linestyle=":")
|
||||
|
||||
ax.set_xlim(-0.3, 4 * (cell_w + gap_x) - gap_x + 0.3)
|
||||
ax.set_ylim(-0.9, 2 * (cell_h + gap_y) - gap_y + 0.9)
|
||||
ax.set_aspect("auto")
|
||||
ax.axis("off")
|
||||
|
||||
_B_title = max(1, cfg.topo.b)
|
||||
_b_title_suffix = f", B={_B_title}" if _B_title > 1 else ""
|
||||
ax.set_title(
|
||||
f"Tensor sharding view (PE {my_pe}) | TP={tp}, CP={cp}"
|
||||
f"{_b_title_suffix}, "
|
||||
f"FFN scope={cfg.topo.ffn_shard_scope} | "
|
||||
f"blue = this PE's shard",
|
||||
fontsize=11, fontweight="bold",
|
||||
)
|
||||
# Legend for line colors
|
||||
ax.plot([], [], color=_TP_LINE_COLOR, linewidth=1.2, linestyle="--",
|
||||
label=f"TP split ({tp} ranks)")
|
||||
ax.plot([], [], color=_CP_LINE_COLOR, linewidth=1.2, linestyle="--",
|
||||
label=f"CP split ({cp} ranks, sequence dim)")
|
||||
ax.plot([], [], color=_FFN_LINE_COLOR, linewidth=1.2, linestyle="--",
|
||||
label=f"FFN scope split ({cfg.ffn_shard_divisor} ranks)")
|
||||
ax.legend(loc="lower center", bbox_to_anchor=(0.5, -0.05),
|
||||
ncol=3, fontsize=8, frameon=False)
|
||||
|
||||
return fig
|
||||
@@ -1,261 +0,0 @@
|
||||
"""Interface + smoke tests for auto_explore.
|
||||
|
||||
Covers:
|
||||
- enumerate_configs yields valid TopologyConfigs with expected pruning
|
||||
- score_config returns a ConfigScore with all fields populated
|
||||
- pareto_frontier is a subset of feasible and is non-empty for a
|
||||
reasonable model+workload
|
||||
- run_auto_explore returns a coherent AutoExploreResult (all_scores
|
||||
sorted by latency, pareto ⊆ feasible ⊆ all_scores)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.analytical_visualization.auto_explore import (
|
||||
ConfigScore,
|
||||
enumerate_configs,
|
||||
pareto_frontier,
|
||||
run_auto_explore,
|
||||
score_config,
|
||||
)
|
||||
from tests.analytical_visualization.autosuggest import auto_suggest
|
||||
from tests.analytical_visualization.model_config import (
|
||||
FullConfig, MachineParams, TopologyConfig,
|
||||
)
|
||||
from tests.analytical_visualization.model_presets import PRESETS
|
||||
|
||||
|
||||
# ── Enumeration ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_enumerate_yields_valid_configs():
|
||||
"""Enumerator produces TopologyConfigs with all 9 knobs set."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
it = enumerate_configs(model, s_kv=8192, mode="decode")
|
||||
first = next(it)
|
||||
for f in ("cp", "tp", "pp", "dp", "kv_shard_mode",
|
||||
"ffn_shard_scope", "tp_placement", "cp_placement",
|
||||
"cp_ring_variant"):
|
||||
assert getattr(first, f) is not None, f"knob {f} not set"
|
||||
assert first.s_kv == 8192
|
||||
assert first.mode == "decode"
|
||||
|
||||
|
||||
def test_enumerate_prunes_pp_beyond_layers():
|
||||
"""PP > model.layers is skipped by the enumerator."""
|
||||
model = PRESETS["Llama 3.1 70B"].model # layers=80
|
||||
for cfg in enumerate_configs(model, s_kv=8192, mode="decode"):
|
||||
assert cfg.pp <= model.layers
|
||||
|
||||
|
||||
def test_enumerate_prunes_ffn_dp_when_dp_is_1():
|
||||
"""ffn_shard_scope containing 'DP' is skipped when dp=1 (redundant)."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
for cfg in enumerate_configs(model, s_kv=8192, mode="decode"):
|
||||
if cfg.dp == 1:
|
||||
assert "DP" not in cfg.ffn_shard_scope
|
||||
|
||||
|
||||
# ── Scoring ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_score_returns_populated_config_score():
|
||||
"""score_config returns a fully-populated ConfigScore."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
topo = next(enumerate_configs(model, s_kv=8192, mode="decode"))
|
||||
machine = MachineParams()
|
||||
cfg = FullConfig(model=model, topo=topo, machine=machine)
|
||||
score = score_config(cfg)
|
||||
assert isinstance(score, ConfigScore)
|
||||
assert score.total_latency_ns > 0
|
||||
assert score.pes_used == topo.total_pes
|
||||
assert 0.0 <= score.efficiency_score <= 1.0
|
||||
assert score.hbm_utilization >= 0.0
|
||||
|
||||
|
||||
# ── Pareto ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_pareto_subset_of_feasible():
|
||||
"""Pareto scores are always feasible (memory + placement)."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
|
||||
for p in res.pareto_scores:
|
||||
assert p.fits_memory
|
||||
assert p.placement_valid
|
||||
|
||||
|
||||
def test_pareto_non_empty_when_feasible_configs_exist():
|
||||
"""When at least one config fits, Pareto must have ≥ 1 entry."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
|
||||
assert res.total_feasible > 0
|
||||
assert len(res.pareto_scores) > 0
|
||||
|
||||
|
||||
def test_pareto_frontier_non_dominated():
|
||||
"""No Pareto entry is dominated by another."""
|
||||
from tests.analytical_visualization.auto_explore import _dominates
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
|
||||
for i, a in enumerate(res.pareto_scores):
|
||||
for j, b in enumerate(res.pareto_scores):
|
||||
if i == j:
|
||||
continue
|
||||
assert not _dominates(b, a), (
|
||||
f"Pareto entry {i} is dominated by entry {j}"
|
||||
)
|
||||
|
||||
|
||||
# ── End-to-end sanity ───────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_run_auto_explore_shape():
|
||||
"""all_scores sorted asc by latency; pareto ⊆ feasible ⊆ all_scores."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
|
||||
|
||||
assert res.total_enumerated == len(res.all_scores)
|
||||
assert res.total_feasible == sum(
|
||||
1 for s in res.all_scores if s.fits_memory and s.placement_valid
|
||||
)
|
||||
assert len(res.pareto_scores) <= res.total_feasible
|
||||
|
||||
latencies = [s.total_latency_ns for s in res.all_scores]
|
||||
assert latencies == sorted(latencies)
|
||||
|
||||
|
||||
def test_pp_does_not_reduce_single_request_latency():
|
||||
"""A single request traverses all layers regardless of PP.
|
||||
Latency-optimal Pareto configs should NOT prefer PP>1 for decode."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
|
||||
# Sort Pareto by latency; the fastest should not need PP>1 to win.
|
||||
fastest = res.pareto_scores[0]
|
||||
assert fastest.pp == 1, (
|
||||
f"expected PP=1 for the fastest decode config; got PP={fastest.pp}"
|
||||
)
|
||||
|
||||
|
||||
# ── Parallelism sensitivity ─────────────────────────────────────────
|
||||
|
||||
|
||||
def test_parallelism_sensitivity_has_all_five_knobs():
|
||||
"""compute_parallelism_sensitivity returns one row per knob."""
|
||||
from tests.analytical_visualization.auto_explore import (
|
||||
compute_parallelism_sensitivity,
|
||||
)
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
|
||||
baseline = res.pareto_scores[0]
|
||||
rows = compute_parallelism_sensitivity(
|
||||
baseline, model, machine, s_kv=8192, mode="decode",
|
||||
)
|
||||
knobs = {r.knob for r in rows}
|
||||
assert knobs == {"cp", "tp", "pp", "dp", "ep"}
|
||||
|
||||
|
||||
def test_attention_only_latency_is_lower_than_full():
|
||||
"""include_ffn=False must produce lower latency than include_ffn=True
|
||||
for the same model+workload (FFN cost is dropped)."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
r_full = run_auto_explore(model, machine, s_kv=8192, mode="decode",
|
||||
include_ffn=True)
|
||||
r_attn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
|
||||
include_ffn=False)
|
||||
best_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns)
|
||||
best_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns)
|
||||
assert best_attn.total_latency_ns < best_full.total_latency_ns, (
|
||||
f"attn-only ({best_attn.latency_us:.2f} us) should be less than "
|
||||
f"full ({best_full.latency_us:.2f} us)"
|
||||
)
|
||||
|
||||
|
||||
def test_attention_only_pareto_non_empty():
|
||||
"""The attention-only sweep must still produce a non-empty Pareto set."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
res = run_auto_explore(model, machine, s_kv=8192, mode="decode",
|
||||
include_attention=True, include_ffn=False)
|
||||
assert res.total_feasible > 0
|
||||
assert len(res.pareto_scores) > 0
|
||||
|
||||
|
||||
def test_ffn_only_latency_less_than_full_and_different_from_attn():
|
||||
"""FFN-only latency must be > 0, < full-model latency, and != attn-only."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
r_full = run_auto_explore(model, machine, s_kv=8192, mode="decode",
|
||||
include_attention=True, include_ffn=True)
|
||||
r_attn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
|
||||
include_attention=True, include_ffn=False)
|
||||
r_ffn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
|
||||
include_attention=False, include_ffn=True)
|
||||
b_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns)
|
||||
b_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns)
|
||||
b_ffn = min(r_ffn.pareto_scores, key=lambda s: s.total_latency_ns)
|
||||
assert b_ffn.total_latency_ns > 0
|
||||
assert b_ffn.total_latency_ns < b_full.total_latency_ns
|
||||
assert abs(b_ffn.total_latency_ns - b_attn.total_latency_ns) > 1.0, (
|
||||
"FFN-only and attention-only latencies should not be identical"
|
||||
)
|
||||
|
||||
|
||||
def test_auto_suggest_cubes_match_default_topology():
|
||||
"""A TopologyConfig built from auto_suggest's returned (cp, tp, pp,
|
||||
cp_placement) with all other fields left at their dataclass defaults
|
||||
must produce the same cubes_used the Suggestion reports. Guards the
|
||||
sidebar apply: any session_state key auto_suggest doesn't explore
|
||||
(tp_placement, kv_shard_mode, cp_ring_variant, ep, ffn_shard_scope)
|
||||
must be reset to the default before rerun, else the applied topology
|
||||
won't match what the memory-min search actually scored.
|
||||
"""
|
||||
machine = MachineParams()
|
||||
matrix = [
|
||||
("Llama 3 70B", 8192),
|
||||
("Llama 3 70B", 32768),
|
||||
("Llama 3 8B", 131072),
|
||||
("Qwen 3 32B", 32768),
|
||||
]
|
||||
for name, skv in matrix:
|
||||
m = PRESETS[name].model
|
||||
for ia, ff in [(True, False), (False, True), (True, True)]:
|
||||
sug = auto_suggest(m, machine, s_kv=skv, mode="decode",
|
||||
include_attention=ia, include_ffn=ff)
|
||||
topo = TopologyConfig(
|
||||
cp=sug.cp, tp=sug.tp, pp=sug.pp,
|
||||
s_kv=skv, mode="decode", b=1,
|
||||
cp_placement=sug.cp_placement,
|
||||
)
|
||||
assert topo.cubes_used == sug.cubes_used, (
|
||||
f"{name} skv={skv} scope=(ia={ia},ff={ff}): "
|
||||
f"topo.cubes_used={topo.cubes_used} != "
|
||||
f"sug.cubes_used={sug.cubes_used}"
|
||||
)
|
||||
|
||||
|
||||
def test_parallelism_sensitivity_includes_baseline_value():
|
||||
"""Each row's values contain the baseline_value."""
|
||||
from tests.analytical_visualization.auto_explore import (
|
||||
compute_parallelism_sensitivity,
|
||||
)
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
machine = MachineParams()
|
||||
res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
|
||||
baseline = res.pareto_scores[0]
|
||||
rows = compute_parallelism_sensitivity(
|
||||
baseline, model, machine, s_kv=8192, mode="decode",
|
||||
)
|
||||
for r in rows:
|
||||
# Baseline value must be sweepable OR mapped to something in values.
|
||||
assert r.baseline_value in r.values or r.baseline_value == 1, (
|
||||
f"knob {r.knob}: baseline_value={r.baseline_value} not in {r.values}"
|
||||
)
|
||||
@@ -1,132 +0,0 @@
|
||||
"""Interface + invariant tests for auto_hardware."""
|
||||
from __future__ import annotations
|
||||
|
||||
from tests.analytical_visualization.auto_hardware import (
|
||||
HardwareCandidate,
|
||||
JointScore,
|
||||
_HW_KNOB_DEFAULTS,
|
||||
enumerate_hardware,
|
||||
joint_explore,
|
||||
)
|
||||
from tests.analytical_visualization.model_presets import PRESETS
|
||||
|
||||
|
||||
# ── Enumeration ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_enumerate_two_stage_yields_one():
|
||||
"""Two-stage depth yields exactly 1 HW candidate (defaults)."""
|
||||
hws = list(enumerate_hardware("two_stage"))
|
||||
assert len(hws) == 1
|
||||
for knob, default_val in _HW_KNOB_DEFAULTS.items():
|
||||
assert getattr(hws[0], knob) == default_val
|
||||
|
||||
|
||||
def test_enumerate_balanced_yields_64():
|
||||
"""Balanced = 2 values × 6 knobs = 2^6 = 64 candidates."""
|
||||
hws = list(enumerate_hardware("balanced"))
|
||||
assert len(hws) == 64
|
||||
|
||||
|
||||
def test_enumerate_coarse_yields_729():
|
||||
"""Coarse = 3 values × 6 knobs = 3^6 = 729 candidates."""
|
||||
hws = list(enumerate_hardware("coarse"))
|
||||
assert len(hws) == 729
|
||||
|
||||
|
||||
def test_default_cost_score_is_6():
|
||||
"""At every knob = its default, cost_score = 6 (one per knob)."""
|
||||
hw = HardwareCandidate(
|
||||
pe_hbm_gb=_HW_KNOB_DEFAULTS["pe_hbm_gb"],
|
||||
bw_hbm_gbs=_HW_KNOB_DEFAULTS["bw_hbm_gbs"],
|
||||
peak_tflops_f16=_HW_KNOB_DEFAULTS["peak_tflops_f16"],
|
||||
bw_intra_gbs=_HW_KNOB_DEFAULTS["bw_intra_gbs"],
|
||||
bw_inter_gbs=_HW_KNOB_DEFAULTS["bw_inter_gbs"],
|
||||
bw_intersip_gbs=_HW_KNOB_DEFAULTS["bw_intersip_gbs"],
|
||||
)
|
||||
assert hw.cost_score == 6.0
|
||||
|
||||
|
||||
# ── Joint explore + Pareto ──────────────────────────────────────────
|
||||
|
||||
|
||||
def test_joint_explore_returns_pareto_non_empty():
|
||||
"""Given a reasonable model+workload, at least one HW+parallelism fits."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
res = joint_explore(model, s_kv=131072, mode="decode", depth="two_stage")
|
||||
assert res.total_joint >= 1
|
||||
assert len(res.pareto_scores) >= 1
|
||||
|
||||
|
||||
def test_pareto_subset_of_all_scores():
|
||||
"""Every Pareto entry is present in all_scores."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
res = joint_explore(model, s_kv=131072, mode="decode", depth="two_stage")
|
||||
pareto_ids = {id(s) for s in res.pareto_scores}
|
||||
all_ids = {id(s) for s in res.all_scores}
|
||||
assert pareto_ids.issubset(all_ids)
|
||||
|
||||
|
||||
def test_pareto_non_dominated():
|
||||
"""No Pareto entry is dominated by another Pareto entry on (lat, cost)."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
res = joint_explore(model, s_kv=131072, mode="decode", depth="balanced")
|
||||
for i, a in enumerate(res.pareto_scores):
|
||||
for j, b in enumerate(res.pareto_scores):
|
||||
if i == j:
|
||||
continue
|
||||
no_worse = (
|
||||
b.total_latency_ns <= a.total_latency_ns
|
||||
and b.cost_score <= a.cost_score
|
||||
)
|
||||
strictly_better = (
|
||||
b.total_latency_ns < a.total_latency_ns
|
||||
or b.cost_score < a.cost_score
|
||||
)
|
||||
assert not (no_worse and strictly_better), (
|
||||
f"Pareto entry {i} is dominated by entry {j}"
|
||||
)
|
||||
|
||||
|
||||
# ── Sensitivity ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_sensitivity_all_knobs_monotone_non_worsening():
|
||||
"""Doubling any HW knob should not slow the sim down (rel_speedup ≥ 0
|
||||
within floating-point tolerance)."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
res = joint_explore(model, s_kv=131072, mode="decode", depth="two_stage")
|
||||
assert len(res.sensitivity) == 6
|
||||
for row in res.sensitivity:
|
||||
# Allow a tiny slop for floating-point but disallow real regressions.
|
||||
assert row.rel_speedup >= -1e-9, (
|
||||
f"knob {row.knob}: doubling slowed latency from "
|
||||
f"{row.baseline_latency_ns} to {row.doubled_latency_ns}"
|
||||
)
|
||||
|
||||
|
||||
def test_joint_explore_attention_only_faster_than_full():
|
||||
"""include_ffn=False produces smaller best-latency than include_ffn=True
|
||||
for the same HW+model."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
r_full = joint_explore(model, s_kv=131072, mode="decode",
|
||||
depth="two_stage", include_ffn=True)
|
||||
r_attn = joint_explore(model, s_kv=131072, mode="decode",
|
||||
depth="two_stage", include_ffn=False)
|
||||
b_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns)
|
||||
b_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns)
|
||||
assert b_attn.total_latency_ns < b_full.total_latency_ns, (
|
||||
f"attn-only ({b_attn.latency_ms:.2f}ms) should be less than "
|
||||
f"full ({b_full.latency_ms:.2f}ms)"
|
||||
)
|
||||
|
||||
|
||||
def test_sensitivity_hbm_bw_dominant_for_llama_decode():
|
||||
"""For Llama 70B decode (memory-bound), HBM BW should be the top
|
||||
sensitivity knob — doubling it gives more speedup than any other knob."""
|
||||
model = PRESETS["Llama 3.1 70B"].model
|
||||
res = joint_explore(model, s_kv=131072, mode="decode", depth="balanced")
|
||||
assert res.sensitivity[0].knob == "bw_hbm_gbs", (
|
||||
f"expected bw_hbm_gbs to top the sensitivity ranking for Llama 70B "
|
||||
f"decode; got {res.sensitivity[0].knob}"
|
||||
)
|
||||
@@ -1,519 +0,0 @@
|
||||
"""Tests for chip_roofline: AI, B*, L*, per-token latency curves."""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.analytical_visualization.chip_roofline import (
|
||||
ai_sensitivity_curve,
|
||||
arithmetic_intensity,
|
||||
balance_context,
|
||||
bound_regime,
|
||||
critical_batch,
|
||||
good_batch,
|
||||
good_context,
|
||||
knee_batch,
|
||||
kv_bytes_per_token,
|
||||
max_batch_within_slo,
|
||||
memory_budget_curve_vs_batch,
|
||||
memory_budget_curve_vs_skv,
|
||||
per_token_latency_curve,
|
||||
size_deployment,
|
||||
step_latency_curve,
|
||||
t_com,
|
||||
t_mem_long,
|
||||
t_mem_short,
|
||||
total_active_params,
|
||||
utilization_at,
|
||||
)
|
||||
from tests.analytical_visualization.model_config import (
|
||||
FullConfig, MachineParams, ModelConfig, TopologyConfig,
|
||||
)
|
||||
from tests.analytical_visualization.model_presets import PRESETS
|
||||
|
||||
|
||||
# ── Arithmetic intensity ────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_arithmetic_intensity_matches_ratio():
|
||||
"""AI = peak_flops / bw_hbm — pure ratio, no util factor."""
|
||||
m = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
|
||||
assert arithmetic_intensity(m) == pytest.approx(8e12 / 256e9)
|
||||
|
||||
|
||||
def test_ai_scales_with_flops_and_bandwidth():
|
||||
m1 = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
|
||||
m2 = MachineParams(peak_tflops_f16=16.0, bw_hbm_gbs=256.0)
|
||||
m3 = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=128.0)
|
||||
assert arithmetic_intensity(m2) == pytest.approx(2 * arithmetic_intensity(m1))
|
||||
assert arithmetic_intensity(m3) == pytest.approx(2 * arithmetic_intensity(m1))
|
||||
|
||||
|
||||
# ── Critical batch B* ───────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_critical_batch_h100_reference():
|
||||
"""H100-class: 989 TFLOPs BF16, 3.35 TB/s HBM3 → B* ≈ 295."""
|
||||
m = MachineParams(peak_tflops_f16=989.0, bw_hbm_gbs=3350.0)
|
||||
model = PRESETS["Llama 3 8B"].model # bf16 (b=2)
|
||||
b_star = critical_batch(m, model)
|
||||
assert b_star == pytest.approx(295, rel=0.05), b_star
|
||||
|
||||
|
||||
def test_critical_batch_scales_with_sparsity():
|
||||
"""MoE 8× sparsity gives 8× B*."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
b_dense = critical_batch(m, model, sparsity=1)
|
||||
b_moe8 = critical_batch(m, model, sparsity=8)
|
||||
assert b_moe8 == pytest.approx(8 * b_dense)
|
||||
|
||||
|
||||
def test_critical_batch_bf16_formula():
|
||||
"""For BF16 (b=2), B* = AI in flops-per-byte units → numerically."""
|
||||
m = MachineParams(peak_tflops_f16=8.0, bw_hbm_gbs=256.0)
|
||||
model = ModelConfig(bytes_per_elem=2)
|
||||
ai = arithmetic_intensity(m)
|
||||
assert critical_batch(m, model) == pytest.approx(ai), (
|
||||
critical_batch(m, model), ai,
|
||||
)
|
||||
|
||||
|
||||
# ── Balance context L* ──────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_balance_context_positive_finite():
|
||||
"""L* for a real model on real machine is a positive finite number."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
l_star = balance_context(m, model)
|
||||
assert 0 < l_star < 1e9, l_star
|
||||
|
||||
|
||||
def test_balance_context_scales_with_bandwidth():
|
||||
"""L* = 2N*W/(C*kv_bpt): doubling HBM BW doubles L*.
|
||||
Slower memory hits the KV wall at a shorter context. Machine-only
|
||||
change so N_active and kv_bpt stay fixed."""
|
||||
fast = MachineParams(bw_hbm_gbs=512.0)
|
||||
slow = MachineParams(bw_hbm_gbs=256.0)
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
assert (balance_context(fast, model)
|
||||
== pytest.approx(2 * balance_context(slow, model)))
|
||||
|
||||
|
||||
# ── Knee ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_knee_equals_b_star_at_short_context():
|
||||
"""S_kv → 0: B_knee → B*."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
b_star = critical_batch(m, model)
|
||||
assert knee_batch(m, model, s_kv=1) == pytest.approx(b_star, rel=1e-3)
|
||||
|
||||
|
||||
def test_knee_diverges_at_balance_context():
|
||||
"""S_kv = L*: knee is infinite; S_kv > L*: no knee (None)."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
l_star = balance_context(m, model)
|
||||
assert knee_batch(m, model, s_kv=int(2 * l_star)) is None
|
||||
# At 0.9 * L*, knee ~= 10 * B*; assert at least 5x to survive rounding.
|
||||
below = knee_batch(m, model, s_kv=int(0.9 * l_star))
|
||||
assert below is not None and below > 5 * critical_batch(m, model)
|
||||
|
||||
|
||||
def test_knee_slides_right_with_context():
|
||||
"""S_kv = L*/2 → B_knee = 2 * B*."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
b_star = critical_batch(m, model)
|
||||
l_star = balance_context(m, model)
|
||||
got = knee_batch(m, model, s_kv=int(l_star / 2))
|
||||
assert got == pytest.approx(2 * b_star, rel=0.01), (got, 2 * b_star)
|
||||
|
||||
|
||||
# ── Per-token latency curve ─────────────────────────────────────────
|
||||
|
||||
|
||||
def test_latency_curve_monotonically_decreasing():
|
||||
"""total_s must strictly decrease as batch increases (weight/B shrinks)."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
pts = per_token_latency_curve(m, model, [1, 2, 4, 8, 16, 32, 64],
|
||||
s_kv=1024)
|
||||
totals = [p.total_s for p in pts]
|
||||
for a, b in zip(totals, totals[1:]):
|
||||
assert b < a, (a, b)
|
||||
|
||||
|
||||
def test_latency_curve_asymptotes_to_compute_plus_kv():
|
||||
"""At very large B, total → compute + KV (weight/B → 0)."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
pts = per_token_latency_curve(m, model, [10_000_000], s_kv=1024)
|
||||
p = pts[0]
|
||||
assert p.total_s == pytest.approx(p.compute_s + p.kv_s, rel=1e-3)
|
||||
|
||||
|
||||
def test_latency_at_b_star_is_roughly_2x_floor():
|
||||
"""At B*, weight = compute (dense, S_kv small), so total ≈ 2 * compute."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
b_star = int(round(critical_batch(m, model)))
|
||||
pts = per_token_latency_curve(m, model, [b_star], s_kv=1)
|
||||
p = pts[0]
|
||||
# weight_s ≈ compute_s at B*
|
||||
assert p.weight_s == pytest.approx(p.compute_s, rel=0.05), (
|
||||
p.weight_s, p.compute_s,
|
||||
)
|
||||
|
||||
|
||||
# ── Regime classification ───────────────────────────────────────────
|
||||
|
||||
|
||||
def test_regime_memory_bound_at_b1():
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
assert bound_regime(m, model, batch=1, s_kv=1024) == "memory-bound"
|
||||
|
||||
|
||||
def test_regime_kv_bound_past_balance_context():
|
||||
"""S_kv > L* with reasonable batch → KV dominates."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
l_star = balance_context(m, model)
|
||||
b_star = int(round(critical_batch(m, model)))
|
||||
assert bound_regime(m, model, batch=b_star * 4,
|
||||
s_kv=int(3 * l_star)) == "kv-bound"
|
||||
|
||||
|
||||
# ── Component sanity ────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_total_active_params_llama3_8b_ballpark():
|
||||
"""Llama 3 8B — attn+FFN alone (no embeddings/LM head) is ~6.9B,
|
||||
HF-reported total 8.03B. Guard the ballpark."""
|
||||
n = total_active_params(PRESETS["Llama 3 8B"].model)
|
||||
assert 6.5e9 < n < 8e9, n
|
||||
|
||||
|
||||
def test_kv_bytes_per_token_llama3_8b():
|
||||
"""Llama 3 8B: 2*8*128*2 bytes/layer/token × 32 layers = 131072 bytes."""
|
||||
kv_bpt = kv_bytes_per_token(PRESETS["Llama 3 8B"].model)
|
||||
assert kv_bpt == 2 * 8 * 128 * 2 * 32
|
||||
|
||||
|
||||
# ── Regime terms + good-batch / good-context recommendations ──────
|
||||
|
||||
|
||||
def test_t_mem_short_shrinks_with_batch():
|
||||
"""t_mem_short = N·b/(W·B) — doubling B halves it."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
assert t_mem_short(m, model, 2) == pytest.approx(
|
||||
t_mem_short(m, model, 1) / 2
|
||||
)
|
||||
|
||||
|
||||
def test_t_mem_long_equals_t_com_at_l_star():
|
||||
"""L* is defined by t_mem_long(L*) == t_com. Cross-check."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
l_star = balance_context(m, model)
|
||||
assert t_mem_long(m, model, int(round(l_star))) == pytest.approx(
|
||||
t_com(m, model), rel=0.01,
|
||||
)
|
||||
|
||||
|
||||
def test_t_com_is_batch_independent():
|
||||
"""t_com = 2·N/C. Constant — doesn't depend on B or S_kv."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
tc = t_com(m, model)
|
||||
assert tc == pytest.approx(2 * total_active_params(model) / m.peak_flops)
|
||||
|
||||
|
||||
def test_good_batch_is_two_b_star():
|
||||
"""Recommendation is 2 × B* (Pope's rule)."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
rec = good_batch(m, model)
|
||||
assert rec.effective == pytest.approx(2 * critical_batch(m, model))
|
||||
assert rec.b_star == pytest.approx(critical_batch(m, model))
|
||||
|
||||
|
||||
def test_good_batch_moe_scales_with_sparsity():
|
||||
"""MoE sparsity=8 → recommended B is 8× dense recommendation."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
r_dense = good_batch(m, model, sparsity=1.0)
|
||||
r_moe = good_batch(m, model, sparsity=8.0)
|
||||
assert r_moe.effective == pytest.approx(8 * r_dense.effective)
|
||||
|
||||
|
||||
def test_good_context_returns_l_star_as_ceiling():
|
||||
"""Recommendation is L*; utilization drops past it."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
rec = good_context(m, model, s_kv=100)
|
||||
assert rec.l_star == pytest.approx(balance_context(m, model))
|
||||
assert rec.max_efficient == pytest.approx(rec.l_star)
|
||||
|
||||
|
||||
def test_utilization_at_l_star_is_half():
|
||||
"""At S_kv = L*, compute and KV read take equal time → 50% util."""
|
||||
l_star = 3000.0
|
||||
assert utilization_at(int(l_star), l_star) == pytest.approx(0.5)
|
||||
|
||||
|
||||
def test_utilization_at_2_l_star_is_one_third():
|
||||
"""At S_kv = 2·L*, util = 1/(1+2) = 33.3%."""
|
||||
l_star = 3000.0
|
||||
assert utilization_at(int(2 * l_star), l_star) == pytest.approx(1 / 3, rel=1e-3)
|
||||
|
||||
|
||||
def test_utilization_at_zero_context_is_100_percent():
|
||||
"""Zero context, no KV to read → all time is compute."""
|
||||
assert utilization_at(0, 3000.0) == pytest.approx(1.0)
|
||||
|
||||
|
||||
def test_utilization_drops_monotonically_with_context():
|
||||
"""util(S_kv) is strictly decreasing."""
|
||||
l_star = 3000.0
|
||||
vals = [utilization_at(s, l_star) for s in [100, 1000, 3000, 6000, 30000]]
|
||||
for a, b in zip(vals, vals[1:]):
|
||||
assert b < a
|
||||
|
||||
|
||||
# ── Step latency (undivided by B) ──────────────────────────────────
|
||||
|
||||
|
||||
def test_step_weight_is_batch_invariant():
|
||||
"""step_weight = N·b/W — same for every B (loaded once per step)."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
pts = step_latency_curve(m, model, [1, 8, 64, 512], s_kv=1024)
|
||||
ws = [p.weight_s for p in pts]
|
||||
assert all(w == pytest.approx(ws[0]) for w in ws), ws
|
||||
|
||||
|
||||
def test_step_compute_linear_in_batch():
|
||||
"""step_compute = 2·N·B/C — doubling B doubles compute."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
pts = step_latency_curve(m, model, [1, 2, 8], s_kv=1024)
|
||||
assert pts[1].compute_s == pytest.approx(2 * pts[0].compute_s)
|
||||
assert pts[2].compute_s == pytest.approx(8 * pts[0].compute_s)
|
||||
|
||||
|
||||
def test_step_kv_linear_in_batch():
|
||||
"""step_kv = B · S_kv · kv_bpt / W — linear in B."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
pts = step_latency_curve(m, model, [1, 4], s_kv=1024)
|
||||
assert pts[1].kv_s == pytest.approx(4 * pts[0].kv_s)
|
||||
|
||||
|
||||
def test_step_total_equals_per_token_times_batch():
|
||||
"""The step and per-token views are just ÷ B / × B of each other."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
B_range = [1, 4, 16, 64]
|
||||
s_kv = 1024
|
||||
step_pts = step_latency_curve(m, model, B_range, s_kv=s_kv)
|
||||
tok_pts = per_token_latency_curve(m, model, B_range, s_kv=s_kv)
|
||||
for sp, tp in zip(step_pts, tok_pts):
|
||||
assert sp.total_s == pytest.approx(tp.total_s * sp.batch)
|
||||
|
||||
|
||||
def test_step_and_per_token_identical_at_b1():
|
||||
"""At B=1, step latency == per-token cost (dividing by 1)."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
sp = step_latency_curve(m, model, [1], s_kv=1024)[0]
|
||||
tp = per_token_latency_curve(m, model, [1], s_kv=1024)[0]
|
||||
assert sp.total_s == pytest.approx(tp.total_s)
|
||||
assert sp.weight_s == pytest.approx(tp.weight_s)
|
||||
assert sp.compute_s == pytest.approx(tp.compute_s)
|
||||
assert sp.kv_s == pytest.approx(tp.kv_s)
|
||||
|
||||
|
||||
# ── PE memory budget sweeps ────────────────────────────────────────
|
||||
|
||||
|
||||
def _budget_cfg():
|
||||
"""Helper: real cfg with enough sharding to leave headroom."""
|
||||
return FullConfig(
|
||||
model=PRESETS["Llama 3 8B"].model,
|
||||
topo=TopologyConfig(cp=2, tp=8, pp=1, b=1, s_kv=8192),
|
||||
machine=MachineParams(),
|
||||
)
|
||||
|
||||
|
||||
def test_kv_grows_linear_with_skv():
|
||||
cfg = _budget_cfg()
|
||||
pts = memory_budget_curve_vs_skv(cfg, s_kv_range=[1024, 2048, 8192],
|
||||
batch=1)
|
||||
assert pts[1].kv_gb == pytest.approx(2 * pts[0].kv_gb, rel=1e-3)
|
||||
assert pts[2].kv_gb == pytest.approx(8 * pts[0].kv_gb, rel=1e-3)
|
||||
|
||||
|
||||
def test_kv_grows_linear_with_batch():
|
||||
cfg = _budget_cfg()
|
||||
pts = memory_budget_curve_vs_batch(cfg, b_range=[1, 2, 8],
|
||||
s_kv=8192)
|
||||
assert pts[1].kv_gb == pytest.approx(2 * pts[0].kv_gb, rel=1e-3)
|
||||
assert pts[2].kv_gb == pytest.approx(8 * pts[0].kv_gb, rel=1e-3)
|
||||
|
||||
|
||||
def test_weights_flat_across_skv_sweep():
|
||||
"""Weight bytes are batch/S_kv-invariant — sharding fixed."""
|
||||
cfg = _budget_cfg()
|
||||
pts = memory_budget_curve_vs_skv(cfg, s_kv_range=[128, 8192, 1_000_000],
|
||||
batch=1)
|
||||
ws = [p.weights_gb for p in pts]
|
||||
assert all(w == pytest.approx(ws[0]) for w in ws), ws
|
||||
|
||||
|
||||
def test_over_budget_flag_trips_past_hbm():
|
||||
"""Push S_kv until KV blows the HBM budget; over_budget must flip."""
|
||||
cfg = _budget_cfg()
|
||||
pts = memory_budget_curve_vs_skv(
|
||||
cfg, s_kv_range=[1024, 10_000_000], batch=64,
|
||||
)
|
||||
# 10M tokens × B=64 KV should overflow the 6 GB budget.
|
||||
assert not pts[0].over_budget, pts[0]
|
||||
assert pts[-1].over_budget, pts[-1]
|
||||
|
||||
|
||||
def test_free_gb_never_negative():
|
||||
cfg = _budget_cfg()
|
||||
pts = memory_budget_curve_vs_skv(
|
||||
cfg, s_kv_range=[128, 10_000_000], batch=32,
|
||||
)
|
||||
for p in pts:
|
||||
assert p.free_gb >= 0, p
|
||||
|
||||
|
||||
# ── AI sensitivity to FLOPs / BW ──────────────────────────────────
|
||||
|
||||
|
||||
def test_ai_scales_linearly_with_flops():
|
||||
"""Doubling FLOPs doubles AI."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
pts = ai_sensitivity_curve(m, model, [1.0, 2.0, 4.0], axis="flops")
|
||||
assert pts[1].ai == pytest.approx(2 * pts[0].ai)
|
||||
assert pts[2].ai == pytest.approx(4 * pts[0].ai)
|
||||
|
||||
|
||||
def test_ai_scales_inversely_with_bw():
|
||||
"""Doubling BW halves AI."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
pts = ai_sensitivity_curve(m, model, [1.0, 2.0], axis="bw")
|
||||
assert pts[1].ai == pytest.approx(pts[0].ai / 2)
|
||||
|
||||
|
||||
def test_ai_sensitivity_b_star_tracks_ai_for_bf16():
|
||||
"""For BF16, B* == AI. Curve should mirror."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model # BF16
|
||||
pts = ai_sensitivity_curve(m, model, [0.5, 1.0, 4.0], axis="flops")
|
||||
for p in pts:
|
||||
assert p.b_star == pytest.approx(p.ai)
|
||||
|
||||
|
||||
def test_ai_sensitivity_bad_axis_raises():
|
||||
with pytest.raises(ValueError):
|
||||
ai_sensitivity_curve(MachineParams(), PRESETS["Llama 3 8B"].model,
|
||||
[1.0], axis="bogus")
|
||||
|
||||
|
||||
# ── GPU sizing (three-axis) ────────────────────────────────────────
|
||||
|
||||
|
||||
def test_capacity_floor_scales_with_model_size():
|
||||
"""A model with 2× params needs ~2× PEs for weights alone (axis A)."""
|
||||
m = MachineParams()
|
||||
small = PRESETS["Llama 3 8B"].model
|
||||
big = PRESETS["Llama 3 70B"].model
|
||||
r_small = size_deployment(m, small, n_users=1, avg_ctx=1024,
|
||||
tpot_slo_s=1.0)
|
||||
r_big = size_deployment(m, big, n_users=1, avg_ctx=1024,
|
||||
tpot_slo_s=1.0)
|
||||
ratio = r_big.pes_axis_a_capacity / r_small.pes_axis_a_capacity
|
||||
# 70B / 8B ≈ 8.75× params — allow generous tolerance
|
||||
assert 5 < ratio < 12, (ratio, r_small, r_big)
|
||||
|
||||
|
||||
def test_kv_headroom_scales_with_users_and_context():
|
||||
"""Doubling n_users OR avg_ctx grows KV load roughly linearly."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
base = size_deployment(m, model, n_users=64, avg_ctx=4096,
|
||||
tpot_slo_s=1.0)
|
||||
twice_users = size_deployment(m, model, n_users=128, avg_ctx=4096,
|
||||
tpot_slo_s=1.0)
|
||||
twice_ctx = size_deployment(m, model, n_users=64, avg_ctx=8192,
|
||||
tpot_slo_s=1.0)
|
||||
# KV bytes doubled either way → axis-B PE count must strictly grow.
|
||||
assert twice_users.pes_axis_b_kv > base.pes_axis_b_kv
|
||||
assert twice_ctx.pes_axis_b_kv > base.pes_axis_b_kv
|
||||
|
||||
|
||||
def test_binding_axis_flips_with_workload():
|
||||
"""Huge N_users + short ctx → throughput binds.
|
||||
Huge ctx + few users → kv (or capacity) binds."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
big_users = size_deployment(m, model, n_users=10_000, avg_ctx=1024,
|
||||
tpot_slo_s=0.05)
|
||||
long_ctx = size_deployment(m, model, n_users=1, avg_ctx=500_000,
|
||||
tpot_slo_s=1.0)
|
||||
assert big_users.binding_axis == "throughput", big_users
|
||||
assert long_ctx.binding_axis in ("kv", "capacity"), long_ctx
|
||||
|
||||
|
||||
def test_max_batch_shrinks_with_shorter_slo():
|
||||
"""Tighter SLO → smaller B fits."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
loose = max_batch_within_slo(m, model, s_kv=1024, tpot_slo_s=1.0)
|
||||
tight = max_batch_within_slo(m, model, s_kv=1024, tpot_slo_s=0.05)
|
||||
assert tight < loose, (tight, loose)
|
||||
|
||||
|
||||
def test_max_batch_returns_zero_when_weight_time_exceeds_slo():
|
||||
"""If weight-fetch alone > SLO, no batch fits."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 70B"].model
|
||||
# 70B in BF16 = 140 GB. Reading at 256 GB/s takes ~547 ms.
|
||||
# SLO of 100 ms cannot be met.
|
||||
assert max_batch_within_slo(m, model, s_kv=100,
|
||||
tpot_slo_s=0.1) == 0
|
||||
|
||||
|
||||
def test_total_pes_equals_per_replica_times_replicas():
|
||||
"""Sanity: total_pes = pes_per_replica × n_replicas."""
|
||||
m = MachineParams()
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
r = size_deployment(m, model, n_users=500, avg_ctx=4096,
|
||||
tpot_slo_s=0.1)
|
||||
assert r.total_pes == r.pes_per_replica * r.n_replicas
|
||||
|
||||
|
||||
# ── App wiring: tab exists on the Streamlit app ────────────────────
|
||||
|
||||
|
||||
def test_roofline_tab_registered_in_app():
|
||||
"""The new 'Chip roofline & B*' tab is listed in st.tabs and gets a
|
||||
corresponding `with tab_roofline:` block."""
|
||||
from pathlib import Path
|
||||
src = (Path(__file__).parent / "app.py").resolve().read_text(
|
||||
encoding="utf-8"
|
||||
)
|
||||
assert src.count('"Chip roofline & B*"') == 1
|
||||
assert "with tab_roofline:" in src
|
||||
@@ -1,101 +0,0 @@
|
||||
"""Tests for draw_pe_layout — CP rank locator visibility across TP tiers.
|
||||
|
||||
The PE-level view has always color-labeled CP ranks when they're packed
|
||||
intra-cube (`cp_placement="pe"`). When TP >= 8, TP fills the cube and CP
|
||||
gets pushed to cube-level, so the intra-cube color branch is skipped —
|
||||
which used to silently drop the CP rank locator entirely. These tests
|
||||
guard the fix: a `[CP rank 0 of N]` cube tag + `CP=0` per-PE label are
|
||||
always present when CP > 1, regardless of placement.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
from tests.analytical_visualization.model_config import (
|
||||
FullConfig, MachineParams, TopologyConfig,
|
||||
)
|
||||
from tests.analytical_visualization.model_presets import PRESETS
|
||||
from tests.analytical_visualization.pe_weight_layout import draw_pe_layout
|
||||
|
||||
|
||||
def _cfg(cp: int, tp: int, cp_placement: str | None = None):
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
kwargs = dict(cp=cp, tp=tp, pp=1, s_kv=8192, mode="decode")
|
||||
if cp_placement is not None:
|
||||
kwargs["cp_placement"] = cp_placement
|
||||
return FullConfig(
|
||||
model=model, topo=TopologyConfig(**kwargs), machine=MachineParams(),
|
||||
)
|
||||
|
||||
|
||||
def _all_text(fig) -> str:
|
||||
"""Concatenate all text artists (labels + title) into one blob."""
|
||||
ax = fig.gca()
|
||||
parts = [t.get_text() for t in ax.texts]
|
||||
parts.append(ax.get_title())
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def test_high_tp_cp_gt_1_shows_rank_locator():
|
||||
"""TP=8, CP=4 -> cp_placement forced to 'cube'; CP rank locator
|
||||
must still appear in cube tag + per-PE header + figure title."""
|
||||
cfg = _cfg(cp=4, tp=8)
|
||||
assert cfg.topo.cp_placement == "cube" # sanity
|
||||
fig = draw_pe_layout(cfg)
|
||||
try:
|
||||
text = _all_text(fig)
|
||||
assert "CP rank 0 of 4" in text, text
|
||||
assert "CP=0" in text, text
|
||||
assert "CP=4" in text, text # figure title dim summary
|
||||
finally:
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def test_intra_cube_cp_placement_still_shows_per_rank_labels():
|
||||
"""cp_placement='pe' (small TP, CP>1) — the old branch already
|
||||
colors + labels each CP rank. Regression check: the new tag does
|
||||
NOT replace those, and each CP=0..CP-1 label is still emitted."""
|
||||
cfg = _cfg(cp=4, tp=2) # 4*2=8 fits intra-cube
|
||||
assert cfg.topo.cp_placement == "cube" # default, but auto/manual can differ
|
||||
# Force pe placement to exercise the intra-cube branch.
|
||||
cfg.topo.cp_placement = "pe"
|
||||
fig = draw_pe_layout(cfg)
|
||||
try:
|
||||
text = _all_text(fig)
|
||||
for r in range(4):
|
||||
assert f"CP={r}" in text, (r, text)
|
||||
# The 'CP rank 0 of N' tag is only for the cube-placement fallback.
|
||||
assert "CP rank 0 of" not in text, text
|
||||
finally:
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def test_cp_1_no_locator_added():
|
||||
"""CP=1: nothing to locate; no CP tag or label should appear."""
|
||||
cfg = _cfg(cp=1, tp=8)
|
||||
fig = draw_pe_layout(cfg)
|
||||
try:
|
||||
text = _all_text(fig)
|
||||
assert "CP rank" not in text, text
|
||||
# The header 'CP=0' should not be added when CP=1.
|
||||
# (Figure title mentions 'CP=1' — that's an OK dim summary,
|
||||
# not a locator per-PE.)
|
||||
assert " | CP=0" not in text, text
|
||||
finally:
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def test_tp_16_still_shows_cp_rank_across_spilled_cubes():
|
||||
"""TP=16 (spills to 2 cubes per CP group), CP=2 -> both cubes
|
||||
for CP rank 0 must carry the '[CP rank 0 of 2]' tag."""
|
||||
cfg = _cfg(cp=2, tp=16)
|
||||
assert cfg.topo.cp_placement == "cube"
|
||||
fig = draw_pe_layout(cfg)
|
||||
try:
|
||||
text = _all_text(fig)
|
||||
assert text.count("CP rank 0 of 2") >= 2, text
|
||||
assert "CP=0" in text, text
|
||||
finally:
|
||||
plt.close(fig)
|
||||
@@ -1,167 +0,0 @@
|
||||
"""Tests for stage_shapes: per-PE shape rows for attention + FFN stages."""
|
||||
from __future__ import annotations
|
||||
|
||||
from tests.analytical_visualization.model_config import (
|
||||
FullConfig, MachineParams, TopologyConfig,
|
||||
)
|
||||
from tests.analytical_visualization.model_presets import PRESETS
|
||||
from tests.analytical_visualization.stage_shapes import (
|
||||
attn_stage_shape_rows,
|
||||
ffn_stage_shape_rows,
|
||||
)
|
||||
|
||||
|
||||
def _cfg(cp=2, tp=8, pp=1, b=1, mode="decode", s_kv=8192):
|
||||
model = PRESETS["Llama 3 8B"].model
|
||||
topo = TopologyConfig(cp=cp, tp=tp, pp=pp, b=b, s_kv=s_kv, mode=mode)
|
||||
return FullConfig(model=model, topo=topo, machine=MachineParams())
|
||||
|
||||
|
||||
# ── Attention shapes ────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_attn_rows_have_required_columns():
|
||||
"""Every attention row must define Stage, Op, Input, Weight, Output."""
|
||||
rows = attn_stage_shape_rows(_cfg())
|
||||
required = {"Stage", "Op", "Input (per PE)",
|
||||
"Weight (per PE)", "Output (per PE)"}
|
||||
assert rows, "empty attention shape rows"
|
||||
for r in rows:
|
||||
missing = required - r.keys()
|
||||
assert not missing, f"row {r.get('Stage')} missing {missing}"
|
||||
for k in required:
|
||||
assert r[k], f"row {r['Stage']} col {k} is empty"
|
||||
|
||||
|
||||
def test_attn_row_prefixes_present_in_order():
|
||||
"""S1..S10 always appear; C1 only when prefill+CP>1; C2 when TP>1;
|
||||
C3 when kv_shard_mode=split and TP>H_kv; S8 only when CP>1."""
|
||||
cfg = _cfg(cp=2, tp=8, mode="decode")
|
||||
prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
|
||||
# decode with CP=2, TP=8 (h_kv=8 for Llama 3 8B) → no C1, no C3; C2 yes; S8 yes.
|
||||
assert "C1" not in prefixes, prefixes
|
||||
assert "C3" not in prefixes, prefixes
|
||||
assert "S8" in prefixes, prefixes
|
||||
assert "C2" in prefixes, prefixes
|
||||
# Ensure S1..S10 sequence appears in relative order.
|
||||
s_order = [p for p in prefixes if p.startswith("S")]
|
||||
assert s_order == sorted(s_order, key=lambda s: int(s[1:])), s_order
|
||||
|
||||
|
||||
def test_attn_cp1_omits_s8_merge():
|
||||
"""CP=1 → no online-softmax merge stage."""
|
||||
cfg = _cfg(cp=1, tp=8, mode="decode")
|
||||
prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
|
||||
assert "S8" not in prefixes
|
||||
|
||||
|
||||
def test_attn_prefill_cp_gt_1_inserts_c1():
|
||||
"""Prefill mode with CP>1 shows the CP-ring stage."""
|
||||
cfg = _cfg(cp=4, tp=2, mode="prefill", s_kv=8192)
|
||||
prefixes = [r["Stage"] for r in attn_stage_shape_rows(cfg)]
|
||||
assert "C1" in prefixes
|
||||
|
||||
|
||||
def test_attn_shape_reflects_batch():
|
||||
"""B appears as the leading dim in each activation shape."""
|
||||
for b in (1, 4, 32):
|
||||
rows = attn_stage_shape_rows(_cfg(b=b))
|
||||
s1 = next(r for r in rows if r["Stage"] == "S1")
|
||||
assert s1["Input (per PE)"].startswith(f"({b},"), s1
|
||||
assert s1["Output (per PE)"].startswith(f"({b},"), s1
|
||||
|
||||
|
||||
def test_attn_s5_shape_uses_tq_and_slocal():
|
||||
"""S5 (Q·K^T) shape must reference both T_q and S_local values."""
|
||||
cfg = _cfg(cp=2, tp=8, mode="decode", s_kv=8192)
|
||||
T_q = cfg.topo.T_q
|
||||
S_local = cfg.topo.s_local
|
||||
rows = attn_stage_shape_rows(cfg)
|
||||
s5 = next(r for r in rows if r["Stage"] == "S5")
|
||||
# Output shape (B, H_q/TP, T_q, S_local)
|
||||
assert str(T_q) in s5["Output (per PE)"], s5
|
||||
assert str(S_local) in s5["Output (per PE)"], s5
|
||||
|
||||
|
||||
# ── FFN shapes ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_ffn_rows_have_required_columns():
|
||||
"""Every FFN row must define Stage, Op, Input, Weight, Output."""
|
||||
rows = ffn_stage_shape_rows(_cfg())
|
||||
required = {"Stage", "Op", "Input (per PE)",
|
||||
"Weight (per PE)", "Output (per PE)"}
|
||||
assert rows, "empty ffn shape rows"
|
||||
for r in rows:
|
||||
missing = required - r.keys()
|
||||
assert not missing, f"row {r.get('Stage')} missing {missing}"
|
||||
|
||||
|
||||
def test_ffn_row_prefixes_f1_to_f5_always_present():
|
||||
"""F1..F5 always emitted; CF1 only when FFN divisor > 1."""
|
||||
rows = ffn_stage_shape_rows(_cfg(cp=2, tp=8))
|
||||
prefixes = [r["Stage"] for r in rows]
|
||||
for p in ("F1", "F2", "F3", "F4", "F5"):
|
||||
assert p in prefixes, prefixes
|
||||
|
||||
|
||||
def test_ffn_gate_up_output_matches_ffn_per_pe():
|
||||
"""F2/F3 output must have ffn_per_pe as the last dim."""
|
||||
cfg = _cfg(cp=2, tp=8)
|
||||
ffn_pe = max(1, cfg.model.ffn_dim
|
||||
// (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
|
||||
rows = ffn_stage_shape_rows(cfg)
|
||||
f2 = next(r for r in rows if r["Stage"] == "F2")
|
||||
assert str(ffn_pe) in f2["Output (per PE)"], (ffn_pe, f2)
|
||||
|
||||
|
||||
# ── App wiring: consolidated shapes section lives on layout tab ───
|
||||
|
||||
|
||||
def test_shapes_section_is_on_physical_layout_tab():
|
||||
"""The 'All tensor shapes (per PE)' expander must live inside
|
||||
tab_layout (Physical layout), not tab_stages. Guards the move so
|
||||
the section doesn't accidentally end up back on the latency tab.
|
||||
"""
|
||||
from pathlib import Path
|
||||
app_path = (
|
||||
Path(__file__).parent / "app.py"
|
||||
).resolve()
|
||||
src = app_path.read_text(encoding="utf-8")
|
||||
# Section appears exactly once.
|
||||
assert src.count('"All tensor shapes (per PE)"') == 1, src.count(
|
||||
'"All tensor shapes (per PE)"'
|
||||
)
|
||||
# And it sits inside tab_layout, before tab_memory begins.
|
||||
section_pos = src.index('"All tensor shapes (per PE)"')
|
||||
tab_memory_pos = src.index("with tab_memory:")
|
||||
tab_stages_pos = src.index("with tab_stages:")
|
||||
assert section_pos < tab_memory_pos, (
|
||||
"shapes section must be inside tab_layout (before tab_memory)"
|
||||
)
|
||||
assert section_pos < tab_stages_pos, (
|
||||
"shapes section must be inside tab_layout (before tab_stages)"
|
||||
)
|
||||
# The old per-stage shape rendering under tab_stages must be gone —
|
||||
# no attn_stage_shape_rows / ffn_stage_shape_rows call may sit
|
||||
# between 'with tab_stages:' and the next 'with tab_' block.
|
||||
stages_block = src[tab_stages_pos:src.index("# ── TAB", tab_stages_pos + 1)]
|
||||
assert "attn_stage_shape_rows(cfg)" not in stages_block, (
|
||||
"per-stage shape table must not remain on Per-stage latency tab"
|
||||
)
|
||||
assert "ffn_stage_shape_rows(cfg)" not in stages_block, (
|
||||
"per-stage FFN shape table must not remain on Per-stage latency tab"
|
||||
)
|
||||
|
||||
|
||||
def test_ttft_tpot_section_wired_on_layout_tab():
|
||||
"""The 'Full-model latency (TTFT & TPOT)' subheader must live on
|
||||
tab_layout, before tab_memory. Guard the section presence + placement."""
|
||||
from pathlib import Path
|
||||
src = (
|
||||
Path(__file__).parent / "app.py"
|
||||
).resolve().read_text(encoding="utf-8")
|
||||
assert src.count('"Full-model latency (TTFT & TPOT)"') == 1
|
||||
pos = src.index('"Full-model latency (TTFT & TPOT)"')
|
||||
tab_memory_pos = src.index("with tab_memory:")
|
||||
assert pos < tab_memory_pos, "TTFT section must live inside tab_layout"
|
||||
@@ -1,626 +0,0 @@
|
||||
"""Draw the SIP topology diagram, color-coded by parallelism group.
|
||||
|
||||
Color scheme:
|
||||
- CP ring: orange arrows (unchanged)
|
||||
- Inter-SIP: red arrows across SIP boundaries
|
||||
- PP stage: cube border color (each PP stage gets a distinct hue)
|
||||
- TP group: PE fill color (each TP group within a cube gets one hue,
|
||||
matches its cube's PP stage but slightly desaturated)
|
||||
- EP group: a small badge on the cube (for MoE)
|
||||
- DP replica: cubes of each DP replica get a hatched pattern
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.patches as patches
|
||||
from matplotlib.lines import Line2D
|
||||
|
||||
from .model_config import FullConfig
|
||||
|
||||
|
||||
MESH_W = 4
|
||||
MESH_H = 4
|
||||
PE_COLS = 4
|
||||
PE_ROWS = 2
|
||||
|
||||
|
||||
def _group_label(cfg, group_idx: int) -> str:
|
||||
"""Label for a cube-level group under the current placement.
|
||||
- cp_placement=cube -> CP rank (or CP/TP if both on cube)
|
||||
- cp_placement=pe, tp_placement=cube -> TP rank
|
||||
- both on pe -> single group ("PE")
|
||||
"""
|
||||
parts = []
|
||||
if cfg.topo.cp_placement == "cube":
|
||||
cp_r = group_idx % max(1, cfg.topo.cp)
|
||||
parts.append(f"CP{cp_r}")
|
||||
if cfg.topo.tp_placement == "cube":
|
||||
div = cfg.topo.cp if cfg.topo.cp_placement == "cube" else 1
|
||||
tp_r = (group_idx // max(1, div)) % max(1, cfg.topo.tp)
|
||||
parts.append(f"TP{tp_r}")
|
||||
return "/".join(parts) if parts else "PE"
|
||||
|
||||
# Fully-distinct palette for (PP stage x CP rank) — 24 unique colors.
|
||||
# Each cube's (pp, cp) pair maps to one entry; wraps if you exceed 24 groups.
|
||||
_GROUP_PALETTE = [
|
||||
"#e6194b", "#3cb44b", "#ffe119", "#4363d8",
|
||||
"#f58231", "#911eb4", "#42d4f4", "#f032e6",
|
||||
"#bfef45", "#fabed4", "#469990", "#dcbeff",
|
||||
"#9a6324", "#fffac8", "#800000", "#aaffc3",
|
||||
"#808000", "#ffd8b1", "#000075", "#a9a9a9",
|
||||
"#004d40", "#c62828", "#4527a0", "#00695c",
|
||||
]
|
||||
|
||||
# Kept for backward-compat / PP-stage arrows
|
||||
_PP_HUES = _GROUP_PALETTE[:8]
|
||||
|
||||
|
||||
def _shade(hue_hex: str, t: float) -> str:
|
||||
"""Mix hue_hex toward white by t in [0, 1] (0 = original, 1 = white)."""
|
||||
h = hue_hex.lstrip("#")
|
||||
r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
|
||||
rn = int(r + (255 - r) * t)
|
||||
gn = int(g + (255 - g) * t)
|
||||
bn = int(b + (255 - b) * t)
|
||||
return f"#{rn:02x}{gn:02x}{bn:02x}"
|
||||
|
||||
|
||||
def _cp_color(pp_stage: int, cp_rank: int, cp_size: int) -> tuple[str, str]:
|
||||
"""Return (border_color, pe_fill_color) for a cube at (pp_stage, cp_rank).
|
||||
|
||||
Each unique (pp, cp) pair gets a **fully distinct color** from the
|
||||
24-entry palette. Border is the pure color; PE fill is a lightened
|
||||
version so the cube border stays visible around the PE grid.
|
||||
"""
|
||||
idx = pp_stage * max(1, cp_size) + cp_rank
|
||||
color = _GROUP_PALETTE[idx % len(_GROUP_PALETTE)]
|
||||
pe_fill = _shade(color, 0.30) # slightly lightened for PE fill
|
||||
return color, pe_fill
|
||||
|
||||
# Inactive (unused) colors.
|
||||
_INACTIVE_CUBE_FACE = "#f7f7f7"
|
||||
_INACTIVE_CUBE_EDGE = "#c8c8c8"
|
||||
_INACTIVE_PE = "#eaeaea"
|
||||
|
||||
|
||||
def _snake_path(n: int, mesh_w: int, mesh_h: int) -> list[int]:
|
||||
path: list[int] = []
|
||||
for r in range(mesh_h):
|
||||
cols = range(mesh_w) if r % 2 == 0 else range(mesh_w - 1, -1, -1)
|
||||
for c in cols:
|
||||
path.append(r * mesh_w + c)
|
||||
if len(path) == n:
|
||||
return path
|
||||
return path[:n]
|
||||
|
||||
|
||||
def _rect_shape(n: int, max_w: int, max_h: int) -> tuple[int, int]:
|
||||
"""Return (rows, cols) for n cubes packed as square as possible in max_wxmax_h."""
|
||||
if n <= 0:
|
||||
return (0, 0)
|
||||
# Prefer near-square shapes; cap at mesh dims.
|
||||
best = None
|
||||
for cols in range(1, min(n, max_w) + 1):
|
||||
rows = (n + cols - 1) // cols
|
||||
if rows > max_h:
|
||||
continue
|
||||
# Score: penalise non-squareness + wasted cells
|
||||
waste = rows * cols - n
|
||||
aspect = abs(rows - cols)
|
||||
score = aspect * 10 + waste
|
||||
if best is None or score < best[0]:
|
||||
best = (score, rows, cols)
|
||||
if best is None:
|
||||
# Fallback: single row
|
||||
return (1, min(n, max_w))
|
||||
return best[1], best[2]
|
||||
|
||||
|
||||
def _pack_groups_2d(items_per_group: int, n_groups: int,
|
||||
mesh_w: int, mesh_h: int,
|
||||
layout_mode: str = "compact") -> list[tuple[int, int]]:
|
||||
"""Return list of (row, col) positions for n_groups × items_per_group
|
||||
cubes in a mesh_w x mesh_h mesh.
|
||||
|
||||
layout_mode:
|
||||
- "compact": each group is a near-square rectangle; groups are also
|
||||
tiled 2D as near-square. Ideal for 2x2, 2x4 etc.
|
||||
- "linear": each group is a single row of cubes; groups tile down
|
||||
as rows (original sequential layout).
|
||||
"""
|
||||
if n_groups == 0 or items_per_group == 0:
|
||||
return []
|
||||
|
||||
if layout_mode == "linear":
|
||||
# Each group is a single horizontal row of cubes.
|
||||
tp_rows, tp_cols = 1, min(items_per_group, mesh_w)
|
||||
else: # compact
|
||||
tp_rows, tp_cols = _rect_shape(items_per_group, mesh_w, mesh_h)
|
||||
if tp_rows == 0:
|
||||
return []
|
||||
|
||||
# Group grid: how many groups per row/col of GROUPS
|
||||
max_group_cols = max(1, mesh_w // tp_cols)
|
||||
max_group_rows = max(1, mesh_h // tp_rows)
|
||||
if layout_mode == "linear":
|
||||
# Row-major: fill each mesh row with groups, then wrap down.
|
||||
g_cols = max_group_cols
|
||||
g_rows = max_group_rows
|
||||
else:
|
||||
g_rows, g_cols = _rect_shape(n_groups, max_group_cols, max_group_rows)
|
||||
|
||||
positions: list[tuple[int, int]] = []
|
||||
for g in range(n_groups):
|
||||
gr = g // g_cols
|
||||
gc = g % g_cols
|
||||
base_row = gr * tp_rows
|
||||
base_col = gc * tp_cols
|
||||
for i in range(items_per_group):
|
||||
r = i // tp_cols
|
||||
c = i % tp_cols
|
||||
row = base_row + r
|
||||
col = base_col + c
|
||||
if row >= mesh_h or col >= mesh_w:
|
||||
fallback = (g * items_per_group + i)
|
||||
row = fallback // mesh_w
|
||||
col = fallback % mesh_w
|
||||
positions.append((row, col))
|
||||
return positions
|
||||
|
||||
|
||||
def _draw_one_sip(ax, cfg: FullConfig, sip_idx: int,
|
||||
sip_x0: float, sip_y0: float,
|
||||
sip_width: float, sip_height: float,
|
||||
cube_pp_cp: dict[int, tuple[int, int]],
|
||||
ring_paths: dict[int, list[int]],
|
||||
tp_group_cubes: dict[int, list[int]]):
|
||||
"""Draw one SIP with PP hue + CP shade + TP group boundaries."""
|
||||
cube_size = min(sip_width / MESH_W, sip_height / MESH_H) * 0.85
|
||||
cube_gap = cube_size * 0.15
|
||||
total_w = MESH_W * (cube_size + cube_gap) - cube_gap
|
||||
total_h = MESH_H * (cube_size + cube_gap) - cube_gap
|
||||
ox = sip_x0 + (sip_width - total_w) / 2
|
||||
oy = sip_y0 + (sip_height - total_h) / 2
|
||||
|
||||
sip_rect = patches.FancyBboxPatch(
|
||||
(sip_x0, sip_y0), sip_width, sip_height,
|
||||
boxstyle="round,pad=0.02", facecolor="#f8f9fa",
|
||||
edgecolor="#212529", linewidth=1.5,
|
||||
)
|
||||
ax.add_patch(sip_rect)
|
||||
ax.text(sip_x0 + sip_width / 2, sip_y0 + sip_height + 0.05,
|
||||
f"SIP {sip_idx}", ha="center", va="bottom",
|
||||
fontsize=10, fontweight="bold")
|
||||
|
||||
def _cube_center(cube_local: int) -> tuple[float, float]:
|
||||
r = cube_local // MESH_W
|
||||
c = cube_local % MESH_W
|
||||
x = ox + c * (cube_size + cube_gap)
|
||||
y = oy + (MESH_H - 1 - r) * (cube_size + cube_gap)
|
||||
return x, y
|
||||
|
||||
for cube_local in range(MESH_W * MESH_H):
|
||||
x, y = _cube_center(cube_local)
|
||||
info = cube_pp_cp.get(cube_local, None)
|
||||
is_used = info is not None
|
||||
|
||||
if is_used:
|
||||
pp_stage, cp_rank = info
|
||||
cube_edge, pe_fill = _cp_color(
|
||||
pp_stage, cp_rank,
|
||||
max(1, cfg.topo.inter_cube_dims),
|
||||
)
|
||||
cube_face = _shade(cube_edge, 0.85)
|
||||
cube_lw = 2.0
|
||||
else:
|
||||
cube_edge = _INACTIVE_CUBE_EDGE
|
||||
cube_face = _INACTIVE_CUBE_FACE
|
||||
cube_lw = 0.6
|
||||
pe_fill = _INACTIVE_PE
|
||||
|
||||
rect = patches.FancyBboxPatch(
|
||||
(x, y), cube_size, cube_size,
|
||||
boxstyle="round,pad=0.01",
|
||||
facecolor=cube_face, edgecolor=cube_edge, linewidth=cube_lw,
|
||||
)
|
||||
ax.add_patch(rect)
|
||||
if is_used:
|
||||
pp_stage, cp_rank = info
|
||||
# Small cube ID above
|
||||
ax.text(x + cube_size / 2, y + cube_size + 0.01,
|
||||
f"cube {cube_local}", ha="center", va="bottom",
|
||||
fontsize=5.5, color=cube_edge, fontweight="bold")
|
||||
# Group label (CP/TP depending on placement)
|
||||
_lbl = _group_label(cfg, cp_rank)
|
||||
ax.text(x + 0.03, y + cube_size - 0.03,
|
||||
_lbl, ha="left", va="top",
|
||||
fontsize=11, color=cube_edge, fontweight="bold",
|
||||
bbox=dict(boxstyle="round,pad=0.15",
|
||||
facecolor="white", edgecolor=cube_edge,
|
||||
linewidth=1.0, alpha=0.9))
|
||||
if cfg.topo.pp > 1:
|
||||
ax.text(x + cube_size - 0.03, y + cube_size - 0.03,
|
||||
f"PP{pp_stage}", ha="right", va="top",
|
||||
fontsize=9, color=cube_edge, fontweight="bold",
|
||||
bbox=dict(boxstyle="round,pad=0.1",
|
||||
facecolor="white", edgecolor=cube_edge,
|
||||
linewidth=0.8, alpha=0.9))
|
||||
else:
|
||||
ax.text(x + cube_size / 2, y + cube_size + 0.008,
|
||||
f"{cube_local}", ha="center", va="bottom",
|
||||
fontsize=6, color=cube_edge)
|
||||
|
||||
# PEs
|
||||
pe_pad = cube_size * 0.08
|
||||
pe_gap = cube_size * 0.02
|
||||
inner_w = cube_size - 2 * pe_pad
|
||||
inner_h = cube_size - 2 * pe_pad
|
||||
pe_w = (inner_w - (PE_COLS - 1) * pe_gap) / PE_COLS
|
||||
pe_h = (inner_h - (PE_ROWS - 1) * pe_gap) / PE_ROWS
|
||||
pes_used = cfg.topo.pes_per_cube_used if is_used else 0
|
||||
# When cp_placement=pe, multiple CP ranks live inside this cube's PEs
|
||||
# (cp_rank = pe_id // tp). Color each PE by its cp_rank so all four
|
||||
# groups are visible; otherwise fall back to the whole-cube pe_fill.
|
||||
_cp_packed = is_used and cfg.topo.cp_placement == "pe" and cfg.topo.cp > 1
|
||||
for pr in range(PE_ROWS):
|
||||
for pc in range(PE_COLS):
|
||||
pe_id = pr * PE_COLS + pc
|
||||
px = x + pe_pad + pc * (pe_w + pe_gap)
|
||||
py = y + pe_pad + (PE_ROWS - 1 - pr) * (pe_h + pe_gap)
|
||||
if pe_id >= pes_used:
|
||||
fill = _INACTIVE_PE
|
||||
elif _cp_packed:
|
||||
_pe_cp_rank = pe_id // cfg.topo.tp
|
||||
_pe_color, _pe_pe_fill = _cp_color(
|
||||
info[0] if is_used else 0,
|
||||
_pe_cp_rank,
|
||||
cfg.topo.cp,
|
||||
)
|
||||
fill = _pe_pe_fill
|
||||
else:
|
||||
fill = pe_fill
|
||||
pe_rect = patches.Rectangle(
|
||||
(px, py), pe_w, pe_h,
|
||||
facecolor=fill, edgecolor="#666", linewidth=0.3,
|
||||
)
|
||||
ax.add_patch(pe_rect)
|
||||
|
||||
if is_used and cfg.topo.ep > 1:
|
||||
ax.text(x + cube_size - 0.02, y + cube_size - 0.02,
|
||||
f"EP{cfg.topo.ep}", ha="right", va="top",
|
||||
fontsize=6, color="#5f0f40", fontweight="bold")
|
||||
|
||||
# (CP ring arrows removed per user request)
|
||||
|
||||
# Draw TP group bounding boxes (RED dashed) when TP > 8
|
||||
for tp_g, cubes in tp_group_cubes.items():
|
||||
if len(cubes) <= 1:
|
||||
continue # TP fits in one cube, redundant with cube border
|
||||
# Bounding box around all cubes in this TP group
|
||||
xs, ys = [], []
|
||||
for cube_local in cubes:
|
||||
x, y = _cube_center(cube_local)
|
||||
xs.extend([x, x + cube_size])
|
||||
ys.extend([y, y + cube_size])
|
||||
pad = 0.03
|
||||
bbox_x = min(xs) - pad
|
||||
bbox_y = min(ys) - pad
|
||||
bbox_w = (max(xs) - min(xs)) + 2 * pad
|
||||
bbox_h = (max(ys) - min(ys)) + 2 * pad
|
||||
tp_rect = patches.Rectangle(
|
||||
(bbox_x, bbox_y), bbox_w, bbox_h,
|
||||
fill=False, edgecolor="#d90429", linewidth=1.8,
|
||||
linestyle="--",
|
||||
)
|
||||
ax.add_patch(tp_rect)
|
||||
ax.text(bbox_x + bbox_w / 2, bbox_y - 0.05,
|
||||
f"TP group {tp_g}", ha="center", va="top",
|
||||
color="#d90429", fontsize=7, fontweight="bold")
|
||||
|
||||
|
||||
def _sip_grid_layout(n_sips: int, topology: str) -> tuple[int, int]:
|
||||
"""(rows, cols) for arranging SIPs on the page.
|
||||
- ring: horizontal chain (1 x N)
|
||||
- mesh2d / torus2d: near-square grid
|
||||
"""
|
||||
if n_sips <= 0:
|
||||
return 0, 0
|
||||
if topology == "ring" or n_sips <= 2:
|
||||
return 1, n_sips
|
||||
import math
|
||||
cols = int(math.ceil(math.sqrt(n_sips)))
|
||||
rows = (n_sips + cols - 1) // cols
|
||||
return rows, cols
|
||||
|
||||
|
||||
def _draw_sip_links(ax, sip_xy: list[tuple[float, float]],
|
||||
topology: str, sip_w: float, sip_h: float,
|
||||
grid_rows: int, grid_cols: int):
|
||||
"""Draw inter-SIP interconnect lines.
|
||||
|
||||
Solid lines: adjacent (grid-neighbor) links.
|
||||
Dashed lines: wrap-around links (ring, torus).
|
||||
"""
|
||||
n = len(sip_xy)
|
||||
if n <= 1:
|
||||
return
|
||||
color = "#c62828"
|
||||
lw = 2.0
|
||||
|
||||
def center(i):
|
||||
x, y = sip_xy[i]
|
||||
return x + sip_w / 2, y + sip_h / 2
|
||||
|
||||
if topology == "ring":
|
||||
# 1D chain (adjacent solid)
|
||||
for i in range(n - 1):
|
||||
x1, y1 = center(i)
|
||||
x2, y2 = center(i + 1)
|
||||
ax.plot([x1 + sip_w / 2, x2 - sip_w / 2], [y1, y2],
|
||||
color=color, linewidth=lw)
|
||||
# wrap-around from last back to first (dashed arc below)
|
||||
if n > 2:
|
||||
x_first, y_first = center(0)
|
||||
x_last, y_last = center(n - 1)
|
||||
arc_y = min(y_first, y_last) - sip_h / 2 - 0.35
|
||||
ax.plot([x_first, x_first], [y_first - sip_h / 2, arc_y],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
ax.plot([x_first, x_last], [arc_y, arc_y],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
ax.plot([x_last, x_last], [arc_y, y_last - sip_h / 2],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
elif n == 2:
|
||||
# 2 SIPs: single link
|
||||
x1, y1 = center(0)
|
||||
x2, y2 = center(1)
|
||||
# wrap link: dashed loop below
|
||||
arc_y = y1 - sip_h / 2 - 0.35
|
||||
ax.plot([x1, x1], [y1 - sip_h / 2, arc_y],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
ax.plot([x1, x2], [arc_y, arc_y],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
ax.plot([x2, x2], [arc_y, y2 - sip_h / 2],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
return
|
||||
|
||||
# mesh2d / torus2d: grid neighbours
|
||||
def idx(r, c):
|
||||
return r * grid_cols + c
|
||||
|
||||
for r in range(grid_rows):
|
||||
for c in range(grid_cols):
|
||||
i = idx(r, c)
|
||||
if i >= n:
|
||||
continue
|
||||
# right
|
||||
if c + 1 < grid_cols and idx(r, c + 1) < n:
|
||||
x1, y1 = center(i)
|
||||
x2, y2 = center(idx(r, c + 1))
|
||||
ax.plot([x1 + sip_w / 2, x2 - sip_w / 2], [y1, y2],
|
||||
color=color, linewidth=lw)
|
||||
# down (visually — grid row+1 is drawn below)
|
||||
if r + 1 < grid_rows and idx(r + 1, c) < n:
|
||||
x1, y1 = center(i)
|
||||
x2, y2 = center(idx(r + 1, c))
|
||||
ax.plot([x1, x2], [y1 - sip_h / 2, y2 + sip_h / 2],
|
||||
color=color, linewidth=lw)
|
||||
|
||||
if topology == "torus2d":
|
||||
# row wrap: last col -> first col in each row (dashed arc BELOW that row)
|
||||
for r in range(grid_rows):
|
||||
left_i = idx(r, 0)
|
||||
right_i = idx(r, grid_cols - 1)
|
||||
if left_i >= n or right_i >= n or left_i == right_i:
|
||||
continue
|
||||
x1, y1 = center(left_i)
|
||||
x2, y2 = center(right_i)
|
||||
arc_y = min(y1, y2) - sip_h / 2 - 0.35
|
||||
ax.plot([x1, x1], [y1 - sip_h / 2, arc_y],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
ax.plot([x1, x2], [arc_y, arc_y],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
ax.plot([x2, x2], [arc_y, y2 - sip_h / 2],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
# col wrap: last row -> first row in each col (dashed arc to the RIGHT)
|
||||
for c in range(grid_cols):
|
||||
top_i = idx(0, c)
|
||||
bot_i = idx(grid_rows - 1, c)
|
||||
if top_i >= n or bot_i >= n or top_i == bot_i:
|
||||
continue
|
||||
x1, y1 = center(top_i) # top row (higher y)
|
||||
x2, y2 = center(bot_i) # bottom row (lower y)
|
||||
arc_x = max(x1, x2) + sip_w / 2 + 0.35
|
||||
ax.plot([x1 + sip_w / 2, arc_x], [y1, y1],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
ax.plot([arc_x, arc_x], [y1, y2],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
ax.plot([arc_x, x2 + sip_w / 2], [y2, y2],
|
||||
color=color, linewidth=lw, linestyle="--")
|
||||
|
||||
|
||||
def draw_topology(cfg: FullConfig, ax=None, layout_mode: str = "linear"):
|
||||
"""Draw one or more SIPs. Each (PP,CP) group gets a unique color.
|
||||
|
||||
SIP grid arrangement + inter-SIP links depend on cfg.topo.sip_topology:
|
||||
"ring" -> 1D horizontal chain, wrap arc below
|
||||
"mesh2d" -> near-square grid, no wrap
|
||||
"torus2d" -> grid + row/col wrap arcs
|
||||
"""
|
||||
sips_used = max(1, cfg.topo.sips_used)
|
||||
total_pes = cfg.topo.total_pes
|
||||
pes_per_stage = cfg.topo.pes_per_stage # CP*TP
|
||||
n_replicas = cfg.topo.dp
|
||||
n_pp_stages = cfg.topo.pp
|
||||
pes_per_cube = cfg.topo.pes_per_cube_hw
|
||||
# Placement-aware: how many cubes per "group" (was: cubes-per-TP).
|
||||
# A group is one cube-level unit (one CP rank if cp on cube, or one TP
|
||||
# rank if only tp on cube, etc.). Cubes per group == intra-cube spill.
|
||||
cubes_per_group = max(
|
||||
1, (cfg.topo.intra_cube_dims + pes_per_cube - 1) // pes_per_cube
|
||||
)
|
||||
n_groups_per_stage = max(1, cfg.topo.inter_cube_dims)
|
||||
cubes_per_stage = cfg.topo.cubes_per_stage
|
||||
cubes_used = cfg.topo.cubes_used
|
||||
cubes_per_sip = MESH_W * MESH_H
|
||||
|
||||
sip_topo = cfg.topo.sip_topology
|
||||
grid_rows, grid_cols = _sip_grid_layout(sips_used, sip_topo)
|
||||
|
||||
sip_w = 4.0
|
||||
sip_h = 4.0
|
||||
sip_gap = 0.6
|
||||
|
||||
fig_w = grid_cols * sip_w + (grid_cols - 1) * sip_gap + 1.5
|
||||
fig_h = grid_rows * sip_h + (grid_rows - 1) * sip_gap + 2.0
|
||||
if ax is None:
|
||||
fig, ax = plt.subplots(figsize=(max(6, fig_w), max(6, fig_h)))
|
||||
else:
|
||||
fig = ax.figure
|
||||
|
||||
# 2D packing: each group's cubes form a rectangle inside the SIP mesh.
|
||||
# If more groups than fit in one SIP, spill to next SIP.
|
||||
tp_rows, tp_cols = _rect_shape(cubes_per_group, MESH_W, MESH_H)
|
||||
if tp_rows == 0:
|
||||
tp_rows, tp_cols = 1, 1
|
||||
groups_per_sip = max(1, (MESH_W // tp_cols) * (MESH_H // tp_rows))
|
||||
|
||||
cube_pp_cp_global: dict[int, tuple[int, int]] = {}
|
||||
ring_globals_by_pp: dict[int, list[int]] = {}
|
||||
tp_group_cubes_global: dict[int, list[int]] = {}
|
||||
|
||||
current_sip = 0
|
||||
for rep in range(n_replicas):
|
||||
for pp_s in range(n_pp_stages):
|
||||
# Chunk cube-level groups into per-SIP batches.
|
||||
for chunk_start in range(0, n_groups_per_stage, groups_per_sip):
|
||||
chunk = min(groups_per_sip, n_groups_per_stage - chunk_start)
|
||||
positions = _pack_groups_2d(cubes_per_group, chunk,
|
||||
MESH_W, MESH_H,
|
||||
layout_mode=layout_mode)
|
||||
for local_gi in range(chunk):
|
||||
group_idx = chunk_start + local_gi
|
||||
tp_gid = ((rep * n_pp_stages) + pp_s) * n_groups_per_stage + group_idx
|
||||
cubes_here: list[int] = []
|
||||
for i in range(cubes_per_group):
|
||||
idx_in_group = local_gi * cubes_per_group + i
|
||||
r, c = positions[idx_in_group]
|
||||
local_cube = r * MESH_W + c
|
||||
gc = current_sip * cubes_per_sip + local_cube
|
||||
cube_pp_cp_global[gc] = (pp_s, group_idx)
|
||||
cubes_here.append(gc)
|
||||
tp_group_cubes_global[tp_gid] = cubes_here
|
||||
ring_globals_by_pp.setdefault(pp_s, []).append(cubes_here[0])
|
||||
current_sip += 1 # next SIP for the next chunk (or (rep, pp))
|
||||
|
||||
# Partition (pp, cp) map by SIP using cube_positions_global
|
||||
cubes_by_sip: list[dict[int, tuple[int, int]]] = [dict() for _ in range(sips_used)]
|
||||
for gc, ppcp in cube_pp_cp_global.items():
|
||||
sip_idx = gc // cubes_per_sip
|
||||
local_cube = gc % cubes_per_sip
|
||||
if sip_idx < sips_used:
|
||||
cubes_by_sip[sip_idx][local_cube] = ppcp
|
||||
|
||||
# Per-SIP snake-path CP rings (organized by pp_stage)
|
||||
ring_by_sip: list[dict[int, list[int]]] = [dict() for _ in range(sips_used)]
|
||||
for pp_s, gc_list in ring_globals_by_pp.items():
|
||||
by_sip_local: dict[int, list[int]] = {}
|
||||
for gc in gc_list:
|
||||
si = gc // cubes_per_sip
|
||||
lc = gc % cubes_per_sip
|
||||
by_sip_local.setdefault(si, []).append(lc)
|
||||
for si, locals_list in by_sip_local.items():
|
||||
if si < sips_used:
|
||||
# Snake through the packed positions in order
|
||||
ring_by_sip[si][pp_s] = locals_list
|
||||
|
||||
# TP group cubes by SIP
|
||||
tp_by_sip: list[dict[int, list[int]]] = [dict() for _ in range(sips_used)]
|
||||
for tp_g, gc_list in tp_group_cubes_global.items():
|
||||
by_sip_local: dict[int, list[int]] = {}
|
||||
for gc in gc_list:
|
||||
si = gc // cubes_per_sip
|
||||
lc = gc % cubes_per_sip
|
||||
by_sip_local.setdefault(si, []).append(lc)
|
||||
for si, locals_list in by_sip_local.items():
|
||||
if si < sips_used:
|
||||
tp_by_sip[si][tp_g] = locals_list
|
||||
|
||||
# Compute (x, y) top-left for each SIP based on grid layout.
|
||||
# Row 0 is drawn at the TOP visually, so y decreases with row index.
|
||||
sip_xy: list[tuple[float, float]] = []
|
||||
for sip_idx in range(sips_used):
|
||||
r = sip_idx // grid_cols
|
||||
c = sip_idx % grid_cols
|
||||
x = c * (sip_w + sip_gap)
|
||||
y = (grid_rows - 1 - r) * (sip_h + sip_gap)
|
||||
sip_xy.append((x, y))
|
||||
|
||||
for sip_idx in range(sips_used):
|
||||
x0, y0 = sip_xy[sip_idx]
|
||||
_draw_one_sip(
|
||||
ax, cfg, sip_idx,
|
||||
sip_x0=x0, sip_y0=y0,
|
||||
sip_width=sip_w, sip_height=sip_h,
|
||||
cube_pp_cp=cubes_by_sip[sip_idx],
|
||||
ring_paths=ring_by_sip[sip_idx],
|
||||
tp_group_cubes=tp_by_sip[sip_idx],
|
||||
)
|
||||
|
||||
# Inter-SIP interconnect links per user-selected topology.
|
||||
_draw_sip_links(ax, sip_xy, sip_topo, sip_w, sip_h, grid_rows, grid_cols)
|
||||
|
||||
total_w = grid_cols * sip_w + (grid_cols - 1) * sip_gap
|
||||
total_h = grid_rows * sip_h + (grid_rows - 1) * sip_gap
|
||||
# Extra room: bottom for wrap arcs, right for torus col-wrap arcs, top for SIP labels
|
||||
ax.set_xlim(-0.3, total_w + 0.9)
|
||||
ax.set_ylim(-0.9, total_h + 0.7)
|
||||
ax.set_aspect("equal")
|
||||
ax.axis("off")
|
||||
|
||||
# Legend — one entry per unique (PP, cube-group) + TP boundary + SIP links
|
||||
legend_handles = []
|
||||
for pp_s in range(n_pp_stages):
|
||||
for gi in range(n_groups_per_stage):
|
||||
color, _ = _cp_color(pp_s, gi, n_groups_per_stage)
|
||||
legend_handles.append(Line2D(
|
||||
[], [], color=color, marker="s", linestyle="",
|
||||
markersize=10,
|
||||
label=(f"PP{pp_s} " if n_pp_stages > 1 else "")
|
||||
+ _group_label(cfg, gi),
|
||||
))
|
||||
if cfg.topo.tp > cfg.topo.pes_per_cube_hw:
|
||||
legend_handles.append(Line2D(
|
||||
[], [], color="#d90429", marker="s", linestyle="--",
|
||||
markersize=10, markerfacecolor="none", markeredgewidth=1.5,
|
||||
label="TP group (dashed red)",
|
||||
))
|
||||
if cfg.topo.ep > 1:
|
||||
legend_handles.append(Line2D(
|
||||
[], [], color="#5f0f40", marker="s", linestyle="",
|
||||
markersize=8, label=f"EP={cfg.topo.ep}",
|
||||
))
|
||||
if sips_used > 1:
|
||||
_topo_label = {
|
||||
"ring": "Inter-SIP: ring (chain + wrap)",
|
||||
"mesh2d": "Inter-SIP: 2D mesh (no wrap)",
|
||||
"torus2d": "Inter-SIP: 2D torus (grid + wrap)",
|
||||
}.get(sip_topo, f"Inter-SIP: {sip_topo}")
|
||||
legend_handles.append(Line2D(
|
||||
[], [], color="#c62828", linewidth=2.0,
|
||||
label=_topo_label,
|
||||
))
|
||||
# Legend below the figure; wrap ncols so it's readable
|
||||
n_items = len(legend_handles)
|
||||
ncol = min(6, max(3, n_items))
|
||||
ax.legend(handles=legend_handles, loc="upper center",
|
||||
bbox_to_anchor=(0.5, -0.02), ncol=ncol,
|
||||
fontsize=8, frameon=False)
|
||||
|
||||
title = (
|
||||
f"Layout: {total_pes} PEs = {n_replicas} DP replica(s) x "
|
||||
f"{n_pp_stages} PP stage(s) x {cubes_per_stage} cubes (CP={cfg.topo.cp}) "
|
||||
f"x TP={cfg.topo.tp} | cubes: {cubes_used}, SIPs: {sips_used}"
|
||||
)
|
||||
ax.set_title(title, fontsize=10)
|
||||
|
||||
return fig
|
||||
@@ -1,193 +0,0 @@
|
||||
"""Batch GQA decode sweep — concurrent independent users on ONE SIP.
|
||||
|
||||
Fixes the target to a single 16-cube SIP and scales the batch B = number
|
||||
of concurrent decode users, each with its own (Q, K, V) placed on a
|
||||
disjoint cube group ``[u*C, u*C + C)`` via ``DPPolicy.cube_start`` and
|
||||
addressed by the kernel's ``cube_base`` local-index parameter. All B
|
||||
launches are deferred (``_defer_wait=True``) so they overlap on the shared
|
||||
discrete-event clock; ``run_bench`` drains them together.
|
||||
|
||||
Per mapping the SIP holds ``16 // C`` users at once (A1=2, A2=4, A4=8,
|
||||
B=16). We sweep B up to that capacity and record aggregate latency,
|
||||
throughput (users / latency), and mean per-user latency, to see which
|
||||
mapping's throughput scales best with batch.
|
||||
|
||||
Metric convention matches the corrected short-context sweeps:
|
||||
``wall = max(t_end) - min(t_start)`` over the op log (deploy excluded).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
|
||||
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
|
||||
|
||||
TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
|
||||
|
||||
D_HEAD = 64
|
||||
DTYPE = "f16"
|
||||
TILE_S_KV = 1024
|
||||
P = 8
|
||||
H_KV = 8
|
||||
CUBES_PER_SIP = 16
|
||||
|
||||
# (mode, kv_per_cube, C); users-per-SIP capacity = CUBES_PER_SIP // C.
|
||||
MODES = [("A1", 1, 8), ("A2", 2, 4), ("A4", 4, 2), ("B", 8, 1)]
|
||||
CONTEXT_LENGTHS = [8 * 1024]
|
||||
|
||||
CSV_OUT = (Path(__file__).resolve().parents[2]
|
||||
/ "docs" / "sweeps" / "decode_batch_sweep.csv")
|
||||
|
||||
|
||||
def _ccl_cfg():
|
||||
return resolve_algorithm_config(
|
||||
load_ccl_config(), name="lrab_hierarchical_allreduce",
|
||||
)
|
||||
|
||||
|
||||
def _engine_factory(t, d):
|
||||
return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
|
||||
|
||||
|
||||
def _run_batch(*, kv_per_cube: int, C: int, S_kv: int, B: int, h_q: int = 64):
|
||||
"""Launch B concurrent decode users on disjoint cube groups of SIP 0."""
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import (
|
||||
_validate_config as _validate_decode_config,
|
||||
gqa_attention_decode_short_kernel,
|
||||
)
|
||||
T_q = 1
|
||||
_validate_decode_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
|
||||
h_q=h_q, h_kv=H_KV, S_kv=S_kv)
|
||||
Q_ROWS = kv_per_cube * T_q
|
||||
Q_COLS = (h_q * D_HEAD) // kv_per_cube
|
||||
k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
def _bench_fn(ctx):
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
# Deploy ALL users' tensors first. Each deploy blocks (advances the
|
||||
# clock), so if a launch were already submitted a later deploy would
|
||||
# drive it to completion and serialize the batch. Creating every
|
||||
# tensor before any launch keeps the deploys from touching kernels.
|
||||
users = []
|
||||
for u in range(B):
|
||||
cs = u * C
|
||||
dp_q = DPPolicy(cube="column_wise", pe="replicate",
|
||||
num_cubes=C, num_pes=P, cube_start=cs)
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
|
||||
num_cubes=C, num_pes=P, cube_start=cs)
|
||||
dp_o = DPPolicy(cube="column_wise", pe="replicate",
|
||||
num_cubes=C, num_pes=P, cube_start=cs)
|
||||
q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
|
||||
name=f"q_u{u}")
|
||||
k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"k_u{u}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"v_u{u}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_u{u}")
|
||||
users.append((u, cs, q, k, v, o))
|
||||
# Now submit every launch deferred (no blocking wait between them),
|
||||
# then drain together so all B run concurrently on the shared clock.
|
||||
pending = []
|
||||
for u, cs, q, k, v, o in users:
|
||||
pending += ctx.launch(
|
||||
f"gqa_decode_u{u}",
|
||||
gqa_attention_decode_short_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube, cs,
|
||||
_auto_dim_remap=False,
|
||||
_defer_wait=True,
|
||||
)
|
||||
for h, _sip_id, meta in pending:
|
||||
ctx.wait(h, _meta=meta)
|
||||
|
||||
return run_bench(
|
||||
topology=topo, bench_fn=_bench_fn,
|
||||
device=resolve_device("sip:0"), engine_factory=_engine_factory,
|
||||
)
|
||||
|
||||
|
||||
def _cube_of(rec) -> int:
|
||||
"""Physical cube index of an op record, or -1 if not a PE op."""
|
||||
cid = rec.component_id or ""
|
||||
if ".cube" in cid:
|
||||
try:
|
||||
return int(cid.split(".cube")[1].split(".")[0])
|
||||
except (ValueError, IndexError):
|
||||
return -1
|
||||
return -1
|
||||
|
||||
|
||||
def _metrics(r, *, mode, kv_per_cube, C, S_kv, B):
|
||||
op_log = r.engine.op_log
|
||||
wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
|
||||
|
||||
# Per-user latency: group ops by which user's cube band they touch.
|
||||
per_user = []
|
||||
for u in range(B):
|
||||
lo, hi = u * C, u * C + C
|
||||
starts = [rec.t_start for rec in op_log if lo <= _cube_of(rec) < hi]
|
||||
ends = [rec.t_end for rec in op_log if lo <= _cube_of(rec) < hi]
|
||||
if starts and ends:
|
||||
per_user.append(max(ends) - min(starts))
|
||||
mean_user_ns = sum(per_user) / len(per_user) if per_user else 0.0
|
||||
|
||||
agg_us = wall_ns / 1000
|
||||
return {
|
||||
"mode": mode,
|
||||
"kv_per_cube": kv_per_cube,
|
||||
"C": C,
|
||||
"S_kv": S_kv,
|
||||
"B": B,
|
||||
"capacity": CUBES_PER_SIP // C,
|
||||
"agg_latency_us": agg_us,
|
||||
"mean_user_latency_us": mean_user_ns / 1000,
|
||||
"throughput_users_per_us": (B / agg_us) if agg_us else 0.0,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_decode_batch_sweep():
|
||||
"""Sweep 4 modes × contexts × B ∈ {1, 2, capacity}; dump batch CSV.
|
||||
|
||||
Throughput is near-linear in B (aggregate latency stays within a few
|
||||
percent of the single-user latency across the batch), so the endpoints
|
||||
plus a low point capture the scaling; the high-B runs of the dense
|
||||
mappings are the expensive ones and full 1..capacity granularity is not
|
||||
needed for the trade-off.
|
||||
"""
|
||||
rows = []
|
||||
for mode, kv_per_cube, C in MODES:
|
||||
capacity = CUBES_PER_SIP // C
|
||||
batch_sizes = sorted({1, 2, capacity})
|
||||
for S_kv in CONTEXT_LENGTHS:
|
||||
for B in batch_sizes:
|
||||
print(f">>> BATCH {mode} S_kv={S_kv//1024}K B={B}/{capacity}",
|
||||
flush=True)
|
||||
r = _run_batch(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, B=B)
|
||||
assert r.completion.ok, (
|
||||
f"BATCH {mode} S_kv={S_kv} B={B}: {r.completion}"
|
||||
)
|
||||
m = _metrics(r, mode=mode, kv_per_cube=kv_per_cube,
|
||||
C=C, S_kv=S_kv, B=B)
|
||||
rows.append(m)
|
||||
print(f" agg={m['agg_latency_us']:.2f}us "
|
||||
f"user={m['mean_user_latency_us']:.2f}us "
|
||||
f"tput={m['throughput_users_per_us']:.3f} u/us",
|
||||
flush=True)
|
||||
|
||||
CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
|
||||
with CSV_OUT.open("w", newline="") as f:
|
||||
w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
w.writeheader()
|
||||
w.writerows(rows)
|
||||
print(f"\nWrote {CSV_OUT} ({len(rows)} rows)")
|
||||
@@ -1,87 +0,0 @@
|
||||
"""Command-count checks for the Case-6 primitive-TILED (streamed) decode kernel.
|
||||
|
||||
The tiled kernel hand-blocks each Q·Kᵀ / P·V matmul into 16×16×16
|
||||
GEMMs, with per-block DMAs of the HBM operand slices and a deferred
|
||||
K-inner sum outside the K loop (ADR-0064 Rev2 D8 single-op FIXED path
|
||||
exercised across DMA, GEMM, and MATH engines).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from math import ceil
|
||||
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import ( # noqa: E501
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _primitive,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16 import ( # noqa: E501
|
||||
gqa_attention_decode_long_ctx_cube_sp_pe_sp_hand_tiled_16x16x16_kernel as _tiled,
|
||||
)
|
||||
from kernbench.common.pe_commands import (
|
||||
DmaReadCmd, GemmCmd, PeCpuOverheadCmd,
|
||||
)
|
||||
from kernbench.common.pe_cost_model import DEFAULT_PE_COST_MODEL
|
||||
from kernbench.triton_emu.tl_context import TLContext, run_kernel
|
||||
|
||||
_C, _P = 8, 8
|
||||
_S_KV = 131_072
|
||||
_D_HEAD, _H_Q, _H_KV, _T_Q = 128, 8, 1, 1
|
||||
_TILE_S_KV = 1024
|
||||
_MAC = 16
|
||||
|
||||
|
||||
def _run(kernel, S_kv: int):
|
||||
tl = TLContext(
|
||||
pe_id=0, num_programs=_P, cost_model=DEFAULT_PE_COST_MODEL,
|
||||
cube_id=6, num_cubes=_C, scratch_base=0x200000, scratch_size=1 << 20,
|
||||
)
|
||||
run_kernel(
|
||||
kernel, tl, 0x1000, 0x2000, 0x3000, 0x4000,
|
||||
_T_Q, S_kv, _H_Q, _H_KV, _D_HEAD, _C, _P,
|
||||
)
|
||||
return tl.commands
|
||||
|
||||
|
||||
def _block_counts():
|
||||
"""Return (n_tiles, qk_blocks, pv_blocks) for the fixed test params."""
|
||||
G = _H_Q // _H_KV
|
||||
S_local = _S_KV // (_C * _P)
|
||||
n_tiles = (S_local + _TILE_S_KV - 1) // _TILE_S_KV
|
||||
tile_s = min(_TILE_S_KV, S_local)
|
||||
M = G * _T_Q
|
||||
qk = ceil(M / _MAC) * ceil(tile_s / _MAC) * ceil(_D_HEAD / _MAC)
|
||||
pv = ceil(M / _MAC) * ceil(_D_HEAD / _MAC) * ceil(tile_s / _MAC)
|
||||
return n_tiles, qk, pv
|
||||
|
||||
|
||||
def test_tiled_gemm_count_matches_block_formula():
|
||||
"""One GemmCmd per 16³ block across both matmuls."""
|
||||
n_tiles, qk, pv = _block_counts()
|
||||
expected = n_tiles * (qk + pv)
|
||||
|
||||
cmds = _run(_tiled, _S_KV)
|
||||
got = sum(1 for c in cmds if isinstance(c, GemmCmd))
|
||||
assert got == expected, f"GemmCmd count: got {got}, expected {expected}"
|
||||
|
||||
|
||||
def test_tiled_dma_count_matches_streamed_formula():
|
||||
"""Streamed operand DMAs — Q·Kᵀ blocks pay 2 DMAs (A+B slice), P·V
|
||||
blocks pay 1 DMA (V slice only; the P side is TCM-resident post-softmax).
|
||||
Each matmul also carries full-shape coarse loads for the GemmCmd
|
||||
operand handles (2 for Q·Kᵀ: Q + K_T, 1 for P·V: V only)."""
|
||||
n_tiles, qk, pv = _block_counts()
|
||||
expected = n_tiles * (2 * qk + 1 * pv + 3)
|
||||
|
||||
cmds = _run(_tiled, _S_KV)
|
||||
got = sum(1 for c in cmds if isinstance(c, DmaReadCmd))
|
||||
assert got == expected, f"DmaReadCmd count: got {got}, expected {expected}"
|
||||
|
||||
|
||||
def test_tiled_amplifies_dispatch_vs_primitive():
|
||||
"""The streamed tiled kernel emits ≫40× more PE_CPU dispatches than
|
||||
the coarse primitive — the whole point of adding this variant."""
|
||||
n_tiled = sum(1 for c in _run(_tiled, _S_KV)
|
||||
if isinstance(c, PeCpuOverheadCmd))
|
||||
n_prim = sum(1 for c in _run(_primitive, _S_KV)
|
||||
if isinstance(c, PeCpuOverheadCmd))
|
||||
assert n_tiled > 40 * n_prim, (
|
||||
f"expected >40× dispatch amplification, got {n_tiled} vs {n_prim}"
|
||||
)
|
||||
@@ -1,17 +1,31 @@
|
||||
"""Short-context GQA kernel tests — unified A1/A2/A4/B (ADR-0070).
|
||||
"""Phase 1 spec test for Phase D: short-context GQA kernels
|
||||
(ADR-0060 §B "Items from the long/short context split", item B.split.2).
|
||||
|
||||
Both ``_gqa_attention_prefill_short.py`` and
|
||||
``_gqa_attention_decode_short.py`` are single unified kernels selected
|
||||
at launch via ``kv_per_cube ∈ {1, 2, 4, 8}``:
|
||||
The long-context kernels (§5.2 decode / §5.5 prefill) shard each KV
|
||||
head row-wise across all CUBEs. At short context (S_kv < 256K), that
|
||||
shard is too thin to feed the engines and the cube-level collective
|
||||
overhead dominates. The short-context kernels drop cube-SP entirely:
|
||||
each CUBE owns ``kv_per_cube`` *whole* KV heads, no S_kv sharding across
|
||||
CUBEs.
|
||||
|
||||
Mode kv_per_cube C group_size
|
||||
---- ----------- ------- ----------
|
||||
A1 1 h_kv P (=8)
|
||||
A2 2 h_kv/2 P/2 (=4)
|
||||
A4 4 h_kv/4 P/4 (=2)
|
||||
B 8 1 1
|
||||
Design (per AskUserQuestion answers in this session):
|
||||
- PE-parallel heads: P PEs split into ``kv_per_cube`` groups, each
|
||||
group does PE-SP across (P/kv_per_cube) PEs for one owned head.
|
||||
- scratch_scope + tl.copy_to discipline mirrors the long kernels.
|
||||
- One short kernel handles kv_per_cube ∈ {1, 2, 4, 8} via a parameter.
|
||||
|
||||
Tests are mode-parametrized over the four (kv_per_cube, C) tuples.
|
||||
Group layout on the 2×4 PE grid:
|
||||
kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
|
||||
kv_per_cube=2, group=4 PEs (one row): row chain only.
|
||||
kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
|
||||
kv_per_cube=8, group=1 PE: no chain.
|
||||
|
||||
After chain reduce-to-group-root, the group's root PE writes its
|
||||
owned head's output. No inter-CUBE reduce.
|
||||
|
||||
Phase 1 (this commit): tests only — production code lands in Phase 2.
|
||||
All tests fail because the short kernels (``_gqa_attention_decode_short.py`` and
|
||||
``_gqa_attention_prefill_short.py``) do not exist yet.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -31,17 +45,6 @@ TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
|
||||
|
||||
D_HEAD = 64
|
||||
DTYPE = "f16"
|
||||
TILE_S_KV = 1024
|
||||
P = 8
|
||||
H_KV = 8
|
||||
|
||||
# (kv_per_cube, C) for the four modes.
|
||||
MODES = [
|
||||
pytest.param(1, 8, id="A1"),
|
||||
pytest.param(2, 4, id="A2"),
|
||||
pytest.param(4, 2, id="A4"),
|
||||
pytest.param(8, 1, id="B"),
|
||||
]
|
||||
|
||||
|
||||
def _ccl_cfg():
|
||||
@@ -58,161 +61,46 @@ def _count(op_log, name: str) -> int:
|
||||
return sum(1 for r in op_log if r.op_name == name)
|
||||
|
||||
|
||||
# ── Prefill ──────────────────────────────────────────────────────────
|
||||
# ── Decode short-context kernel ──────────────────────────────────────
|
||||
|
||||
|
||||
def _run_prefill(*, kv_per_cube: int, C: int, T_q: int, S_kv: int,
|
||||
h_q: int = 8):
|
||||
"""Run the unified prefill kernel in the given mode."""
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import (
|
||||
_validate_config as _validate_prefill_config,
|
||||
gqa_attention_prefill_short_kernel,
|
||||
)
|
||||
_validate_prefill_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
|
||||
h_q=h_q, h_kv=H_KV, S_kv=S_kv)
|
||||
n_tiles = S_kv // TILE_S_KV
|
||||
Q_ROWS = kv_per_cube * T_q
|
||||
Q_COLS = (h_q * D_HEAD) // kv_per_cube
|
||||
def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int,
|
||||
C: int, P: int, S_kv: int):
|
||||
"""Run the short-context decode kernel with PE-parallel heads.
|
||||
|
||||
Layout (after design iteration — see Phase D failure-recovery):
|
||||
Q: (T_q, h_q·D_HEAD) replicated; kernel reshapes byte-conservingly.
|
||||
K, V: (h_kv·S_kv, D_HEAD) head-stacked, ``cube=row_wise, pe=row_wise``
|
||||
so each PE gets a contiguous (S_local, D_HEAD) chunk at its own
|
||||
addressable shard (no partial reads needed).
|
||||
O: replicated; each group root writes the full byte-conserving
|
||||
(h_q·T_q, D_HEAD) result.
|
||||
"""
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import gqa_attention_decode_short_kernel # Phase 2
|
||||
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
def _bench_fn(ctx):
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
|
||||
name=f"q_pre_kv{kv_per_cube}")
|
||||
k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"k_pre_kv{kv_per_cube}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_pre_kv{kv_per_cube}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_pre_kv{kv_per_cube}")
|
||||
dp_full = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=C, num_pes=P)
|
||||
# Head-stacked KV with row_wise sharding so each PE's chunk is
|
||||
# contiguous and exactly (S_local, D_HEAD) addressable.
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
|
||||
num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((1, h_q * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_full, name=f"q_short_{kv_per_cube}_{C}")
|
||||
k = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"k_short_{kv_per_cube}_{C}")
|
||||
v = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_short_{kv_per_cube}_{C}")
|
||||
o = ctx.empty((1, h_q * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}")
|
||||
ctx.launch(
|
||||
f"gqa_prefill_kv{kv_per_cube}",
|
||||
gqa_attention_prefill_short_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
return run_bench(
|
||||
topology=topo, bench_fn=_bench_fn,
|
||||
device=resolve_device(None), engine_factory=_engine_factory,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_smoke(kv_per_cube, C):
|
||||
"""All four prefill modes complete on a single-tile mini config."""
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok, (
|
||||
f"prefill kv_per_cube={kv_per_cube}: {r.completion}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_qtile_split_dma_writes(kv_per_cube, C):
|
||||
"""Every PE in every active group stores its q-tile output:
|
||||
dma_writes = group_size · kv_per_cube · C."""
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok
|
||||
group_size = P // kv_per_cube
|
||||
expected = group_size * kv_per_cube * C
|
||||
n_writes = _count(r.engine.op_log, "dma_write")
|
||||
assert n_writes == expected, (
|
||||
f"prefill kv_per_cube={kv_per_cube}: expected {expected} "
|
||||
f"dma_writes; got {n_writes}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_kv_hbm_read_collapsed(kv_per_cube, C):
|
||||
"""IPCQ KV broadcast collapses HBM K/V reads to one per group per tile:
|
||||
dma_reads = (P · C) Q-reads + (kv_per_cube · 2 · n_tiles · C) KV-reads.
|
||||
"""
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok
|
||||
n_tiles = 1024 // TILE_S_KV
|
||||
expected = (P * C) + (kv_per_cube * 2 * n_tiles * C)
|
||||
n_reads = _count(r.engine.op_log, "dma_read")
|
||||
assert n_reads == expected, (
|
||||
f"prefill kv_per_cube={kv_per_cube}: expected {expected} dma_reads "
|
||||
f"(Q + IPCQ-broadcasted KV); got {n_reads}"
|
||||
)
|
||||
|
||||
|
||||
def test_prefill_no_ring_KV_traffic():
|
||||
"""Short prefill never uses Ring KV — no inter-CUBE E/W IPCQ."""
|
||||
r = _run_prefill(kv_per_cube=2, C=4, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok
|
||||
inter_cube_ipcq = sum(
|
||||
1 for rec in r.engine.op_log
|
||||
if rec.op_name == "ipcq_copy"
|
||||
and rec.params.get("direction") in ("E", "W")
|
||||
)
|
||||
assert inter_cube_ipcq == 0, (
|
||||
f"short prefill must have no inter-CUBE E/W IPCQ; "
|
||||
f"got {inter_cube_ipcq}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_multitile_scaling(kv_per_cube, C):
|
||||
"""Prefill dma_read scales linearly with n_tiles in every mode.
|
||||
|
||||
Per-mode formula (IPCQ broadcast: cube's group-root PE reads K/V):
|
||||
dma_reads = P·C + 2·kv_per_cube·C·n_tiles
|
||||
= (Q: 1 per PE) + (K + V: 1 per group per tile)
|
||||
"""
|
||||
for S_kv in (1024, 2048, 4096):
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=S_kv)
|
||||
assert r.completion.ok, f"kv={kv_per_cube} S_kv={S_kv}"
|
||||
n_tiles = S_kv // TILE_S_KV
|
||||
expected = (P * C) + (2 * kv_per_cube * C * n_tiles)
|
||||
n_reads = _count(r.engine.op_log, "dma_read")
|
||||
assert n_reads == expected, (
|
||||
f"prefill kv={kv_per_cube} n_tiles={n_tiles}: expected "
|
||||
f"{expected} dma_reads; got {n_reads}"
|
||||
)
|
||||
|
||||
|
||||
# ── Decode ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _run_decode(*, kv_per_cube: int, C: int, S_kv: int, h_q: int = 8):
|
||||
"""Run the unified decode kernel in the given mode."""
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import (
|
||||
_validate_config as _validate_decode_config,
|
||||
gqa_attention_decode_short_kernel,
|
||||
)
|
||||
T_q = 1
|
||||
_validate_decode_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
|
||||
h_q=h_q, h_kv=H_KV, S_kv=S_kv)
|
||||
Q_ROWS = kv_per_cube * T_q
|
||||
Q_COLS = (h_q * D_HEAD) // kv_per_cube
|
||||
# K total elements = h_kv · S_kv · d_head; mode-invariant deploy.
|
||||
k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
def _bench_fn(ctx):
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P)
|
||||
dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
|
||||
name=f"q_dec_kv{kv_per_cube}")
|
||||
k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"k_dec_kv{kv_per_cube}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"v_dec_kv{kv_per_cube}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_dec_kv{kv_per_cube}")
|
||||
ctx.launch(
|
||||
f"gqa_decode_kv{kv_per_cube}",
|
||||
f"gqa_decode_short_{kv_per_cube}_{C}",
|
||||
gqa_attention_decode_short_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
|
||||
1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
@@ -222,200 +110,117 @@ def _run_decode(*, kv_per_cube: int, C: int, S_kv: int, h_q: int = 8):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_smoke(kv_per_cube, C):
|
||||
"""All four decode modes complete at S_kv = 8K (min multi-tile-aligned
|
||||
config that satisfies every mode's group-size constraint)."""
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
|
||||
assert r.completion.ok, (
|
||||
f"decode kv_per_cube={kv_per_cube}: {r.completion}"
|
||||
def test_short_decode_smoke_kv_per_cube_2_C_4():
|
||||
"""ADR-0060 §B.split.2 headline: kv_per_cube=2, C=4 — each CUBE
|
||||
owns 2 heads (8 heads / 4 CUBEs). PE-parallel heads splits the 8
|
||||
PEs into 2 groups of 4, each group does PE-SP for one owned head.
|
||||
Smoke: kernel completes."""
|
||||
result = _run_decode_short(
|
||||
h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok, (
|
||||
f"short decode kv_per_cube=2 must complete; got {result.completion}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_chain_reduce_store_count(kv_per_cube, C):
|
||||
"""Decode chain-reduce: only group root (PE 0 per group) stores —
|
||||
dma_writes = kv_per_cube · C (one per group root × cubes)."""
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
|
||||
assert r.completion.ok
|
||||
expected = kv_per_cube * C
|
||||
n_writes = _count(r.engine.op_log, "dma_write")
|
||||
assert n_writes == expected, (
|
||||
f"decode kv_per_cube={kv_per_cube}: expected {expected} dma_writes "
|
||||
f"(1 per group root); got {n_writes}"
|
||||
def test_short_decode_smoke_kv_per_cube_4_C_2():
|
||||
"""kv_per_cube=4, C=2 — half the CUBEs participate, each owns 4
|
||||
heads, 4 PE groups of 2 PEs each."""
|
||||
result = _run_decode_short(
|
||||
h_q=8, h_kv=8, kv_per_cube=4, C=2, P=8, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok, (
|
||||
f"short decode kv_per_cube=4 must complete; got {result.completion}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_multitile_per_pe(kv_per_cube, C):
|
||||
"""Decode multi-tile sweep path exercised in every mode.
|
||||
def test_short_decode_smoke_kv_per_cube_8_C_1():
|
||||
"""kv_per_cube=8, C=1 — all heads on one CUBE. 8 PE groups of 1 PE
|
||||
each → no chain reduce, each PE writes its head's output."""
|
||||
result = _run_decode_short(
|
||||
h_q=8, h_kv=8, kv_per_cube=8, C=1, P=8, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok, (
|
||||
f"short decode kv_per_cube=8 must complete; got {result.completion}"
|
||||
)
|
||||
|
||||
S_kv is chosen so each PE owns 2 tiles (n_tiles_per_pe == 2),
|
||||
forcing the post-tile-0 sweep loop to run.
|
||||
|
||||
def test_short_decode_dma_write_count_equals_h_kv():
|
||||
"""ADR-0060 §B.split.2: each owned head produces exactly one
|
||||
output (no inter-CUBE reduce). Total dma_writes = h_kv across all
|
||||
participating CUBEs and groups.
|
||||
|
||||
For h_kv=8, kv_per_cube=2, C=4: each CUBE writes 2 outputs →
|
||||
4 × 2 = 8 dma_writes total.
|
||||
"""
|
||||
group_size = P // kv_per_cube
|
||||
S_kv = group_size * 2 * TILE_S_KV
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv)
|
||||
assert r.completion.ok, (
|
||||
f"decode kv={kv_per_cube} S_kv={S_kv}: {r.completion}"
|
||||
result = _run_decode_short(
|
||||
h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok
|
||||
n_writes = _count(result.engine.op_log, "dma_write")
|
||||
assert n_writes == 8, (
|
||||
f"short decode: expected 8 dma_writes (h_kv); got {n_writes}"
|
||||
)
|
||||
|
||||
|
||||
# ── ADR-0011 D-VA1 contract regression tests ───────────────────────
|
||||
def test_short_decode_no_inter_cube_traffic():
|
||||
"""ADR-0060 §B.split.2: each head is fully owned by one CUBE → no
|
||||
inter-CUBE reduce. The kernel must not invoke CUBE-level E/W IPCQ.
|
||||
|
||||
|
||||
def _per_cube_dma_busy(op_log):
|
||||
"""Sum of dma_read t_end-t_start grouped by cube index."""
|
||||
from collections import defaultdict
|
||||
busy = defaultdict(float)
|
||||
for rec in op_log:
|
||||
if rec.op_name != "dma_read":
|
||||
continue
|
||||
cid = rec.component_id or ""
|
||||
for part in cid.split("."):
|
||||
if part.startswith("cube"):
|
||||
try:
|
||||
busy[int(part[4:])] += rec.t_end - rec.t_start
|
||||
except ValueError:
|
||||
pass
|
||||
break
|
||||
return dict(busy)
|
||||
|
||||
|
||||
def _per_cube_large_read_addrs(op_log, *, min_nbytes: int = 1024):
|
||||
"""Group dma_read src_addrs (>= min_nbytes) by issuing cube."""
|
||||
from collections import defaultdict
|
||||
per_cube: dict[int, set[int]] = defaultdict(set)
|
||||
for rec in op_log:
|
||||
if rec.op_name != "dma_read":
|
||||
continue
|
||||
if rec.params.get("nbytes", 0) < min_nbytes:
|
||||
continue
|
||||
cid = rec.component_id or ""
|
||||
for part in cid.split("."):
|
||||
if part.startswith("cube"):
|
||||
try:
|
||||
per_cube[int(part[4:])].add(rec.params["src_addr"])
|
||||
except ValueError:
|
||||
pass
|
||||
break
|
||||
return per_cube
|
||||
|
||||
|
||||
def _assert_cubes_disjoint(per_cube_addrs, label):
|
||||
seen = sorted(per_cube_addrs.items())
|
||||
for i, (ci, ai) in enumerate(seen):
|
||||
for cj, aj in seen[i + 1:]:
|
||||
assert ai.isdisjoint(aj), (
|
||||
f"{label}: cube{ci} and cube{cj} share {len(ai & aj)} "
|
||||
f"dma_read src_addrs; cubes must target disjoint HBM regions."
|
||||
)
|
||||
|
||||
|
||||
def _assert_cube_balanced(busy, label, max_ratio=1.1):
|
||||
if len(busy) <= 1:
|
||||
return # single-cube modes auto-pass
|
||||
vals = list(busy.values())
|
||||
ratio = max(vals) / min(vals) if min(vals) > 0 else 0
|
||||
assert ratio < max_ratio, (
|
||||
f"{label}: per-cube DMA_READ unbalanced: max/min = {ratio:.2f}× "
|
||||
f"(expected < {max_ratio}×). Per-cube busy (μs): "
|
||||
f"{[f'cube{c}={v/1000:.1f}' for c,v in sorted(busy.items())]}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_per_cube_disjoint_src_addrs(kv_per_cube, C):
|
||||
"""Decode: cubes must read from disjoint HBM regions (per ADR-0011
|
||||
D-VA1, kernel computes cube/PE shard offset from program_id).
|
||||
|
||||
Regression guard: pre-fix all 64 PEs read offset 0 → all cubes
|
||||
share the same src_addrs.
|
||||
Today's long-context kernel at C=4 emits ~12 inter-CUBE ipcq_copy
|
||||
via direction "E"/"W". The short kernel must emit zero of those,
|
||||
keeping all IPCQ traffic on the ``intra_*`` directions.
|
||||
"""
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
|
||||
assert r.completion.ok
|
||||
per_cube = _per_cube_large_read_addrs(r.engine.op_log)
|
||||
assert len(per_cube) == C, (
|
||||
f"decode kv={kv_per_cube}: expected {C} cubes; got {len(per_cube)}"
|
||||
result = _run_decode_short(
|
||||
h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
|
||||
)
|
||||
assert result.completion.ok
|
||||
# Count IPCQ traffic that targeted the CUBE-level "E"/"W" directions.
|
||||
inter_cube_ipcq = sum(
|
||||
1 for r in result.engine.op_log
|
||||
if r.op_name == "ipcq_copy"
|
||||
and r.params.get("direction") in ("E", "W")
|
||||
)
|
||||
assert inter_cube_ipcq == 0, (
|
||||
f"short decode must have no inter-CUBE E/W IPCQ; got {inter_cube_ipcq}"
|
||||
)
|
||||
_assert_cubes_disjoint(per_cube, f"decode kv={kv_per_cube}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_per_cube_dma_balanced(kv_per_cube, C):
|
||||
"""Decode: per-cube DMA_READ busy must be ~equal (cubes operate on
|
||||
independent local HBM). max/min < 1.1× for multi-cube modes.
|
||||
# ── Prefill short-context kernel ─────────────────────────────────────
|
||||
|
||||
Regression guard: pre-fix 11.5× ratio (cross-cube traffic to cube0).
|
||||
|
||||
def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
|
||||
C: int, P: int, T_q: int, S_kv: int):
|
||||
"""Run the short-context prefill kernel with PE-parallel heads.
|
||||
|
||||
Layout: same head-stacked K/V scheme as decode short, with Q/O
|
||||
replicated and byte-conserving reshape inside the kernel.
|
||||
"""
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
|
||||
assert r.completion.ok
|
||||
busy = _per_cube_dma_busy(r.engine.op_log)
|
||||
assert len(busy) == C, (
|
||||
f"decode kv={kv_per_cube}: expected {C} cubes; got {sorted(busy.keys())}"
|
||||
)
|
||||
_assert_cube_balanced(busy, f"decode kv={kv_per_cube}")
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel # Phase 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_per_cube_disjoint_src_addrs(kv_per_cube, C):
|
||||
"""Prefill: cubes must read from disjoint HBM regions."""
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok
|
||||
per_cube = _per_cube_large_read_addrs(r.engine.op_log)
|
||||
assert len(per_cube) == C, (
|
||||
f"prefill kv={kv_per_cube}: expected {C} cubes; got {len(per_cube)}"
|
||||
)
|
||||
_assert_cubes_disjoint(per_cube, f"prefill kv={kv_per_cube}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_per_cube_dma_balanced(kv_per_cube, C):
|
||||
"""Prefill: per-cube DMA_READ busy must be ~equal (cube-parallel)."""
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok
|
||||
busy = _per_cube_dma_busy(r.engine.op_log)
|
||||
assert len(busy) == C, (
|
||||
f"prefill kv={kv_per_cube}: expected {C} cubes; got {sorted(busy.keys())}"
|
||||
)
|
||||
_assert_cube_balanced(busy, f"prefill kv={kv_per_cube}")
|
||||
|
||||
|
||||
# ── Composite (second-level) variant smoke tests ────────────────────
|
||||
|
||||
|
||||
def _run_prefill_composite(*, kv_per_cube: int, C: int, T_q: int,
|
||||
S_kv: int, h_q: int = 8):
|
||||
"""Same helper as _run_prefill but for the composite-GEMM kernel."""
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short_composite import (
|
||||
_validate_config as _validate_prefill_cmp_config,
|
||||
gqa_attention_prefill_short_composite_kernel,
|
||||
)
|
||||
_validate_prefill_cmp_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
|
||||
C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
|
||||
n_tiles = S_kv // TILE_S_KV
|
||||
Q_ROWS = kv_per_cube * T_q
|
||||
Q_COLS = (h_q * D_HEAD) // kv_per_cube
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
def _bench_fn(ctx):
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
|
||||
name=f"q_pre_cmp_kv{kv_per_cube}")
|
||||
k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"k_pre_cmp_kv{kv_per_cube}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_pre_cmp_kv{kv_per_cube}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_pre_cmp_kv{kv_per_cube}")
|
||||
dp_q = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
|
||||
num_cubes=C, num_pes=P)
|
||||
dp_o = DPPolicy(cube="replicate", pe="replicate",
|
||||
num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((T_q, h_kv * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_q, name=f"q_pre_short_{kv_per_cube}_{C}")
|
||||
k = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"k_pre_short_{kv_per_cube}_{C}")
|
||||
v = ctx.zeros((h_kv * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_pre_short_{kv_per_cube}_{C}")
|
||||
o = ctx.empty((T_q, h_kv * D_HEAD),
|
||||
dtype=DTYPE, dp=dp_o, name=f"o_pre_short_{kv_per_cube}_{C}")
|
||||
ctx.launch(
|
||||
f"gqa_prefill_cmp_kv{kv_per_cube}",
|
||||
gqa_attention_prefill_short_composite_kernel,
|
||||
f"gqa_prefill_short_{kv_per_cube}_{C}",
|
||||
gqa_attention_prefill_short_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
|
||||
T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
@@ -425,184 +230,31 @@ def _run_prefill_composite(*, kv_per_cube: int, C: int, T_q: int,
|
||||
)
|
||||
|
||||
|
||||
def _run_decode_composite(*, kv_per_cube: int, C: int, S_kv: int,
|
||||
h_q: int = 8):
|
||||
"""Same helper as _run_decode but for the composite-GEMM kernel."""
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short_composite import (
|
||||
_validate_config as _validate_decode_cmp_config,
|
||||
gqa_attention_decode_short_composite_kernel,
|
||||
def test_short_prefill_smoke_kv_per_cube_2_C_4():
|
||||
"""Short prefill kv_per_cube=2, C=4. Smoke: completes."""
|
||||
result = _run_prefill_short(
|
||||
h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
|
||||
)
|
||||
T_q = 1
|
||||
_validate_decode_cmp_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
|
||||
C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
|
||||
Q_ROWS = kv_per_cube * T_q
|
||||
Q_COLS = (h_q * D_HEAD) // kv_per_cube
|
||||
k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
def _bench_fn(ctx):
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P)
|
||||
dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
|
||||
name=f"q_dec_cmp_kv{kv_per_cube}")
|
||||
k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"k_dec_cmp_kv{kv_per_cube}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"v_dec_cmp_kv{kv_per_cube}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_dec_cmp_kv{kv_per_cube}")
|
||||
ctx.launch(
|
||||
f"gqa_decode_cmp_kv{kv_per_cube}",
|
||||
gqa_attention_decode_short_composite_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
return run_bench(
|
||||
topology=topo, bench_fn=_bench_fn,
|
||||
device=resolve_device(None), engine_factory=_engine_factory,
|
||||
assert result.completion.ok, (
|
||||
f"short prefill kv_per_cube=2 must complete; got {result.completion}"
|
||||
)
|
||||
|
||||
|
||||
def _run_prefill_composite_fused(*, kv_per_cube: int, C: int, T_q: int,
|
||||
S_kv: int, h_q: int = 8):
|
||||
"""Variant (3): GEMM-only composite + softmax_merge prologue fusion."""
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short_composite_fused import (
|
||||
_validate_config as _validate_prefill_fuse_config,
|
||||
gqa_attention_prefill_short_composite_fused_kernel,
|
||||
def test_short_prefill_no_ring_KV_traffic():
|
||||
"""ADR-0060 §B.split.2: short prefill DOES NOT use Ring KV — each
|
||||
CUBE owns its KV heads fully, no rotation. The kernel must not
|
||||
emit any KV-rotation IPCQ traffic at the CUBE level.
|
||||
"""
|
||||
result = _run_prefill_short(
|
||||
h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
|
||||
)
|
||||
_validate_prefill_fuse_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
|
||||
C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
|
||||
n_tiles = S_kv // TILE_S_KV
|
||||
Q_ROWS = kv_per_cube * T_q
|
||||
Q_COLS = (h_q * D_HEAD) // kv_per_cube
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
def _bench_fn(ctx):
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
|
||||
name=f"q_pre_fuse_kv{kv_per_cube}")
|
||||
k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"k_pre_fuse_kv{kv_per_cube}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD),
|
||||
dtype=DTYPE, dp=dp_kv, name=f"v_pre_fuse_kv{kv_per_cube}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_pre_fuse_kv{kv_per_cube}")
|
||||
ctx.launch(
|
||||
f"gqa_prefill_fuse_kv{kv_per_cube}",
|
||||
gqa_attention_prefill_short_composite_fused_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
return run_bench(
|
||||
topology=topo, bench_fn=_bench_fn,
|
||||
device=resolve_device(None), engine_factory=_engine_factory,
|
||||
assert result.completion.ok
|
||||
inter_cube_ipcq = sum(
|
||||
1 for r in result.engine.op_log
|
||||
if r.op_name == "ipcq_copy"
|
||||
and r.params.get("direction") in ("E", "W")
|
||||
)
|
||||
|
||||
|
||||
def _run_decode_composite_fused(*, kv_per_cube: int, C: int, S_kv: int,
|
||||
h_q: int = 8):
|
||||
"""Variant (3): GEMM-only composite + softmax_merge prologue fusion."""
|
||||
from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short_composite_fused import (
|
||||
_validate_config as _validate_decode_fuse_config,
|
||||
gqa_attention_decode_short_composite_fused_kernel,
|
||||
)
|
||||
T_q = 1
|
||||
_validate_decode_fuse_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
|
||||
C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
|
||||
Q_ROWS = kv_per_cube * T_q
|
||||
Q_COLS = (h_q * D_HEAD) // kv_per_cube
|
||||
k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
def _bench_fn(ctx):
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P)
|
||||
dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
|
||||
q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
|
||||
name=f"q_dec_fuse_kv{kv_per_cube}")
|
||||
k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"k_dec_fuse_kv{kv_per_cube}")
|
||||
v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
|
||||
name=f"v_dec_fuse_kv{kv_per_cube}")
|
||||
o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
|
||||
name=f"o_dec_fuse_kv{kv_per_cube}")
|
||||
ctx.launch(
|
||||
f"gqa_decode_fuse_kv{kv_per_cube}",
|
||||
gqa_attention_decode_short_composite_fused_kernel,
|
||||
q, k, v, o,
|
||||
T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
return run_bench(
|
||||
topology=topo, bench_fn=_bench_fn,
|
||||
device=resolve_device(None), engine_factory=_engine_factory,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_composite_smoke(kv_per_cube, C):
|
||||
"""Variant (2) GEMM-only composite prefill — all 4 modes."""
|
||||
r = _run_prefill_composite(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
|
||||
assert r.completion.ok, (
|
||||
f"prefill-composite kv_per_cube={kv_per_cube}: {r.completion}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_composite_smoke(kv_per_cube, C):
|
||||
"""Variant (2) GEMM-only composite decode — all 4 modes."""
|
||||
r = _run_decode_composite(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
|
||||
assert r.completion.ok, (
|
||||
f"decode-composite kv_per_cube={kv_per_cube}: {r.completion}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_composite_fused_smoke(kv_per_cube, C):
|
||||
"""Variant (3) composite + softmax_merge fused decode — all 4 modes."""
|
||||
r = _run_decode_composite_fused(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
|
||||
assert r.completion.ok, (
|
||||
f"decode-composite-fused kv_per_cube={kv_per_cube}: {r.completion}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_decode_composite_fused_multitile(kv_per_cube, C):
|
||||
"""Fused decode with n_tiles_per_pe == 2 in EVERY mode, so the tile-1+
|
||||
fused composite loop actually runs.
|
||||
|
||||
The smoke test at S_kv=8192 gives A1 (group_size=8) n_tiles_per_pe=1,
|
||||
which never enters the fused path. S_kv = group_size·2·TILE_S_KV forces
|
||||
2 tiles per PE for all modes (mirrors test_decode_multitile_per_pe)."""
|
||||
group_size = P // kv_per_cube
|
||||
S_kv = group_size * 2 * TILE_S_KV
|
||||
r = _run_decode_composite_fused(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv)
|
||||
assert r.completion.ok, (
|
||||
f"decode-composite-fused-multitile kv={kv_per_cube} S_kv={S_kv}: "
|
||||
f"{r.completion}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("kv_per_cube,C", MODES)
|
||||
def test_prefill_composite_fused_smoke(kv_per_cube, C):
|
||||
"""Variant (3) composite + softmax_merge fused prefill — all 4 modes.
|
||||
|
||||
S_kv=2048 (n_tiles=2) so the tile-1+ fused composite loop actually
|
||||
runs; S_kv=1024 gives n_tiles=1 and never enters the fused path.
|
||||
Multi-cube (A1/A2/A4) works now that recv'd K/V slots are pinned."""
|
||||
r = _run_prefill_composite_fused(kv_per_cube=kv_per_cube, C=C,
|
||||
T_q=8, S_kv=2048)
|
||||
assert r.completion.ok, (
|
||||
f"prefill-composite-fused kv_per_cube={kv_per_cube}: {r.completion}"
|
||||
assert inter_cube_ipcq == 0, (
|
||||
f"short prefill must have no Ring KV (no inter-CUBE E/W IPCQ); "
|
||||
f"got {inter_cube_ipcq}"
|
||||
)
|
||||
|
||||
@@ -1,131 +0,0 @@
|
||||
"""Decode mode-comparison sweep across context lengths.
|
||||
|
||||
Measures per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}):
|
||||
1. Wall clock latency (μs)
|
||||
2. Hardware utilization
|
||||
- GEMM engine utilization (Σ gemm time / wall × n_pe)
|
||||
- MATH op count
|
||||
3. HBM bandwidth utilization
|
||||
(Σ dma_read+write bytes / (wall × n_pe × peak_pe_hbm_bw))
|
||||
4. Communication cost (IPCQ bytes — chain reduce)
|
||||
5. KV cache space cost (per-cube bytes)
|
||||
|
||||
Output: ``docs/sweeps/short_context_decode_sweep.csv`` (16 rows).
|
||||
|
||||
Run: ``pytest tests/attention/test_gqa_short_context_sweep_decode.py -m slow``
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.attention.test_gqa_short_context import (
|
||||
D_HEAD, H_KV, P, TILE_S_KV, _run_decode,
|
||||
)
|
||||
|
||||
PEAK_PE_HBM_BPS = 8 * 32.0 # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
|
||||
|
||||
MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
|
||||
"maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
|
||||
"max_elem"}
|
||||
|
||||
MODES = [
|
||||
("A1", 1, 8),
|
||||
("A2", 2, 4),
|
||||
("A4", 4, 2),
|
||||
("B", 8, 1),
|
||||
]
|
||||
|
||||
CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
|
||||
|
||||
CSV_OUT = (Path(__file__).resolve().parents[2]
|
||||
/ "docs" / "sweeps" / "short_context_decode_sweep.csv")
|
||||
|
||||
|
||||
def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
|
||||
op_log = r.engine.op_log
|
||||
wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
|
||||
|
||||
active_pes = set()
|
||||
for rec in op_log:
|
||||
cid = rec.component_id or ""
|
||||
if cid and ".cube" in cid and ".pe" in cid:
|
||||
parts = cid.split(".")
|
||||
active_pes.add(".".join(parts[:3]))
|
||||
n_pe = len(active_pes)
|
||||
|
||||
gemm_time_ns = 0.0
|
||||
math_count = 0
|
||||
dma_read_bytes = 0
|
||||
dma_write_bytes = 0
|
||||
ipcq_bytes = 0
|
||||
for rec in op_log:
|
||||
name = rec.op_name
|
||||
nbytes = rec.params.get("nbytes", 0) or 0
|
||||
if name == "gemm_f16":
|
||||
gemm_time_ns += rec.t_end - rec.t_start
|
||||
elif name in MATH_OPS:
|
||||
math_count += 1
|
||||
elif name == "dma_read":
|
||||
dma_read_bytes += nbytes
|
||||
elif name == "dma_write":
|
||||
dma_write_bytes += nbytes
|
||||
elif name == "ipcq_copy":
|
||||
ipcq_bytes += nbytes
|
||||
|
||||
gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
|
||||
hbm_bw_util = (
|
||||
(dma_read_bytes + dma_write_bytes)
|
||||
/ (wall_ns * n_pe * PEAK_PE_HBM_BPS)
|
||||
) if n_pe else 0.0
|
||||
|
||||
kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
|
||||
|
||||
return {
|
||||
"mode": mode,
|
||||
"kv_per_cube": kv_per_cube,
|
||||
"C": C,
|
||||
"S_kv": S_kv,
|
||||
"wall_us": wall_ns / 1000,
|
||||
"n_pe": n_pe,
|
||||
"gemm_util": gemm_util,
|
||||
"math_count": math_count,
|
||||
"hbm_bw_util": hbm_bw_util,
|
||||
"hbm_read_mb": dma_read_bytes / (1024 * 1024),
|
||||
"hbm_write_kb": dma_write_bytes / 1024,
|
||||
"ipcq_kb": ipcq_bytes / 1024,
|
||||
"kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_decode_mode_context_sweep():
|
||||
"""Run 4 modes × 4 context lengths, dump 5 metrics to CSV."""
|
||||
rows = []
|
||||
for mode, kv_per_cube, C in MODES:
|
||||
for S_kv in CONTEXT_LENGTHS:
|
||||
print(f">>> DECODE {mode} S_kv={S_kv//1024}K "
|
||||
f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
|
||||
r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, h_q=64)
|
||||
assert r.completion.ok, (
|
||||
f"DECODE {mode} S_kv={S_kv}: {r.completion}"
|
||||
)
|
||||
m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
|
||||
C=C, S_kv=S_kv)
|
||||
rows.append(m)
|
||||
print(f" wall={m['wall_us']:.1f}μs "
|
||||
f"gemm_util={m['gemm_util']:.3f} "
|
||||
f"hbm_bw_util={m['hbm_bw_util']:.3f} "
|
||||
f"ipcq={m['ipcq_kb']:.1f}KB "
|
||||
f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
|
||||
flush=True)
|
||||
|
||||
CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
|
||||
with CSV_OUT.open("w", newline="") as f:
|
||||
w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
w.writeheader()
|
||||
for row in rows:
|
||||
w.writerow(row)
|
||||
print(f"\nWrote {len(rows)} rows to {CSV_OUT}")
|
||||
@@ -1,143 +0,0 @@
|
||||
"""Decode (variant 2) GEMM-only composite mode-comparison sweep.
|
||||
|
||||
Per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}) the
|
||||
same 5 metrics as the first-level sweep are collected:
|
||||
1. Wall clock latency (μs)
|
||||
2. Hardware utilization (GEMM util, MATH op count)
|
||||
3. HBM bandwidth utilization
|
||||
4. Communication cost (IPCQ bytes — chain reduce)
|
||||
5. KV cache space cost (per-cube bytes)
|
||||
|
||||
Kernel: ``_gqa_attention_decode_short_composite.py`` (variant 2,
|
||||
GEMM-only composite; Q·Kᵀ via tl.composite, P·V via tl.dot).
|
||||
|
||||
Output: ``docs/sweeps/short_context_decode_composite_sweep.csv``.
|
||||
|
||||
Run: ``pytest tests/attention/test_gqa_short_context_sweep_decode_composite.py -m slow``
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.attention.test_gqa_short_context import (
|
||||
D_HEAD, H_KV, P, TILE_S_KV, _run_decode_composite,
|
||||
)
|
||||
|
||||
PEAK_PE_HBM_BPS = 8 * 32.0 # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
|
||||
|
||||
MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
|
||||
"maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
|
||||
"max_elem"}
|
||||
|
||||
MODES = [
|
||||
("A1", 1, 8),
|
||||
("A2", 2, 4),
|
||||
("A4", 4, 2),
|
||||
("B", 8, 1),
|
||||
]
|
||||
|
||||
CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
|
||||
|
||||
CSV_OUT = (Path(__file__).resolve().parents[2]
|
||||
/ "docs" / "sweeps" / "short_context_decode_composite_sweep.csv")
|
||||
|
||||
|
||||
def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
|
||||
op_log = r.engine.op_log
|
||||
wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
|
||||
|
||||
active_pes = set()
|
||||
for rec in op_log:
|
||||
cid = rec.component_id or ""
|
||||
if cid and ".cube" in cid and ".pe" in cid:
|
||||
parts = cid.split(".")
|
||||
active_pes.add(".".join(parts[:3]))
|
||||
n_pe = len(active_pes)
|
||||
|
||||
gemm_time_ns = 0.0
|
||||
math_count = 0
|
||||
math_pipeline_time_ns = 0.0
|
||||
dma_read_bytes = 0
|
||||
dma_write_bytes = 0
|
||||
ipcq_bytes = 0
|
||||
for rec in op_log:
|
||||
name = rec.op_name
|
||||
nbytes = rec.params.get("nbytes", 0) or 0
|
||||
# Composite/fused variants emit GEMM/MATH via two paths:
|
||||
# - non-pipeline ``tl.dot`` / primitives → ``gemm_f16`` / MATH_OPS
|
||||
# - pipeline ``tl.composite("gemm")`` → ``TileToken/GEMM`` (and
|
||||
# ``TileToken/MATH`` when a recipe prologue/epilogue is fused)
|
||||
# Count both so GEMM util / MATH op count reflect the real work.
|
||||
if name in ("gemm_f16", "TileToken/GEMM"):
|
||||
gemm_time_ns += rec.t_end - rec.t_start
|
||||
elif name == "TileToken/MATH":
|
||||
math_count += 1
|
||||
math_pipeline_time_ns = math_pipeline_time_ns + (rec.t_end - rec.t_start)
|
||||
elif name in MATH_OPS:
|
||||
math_count += 1
|
||||
elif name == "dma_read":
|
||||
dma_read_bytes += nbytes
|
||||
elif name == "dma_write":
|
||||
dma_write_bytes += nbytes
|
||||
elif name == "ipcq_copy":
|
||||
ipcq_bytes += nbytes
|
||||
|
||||
gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
|
||||
hbm_bw_util = (
|
||||
(dma_read_bytes + dma_write_bytes)
|
||||
/ (wall_ns * n_pe * PEAK_PE_HBM_BPS)
|
||||
) if n_pe else 0.0
|
||||
kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
|
||||
|
||||
return {
|
||||
"variant": "composite",
|
||||
"mode": mode,
|
||||
"kv_per_cube": kv_per_cube,
|
||||
"C": C,
|
||||
"S_kv": S_kv,
|
||||
"wall_us": wall_ns / 1000,
|
||||
"n_pe": n_pe,
|
||||
"gemm_util": gemm_util,
|
||||
"math_count": math_count,
|
||||
"math_pipeline_us": math_pipeline_time_ns / 1000,
|
||||
"hbm_bw_util": hbm_bw_util,
|
||||
"hbm_read_mb": dma_read_bytes / (1024 * 1024),
|
||||
"hbm_write_kb": dma_write_bytes / 1024,
|
||||
"ipcq_kb": ipcq_bytes / 1024,
|
||||
"kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_decode_composite_mode_context_sweep():
|
||||
"""Run 4 modes × 5 context lengths (variant 2), dump 5 metrics to CSV."""
|
||||
rows = []
|
||||
for mode, kv_per_cube, C in MODES:
|
||||
for S_kv in CONTEXT_LENGTHS:
|
||||
print(f">>> DECODE-cmp {mode} S_kv={S_kv//1024}K "
|
||||
f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
|
||||
r = _run_decode_composite(kv_per_cube=kv_per_cube, C=C,
|
||||
S_kv=S_kv, h_q=64)
|
||||
assert r.completion.ok, (
|
||||
f"DECODE-cmp {mode} S_kv={S_kv}: {r.completion}"
|
||||
)
|
||||
m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
|
||||
C=C, S_kv=S_kv)
|
||||
rows.append(m)
|
||||
print(f" wall={m['wall_us']:.1f}μs "
|
||||
f"gemm_util={m['gemm_util']:.3f} "
|
||||
f"hbm_bw_util={m['hbm_bw_util']:.3f} "
|
||||
f"ipcq={m['ipcq_kb']:.1f}KB "
|
||||
f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
|
||||
flush=True)
|
||||
|
||||
CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
|
||||
with CSV_OUT.open("w", newline="") as f:
|
||||
w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
w.writeheader()
|
||||
for row in rows:
|
||||
w.writerow(row)
|
||||
print(f"\nWrote {len(rows)} rows to {CSV_OUT}")
|
||||
@@ -1,143 +0,0 @@
|
||||
"""Decode (variant 3) composite + softmax_merge fused mode-comparison sweep.
|
||||
|
||||
Per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}) the
|
||||
same 5 metrics as the first-level sweep are collected:
|
||||
1. Wall clock latency (μs)
|
||||
2. Hardware utilization (GEMM util, MATH op count)
|
||||
3. HBM bandwidth utilization
|
||||
4. Communication cost (IPCQ bytes — chain reduce)
|
||||
5. KV cache space cost (per-cube bytes)
|
||||
|
||||
Kernel: ``_gqa_attention_decode_short_composite.py`` (variant 2,
|
||||
GEMM-only composite; Q·Kᵀ via tl.composite, P·V via tl.dot).
|
||||
|
||||
Output: ``docs/sweeps/short_context_decode_composite_fused_sweep.csv``.
|
||||
|
||||
Run: ``pytest tests/attention/test_gqa_short_context_sweep_decode_composite_fused.py -m slow``
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.attention.test_gqa_short_context import (
|
||||
D_HEAD, H_KV, P, TILE_S_KV, _run_decode_composite_fused,
|
||||
)
|
||||
|
||||
PEAK_PE_HBM_BPS = 8 * 32.0 # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
|
||||
|
||||
MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
|
||||
"maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
|
||||
"max_elem"}
|
||||
|
||||
MODES = [
|
||||
("A1", 1, 8),
|
||||
("A2", 2, 4),
|
||||
("A4", 4, 2),
|
||||
("B", 8, 1),
|
||||
]
|
||||
|
||||
CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
|
||||
|
||||
CSV_OUT = (Path(__file__).resolve().parents[2]
|
||||
/ "docs" / "sweeps" / "short_context_decode_composite_fused_sweep.csv")
|
||||
|
||||
|
||||
def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
|
||||
op_log = r.engine.op_log
|
||||
wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
|
||||
|
||||
active_pes = set()
|
||||
for rec in op_log:
|
||||
cid = rec.component_id or ""
|
||||
if cid and ".cube" in cid and ".pe" in cid:
|
||||
parts = cid.split(".")
|
||||
active_pes.add(".".join(parts[:3]))
|
||||
n_pe = len(active_pes)
|
||||
|
||||
gemm_time_ns = 0.0
|
||||
math_count = 0
|
||||
math_pipeline_time_ns = 0.0
|
||||
dma_read_bytes = 0
|
||||
dma_write_bytes = 0
|
||||
ipcq_bytes = 0
|
||||
for rec in op_log:
|
||||
name = rec.op_name
|
||||
nbytes = rec.params.get("nbytes", 0) or 0
|
||||
# Composite/fused variants emit GEMM/MATH via two paths:
|
||||
# - non-pipeline ``tl.dot`` / primitives → ``gemm_f16`` / MATH_OPS
|
||||
# - pipeline ``tl.composite("gemm")`` → ``TileToken/GEMM`` (and
|
||||
# ``TileToken/MATH`` when a recipe prologue/epilogue is fused)
|
||||
# Count both so GEMM util / MATH op count reflect the real work.
|
||||
if name in ("gemm_f16", "TileToken/GEMM"):
|
||||
gemm_time_ns += rec.t_end - rec.t_start
|
||||
elif name == "TileToken/MATH":
|
||||
math_count += 1
|
||||
math_pipeline_time_ns = math_pipeline_time_ns + (rec.t_end - rec.t_start)
|
||||
elif name in MATH_OPS:
|
||||
math_count += 1
|
||||
elif name == "dma_read":
|
||||
dma_read_bytes += nbytes
|
||||
elif name == "dma_write":
|
||||
dma_write_bytes += nbytes
|
||||
elif name == "ipcq_copy":
|
||||
ipcq_bytes += nbytes
|
||||
|
||||
gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
|
||||
hbm_bw_util = (
|
||||
(dma_read_bytes + dma_write_bytes)
|
||||
/ (wall_ns * n_pe * PEAK_PE_HBM_BPS)
|
||||
) if n_pe else 0.0
|
||||
kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
|
||||
|
||||
return {
|
||||
"variant": "composite_fused",
|
||||
"mode": mode,
|
||||
"kv_per_cube": kv_per_cube,
|
||||
"C": C,
|
||||
"S_kv": S_kv,
|
||||
"wall_us": wall_ns / 1000,
|
||||
"n_pe": n_pe,
|
||||
"gemm_util": gemm_util,
|
||||
"math_count": math_count,
|
||||
"math_pipeline_us": math_pipeline_time_ns / 1000,
|
||||
"hbm_bw_util": hbm_bw_util,
|
||||
"hbm_read_mb": dma_read_bytes / (1024 * 1024),
|
||||
"hbm_write_kb": dma_write_bytes / 1024,
|
||||
"ipcq_kb": ipcq_bytes / 1024,
|
||||
"kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_decode_composite_fused_mode_context_sweep():
|
||||
"""Run 4 modes × 5 context lengths (variant 2), dump 5 metrics to CSV."""
|
||||
rows = []
|
||||
for mode, kv_per_cube, C in MODES:
|
||||
for S_kv in CONTEXT_LENGTHS:
|
||||
print(f">>> DECODE-fuse {mode} S_kv={S_kv//1024}K "
|
||||
f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
|
||||
r = _run_decode_composite_fused(kv_per_cube=kv_per_cube, C=C,
|
||||
S_kv=S_kv, h_q=64)
|
||||
assert r.completion.ok, (
|
||||
f"DECODE-fuse {mode} S_kv={S_kv}: {r.completion}"
|
||||
)
|
||||
m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
|
||||
C=C, S_kv=S_kv)
|
||||
rows.append(m)
|
||||
print(f" wall={m['wall_us']:.1f}μs "
|
||||
f"gemm_util={m['gemm_util']:.3f} "
|
||||
f"hbm_bw_util={m['hbm_bw_util']:.3f} "
|
||||
f"ipcq={m['ipcq_kb']:.1f}KB "
|
||||
f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
|
||||
flush=True)
|
||||
|
||||
CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
|
||||
with CSV_OUT.open("w", newline="") as f:
|
||||
w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
w.writeheader()
|
||||
for row in rows:
|
||||
w.writerow(row)
|
||||
print(f"\nWrote {len(rows)} rows to {CSV_OUT}")
|
||||
@@ -1,145 +0,0 @@
|
||||
"""Prefill mode-comparison sweep across context lengths.
|
||||
|
||||
Measures per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}):
|
||||
1. Wall clock latency (μs)
|
||||
2. Hardware utilization
|
||||
- GEMM engine utilization (Σ gemm time / wall × n_pe)
|
||||
- MATH op count
|
||||
3. HBM bandwidth utilization
|
||||
(Σ dma_read+write bytes / (wall × n_pe × peak_pe_hbm_bw))
|
||||
4. Communication cost (IPCQ bytes — broadcast)
|
||||
5. KV cache space cost (per-cube bytes)
|
||||
|
||||
Sliced-prefill convention: T_q = 8.
|
||||
|
||||
Output: ``docs/sweeps/prefill_sweep.csv`` (16 rows).
|
||||
|
||||
Run: ``pytest tests/attention/test_gqa_short_context_sweep_prefill.py -m slow``
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.attention.test_gqa_short_context import (
|
||||
D_HEAD, H_KV, P, TILE_S_KV, _run_prefill,
|
||||
)
|
||||
|
||||
# Peak HBM BW per PE — derived from topology (cube total / PE count).
|
||||
# topology.yaml: hbm_total_bw_gbs = 1024.0 per cube, 8 PE per cube.
|
||||
PEAK_PE_HBM_BPS = 8 * 32.0 # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
|
||||
|
||||
MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
|
||||
"maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
|
||||
"max_elem"}
|
||||
|
||||
MODES = [
|
||||
("A1", 1, 8),
|
||||
("A2", 2, 4),
|
||||
("A4", 4, 2),
|
||||
("B", 8, 1),
|
||||
]
|
||||
|
||||
CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
|
||||
|
||||
T_Q_PREFILL = 8
|
||||
|
||||
CSV_OUT = (Path(__file__).resolve().parents[2]
|
||||
/ "docs" / "sweeps" / "short_context_prefill_sweep.csv")
|
||||
|
||||
|
||||
def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
|
||||
op_log = r.engine.op_log
|
||||
wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
|
||||
|
||||
# Per-PE detection (cube.pe)
|
||||
active_pes = set()
|
||||
for rec in op_log:
|
||||
cid = rec.component_id or ""
|
||||
if cid and ".cube" in cid and ".pe" in cid:
|
||||
parts = cid.split(".")
|
||||
active_pes.add(".".join(parts[:3]))
|
||||
n_pe = len(active_pes)
|
||||
|
||||
# Stage time / count / bytes
|
||||
gemm_time_ns = 0.0
|
||||
math_count = 0
|
||||
dma_read_bytes = 0
|
||||
dma_write_bytes = 0
|
||||
ipcq_bytes = 0
|
||||
for rec in op_log:
|
||||
name = rec.op_name
|
||||
nbytes = rec.params.get("nbytes", 0) or 0
|
||||
if name == "gemm_f16":
|
||||
gemm_time_ns += rec.t_end - rec.t_start
|
||||
elif name in MATH_OPS:
|
||||
math_count += 1
|
||||
elif name == "dma_read":
|
||||
dma_read_bytes += nbytes
|
||||
elif name == "dma_write":
|
||||
dma_write_bytes += nbytes
|
||||
elif name == "ipcq_copy":
|
||||
ipcq_bytes += nbytes
|
||||
|
||||
gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
|
||||
hbm_bw_util = (
|
||||
(dma_read_bytes + dma_write_bytes)
|
||||
/ (wall_ns * n_pe * PEAK_PE_HBM_BPS)
|
||||
) if n_pe else 0.0
|
||||
|
||||
# KV cache space — per-cube bytes (K + V, f16)
|
||||
# Per cube owns kv_per_cube heads × S_kv tokens × d_head × 2 (K+V) × 2 (f16)
|
||||
kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
|
||||
|
||||
return {
|
||||
"mode": mode,
|
||||
"kv_per_cube": kv_per_cube,
|
||||
"C": C,
|
||||
"S_kv": S_kv,
|
||||
"wall_us": wall_ns / 1000,
|
||||
"n_pe": n_pe,
|
||||
"gemm_util": gemm_util,
|
||||
"math_count": math_count,
|
||||
"hbm_bw_util": hbm_bw_util,
|
||||
"hbm_read_mb": dma_read_bytes / (1024 * 1024),
|
||||
"hbm_write_kb": dma_write_bytes / 1024,
|
||||
"ipcq_kb": ipcq_bytes / 1024,
|
||||
"kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_prefill_mode_context_sweep():
|
||||
"""Run 4 modes × 4 context lengths, dump 5 metrics to CSV.
|
||||
|
||||
Smoke-asserts each run completes; the measurement output is the CSV.
|
||||
"""
|
||||
rows = []
|
||||
for mode, kv_per_cube, C in MODES:
|
||||
for S_kv in CONTEXT_LENGTHS:
|
||||
print(f">>> PREFILL {mode} S_kv={S_kv//1024}K "
|
||||
f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
|
||||
r = _run_prefill(kv_per_cube=kv_per_cube, C=C,
|
||||
T_q=T_Q_PREFILL, S_kv=S_kv, h_q=64)
|
||||
assert r.completion.ok, (
|
||||
f"PREFILL {mode} S_kv={S_kv}: {r.completion}"
|
||||
)
|
||||
m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
|
||||
C=C, S_kv=S_kv)
|
||||
rows.append(m)
|
||||
print(f" wall={m['wall_us']:.1f}μs "
|
||||
f"gemm_util={m['gemm_util']:.3f} "
|
||||
f"hbm_bw_util={m['hbm_bw_util']:.3f} "
|
||||
f"ipcq={m['ipcq_kb']:.1f}KB "
|
||||
f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
|
||||
flush=True)
|
||||
|
||||
CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
|
||||
with CSV_OUT.open("w", newline="") as f:
|
||||
w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
w.writeheader()
|
||||
for row in rows:
|
||||
w.writerow(row)
|
||||
print(f"\nWrote {len(rows)} rows to {CSV_OUT}")
|
||||
@@ -1,143 +0,0 @@
|
||||
"""Prefill (variant 2) GEMM-only composite mode-comparison sweep.
|
||||
|
||||
Per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}) the
|
||||
same 5 metrics as the first-level sweep are collected:
|
||||
1. Wall clock latency (μs)
|
||||
2. Hardware utilization (GEMM util, MATH op count)
|
||||
3. HBM bandwidth utilization
|
||||
4. Communication cost (IPCQ bytes — broadcast)
|
||||
5. KV cache space cost (per-cube bytes)
|
||||
|
||||
Kernel: ``_gqa_attention_prefill_short_composite.py`` (variant 2,
|
||||
GEMM-only composite; Q·Kᵀ via tl.composite, P·V via tl.dot).
|
||||
|
||||
Output: ``docs/sweeps/short_context_prefill_composite_sweep.csv``.
|
||||
|
||||
Run: ``pytest tests/attention/test_gqa_short_context_sweep_prefill_composite.py -m slow``
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.attention.test_gqa_short_context import (
|
||||
D_HEAD, H_KV, P, TILE_S_KV, _run_prefill_composite,
|
||||
)
|
||||
|
||||
PEAK_PE_HBM_BPS = 8 * 32.0 # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
|
||||
|
||||
MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
|
||||
"maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
|
||||
"max_elem"}
|
||||
|
||||
MODES = [
|
||||
("A1", 1, 8),
|
||||
("A2", 2, 4),
|
||||
("A4", 4, 2),
|
||||
("B", 8, 1),
|
||||
]
|
||||
|
||||
CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
|
||||
|
||||
CSV_OUT = (Path(__file__).resolve().parents[2]
|
||||
/ "docs" / "sweeps" / "short_context_prefill_composite_sweep.csv")
|
||||
|
||||
|
||||
def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
|
||||
op_log = r.engine.op_log
|
||||
wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
|
||||
|
||||
active_pes = set()
|
||||
for rec in op_log:
|
||||
cid = rec.component_id or ""
|
||||
if cid and ".cube" in cid and ".pe" in cid:
|
||||
parts = cid.split(".")
|
||||
active_pes.add(".".join(parts[:3]))
|
||||
n_pe = len(active_pes)
|
||||
|
||||
gemm_time_ns = 0.0
|
||||
math_count = 0
|
||||
math_pipeline_time_ns = 0.0
|
||||
dma_read_bytes = 0
|
||||
dma_write_bytes = 0
|
||||
ipcq_bytes = 0
|
||||
for rec in op_log:
|
||||
name = rec.op_name
|
||||
nbytes = rec.params.get("nbytes", 0) or 0
|
||||
# Composite/fused variants emit GEMM/MATH via two paths:
|
||||
# - non-pipeline ``tl.dot`` / primitives → ``gemm_f16`` / MATH_OPS
|
||||
# - pipeline ``tl.composite("gemm")`` → ``TileToken/GEMM`` (and
|
||||
# ``TileToken/MATH`` when a recipe prologue/epilogue is fused)
|
||||
# Count both so GEMM util / MATH op count reflect the real work.
|
||||
if name in ("gemm_f16", "TileToken/GEMM"):
|
||||
gemm_time_ns += rec.t_end - rec.t_start
|
||||
elif name == "TileToken/MATH":
|
||||
math_count += 1
|
||||
math_pipeline_time_ns = math_pipeline_time_ns + (rec.t_end - rec.t_start)
|
||||
elif name in MATH_OPS:
|
||||
math_count += 1
|
||||
elif name == "dma_read":
|
||||
dma_read_bytes += nbytes
|
||||
elif name == "dma_write":
|
||||
dma_write_bytes += nbytes
|
||||
elif name == "ipcq_copy":
|
||||
ipcq_bytes += nbytes
|
||||
|
||||
gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
|
||||
hbm_bw_util = (
|
||||
(dma_read_bytes + dma_write_bytes)
|
||||
/ (wall_ns * n_pe * PEAK_PE_HBM_BPS)
|
||||
) if n_pe else 0.0
|
||||
kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
|
||||
|
||||
return {
|
||||
"variant": "composite",
|
||||
"mode": mode,
|
||||
"kv_per_cube": kv_per_cube,
|
||||
"C": C,
|
||||
"S_kv": S_kv,
|
||||
"wall_us": wall_ns / 1000,
|
||||
"n_pe": n_pe,
|
||||
"gemm_util": gemm_util,
|
||||
"math_count": math_count,
|
||||
"math_pipeline_us": math_pipeline_time_ns / 1000,
|
||||
"hbm_bw_util": hbm_bw_util,
|
||||
"hbm_read_mb": dma_read_bytes / (1024 * 1024),
|
||||
"hbm_write_kb": dma_write_bytes / 1024,
|
||||
"ipcq_kb": ipcq_bytes / 1024,
|
||||
"kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_prefill_composite_mode_context_sweep():
|
||||
"""Run 4 modes × 5 context lengths (variant 2 prefill), dump 5 metrics to CSV."""
|
||||
rows = []
|
||||
for mode, kv_per_cube, C in MODES:
|
||||
for S_kv in CONTEXT_LENGTHS:
|
||||
print(f">>> PREFILL-cmp {mode} S_kv={S_kv//1024}K "
|
||||
f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
|
||||
r = _run_prefill_composite(kv_per_cube=kv_per_cube, C=C,
|
||||
T_q=8, S_kv=S_kv, h_q=64)
|
||||
assert r.completion.ok, (
|
||||
f"PREFILL-cmp {mode} S_kv={S_kv}: {r.completion}"
|
||||
)
|
||||
m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
|
||||
C=C, S_kv=S_kv)
|
||||
rows.append(m)
|
||||
print(f" wall={m['wall_us']:.1f}μs "
|
||||
f"gemm_util={m['gemm_util']:.3f} "
|
||||
f"hbm_bw_util={m['hbm_bw_util']:.3f} "
|
||||
f"ipcq={m['ipcq_kb']:.1f}KB "
|
||||
f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
|
||||
flush=True)
|
||||
|
||||
CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
|
||||
with CSV_OUT.open("w", newline="") as f:
|
||||
w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
w.writeheader()
|
||||
for row in rows:
|
||||
w.writerow(row)
|
||||
print(f"\nWrote {len(rows)} rows to {CSV_OUT}")
|
||||
@@ -1,143 +0,0 @@
|
||||
"""Prefill (variant 3) composite + softmax_merge fused mode-comparison sweep.
|
||||
|
||||
Per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}) the
|
||||
same 5 metrics as the first-level sweep are collected:
|
||||
1. Wall clock latency (μs)
|
||||
2. Hardware utilization (GEMM util, MATH op count)
|
||||
3. HBM bandwidth utilization
|
||||
4. Communication cost (IPCQ bytes — broadcast)
|
||||
5. KV cache space cost (per-cube bytes)
|
||||
|
||||
Kernel: ``_gqa_attention_prefill_short_composite_fused.py`` (variant 2,
|
||||
GEMM-only composite; Q·Kᵀ via tl.composite, P·V via tl.dot).
|
||||
|
||||
Output: ``docs/sweeps/short_context_prefill_composite_fused_sweep.csv``.
|
||||
|
||||
Run: ``pytest tests/attention/test_gqa_short_context_sweep_prefill_composite_fused.py -m slow``
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.attention.test_gqa_short_context import (
|
||||
D_HEAD, H_KV, P, TILE_S_KV, _run_prefill_composite_fused,
|
||||
)
|
||||
|
||||
PEAK_PE_HBM_BPS = 8 * 32.0 # 8 channels/PE x 32 GB/s = 256 bytes/ns per PE
|
||||
|
||||
MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
|
||||
"maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
|
||||
"max_elem"}
|
||||
|
||||
MODES = [
|
||||
("A1", 1, 8),
|
||||
("A2", 2, 4),
|
||||
("A4", 4, 2),
|
||||
("B", 8, 1),
|
||||
]
|
||||
|
||||
CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
|
||||
|
||||
CSV_OUT = (Path(__file__).resolve().parents[2]
|
||||
/ "docs" / "sweeps" / "short_context_prefill_composite_fused_sweep.csv")
|
||||
|
||||
|
||||
def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
|
||||
op_log = r.engine.op_log
|
||||
wall_ns = max(rec.t_end for rec in op_log) - min(rec.t_start for rec in op_log)
|
||||
|
||||
active_pes = set()
|
||||
for rec in op_log:
|
||||
cid = rec.component_id or ""
|
||||
if cid and ".cube" in cid and ".pe" in cid:
|
||||
parts = cid.split(".")
|
||||
active_pes.add(".".join(parts[:3]))
|
||||
n_pe = len(active_pes)
|
||||
|
||||
gemm_time_ns = 0.0
|
||||
math_count = 0
|
||||
math_pipeline_time_ns = 0.0
|
||||
dma_read_bytes = 0
|
||||
dma_write_bytes = 0
|
||||
ipcq_bytes = 0
|
||||
for rec in op_log:
|
||||
name = rec.op_name
|
||||
nbytes = rec.params.get("nbytes", 0) or 0
|
||||
# Composite/fused variants emit GEMM/MATH via two paths:
|
||||
# - non-pipeline ``tl.dot`` / primitives → ``gemm_f16`` / MATH_OPS
|
||||
# - pipeline ``tl.composite("gemm")`` → ``TileToken/GEMM`` (and
|
||||
# ``TileToken/MATH`` when a recipe prologue/epilogue is fused)
|
||||
# Count both so GEMM util / MATH op count reflect the real work.
|
||||
if name in ("gemm_f16", "TileToken/GEMM"):
|
||||
gemm_time_ns += rec.t_end - rec.t_start
|
||||
elif name == "TileToken/MATH":
|
||||
math_count += 1
|
||||
math_pipeline_time_ns = math_pipeline_time_ns + (rec.t_end - rec.t_start)
|
||||
elif name in MATH_OPS:
|
||||
math_count += 1
|
||||
elif name == "dma_read":
|
||||
dma_read_bytes += nbytes
|
||||
elif name == "dma_write":
|
||||
dma_write_bytes += nbytes
|
||||
elif name == "ipcq_copy":
|
||||
ipcq_bytes += nbytes
|
||||
|
||||
gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
|
||||
hbm_bw_util = (
|
||||
(dma_read_bytes + dma_write_bytes)
|
||||
/ (wall_ns * n_pe * PEAK_PE_HBM_BPS)
|
||||
) if n_pe else 0.0
|
||||
kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
|
||||
|
||||
return {
|
||||
"variant": "composite_fused",
|
||||
"mode": mode,
|
||||
"kv_per_cube": kv_per_cube,
|
||||
"C": C,
|
||||
"S_kv": S_kv,
|
||||
"wall_us": wall_ns / 1000,
|
||||
"n_pe": n_pe,
|
||||
"gemm_util": gemm_util,
|
||||
"math_count": math_count,
|
||||
"math_pipeline_us": math_pipeline_time_ns / 1000,
|
||||
"hbm_bw_util": hbm_bw_util,
|
||||
"hbm_read_mb": dma_read_bytes / (1024 * 1024),
|
||||
"hbm_write_kb": dma_write_bytes / 1024,
|
||||
"ipcq_kb": ipcq_bytes / 1024,
|
||||
"kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_prefill_composite_fused_mode_context_sweep():
|
||||
"""Run 4 modes × 5 context lengths (variant 2 prefill), dump 5 metrics to CSV."""
|
||||
rows = []
|
||||
for mode, kv_per_cube, C in MODES:
|
||||
for S_kv in CONTEXT_LENGTHS:
|
||||
print(f">>> PREFILL-fuse {mode} S_kv={S_kv//1024}K "
|
||||
f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
|
||||
r = _run_prefill_composite_fused(kv_per_cube=kv_per_cube, C=C,
|
||||
T_q=8, S_kv=S_kv, h_q=64)
|
||||
assert r.completion.ok, (
|
||||
f"PREFILL-fuse {mode} S_kv={S_kv}: {r.completion}"
|
||||
)
|
||||
m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
|
||||
C=C, S_kv=S_kv)
|
||||
rows.append(m)
|
||||
print(f" wall={m['wall_us']:.1f}μs "
|
||||
f"gemm_util={m['gemm_util']:.3f} "
|
||||
f"hbm_bw_util={m['hbm_bw_util']:.3f} "
|
||||
f"ipcq={m['ipcq_kb']:.1f}KB "
|
||||
f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
|
||||
flush=True)
|
||||
|
||||
CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
|
||||
with CSV_OUT.open("w", newline="") as f:
|
||||
w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
|
||||
w.writeheader()
|
||||
for row in rows:
|
||||
w.writerow(row)
|
||||
print(f"\nWrote {len(rows)} rows to {CSV_OUT}")
|
||||
@@ -1,75 +0,0 @@
|
||||
"""Regression: a recv handle used as a composite ``b`` operand must be read
|
||||
in place (on-chip), NOT DMA_READ from its non-HBM slot address.
|
||||
|
||||
Before the recv-pinning fix, recv handles were unpinned → the composite
|
||||
scheduler emitted DMA_READ(B) from the IPCQ slot address (non-HBM) →
|
||||
PhysAddr.decode raised PhysAddrError. On-chip (TCM/SRAM) recv slots are
|
||||
now pinned (ADR-0065 D4), so a composite operand reads them in place.
|
||||
|
||||
Isolates operand ``b``: ``a`` is a pinned ``tl.load``; only ``b`` is the
|
||||
recv handle. Mirrors the decode row-chain IPCQ directions (intra_E recv /
|
||||
intra_W send) so the recv actually lands without deadlock.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
|
||||
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
|
||||
|
||||
TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[1] / "topology.yaml"
|
||||
P = 8
|
||||
|
||||
|
||||
def _ccl_cfg():
|
||||
return resolve_algorithm_config(
|
||||
load_ccl_config(), name="lrab_hierarchical_allreduce",
|
||||
)
|
||||
|
||||
|
||||
def _engine_factory(t, d):
|
||||
return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)
|
||||
|
||||
|
||||
def _recv_operand_kernel(a_ptr, b_ptr, o_ptr, M, K, N, P, *, tl):
|
||||
pe_id = tl.program_id(axis=0)
|
||||
pe_in_group = pe_id % P
|
||||
group_cols = min(4, P)
|
||||
pe_col = pe_in_group % group_cols
|
||||
|
||||
A = tl.load(a_ptr, shape=(M, K), dtype="f16") # pinned (no DMA)
|
||||
if pe_col < group_cols - 1:
|
||||
B = tl.recv(dir="intra_E", shape=(K, N), dtype="f16") # recv handle
|
||||
out = tl.composite("gemm", a=A, b=B) # ← b = recv slot
|
||||
tl.store(o_ptr, out)
|
||||
if pe_col > 0:
|
||||
B_send = tl.load(b_ptr, shape=(K, N), dtype="f16")
|
||||
tl.send(dir="intra_W", src=B_send)
|
||||
|
||||
|
||||
def test_recv_handle_as_composite_operand_reads_in_place():
|
||||
M = K = N = 64
|
||||
topo = resolve_topology(str(TOPOLOGY_DEFAULT))
|
||||
|
||||
def _bench_fn(ctx):
|
||||
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
|
||||
dp = DPPolicy(cube="replicate", pe="replicate", num_cubes=1, num_pes=P)
|
||||
a = ctx.zeros((M, K), dtype="f16", dp=dp, name="recv_operand_a")
|
||||
b = ctx.zeros((K, N), dtype="f16", dp=dp, name="recv_operand_b")
|
||||
o = ctx.empty((M, N), dtype="f16", dp=dp, name="recv_operand_o")
|
||||
ctx.launch(
|
||||
"recv_operand_composite", _recv_operand_kernel,
|
||||
a, b, o, M, K, N, P,
|
||||
_auto_dim_remap=False,
|
||||
)
|
||||
|
||||
r = run_bench(
|
||||
topology=topo, bench_fn=_bench_fn,
|
||||
device=resolve_device(None), engine_factory=_engine_factory,
|
||||
)
|
||||
assert r.completion.ok, f"recv-operand composite completion: {r.completion}"
|
||||
@@ -1,7 +1,6 @@
|
||||
|
||||
system:
|
||||
ns_per_mm: 0.01 # wire propagation delay: 10 ps/mm (on-chip silicon)
|
||||
flit_bytes: 4096 # wire chunk size (ADR-0033 D1: multiple of hbm_ctrl.burst_bytes=256)
|
||||
|
||||
sips:
|
||||
count: 2
|
||||
|
||||
Reference in New Issue
Block a user