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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user