New ``test_attention_8kv_groups_diag.py`` probes the path from
validation scale to the GQA Llama-70B 1-Q-head-per-cube headline
(64 cubes = 8 KV-groups x 8 cubes/group on a 4-SIP topology) in
three incremental steps:
step_1_single_kv_group_at_full_breadth
ONE 2x4 multi_user_decode launch (8 cubes) via the 2D mesh-mlo
kernel. Verifies the per-KV-group 2D AllReduce works at full
breadth.
step_2_four_kv_groups_one_per_sip
Four sequential 2x4 launches, one per SIP. Uses the CCL
milestone's set_device pattern (milestone_1h_ccl.py:283-292):
ONE run_bench with ``target_device="all"`` and
``ctx.ahbm.set_device(sip)`` between launches — not four
separate run_bench calls with ``DeviceSelector("sip:N")``
(which would misuse DeviceSelector and hit a
``DPPolicy x target_device`` allocator mismatch).
step_3_eight_kv_groups_two_per_sip
Two 2x4 launches per SIP x 4 SIPs = 8 KV-groups (64 cubes
total). Second launch per SIP uses ``cube_start=8`` to address
cubes 8..15 — the disjoint sub-mesh that was unreachable before
``DPPolicy.cube_start`` landed. Demonstrates the headline-enabling
use of cube_start end-to-end (DPPolicy + kernel kwarg).
All three steps pass at validation dims (S_q=1, S_kv=16, h=1,
d_head=64). Headline-scale dims, multi_user_prefill 2D, and the B=8
"Batch on batch" PE parallelism remain follow-up work.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
mukesh
39fc2e953f