ywkang
cd6f0ed91d
Add two command-form variants of the Case-6 (Cube-SP x PE-SP) long-context
decode kernel alongside the primitive baseline:
- composite: per-tile GEMMs issued as one coarse tl.composite over the
full S_local (PE_SCHEDULER tiles/streams K,V) -- O(1) PE_CPU commands
vs the primitive kernel's O(n_tiles).
- composite_extended: Q.K^T composite + softmax_merge recipe composite
(ADR-0065), folding the per-tile online merge + P.V.
Extract the shared two-level (m,l,O) reduce into _gqa_mlo_reduce and
refactor the baseline to use it (byte-equal, guarded by a 64-rank
command-stream digest test).
Engine fix: _make_compute_out omitted pinned=True, so an on-chip TCM
compute result consumed as a composite GEMM operand was scheduled as an
HBM DMA_READ of its bit-61 scratch address -> PhysAddrError in data mode.
Mark compute outputs pinned (resident), matching the composite
auto-output and recipe scratch handles. .pinned is read only by the
tiling DMA gate, so this is byte-equal for existing benches.
Add the composite sweep (emit-level PE_CPU dispatch to 1M + data-mode
latency to 128K), its plot, umbrella wiring (GQA_1H_SWEEPS=composite),
and tests (refactor guard, dispatch saturation, engine-fix, e2e).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
60 lines
2.5 KiB
Python
60 lines
2.5 KiB
Python
"""Engine-fix spec test — composite GEMM with an on-chip (TCM) operand.
|
|
|
|
Root cause: ``_make_compute_out`` (tl_context) builds compute-result
|
|
handles with ``space="tcm"`` but omits ``pinned=True``. The PE_SCHEDULER
|
|
tiling gate (tiling.py:306) skips the operand DMA_READ only when
|
|
``pinned`` is set, so a compute result fed as a composite head-GEMM
|
|
operand ``a`` (e.g. the attention probabilities into P·V) is wrongly
|
|
scheduled as an HBM ``DMA_READ`` of its bit-61 scratch address →
|
|
``PhysAddrError`` in data mode.
|
|
|
|
Every other TCM-resident handle in tl_context is already marked
|
|
``pinned=True`` for exactly this reason (composite auto-output, recipe
|
|
scratch / primary-out). ``_make_compute_out`` is the one site that
|
|
forgot. ``.pinned`` is read ONLY by the tiling DMA gate, so setting it on
|
|
compute outputs is side-effect-free and byte-equal for existing benches
|
|
(which never feed an unpinned compute output into a composite operand).
|
|
|
|
Phase 1: this test FAILS until the fix (a DMA_READ for operand A is
|
|
emitted for the on-chip ``a``).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from kernbench.common.pe_commands import CompositeCmd
|
|
from kernbench.components.builtin.pe_types import StageType
|
|
from kernbench.components.builtin.tiling import generate_plan_from_ops
|
|
from kernbench.triton_emu.tl_context import TLContext, run_kernel
|
|
|
|
|
|
def _pv_composite_kernel(*, tl) -> None:
|
|
"""P·V composite whose ``a`` is a real on-chip compute result."""
|
|
x = tl.load(0x1000, shape=(8, 1024), dtype="f16")
|
|
a = tl.exp(x) # compute output → TCM
|
|
V = tl.ref(0x3000, shape=(1024, 128), dtype="f16") # HBM, streamed
|
|
O = tl.zeros((8, 128), dtype="f16")
|
|
tl.composite(op="gemm", a=a, b=V, out=O)
|
|
|
|
|
|
def _pv_plan():
|
|
tl = TLContext(pe_id=0, num_programs=1,
|
|
scratch_base=0x200000, scratch_size=1 << 20)
|
|
run_kernel(_pv_composite_kernel, tl)
|
|
comp = next(c for c in tl.commands if isinstance(c, CompositeCmd))
|
|
return generate_plan_from_ops(
|
|
ops=comp.ops, tile_m=32, tile_k=64, tile_n=32,
|
|
bytes_per_element=2, pe_prefix="sip0.cube0.pe0",
|
|
)
|
|
|
|
|
|
def test_onchip_operand_a_not_dma_read():
|
|
"""The on-chip ``a`` (a compute result) must NOT be streamed via a
|
|
DMA_READ; only the HBM ref ``b`` (V) is."""
|
|
plan = _pv_plan()
|
|
read_operands = [
|
|
s.params.get("operand")
|
|
for t in plan.tiles for s in t.stages
|
|
if s.stage_type == StageType.DMA_READ
|
|
]
|
|
assert "A" not in read_operands, read_operands
|
|
assert "B" in read_operands, read_operands
|