test short context length attention kernel

tests/attention/test_gqa_short_context_sweep_decode.py

@@ -0,0 +1,131 @@
+"""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}")