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

gqa(adr): ADR-0064 Rev2 (structural dispatch cost) + ADR-0065/DDD-0065 (flat-ops composite + softmax_merge recipe)

Browse Source 
ADR-0064 Revision 2: replace per-op-type calibration table with a
structural formula `FIXED_PER_CMD + cmd.logical_bytes × R`, defaults
anchored at typical composite ≈ 45 ns. Topology yaml override under
`pe_cost_model:` block; logical_bytes property per PE command.

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

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

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook Kang
2026-06-10 16:11:20 -07:00
parent 91cdfebb67
commit a8c50238c6
5 changed files with 1635 additions and 243 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/adr-proposed/ADR-0064-perf-cpu-issue-cost-model.md
+202 -123
View File
@@ -1,179 +1,258 @@
# ADR-0064: Per-op-type CPU issue cost model (command construct + dispatch)
# ADR-0064: Structural CPU dispatch cost model (`logical_bytes` + FIXED + R)
## Status
Proposed
Proposed (Revision 2)
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention). The hybrid
> decision there (GEMMs via `tl.composite`, softmax merge in the kernel)
> wins by **offloading tiling to PE_SCHEDULER so the CPU issues coarse
> descriptors and runs ahead, keeping the engines saturated**. That win is
> currently **invisible in the simulator** because per-op CPU issue cost is
> zero. This ADR makes the issue cost real and **op-type-differentiated**,
> so composite-vs-primitive trade-offs (and CPU saturation) are measurable.
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention) and **ADR-0065**
> (flat-ops composite + first stateful recipe). The hybrid decision there
> (GEMMs via `tl.composite`, softmax merge in the kernel) wins by
> **offloading tiling to PE_SCHEDULER so the CPU issues coarse descriptors
> and runs ahead, keeping the engines saturated**. That win is currently
> **invisible in the simulator** because per-op CPU issue cost is zero.
>
> Revision 2 replaces the **op-type calibration table** (the original
> proposal) with a **structural formula** derived from each command's
> `logical_bytes` — no per-op-type calibration needed; new op kinds are
> covered automatically.
## Context
### What exists today
- Every `tl.*` op calls `_emit_dispatch_overhead()` before emitting its
command (`tl_context.py:123-125, 190, 227, 235, 612, …`), which emits
command (`tl_context.py:196-212`), which emits
`PeCpuOverheadCmd(cycles=dispatch_cycles)` **only if** `dispatch_cycles
> 0`. So the issue cost is a **single uniform knob** applied identically
to `tl.load`, `tl.dot`, a MATH op, and `tl.composite`.
- That knob is hardcoded to **0** in both live execution paths
(`pe_cpu.py:101` greenlet runner, `:195` legacy replay). ⇒ issuing a
command — *constructing the descriptor and pushing it to the scheduler
queue* — currently costs **0 ns** on PE_CPU.
> 0`. The knob is **uniform** across op kinds and hardcoded to **0**
in both live execution paths (`pe_cpu.py:101` greenlet, `:195` replay).
- ⇒ issuing a command — *constructing the descriptor and pushing it to
the scheduler queue* — currently costs **0 ns** on PE_CPU.
- `PeCpuOverheadCmd` is consumed as `yield env.timeout(cmd.cycles)` on
PE_CPU (`kernel_runner.py:131-132`) and on PE_SCHEDULER
(`pe_scheduler.py:97-100`); a manual `tl.cycles(n)` also exists
(`tl_context.py:695`).
### Two findings that frame this ADR (from ADR-0060 review)
- **Q1 — composite issue cost.** Constructing + pushing a `tl.composite`
is expected to cost on the order of **~40 ns** of CPU time (descriptor
build + queue push). Today it is **0**. The hook exists; only the value
(and its per-op-type differentiation) is missing.
- **Q2 — scheduler dispatch vs DMA latency.** `PE_SCHEDULER`'s composite
dispatch is **non-blocking**: `_dispatch_composite` generates the tile
plan and enqueues to the feeder, returning immediately
(`pe_scheduler.py:104-121`). The actual **DMA latency is charged on
PE_DMA** as each tile flows through (`drain_ns = compute_drain_ns(path,
nbytes)`, `pe_dma.py:89`), **not** lumped into the scheduler dispatch ⇒
**no double-counting**, and DMA stays on its modelled component
(SPEC §0.1). The scheduler's own plan-generation currently costs **0**
sim time.
PE_CPU (`kernel_runner.py:131-132`).
### Why uniform-and-zero is wrong for the hybrid
The hybrid's whole argument is that **one** composite descriptor offloads
ADR-0060 §1's argument is that **one** composite descriptor offloads
`N_tiles` worth of GEMM tiling, so the CPU issues `O(1)` coarse commands
instead of `O(N_tiles × ops/tile)` fine ones. With issue cost = 0 (and
uniform), the model cannot show:
instead of `O(N_tiles × ops/tile)` fine ones. With issue cost = 0, the
model cannot show:
- that the primitive path may **fail to saturate** the engines when the
CPU cannot push fast enough (the core ADR-0060 §1 claim), nor
- that a composite **costs more to construct** than a single primitive but
far less than the many primitives it replaces.
CPU cannot push fast enough, nor
- that a composite **costs more to construct** than a single primitive
but far less than the many primitives it replaces.
A single uniform `dispatch_cycles` cannot express this: `tl.load` ≪
`tl.composite` in real construct cost.
### Why per-op-type calibration (Revision 1) was over-shaped
The original proposal had a `cost_table[kind]` keyed by op kind
(`composite`, `load`, `dot`, `math`, …). That required:
- a value per kind (calibration cost ≥ |kinds|),
- a new entry every time a new kind appears,
- and yet the *ratio* it tried to capture — "composite ≫ primitive,
but ≪ the primitives it replaces" — is structurally a function of
**how many fields the command carries**, not of the op kind.
A command's *byte footprint* is the natural proxy: a composite carrying
N OpSpecs has ~N× the bytes of a primitive op with one OpSpec. The
*fixed* part (queue head update, completion register, MMIO-class
latency) is per-command. The two together compose: `FIXED + bytes × R`.
## Decision
### D1. Per-op-type issue-cost table on PE_CPU
### D1. Structural dispatch cost formula
Replace the single `dispatch_cycles` scalar with a **cost table keyed by
command type**, charged on PE_CPU at issue (before the command is
dispatched to the scheduler). The greenlet pays this time, which is
exactly what gates how fast the CPU can push work — the saturation lever.
Each PE command going to PE_SCHEDULER incurs PE_CPU dispatch cycles:
Starting estimates (calibrate later — see review items):
```
dispatch_cycles(cmd) = FIXED_PER_CMD + cmd.logical_bytes × R
```
| Issued op | CPU issue cost (construct + push) |
where:
- `FIXED_PER_CMD` (cycles per command) models queue-tail update, MMIO-class
RTT, completion-event registration — fixed per command regardless of size.
- `R` (cycles per byte) models the queue-write bandwidth — bytes of the
command serialized into the scheduler queue.
- `cmd.logical_bytes` (int) is each command's *HW-logical* byte size,
computed from D2 below — not Python's `sys.getsizeof`.
PE_CPU emits `PeCpuOverheadCmd(cycles=dispatch_cycles(cmd))` before
dispatching, exactly as the existing hook (`tl_context.py:_emit_dispatch_
overhead`) — only the cycle value changes.
### D2. `logical_bytes` rule
Each PE command dataclass exposes `logical_bytes: int` (property). The
counting rule (HW-friendly, ignores Python overhead):
| Field kind | Bytes |
|---|---|
| `tl.composite` (GEMM/MATH descriptor) | ~40 ns (Q1 estimate) |
| `tl.load` / `tl.store` (DMA descriptor) | small (a few ns) |
| `tl.dot` / MATH / IPCQ send/recv | small (a few ns) |
| Command framing (cmd-type discriminator + completion id ref) | 4 |
| Opcode (op kind enum) | 1 |
| Enum (scope, etc.) | 1 |
| `TensorHandle` reference (address only — shape/dtype assumed in descriptor table) | 8 |
| Scalar (int/float) | 4 |
| Tuple length marker | 1 |
The point is the **ratio**: a composite construct ≫ a primitive issue, but
≪ the sum of the primitive issues it replaces. Exact ns are configurable.
`CompositeCmd` recursively sums its `ops` and `rw_handles`:
### D2. Execution latency stays on the engines (Q2 — no change)
```python
@property
def logical_bytes(self) -> int:
return (
4 # framing
+ 1 + sum(op.logical_bytes for op in self.ops)
+ 1 + 8 * len(self.rw_handles)
)
```
DMA/GEMM/MATH execution latency remains charged on PE_DMA / PE_GEMM /
PE_MATH as today. The issue cost (D1) is **CPU-side only** and additive; it
does **not** touch `drain_ns`, so there is no double-count. This preserves
SPEC §0.1 (latency on modelled components).
`OpSpec`:
### D3. Scheduler plan-generation cost — start at 0, revisit
```python
@property
def logical_bytes(self) -> int:
return (
1 + 1 # opcode + scope
+ 1 + 8 * len(self.operands) # named operand handles
+ (8 if self.out is not None else 0) # out handle
+ 1 + sum(4 for _ in self.extra.values()) # extra scalars
)
```
`PE_SCHEDULER`'s tile-plan generation costs 0 sim time today. Keep that for
now (the dominant lever is CPU issue cost, D1); expose it later via the
existing `overhead_ns` node attr if a scheduler-side cost proves material.
Marked as a review item, not decided here.
(Identical rule for `DmaReadCmd`, `MathCmd`, etc. — one property per
dataclass, ~3 lines each.)
### D4. Configurable values; goldens regenerate
### D3. Defaults — anchored on a typical composite ≈ 45 ns
The cost table is configurable (per-topology / node attrs, default to the
D1 estimates). Turning issue cost non-zero changes **every** bench's
latency, so golden latencies are **regenerated** once — a model-fidelity
improvement, same posture as ADR-0062's global lazy-load regression.
Anchor: a typical 1-op DMA→GEMM→DMA composite has `logical_bytes ≈ 52`
(framing 4 + GEMM OpSpec 39 + rw_handles 9). Target dispatch = 45 ns.
On-die producer→consumer queue assumed at 4 bytes/cycle.
```
clock = 1 GHz # 1 cycle = 1 ns
FIXED_PER_CMD = 32 cycles
R = 0.25 cycles/byte
```
Verification: `32 + 52 × 0.25 = 45 cycles ≈ 45 ns` ✓
### D4. Topology config override
Defaults are baked into `pe_cpu.py`. Topology yaml may override under a
`pe_cost_model:` section at the PE node attrs:
```yaml
pe:
attrs:
pe_cost_model:
fixed_per_cmd_cycles: 32
byte_cycles_recip: 0.25
clock_freq_ghz: 1.0
```
Missing keys fall back to defaults. The dispatch formula reads from
`node.attrs["pe_cost_model"]` at PE_CPU init.
### D5. Scope — what does and does not pay
| Path | Pays dispatch cost? |
|---|---|
| PE_CPU → PE_SCHEDULER for any `PeCommand` | **Yes** |
| `PeCpuOverheadCmd` itself (already cycles-explicit) | **No** (formula bypass) |
| Stages auto-generated by PE_SCHEDULER (DMA_READ/WRITE/FETCH/STORE) | **No** (PE_SCHEDULER-internal) |
| Engine compute latency (DMA `drain_ns`, GEMM/MATH `_compute_ns`) | **No change** — stays on engines (SPEC §0.1) |
This preserves the "latency on modelled components" invariant — dispatch
cost is *additional* CPU-side time, not folded into engine times.
### D6. Configurable values; goldens regenerate
Turning issue cost non-zero changes **every** bench's latency. Golden
latencies are **regenerated once** when this ADR lands — same posture as
ADR-0062 D3 lazy-load. After regeneration, the same calibration is in
effect for ADR-0065 opt2 measurement.
## Alternatives
### A1. Keep a single uniform `dispatch_cycles > 0`
### A1. Keep Revision 1's op-type calibration table
Rejected: op construct costs differ by an order of magnitude (`tl.load` vs
`tl.composite`); a uniform value either over-charges primitives or
under-charges composites, and in both cases misrepresents the
composite-vs-primitive trade-off the hybrid depends on.
Rejected: calibration cost scales with |kinds|, and the *ratio* the
table tried to capture is structurally a function of cmd size. The
structural formula reaches the same qualitative behaviour with two
calibratable numbers instead of N.
### A2. Charge the issue cost on PE_SCHEDULER instead of PE_CPU
### A2. Byte-only formula (no FIXED term)
Rejected. With FIXED = 0, opt2 (Option Y per ADR-0065) does **not** win
over opt3 — the total *bytes* dispatched per tile are similar (opt3 ≈ 232,
opt2 ≈ 380); the win is entirely in *fewer per-cmd fixed costs*. A
byte-only formula erases the very signal the model needs to expose.
### A3. Charge dispatch on PE_SCHEDULER instead of PE_CPU
Rejected: the saturation question is *"can the CPU push descriptors fast
enough to keep the engines busy?"* — that is a **PE_CPU** issue-bandwidth
property. Charging it on the scheduler would not model CPU back-pressure.
property. Charging on the scheduler would not model CPU back-pressure.
### A3. Model DMA program/setup time as a separate fixed per-descriptor cost
### A4. Model DMA program/setup time as a separate fixed per-descriptor cost
Real HW pays a DMA descriptor program cost distinct from transfer time.
Deferred: initially fold the descriptor-program cost into the **issuing
op's** D1 cost (a `tl.load` / composite that triggers DMA). Split it out to
a PE_DMA fixed setup only if calibration shows it matters. Review item.
op's** dispatch cost. Split it out to a PE_DMA fixed setup only if
calibration shows it matters.
## Consequences
### Positive
- The hybrid's CPU-offload / saturation win (ADR-0060 §1) becomes
**measurable**, not just structural.
- Composite-vs-primitive and tiling-granularity trade-offs are visible in
latency, enabling the eval narrative.
- More faithful to hardware (issue bandwidth is a real bottleneck).
- Hybrid's CPU-offload / saturation win (ADR-0060 §1) becomes
**measurable**, with a structurally honest model (no calibration table).
- Adding new op kinds (e.g., ADR-0065's `softmax_merge` 8-step recipe)
costs zero — they fit the same formula automatically.
- More faithful to hardware (queue-head MMIO RTT + queue-write bandwidth).
### Negative
- **All** bench goldens shift → one-time regeneration (D4); CI golden
- **All** bench goldens shift → one-time regeneration (D6); CI golden
fixtures update.
- Per-op-type values need calibration; the ~40 ns composite figure is an
estimate, primitives are unspecified — results are only as good as the
numbers (do not over-claim absolute latencies until calibrated).
- Adds a cost-table lookup on the issue path (negligible runtime).
- Two calibration knobs (FIXED, R) need values; defaults are anchored on
a documented assumption — treat absolute latencies as provisional until
a reference exists; keep the **ratios** defensible.
- Adds a small `logical_bytes` property to each PE command dataclass.
## Open review items (decided autonomously; revise on review)
## Open review items
1. **Calibration source for the per-op-type values.** Composite ≈ 40 ns is
a working estimate; primitive issue costs are placeholders. *Recommend:*
pick values from a documented assumption (instruction-issue + queue
push) and treat absolute latency as provisional until a real reference
exists; keep the **ratios** defensible.
2. **Scheduler plan-gen cost (D3).** *Recommend:* keep 0 initially; expose
via `overhead_ns` if a workload shows scheduler-bound behaviour.
3. **DMA program time (A3).** *Recommend:* fold into the issuing op's cost
first; split to PE_DMA setup only if needed.
4. **Where the table lives.** *Recommend:* a small central cost module
keyed by command type, overridable per topology — not scattered node
attrs — so the values are reviewable in one place.
5. **Regression rollout.** *Recommend:* land lazy-load (ADR-0062) and this
ADR's cost model in one golden-regeneration pass to avoid two churns.
6. **Interaction with the legacy replay path** (`pe_cpu.py:_execute_legacy`)
— ensure both the greenlet and replay paths read the same cost table so
results match. Verify.
1. **Calibration source for FIXED and R.** Defaults from "typical
composite = 45 ns + on-die queue 4 bytes/cycle"; reasonable for an
on-die producer→consumer queue. Revisit when a HW reference appears.
2. **Scheduler plan-gen cost.** Stays 0 — D5 keeps PE_SCHEDULER's
plan-generation outside the dispatch formula. Expose via existing
`overhead_ns` if a workload shows scheduler-bound behaviour.
3. **Where the override lives.** `pe_cost_model:` block under PE node
attrs in topology yaml — keeps all knobs in one place, reviewable.
4. **Path parity.** Both greenlet (`_execute_legacy` and `kernel_runner`)
and replay paths must read the same cost model. Verify.
## Test Requirements
1. **Composite charges once; primitives charge per-op.** A kernel issuing
one `tl.composite` over `N_tiles` charges one composite issue cost on
PE_CPU; the equivalent `N_tiles × ops` primitive kernel charges the
per-op cost `N_tiles × ops` times. Assert the PE_CPU busy time differs
accordingly.
2. **DMA latency unchanged (Q2).** For a fixed transfer, PE_DMA `drain_ns`
is identical to the pre-ADR value — the issue cost is additive on
PE_CPU, not folded into DMA (no double-count).
3. **Saturation is observable.** With non-zero per-op issue cost, a
many-tile primitive sweep shows GEMM-engine idle (CPU-bound issue)
whereas the composite sweep keeps it busy — the ADR-0060 §1 lever.
4. **Determinism:** identical inputs → identical op_log + latency
1. **Anchor preservation.** A typical DMA→GEMM→DMA composite (1 op,
`logical_bytes ≈ 52`) dispatches in 45 ns at the default values.
2. **Structural ratio.** opt3 vs opt2 dispatch (per ADR-0065 §verification):
`opt3 / opt2 ≈ 2.4×` at default calibration.
3. **Override path.** Topology yaml `pe_cost_model:` block changes the
per-PE dispatch cost; default is recovered when block is missing.
4. **`PeCpuOverheadCmd` bypass.** Manual `tl.cycles(n)` issues exactly `n`
cycles, not `n + dispatch_cycles(...)`.
5. **No double-count.** PE_DMA `drain_ns`, PE_GEMM/MATH `_compute_ns`
identical to pre-ADR values.
6. **Determinism.** Identical inputs → identical op_log + latency
(SPEC §0.1).
5. **Path parity:** greenlet and legacy-replay paths produce identical
issue-cost accounting for the same kernel.
7. **Path parity.** Greenlet and replay paths produce identical
dispatch-cycle accounting for the same kernel.
## Migration
ADR-0064 Revision 2 lands as a single PR with:
- `logical_bytes` property on each `PeCommand` dataclass
- formula application in `pe_cpu.py` dispatch path
- `pe_cost_model:` override read at PE_CPU init
- one-time goldens regeneration
After this lands, ADR-0065 builds on top with no further goldens churn
in existing benches (ADR-0065 is a meaning-preserving refactor of
`CompositeCmd` for the existing path; only opt2 is a new bench).