gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model
Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.
ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
(m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.
ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
owns whole KV heads; PE-parallel heads with intra-group chain
reduce. Prefill has no Ring KV (each head fully resident).
ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).
ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.
ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
consults table with dispatch_cycles fallback (ADR-0046 §D6
back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
measurable end-to-end.
P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).
Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
_attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
(EN + KO kept in sync).
Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,52 @@
|
||||
"""Per-op-type CPU issue cost table (ADR-0064 D1).
|
||||
|
||||
Replaces the single uniform ``dispatch_cycles`` scalar with a cost table
|
||||
keyed by command kind. Charged on PE_CPU at issue time (before the command
|
||||
is dispatched to PE_SCHEDULER) so the hybrid's CPU-saturation win
|
||||
(ADR-0060 §1) becomes measurable.
|
||||
|
||||
The table is consulted by ``TLContext._emit_dispatch_overhead(kind)``;
|
||||
live PE_CPU paths (greenlet via ``kernel_runner.py``, legacy replay via
|
||||
``pe_cpu.py:_execute_legacy``) construct TLContext with
|
||||
``issue_cost_table=DEFAULT_CPU_ISSUE_COST`` so all benches see the cost.
|
||||
|
||||
Absolute ns values are provisional (ADR-0064 review item #1). The
|
||||
defensible claim is the **ratio** — composite ≫ primitive.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal
|
||||
|
||||
OpKind = Literal[
|
||||
"composite",
|
||||
"load",
|
||||
"store",
|
||||
"dot",
|
||||
"math",
|
||||
"ipcq_send",
|
||||
"ipcq_recv",
|
||||
"copy_to",
|
||||
]
|
||||
|
||||
|
||||
DEFAULT_CPU_ISSUE_COST: dict[str, int] = {
|
||||
"composite": 40,
|
||||
"load": 5,
|
||||
"store": 5,
|
||||
"dot": 5,
|
||||
"math": 5,
|
||||
"ipcq_send": 5,
|
||||
"ipcq_recv": 5,
|
||||
"copy_to": 5,
|
||||
}
|
||||
|
||||
|
||||
def get_issue_cost(kind: str, table: dict[str, int] | None = None) -> int:
|
||||
"""Return per-op-type CPU issue cost in ns.
|
||||
|
||||
Unknown kinds return 0 (no charge) so adding a new ``tl.*`` op kind
|
||||
doesn't accidentally over-charge before the table is updated.
|
||||
"""
|
||||
if table is None:
|
||||
table = DEFAULT_CPU_ISSUE_COST
|
||||
return table.get(kind, 0)
|
||||
@@ -73,6 +73,12 @@ class TensorHandle:
|
||||
data: object = None # reserved for validate mode
|
||||
space: str = "tcm" # MemoryStore space ("tcm" | "hbm" | "sram")
|
||||
pinned: bool = False # operand already DMA-staged in TCM (via tl.load)
|
||||
# ADR-0062 §D2: lazy tl.load attaches a LoadFuture here. None for
|
||||
# handles that have no in-flight DMA (constants, math outputs, etc.).
|
||||
# Consumer ops call _await_pending() to yield on the future before
|
||||
# emitting their own command. Excluded from eq/hash/repr so handle
|
||||
# identity is unaffected by pending state.
|
||||
pending: object = field(default=None, compare=False, hash=False, repr=False)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -140,6 +146,22 @@ class MathCmd:
|
||||
data_op: bool = True
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyCmd:
|
||||
"""TCM-to-TCM byte copy (ADR-0063 §D3.1).
|
||||
|
||||
Emitted by ``tl.copy_to`` to persist a scoped result's bytes to an
|
||||
outside-``scratch_scope`` (persistent) address — the two-arena
|
||||
pattern for tiled flash attention. Runs on the vector engine;
|
||||
op_log classifies as ``op_kind="math"``, ``op_name="copy"``.
|
||||
"""
|
||||
|
||||
src: TensorHandle
|
||||
dst: TensorHandle
|
||||
nbytes: int
|
||||
data_op: bool = True
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CompositeCmd:
|
||||
"""Composite command: tiled pipeline of DMA_READ + COMPUTE + DMA_WRITE.
|
||||
@@ -178,7 +200,7 @@ class PeCpuOverheadCmd:
|
||||
|
||||
# Union type for all PE commands
|
||||
PeCommand = (
|
||||
DmaReadCmd | DmaWriteCmd | GemmCmd | MathCmd
|
||||
DmaReadCmd | DmaWriteCmd | GemmCmd | MathCmd | CopyCmd
|
||||
| CompositeCmd | WaitCmd | PeCpuOverheadCmd
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user