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