test short context length attention kernel

tests/attention/test_gqa_short_context.py

@@ -1,31 +1,17 @@
-"""Phase 1 spec test for Phase D: short-context GQA kernels
+"""Short-context GQA kernel tests — unified A1/A2/A4/B (ADR-0070).
-(ADR-0060 §B "Items from the long/short context split", item B.split.2).
 
-The long-context kernels (§5.2 decode / §5.5 prefill) shard each KV
+Both ``_gqa_attention_prefill_short.py`` and
-head row-wise across all CUBEs. At short context (S_kv < 256K), that
+``_gqa_attention_decode_short.py`` are single unified kernels selected
-shard is too thin to feed the engines and the cube-level collective
+at launch via ``kv_per_cube ∈ {1, 2, 4, 8}``:
-overhead dominates. The short-context kernels drop cube-SP entirely:
-each CUBE owns ``kv_per_cube`` *whole* KV heads, no S_kv sharding across
-CUBEs.
 
-Design (per AskUserQuestion answers in this session):
+    Mode  kv_per_cube  C        group_size
-  - PE-parallel heads: P PEs split into ``kv_per_cube`` groups, each
+    ----  -----------  -------  ----------
-    group does PE-SP across (P/kv_per_cube) PEs for one owned head.
+    A1            1    h_kv      P (=8)
-  - scratch_scope + tl.copy_to discipline mirrors the long kernels.
+    A2            2    h_kv/2    P/2 (=4)
-  - One short kernel handles kv_per_cube ∈ {1, 2, 4, 8} via a parameter.
+    A4            4    h_kv/4    P/4 (=2)
+    B             8    1         1
 
-Group layout on the 2×4 PE grid:
+Tests are mode-parametrized over the four (kv_per_cube, C) tuples.
-  kv_per_cube=1, group=8 PEs (full 2×4): row chain + col bridge.
-  kv_per_cube=2, group=4 PEs (one row):  row chain only.
-  kv_per_cube=4, group=2 PEs (adj cols): 1-step chain.
-  kv_per_cube=8, group=1 PE:             no chain.
-
-After chain reduce-to-group-root, the group's root PE writes its
-owned head's output. No inter-CUBE reduce.
-
-Phase 1 (this commit): tests only — production code lands in Phase 2.
-All tests fail because the short kernels (``_gqa_attention_decode_short.py`` and
-``_gqa_attention_prefill_short.py``) do not exist yet.
 """
 from __future__ import annotations
 
@@ -45,6 +31,17 @@ TOPOLOGY_DEFAULT = Path(__file__).resolve().parents[2] / "topology.yaml"
 
 D_HEAD = 64
 DTYPE = "f16"
+TILE_S_KV = 1024
+P = 8
+H_KV = 8
+
+# (kv_per_cube, C) for the four modes.
+MODES = [
+    pytest.param(1, 8, id="A1"),
+    pytest.param(2, 4, id="A2"),
+    pytest.param(4, 2, id="A4"),
+    pytest.param(8, 1, id="B"),
+]
 
 
 def _ccl_cfg():
@@ -61,166 +58,41 @@ def _count(op_log, name: str) -> int:
     return sum(1 for r in op_log if r.op_name == name)
 
 
-# ── Decode short-context kernel ──────────────────────────────────────
+# ── Prefill ──────────────────────────────────────────────────────────
 
 
-def _run_decode_short(*, h_q: int, h_kv: int, kv_per_cube: int,
+def _run_prefill(*, kv_per_cube: int, C: int, T_q: int, S_kv: int,
-                     C: int, P: int, S_kv: int):
+                 h_q: int = 8):
-    """Run the short-context decode kernel with PE-parallel heads.
+    """Run the unified prefill kernel in the given mode."""
-
+    from kernbench.benches._gqa_attention_prefill_short import (
-    Layout (after design iteration — see Phase D failure-recovery):
+        _validate_config as _validate_prefill_config,
-      Q: (T_q, h_q·D_HEAD) replicated; kernel reshapes byte-conservingly.
+        gqa_attention_prefill_short_kernel,
-      K, V: (h_kv·S_kv, D_HEAD) head-stacked, ``cube=row_wise, pe=row_wise``
+    )
-          so each PE gets a contiguous (S_local, D_HEAD) chunk at its own
+    _validate_prefill_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
-          addressable shard (no partial reads needed).
+                             h_q=h_q, h_kv=H_KV, S_kv=S_kv)
-      O: replicated; each group root writes the full byte-conserving
+    n_tiles = S_kv // TILE_S_KV
-          (h_q·T_q, D_HEAD) result.
+    Q_ROWS = kv_per_cube * T_q
-    """
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
-    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short import gqa_attention_decode_short_kernel  # Phase 2
-
     topo = resolve_topology(str(TOPOLOGY_DEFAULT))
 
     def _bench_fn(ctx):
         configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
-        dp_full = DPPolicy(cube="replicate", pe="replicate",
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
-                           num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P)
-        # Head-stacked KV with row_wise sharding so each PE's chunk is
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
-        # contiguous and exactly (S_local, D_HEAD) addressable.
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
-        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
+                      name=f"q_pre_kv{kv_per_cube}")
-                         num_cubes=C, num_pes=P)
+        k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV),
-        q = ctx.zeros((1, h_q * D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_pre_kv{kv_per_cube}")
-                      dtype=DTYPE, dp=dp_full, name=f"q_short_{kv_per_cube}_{C}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD),
-        k = ctx.zeros((h_kv * S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_kv{kv_per_cube}")
-                      dtype=DTYPE, dp=dp_kv, name=f"k_short_{kv_per_cube}_{C}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
-        v = ctx.zeros((h_kv * S_kv, D_HEAD),
+                      name=f"o_pre_kv{kv_per_cube}")
-                      dtype=DTYPE, dp=dp_kv, name=f"v_short_{kv_per_cube}_{C}")
-        o = ctx.empty((1, h_q * D_HEAD),
-                      dtype=DTYPE, dp=dp_full, name=f"o_short_{kv_per_cube}_{C}")
         ctx.launch(
-            f"gqa_decode_short_{kv_per_cube}_{C}",
+            f"gqa_prefill_kv{kv_per_cube}",
-            gqa_attention_decode_short_kernel,
-            q, k, v, o,
-            1, S_kv, h_q, h_kv, D_HEAD, C, P, kv_per_cube,
-            _auto_dim_remap=False,
-        )
-
-    return run_bench(
-        topology=topo, bench_fn=_bench_fn,
-        device=resolve_device(None), engine_factory=_engine_factory,
-    )
-
-
-def test_short_decode_smoke_kv_per_cube_2_C_4():
-    """ADR-0060 §B.split.2 headline: kv_per_cube=2, C=4 — each CUBE
-    owns 2 heads (8 heads / 4 CUBEs). PE-parallel heads splits the 8
-    PEs into 2 groups of 4, each group does PE-SP for one owned head.
-    Smoke: kernel completes."""
-    result = _run_decode_short(
-        h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
-    )
-    assert result.completion.ok, (
-        f"short decode kv_per_cube=2 must complete; got {result.completion}"
-    )
-
-
-def test_short_decode_smoke_kv_per_cube_4_C_2():
-    """kv_per_cube=4, C=2 — half the CUBEs participate, each owns 4
-    heads, 4 PE groups of 2 PEs each."""
-    result = _run_decode_short(
-        h_q=8, h_kv=8, kv_per_cube=4, C=2, P=8, S_kv=64,
-    )
-    assert result.completion.ok, (
-        f"short decode kv_per_cube=4 must complete; got {result.completion}"
-    )
-
-
-def test_short_decode_smoke_kv_per_cube_8_C_1():
-    """kv_per_cube=8, C=1 — all heads on one CUBE. 8 PE groups of 1 PE
-    each → no chain reduce, each PE writes its head's output."""
-    result = _run_decode_short(
-        h_q=8, h_kv=8, kv_per_cube=8, C=1, P=8, S_kv=64,
-    )
-    assert result.completion.ok, (
-        f"short decode kv_per_cube=8 must complete; got {result.completion}"
-    )
-
-
-def test_short_decode_dma_write_count_equals_h_kv():
-    """ADR-0060 §B.split.2: each owned head produces exactly one
-    output (no inter-CUBE reduce). Total dma_writes = h_kv across all
-    participating CUBEs and groups.
-
-    For h_kv=8, kv_per_cube=2, C=4: each CUBE writes 2 outputs →
-    4 × 2 = 8 dma_writes total.
-    """
-    result = _run_decode_short(
-        h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
-    )
-    assert result.completion.ok
-    n_writes = _count(result.engine.op_log, "dma_write")
-    assert n_writes == 8, (
-        f"short decode: expected 8 dma_writes (h_kv); got {n_writes}"
-    )
-
-
-def test_short_decode_no_inter_cube_traffic():
-    """ADR-0060 §B.split.2: each head is fully owned by one CUBE → no
-    inter-CUBE reduce. The kernel must not invoke CUBE-level E/W IPCQ.
-
-    Today's long-context kernel at C=4 emits ~12 inter-CUBE ipcq_copy
-    via direction "E"/"W". The short kernel must emit zero of those,
-    keeping all IPCQ traffic on the ``intra_*`` directions.
-    """
-    result = _run_decode_short(
-        h_q=8, h_kv=8, kv_per_cube=2, C=4, P=8, S_kv=64,
-    )
-    assert result.completion.ok
-    # Count IPCQ traffic that targeted the CUBE-level "E"/"W" directions.
-    inter_cube_ipcq = sum(
-        1 for r in result.engine.op_log
-        if r.op_name == "ipcq_copy"
-        and r.params.get("direction") in ("E", "W")
-    )
-    assert inter_cube_ipcq == 0, (
-        f"short decode must have no inter-CUBE E/W IPCQ; got {inter_cube_ipcq}"
-    )
-
-
-# ── Prefill short-context kernel ─────────────────────────────────────
-
-
-def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
-                      C: int, P: int, T_q: int, S_kv: int):
-    """Run the short-context prefill kernel with PE-parallel heads.
-
-    Layout: same head-stacked K/V scheme as decode short, with Q/O
-    replicated and byte-conserving reshape inside the kernel.
-    """
-    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import gqa_attention_prefill_short_kernel  # Phase 2
-
-    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
-
-    def _bench_fn(ctx):
-        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
-        dp_q = DPPolicy(cube="replicate", pe="replicate",
-                        num_cubes=C, num_pes=P)
-        dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
-                         num_cubes=C, num_pes=P)
-        dp_o = DPPolicy(cube="replicate", pe="replicate",
-                        num_cubes=C, num_pes=P)
-        q = ctx.zeros((T_q, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_q, name=f"q_pre_short_{kv_per_cube}_{C}")
-        k = ctx.zeros((h_kv * S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"k_pre_short_{kv_per_cube}_{C}")
-        v = ctx.zeros((h_kv * S_kv, D_HEAD),
-                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_short_{kv_per_cube}_{C}")
-        o = ctx.empty((T_q, h_kv * D_HEAD),
-                      dtype=DTYPE, dp=dp_o, name=f"o_pre_short_{kv_per_cube}_{C}")
-        ctx.launch(
-            f"gqa_prefill_short_{kv_per_cube}_{C}",
             gqa_attention_prefill_short_kernel,
             q, k, v, o,
-            T_q, S_kv, h_kv, D_HEAD, C, P, kv_per_cube,
+            T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
             _auto_dim_remap=False,
         )
 
@@ -230,31 +102,384 @@ def _run_prefill_short(*, h_kv: int, kv_per_cube: int,
     )
 
 
-def test_short_prefill_smoke_kv_per_cube_2_C_4():
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
-    """Short prefill kv_per_cube=2, C=4. Smoke: completes."""
+def test_prefill_smoke(kv_per_cube, C):
-    result = _run_prefill_short(
+    """All four prefill modes complete on a single-tile mini config."""
-        h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
-    )
+    assert r.completion.ok, (
-    assert result.completion.ok, (
+        f"prefill kv_per_cube={kv_per_cube}: {r.completion}"
-        f"short prefill kv_per_cube=2 must complete; got {result.completion}"
     )
 
 
-def test_short_prefill_no_ring_KV_traffic():
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
-    """ADR-0060 §B.split.2: short prefill DOES NOT use Ring KV — each
+def test_prefill_qtile_split_dma_writes(kv_per_cube, C):
-    CUBE owns its KV heads fully, no rotation. The kernel must not
+    """Every PE in every active group stores its q-tile output:
-    emit any KV-rotation IPCQ traffic at the CUBE level.
+    dma_writes = group_size · kv_per_cube · C."""
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
+    assert r.completion.ok
+    group_size = P // kv_per_cube
+    expected = group_size * kv_per_cube * C
+    n_writes = _count(r.engine.op_log, "dma_write")
+    assert n_writes == expected, (
+        f"prefill kv_per_cube={kv_per_cube}: expected {expected} "
+        f"dma_writes; got {n_writes}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_kv_hbm_read_collapsed(kv_per_cube, C):
+    """IPCQ KV broadcast collapses HBM K/V reads to one per group per tile:
+    dma_reads = (P · C) Q-reads + (kv_per_cube · 2 · n_tiles · C) KV-reads.
     """
-    result = _run_prefill_short(
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
-        h_kv=8, kv_per_cube=2, C=4, P=8, T_q=4, S_kv=64,
+    assert r.completion.ok
+    n_tiles = 1024 // TILE_S_KV
+    expected = (P * C) + (kv_per_cube * 2 * n_tiles * C)
+    n_reads = _count(r.engine.op_log, "dma_read")
+    assert n_reads == expected, (
+        f"prefill kv_per_cube={kv_per_cube}: expected {expected} dma_reads "
+        f"(Q + IPCQ-broadcasted KV); got {n_reads}"
     )
-    assert result.completion.ok
+
+
+def test_prefill_no_ring_KV_traffic():
+    """Short prefill never uses Ring KV — no inter-CUBE E/W IPCQ."""
+    r = _run_prefill(kv_per_cube=2, C=4, T_q=8, S_kv=1024)
+    assert r.completion.ok
     inter_cube_ipcq = sum(
-        1 for r in result.engine.op_log
+        1 for rec in r.engine.op_log
-        if r.op_name == "ipcq_copy"
+        if rec.op_name == "ipcq_copy"
-        and r.params.get("direction") in ("E", "W")
+        and rec.params.get("direction") in ("E", "W")
     )
     assert inter_cube_ipcq == 0, (
-        f"short prefill must have no Ring KV (no inter-CUBE E/W IPCQ); "
+        f"short prefill must have no inter-CUBE E/W IPCQ; "
         f"got {inter_cube_ipcq}"
     )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_multitile_scaling(kv_per_cube, C):
+    """Prefill dma_read scales linearly with n_tiles in every mode.
+
+    Per-mode formula (IPCQ broadcast: cube's group-root PE reads K/V):
+        dma_reads = P·C + 2·kv_per_cube·C·n_tiles
+                  = (Q: 1 per PE)  +  (K + V: 1 per group per tile)
+    """
+    for S_kv in (1024, 2048, 4096):
+        r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=S_kv)
+        assert r.completion.ok, f"kv={kv_per_cube} S_kv={S_kv}"
+        n_tiles = S_kv // TILE_S_KV
+        expected = (P * C) + (2 * kv_per_cube * C * n_tiles)
+        n_reads = _count(r.engine.op_log, "dma_read")
+        assert n_reads == expected, (
+            f"prefill kv={kv_per_cube} n_tiles={n_tiles}: expected "
+            f"{expected} dma_reads; got {n_reads}"
+        )
+
+
+# ── Decode ───────────────────────────────────────────────────────────
+
+
+def _run_decode(*, kv_per_cube: int, C: int, S_kv: int, h_q: int = 8):
+    """Run the unified decode kernel in the given mode."""
+    from kernbench.benches._gqa_attention_decode_short import (
+        _validate_config as _validate_decode_config,
+        gqa_attention_decode_short_kernel,
+    )
+    T_q = 1
+    _validate_decode_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P, C=C,
+                            h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
+    # K total elements = h_kv · S_kv · d_head; mode-invariant deploy.
+    k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P)
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                      name=f"q_dec_kv{kv_per_cube}")
+        k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
+                      name=f"k_dec_kv{kv_per_cube}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"v_dec_kv{kv_per_cube}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                      name=f"o_dec_kv{kv_per_cube}")
+        ctx.launch(
+            f"gqa_decode_kv{kv_per_cube}",
+            gqa_attention_decode_short_kernel,
+            q, k, v, o,
+            T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_smoke(kv_per_cube, C):
+    """All four decode modes complete at S_kv = 8K (min multi-tile-aligned
+    config that satisfies every mode's group-size constraint)."""
+    r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok, (
+        f"decode kv_per_cube={kv_per_cube}: {r.completion}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_chain_reduce_store_count(kv_per_cube, C):
+    """Decode chain-reduce: only group root (PE 0 per group) stores —
+    dma_writes = kv_per_cube · C (one per group root × cubes)."""
+    r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok
+    expected = kv_per_cube * C
+    n_writes = _count(r.engine.op_log, "dma_write")
+    assert n_writes == expected, (
+        f"decode kv_per_cube={kv_per_cube}: expected {expected} dma_writes "
+        f"(1 per group root); got {n_writes}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_multitile_per_pe(kv_per_cube, C):
+    """Decode multi-tile sweep path exercised in every mode.
+
+    S_kv is chosen so each PE owns 2 tiles (n_tiles_per_pe == 2),
+    forcing the post-tile-0 sweep loop to run.
+    """
+    group_size = P // kv_per_cube
+    S_kv = group_size * 2 * TILE_S_KV
+    r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv)
+    assert r.completion.ok, (
+        f"decode kv={kv_per_cube} S_kv={S_kv}: {r.completion}"
+    )
+
+
+# ── ADR-0011 D-VA1 contract regression tests ───────────────────────
+
+
+def _per_cube_dma_busy(op_log):
+    """Sum of dma_read t_end-t_start grouped by cube index."""
+    from collections import defaultdict
+    busy = defaultdict(float)
+    for rec in op_log:
+        if rec.op_name != "dma_read":
+            continue
+        cid = rec.component_id or ""
+        for part in cid.split("."):
+            if part.startswith("cube"):
+                try:
+                    busy[int(part[4:])] += rec.t_end - rec.t_start
+                except ValueError:
+                    pass
+                break
+    return dict(busy)
+
+
+def _per_cube_large_read_addrs(op_log, *, min_nbytes: int = 1024):
+    """Group dma_read src_addrs (>= min_nbytes) by issuing cube."""
+    from collections import defaultdict
+    per_cube: dict[int, set[int]] = defaultdict(set)
+    for rec in op_log:
+        if rec.op_name != "dma_read":
+            continue
+        if rec.params.get("nbytes", 0) < min_nbytes:
+            continue
+        cid = rec.component_id or ""
+        for part in cid.split("."):
+            if part.startswith("cube"):
+                try:
+                    per_cube[int(part[4:])].add(rec.params["src_addr"])
+                except ValueError:
+                    pass
+                break
+    return per_cube
+
+
+def _assert_cubes_disjoint(per_cube_addrs, label):
+    seen = sorted(per_cube_addrs.items())
+    for i, (ci, ai) in enumerate(seen):
+        for cj, aj in seen[i + 1:]:
+            assert ai.isdisjoint(aj), (
+                f"{label}: cube{ci} and cube{cj} share {len(ai & aj)} "
+                f"dma_read src_addrs; cubes must target disjoint HBM regions."
+            )
+
+
+def _assert_cube_balanced(busy, label, max_ratio=1.1):
+    if len(busy) <= 1:
+        return  # single-cube modes auto-pass
+    vals = list(busy.values())
+    ratio = max(vals) / min(vals) if min(vals) > 0 else 0
+    assert ratio < max_ratio, (
+        f"{label}: per-cube DMA_READ unbalanced: max/min = {ratio:.2f}× "
+        f"(expected < {max_ratio}×). Per-cube busy (μs): "
+        f"{[f'cube{c}={v/1000:.1f}' for c,v in sorted(busy.items())]}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_per_cube_disjoint_src_addrs(kv_per_cube, C):
+    """Decode: cubes must read from disjoint HBM regions (per ADR-0011
+    D-VA1, kernel computes cube/PE shard offset from program_id).
+
+    Regression guard: pre-fix all 64 PEs read offset 0 → all cubes
+    share the same src_addrs.
+    """
+    r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok
+    per_cube = _per_cube_large_read_addrs(r.engine.op_log)
+    assert len(per_cube) == C, (
+        f"decode kv={kv_per_cube}: expected {C} cubes; got {len(per_cube)}"
+    )
+    _assert_cubes_disjoint(per_cube, f"decode kv={kv_per_cube}")
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_per_cube_dma_balanced(kv_per_cube, C):
+    """Decode: per-cube DMA_READ busy must be ~equal (cubes operate on
+    independent local HBM). max/min < 1.1× for multi-cube modes.
+
+    Regression guard: pre-fix 11.5× ratio (cross-cube traffic to cube0).
+    """
+    r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok
+    busy = _per_cube_dma_busy(r.engine.op_log)
+    assert len(busy) == C, (
+        f"decode kv={kv_per_cube}: expected {C} cubes; got {sorted(busy.keys())}"
+    )
+    _assert_cube_balanced(busy, f"decode kv={kv_per_cube}")
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_per_cube_disjoint_src_addrs(kv_per_cube, C):
+    """Prefill: cubes must read from disjoint HBM regions."""
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
+    assert r.completion.ok
+    per_cube = _per_cube_large_read_addrs(r.engine.op_log)
+    assert len(per_cube) == C, (
+        f"prefill kv={kv_per_cube}: expected {C} cubes; got {len(per_cube)}"
+    )
+    _assert_cubes_disjoint(per_cube, f"prefill kv={kv_per_cube}")
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_per_cube_dma_balanced(kv_per_cube, C):
+    """Prefill: per-cube DMA_READ busy must be ~equal (cube-parallel)."""
+    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
+    assert r.completion.ok
+    busy = _per_cube_dma_busy(r.engine.op_log)
+    assert len(busy) == C, (
+        f"prefill kv={kv_per_cube}: expected {C} cubes; got {sorted(busy.keys())}"
+    )
+    _assert_cube_balanced(busy, f"prefill kv={kv_per_cube}")
+
+
+# ── Composite (second-level) variant smoke tests ────────────────────
+
+
+def _run_prefill_composite(*, kv_per_cube: int, C: int, T_q: int,
+                           S_kv: int, h_q: int = 8):
+    """Same helper as _run_prefill but for the composite-GEMM kernel."""
+    from kernbench.benches._gqa_attention_prefill_short_composite import (
+        _validate_config as _validate_prefill_cmp_config,
+        gqa_attention_prefill_short_composite_kernel,
+    )
+    _validate_prefill_cmp_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
+                                 C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    n_tiles = S_kv // TILE_S_KV
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                      name=f"q_pre_cmp_kv{kv_per_cube}")
+        k = ctx.zeros((H_KV * n_tiles * D_HEAD, TILE_S_KV),
+                      dtype=DTYPE, dp=dp_kv, name=f"k_pre_cmp_kv{kv_per_cube}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD),
+                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_cmp_kv{kv_per_cube}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                      name=f"o_pre_cmp_kv{kv_per_cube}")
+        ctx.launch(
+            f"gqa_prefill_cmp_kv{kv_per_cube}",
+            gqa_attention_prefill_short_composite_kernel,
+            q, k, v, o,
+            T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+def _run_decode_composite(*, kv_per_cube: int, C: int, S_kv: int,
+                          h_q: int = 8):
+    """Same helper as _run_decode but for the composite-GEMM kernel."""
+    from kernbench.benches._gqa_attention_decode_short_composite import (
+        _validate_config as _validate_decode_cmp_config,
+        gqa_attention_decode_short_composite_kernel,
+    )
+    T_q = 1
+    _validate_decode_cmp_config(kv_per_cube=kv_per_cube, T_q=T_q, P=P,
+                                C=C, h_q=h_q, h_kv=H_KV, S_kv=S_kv)
+    Q_ROWS = kv_per_cube * T_q
+    Q_COLS = (h_q * D_HEAD) // kv_per_cube
+    k_rows = (H_KV * S_kv * D_HEAD) // TILE_S_KV
+    topo = resolve_topology(str(TOPOLOGY_DEFAULT))
+
+    def _bench_fn(ctx):
+        configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
+        dp_q = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        dp_kv = DPPolicy(cube="row_wise", pe="row_wise", num_cubes=C, num_pes=P)
+        dp_o = DPPolicy(cube="column_wise", pe="replicate", num_cubes=C, num_pes=P)
+        q = ctx.zeros((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_q,
+                      name=f"q_dec_cmp_kv{kv_per_cube}")
+        k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
+                      name=f"k_dec_cmp_kv{kv_per_cube}")
+        v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
+                      name=f"v_dec_cmp_kv{kv_per_cube}")
+        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
+                      name=f"o_dec_cmp_kv{kv_per_cube}")
+        ctx.launch(
+            f"gqa_decode_cmp_kv{kv_per_cube}",
+            gqa_attention_decode_short_composite_kernel,
+            q, k, v, o,
+            T_q, S_kv, h_q, H_KV, D_HEAD, C, P, kv_per_cube,
+            _auto_dim_remap=False,
+        )
+
+    return run_bench(
+        topology=topo, bench_fn=_bench_fn,
+        device=resolve_device(None), engine_factory=_engine_factory,
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_prefill_composite_smoke(kv_per_cube, C):
+    """Composite-GEMM prefill completes in all 4 modes."""
+    r = _run_prefill_composite(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
+    assert r.completion.ok, (
+        f"prefill-composite kv_per_cube={kv_per_cube}: {r.completion}"
+    )
+
+
+@pytest.mark.parametrize("kv_per_cube,C", MODES)
+def test_decode_composite_smoke(kv_per_cube, C):
+    """Composite-GEMM decode completes in all 4 modes."""
+    r = _run_decode_composite(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
+    assert r.completion.ok, (
+        f"decode-composite kv_per_cube={kv_per_cube}: {r.completion}"
+    )

tests/attention/test_gqa_short_context_sweep_decode.py

@@ -0,0 +1,131 @@
+"""Decode mode-comparison sweep across context lengths.
+
+Measures per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}):
+  1. Wall clock latency (μs)
+  2. Hardware utilization
+       - GEMM engine utilization (Σ gemm time / wall × n_pe)
+       - MATH op count
+  3. HBM bandwidth utilization
+       (Σ dma_read+write bytes / (wall × n_pe × peak_pe_hbm_bw))
+  4. Communication cost (IPCQ bytes — chain reduce)
+  5. KV cache space cost (per-cube bytes)
+
+Output: ``docs/sweeps/short_context_decode_sweep.csv`` (16 rows).
+
+Run: ``pytest tests/attention/test_gqa_short_context_sweep_decode.py -m slow``
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+from tests.attention.test_gqa_short_context import (
+    D_HEAD, H_KV, P, TILE_S_KV, _run_decode,
+)
+
+PEAK_PE_HBM_BPS = 1024.0 / 8  # 128 bytes/ns
+
+MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
+            "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
+            "max_elem"}
+
+MODES = [
+    ("A1", 1, 8),
+    ("A2", 2, 4),
+    ("A4", 4, 2),
+    ("B",  8, 1),
+]
+
+CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
+
+CSV_OUT = (Path(__file__).resolve().parents[2]
+           / "docs" / "sweeps" / "short_context_decode_sweep.csv")
+
+
+def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
+    op_log = r.engine.op_log
+    wall_ns = max(rec.t_end for rec in op_log)
+
+    active_pes = set()
+    for rec in op_log:
+        cid = rec.component_id or ""
+        if cid and ".cube" in cid and ".pe" in cid:
+            parts = cid.split(".")
+            active_pes.add(".".join(parts[:3]))
+    n_pe = len(active_pes)
+
+    gemm_time_ns = 0.0
+    math_count = 0
+    dma_read_bytes = 0
+    dma_write_bytes = 0
+    ipcq_bytes = 0
+    for rec in op_log:
+        name = rec.op_name
+        nbytes = rec.params.get("nbytes", 0) or 0
+        if name == "gemm_f16":
+            gemm_time_ns += rec.t_end - rec.t_start
+        elif name in MATH_OPS:
+            math_count += 1
+        elif name == "dma_read":
+            dma_read_bytes += nbytes
+        elif name == "dma_write":
+            dma_write_bytes += nbytes
+        elif name == "ipcq_copy":
+            ipcq_bytes += nbytes
+
+    gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
+    hbm_bw_util = (
+        (dma_read_bytes + dma_write_bytes)
+        / (wall_ns * n_pe * PEAK_PE_HBM_BPS)
+    ) if n_pe else 0.0
+
+    kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
+
+    return {
+        "mode": mode,
+        "kv_per_cube": kv_per_cube,
+        "C": C,
+        "S_kv": S_kv,
+        "wall_us": wall_ns / 1000,
+        "n_pe": n_pe,
+        "gemm_util": gemm_util,
+        "math_count": math_count,
+        "hbm_bw_util": hbm_bw_util,
+        "hbm_read_mb": dma_read_bytes / (1024 * 1024),
+        "hbm_write_kb": dma_write_bytes / 1024,
+        "ipcq_kb": ipcq_bytes / 1024,
+        "kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
+    }
+
+
+@pytest.mark.slow
+def test_decode_mode_context_sweep():
+    """Run 4 modes × 4 context lengths, dump 5 metrics to CSV."""
+    rows = []
+    for mode, kv_per_cube, C in MODES:
+        for S_kv in CONTEXT_LENGTHS:
+            print(f">>> DECODE {mode} S_kv={S_kv//1024}K "
+                  f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
+            r = _run_decode(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv, h_q=64)
+            assert r.completion.ok, (
+                f"DECODE {mode} S_kv={S_kv}: {r.completion}"
+            )
+            m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
+                                 C=C, S_kv=S_kv)
+            rows.append(m)
+            print(f"    wall={m['wall_us']:.1f}μs  "
+                  f"gemm_util={m['gemm_util']:.3f}  "
+                  f"hbm_bw_util={m['hbm_bw_util']:.3f}  "
+                  f"ipcq={m['ipcq_kb']:.1f}KB  "
+                  f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
+                  flush=True)
+
+    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
+    with CSV_OUT.open("w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader()
+        for row in rows:
+            w.writerow(row)
+    print(f"\nWrote {len(rows)} rows to {CSV_OUT}")

tests/attention/test_gqa_short_context_sweep_prefill.py

@@ -0,0 +1,145 @@
+"""Prefill mode-comparison sweep across context lengths.
+
+Measures per (mode ∈ {A1, A2, A4, B}) × (S_kv ∈ {8K, 16K, 32K, 64K}):
+  1. Wall clock latency (μs)
+  2. Hardware utilization
+       - GEMM engine utilization (Σ gemm time / wall × n_pe)
+       - MATH op count
+  3. HBM bandwidth utilization
+       (Σ dma_read+write bytes / (wall × n_pe × peak_pe_hbm_bw))
+  4. Communication cost (IPCQ bytes — broadcast)
+  5. KV cache space cost (per-cube bytes)
+
+Sliced-prefill convention: T_q = 8.
+
+Output: ``docs/sweeps/prefill_sweep.csv`` (16 rows).
+
+Run: ``pytest tests/attention/test_gqa_short_context_sweep_prefill.py -m slow``
+"""
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+from tests.attention.test_gqa_short_context import (
+    D_HEAD, H_KV, P, TILE_S_KV, _run_prefill,
+)
+
+# Peak HBM BW per PE — derived from topology (cube total / PE count).
+# topology.yaml: hbm_total_bw_gbs = 1024.0 per cube, 8 PE per cube.
+PEAK_PE_HBM_BPS = 1024.0 / 8  # 128 bytes/ns
+
+MATH_OPS = {"max", "sum", "exp", "sub", "mul_bcast", "add", "div",
+            "maximum", "minimum", "neg", "rmax", "rsum", "exp_diff",
+            "max_elem"}
+
+MODES = [
+    ("A1", 1, 8),
+    ("A2", 2, 4),
+    ("A4", 4, 2),
+    ("B",  8, 1),
+]
+
+CONTEXT_LENGTHS = [8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024]
+
+T_Q_PREFILL = 8
+
+CSV_OUT = (Path(__file__).resolve().parents[2]
+           / "docs" / "sweeps" / "short_context_prefill_sweep.csv")
+
+
+def _compute_metrics(r, *, mode, kv_per_cube, C, S_kv):
+    op_log = r.engine.op_log
+    wall_ns = max(rec.t_end for rec in op_log)
+
+    # Per-PE detection (cube.pe)
+    active_pes = set()
+    for rec in op_log:
+        cid = rec.component_id or ""
+        if cid and ".cube" in cid and ".pe" in cid:
+            parts = cid.split(".")
+            active_pes.add(".".join(parts[:3]))
+    n_pe = len(active_pes)
+
+    # Stage time / count / bytes
+    gemm_time_ns = 0.0
+    math_count = 0
+    dma_read_bytes = 0
+    dma_write_bytes = 0
+    ipcq_bytes = 0
+    for rec in op_log:
+        name = rec.op_name
+        nbytes = rec.params.get("nbytes", 0) or 0
+        if name == "gemm_f16":
+            gemm_time_ns += rec.t_end - rec.t_start
+        elif name in MATH_OPS:
+            math_count += 1
+        elif name == "dma_read":
+            dma_read_bytes += nbytes
+        elif name == "dma_write":
+            dma_write_bytes += nbytes
+        elif name == "ipcq_copy":
+            ipcq_bytes += nbytes
+
+    gemm_util = gemm_time_ns / (wall_ns * n_pe) if n_pe else 0.0
+    hbm_bw_util = (
+        (dma_read_bytes + dma_write_bytes)
+        / (wall_ns * n_pe * PEAK_PE_HBM_BPS)
+    ) if n_pe else 0.0
+
+    # KV cache space — per-cube bytes (K + V, f16)
+    # Per cube owns kv_per_cube heads × S_kv tokens × d_head × 2 (K+V) × 2 (f16)
+    kv_cache_per_cube_bytes = kv_per_cube * S_kv * D_HEAD * 2 * 2
+
+    return {
+        "mode": mode,
+        "kv_per_cube": kv_per_cube,
+        "C": C,
+        "S_kv": S_kv,
+        "wall_us": wall_ns / 1000,
+        "n_pe": n_pe,
+        "gemm_util": gemm_util,
+        "math_count": math_count,
+        "hbm_bw_util": hbm_bw_util,
+        "hbm_read_mb": dma_read_bytes / (1024 * 1024),
+        "hbm_write_kb": dma_write_bytes / 1024,
+        "ipcq_kb": ipcq_bytes / 1024,
+        "kv_cache_per_cube_mb": kv_cache_per_cube_bytes / (1024 * 1024),
+    }
+
+
+@pytest.mark.slow
+def test_prefill_mode_context_sweep():
+    """Run 4 modes × 4 context lengths, dump 5 metrics to CSV.
+
+    Smoke-asserts each run completes; the measurement output is the CSV.
+    """
+    rows = []
+    for mode, kv_per_cube, C in MODES:
+        for S_kv in CONTEXT_LENGTHS:
+            print(f">>> PREFILL {mode} S_kv={S_kv//1024}K "
+                  f"(kv_per_cube={kv_per_cube}, C={C})", flush=True)
+            r = _run_prefill(kv_per_cube=kv_per_cube, C=C,
+                             T_q=T_Q_PREFILL, S_kv=S_kv, h_q=64)
+            assert r.completion.ok, (
+                f"PREFILL {mode} S_kv={S_kv}: {r.completion}"
+            )
+            m = _compute_metrics(r, mode=mode, kv_per_cube=kv_per_cube,
+                                 C=C, S_kv=S_kv)
+            rows.append(m)
+            print(f"    wall={m['wall_us']:.1f}μs  "
+                  f"gemm_util={m['gemm_util']:.3f}  "
+                  f"hbm_bw_util={m['hbm_bw_util']:.3f}  "
+                  f"ipcq={m['ipcq_kb']:.1f}KB  "
+                  f"kv_cache={m['kv_cache_per_cube_mb']:.1f}MB/cube",
+                  flush=True)
+
+    CSV_OUT.parent.mkdir(parents=True, exist_ok=True)
+    with CSV_OUT.open("w", newline="") as f:
+        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
+        w.writeheader()
+        for row in rows:
+            w.writerow(row)
+    print(f"\nWrote {len(rows)} rows to {CSV_OUT}")