gqa(adr-0064/0065): flat-ops CompositeCmd (P1) + structural dispatch cost (ADR-0064 Rev2); promote ADR-0064

ADR-0065 P1: CompositeCmd -> flat ordered ops list (drop legacy op/a/b/out_addr fields); OpSpec.operands dict + out handle. Meaning-preserving (op_log byte-equal); pe_scheduler + op_log read the head op.

ADR-0064 Rev2: replace Rev1 per-op cost table with structural FIXED + logical_bytes*R formula. logical_bytes on every PeCommand; new common/pe_cost_model.py; cost centralized in TLContext._emit (load/recv_async charge explicitly); pe_cpu/kernel_runner wire the per-PE model + clock. D7: cap exceeded -> ValueError (no auto-segmentation). Remove Rev1 cpu_issue_cost.py + its tests. No goldens churn.

Promote ADR-0064 Rev2 Proposed->Accepted (docs/adr/ + docs/adr-ko/); amend D7 (error not segmentation) + record P1-before-P0 ordering in ADR-0064/0065 Migration notes (EN+KO).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-10 19:18:04 -07:00
parent 79ddb12b42
commit 47e2c78c66
18 changed files with 703 additions and 550 deletions
+12 -2
View File
@@ -66,6 +66,14 @@ class PeCpuComponent(ComponentBase):
| (self._cube_idx << 32)
| (self._pe_idx << 24)
)
# ADR-0064 Rev2: structural dispatch cost model. Reads an optional
# `pe_cost_model:` block from this PE_CPU node's attrs (D4); missing
# keys fall back to defaults. Clock for cycle→ns is this node's
# `clock_freq_ghz` (D3), same attr PE_GEMM/PE_MATH use.
from kernbench.common.pe_cost_model import from_node_attrs
self._cost_model = from_node_attrs(node.attrs)
self._clock_freq_ghz = float(node.attrs.get("clock_freq_ghz", 1.0))
def _find_shard(self, shards: tuple) -> Any:
"""Find shard matching this PE's (sip, cube, pe). Fallback to positional index."""
@@ -176,6 +184,8 @@ class PeCpuComponent(ComponentBase):
store=store,
scratch_base=self._tl_scratch_base,
scratch_size=self._tl_scratch_size,
cost_model=self._cost_model,
clock_freq_ghz=self._clock_freq_ghz,
)
yield from runner.run(env, kernel_fn, kernel_args, num_programs)
return getattr(runner, "_composite_results", [])
@@ -184,7 +194,6 @@ class PeCpuComponent(ComponentBase):
self, env, kernel_fn, kernel_args, num_programs, scheduler_id,
) -> Generator:
"""Legacy Phase 0 + replay: generate command list, then dispatch."""
from kernbench.common.cpu_issue_cost import DEFAULT_CPU_ISSUE_COST
from kernbench.common.pe_commands import (
CompositeCmd, PeCpuOverheadCmd, PeInternalTxn, WaitCmd,
)
@@ -194,7 +203,8 @@ class PeCpuComponent(ComponentBase):
pe_id=self._pe_idx, num_programs=num_programs,
cube_id=self._cube_idx, num_cubes=self._num_cubes,
dispatch_cycles=0,
issue_cost_table=DEFAULT_CPU_ISSUE_COST,
cost_model=self._cost_model,
clock_freq_ghz=self._clock_freq_ghz,
)
run_kernel(kernel_fn, tl, *kernel_args)
commands = tl.commands
@@ -156,19 +156,21 @@ class PeSchedulerComponent(ComponentBase):
pp = self._pe_prefix
bpe = 2 # default bytes per element (f16)
if cmd.op == "gemm" and cmd.b is not None:
a = cmd.a
b = cmd.b
# Flat-ops (ADR-0065 D1): ops[0] is the head, ops[1:] are epilogue
# specs placed by scope. The head's kind selects the engine path.
head = cmd.ops[0]
epi_specs = tuple(cmd.ops[1:])
if head.kind == "gemm" and "b" in head.operands:
a = head.operands["a"]
b = head.operands["b"]
M, K = a.shape[-2], a.shape[-1]
N = b.shape[-1]
# When CompositeCmd.ops is populated, ops[0] is the head and
# ops[1:] is the epilogue spec list. Empty ops → legacy path.
epi_specs = tuple(cmd.ops[1:]) if cmd.ops else ()
return generate_gemm_plan(
M=M, K=K, N=N,
tile_m=self.TILE_M, tile_k=self.TILE_K, tile_n=self.TILE_N,
bytes_per_element=bpe,
A_addr=a.addr, B_addr=b.addr, C_addr=cmd.out_addr,
A_addr=a.addr, B_addr=b.addr, C_addr=head.out.addr,
pe_prefix=pp,
a_pinned=getattr(a, "pinned", False),
b_pinned=getattr(b, "pinned", False),
@@ -176,14 +178,14 @@ class PeSchedulerComponent(ComponentBase):
)
else:
# Math composite
a = cmd.a
a = head.operands["a"]
M = a.shape[-2] if len(a.shape) >= 2 else a.shape[0]
N = a.shape[-1] if len(a.shape) >= 2 else 1
return generate_math_plan(
M=M, N=N,
tile_m=self.TILE_M, tile_n=self.TILE_N,
bytes_per_element=bpe,
math_op=cmd.math_op or "identity",
src_addr=a.addr, dst_addr=cmd.out_addr,
math_op=head.extra.get("math_op") or "identity",
src_addr=a.addr, dst_addr=head.out.addr,
pe_prefix=pp,
)