gqa(decode-4cases): Case 4 anchor in dedicated bench (5C.D)

First milestone of the decode 4-cases comparative study per
GQA_full_deck.pptx slides 11-17. Case 4 (Cube-SP × PE-SP, the optimal
case per slide 11) is structurally the existing _gqa_attention_decode_long
at sub_w=4 (Increment 2's lrab-adapted center-root reduce). This commit
wires it into a dedicated bench so the remaining cases land alongside.

Changes:
- New bench: src/kernbench/benches/milestone_gqa_decode_4cases.py
  Houses the 4 case panels under one milestone-gqa-decode-4cases entry
  (gated by GQA_DECODE_4CASES_RUN=1; output to
   1H_milestone_output/gqa_decode_4cases/sweep.json). Cases 1-3 are
  TBD in subsequent sub-increments (5C.A/B/C).
- New panel: single_kv_group_decode_gqa_cube_sp_pe_sp
  C=8, P=8, sub_w=4, T_q=1, S_kv=131_072, d_head=128, h_q=8, h_kv=1.
- src/kernbench/benches/milestone_gqa_headline.py: _run_decode_panel
  extended with keyword-only sub_w/T_q/d_head/h_q/h_kv overrides
  (defaults preserve existing-panel behaviour).
- tests/attention/test_milestone_gqa_decode_4cases.py: 4 new tests
  asserting registration, smoke completion, reduce-to-root at the lrab
  center cube (cube 6), and the predicted 189-ipcq Case-4 traffic
  pattern (168 intra-CUBE + 21 inter-CUBE lrab Phase 1+2).
- tests/attention/test_milestone_gqa_headline.py: rename
  test_sweep_json_has_four_panels -> test_sweep_json_has_expected_panels
  and switch hardcoded 4 to len(PANELS) (the panel set grew to 5
  with Increment 5's single_kv_group_prefill_gqa_c8_p8).

Deviation noted: slide 13 prescribes AllReduce on (m,ℓ,O); our kernel
does reduce-to-root (only the lrab center cube has the answer) per
ADR-0060 §4. Treated as the kernbench Case-4 baseline.

Verification: all 4 new tests pass; 90 regression tests pass; the
previously-failing test_sweep_json_has_four_panels now passes under
its renamed form.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
2026-06-15 14:31:44 -07:00
parent 5b4d9cb597
commit c164645aee
4 changed files with 356 additions and 8 deletions
@@ -143,24 +143,31 @@ def _run_prefill_panel(
)
def _run_decode_panel(ctx, *, panel: str, C: int, P: int, S_kv: int) -> None:
def _run_decode_panel(
ctx, *, panel: str, C: int, P: int, S_kv: int,
sub_w: int = 0,
T_q: int = _T_Q_DECODE,
d_head: int = _D_HEAD,
h_q: int = _H_Q_DECODE,
h_kv: int = _H_KV_DECODE,
) -> None:
configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
dp_full = DPPolicy(cube="replicate", pe="replicate",
num_cubes=C, num_pes=P)
dp_kv = DPPolicy(cube="row_wise" if C > 1 else "replicate",
pe="row_wise", num_cubes=C, num_pes=P)
q = ctx.zeros((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
q = ctx.zeros((T_q, h_q * d_head),
dtype=_DTYPE, dp=dp_full, name=f"{panel}_q")
k = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
k = ctx.zeros((S_kv, h_kv * d_head),
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_k")
v = ctx.zeros((S_kv, _H_KV_DECODE * _D_HEAD),
v = ctx.zeros((S_kv, h_kv * d_head),
dtype=_DTYPE, dp=dp_kv, name=f"{panel}_v")
o = ctx.empty((_T_Q_DECODE, _H_Q_DECODE * _D_HEAD),
o = ctx.empty((T_q, h_q * d_head),
dtype=_DTYPE, dp=dp_full, name=f"{panel}_o")
ctx.launch(
panel, gqa_attention_decode_long_kernel,
q, k, v, o,
_T_Q_DECODE, S_kv, _H_Q_DECODE, _H_KV_DECODE, _D_HEAD, C, P,
T_q, S_kv, h_q, h_kv, d_head, C, P, sub_w,
_auto_dim_remap=False,
)