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

gqa(adr-0065): D8 — composite returns output handle + output-space DMA

Browse Source 
tl.composite now returns the output TensorHandle (not a CompletionHandle) so its result chains like tl.dot's; the handle carries the completion in a CompositeFuture pending so downstream ops and tl.wait auto-await it. out is a handle: out=tl.ref(addr,shape) (HBM, DMA_WRITE inside the composite) or an in-place TCM handle (STORE only); omitted -> TLContext auto-allocates a TCM scratch. out_ptr kept as HBM shorthand (= out=tl.ref(out_ptr, shape)) to avoid churning ~30 existing call sites.

tl.ref now returns space=hbm (it references HBM data; operand-input DMA stays pinned-based per D4 so input streaming is unchanged). tiling: the tile loop's DMA_WRITE is gated on out.space==hbm (out analog of the operand pinned rule) — a TCM output stays on-chip (chainable) and its high-bit scratch address no longer hits the DMA PA decoder.

Fixes the opt2 data-mode crash: the recipe accumulator O is TCM -> no DMA_WRITE -> opt2 now RUNS end-to-end in data mode (enable_data=True). Numeric parity of the recipe MATH ops is the next step. Suite 812 pass / 3 pre-existing fail; existing composite benches use out_ptr->hbm->DMA_WRITE unchanged (byte-equal).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Yangwook
2026-06-10 22:29:10 -07:00
parent dd525bfcb7
commit e8d6c283d8
7 changed files with 267 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/benches/_gqa_attention_decode_opt2.py
+3 -3
View File
@@ -67,9 +67,9 @@ def gqa_attention_decode_opt2_kernel(
shape=(d_head, rest), dtype="f16")
V1 = tl.ref(v_ptr + half * KV_ROW_BYTES, shape=(rest, d_head),
dtype="f16")
# #1: Q·Kᵀ composite → score tile (TCM-resident, consumed by #2).
scores1 = tl.zeros((G * T_q, rest), dtype="f16")
tl.composite(op="gemm", a=Q, b=K_T1, out=scores1)
# #1: Q·Kᵀ composite → score tile. No `out` → auto-allocated TCM
# scratch returned as a handle (ADR-0065 D8), consumed by #2.
scores1 = tl.composite(op="gemm", a=Q, b=K_T1)
# #2: softmax_merge recipe (online merge of (m,l,O)) + P·V + add.
tl.composite(
prologue=[{"op": "softmax_merge", "s": scores1,
src/kernbench/components/builtin/tiling.py
+14 -7
View File
@@ -24,6 +24,7 @@ def generate_gemm_plan(
a_pinned: bool = False,
b_pinned: bool = False,
epilogue_specs: tuple = (),
out_space: str = "hbm",
) -> PipelinePlan:
"""Generate GEMM tile plan: M→N→K order.
@@ -152,13 +153,18 @@ def generate_gemm_plan(
"nbytes": out_bytes,
},
))
stages.append(Stage(
stage_type=StageType.DMA_WRITE,
component=dma_id,
params={
"dst_addr": c_addr, "nbytes": out_bytes,
},
))
# DMA_WRITE only when the output is HBM-resident
# (ADR-0065 D8). A TCM output stays on-chip (chainable);
# writing its high-bit scratch address as an HBM PA would
# also fail the DMA decoder.
if out_space == "hbm":
stages.append(Stage(
stage_type=StageType.DMA_WRITE,
component=dma_id,
params={
"dst_addr": c_addr, "nbytes": out_bytes,
},
))
tiles.append(TilePlan(tile_id=tile_id, stages=tuple(stages)))
tile_id += 1
@@ -281,6 +287,7 @@ def generate_plan_from_ops(
a_pinned=getattr(a, "pinned", False),
b_pinned=getattr(b, "pinned", False),
epilogue_specs=tuple(post_ops),
out_space=getattr(head.out, "space", "hbm"),
)
# Prologue (pre-GEMM KERNEL) + post-loop KERNEL ops become single-shot
src/kernbench/triton_emu/tl_context.py
+73 -14
View File
@@ -64,6 +64,23 @@ class LoadFuture:
self.data: object = None
class CompositeFuture:
"""Pending marker on a composite's output handle (ADR-0065 D8).
Carries the composite's ``CompletionHandle`` so a downstream op (via
``_await_pending``) or ``tl.wait`` can await it. Distinct from
``LoadFuture``: composite completion is tracked by completion id in the
runner's pending dict, so awaiting just re-issues a ``WaitCmd`` (a
second wait on an already-drained completion is a harmless no-op).
"""
__slots__ = ("completion", "resolved")
def __init__(self, completion: "CompletionHandle") -> None:
self.completion = completion
self.resolved: bool = False
class _ScratchScope:
"""Context manager that recycles per-tile scratch (ADR-0063 D1).
@@ -256,9 +273,14 @@ class TLContext:
"""Create a TensorHandle referencing HBM data without issuing DMA.
Used when the scheduler will stream data per-tile (e.g., tensor b
in a composite GEMM). No command is generated.
in a composite GEMM) or as a composite's HBM output handle
(ADR-0065 D8). The handle's ``space="hbm"`` — it references HBM
data. No command is generated. (Operand-input DMA stays
``pinned``-based per ADR-0065 D4, so this label does not change
input streaming.)
"""
return self._make_handle(addr=ptr, shape=shape, dtype=dtype)
return self._make_handle(addr=ptr, shape=shape, dtype=dtype,
space="hbm")
# ── Data Movement (blocking, DMA engine) ──────────────────────
@@ -333,7 +355,12 @@ class TLContext:
pending = getattr(h, "pending", None)
if pending is None or pending.resolved:
continue
self._runner.switch_to_simpy(("load_await", h))
if isinstance(pending, CompositeFuture):
# Composite output: wait on its completion (ADR-0065 D8).
self._emit(WaitCmd(handle=pending.completion))
pending.resolved = True
else:
self._runner.switch_to_simpy(("load_await", h))
def store(self, ptr: int, handle: TensorHandle) -> None:
"""Store tensor from TCM to HBM."""
@@ -737,26 +764,33 @@ class TLContext:
op: Literal["gemm", "math"],
a: TensorHandle | None = None,
b: TensorHandle | None = None,
out_ptr: int = 0,
math_op: str | None = None,
*,
out: TensorHandle | None = None,
out_ptr: int | None = None,
prologue: list[dict] | None = None,
epilogue: list[dict] | None = None,
acc_dtype: str | None = None,
tile_shape: tuple[int, int, int] | None = None,
) -> CompletionHandle:
) -> TensorHandle:
"""Submit a composite command (non-blocking, tiled pipeline).
``prologue`` (ADR-0065 D5) is an ordered list of recipe dicts — each
with an ``"op"`` naming a ``RECIPE_DESCRIPTORS`` entry plus its
operands. TLContext expands each into flat MATH OpSpecs that run
before the head op and auto-binds the recipe's primary output into
the head GEMM's ``a`` operand (D6.6). ``out`` (TensorHandle) is the
explicit write-back handle, preferred over ``out_ptr``. ``epilogue``
is an ordered list of ``EPILOGUE_OPS`` dicts.
the head GEMM's ``a`` operand (D6.6). ``epilogue`` is an ordered list
of ``EPILOGUE_OPS`` dicts.
Returns CompletionHandle for use with wait().
``out`` (ADR-0065 D8) is a ``TensorHandle``: an HBM ref
(``tl.ref(addr, shape)`` → ``space="hbm"`` → DMA_WRITE) or an
in-place TCM handle (e.g. a recipe accumulator → STORE only).
``out_ptr`` (int) is shorthand for an HBM output — equivalent to
``out=tl.ref(out_ptr, <output shape>)``. If neither is given, a TCM
scratch is auto-allocated (chainable, like ``tl.dot``).
Returns the **output TensorHandle** — auto-awaited by downstream ops
(and by ``tl.wait(handle)``) via its ``pending`` field.
"""
# ADR-0062: composite operand DMA paths still need their inputs
# to be resolved before the composite reads them via PE_SCHEDULER.
@@ -798,14 +832,25 @@ class TLContext:
out_shape = a.shape
out_nbytes = a.nbytes
# Write-back handle: explicit `out` handle wins over `out_ptr`.
# Write-back handle (ADR-0065 D8). Explicit `out` (HBM ref or an
# in-place TCM accumulator) wins; `out_ptr` is HBM shorthand; else
# auto-allocate a TCM scratch — the chainable default, like tl.dot.
# pinned=True on the auto output: it is TCM-resident, so a downstream
# op reading it needs no DMA_READ.
if out is not None:
out_handle = out
else:
elif out_ptr is not None:
out_handle = TensorHandle(
id=self._next_handle_id(),
addr=out_ptr, shape=out_shape, dtype=out_dtype,
nbytes=out_nbytes, space="tcm",
nbytes=out_nbytes, space="hbm",
)
else:
out_handle = TensorHandle(
id=self._next_handle_id(),
addr=self._scratch_alloc(out_nbytes),
shape=out_shape, dtype=out_dtype, nbytes=out_nbytes,
space="tcm", pinned=True,
)
# Head op (flat-ops, ADR-0065 D2): GEMM takes named a/b and carries
@@ -836,7 +881,12 @@ class TLContext:
self._emit(CompositeCmd(
completion=completion, ops=ops_tuple, rw_handles=rw_handles,
))
return completion
# Attach the completion to the output handle so downstream ops (and
# tl.wait) auto-await it (ADR-0065 D8 / ADR-0062 lazy pattern), and
# return it so the result chains like tl.dot's (frozen dataclass →
# set pending via object.__setattr__).
object.__setattr__(out_handle, "pending", CompositeFuture(completion))
return out_handle
def _expand_prologue(
self, prologue: list[dict],
@@ -1041,7 +1091,16 @@ class TLContext:
handle.result = th
return th
# Composite path (existing behaviour)
# Composite output handle (ADR-0065 D8): wait on its completion.
if isinstance(handle, TensorHandle):
pending = getattr(handle, "pending", None)
if isinstance(pending, CompositeFuture):
if not pending.resolved:
self._emit(WaitCmd(handle=pending.completion))
pending.resolved = True
return None
# Composite path (CompletionHandle or None — existing behaviour).
self._emit(WaitCmd(handle=handle))
return None