17fb94086eb08917b4d9f59d8e16e501f5d6f166
9 Commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
 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
|
91cdfebb67 |
gqa: S_kv tile sweep for decode_long/decode_short/prefill_short (ADR-0063 §A.2)
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> |
||
|
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>
|
||
|
mukesh
|
39fc2e953f |
attention: add 8-KV-group diag harness (CCL pattern + cube_start sub-meshes)
New ``test_attention_8kv_groups_diag.py`` probes the path from
validation scale to the GQA Llama-70B 1-Q-head-per-cube headline
(64 cubes = 8 KV-groups x 8 cubes/group on a 4-SIP topology) in
three incremental steps:
step_1_single_kv_group_at_full_breadth
ONE 2x4 multi_user_decode launch (8 cubes) via the 2D mesh-mlo
kernel. Verifies the per-KV-group 2D AllReduce works at full
breadth.
step_2_four_kv_groups_one_per_sip
Four sequential 2x4 launches, one per SIP. Uses the CCL
milestone's set_device pattern (milestone_1h_ccl.py:283-292):
ONE run_bench with ``target_device="all"`` and
``ctx.ahbm.set_device(sip)`` between launches — not four
separate run_bench calls with ``DeviceSelector("sip:N")``
(which would misuse DeviceSelector and hit a
``DPPolicy x target_device`` allocator mismatch).
step_3_eight_kv_groups_two_per_sip
Two 2x4 launches per SIP x 4 SIPs = 8 KV-groups (64 cubes
total). Second launch per SIP uses ``cube_start=8`` to address
cubes 8..15 — the disjoint sub-mesh that was unreachable before
``DPPolicy.cube_start`` landed. Demonstrates the headline-enabling
use of cube_start end-to-end (DPPolicy + kernel kwarg).
All three steps pass at validation dims (S_q=1, S_kv=16, h=1,
d_head=64). Headline-scale dims, multi_user_prefill 2D, and the B=8
"Batch on batch" PE parallelism remain follow-up work.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
||
|
mukesh
|
4859149392 |
attention: add 2D row-then-col AllReduce-mlo decode kernel (C2)
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> |
||
|
mukesh
|
e748a62264 |
attention: land milestone-gqa-llama70b 4-panel sweep bench (ADR-0057 v1)
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> |
||
|
mukesh
|
222815d374 |
attention: add rank_axis kwarg to mesh kernels for multi_user cube ring
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>
|
||
|
mukesh
|
d9e767d048 |
runtime_api: ctx.launch honors DPPolicy.num_cubes + adds _auto_dim_remap opt-out
Two compounding bugs in ctx.launch's dim-translation path surfaced by multi_user_* panels of milestone-gqa-llama70b (sub-cycle 4c step 2): Bug A: _compute_local_shape divided by self._num_cubes (the topology's cube count, 16 in default topology.yaml) instead of the DPPolicy's effective num_cubes (4 for validation-scale multi_user). The tensor allocator at context.py:471-484 already honored dp.num_cubes; the parallel computation inside launch was out of sync. Fix mirrors the allocator's eff_num_cubes precedence pattern. Bug B: dim_map was keyed by value, so any scalar whose value coincidentally equaled a global tensor dim got rewritten to that dim's local value — e.g. d_head=64 colliding with K's global M=64 in multi_user mode. Legacy bench kernels (va_offset etc.) rely on this remap, so the fix is opt-out: ctx.launch(..., _auto_dim_remap=False) preserves scalars exactly as passed. Default remains True. Tests: 3 new dim-translation tests + 4-panel diag harness covers single_user_* (PASS) and multi_user_* (advances to new SFR/axis layer failure, tracked separately). va_offset + full attention spec suite unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> |