attention: docstring cleanup — clarify single/multi user per GQA study

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>

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:
+    ``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).
+    matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction
-        1 — multi_user_* panels: ring is at the cube level. Only PE 0 in each
+    (`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`):
-            cube participates; the other 7 hold KV replicas but stay silent.
+
+      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:

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:
+    ``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).
+    matching the GQA Llama-70B sharding study's TL/BL vs TR/BR distinction
-        1 — multi_user_* panels: ring is at the cube level. Only PE 0 in each
+    (`llm_paper_review/notes/GQA_MHA_sharding/scripts/_gen_llama70b_1M_4cases.py`):
-            cube participates; the other 7 hold KV replicas but stay silent.
+
+      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:

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