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:
@@ -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 없음
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"),
|
||||
),
|
||||
),
|
||||
}
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user