From 184f6542956a62478966ea519d4e05f8e740e7ec Mon Sep 17 00:00:00 2001 From: Yangwook Date: Wed, 10 Jun 2026 19:48:29 -0700 Subject: [PATCH] =?UTF-8?q?gqa(adr-0065):=20P2=20=E2=80=94=20softmax=5Fmer?= =?UTF-8?q?ge=20recipe=20+=20TLContext=20prologue=20lowering?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../ADR-0064-perf-cpu-issue-cost-model.md | 55 ++--- .../adr/ADR-0064-perf-cpu-issue-cost-model.md | 63 +++--- src/kernbench/triton_emu/tl_context.py | 193 ++++++++++++++++-- src/kernbench/triton_emu/tl_recipes.py | 84 ++++++++ tests/test_tl_composite_validation.py | 87 ++++++++ tests/test_tl_recipe_softmax_merge.py | 111 ++++++++++ 6 files changed, 508 insertions(+), 85 deletions(-) create mode 100644 src/kernbench/triton_emu/tl_recipes.py create mode 100644 tests/test_tl_composite_validation.py create mode 100644 tests/test_tl_recipe_softmax_merge.py diff --git a/docs/adr-ko/ADR-0064-perf-cpu-issue-cost-model.md b/docs/adr-ko/ADR-0064-perf-cpu-issue-cost-model.md index eca6000..fee0146 100644 --- a/docs/adr-ko/ADR-0064-perf-cpu-issue-cost-model.md +++ b/docs/adr-ko/ADR-0064-perf-cpu-issue-cost-model.md @@ -2,7 +2,7 @@ ## Status -Proposed (Revision 2) +Accepted (Revision 2) > **ADR-0060** (AHBM GQA Fused Attention) 와 **ADR-0065** (flat-ops > 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 와 같은 패턴. 재생성 후 동일 calibration 이 ADR-0065 opt2 측정에 적용. -### D7. Composite 크기 cap (deterministic segmentation) +### D7. Composite 크기 cap (하드 제한 — validation error) -각 `CompositeCmd` 의 `logical_bytes` 가 **`MAX_COMPOSITE_LOGICAL_BYTES`** -(default **1024 bytes**, D4 로 override) 로 제한. 초과하는 cmd 는 -*emitter* (host-side TLContext, ADR-0065 D5) 가 N 개 연속 `CompositeCmd` -로 deterministic 하게 **segment** — 각 ≤ cap. - -- 각 segment 가 자기 `completion: CompletionHandle` 보유. -- 각 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 가 아님. +각 `CompositeCmd` 의 `logical_bytes` 가 **`max_composite_logical_bytes`** +(default **1024 bytes**, D4 로 override) 로 제한. `logical_bytes` 가 cap +을 초과하는 composite 는 host-side TLContext 가 emit 시점에 `ValueError` +로 **거부** — **자동 segmentation 없음**. 커널 작성자가 각 `CompositeCmd` +가 descriptor capacity 에 맞도록 recipe 를 재구성(예: 여러 개의 더 작은 +composite 로 명시적으로 분할)해야 함. **cap 이 필요한 이유.** 실제 하드웨어는 descriptor queue entry 크기, scheduler parser buffer, command SRAM, firmware 입력의 하드 제한 보유. @@ -248,18 +242,15 @@ composite (가장 큰 것이 decode opt2 의 `#2` ~322 bytes) 보다 훨씬 위 이 default 의 역할은 cap 을 *원칙으로 존재시키는 것*, 특정 HW 에 맞추는 것이 아님. -**segment 간 순서 — strict FIFO 가 주체, `rw_handles` 아님.** segment 들은 -emit 순서로 PE_CPU → PE_SCHEDULER 큐에 들어감. strict-FIFO dispatch -(ADR-0065 D6.3) 가 *순서의 주체* — segment 가 큐 진입 순서로 실행. 각 -segment 의 `rw_handles` block 은 cross-composite 해저드 추적기의 **의존성 -메타데이터** — 그 자체로 inter-segment 순서를 보장하지 **않음**. 미래 -scheduler 가 FIFO 를 완화 (예: RW-aware reorder, ADR-0065 A4) 한다면 -segmenter 는 (a) segment 를 단일 CompositeCmd 로 merge 또는 -(b) 명시적 completion-handle 의존 체인 도입 필요. Phase 1 에선 무관 — -strict FIFO 가 유효. +**자동 segmentation 대신 하드 에러를 택한 근거.** 초과 composite 를 자동 +분할(원래 Revision 2 제안)하는 것은 inter-segment 순서, `rw_handles` +공유, completion 체이닝 등 emitter 복잡도를 추가하면서, 실질적으로는 +하드웨어가 허용하는 것보다 큰 descriptor 를 요청한 커널을 덮어주는 +것일 뿐. 명시적 에러로 표면화하면 emitter 가 단순해지고, HW 제약이 +작업을 어떻게 분할할지 가장 잘 아는 커널 작성자에게 보임. Decode opt2 의 `#2` composite (10 ops, ~322 bytes) 는 1024 cap 안에 -편안히 — GQA workload 에 segmentation 없음. +편안히 — GQA workload 에 에러 없음. ## Alternatives @@ -350,10 +341,8 @@ calibration 이 바뀌거나 OpSpec/CompositeCmd 필드가 추가되어도 유 7. **경로 parity.** greenlet 과 replay 가 동일 커널에 대해 동일 dispatch-cycle 회계. 8. **Composite 크기 cap (D7).** `logical_bytes > - MAX_COMPOSITE_LOGICAL_BYTES` 가 될 recipe 가 N 개 연속 `CompositeCmd` - 로 segment; `sum(segment.logical_bytes) == original logical_bytes`; - 총 dispatch = segment dispatch 의 합 (FIXED 가 segment 당 지불); - inter-segment 순서 strict FIFO 로 보존 (`rw_handles` 단독 아님). + max_composite_logical_bytes` 가 될 composite 는 emit 시점에 `ValueError` + (segmentation 없음). cap 이내 composite 는 정상 emit. 9. **민감도 (정성적).** `R ∈ {0.25, 0.0625, 0.03125}` cycles/byte 에서 3 지점 모두 `opt3 > opt2`. 방향 (R 감소 시 ratio 단조 증가) 도 단언, 단 절대 ratio 값은 *불필요*. @@ -366,11 +355,11 @@ ADR-0064 Revision 2 는 단일 PR 로 land: - `pe_cpu.py` dispatch 경로의 공식 적용 - PE_CPU init 의 `pe_cost_model:` override read (cycle-domain knob + `max_composite_logical_bytes`) -- **composite 크기 cap (D7)** — TLContext-side segmentation 로직, - `MAX_COMPOSITE_LOGICAL_BYTES` default 1024; 기존 bench 에 필요 없음 - (현재 최대 composite 가 200 bytes 훨씬 아래), 하지만 ADR-0065 의 10-op - decode-opt2 composite (~310 bytes) + 더 큰 recipe 가 나타날 때를 위해 - 메커니즘 land. +- **composite 크기 cap (D7)** — TLContext-side 하드 cap, + `max_composite_logical_bytes` default 1024: cap 초과 composite 는 emit + 시점에 `ValueError` (자동 segmentation 없음). 기존 bench 는 트리거 + 안 됨 (현재 최대 composite 가 ~322 bytes), 하지만 recipe 가 커질 때 + descriptor-capacity 제한이 강제되도록 체크 land. - 일회성 골든 재생성 land 후 ADR-0065 가 위에 쌓임 — 기존 bench 추가 골든 churn 없음 diff --git a/docs/adr/ADR-0064-perf-cpu-issue-cost-model.md b/docs/adr/ADR-0064-perf-cpu-issue-cost-model.md index 47ed095..069284a 100644 --- a/docs/adr/ADR-0064-perf-cpu-issue-cost-model.md +++ b/docs/adr/ADR-0064-perf-cpu-issue-cost-model.md @@ -2,7 +2,7 @@ ## Status -Proposed (Revision 2) +Accepted (Revision 2) > Supporting ADR for **ADR-0060** (AHBM GQA Fused Attention) and **ADR-0065** > (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 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 -**`MAX_COMPOSITE_LOGICAL_BYTES`** (default **1024 bytes**, overridable -per D4). Oversized commands are deterministically **segmented** by the -*emitter* (host-side TLContext, ADR-0065 D5) into N consecutive -`CompositeCmd`s, each ≤ cap. - -- Each segment carries its own `completion: CompletionHandle`. -- 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. +**`max_composite_logical_bytes`** (default **1024 bytes**, overridable +per D4). A composite whose `logical_bytes` exceeds the cap is **rejected +at emit time** by the host-side TLContext with a `ValueError` — there is +**no automatic segmentation**. The kernel author must restructure the +recipe (e.g., split it into multiple smaller composites explicitly) so +each `CompositeCmd` fits within the descriptor capacity. **Why a cap is needed.** Real hardware imposes hard limits — descriptor 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 cap *exist as a discipline*, not to fit a specific HW. -**Ordering of segments — driven by strict FIFO, not `rw_handles`.** -Segments are emitted into the PE_CPU → PE_SCHEDULER queue in their -emit order. Strict-FIFO dispatch (ADR-0065 D6.3) is the *ordering -source*: segments execute in the order they enter the queue. The -`rw_handles` block on each segment is **dependency metadata** for the -cross-composite hazard tracker — it does **not** by itself guarantee -inter-segment ordering. If a future scheduler relaxed FIFO (e.g., to -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. +**Rationale for a hard error over auto-segmentation.** Auto-splitting an +oversized composite (the original Revision 2 proposal) added emitter +complexity — inter-segment ordering, shared `rw_handles`, completion +chaining — to paper over what is, in practice, a kernel that asked for a +descriptor larger than the hardware allows. Surfacing it as an explicit +error keeps the emitter simple and makes the HW constraint visible to the +kernel author, who is best placed to decide how to split the work. 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 @@ -379,12 +368,9 @@ OpSpec/CompositeCmd fields are added. (SPEC §0.1). 7. **Path parity.** Greenlet and replay paths produce identical dispatch-cycle accounting for the same kernel. -8. **Composite size cap (D7).** A recipe that would emit `logical_bytes - > MAX_COMPOSITE_LOGICAL_BYTES` is segmented into N consecutive - `CompositeCmd`s; `sum(segment.logical_bytes) == original - logical_bytes`; total dispatch = sum of segment dispatches (FIXED - paid per segment); inter-segment ordering preserved by strict FIFO - (not by `rw_handles` alone). +8. **Composite size cap (D7).** A composite that would emit `logical_bytes + > max_composite_logical_bytes` raises a `ValueError` at emit time (no + segmentation). A composite within the cap emits normally. 9. **Sensitivity (qualitative).** At `R ∈ {0.25, 0.0625, 0.03125}` cycles/byte, `opt3 > opt2` at *all* three points. Direction (ratio 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 - `pe_cost_model:` override read at PE_CPU init (cycle-domain knobs + `max_composite_logical_bytes`) -- **composite size cap (D7)** — TLContext-side segmentation logic with - `MAX_COMPOSITE_LOGICAL_BYTES` default 1024; not needed for any - existing bench (largest current composite is well under 200 bytes), - but the mechanism lands so it is in place when ADR-0065's - 10-op decode-opt2 composite (~310 bytes) and larger recipes appear. +- **composite size cap (D7)** — TLContext-side hard cap with + `max_composite_logical_bytes` default 1024: a composite over the cap + raises a `ValueError` at emit time (no auto-segmentation). Not + triggered by any existing bench (largest current composite is ~322 + bytes), but the check lands so the descriptor-capacity limit is + enforced as recipes grow. - one-time goldens regeneration After this lands, ADR-0065 builds on top with no further goldens churn diff --git a/src/kernbench/triton_emu/tl_context.py b/src/kernbench/triton_emu/tl_context.py index 2bd7c4b..ca6d36b 100644 --- a/src/kernbench/triton_emu/tl_context.py +++ b/src/kernbench/triton_emu/tl_context.py @@ -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: diff --git a/src/kernbench/triton_emu/tl_recipes.py b/src/kernbench/triton_emu/tl_recipes.py new file mode 100644 index 0000000..21265ac --- /dev/null +++ b/src/kernbench/triton_emu/tl_recipes.py @@ -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"), + ), + ), +} diff --git a/tests/test_tl_composite_validation.py b/tests/test_tl_composite_validation.py new file mode 100644 index 0000000..0b25c74 --- /dev/null +++ b/tests/test_tl_composite_validation.py @@ -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 diff --git a/tests/test_tl_recipe_softmax_merge.py b/tests/test_tl_recipe_softmax_merge.py new file mode 100644 index 0000000..d05ccad --- /dev/null +++ b/tests/test_tl_recipe_softmax_merge.py @@ -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