ywkang
24c705419e
Complete the composite-command study across both attention regimes and write up the result, resolving when the composite command helps latency. Decode (memory-bound, T_q=1, M=8): the command form is latency-neutral — the kernel is bound by KV streaming and the MAC array is near-idle, so the composite's only benefit here is host-issue offload (O(n_tiles) -> O(1) PE_CPU commands). Figure: gqa_decode_long_ctx_composite (latency flat, command count saturates). Prefill (compute-bound, large M=G*T_q tile-filling): add three command-form variants of a single-rank FlashAttention prefill kernel (_gqa_prefill_compute_bound: primitive / composite / composite_extended). Here the composite's scheduler-internal per-tile DMA<->compute pipeline keeps the MAC array fed while the primitive's blocking tl.dot starves it, so composite wins on MAC utilization (68% flat -> 83%) and wall-clock (19% faster at ctx=1024), with the margin growing with context (deeper P.V reduction = more tiles to pipeline) — the compute-bound mirror of the GEMM result in section 3. Figure: gqa_prefill_compute_bound (latency + MAC util). Sweep wired as GQA_1H_SWEEPS=prefill_cb; tests cover structure, e2e, and composite<primitive in the compute-bound regime. The synthesis: composite has two benefits set by roofline position — host-issue offload (always) and MAC-array feeding (compute-bound only). Decode exercises the first, prefill both. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
62 lines
2.5 KiB
Python
62 lines
2.5 KiB
Python
"""Tests for the compute-bound prefill composite-command study.
|
|
|
|
Unlike memory-bound decode (where command form is latency-neutral), in
|
|
compute-bound prefill the composite command keeps the MAC array fed by
|
|
pipelining DMA↔compute per HW tile, so it wins on wall-clock latency. The
|
|
rigorous claim here is therefore the *opposite* of the decode study:
|
|
``composite_latency < primitive_latency`` at a compute-bound context.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from kernbench.common.pe_commands import CompositeCmd
|
|
from kernbench.triton_emu.tl_context import TLContext, run_kernel
|
|
|
|
from kernbench.benches.gqa_helpers.long_ctx._gqa_prefill_compute_bound import (
|
|
gqa_prefill_composite_ext_kernel,
|
|
gqa_prefill_composite_kernel,
|
|
gqa_prefill_primitive_kernel,
|
|
)
|
|
|
|
_TOPOLOGY = Path(__file__).resolve().parents[2] / "topology.yaml"
|
|
_H_Q, _H_KV, _D_HEAD = 8, 1, 128
|
|
|
|
|
|
def _n_composites(kernel, T_q: int, S_kv: int) -> int:
|
|
tl = TLContext(pe_id=0, num_programs=1, scratch_base=1 << 61,
|
|
scratch_size=1 << 20)
|
|
run_kernel(kernel, tl, 0x1000, 0x2000, 0x3000, 0x4000,
|
|
T_q, S_kv, _H_Q, _H_KV, _D_HEAD, 1, 1)
|
|
return sum(1 for c in tl.commands if isinstance(c, CompositeCmd))
|
|
|
|
|
|
def test_command_form_composite_counts():
|
|
"""One Q-block (T_q=16 → M=128): the primitive issues no composites;
|
|
each composite form issues exactly two (Q·Kᵀ and P·V / recipe)."""
|
|
assert _n_composites(gqa_prefill_primitive_kernel, 16, 256) == 0
|
|
assert _n_composites(gqa_prefill_composite_kernel, 16, 256) == 2
|
|
assert _n_composites(gqa_prefill_composite_ext_kernel, 16, 256) == 2
|
|
|
|
|
|
def _latency_ns(variant: str, ctx_len: int) -> float:
|
|
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_compute_bound import (
|
|
_run_panel,
|
|
)
|
|
return _run_panel(variant, ctx_len, str(_TOPOLOGY))["latency_ns"]
|
|
|
|
|
|
def test_three_variants_complete_in_data_mode():
|
|
"""All three prefill forms run end-to-end through the engine."""
|
|
for variant in ("primitive", "composite", "composite_extended"):
|
|
assert _latency_ns(variant, 256) > 0
|
|
|
|
|
|
def test_composite_faster_in_compute_bound_prefill():
|
|
"""At a compute-bound context (512), the composite command's DMA↔compute
|
|
pipelining beats the primitive's serial load→dot — the opposite of the
|
|
memory-bound decode study, where the two tie."""
|
|
prim = _latency_ns("primitive", 512)
|
|
comp = _latency_ns("composite", 512)
|
|
assert comp < prim, f"composite {comp} not < primitive {prim}"
|