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:
2026-06-10 19:48:29 -07:00
parent 47e2c78c66
commit 184f654295
6 changed files with 508 additions and 85 deletions
@@ -2,7 +2,7 @@
## Status ## Status
Proposed (Revision 2) Accepted (Revision 2)
> **ADR-0060** (AHBM GQA Fused Attention) 와 **ADR-0065** (flat-ops > **ADR-0060** (AHBM GQA Fused Attention) 와 **ADR-0065** (flat-ops
> composite + 첫 stateful recipe) 의 보조 ADR. 그 hybrid 의 핵심 — "GEMM 은 > composite + 첫 stateful recipe) 의 보조 ADR. 그 hybrid 의 핵심 — "GEMM 은
@@ -219,20 +219,14 @@ issue cost 를 0 → non-zero 로 바꾸면 **모든** bench 의 latency 변화.
이 ADR land 시 골든 latency **한 번** 재생성 — ADR-0062 D3 lazy-load 와 이 ADR land 시 골든 latency **한 번** 재생성 — ADR-0062 D3 lazy-load 와
같은 패턴. 재생성 후 동일 calibration 이 ADR-0065 opt2 측정에 적용. 같은 패턴. 재생성 후 동일 calibration 이 ADR-0065 opt2 측정에 적용.
### D7. Composite 크기 cap (deterministic segmentation) ### D7. Composite 크기 cap (하드 제한 — validation error)
`CompositeCmd``logical_bytes`**`MAX_COMPOSITE_LOGICAL_BYTES`** `CompositeCmd``logical_bytes`**`max_composite_logical_bytes`**
(default **1024 bytes**, D4 로 override) 로 제한. 초과하는 cmd 는 (default **1024 bytes**, D4 로 override) 로 제한. `logical_bytes` 가 cap
*emitter* (host-side TLContext, ADR-0065 D5) 가 N 개 연속 `CompositeCmd` 을 초과하는 composite 는 host-side TLContext 가 emit 시점에 `ValueError`
deterministic 하게 **segment** — 각 ≤ cap. **거부** **자동 segmentation 없음**. 커널 작성자가 각 `CompositeCmd`
가 descriptor capacity 에 맞도록 recipe 를 재구성(예: 여러 개의 더 작은
- 각 segment 가 자기 `completion: CompletionHandle` 보유. composite 로 명시적으로 분할)해야 함.
- 각 segment 가 자기 dispatch cost — `total = sum(FIXED + bytes_i × R) =
N × FIXED + total_bytes × R`. FIXED 항이 segment 당.
- segment 가 `rw_handles` 공유 — **strict FIFO** (ADR-0065 D6.3) 가
write-after-write 순서 자동 보존.
- segmentation 알고리즘 결정적 (emit 순서로 op 인덱스 greedy); host
emitter 의 일부, PE_SCHEDULER 가 아님.
**cap 이 필요한 이유.** 실제 하드웨어는 descriptor queue entry 크기, **cap 이 필요한 이유.** 실제 하드웨어는 descriptor queue entry 크기,
scheduler parser buffer, command SRAM, firmware 입력의 하드 제한 보유. scheduler parser buffer, command SRAM, firmware 입력의 하드 제한 보유.
@@ -248,18 +242,15 @@ composite (가장 큰 것이 decode opt2 의 `#2` ~322 bytes) 보다 훨씬 위
이 default 의 역할은 cap 을 *원칙으로 존재시키는 것*, 특정 HW 에 맞추는 이 default 의 역할은 cap 을 *원칙으로 존재시키는 것*, 특정 HW 에 맞추는
것이 아님. 것이 아님.
**segment 간 순서 — strict FIFO 가 주체, `rw_handles` 아님.** segment 들은 **자동 segmentation 대신 하드 에러를 택한 근거.** 초과 composite 를 자동
emit 순서로 PE_CPU → PE_SCHEDULER 큐에 들어감. strict-FIFO dispatch 분할(원래 Revision 2 제안)하는 것은 inter-segment 순서, `rw_handles`
(ADR-0065 D6.3) 가 *순서의 주체* — segment 가 큐 진입 순서로 실행. 각 공유, completion 체이닝 등 emitter 복잡도를 추가하면서, 실질적으로는
segment 의 `rw_handles` block 은 cross-composite 해저드 추적기의 **의존성 하드웨어가 허용하는 것보다 큰 descriptor 를 요청한 커널을 덮어주는
메타데이터** — 그 자체로 inter-segment 순서를 보장하지 **않음**. 미래 것일 뿐. 명시적 에러로 표면화하면 emitter 가 단순해지고, HW 제약이
scheduler 가 FIFO 를 완화 (예: RW-aware reorder, ADR-0065 A4) 한다면 작업을 어떻게 분할할지 가장 잘 아는 커널 작성자에게 보임.
segmenter 는 (a) segment 를 단일 CompositeCmd 로 merge 또는
(b) 명시적 completion-handle 의존 체인 도입 필요. Phase 1 에선 무관 —
strict FIFO 가 유효.
Decode opt2 의 `#2` composite (10 ops, ~322 bytes) 는 1024 cap 안에 Decode opt2 의 `#2` composite (10 ops, ~322 bytes) 는 1024 cap 안에
편안히 — GQA workload 에 segmentation 없음. 편안히 — GQA workload 에 에러 없음.
## Alternatives ## Alternatives
@@ -350,10 +341,8 @@ calibration 이 바뀌거나 OpSpec/CompositeCmd 필드가 추가되어도 유
7. **경로 parity.** greenlet 과 replay 가 동일 커널에 대해 동일 7. **경로 parity.** greenlet 과 replay 가 동일 커널에 대해 동일
dispatch-cycle 회계. dispatch-cycle 회계.
8. **Composite 크기 cap (D7).** `logical_bytes > 8. **Composite 크기 cap (D7).** `logical_bytes >
MAX_COMPOSITE_LOGICAL_BYTES` 가 될 recipe 가 N 개 연속 `CompositeCmd` max_composite_logical_bytes` 가 될 composite 는 emit 시점에 `ValueError`
로 segment; `sum(segment.logical_bytes) == original logical_bytes`; (segmentation 없음). cap 이내 composite 는 정상 emit.
총 dispatch = segment dispatch 의 합 (FIXED 가 segment 당 지불);
inter-segment 순서 strict FIFO 로 보존 (`rw_handles` 단독 아님).
9. **민감도 (정성적).** `R ∈ {0.25, 0.0625, 0.03125}` cycles/byte 에서 9. **민감도 (정성적).** `R ∈ {0.25, 0.0625, 0.03125}` cycles/byte 에서
3 지점 모두 `opt3 > opt2`. 방향 (R 감소 시 ratio 단조 증가) 도 단언, 3 지점 모두 `opt3 > opt2`. 방향 (R 감소 시 ratio 단조 증가) 도 단언,
단 절대 ratio 값은 *불필요*. 단 절대 ratio 값은 *불필요*.
@@ -366,11 +355,11 @@ ADR-0064 Revision 2 는 단일 PR 로 land:
- `pe_cpu.py` dispatch 경로의 공식 적용 - `pe_cpu.py` dispatch 경로의 공식 적용
- PE_CPU init 의 `pe_cost_model:` override read (cycle-domain knob + - PE_CPU init 의 `pe_cost_model:` override read (cycle-domain knob +
`max_composite_logical_bytes`) `max_composite_logical_bytes`)
- **composite 크기 cap (D7)** — TLContext-side segmentation 로직, - **composite 크기 cap (D7)** — TLContext-side 하드 cap,
`MAX_COMPOSITE_LOGICAL_BYTES` default 1024; 기존 bench 에 필요 없음 `max_composite_logical_bytes` default 1024: cap 초과 composite 는 emit
(현재 최대 composite 가 200 bytes 훨씬 아래), 하지만 ADR-0065 의 10-op 시점에 `ValueError` (자동 segmentation 없음). 기존 bench 는 트리거
decode-opt2 composite (~310 bytes) + 더 큰 recipe 가 나타날 때를 위해 안 됨 (현재 최대 composite 가 ~322 bytes), 하지만 recipe 가 커질 때
메커니즘 land. descriptor-capacity 제한이 강제되도록 체크 land.
- 일회성 골든 재생성 - 일회성 골든 재생성
land 후 ADR-0065 가 위에 쌓임 — 기존 bench 추가 골든 churn 없음 land 후 ADR-0065 가 위에 쌓임 — 기존 bench 추가 골든 churn 없음
+25 -38
View File
@@ -2,7 +2,7 @@
## Status ## Status
Proposed (Revision 2) Accepted (Revision 2)
> Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention) and **ADR-0065** > Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention) and **ADR-0065**
> (flat-ops composite + first stateful recipe). The hybrid decision there > (flat-ops composite + first stateful recipe). The hybrid decision there
@@ -234,22 +234,15 @@ latencies are **regenerated once** when this ADR lands — same posture as
ADR-0062 D3 lazy-load. After regeneration, the same calibration is in ADR-0062 D3 lazy-load. After regeneration, the same calibration is in
effect for ADR-0065 opt2 measurement. effect for ADR-0065 opt2 measurement.
### D7. Composite size cap (deterministic segmentation) ### D7. Composite size cap (hard limit — validation error)
Each `CompositeCmd`'s `logical_bytes` is capped at Each `CompositeCmd`'s `logical_bytes` is capped at
**`MAX_COMPOSITE_LOGICAL_BYTES`** (default **1024 bytes**, overridable **`max_composite_logical_bytes`** (default **1024 bytes**, overridable
per D4). Oversized commands are deterministically **segmented** by the per D4). A composite whose `logical_bytes` exceeds the cap is **rejected
*emitter* (host-side TLContext, ADR-0065 D5) into N consecutive at emit time** by the host-side TLContext with a `ValueError` — there is
`CompositeCmd`s, each ≤ cap. **no automatic segmentation**. The kernel author must restructure the
recipe (e.g., split it into multiple smaller composites explicitly) so
- Each segment carries its own `completion: CompletionHandle`. each `CompositeCmd` fits within the descriptor capacity.
- Each segment incurs its own dispatch cost — `total =
sum(FIXED + bytes_i × R) = N × FIXED + total_bytes × R`. The FIXED
term is paid per segment.
- Segments share `rw_handles` where applicable; **strict FIFO**
(ADR-0065 D6.3) preserves write-after-write ordering automatically.
- The segmentation algorithm is deterministic (greedy by op index in
emit order); it is part of the host emitter, not PE_SCHEDULER.
**Why a cap is needed.** Real hardware imposes hard limits — descriptor **Why a cap is needed.** Real hardware imposes hard limits — descriptor
queue entry size, scheduler parser buffer, command SRAM, firmware queue entry size, scheduler parser buffer, command SRAM, firmware
@@ -266,20 +259,16 @@ overridable per topology (D4); when a real HW reference appears, the
value should be recalibrated. The role of this default is to make the value should be recalibrated. The role of this default is to make the
cap *exist as a discipline*, not to fit a specific HW. cap *exist as a discipline*, not to fit a specific HW.
**Ordering of segments — driven by strict FIFO, not `rw_handles`.** **Rationale for a hard error over auto-segmentation.** Auto-splitting an
Segments are emitted into the PE_CPU → PE_SCHEDULER queue in their oversized composite (the original Revision 2 proposal) added emitter
emit order. Strict-FIFO dispatch (ADR-0065 D6.3) is the *ordering complexity — inter-segment ordering, shared `rw_handles`, completion
source*: segments execute in the order they enter the queue. The chaining — to paper over what is, in practice, a kernel that asked for a
`rw_handles` block on each segment is **dependency metadata** for the descriptor larger than the hardware allows. Surfacing it as an explicit
cross-composite hazard tracker — it does **not** by itself guarantee error keeps the emitter simple and makes the HW constraint visible to the
inter-segment ordering. If a future scheduler relaxed FIFO (e.g., to kernel author, who is best placed to decide how to split the work.
RW-aware reorder, ADR-0065 A4), the segmenter would need to either
(a) merge segments into a single CompositeCmd, or (b) introduce an
explicit completion-handle dependency chain. For Phase 1 this does
not arise: strict FIFO is in effect.
Decode opt2's `#2` composite (10 ops, ~322 bytes) sits comfortably Decode opt2's `#2` composite (10 ops, ~322 bytes) sits comfortably
inside the 1024 cap — no segmentation for the GQA workload. inside the 1024 cap — no error for the GQA workload.
## Alternatives ## Alternatives
@@ -379,12 +368,9 @@ OpSpec/CompositeCmd fields are added.
(SPEC §0.1). (SPEC §0.1).
7. **Path parity.** Greenlet and replay paths produce identical 7. **Path parity.** Greenlet and replay paths produce identical
dispatch-cycle accounting for the same kernel. dispatch-cycle accounting for the same kernel.
8. **Composite size cap (D7).** A recipe that would emit `logical_bytes 8. **Composite size cap (D7).** A composite that would emit `logical_bytes
> MAX_COMPOSITE_LOGICAL_BYTES` is segmented into N consecutive > max_composite_logical_bytes` raises a `ValueError` at emit time (no
`CompositeCmd`s; `sum(segment.logical_bytes) == original segmentation). A composite within the cap emits normally.
logical_bytes`; total dispatch = sum of segment dispatches (FIXED
paid per segment); inter-segment ordering preserved by strict FIFO
(not by `rw_handles` alone).
9. **Sensitivity (qualitative).** At `R ∈ {0.25, 0.0625, 0.03125}` 9. **Sensitivity (qualitative).** At `R ∈ {0.25, 0.0625, 0.03125}`
cycles/byte, `opt3 > opt2` at *all* three points. Direction (ratio cycles/byte, `opt3 > opt2` at *all* three points. Direction (ratio
monotonically increases as R decreases) is also asserted, but monotonically increases as R decreases) is also asserted, but
@@ -398,11 +384,12 @@ ADR-0064 Revision 2 lands as a single PR with:
- formula application in `pe_cpu.py` dispatch path - formula application in `pe_cpu.py` dispatch path
- `pe_cost_model:` override read at PE_CPU init (cycle-domain knobs + - `pe_cost_model:` override read at PE_CPU init (cycle-domain knobs +
`max_composite_logical_bytes`) `max_composite_logical_bytes`)
- **composite size cap (D7)** — TLContext-side segmentation logic with - **composite size cap (D7)** — TLContext-side hard cap with
`MAX_COMPOSITE_LOGICAL_BYTES` default 1024; not needed for any `max_composite_logical_bytes` default 1024: a composite over the cap
existing bench (largest current composite is well under 200 bytes), raises a `ValueError` at emit time (no auto-segmentation). Not
but the mechanism lands so it is in place when ADR-0065's triggered by any existing bench (largest current composite is ~322
10-op decode-opt2 composite (~310 bytes) and larger recipes appear. bytes), but the check lands so the descriptor-capacity limit is
enforced as recipes grow.
- one-time goldens regeneration - one-time goldens regeneration
After this lands, ADR-0065 builds on top with no further goldens churn After this lands, ADR-0065 builds on top with no further goldens churn
+179 -14
View File
@@ -735,28 +735,58 @@ class TLContext:
def composite( def composite(
self, self,
op: Literal["gemm", "math"], op: Literal["gemm", "math"],
a: TensorHandle, a: TensorHandle | None = None,
b: TensorHandle | None = None, b: TensorHandle | None = None,
out_ptr: int = 0, out_ptr: int = 0,
math_op: str | None = None, math_op: str | None = None,
*, *,
out: TensorHandle | None = None,
prologue: list[dict] | None = None,
epilogue: list[dict] | None = None, epilogue: list[dict] | None = None,
acc_dtype: str | None = None, acc_dtype: str | None = None,
tile_shape: tuple[int, int, int] | None = None, tile_shape: tuple[int, int, int] | None = None,
) -> CompletionHandle: ) -> CompletionHandle:
"""Submit a composite command (non-blocking, tiled pipeline). """Submit a composite command (non-blocking, tiled pipeline).
Optional ``epilogue`` is an ordered list of dicts; each dict has a ``prologue`` (ADR-0065 D5) is an ordered list of recipe dicts each
required ``"op"`` key (one of ``EPILOGUE_OPS``) plus op-specific with an ``"op"`` naming a ``RECIPE_DESCRIPTORS`` entry plus its
fields and an optional ``"scope"``. Validation happens here so operands. TLContext expands each into flat MATH OpSpecs that run
typos fail before the command is emitted. 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(). Returns CompletionHandle for use with wait().
""" """
# ADR-0062: composite operand DMA paths still need their inputs # ADR-0062: composite operand DMA paths still need their inputs
# to be resolved before the composite reads them via PE_SCHEDULER. # to be resolved before the composite reads them via PE_SCHEDULER.
self._await_pending(a, b) 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: if op == "gemm" and b is not None:
m, k = a.shape[-2], a.shape[-1] m, k = a.shape[-2], a.shape[-1]
n = b.shape[-1] n = b.shape[-1]
@@ -768,12 +798,15 @@ class TLContext:
out_shape = a.shape out_shape = a.shape
out_nbytes = a.nbytes out_nbytes = a.nbytes
# Head op's write-back handle carries the output address (ADR-0065 D1). # Write-back handle: explicit `out` handle wins over `out_ptr`.
out_handle = TensorHandle( if out is not None:
id=self._next_handle_id(), out_handle = out
addr=out_ptr, shape=out_shape, dtype=out_dtype, else:
nbytes=out_nbytes, space="tcm", 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 # 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. # 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) epi_specs = tuple(self._build_epilogue_spec(e, i)
for i, e in enumerate(epilogue or [])) 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()) 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 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 @staticmethod
def _build_epilogue_spec(entry: dict, idx: int) -> OpSpec: def _build_epilogue_spec(entry: dict, idx: int) -> OpSpec:
if not isinstance(entry, dict) or "op" not in entry: if not isinstance(entry, dict) or "op" not in entry:
+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"),
),
),
}
+87
View File
@@ -0,0 +1,87 @@
"""Phase 1 spec tests for ADR-0065 P2 — composite emit-time validation.
Covers the two reachable invariants enforced by the TLContext lowering pass
(ADR-0065 D6.6 / D6.7):
- **Auto-bind conflict (D6.6).** If a prologue recipe declares a
`primary_out` (which auto-binds into the head GEMM's `a`) AND the kernel
also passes `a` explicitly, lowering raises — ambiguous which value wins.
- **MATH operand TCM-only (D6.7).** Every operand of a *prologue recipe*
MATH op must be TCM-resident; passing an HBM handle as a recipe operand
raises at emit time. (The head op — gemm *or* math — keeps the existing
DMA-staged-from-HBM behavior; D6.7 does not apply to it.)
(D6.1 GEMM-count ≤ 1 is implemented as a defensive guard but is not
reachable through the public API while `softmax_merge` is MATH-only and the
head carries a single GEMM — its dedicated test is deferred until a
GEMM-emitting recipe exists.)
Phase 1 (this commit): tests only. FAIL until P2.
"""
from __future__ import annotations
import math
import pytest
from kernbench.common.pe_commands import TensorHandle
from kernbench.triton_emu.tl_context import TLContext
G, TILE, D = 8, 64, 128
def _tcm(addr: int, shape: tuple[int, ...]) -> TensorHandle:
return TensorHandle(
id=f"h{addr:x}", addr=addr, shape=shape, dtype="f16",
nbytes=2 * math.prod(shape), space="tcm",
)
def _hbm(addr: int, shape: tuple[int, ...]) -> TensorHandle:
return TensorHandle(
id=f"hb{addr:x}", addr=addr, shape=shape, dtype="f16",
nbytes=2 * math.prod(shape), space="hbm",
)
def test_autobind_conflict_raises(monkeypatch=None):
"""D6.6 — explicit `a` + a primary_out recipe is ambiguous."""
tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000)
A = _tcm(0x9000, (G, TILE))
s = _tcm(0x1000, (G, TILE))
m = _tcm(0x2000, (G,))
l = _tcm(0x3000, (G,))
O = _tcm(0x4000, (G, D))
V = tl.ref(0x5000, shape=(TILE, D), dtype="f16")
with pytest.raises(ValueError, match="auto-bind"):
tl.composite(
op="gemm", a=A, b=V, out=O,
prologue=[{"op": "softmax_merge", "s": s, "m": m, "l": l, "O": O}],
)
def test_math_operand_must_be_tcm_raises():
"""D6.7 — an HBM handle reaching a prologue recipe MATH op is rejected.
softmax_merge's first op (rmax) reads ``s``; an HBM ``s`` makes the
expanded MATH op reference HBM, which the recipe prohibits.
"""
tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000)
s_hbm = _hbm(0x1000, (G, TILE)) # Sj as HBM — illegal recipe operand
m = _tcm(0x2000, (G,))
l = _tcm(0x3000, (G,))
O = _tcm(0x4000, (G, D))
V = tl.ref(0x5000, shape=(TILE, D), dtype="f16")
with pytest.raises(ValueError, match="TCM"):
tl.composite(
op="gemm", b=V, out=O,
prologue=[{"op": "softmax_merge", "s": s_hbm, "m": m, "l": l, "O": O}],
)
def test_math_head_from_hbm_is_allowed():
"""The head op (op="math") keeps the existing HBM-streamed behavior —
D6.7 does not apply to it (only to prologue recipe ops)."""
tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000)
a_hbm = _hbm(0x1000, (G, TILE))
tl.composite(op="math", a=a_hbm, math_op="exp", out_ptr=0x6000) # no raise
+111
View File
@@ -0,0 +1,111 @@
"""Phase 1 spec tests for ADR-0065 P2 — `softmax_merge` recipe lowering.
P2 adds `tl_recipes.py` (RECIPE_DESCRIPTORS) and a TLContext lowering pass
so that
tl.composite(
op="gemm", b=V_ref, out=O,
prologue=[{"op": "softmax_merge", "s": Sj, "m": m, "l": l, "O": O}],
epilogue=[{"op": "add", "other": O}],
)
expands the single `softmax_merge` recipe into 8 flat MATH OpSpecs, binds the
recipe's primary output `P` into the head GEMM's `a` operand, and emits one
`CompositeCmd` with 10 ops + `rw_handles == (m, l, O)` (ADR-0065 D5 / DDD §6).
This is host-side lowering only — PE_SCHEDULER consumption of the pre-GEMM
KERNEL-scope ops is P3, so these tests inspect the lowered CompositeCmd
object and do NOT run it through the scheduler.
Phase 1 (this commit): tests only. FAIL until P2:
- `kernbench.triton_emu.tl_recipes` does not exist,
- `TLContext.composite` has no `prologue=` / `out=TensorHandle` kwargs.
"""
from __future__ import annotations
import math
from kernbench.common.pe_commands import CompositeCmd, TensorHandle
from kernbench.triton_emu.tl_context import TLContext
G, TILE, D = 8, 64, 128
_EXPECTED_KINDS = [
"rmax", "max_elem", "exp_diff", "exp_diff", "rsum", "fma",
"mul_bcast", "copy", # 8 MATH (softmax_merge)
"gemm", # head GEMM (P · V)
"add", # epilogue
]
def _tcm(addr: int, shape: tuple[int, ...]) -> TensorHandle:
return TensorHandle(
id=f"h{addr:x}", addr=addr, shape=shape, dtype="f16",
nbytes=2 * math.prod(shape), space="tcm",
)
def _lower():
"""Lower decode-opt2 #2 and return (tl, cmd, handles)."""
tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000,
scratch_size=1 << 20)
s = _tcm(0x1000, (G, TILE)) # Sj — #1's output, TCM-resident
m = _tcm(0x2000, (G,))
l = _tcm(0x3000, (G,))
O = _tcm(0x4000, (G, D))
V = tl.ref(0x5000, shape=(TILE, D), dtype="f16")
tl.composite(
op="gemm", b=V, out=O,
prologue=[{"op": "softmax_merge", "s": s, "m": m, "l": l, "O": O}],
epilogue=[{"op": "add", "other": O}],
)
cmd = [c for c in tl.commands if isinstance(c, CompositeCmd)][-1]
return tl, cmd, dict(s=s, m=m, l=l, O=O, V=V)
# ── recipe shape (ADR-0065 Test #2) ──────────────────────────────────
def test_recipe_lowers_to_ten_ops_in_order():
_, cmd, _ = _lower()
assert [o.kind for o in cmd.ops] == _EXPECTED_KINDS, [o.kind for o in cmd.ops]
def test_recipe_rw_handles_are_m_l_O():
_, cmd, h = _lower()
assert cmd.rw_handles == (h["m"], h["l"], h["O"]), cmd.rw_handles
def test_head_gemm_autobinds_primary_out_and_geometry():
_, cmd, h = _lower()
gemm = cmd.ops[8]
assert gemm.kind == "gemm"
# b is V (the kernel-provided operand); a is the recipe's primary out P
assert gemm.operands["b"] == h["V"]
assert "a" in gemm.operands and gemm.operands["a"].shape == (G, TILE)
assert gemm.out is not None and gemm.out.addr == h["O"].addr
assert gemm.extra.get("m") == G
assert gemm.extra.get("k") == TILE
assert gemm.extra.get("n") == D
def test_recipe_scratch_addresses_are_concrete_and_distinct():
_, cmd, _ = _lower()
# Collect every operand/out handle address across the 8 MATH ops.
addrs = set()
for op in cmd.ops[:8]:
for hd in op.operands.values():
assert hd.addr is not None
addrs.add(hd.addr)
if op.out is not None:
addrs.add(op.out.addr)
# The 5 scratch slots (m_loc, m_new, corr, P, l_loc) must have been
# allocated to distinct, non-zero addresses (scratch_base supplied).
scratch_addrs = {a for a in addrs if a >= 0x100000}
assert len(scratch_addrs) >= 5, scratch_addrs
def test_recipe_composite_within_size_cap():
"""ADR-0065 Test #9 — the 10-op composite fits the default 1024 cap."""
_, cmd, _ = _lower()
assert cmd.logical_bytes < 1024, cmd.logical_bytes