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:
@@ -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