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
+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