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:
@@ -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).
|
||||
|
||||
Reference in New Issue
Block a user