gqa(adr-0065): P2 — softmax_merge recipe + TLContext prologue lowering
New triton_emu/tl_recipes.py: RecipeDescriptor/EngineOp/PrimaryOutSpec + RECIPE_DESCRIPTORS['softmax_merge'] (8-step engine_seq). PE_SCHEDULER does not import it (ADR-0065 D5 boundary). TLContext.composite(): add prologue=[...] + out=TensorHandle kwargs, a optional. _expand_prologue lowers a recipe into flat MATH OpSpecs (scope=KERNEL), allocates TCM scratch, derives the primary-out slot 'P' and auto-binds it into the head GEMM a (D6.6 conflict check); rw_handles=(m,l,O). decode-opt2 #2 lowers to 10 ops [rmax,max_elem,exp_diff,exp_diff,rsum,fma,mul_bcast,copy,gemm,add]. D6.7 (MATH operand TCM-only) scoped to prologue recipe ops only — the head op (gemm or math) keeps existing DMA-staged-from-HBM behavior. D6.1 (GEMM count <=1) on the whole composite. Host-side lowering only; PE_SCHEDULER position-scan is P3. Also commit the ADR-0064 Rev2 promotion content that the prior commit's git mv dropped: Status Proposed->Accepted + D7 amended to hard-cap ValueError (no segmentation), EN+KO. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -735,28 +735,58 @@ class TLContext:
|
||||
def composite(
|
||||
self,
|
||||
op: Literal["gemm", "math"],
|
||||
a: TensorHandle,
|
||||
a: TensorHandle | None = None,
|
||||
b: TensorHandle | None = None,
|
||||
out_ptr: int = 0,
|
||||
math_op: str | None = None,
|
||||
*,
|
||||
out: TensorHandle | 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:
|
||||
"""Submit a composite command (non-blocking, tiled pipeline).
|
||||
|
||||
Optional ``epilogue`` is an ordered list of dicts; each dict has a
|
||||
required ``"op"`` key (one of ``EPILOGUE_OPS``) plus op-specific
|
||||
fields and an optional ``"scope"``. Validation happens here so
|
||||
typos fail before the command is emitted.
|
||||
``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.
|
||||
|
||||
Returns CompletionHandle for use with wait().
|
||||
"""
|
||||
# ADR-0062: composite operand DMA paths still need their inputs
|
||||
# to be resolved before the composite reads them via PE_SCHEDULER.
|
||||
self._await_pending(a, b)
|
||||
# Compute output geometry based on op.
|
||||
|
||||
# Prologue recipes → flat MATH OpSpecs + primary-output handle (D5).
|
||||
prologue_ops: tuple[OpSpec, ...] = ()
|
||||
primary_out: TensorHandle | None = None
|
||||
rw_handles: tuple[TensorHandle, ...] = ()
|
||||
if prologue:
|
||||
prologue_ops, primary_out, rw_handles = self._expand_prologue(prologue)
|
||||
|
||||
# Auto-bind (D6.6): the recipe's primary output fills the head GEMM's
|
||||
# `a` unless the kernel supplied `a` explicitly (then it's ambiguous).
|
||||
if op == "gemm" and primary_out is not None:
|
||||
if a is not None:
|
||||
raise ValueError(
|
||||
"auto-bind conflict: head GEMM operand 'a' was provided "
|
||||
"explicitly but the prologue recipe declares a primary_out "
|
||||
"(ADR-0065 D6.6) — omit 'a' to let the recipe bind it"
|
||||
)
|
||||
a = primary_out
|
||||
|
||||
if a is None:
|
||||
raise ValueError(
|
||||
"composite requires operand 'a' (or a prologue recipe with a "
|
||||
"primary_out)"
|
||||
)
|
||||
|
||||
# Output geometry.
|
||||
if op == "gemm" and b is not None:
|
||||
m, k = a.shape[-2], a.shape[-1]
|
||||
n = b.shape[-1]
|
||||
@@ -768,12 +798,15 @@ class TLContext:
|
||||
out_shape = a.shape
|
||||
out_nbytes = a.nbytes
|
||||
|
||||
# Head op's write-back handle carries the output address (ADR-0065 D1).
|
||||
out_handle = TensorHandle(
|
||||
id=self._next_handle_id(),
|
||||
addr=out_ptr, shape=out_shape, dtype=out_dtype,
|
||||
nbytes=out_nbytes, space="tcm",
|
||||
)
|
||||
# Write-back handle: explicit `out` handle wins over `out_ptr`.
|
||||
if out is not None:
|
||||
out_handle = out
|
||||
else:
|
||||
out_handle = TensorHandle(
|
||||
id=self._next_handle_id(),
|
||||
addr=out_ptr, shape=out_shape, dtype=out_dtype,
|
||||
nbytes=out_nbytes, space="tcm",
|
||||
)
|
||||
|
||||
# Head op (flat-ops, ADR-0065 D2): GEMM takes named a/b and carries
|
||||
# m/k/n in extra; MATH takes named a and carries math_op in extra.
|
||||
@@ -795,12 +828,144 @@ class TLContext:
|
||||
)
|
||||
epi_specs = tuple(self._build_epilogue_spec(e, i)
|
||||
for i, e in enumerate(epilogue or []))
|
||||
ops_tuple = (head_spec, *epi_specs)
|
||||
ops_tuple = (*prologue_ops, head_spec, *epi_specs)
|
||||
|
||||
self._validate_composite_ops(ops_tuple, prologue_ops)
|
||||
|
||||
completion = CompletionHandle(id=self._next_completion_id())
|
||||
self._emit(CompositeCmd(completion=completion, ops=ops_tuple))
|
||||
self._emit(CompositeCmd(
|
||||
completion=completion, ops=ops_tuple, rw_handles=rw_handles,
|
||||
))
|
||||
return completion
|
||||
|
||||
def _expand_prologue(
|
||||
self, prologue: list[dict],
|
||||
) -> tuple[tuple[OpSpec, ...], TensorHandle | None, tuple[TensorHandle, ...]]:
|
||||
"""Expand prologue recipe items into flat MATH OpSpecs (ADR-0065 D5).
|
||||
|
||||
Returns ``(math_ops, primary_out_handle, rw_handles)``. PE_SCHEDULER
|
||||
never sees the recipe — only the flat ops.
|
||||
"""
|
||||
from kernbench.triton_emu.tl_recipes import (
|
||||
PRIMARY_OUT_SLOT, RECIPE_DESCRIPTORS,
|
||||
)
|
||||
|
||||
math_ops: list[OpSpec] = []
|
||||
primary_out: TensorHandle | None = None
|
||||
rw_handles: list[TensorHandle] = []
|
||||
|
||||
for item in prologue:
|
||||
if not isinstance(item, dict) or "op" not in item:
|
||||
raise ValueError("prologue entry must be a dict with an 'op' key")
|
||||
name = item["op"]
|
||||
recipe = RECIPE_DESCRIPTORS.get(name)
|
||||
if recipe is None:
|
||||
known = ", ".join(sorted(RECIPE_DESCRIPTORS))
|
||||
raise ValueError(
|
||||
f"unknown prologue recipe {name!r} (known: {known})"
|
||||
)
|
||||
|
||||
# Bind recipe operands from the item.
|
||||
slots: dict[str, TensorHandle] = {}
|
||||
for opd_name in recipe.operands:
|
||||
if opd_name not in item:
|
||||
raise ValueError(
|
||||
f"prologue recipe {name!r} missing operand {opd_name!r}"
|
||||
)
|
||||
slots[opd_name] = item[opd_name]
|
||||
|
||||
# Expand engine_seq sequentially — later ops read earlier dsts.
|
||||
for eop in recipe.engine_seq:
|
||||
operands: dict[str, Any] = {}
|
||||
for key in ("src", "src_a", "src_b", "src_c"):
|
||||
slot_name = getattr(eop, key)
|
||||
if slot_name is None:
|
||||
continue
|
||||
if slot_name not in slots:
|
||||
raise ValueError(
|
||||
f"recipe {name!r} op {eop.op_kind!r} reads unbound "
|
||||
f"slot {slot_name!r}"
|
||||
)
|
||||
operands[key] = slots[slot_name]
|
||||
|
||||
dst_handle = self._resolve_recipe_dst(eop, recipe, slots, operands)
|
||||
extra: dict[str, Any] = {}
|
||||
if eop.reduce_axis is not None:
|
||||
extra["reduce_axis"] = eop.reduce_axis
|
||||
if eop.bcast_axis is not None:
|
||||
extra["bcast_axis"] = eop.bcast_axis
|
||||
|
||||
math_ops.append(OpSpec(
|
||||
kind=eop.op_kind, scope=Scope.KERNEL,
|
||||
operands=operands, extra=extra, out=dst_handle,
|
||||
))
|
||||
slots[eop.dst] = dst_handle
|
||||
|
||||
if recipe.primary_out is not None:
|
||||
primary_out = slots.get(PRIMARY_OUT_SLOT)
|
||||
for opd_name, rw in recipe.operands.items():
|
||||
if rw == "RW":
|
||||
rw_handles.append(slots[opd_name])
|
||||
|
||||
return tuple(math_ops), primary_out, tuple(rw_handles)
|
||||
|
||||
def _resolve_recipe_dst(
|
||||
self, eop: Any, recipe: Any, slots: dict, operands: dict,
|
||||
) -> TensorHandle:
|
||||
"""Resolve an EngineOp's dst to a handle: reuse an existing slot
|
||||
(RW operand / already-allocated), else allocate fresh TCM scratch
|
||||
with an inferred shape."""
|
||||
from kernbench.triton_emu.tl_recipes import PRIMARY_OUT_SLOT
|
||||
|
||||
dst_name = eop.dst
|
||||
if dst_name in slots:
|
||||
return slots[dst_name] # write in place
|
||||
|
||||
if recipe.primary_out is not None and dst_name == PRIMARY_OUT_SLOT:
|
||||
ref = slots[recipe.primary_out.from_shape]
|
||||
shape, dtype = ref.shape, ref.dtype
|
||||
elif eop.reduce_axis is not None:
|
||||
ref = next(iter(operands.values()))
|
||||
shape = ref.shape[:-1] if len(ref.shape) > 1 else ref.shape
|
||||
dtype = ref.dtype
|
||||
else:
|
||||
ref = next(iter(operands.values()))
|
||||
shape, dtype = ref.shape, ref.dtype
|
||||
|
||||
nbytes = self._nbytes(shape, dtype)
|
||||
addr = self._scratch_alloc(nbytes)
|
||||
return TensorHandle(
|
||||
id=self._next_handle_id(),
|
||||
addr=addr, shape=shape, dtype=dtype, nbytes=nbytes, space="tcm",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _validate_composite_ops(
|
||||
ops: tuple[OpSpec, ...], prologue_ops: tuple[OpSpec, ...],
|
||||
) -> None:
|
||||
"""Emit-time invariants (ADR-0065 D6.1 / D6.7).
|
||||
|
||||
D6.1 (GEMM count ≤ 1) applies to the whole composite. D6.7 (MATH
|
||||
operands TCM-only) applies to **prologue recipe ops only** — the
|
||||
head op (gemm *or* math) and epilogue ops keep the existing
|
||||
DMA-staged-from-HBM behavior (a math head's `a` is streamed by
|
||||
PE_SCHEDULER, like a GEMM operand).
|
||||
"""
|
||||
gemm_count = sum(1 for o in ops if o.kind == "gemm")
|
||||
if gemm_count > 1:
|
||||
raise ValueError(
|
||||
f"composite must carry at most 1 GEMM op (ADR-0065 D6.1); "
|
||||
f"got {gemm_count}"
|
||||
)
|
||||
for o in prologue_ops:
|
||||
for h in o.operands.values():
|
||||
if getattr(h, "space", "tcm") != "tcm":
|
||||
raise ValueError(
|
||||
f"MATH operand must be TCM-resident (ADR-0065 D6.7); "
|
||||
f"prologue op {o.kind!r} operand has space="
|
||||
f"{getattr(h, 'space', None)!r}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_epilogue_spec(entry: dict, idx: int) -> OpSpec:
|
||||
if not isinstance(entry, dict) or "op" not in entry:
|
||||
|
||||
Reference in New Issue
Block a user