feat(gqa): Case-6 composite-command decode variants + on-chip-operand DMA fix
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>
This commit is contained in:
@@ -8,12 +8,13 @@ Currently exercises (long-context only — short-context panels are
|
||||
future work):
|
||||
- 4-cases prefill comparative study (gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases)
|
||||
- 4-cases decode comparative study (gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases)
|
||||
- Case-6 composite-command study (gqa_helpers.long_ctx.gqa_decode_long_ctx_composite)
|
||||
|
||||
Each sub-sweep writes its own ``sweep_{prefill,decode}.json`` into the
|
||||
shared output dir ``benches/1H_milestone_output/gqa/gqa_long_ctx/``.
|
||||
Each sub-sweep writes its own ``sweep_*.json`` into the shared output
|
||||
dir ``benches/1H_milestone_output/gqa/gqa_long_ctx/``.
|
||||
Selection via the env var ``GQA_1H_SWEEPS=prefill,decode`` (default
|
||||
runs both). Toggle individual sweeps with ``GQA_1H_SWEEPS=prefill``
|
||||
or ``GQA_1H_SWEEPS=decode``.
|
||||
runs prefill+decode). The composite study is opt-in (it sweeps the
|
||||
data-mode engine over several S_kv points): ``GQA_1H_SWEEPS=composite``.
|
||||
|
||||
Gated by ``GQA_1H_RUN=1`` to keep CI fast.
|
||||
"""
|
||||
@@ -24,6 +25,9 @@ import os
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_4cases import (
|
||||
run_sweep as _run_decode_sweep,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_decode_long_ctx_composite import (
|
||||
run_sweep as _run_composite_sweep,
|
||||
)
|
||||
from kernbench.benches.gqa_helpers.long_ctx.gqa_prefill_long_ctx_4cases import (
|
||||
run_sweep as _run_prefill_sweep,
|
||||
)
|
||||
@@ -57,6 +61,7 @@ def run(torch) -> None:
|
||||
runners = {
|
||||
"prefill": _run_prefill_sweep,
|
||||
"decode": _run_decode_sweep,
|
||||
"composite": _run_composite_sweep,
|
||||
}
|
||||
unknown = [s for s in sweeps if s not in runners]
|
||||
if unknown:
|
||||
|
||||
Reference in New Issue
Block a user