From d0f904c57eac0d95d7a0e48ecd4c080fdce7872e Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Tue, 2 Jun 2026 14:38:30 -0700 Subject: [PATCH] =?UTF-8?q?attention:=20docstring=20cleanup=20=E2=80=94=20?= =?UTF-8?q?clarify=20single/multi=20user=20per=20GQA=20study?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- src/kernbench/benches/_attention_mesh_kv.py | 21 ++++++++++++--- src/kernbench/benches/_attention_mesh_mlo.py | 24 ++++++++++++++--- .../benches/milestone_gqa_llama70b.py | 26 +++++++++++++++++-- 3 files changed, 61 insertions(+), 10 deletions(-) diff --git a/src/kernbench/benches/_attention_mesh_kv.py b/src/kernbench/benches/_attention_mesh_kv.py index df03fa9..cc61dd0 100644 --- a/src/kernbench/benches/_attention_mesh_kv.py +++ b/src/kernbench/benches/_attention_mesh_kv.py @@ -82,10 +82,23 @@ def attention_mesh_kv_kernel( ) -> None: """Mesh-native bidirectional Ring-K/V attention — see module docstring. - ``rank_axis`` selects which program-id dimension carries the ring rank: - 0 — single_user_* panels: rank == tl.program_id(axis=0) (PE id in cube). - 1 — multi_user_* panels: ring is at the cube level. Only PE 0 in each - cube participates; the other 7 hold KV replicas but stay silent. + ``rank_axis`` selects which program-id dimension carries the ring rank, + matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction + (`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`): + + 0 — single_user_* panels (TL/BL): rank == tl.program_id(axis=0) (PE + id in cube). KV is split @ PEs **intra-cube**; ring runs over + the 8 PEs of one cube (NOC-only). At Llama-70B headline scale + this kernel launches once per cube; 64 such cubes run in + parallel for one user. Each cube's PE-level ring is independent. + + 1 — multi_user_* panels (TR/BR): rank == tl.program_id(axis=1) + (cube id). KV is split @ cubes **inter-cube**; ring runs over + the cubes of one KV-group. The kernel gates ``pe_id != 0`` to + return early — a v1 simplification: at headline scale (B=8) the + study's "Batch on batch" pattern would have all 8 PEs each handle + one user's batch element instead of staying silent. Adding the + per-cube batch dimension is sub-cycle 4c headline work. """ # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring. if rank_axis != 0 and tl.program_id(axis=0) != 0: diff --git a/src/kernbench/benches/_attention_mesh_mlo.py b/src/kernbench/benches/_attention_mesh_mlo.py index 2474626..1e1b2bb 100644 --- a/src/kernbench/benches/_attention_mesh_mlo.py +++ b/src/kernbench/benches/_attention_mesh_mlo.py @@ -53,10 +53,26 @@ def attention_mesh_mlo_kernel( ) -> None: """Mesh-native bidirectional AllReduce-mlo — see module docstring. - ``rank_axis`` selects which program-id dimension carries the ring rank: - 0 — single_user_* panels: rank == tl.program_id(axis=0) (PE id in cube). - 1 — multi_user_* panels: ring is at the cube level. Only PE 0 in each - cube participates; the other 7 hold KV replicas but stay silent. + ``rank_axis`` selects which program-id dimension carries the ring rank, + matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction + (`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`): + + 0 — single_user_* panels (TL/BL): rank == tl.program_id(axis=0) (PE + id in cube). KV is split @ PEs **intra-cube**; ring runs over + the 8 PEs of one cube (NOC-only). At Llama-70B headline scale + this kernel launches once per cube; 64 such cubes run in + parallel for one user (1 Q-head per cube × 8 cubes per KV-group + × 8 KV-groups). The PE-level ring inside each cube is + independent of the others. + + 1 — multi_user_* panels (TR/BR): rank == tl.program_id(axis=1) + (cube id). KV is split @ cubes **inter-cube**; ring runs over + the cubes of one KV-group. The kernel gates ``pe_id != 0`` to + return early — a v1 simplification: at headline scale (B=8) the + study's "Batch on batch" pattern would have all 8 PEs each handle + one user's batch element instead of staying silent. Validation + shipped with B=1 to focus on the inter-cube ring's correctness; + adding the per-cube batch dimension is sub-cycle 4c headline work. """ # For multi_user (rank_axis=1) only PE 0 in each cube runs the ring. if rank_axis != 0 and tl.program_id(axis=0) != 0: diff --git a/src/kernbench/benches/milestone_gqa_llama70b.py b/src/kernbench/benches/milestone_gqa_llama70b.py index 86dbf54..b3a38cb 100644 --- a/src/kernbench/benches/milestone_gqa_llama70b.py +++ b/src/kernbench/benches/milestone_gqa_llama70b.py @@ -14,12 +14,34 @@ v1 (sub-cycle 4a + 4c.0) covers all four panels at validation scale: single_user_decode BL configure_sfr_intracube_pe_ring multi_user_decode BR configure_sfr_intercube_multisip +Per the GQA sharding study (`llm_paper_review/notes/GQA_MHA_sharding/scripts +/_gen_llama70b_1M_4cases.py`): + + Single User (B=1) panels — TL prefill, BL decode: + "n_cubes: 8 (1 KV-group)", KV split @ PEs intra-cube. Each cube does + its own 8-PE ring with no cube-to-cube attention traffic. At Llama-70B + headline scale this is 64 cubes (8 KV-groups × 8 cubes/group), each + independently running the kernel for one Q-head; 1 user spans all 64. + + Multi User (B=8) panels — TR prefill, BR decode: + "8 cubes / KV-group", KV split @ cubes inter-cube. The 8 cubes of a + KV-group form a ring; "Inside each cube: 8 PEs each handle 1 different + user → Batch on batch, batch = 8/cube." At headline scale 8 KV-groups + run in parallel = 64 cubes serving 8 users. + Kernels use the mesh-native variants (ADR-0059), invoked with the ``rank_axis`` kwarg (0 for single_user PE-level rings, 1 for multi_user -cube-level rings — the latter also gates 7 of every 8 PEs to silence). +cube-level rings). The v1 multi_user kernel gates ``pe_id != 0`` to return, +which simplifies B=8 → B=1 — that's a validation simplification; the +per-cube "Batch on batch" parallelism is sub-cycle 4c headline work. Validation-scale config (ADR-0057 D4) — kept small so the simulator's -1 MB per-PE TCM scratch budget is not exhausted across n_ranks ring steps. +1 MB per-PE TCM scratch budget is not exhausted across n_ranks ring steps: +``S_q_prefill = S_kv_per_rank = 16``, ``h_q = h_kv = 1``, ``d_head = 64``, +``n_ranks_single_user = 8`` (8 PEs of one cube), ``n_ranks_multi_user = 4`` +(half a KV-group, vs the study's 8). Headline-scale dims (``S_q = 1M``, +``S_kv = 1M``, ``h_q = 8 / h_kv = 1`` GQA, ``d_head = 128``, ``B = 8``) are +also deferred. """ from __future__ import annotations