ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gqa(adr-0065): P2 — softmax_merge recipe + TLContext prologue lowering

Browse Source 
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:
Yangwook
2026-06-10 19:48:29 -07:00
parent 47e2c78c66
commit 184f654295
6 changed files with 508 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernbench/triton_emu/tl_context.py
+179 -14
View File
@@ -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:
src/kernbench/triton_emu/tl_recipes.py
+84
View File
@@ -0,0 +1,84 @@
"""TLContext recipe descriptors (ADR-0065 D5).
A *recipe* is a named macro op (e.g. ``softmax_merge``) that TLContext
expands, at composite-emit time, into a sequence of flat ``OpSpec``s with
all addresses / sizes filled. The recipe table lives here, beside
``tl_context.py`` (the compiler analog). **PE_SCHEDULER never imports this
module** — it sees only the flat ops list (ADR-0065 D5 / D7 boundary).
Slot names used in ``EngineOp`` are either recipe operand names
(``RecipeDescriptor.operands``) or internal scratch slots produced by an
earlier op in ``engine_seq``. The scratch slot named ``"P"`` is, by
convention, the recipe's *primary output* — auto-bound into the head GEMM's
first matrix operand by the lowering pass.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Callable, Literal
PRIMARY_OUT_SLOT = "P"
@dataclass(frozen=True)
class PrimaryOutSpec:
"""Spec for a recipe's implicit primary output (slot ``"P"``)."""
from_shape: str # operand name to copy shape from
from_dtype: str # operand name to copy dtype from
transform: str # "identity" | "trans" (forward-compat)
@dataclass(frozen=True)
class EngineOp:
"""One micro-op in a recipe's ``engine_seq``.
Slot fields (``src``, ``src_a``, ``src_b``, ``src_c``, ``dst``) name
either ``RecipeDescriptor.operands`` or internal scratch slots.
"""
engine: Literal["MATH", "GEMM", "DMA"]
op_kind: str # PE_MATH op_kind label (opaque to the engine)
src: str | None = None
src_a: str | None = None
src_b: str | None = None
src_c: str | None = None
dst: str | None = None
reduce_axis: int | None = None
bcast_axis: int | None = None
@dataclass(frozen=True)
class RecipeDescriptor:
operands: dict[str, str] # name → "R" | "RW"
primary_out: PrimaryOutSpec | None # implicit output (slot "P")
tile_alignment: Literal["single_shot", "tile_aligned"]
internal_scratch_bytes_fn: Callable[..., int]
engine_seq: tuple[EngineOp, ...]
RECIPE_DESCRIPTORS: dict[str, RecipeDescriptor] = {
# Online-softmax merge (flash-attention accumulator update). Reads the
# current score tile ``s`` and the running ``(m, l, O)``; emits the
# rescaled probabilities ``P`` (slot "P") for the downstream P·V GEMM.
"softmax_merge": RecipeDescriptor(
operands={"s": "R", "m": "RW", "l": "RW", "O": "RW"},
primary_out=PrimaryOutSpec(
from_shape="s", from_dtype="s", transform="identity",
),
tile_alignment="single_shot",
internal_scratch_bytes_fn=lambda G, TILE, d, bpe: bpe * (
G + G + G + G * TILE + G # m_loc + m_new + corr + P + l_loc
),
engine_seq=(
EngineOp("MATH", "rmax", src="s", dst="m_loc", reduce_axis=-1),
EngineOp("MATH", "max_elem", src_a="m", src_b="m_loc", dst="m_new"),
EngineOp("MATH", "exp_diff", src_a="m", src_b="m_new", dst="corr"),
EngineOp("MATH", "exp_diff", src_a="s", src_b="m_new", dst="P", bcast_axis=0),
EngineOp("MATH", "rsum", src="P", dst="l_loc", reduce_axis=-1),
EngineOp("MATH", "fma", src_a="l", src_b="corr", src_c="l_loc", dst="l"),
EngineOp("MATH", "mul_bcast", src_a="O", src_b="corr", dst="O", bcast_axis=1),
EngineOp("MATH", "copy", src="m_new", dst="m"),
),
),
}