Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 80510b89a9 | |||
| d0f904c57e |
@@ -82,10 +82,23 @@ def attention_mesh_kv_kernel(
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""Mesh-native bidirectional Ring-K/V attention — see module docstring.
|
"""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.
|
# 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:
|
if rank_axis != 0 and tl.program_id(axis=0) != 0:
|
||||||
|
|||||||
@@ -53,10 +53,26 @@ def attention_mesh_mlo_kernel(
|
|||||||
) -> None:
|
) -> None:
|
||||||
"""Mesh-native bidirectional AllReduce-mlo — see module docstring.
|
"""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.
|
# 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:
|
if rank_axis != 0 and tl.program_id(axis=0) != 0:
|
||||||
|
|||||||
@@ -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
|
single_user_decode BL configure_sfr_intracube_pe_ring
|
||||||
multi_user_decode BR configure_sfr_intercube_multisip
|
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
|
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
|
``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
|
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
|
from __future__ import annotations
|
||||||
|
|
||||||
|
|||||||
@@ -161,3 +161,79 @@ def configure_sfr_intercube_multisip(
|
|||||||
algo_module=mock_module,
|
algo_module=mock_module,
|
||||||
rank_to_pe=pe_idx_to_pe,
|
rank_to_pe=pe_idx_to_pe,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
# ── Intra-cube PE ring (ADR-0058 Proposed) ─────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def configure_sfr_intracube_pe_ring(
|
||||||
|
engine: Any,
|
||||||
|
spec: dict,
|
||||||
|
cfg: dict,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""Install an 8-PE logical ring inside every cube on every SIP.
|
||||||
|
|
||||||
|
Per cube ``c`` on every SIP, every PE ``i`` ∈ [0, pes_per_cube) gets:
|
||||||
|
|
||||||
|
nbrs[i]["E"] = pe ((i + 1) % pes_per_cube) on the same cube and sip
|
||||||
|
nbrs[i]["W"] = pe ((i - 1) % pes_per_cube) on the same cube and sip
|
||||||
|
|
||||||
|
No intercube or inter-SIP edges are installed by this function — the
|
||||||
|
single_user_* attention panels (ADR-0057) operate inside one cube,
|
||||||
|
so cross-cube traffic is architecturally not part of their kernel.
|
||||||
|
|
||||||
|
The ``E``/``W`` direction namespace is shared with
|
||||||
|
``configure_sfr_intercube_multisip`` (which writes cube-mesh
|
||||||
|
same-lane edges to the same names). The two installs are
|
||||||
|
**mutually exclusive on the same engine** — the bench is responsible
|
||||||
|
for picking one per panel-run (ADR-0058 D2). v1 does not add a
|
||||||
|
runtime guard.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
engine: GraphEngine with ``_components``.
|
||||||
|
spec: topology spec dict (from topology.yaml).
|
||||||
|
cfg: merged algorithm config (from ``resolve_algorithm_config``).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
The install plan dict from ``install_ipcq`` (same shape as
|
||||||
|
``configure_sfr_intercube_multisip``).
|
||||||
|
"""
|
||||||
|
cm = spec["sip"]["cube_mesh"]
|
||||||
|
n_cubes = int(cm["w"]) * int(cm["h"])
|
||||||
|
n_sips = int(spec.get("system", {}).get("sips", {}).get("count", 1))
|
||||||
|
pl = spec["cube"]["pe_layout"]
|
||||||
|
pes_per_cube = int(pl["pe_per_corner"]) * len(pl["corners"])
|
||||||
|
|
||||||
|
world_size = n_sips * n_cubes * pes_per_cube
|
||||||
|
pe_idx_to_pe: list[tuple[int, int, int]] = [
|
||||||
|
(sip, cube, pe)
|
||||||
|
for sip in range(n_sips)
|
||||||
|
for cube in range(n_cubes)
|
||||||
|
for pe in range(pes_per_cube)
|
||||||
|
]
|
||||||
|
|
||||||
|
def _pe_idx(sip: int, cube: int, pe: int) -> int:
|
||||||
|
return (sip * n_cubes + cube) * pes_per_cube + pe
|
||||||
|
|
||||||
|
def _neighbors(pe_idx: int, ws: int, _base: dict) -> dict[str, int]:
|
||||||
|
tmp = pe_idx
|
||||||
|
pe = tmp % pes_per_cube
|
||||||
|
tmp //= pes_per_cube
|
||||||
|
cube = tmp % n_cubes
|
||||||
|
sip = tmp // n_cubes
|
||||||
|
return {
|
||||||
|
"E": _pe_idx(sip, cube, (pe + 1) % pes_per_cube),
|
||||||
|
"W": _pe_idx(sip, cube, (pe - 1) % pes_per_cube),
|
||||||
|
}
|
||||||
|
|
||||||
|
mock_module = types.SimpleNamespace(neighbors=_neighbors)
|
||||||
|
|
||||||
|
cfg_copy = dict(cfg)
|
||||||
|
cfg_copy["world_size"] = world_size
|
||||||
|
cfg_copy["topology"] = "none"
|
||||||
|
|
||||||
|
return install_ipcq(
|
||||||
|
engine, spec, cfg_copy,
|
||||||
|
algo_module=mock_module,
|
||||||
|
rank_to_pe=pe_idx_to_pe,
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user