From e8d6c283d8b506e54e2245c4550016edc1922287 Mon Sep 17 00:00:00 2001 From: Yangwook Date: Wed, 10 Jun 2026 22:29:10 -0700 Subject: [PATCH] =?UTF-8?q?gqa(adr-0065):=20D8=20=E2=80=94=20composite=20r?= =?UTF-8?q?eturns=20output=20handle=20+=20output-space=20DMA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- ...rog-flat-ops-composite-softmax-merge-ko.md | 35 +++++++ ...5-prog-flat-ops-composite-softmax-merge.md | 40 ++++++++ .../benches/_gqa_attention_decode_opt2.py | 6 +- src/kernbench/components/builtin/tiling.py | 21 ++-- src/kernbench/triton_emu/tl_context.py | 87 ++++++++++++++--- tests/test_composite_output_handle.py | 96 +++++++++++++++++++ tests/test_triton_emu.py | 8 +- 7 files changed, 267 insertions(+), 26 deletions(-) create mode 100644 tests/test_composite_output_handle.py diff --git a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md index 172ce6f..33a8a05 100644 --- a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md +++ b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge-ko.md @@ -241,6 +241,41 @@ TLContext (compiler analog) PE_SCHEDULER (scheduler/dispatcher) imports: recipe 관련 없음 ``` +### D8. Composite 출력 handle + 출력-space DMA + +`tl.composite(...)` 는 **출력 `TensorHandle`** 을 반환(단순 `CompletionHandle` +아님) — composite 결과를 `tl.dot` 결과처럼 downstream op 에 먹일 수 있음. +handle 의 `pending` 필드가 composite 완료를 참조; downstream consumer 가 +auto-await(`_await_pending`, ADR-0062 lazy 패턴), `tl.wait(handle)` 동작. + +**`out` 은 항상 `TensorHandle`** (raw 주소 아님) — 위치는 커널의 선택, 추측 안 함: + +- **HBM 출력** → `out=tl.ref(addr, shape)`. `tl.ref` 가 `space="hbm"` handle + 반환 → composite 가 DMA_WRITE 로 write-back. STORE + DMA_WRITE 는 composite + 파이프라인 **내부**에 유지 — fused op 의 일부이지 별도 `tl.store` 아님. +- **In-place TCM** → `out=<기존 TCM handle>` (예: recipe 누적기 `O`). + `space="tcm"` → STORE 만, 결과 TCM 에 남음 (chainable). +- **미지정** → TLContext 가 **TCM scratch 자동 할당** (`tl.dot` 과 동일) 후 반환. + +편의상 `out_ptr: int` 는 **HBM 출력 shorthand** 로 허용 — +`out=tl.ref(out_ptr, <출력 shape>)` 와 동치(컴파일러가 `space="hbm"` handle 로 +감쌈). 출력 미지정 시 scratch 주소는 커널이 아니라 컴파일러가 소유. +(`tl.ref` 는 `space="hbm"` 반환 — HBM 데이터 참조이므로. operand-**입력** DMA +결정은 D4 대로 `pinned` 기반이라 이 라벨이 입력 스트리밍에 영향 없음.) + +출력 handle 의 **`space` 가 타일 루프의 write-back 을 결정** (D4 의 operand +`pinned` 규칙의 출력판): + +| `out.space` | 타일 루프 write-back | +| --- | --- | +| `"tcm"` | STORE 만 — TCM 에 남음 (chainable; **DMA_WRITE 없음**) | +| `"hbm"` | STORE + DMA_WRITE — 최종 HBM 출력 | + +이로써 출력이 TCM 에 남는 composite(예: recipe 누적기)가 합법이 되고, +TCM-scratch 주소가 DMA 엔진의 PA 디코더(고비트 scratch 주소를 거부)로 가는 +일을 막음. 기존 벤치는 HBM 출력을 `tl.ref` 로 감싸므로 write-back +불변(byte-equal). + ## Alternatives ### A1. tile 당 3 composite (prologue 개념 없이) diff --git a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md index 8514fdc..62b2be5 100644 --- a/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md +++ b/docs/adr-proposed/ADR-0065-prog-flat-ops-composite-softmax-merge.md @@ -267,6 +267,46 @@ TLContext (compiler analog) PE_SCHEDULER (scheduler/dispatcher) imports: nothing recipe-related ``` +### D8. Composite output handle + output-space DMA + +`tl.composite(...)` returns the **output `TensorHandle`** (not a bare +`CompletionHandle`), so a composite's result can feed a downstream op the +same way `tl.dot`'s result does. The handle's `pending` field references +the composite's completion; downstream consumers auto-await it +(`_await_pending`, ADR-0062 lazy pattern) and `tl.wait(handle)` works. + +**`out` is always a `TensorHandle`** (never a raw address) — its location is +the kernel's choice, never guessed: + +- **HBM output** → `out=tl.ref(addr, shape)`. `tl.ref` returns a handle with + `space="hbm"`, so the composite writes it back via DMA_WRITE. The + STORE + DMA_WRITE stay **inside** the composite pipeline — they are part of + the fused op, *not* a separate `tl.store`. +- **In-place TCM** → `out=` (e.g. the recipe accumulator + `O`). `space="tcm"` → STORE only, result stays in TCM (chainable). +- **Omitted** → TLContext **auto-allocates a TCM scratch** (like `tl.dot`) + and returns it. + +For ergonomics, `out_ptr: int` is accepted as **shorthand for an HBM +output** — equivalent to `out=tl.ref(out_ptr, )` (the compiler +wraps it into an `space="hbm"` handle). When no output is specified the +compiler owns the scratch address, never the kernel. (`tl.ref` returns +`space="hbm"` — it references HBM data. The operand-**input** DMA decision +stays `pinned`-based per D4, so this label does not affect input streaming.) + +The output handle's **`space` drives the tile loop's write-back** (the +output analog of the operand `pinned` rule in D4): + +| `out.space` | tile-loop write-back | +| --- | --- | +| `"tcm"` | STORE only — result stays in TCM (chainable; **no DMA_WRITE**) | +| `"hbm"` | STORE + DMA_WRITE — final HBM output | + +This makes a TCM-resident composite output (e.g. a recipe accumulator) legal +**and** keeps a TCM-scratch address out of the DMA engine's PA decoder +(which would otherwise reject the high-bit scratch address). Existing benches +wrap their HBM output via `tl.ref` → write-back unchanged (byte-equal). + ## Alternatives ### A1. Three composites per tile (no prologue concept) diff --git a/src/kernbench/benches/_gqa_attention_decode_opt2.py b/src/kernbench/benches/_gqa_attention_decode_opt2.py index 4be0f03..669564d 100644 --- a/src/kernbench/benches/_gqa_attention_decode_opt2.py +++ b/src/kernbench/benches/_gqa_attention_decode_opt2.py @@ -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, diff --git a/src/kernbench/components/builtin/tiling.py b/src/kernbench/components/builtin/tiling.py index 2401a64..6b7f8f2 100644 --- a/src/kernbench/components/builtin/tiling.py +++ b/src/kernbench/components/builtin/tiling.py @@ -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 diff --git a/src/kernbench/triton_emu/tl_context.py b/src/kernbench/triton_emu/tl_context.py index c25cf2b..e30e97d 100644 --- a/src/kernbench/triton_emu/tl_context.py +++ b/src/kernbench/triton_emu/tl_context.py @@ -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, )``. 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 diff --git a/tests/test_composite_output_handle.py b/tests/test_composite_output_handle.py new file mode 100644 index 0000000..69bdd6a --- /dev/null +++ b/tests/test_composite_output_handle.py @@ -0,0 +1,96 @@ +"""Phase 1 spec tests for ADR-0065 D8 — composite output handle + output-space DMA. + +`tl.composite(...)` returns the **output TensorHandle** (not a bare +CompletionHandle) so its result chains into downstream ops like `tl.dot`'s. +`out` is always a handle: `out=tl.ref(addr, shape)` for HBM output (DMA_WRITE +inside the composite), `out=` for in-place, or omitted → TLContext +auto-allocates a TCM scratch. The output handle's `space` decides the tile +loop's write-back (tcm → STORE only; hbm → STORE + DMA_WRITE). `tl.ref` +returns `space="hbm"`. + +Phase 1 (this commit): tests only. FAIL until D8: + - `tl.ref` returns space="tcm"; `composite` returns CompletionHandle; + output always DMA_WRITEs; `tl.wait` rejects a TensorHandle. +""" +from __future__ import annotations + +from kernbench.common.pe_commands import OpSpec, Scope, TensorHandle +from kernbench.triton_emu.tl_context import TLContext + + +def _tl() -> TLContext: + return TLContext(pe_id=0, num_programs=1, scratch_base=0x200000, + scratch_size=1 << 20) + + +def test_ref_returns_hbm_space(): + h = _tl().ref(0x1000, shape=(4, 4), dtype="f16") + assert h.space == "hbm" + + +def test_composite_returns_handle_auto_tcm(): + tl = _tl() + a = tl.ref(0x1000, shape=(8, 16), dtype="f16") + b = tl.ref(0x2000, shape=(16, 8), dtype="f16") + out = tl.composite(op="gemm", a=a, b=b) # out omitted → auto TCM + assert isinstance(out, TensorHandle) + assert out.space == "tcm" + assert out.addr >= 0x200000 # auto-allocated scratch + assert out.shape == (8, 8) + + +def test_composite_returns_the_explicit_out_handle(): + tl = _tl() + a = tl.ref(0x1000, shape=(8, 16), dtype="f16") + b = tl.ref(0x2000, shape=(16, 8), dtype="f16") + O = TensorHandle(id="O", addr=0x5000, shape=(8, 8), dtype="f16", + nbytes=128, space="tcm", pinned=True) + out = tl.composite(op="gemm", a=a, b=b, out=O) + assert out is O + + +def test_composite_output_chains_into_next_op(): + tl = _tl() + a = tl.ref(0x1000, shape=(8, 16), dtype="f16") + b = tl.ref(0x2000, shape=(16, 8), dtype="f16") + scores = tl.composite(op="gemm", a=a, b=b) # handle + out2 = tl.composite(op="math", a=scores, math_op="exp") + assert isinstance(out2, TensorHandle) + + +def test_wait_accepts_tensor_handle(): + tl = _tl() + a = tl.ref(0x1000, shape=(8, 16), dtype="f16") + b = tl.ref(0x2000, shape=(16, 8), dtype="f16") + out = tl.composite(op="gemm", a=a, b=b) + tl.wait(out) # must not raise + + +def test_output_space_drives_dma_write(): + from kernbench.components.builtin.pe_types import StageType + from kernbench.components.builtin.tiling import generate_plan_from_ops + + a = TensorHandle(id="a", addr=0x1000, shape=(64, 64), dtype="f16", + nbytes=8192, space="tcm", pinned=True) + b = TensorHandle(id="b", addr=0x2000, shape=(64, 64), dtype="f16", + nbytes=8192, space="hbm") + + def _gemm(out): + return OpSpec(kind="gemm", scope=Scope.OUTPUT_TILE, + operands={"a": a, "b": b}, out=out, + extra={"m": 64, "k": 64, "n": 64}) + + o_tcm = TensorHandle(id="ot", addr=0x300000, shape=(64, 64), dtype="f16", + nbytes=8192, space="tcm", pinned=True) + o_hbm = TensorHandle(id="oh", addr=0x9000, shape=(64, 64), dtype="f16", + nbytes=8192, space="hbm") + + def _has_write(ops): + plan = generate_plan_from_ops(ops, tile_m=64, tile_k=64, tile_n=64, + bytes_per_element=2, + pe_prefix="sip0.cube0.pe0") + return any(s.stage_type == StageType.DMA_WRITE + for s in plan.tiles[0].stages) + + assert not _has_write((_gemm(o_tcm),)), "tcm out → no DMA_WRITE" + assert _has_write((_gemm(o_hbm),)), "hbm out → DMA_WRITE" diff --git a/tests/test_triton_emu.py b/tests/test_triton_emu.py index 85ba537..9cb1e74 100644 --- a/tests/test_triton_emu.py +++ b/tests/test_triton_emu.py @@ -175,7 +175,10 @@ def test_tl_composite_nonblocking(): a = tl.load(0x1000, shape=(32, 64), dtype="f16") b = tl.load(0x2000, shape=(64, 32), dtype="f16") h = tl.composite(op="gemm", a=a, b=b, out_ptr=0x3000) - assert isinstance(h, CompletionHandle) + # ADR-0065 D8: composite returns the output TensorHandle (chainable), + # carrying the completion in its `pending` field. + assert isinstance(h, TensorHandle) + assert h.addr == 0x3000 and h.space == "hbm" comp_cmds = [c for c in tl.commands if isinstance(c, CompositeCmd)] assert len(comp_cmds) == 1 # Flat-ops shape (ADR-0065 D1): head op carries kind + write-back handle. @@ -195,7 +198,8 @@ def test_tl_wait_specific(): tl.wait(h) wait_cmds = [c for c in tl.commands if isinstance(c, WaitCmd)] assert len(wait_cmds) == 1 - assert wait_cmds[0].handle == h + # ADR-0065 D8: h is the output handle; tl.wait targets its completion. + assert wait_cmds[0].handle == h.pending.completion # ── 9. tl.wait() → WaitCmd(handle=None) ──────────────────────────