gqa(adr-0065): P5(B) — decode opt2 two-composite kernel + dispatch measurement
New _gqa_attention_decode_opt2.py: per-tile attention as #1 Q.Kt GEMM composite + #2 softmax_merge recipe composite (online merge + P.V + add). Runnable in op_log mode; full data-mode numeric parity is a deferred follow-up (P5-numerics). Measured per-tile PE_CPU dispatch (ADR-0064 Rev2 default): opt3=960ns vs opt2=184ns = 5.22x (gate >2x), FIXED-dominated (command-count reduction) — the ADR-0060/0064 CPU-offload win. opt2<opt3 across R in {0.25,0.0625,0.03125}. K-before-V: softmax_merge prologue carries no DMA; V (ref) streams only in the GEMM tile loop. Fix latent P3 bug surfaced by the first e2e recipe run: prologue MATH stages were FOLDED into the first GEMM tile, but a folded MATH->DMA_READ boundary needs a PE_MATH->PE_DMA token route the pipeline never wires (KeyError pe_dma). Now prologue/post-loop stages are fed as standalone 1-stage tiles (each completes on its component); completion count includes them. Existing benches have no prologue -> tiles + feed order unchanged (byte-equal); full suite 806 pass / 3 pre-existing fail. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,84 @@
|
||||
"""GQA decode opt2 variant — two-composite inner attention (ADR-0065).
|
||||
|
||||
opt2 of decode (ADR-0060 §8 item 4 / ADR-0065): the per-tile attention is
|
||||
expressed as **two composites** instead of opt3's long primitive sequence —
|
||||
|
||||
#1 Q·Kᵀ → a GEMM composite writing the score tile ``scores``.
|
||||
#2 softmax → a ``softmax_merge`` recipe composite (online-softmax merge
|
||||
+ P·V of ``(m, l, O)``) whose head GEMM is P·V, with an ``add``
|
||||
epilogue folding the P·V contribution into ``O``.
|
||||
|
||||
Fewer PE_CPU-issued commands ⇒ lower dispatch cost under the ADR-0064 Rev2
|
||||
structural model (the CPU-offload win). Tile 0 establishes the running
|
||||
state with primitives (kernbench has no scratch-backed ``-inf`` initializer,
|
||||
same limitation as the opt3 kernel); the next tile exercises the recipe.
|
||||
|
||||
**Scope (P5 B):** runnable in op_log mode (latency only). Full data-mode
|
||||
numeric parity (computing the recipe's 8 MATH ops in the DataExecutor) is a
|
||||
separate follow-up — see DDD-0065 / the P5-numerics note.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def gqa_attention_decode_opt2_kernel(
|
||||
q_ptr: int,
|
||||
k_ptr: int,
|
||||
v_ptr: int,
|
||||
o_ptr: int,
|
||||
T_q: int,
|
||||
S_kv: int,
|
||||
h_q: int,
|
||||
h_kv: int,
|
||||
d_head: int,
|
||||
C: int,
|
||||
P: int,
|
||||
*,
|
||||
tl,
|
||||
) -> None:
|
||||
"""Single-rank (C=P=1) GQA decode using the opt2 two-composite form.
|
||||
|
||||
Layout mirrors ``gqa_attention_decode_long_kernel`` (M-fold Q, K loaded
|
||||
as ``(d_head, S_local)``). The KV slice is split into two sub-tiles so
|
||||
the second one drives the ``softmax_merge`` recipe composite.
|
||||
"""
|
||||
G = h_q // h_kv
|
||||
S_local = S_kv // (C * P)
|
||||
KV_ROW_BYTES = d_head * 2 # f16
|
||||
|
||||
Q = tl.load(q_ptr, shape=(G * T_q, d_head), dtype="f16")
|
||||
|
||||
# Split the local KV slice in two: tile 0 establishes the running state,
|
||||
# tile 1 merges via the opt2 two-composite path.
|
||||
half = max(1, S_local // 2)
|
||||
rest = S_local - half
|
||||
|
||||
# ── Tile 0: establish running (m_local, l_local, O_local) ──
|
||||
K_T0 = tl.load(k_ptr, shape=(d_head, half), dtype="f16")
|
||||
V0 = tl.load(v_ptr, shape=(half, d_head), dtype="f16")
|
||||
scores0 = tl.dot(Q, K_T0)
|
||||
m_local = tl.max(scores0, axis=-1)
|
||||
exp0 = tl.exp(scores0 - m_local)
|
||||
l_local = tl.sum(exp0, axis=-1)
|
||||
O_local = tl.dot(exp0, V0)
|
||||
|
||||
# ── Tile 1: opt2 two-composite merge (skipped when S_local == 1) ──
|
||||
if rest > 0:
|
||||
K_T1 = tl.load(k_ptr + half * KV_ROW_BYTES,
|
||||
shape=(d_head, rest), dtype="f16")
|
||||
V1 = tl.ref(v_ptr + half * KV_ROW_BYTES, shape=(rest, d_head),
|
||||
dtype="f16")
|
||||
# #1: Q·Kᵀ composite → score tile (TCM-resident, consumed by #2).
|
||||
scores1 = tl.zeros((G * T_q, rest), dtype="f16")
|
||||
tl.composite(op="gemm", a=Q, b=K_T1, out=scores1)
|
||||
# #2: softmax_merge recipe (online merge of (m,l,O)) + P·V + add.
|
||||
tl.composite(
|
||||
prologue=[{"op": "softmax_merge", "s": scores1,
|
||||
"m": m_local, "l": l_local, "O": O_local}],
|
||||
op="gemm", b=V1, out=O_local,
|
||||
epilogue=[{"op": "add", "other": O_local}],
|
||||
)
|
||||
|
||||
# ── Final normalise + store (root only) ──
|
||||
if tl.program_id(axis=0) == 0 and tl.program_id(axis=1) == 0:
|
||||
O_final = O_local / l_local
|
||||
tl.store(o_ptr, O_final)
|
||||
@@ -170,9 +170,13 @@ class PeSchedulerComponent(ComponentBase):
|
||||
plan = self._generate_plan(cmd)
|
||||
|
||||
self._pipeline_counter += 1
|
||||
# Prologue / post-loop single-shot stages each count as a tile for
|
||||
# completion (ADR-0065 D3). Empty for legacy composites → unchanged.
|
||||
n_pre = len(getattr(plan, "prologue_stages", ()))
|
||||
n_post = len(getattr(plan, "epilogue_stages", ()))
|
||||
ctx = PipelineContext(
|
||||
id=f"p{self._pipeline_counter}",
|
||||
total_tiles=len(plan.tiles),
|
||||
total_tiles=len(plan.tiles) + n_pre + n_post,
|
||||
done_event=pe_txn.done,
|
||||
)
|
||||
|
||||
@@ -198,6 +202,13 @@ class PeSchedulerComponent(ComponentBase):
|
||||
assert self._pending_feeds is not None
|
||||
while True:
|
||||
plan, ctx = yield self._pending_feeds.get()
|
||||
tid = 0
|
||||
# Prologue single-shot stages first (ADR-0065 D3) — each a
|
||||
# standalone 1-stage tile that runs on its component and
|
||||
# completes (no inter-component routing).
|
||||
for st in getattr(plan, "prologue_stages", ()):
|
||||
yield from self._feed_single_stage(st, ctx, tid)
|
||||
tid += 1
|
||||
for tile in plan.tiles:
|
||||
first_stage = tile.stages[0]
|
||||
token = TileToken(
|
||||
@@ -208,6 +219,20 @@ class PeSchedulerComponent(ComponentBase):
|
||||
params=first_stage.params,
|
||||
)
|
||||
yield self.out_ports[first_stage.component].put(token)
|
||||
for st in getattr(plan, "epilogue_stages", ()):
|
||||
yield from self._feed_single_stage(st, ctx, tid)
|
||||
tid += 1
|
||||
|
||||
def _feed_single_stage(self, stage: Any, ctx: Any, tile_id: int) -> Generator:
|
||||
"""Feed one standalone stage as a 1-stage tile (prologue/post-loop)."""
|
||||
from kernbench.components.builtin.pe_types import TilePlan, TileToken
|
||||
|
||||
token = TileToken(
|
||||
tile_id=tile_id, pipeline_ctx=ctx,
|
||||
plan=TilePlan(tile_id=tile_id, stages=(stage,)),
|
||||
stage_idx=0, params=stage.params,
|
||||
)
|
||||
yield self.out_ports[stage.component].put(token)
|
||||
|
||||
def _generate_plan(self, cmd: Any) -> Any:
|
||||
"""Generate a PipelinePlan from a flat-ops CompositeCmd (ADR-0065 D3).
|
||||
|
||||
@@ -283,22 +283,13 @@ def generate_plan_from_ops(
|
||||
epilogue_specs=tuple(post_ops),
|
||||
)
|
||||
|
||||
pre_stages = tuple(_math_stage(o, pe_prefix) for o in pre_ops)
|
||||
post_stages = tuple(_math_stage(o, pe_prefix) for o in post_ops
|
||||
if o.scope == _Scope.KERNEL)
|
||||
plan.prologue_stages = pre_stages
|
||||
plan.epilogue_stages = post_stages
|
||||
|
||||
# Fold prologue/post-loop stages into the first/last tile so the feeder
|
||||
# and completion counting are untouched (existing benches have neither,
|
||||
# so their tiles are unchanged → byte-equal op_log).
|
||||
if pre_stages and plan.tiles:
|
||||
t0 = plan.tiles[0]
|
||||
plan.tiles[0] = TilePlan(tile_id=t0.tile_id,
|
||||
stages=(*pre_stages, *t0.stages))
|
||||
if post_stages and plan.tiles:
|
||||
tl = plan.tiles[-1]
|
||||
plan.tiles[-1] = TilePlan(tile_id=tl.tile_id,
|
||||
stages=(*tl.stages, *post_stages))
|
||||
|
||||
# Prologue (pre-GEMM KERNEL) + post-loop KERNEL ops become single-shot
|
||||
# MATH stages. They are fed as standalone 1-stage tiles by the scheduler
|
||||
# (PE_SCHEDULER._feed_loop) — NOT folded into the GEMM tiles, because a
|
||||
# folded MATH→DMA_READ boundary would require a PE_MATH→PE_DMA token
|
||||
# route that the pipeline does not wire. Existing benches have neither →
|
||||
# their tiles + feed order are unchanged (byte-equal op_log).
|
||||
plan.prologue_stages = tuple(_math_stage(o, pe_prefix) for o in pre_ops)
|
||||
plan.epilogue_stages = tuple(_math_stage(o, pe_prefix) for o in post_ops
|
||||
if o.scope == _Scope.KERNEL)
|
||||
return plan
|
||||
|
||||
Reference in New Issue
Block a user