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>
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user