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>
This commit is contained in:
2026-06-10 22:29:10 -07:00
parent dd525bfcb7
commit e8d6c283d8
7 changed files with 267 additions and 26 deletions
@@ -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,