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>
Replace each kernel's local one-shot partial with a Tile-0 bootstrap + a
``scratch_scope``-wrapped merge loop. Per-rank scratch is now bounded by
``TILE_S_KV = 1024`` regardless of ``S_local``, lifting the long-context S
ceiling. Backward compatible at ``S_local <= TILE_S_KV`` (loop body is
empty; op_log structurally identical to today).
Framework: extend ``memory_store.read`` to support partial reads at
offsets within a stored region. Models real Triton ``tl.load(ptr+offset,
shape=...)`` — needed because each tile loads ``k_ptr + tile_start*row_bytes``
for its sub-slice of the per-rank KV region.
prefill_long is intentionally left untiled. Its ring loop carries
full-slice ``Kc``/``Vc`` for ``tl.send`` between CUBEs, so tile-sweeping
step 0 wouldn't shrink the kernel's actual scratch footprint. Lifting
prefill_long's ceiling requires a tile-granular ring rewrite — separate
phase. Docstring + inline comments cleaned to reflect the current shape.
Docstrings across all 4 GQA kernels: drop P1a/P2a/P2b/P6a/P6b lineage
paragraphs and historical deviation lists; describe each kernel by what
it does, not how it got here. Add explicit ``# Local attention`` /
``# Communication`` section headers.
Tests: new ``tests/attention/test_gqa_tile_sweep.py`` with 5 tests
covering single-tile op_log stability, multi-tile ``copy_to`` emission
across all 3 refactored kernels, and the headline ceiling-lift
(decode_long at S_kv=256K). Full focused regression green: 85/85 across
``tests/attention/`` + Phase E + TL discipline tests.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
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>
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>
New ``_attention_mesh_mlo_2d.py`` decomposes a
``(mesh_rows x mesh_cols)`` cube sub-mesh into two stages of
bidirectional AllReduce-mlo:
Stage 1 — row reduce (E/W edges, mesh_cols-1 steps)
Stage 2 — col reduce (N/S edges, mesh_rows-1 steps)
After both stages every cube holds the same final ``(m, l, o)`` and
writes the normalized output. The online-softmax mlo merge is
associative, so row-then-col partitioning is mathematically
equivalent to a 1D ring AllReduce-mlo over all
``mesh_rows * mesh_cols`` cubes but uses fewer hops:
- 2x4 (8 cubes / KV-group): 4 steps vs 7 (1.75x faster)
- 4x4 (16 cubes / full SIP): 6 steps vs 15 (2.5x faster)
Motivation: the original 1D ring kernel ``_attention_mesh_mlo.py``
hit ``IpcqInvalidDirection`` at cube 4 when ``n_ranks=8`` on the
4x4 cube mesh — cube 4 has no W neighbor at the row 0/1 boundary.
N/S edges are already installed by ``configure_sfr_intercube_multisip``
so the 2D kernel runs on existing wiring without SFR changes.
The kernel accepts ``cube_start: int = 0`` and subtracts it from
``program_id(axis=1)`` so the ring math uses launch-local rank. This
matters because kernbench's ``program_id(axis=1)`` returns the
physical cube id (ADR-0022), so a launch starting at cube 8 would
otherwise compute ``my_row = 8//4 = 2`` (out of sub-mesh bounds) and
deadlock. Default ``cube_start=0`` keeps the existing
multi_user_decode validation behavior bit-for-bit.
Bench dispatch: ``multi_user_decode`` in milestone-gqa-llama70b now
uses the 2D kernel via a new ``mesh_shape`` column in
``_PANEL_DISPATCH``. At validation ``N_RANKS_MULTI_USER=4``, the
shape is ``(1, 4)`` — a degenerate single-row mesh, equivalent in
step count and op_log structure to the prior 1D ring at n_ranks=4.
The other three panels keep their 1D kernels.
Tests: 4 new unit tests in ``test_mesh_mlo_2d_correctness.py`` —
1x4 (degenerate row), 2x4 (8-KV-group target), 4x4 (full SIP), and
2x4 at cube_start=8 (the second sub-mesh per SIP). Existing
milestone (12 tests) and mesh-kernels-rank-axis (7 tests) suites
stay green — no regression.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Both mesh kernels and the milestone-gqa bench described single_user vs
multi_user in ways that were architecturally misleading. Source of truth
is the GQA Llama-70B sharding study at
llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py:
Single User (B=1) — TL prefill, BL decode:
KV split @ PEs intra-cube. Each cube does its own 8-PE ring,
independent of other cubes. Headline = 64 cubes serving one user.
Multi User (B=8) — TR prefill, BR decode:
KV split @ cubes inter-cube. 8 cubes/KV-group form the ring; inside
each cube, "8 PEs each handle 1 different user → Batch on batch."
Updates the rank_axis docstring on both mesh kernels and the bench's
module docstring to make this accurate, including the v1 simplification
of the multi_user kernel (gating pe_id != 0 collapses B=8 → B=1; the
per-cube batch parallelism is deferred to sub-cycle 4c headline).
No behavior change; existing tests unchanged.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Add 5 of the 6 figure renderers ADR-0057 D3 sub-cycle 4c specifies:
- gqa_op_log_{panel}.png × 4 — per-panel bar chart of the 5 op_log
counts (gemm, ipcq_send, ipcq_recv, dma_read, dma_write).
- gqa_comparison.png — cross-panel grouped bars over the same 5 series.
Sixth figure (gqa_scaling.png) depends on sub-cycle 4b's Q/cube ∈
{1, 2, 4} sweep on multi_user_* panels and is deferred until that
data exists; emit_all_gqa_plots returns just the 5 in-scope paths.
Add MILESTONE_FAST=1 mode to run(): skip the panel sweep, reuse the
committed sweep.json, render figures only. Validation mode unchanged.
The runtime errors clearly when neither env var is set, listing the
two supported modes.
Renderers live in the bench module (the milestone-1h-gemm pattern);
tests/gqa/_gqa_plot_helpers.py re-exports them for figure tests.
Tests: tests/gqa/test_plot_gqa_figures.py — 7 tests, all green:
- 4 parametrized per-panel emit assertions
- 1 comparison emit assertion
- 1 emit_all returns exactly 5 PNG paths
- 1 default out_dir matches the bench _OUTPUT_DIR
Commits the 5 PNG baselines under the bench output dir alongside
sweep.json, mirroring milestone-1h-gemm's committed-figures pattern.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Self-contained eval bench (ADR-0054) that drives the four GQA Llama-70B
panels through run_bench with enable_data=True at validation scale and
emits sweep.json with the v1 schema (ADR-0057 D7).
Panel dispatch table maps each panel to (kernel, SFR install, S_q,
n_ranks, rank_axis):
single_user_prefill mesh_kv_kernel, intracube_pe_ring, S_q=16, n=8, rank_axis=0
multi_user_prefill mesh_kv_kernel, intercube_multisip, S_q=16, n=4, rank_axis=1
single_user_decode mesh_mlo_kernel, intracube_pe_ring, S_q=1, n=8, rank_axis=0
multi_user_decode mesh_mlo_kernel, intercube_multisip, S_q=1, n=4, rank_axis=1
multi_user panels pass _auto_dim_remap=False (avoid d_head=64
colliding with K's global M=64) and rank_axis=1 (cube-level ring,
gates 7 of every 8 PEs to silence).
Each panel runs on a fresh per-config GraphEngine, then op_log is
summarized into gemm/dma/ipcq counts. Both decode panels emit exactly
2*n_ranks GEMMs (one-shot partial attention per rank, ADR-0056 D3).
v1 supports GQA_VALIDATION=1 only; headline mode + figures deferred to
sub-cycles 4b/4c. Sentinel tensor satisfies the run_bench
"at least one request" contract (ADR-0045 D4 / ADR-0054 D2 carve-out).
Tests: tests/attention/test_milestone_gqa_llama70b.py — all 12 pass.
Includes committed sweep.json baseline at the bench's _OUTPUT_DIR so
subsequent test runs reuse it instead of re-simulating.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
ADR-0059 single_user_* panels run the ring across PEs in one cube
(rank == tl.program_id(axis=0)). multi_user_* panels run the ring
across cubes — rank should be cube_id (axis=1), and 7 of every 8 PEs
in each cube must stay silent because the cube-level SFR install only
gives the cube-coordinate PE 0 an E/W neighbor.
Add ``rank_axis: int = 0`` kwarg to both ``attention_mesh_mlo_kernel``
and ``attention_mesh_kv_kernel``:
- 0 (default): rank == tl.program_id(axis=0). Existing single_user
behavior, all spec tests unchanged.
- 1: gate ``if tl.program_id(axis=0) != 0: return`` at kernel start,
then ``rank = tl.program_id(axis=1)``. multi_user_* panels pass
this to the kernel via ctx.launch positional arg.
Also brings in _attention_mesh_kv.py and _attention_mesh_mlo.py as
the committed home of the ADR-0059 kernels (previously living
uncommitted in the working tree from sub-cycle 4b).
Tests: 7-test rank_axis spec file (default-path + rank_axis=1 gating
and cube-id semantics, both kernels); 4-panel diag harness now green
end-to-end (single_user_prefill/decode + multi_user_prefill/decode);
763-test wider sweep clean.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Per request, the milestone bench output is now tracked in git instead of
gitignored, so the figures/results are viewable on the remote:
- src/kernbench/benches/1H_milestone_output/gemm/ (3 PNGs + gemm_sweep.json)
- src/kernbench/benches/1H_milestone_output/ccl/ (3 per-topology PNGs,
buffer-kind PNG+CSV, FSIM comparison PNG, topology.png, summary.csv)
Drop the .gitignore rule; update ADR-0054 D3 + Negative (EN+KO) to say the
output is committed (regenerable by rerunning the bench). Artifacts produced
by full bench runs (milestone-1h-gemm non-FAST, milestone-1h-ccl).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Move the GEMM + allreduce sweep/render logic out of scripts/ and tests/
into two self-contained eval benches so a user can regenerate every
result + figure with one command:
kernbench run --bench milestone-1h-gemm (MILESTONE_FAST=1 reuses JSON)
kernbench run --bench milestone-1h-ccl
- benches/milestone_1h_{gemm,ccl}.py: single home for each domain; the
run(torch) entry drives the sweeps and writes figures into
benches/1H_milestone_output/{gemm,ccl}/ (gitignored), then submits a
sentinel tensor to satisfy the run_bench contract.
- tests/gemm + tests/sccl helpers and scripts/gemm_sweep.py become thin
re-export/wrapper shims over the benches (single source preserved); the
pytest-only param builders + _run_distributed wrapper stay in the shim.
- eval-bench pattern: a bench may drive many configs + build its own
per-config engines (extends ADR-0045 D5; reverses ADR-0044 D1/D2).
ADR-0054 (EN+KO) records the design; ADR-0043/0044/0045 + CLAUDE.md CLI
Semantics amended; ADR INDEX regenerated. Verified: milestone benches run
clean (ok=True, all artifacts), full suite 67 passed, lang-pairs OK.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Move benches/ -> src/kernbench/benches/ and src/kernbench/cli/probe.py ->
src/kernbench/probes/probe.py. Each bench self-registers via
@bench(name=..., description=...); kernbench list enumerates benches
with auto-assigned indices, --bench accepts kebab-case name or numeric
index. Audit at package-import time fails if any non-underscore module
forgets the decorator. ADR-0010 (EN + KO) updated to reflect the new
resolver path, list subcommand, and probes package separation.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
mukesh
ywkang