test: 3-variant sweep (baseline + composite + composite-fused) for prefill/decode
This commit is contained in:
eusonice
@@ -0,0 +1,143 @@
|
||||
"""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 = 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_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)
|
||||
|
||||
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}")
|
||||
@@ -0,0 +1,143 @@
|
||||
"""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 = 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_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)
|
||||
|
||||
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}")
|
||||
@@ -0,0 +1,143 @@
|
||||
"""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 = 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_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)
|
||||
|
||||
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}")
|
||||
@@ -0,0 +1,143 @@
|
||||
"""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 = 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_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)
|
||||
|
||||
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}")
|
||||
Reference in New Issue
Block a user