c164645aee1e70b2882e45a93ca6fd01aebbb093
6 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 mukesh
|
c164645aee |
gqa(decode-4cases): Case 4 anchor in dedicated bench (5C.D)
First milestone of the decode 4-cases comparative study per GQA_full_deck.pptx slides 11-17. Case 4 (Cube-SP × PE-SP, the optimal case per slide 11) is structurally the existing _gqa_attention_decode_long at sub_w=4 (Increment 2's lrab-adapted center-root reduce). This commit wires it into a dedicated bench so the remaining cases land alongside. Changes: - New bench: src/kernbench/benches/milestone_gqa_decode_4cases.py Houses the 4 case panels under one milestone-gqa-decode-4cases entry (gated by GQA_DECODE_4CASES_RUN=1; output to 1H_milestone_output/gqa_decode_4cases/sweep.json). Cases 1-3 are TBD in subsequent sub-increments (5C.A/B/C). - New panel: single_kv_group_decode_gqa_cube_sp_pe_sp C=8, P=8, sub_w=4, T_q=1, S_kv=131_072, d_head=128, h_q=8, h_kv=1. - src/kernbench/benches/milestone_gqa_headline.py: _run_decode_panel extended with keyword-only sub_w/T_q/d_head/h_q/h_kv overrides (defaults preserve existing-panel behaviour). - tests/attention/test_milestone_gqa_decode_4cases.py: 4 new tests asserting registration, smoke completion, reduce-to-root at the lrab center cube (cube 6), and the predicted 189-ipcq Case-4 traffic pattern (168 intra-CUBE + 21 inter-CUBE lrab Phase 1+2). - tests/attention/test_milestone_gqa_headline.py: rename test_sweep_json_has_four_panels -> test_sweep_json_has_expected_panels and switch hardcoded 4 to len(PANELS) (the panel set grew to 5 with Increment 5's single_kv_group_prefill_gqa_c8_p8). Deviation noted: slide 13 prescribes AllReduce on (m,ℓ,O); our kernel does reduce-to-root (only the lrab center cube has the answer) per ADR-0060 §4. Treated as the kernbench Case-4 baseline. Verification: all 4 new tests pass; 90 regression tests pass; the previously-failing test_sweep_json_has_four_panels now passes under its renamed form. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> |
||
|
mukesh
|
5b4d9cb597 |
bench(milestone-gqa-headline): scale single_kv_group prefill panel T_q=S_kv=1K (scratch budget)
The headline T_q=S_kv=32K target overflows the 1 MB per-PE scratch pool: at T_q_local=4K and tile_s=1024 the scores matrix alone is 8 MB. The prefill kernel's bootstrap section also leaves K_t/V_t/scores/ exp_scores persistent (outside tl.scratch_scope), inflating baseline. Scale-down to T_q=S_kv=1K (T_q_local=128, fits comfortably) preserves the C=8 + P=8 architecture demonstration; the true LLaMA 32K headline awaits a future increment to add Q-axis tiling and tighten bootstrap scratch discipline. Verified end-to-end: kernbench run --bench milestone-gqa-headline now produces sweep.json with all 5 panels. The new panel shows ipcq_copy=896 (matches (C-1)·n_tiles·2·C·P = 7·1·2·8·8) and dma_write=64 (one per PE, head-parallel + intra-CUBE PE-SP). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> |
||
|
mukesh
|
9e1242039b |
gqa: single-KV-group LLaMA-3.1-70B prefill milestone (Increments 1-5)
End-to-end wires C=8 P=8 d_head=128 prefill with snake-ring inter-CUBE SFR, intra-CUBE PE-SP (all 64 ranks active), and the milestone bench panel. Decode kernel gains lrab-adapted center-root reduce for the 2×4 sub-mesh per ADR-0060 §4.2. Increment 1 — SFR multi-row snake src/kernbench/ccl/sfr_config.py: configure_sfr_intercube_ring gains submesh_shape / submesh_origin kwargs; installs a Hamiltonian snake ring through a rectangular sub-mesh (every hop is 1-hop physical neighbour). Backward-compat: 1D-row behaviour preserved when submesh_shape is None. tests/test_intercube_snake_ring.py (12 tests) Increment 2 — Decode lrab-adapted center-root reduce src/kernbench/benches/_gqa_attention_decode_long.py: new sub_w param (default 0 = existing 1D-chain). sub_w >= 2 selects the ADR-0060 §4.2 prescribed lrab-adapted Phase 1+2 reduce (bidirectional row + bidirectional col converge to the center cube), with log-sum-exp _merge_running replacing the plain + of lrab. tests/attention/test_gqa_decode_long_2d_reduce.py (4 tests) Increment 3 — Prefill kernel at C=8 (no production change) Verified by inspection that the existing prefill_long kernel + Increment 1's snake SFR already work at C=8 without any kernel edit. The kernel speaks logical W/E; the snake routes it. tests/attention/test_gqa_prefill_long_c8_snake.py (3 tests) Increment 4 — Intra-CUBE PE-SP in prefill (all 64 ranks) src/kernbench/benches/_gqa_attention_prefill_long.py: new P param (default 1 = existing PE-0-only). P > 1 splits T_q query-axis-wise across the P PEs of each CUBE; output rows are disjoint per PE so no intra-CUBE reduce is needed; each PE drives its own same-lane ring (P parallel rings). tests/attention/test_gqa_prefill_long_pe_sp.py (5 tests) Increment 5 — LLaMA-scale milestone bench panel src/kernbench/benches/milestone_gqa_headline.py: new panel single_kv_group_prefill_gqa_c8_p8 (C=8, P=8, T_q=S_kv=32K, d_head=128). _run_prefill_panel extended with P/T_q/d_head defaults; routes snake SFR when C > mesh_w. tests/attention/test_milestone_gqa_single_kv_group_prefill_panel.py (3 tests) Total: 4 production files modified, 5 new test files, 27 new tests. Followed the Phase 1/2 protocol per CLAUDE.md throughout. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> |
||
|
mukesh
|
7fad0371c5 |
gqa: tile-granular Ring KV (P3c) + rename to gqa_attention_* + ADR-0060/62/63/64 → Accepted
Three logically distinct changes, bundled for atomic test green:
1. **P3c — prefill_long tile-granular Ring KV** (ADR-0060 §5.5.1 amendment).
Convert the ring from slice-granular (one full ``(d_head, S_local)``
KV slice per step) to tile-granular (``n_tiles`` tiles of
``TILE_S_KV`` per step). Nested loop with outer tile, inner ring step:
each tile propagates through all C ring positions before the next
tile starts, so IPCQ in-flight depth stays at 1 per direction.
Bootstrap at ``(t=0, k=0)`` outside the scratch_scope establishes the
persistent ``(m, ℓ, O)``; every other iteration scope-wraps + persists
via ``copy_to``. Per-rank persistent scratch shrinks to ~1 KB; per-tile
scope bounded by TILE_S_KV regardless of S_local. Headline:
prefill_long now completes at S_kv=128K (previously overflowed).
New: ``tests/attention/test_gqa_prefill_long_tile_ring.py``
(3 tests — ceiling-lift + tile-granular ipcq_copy count +
per-CUBE distributed output regression guard).
2. **Rename ``gqa_*`` → ``gqa_attention_*``** across kernel files,
function names, and importers. The "attention" name makes the role
explicit (GQA is grouped-query attention) and matches upstream Triton
FlashAttention naming conventions. Renames:
_gqa_decode_long.py -> _gqa_attention_decode_long.py
_gqa_decode_short.py -> _gqa_attention_decode_short.py
_gqa_prefill_long.py -> _gqa_attention_prefill_long.py
_gqa_prefill_short.py -> _gqa_attention_prefill_short.py
And function names ``gqa_<phase>_<context>_kernel`` →
``gqa_attention_<phase>_<context>_kernel``. Updated 1 bench file
(milestone_gqa_headline.py) and 10 test files.
3. **ADR-0060 / 0062 / 0063 / 0064: Proposed → Accepted**.
All four are reflected in production code and covered by tests:
- ADR-0060 (GQA fused attention): 4 kernels deployed; §5.5.1
amendment added for the tile-granular Ring KV introduced by P3c
(EN + KO mirror).
- ADR-0062 (lazy tl.load): LoadFuture + _await_pending live in
tl_context.py.
- ADR-0063 (tl.scratch_scope + tl.copy_to): used in every chain
reduce + tile sweep + ring step. EN-only previously; KO
translation authored as part of this commit (CLAUDE.md
bidirectional rule).
- ADR-0064 (per-op-type CPU issue cost): cpu_issue_cost.py +
issue_cost_table wiring in tl_context.py (Phase E).
Files git mv'd from docs/adr-proposed/ to docs/adr/ (EN) and
docs/adr-ko/ (KO). ADR-0061 (tl.broadcast) stays Proposed — no
implementation; documented as optional convenience primitive in
the ADR itself.
Tests: 88/88 focused regression green
(tests/attention/ + Phase E + TL discipline).
ADR pair verification: ``python tools/verify_adr_lang_pairs.py`` OK.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
|
||
|
mukesh
|
5a76ed4f6a |
gqa: rename long-context kernels to *_long for symmetry with *_short
Make the file + function naming symmetric:
_gqa_decode.py -> _gqa_decode_long.py
_gqa_prefill.py -> _gqa_prefill_long.py
gqa_decode_kernel -> gqa_decode_long_kernel
gqa_prefill_kernel -> gqa_prefill_long_kernel
Mirrors the existing _gqa_{decode,prefill}_short.py naming. Updates the
two imports + two call sites in milestone_gqa_headline.py and the 9
attention tests that import the kernels.
Tests: 72/72 focused regression green (tests/attention/ + Phase E + TL
discipline). milestone-gqa-headline bench passes its 7 panel/schema
assertions.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
|
||
|
mukesh
|
d282144339 |
gqa: ADR-0060/0062/0063/0064 unified GQA kernels + CPU cost model
Land the new GQA fused-attention kernels (ADR-0060) for prefill/decode
across long and short context, the TL discipline primitives they depend
on (ADR-0062 lazy load, ADR-0063 scratch_scope + copy_to), and the
per-op-type CPU issue cost model (ADR-0064). Remove the pre-ADR-0060
mesh-attention baseline now that the unified kernels supersede it.
ADR-0060 (long context)
- _gqa_decode.py: M-fold + 2-level chain reduce-to-root (Level-2
intra-CUBE row-then-col + Level-1 inter-CUBE) — root-only output.
- _gqa_prefill.py: head-parallel + Ring KV rotation around C CUBEs,
online-softmax merge per ring step, per-CUBE distributed output.
- Each merge stage wraps in scratch_scope() and persists running
(m, l, O) via copy_to() to lift the 1 MiB scratch ceiling.
ADR-0060 §B.split.2 (short context, kv_per_cube in {1,2,4,8})
- _gqa_decode_short.py / _gqa_prefill_short.py: no cube-SP; each CUBE
owns whole KV heads; PE-parallel heads with intra-group chain
reduce. Prefill has no Ring KV (each head fully resident).
ADR-0062 (lazy tl.load): future-bearing TensorHandle, auto-wait at
first consuming op (dot/MATH/store/send/copy_to/composite).
ADR-0063 (tl.scratch_scope + tl.copy_to): scoped per-tile arena with
copy_to writeback primitive for persistent running state.
ADR-0064 (CPU issue cost model)
- common/cpu_issue_cost.py: per-op-type table (composite=40 ns,
primitives=5 ns); ratios are load-bearing per D1.
- TLContext: issue_cost_table param; _emit_dispatch_overhead(kind)
consults table with dispatch_cycles fallback (ADR-0046 §D6
back-compat).
- Live PE_CPU paths (greenlet + legacy) construct TLContext with
DEFAULT_CPU_ISSUE_COST so saturation lever (ADR-0060 §1) is
measurable end-to-end.
P7 headline bench: milestone-gqa-headline writes per-panel
op_log_summary to 1H_milestone_output/gqa_headline/sweep.json. No
figure renderers yet (deferred).
Removals (pre-ADR-0060 baseline now superseded):
- benches: _attention_mesh_kv.py, _attention_mesh_mlo.py,
_attention_mesh_mlo_2d.py, milestone_gqa_llama70b.py
- tests: test_attention_*, test_mesh_*, test_milestone_gqa_llama70b
- topology: llama70b_4sip.yaml (only consumer was the deleted diag)
- artifacts: 1H_milestone_output/gqa/ (sweep.json + 5 PNGs)
- tests/gqa/ plot helper + test (broken on Windows Tcl/Tkinter)
- ADR-0060/0061 references to deleted file paths cleaned up
(EN + KO kept in sync).
Tests: 124/124 focused regression green (attention + Phase E + TL
discipline + triton_emu + pe_components). Full regression: 764 pass,
2 pre-existing test_bench_registry failures (stale EXPECTED_NAMES
across multiple benches, not introduced here).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
|