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

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>

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 개념 없이)

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=<existing TCM handle>` (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, <output shape>)` (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)

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,

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,6 +153,11 @@ def generate_gemm_plan(
                             "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,
@@ -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

@@ -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,6 +355,11 @@ class TLContext:
             pending = getattr(h, "pending", None)
             if pending is None or pending.resolved:
                 continue
+            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:
@@ -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
 

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=<TCM handle>` 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"

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) ──────────────────────────