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

gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model

Browse Source 
Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.

ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
  intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
  online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
  (m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.

ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
  owns whole KV heads; PE-parallel heads with intra-group chain
  reduce. Prefill has no Ring KV (each head fully resident).

ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).

ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.

ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
  primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
  consults table with dispatch_cycles fallback (ADR-0046 §D6
  back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
  DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
  measurable end-to-end.

P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).

Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
  _attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
  (EN + KO kept in sync).

Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-06-09 18:15:59 -07:00
parent b3730a33eb
commit d282144339
52 changed files with 3952 additions and 2526 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/triton_emu/tl_context.py
+204 -37
View File
@@ -22,6 +22,7 @@ from kernbench.common.pe_commands import (
EPILOGUE_OPS,
CompletionHandle,
CompositeCmd,
CopyCmd,
DmaReadCmd,
DmaWriteCmd,
GemmCmd,
@@ -42,13 +43,67 @@ _DTYPE_BYTES: dict[str, int] = {
}
class LoadFuture:
"""Lazy ``tl.load`` future (ADR-0062 §D2).
Mirrors ``RecvFuture`` for IPCQ, generalised to HBM loads. Carries
the originating ``DmaReadCmd``, the SimPy completion event (set by
the runner once the DMA has been issued), and a resolved flag.
Consumer ops auto-wait via ``TLContext._await_pending(handle)``,
which yields ``event`` if ``resolved`` is False, then reads the
DMA's bytes into ``handle.data`` and marks the future resolved.
"""
__slots__ = ("cmd", "event", "resolved", "data")
def __init__(self, cmd: DmaReadCmd) -> None:
self.cmd = cmd
self.event: object | None = None # simpy.Event set by runner
self.resolved: bool = False
self.data: object = None
class _ScratchScope:
"""Context manager that recycles per-tile scratch (ADR-0063 D1).
``__enter__`` snapshots ``_scratch_cursor``; ``__exit__`` restores it,
so every handle allocated inside the ``with``-block has its address
freed for the next iteration. Persistent state (running ``(m, ℓ, O)``,
prefetch buffers) lives outside the scope per ADR-0063 D3.
"""
def __init__(self, ctx: "TLContext") -> None:
self._ctx = ctx
self._save: int | None = None
def __enter__(self) -> "_ScratchScope":
self._save = self._ctx._scratch_cursor
return self
def __exit__(self, *exc_info: object) -> bool:
if self._save is not None:
self._ctx._scratch_cursor = self._save
return False
class TLContext:
"""Fake Triton Language context.
Args:
pe_id: program instance index (returned by program_id).
num_programs: total number of program instances.
dispatch_cycles: PE_CPU overhead per tl API call (auto-inserted).
dispatch_cycles: uniform PE_CPU overhead per tl API call. Used as
a fallback when ``issue_cost_table`` is None (ADR-0046 §D6
back-compat). When ``issue_cost_table`` is provided, the
per-kind table value is used instead.
issue_cost_table: optional per-op-type CPU issue cost table
(ADR-0064 D1). When provided, each ``tl.*`` call charges the
table value keyed by op kind ("composite", "load", "store",
"dot", "math", "ipcq_send", "ipcq_recv", "copy_to"). Unknown
kinds fall back to ``dispatch_cycles``. Live PE_CPU paths
construct TLContext with ``DEFAULT_CPU_ISSUE_COST`` so the
hybrid's CPU-saturation lever is measurable.
"""
def __init__(
@@ -61,12 +116,14 @@ class TLContext:
num_cubes: int = 1,
scratch_base: int = 0,
scratch_size: int = 1 << 20, # 1 MiB per kernel invocation
issue_cost_table: dict[str, int] | None = None,
) -> None:
self._pe_id = pe_id
self._num_programs = num_programs
self._cube_id = cube_id
self._num_cubes = num_cubes
self._dispatch_cycles = dispatch_cycles
self._issue_cost_table = issue_cost_table
self._commands: list[PeCommand] = []
self._handle_counter = 0
self._completion_counter = 0
@@ -79,6 +136,22 @@ class TLContext:
self._scratch_size = scratch_size
self._scratch_cursor = 0
def scratch_scope(self) -> _ScratchScope:
"""Per-tile scratch recycling context manager (ADR-0063).
Usage:
with tl.scratch_scope():
s = tl.dot(q, k_t) # per-tile temporaries —
p = tl.softmax(s) # their scratch is rewound
o_j = tl.dot(p, v) # on __exit__
Persistent state (running ``(m, ℓ, O)``, lazy-load prefetch
buffers) must be allocated **outside** the scope; only handles
allocated inside are recycled.
"""
return _ScratchScope(self)
def _scratch_alloc(self, nbytes: int) -> int:
"""Allocate a unique scratch address for an output TensorHandle.
@@ -120,9 +193,23 @@ class TLContext:
def _nbytes(self, shape: tuple[int, ...], dtype: str) -> int:
return math.prod(shape) * self._dtype_bytes(dtype)
def _emit_dispatch_overhead(self) -> None:
if self._dispatch_cycles > 0:
self._emit(PeCpuOverheadCmd(cycles=self._dispatch_cycles))
def _emit_dispatch_overhead(self, kind: str | None = None) -> None:
"""Charge per-op-type CPU issue cost (ADR-0064 D1).
When ``issue_cost_table`` was provided, look up the per-kind cost
and emit ``PeCpuOverheadCmd(cycles=N)`` if N > 0. Unknown kinds
fall back to the uniform ``dispatch_cycles`` for forward-compat
when a new ``tl.*`` op is added before the table is updated.
When ``issue_cost_table`` is None, preserve the ADR-0046 §D6
contract: emit ``PeCpuOverheadCmd(dispatch_cycles)`` if positive.
"""
if self._issue_cost_table is not None and kind is not None:
cycles = self._issue_cost_table.get(kind, self._dispatch_cycles)
else:
cycles = self._dispatch_cycles
if cycles > 0:
self._emit(PeCpuOverheadCmd(cycles=cycles))
def _make_handle(
self, addr: int, shape: tuple[int, ...], dtype: str,
@@ -177,37 +264,103 @@ class TLContext:
def load(
self, ptr: int, shape: tuple[int, ...], dtype: str = "f16",
) -> TensorHandle:
"""Load tensor from HBM. Returns TensorHandle pointing at HBM[ptr].
"""Load tensor from HBM — **lazy** (ADR-0062 §D1/§D2).
In greenlet mode: returns TensorHandle with actual numpy data.
In command-list mode: returns TensorHandle with data=None.
Posts the ``DmaReadCmd`` to PE_DMA and returns immediately with a
``TensorHandle`` whose ``pending`` field references a fresh
``LoadFuture``. The actual DMA completion is awaited at the first
consuming op (``tl.dot``, MATH, ``tl.store``, ``tl.send``,
``tl.copy_to``, ``tl.composite``) via ``_await_pending``.
The returned handle's ``space`` is "hbm" so subsequent ops (math,
send, store) using this handle as a source resolve via MemoryStore
at ``(hbm, ptr)`` — which is where the load's underlying data
actually lives in Phase 2 storage.
Command-list mode: emits the DmaReadCmd to ``self._commands``,
attaches a LoadFuture for structural compatibility (its event
stays None — no engine, no SimPy event to wait on).
"""
self._emit_dispatch_overhead()
handle = self._make_handle(
addr=ptr, shape=shape, dtype=dtype, space="hbm", pinned=True,
self._emit_dispatch_overhead("load")
nbytes = self._nbytes(shape, dtype)
# LoadFuture is mutable; create it first, attach to the handle,
# then point its ``cmd`` at the DmaReadCmd that references the
# *final* handle. This guarantees ``handle.pending.cmd.handle is handle``.
future = LoadFuture.__new__(LoadFuture)
future.event = None
future.resolved = False
future.data = None
handle = TensorHandle(
id=self._next_handle_id(),
addr=ptr, shape=shape, dtype=dtype, nbytes=nbytes,
data=None, space="hbm", pinned=True, pending=future,
)
cmd = DmaReadCmd(handle=handle, src_addr=ptr, nbytes=handle.nbytes)
data = self._emit(cmd)
if data is not None:
# Greenlet mode: attach real data to handle (preserve space + pinned)
return TensorHandle(
id=handle.id, addr=handle.addr, shape=handle.shape,
dtype=handle.dtype, nbytes=handle.nbytes, data=data,
space=handle.space, pinned=handle.pinned,
)
cmd = DmaReadCmd(handle=handle, src_addr=ptr, nbytes=nbytes)
future.cmd = cmd
if self._runner is not None:
# Lazy: runner posts the DmaReadCmd, sets future.event, then
# switches back immediately. No yield on completion here.
self._runner.switch_to_simpy(("load_issue", future))
else:
self._commands.append(cmd)
return handle
def _await_pending(self, *handles: TensorHandle | None) -> None:
"""Auto-wait at first use (ADR-0062 §D2).
For each handle carrying an unresolved ``LoadFuture``, yield on
the DMA completion event (greenlet → runner). Command-list mode
is a no-op.
"""
if self._runner is None:
return
for h in handles:
if h is None:
continue
pending = getattr(h, "pending", None)
if pending is None or pending.resolved:
continue
self._runner.switch_to_simpy(("load_await", h))
def store(self, ptr: int, handle: TensorHandle) -> None:
"""Store tensor from TCM to HBM."""
self._emit_dispatch_overhead()
self._await_pending(handle)
self._emit_dispatch_overhead("store")
cmd = DmaWriteCmd(handle=handle, dst_addr=ptr, nbytes=handle.nbytes)
self._emit(cmd)
def copy_to(self, dst: TensorHandle, src: TensorHandle) -> None:
"""TCM-to-TCM byte copy (ADR-0063 §D3.1).
Writes ``src``'s bytes into ``dst``'s address. Both handles must
live in TCM and have matching shape and dtype. Symmetric to
``tl.store`` (which targets HBM) but stays on-chip so it doesn't
emit a DMA entry into op_log.
Used inside ``tl.scratch_scope()`` to persist a scoped result —
typically an updated running ``(m, ℓ, O)`` — to an outside-scope
(persistent) handle so its bytes survive the scope's ``__exit__``
cursor rewind (ADR-0063 §D3 two-arena pattern).
"""
if src.shape != dst.shape:
raise ValueError(
f"tl.copy_to: shape mismatch — src.shape={src.shape} "
f"vs dst.shape={dst.shape}"
)
if src.dtype != dst.dtype:
raise ValueError(
f"tl.copy_to: dtype mismatch — src.dtype={src.dtype!r} "
f"vs dst.dtype={dst.dtype!r}"
)
if dst.space != "tcm":
raise ValueError(
f"tl.copy_to: dst must be in TCM (got space={dst.space!r}); "
"writes to HBM go through tl.store"
)
if src.space != "tcm":
raise ValueError(
f"tl.copy_to: src must be in TCM (got space={src.space!r}); "
"reads from HBM go through tl.load"
)
self._await_pending(src)
self._emit_dispatch_overhead("copy_to")
self._emit(CopyCmd(src=src, dst=dst, nbytes=src.nbytes))
# ── GEMM Engine (blocking) ────────────────────────────────────
def dot(self, a: TensorHandle, b: TensorHandle) -> TensorHandle:
@@ -224,7 +377,8 @@ class TLContext:
out_shape = (*a.shape[:-2], m, n)
out_dtype = a.dtype
out = self._make_compute_out(shape=out_shape, dtype=out_dtype)
self._emit_dispatch_overhead()
self._await_pending(a, b)
self._emit_dispatch_overhead("dot")
self._emit(GemmCmd(a=a, b=b, out=out, m=m, k=k, n=n))
return out
@@ -232,7 +386,8 @@ class TLContext:
def _unary_math(self, op: str, x: TensorHandle) -> TensorHandle:
out = self._make_compute_out(shape=x.shape, dtype=x.dtype)
self._emit_dispatch_overhead()
self._await_pending(x)
self._emit_dispatch_overhead("math")
self._emit(MathCmd(op=op, inputs=(x,), out=out))
return out
@@ -265,7 +420,8 @@ class TLContext:
out_shape = list(x.shape)
out_shape[axis] = 1
out = self._make_compute_out(shape=tuple(out_shape), dtype=x.dtype)
self._emit_dispatch_overhead()
self._await_pending(x)
self._emit_dispatch_overhead("math")
self._emit(MathCmd(op=op, inputs=(x,), out=out, axis=axis))
return out
@@ -284,7 +440,8 @@ class TLContext:
self, op: str, a: TensorHandle, b: TensorHandle,
) -> TensorHandle:
out = self._make_compute_out(shape=a.shape, dtype=a.dtype)
self._emit_dispatch_overhead()
self._await_pending(a, b)
self._emit_dispatch_overhead("math")
self._emit(MathCmd(op=op, inputs=(a, b), out=out))
return out
@@ -292,7 +449,8 @@ class TLContext:
self, cond: TensorHandle, a: TensorHandle, b: TensorHandle,
) -> TensorHandle:
out = self._make_compute_out(shape=a.shape, dtype=a.dtype)
self._emit_dispatch_overhead()
self._await_pending(cond, a, b)
self._emit_dispatch_overhead("math")
self._emit(MathCmd(op="where", inputs=(cond, a, b), out=out))
return out
@@ -309,7 +467,8 @@ class TLContext:
) -> TensorHandle:
"""Fused multiply-add: a * b + c (real Triton: tl.fma)."""
out = self._make_compute_out(shape=a.shape, dtype=a.dtype)
self._emit_dispatch_overhead()
self._await_pending(a, b, c)
self._emit_dispatch_overhead("math")
self._emit(MathCmd(op="fma", inputs=(a, b, c), out=out))
return out
@@ -321,7 +480,8 @@ class TLContext:
) -> TensorHandle:
"""Clamp x to [min, max] (real Triton: tl.clamp)."""
out = self._make_compute_out(shape=x.shape, dtype=x.dtype)
self._emit_dispatch_overhead()
self._await_pending(x, min, max)
self._emit_dispatch_overhead("math")
self._emit(MathCmd(op="clamp", inputs=(x, min, max), out=out))
return out
@@ -333,7 +493,8 @@ class TLContext:
canonical (x - max) → exp → sum → div sequence.
"""
out = self._make_compute_out(shape=x.shape, dtype=x.dtype)
self._emit_dispatch_overhead()
self._await_pending(x)
self._emit_dispatch_overhead("math")
self._emit(MathCmd(op="softmax", inputs=(x,), out=out, axis=axis))
return out
@@ -427,13 +588,16 @@ class TLContext:
space = getattr(src, "space", space)
if src_addr is None or nbytes is None or shape is None:
raise ValueError("tl.send: provide either a TensorHandle or src_addr/nbytes/shape")
# ADR-0062: if the source is a lazy-loaded handle, await first so
# the data snapshot below sees the real bytes.
self._await_pending(src)
# Carry the handle's .data snapshot (if available). When the source
# is a recv slot, .data holds the numpy array that was read from
# MemoryStore at recv-time. This prevents a Phase 1 race where a
# later IPCQ inbound overwrites the slot before the outbound
# PE_DMA reads it.
handle_data = getattr(src, "data", None) if src is not None else None
self._emit_dispatch_overhead()
self._emit_dispatch_overhead("ipcq_send")
cmd = IpcqSendCmd(
direction=dir,
src_addr=src_addr, src_space=space,
@@ -467,7 +631,7 @@ class TLContext:
arrived. In greenlet/runner mode, ``handle.data`` carries the
actual ndarray; in command-list mode the handle is a placeholder.
"""
self._emit_dispatch_overhead()
self._emit_dispatch_overhead("ipcq_recv")
if dst_addr is not None and dst_space is not None:
cmd = IpcqRecvCmd(
direction=dir,
@@ -518,7 +682,7 @@ class TLContext:
they receive. This API is segregated from ``tl.recv`` so the
diagnostic flag can never accidentally be set in real workloads.
"""
self._emit_dispatch_overhead()
self._emit_dispatch_overhead("ipcq_recv")
cmd = IpcqRecvCmd(
direction=dir,
shape=shape, dtype=dtype,
@@ -547,7 +711,7 @@ class TLContext:
dtype: str = "f16",
) -> "RecvFuture":
"""Non-blocking recv. Returns a future to pass into ``tl.wait``."""
self._emit_dispatch_overhead()
self._emit_dispatch_overhead("ipcq_recv")
cmd = IpcqRecvCmd(
direction=dir,
shape=shape, dtype=dtype,
@@ -582,6 +746,9 @@ class TLContext:
Returns CompletionHandle for use with wait().
"""
# ADR-0062: composite operand DMA paths still need their inputs
# to be resolved before the composite reads them via PE_SCHEDULER.
self._await_pending(a, b)
# Compute output size based on op
if op == "gemm" and b is not None:
m, k = a.shape[-2], a.shape[-1]
@@ -609,7 +776,7 @@ class TLContext:
ops_tuple = (head_spec, *epi_specs)
completion = CompletionHandle(id=self._next_completion_id())
self._emit_dispatch_overhead()
self._emit_dispatch_overhead("composite")
self._emit(CompositeCmd(
completion=completion, op=op,
a=a, b=b, out_addr=out_ptr, out_nbytes=out_nbytes,