eusonice
132 lines
4.2 KiB
Python
132 lines
4.2 KiB
Python
"""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 = 1024.0 / 8 # 128 bytes/ns
|
||
|
||
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)
|
||
|
||
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}")
|