Commit Graph

3 Commits

Author SHA1 Message Date
ywkang 24c705419e paper(gqa): two-regime "use of composite commands" — decode + compute-bound prefill
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>
2026-06-23 17:04:26 -07:00
ywkang cd6f0ed91d 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>
2026-06-23 16:07:41 -07:00
mukesh 77eece81c4 gqa(milestone-1h-gqa): add umbrella bench + first sweep outputs
Single @bench entry that drives the prefill + decode long-context
4-cases sweeps in one invocation, mirroring milestone-1h-gemm. The
individual gqa_helpers/long_ctx/gqa_{prefill,decode}_long_ctx_4cases
modules are now pure helpers (no @bench), reached only via this
umbrella.

Env var contract:
  GQA_1H_RUN=1                    (required gate)
  GQA_1H_SWEEPS=prefill,decode    (default: both; pick subset to skip)
  GQA_1H_TOPOLOGY=topology.yaml   (override)

Output layout:
  benches/1H_milestone_output/gqa/long_ctx/
    sweep_prefill.json
    sweep_decode.json
    gqa_prefill_long_ctx_4cases_{latency,traffic,memory,parallelism}.png
    gqa_decode_long_ctx_4cases_{latency,traffic,memory,parallelism}.png

The decode bench config drops S_kv from 131_072 to 8_192 so the umbrella
finishes in minutes — the comparative-story bars are the same shape at
8K. The 131K production headline number is recoverable by overriding
the dispatch; the helper docstring notes this.

Outputs (2 sweep JSONs + 8 PNGs) ship in this commit alongside the
umbrella that produced them, following the milestone-1h-gemm pattern
where derived artifacts live with the bench that emits them.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-06-16 13:06:33 -07:00