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:
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user