"""Short-context GQA kernel tests — unified A1/A2/A4/B (ADR-0070).

Both ``_gqa_attention_prefill_short.py`` and
``_gqa_attention_decode_short.py`` are single unified kernels selected
at launch via ``kv_per_cube ∈ {1, 2, 4, 8}``:

    Mode  kv_per_cube  C        group_size
    ----  -----------  -------  ----------
    A1            1    h_kv      P (=8)
    A2            2    h_kv/2    P/2 (=4)
    A4            4    h_kv/4    P/4 (=2)
    B             8    1         1

Tests are mode-parametrized over the four (kv_per_cube, C) tuples.
"""
from __future__ import annotations

from pathlib import Path

import pytest

from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config
from kernbench.ccl.sfr_config import configure_sfr_intercube_multisip
from kernbench.policy.placement.dp import DPPolicy
from kernbench.runtime_api.bench_runner import run_bench
from kernbench.runtime_api.types import resolve_device
from kernbench.sim_engine.engine import GraphEngine
from kernbench.topology.builder import resolve_topology

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():
    return resolve_algorithm_config(
        load_ccl_config(), name="lrab_hierarchical_allreduce",
    )


def _engine_factory(t, d):
    return GraphEngine(getattr(t, "topology_obj", t), enable_data=True)


def _count(op_log, name: str) -> int:
    return sum(1 for r in op_log if r.op_name == name)


# ── Prefill ──────────────────────────────────────────────────────────


def _run_prefill(*, kv_per_cube: int, C: int, T_q: int, S_kv: int,
                 h_q: int = 8):
    """Run the unified prefill kernel in the given mode."""
    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short import (
        _validate_config as _validate_prefill_config,
        gqa_attention_prefill_short_kernel,
    )
    _validate_prefill_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_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_kv{kv_per_cube}")
        v = ctx.zeros((H_KV * S_kv, D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_kv{kv_per_cube}")
        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
                      name=f"o_pre_kv{kv_per_cube}")
        ctx.launch(
            f"gqa_prefill_kv{kv_per_cube}",
            gqa_attention_prefill_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_prefill_smoke(kv_per_cube, C):
    """All four prefill modes complete on a single-tile mini config."""
    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
    assert r.completion.ok, (
        f"prefill kv_per_cube={kv_per_cube}: {r.completion}"
    )


@pytest.mark.parametrize("kv_per_cube,C", MODES)
def test_prefill_qtile_split_dma_writes(kv_per_cube, C):
    """Every PE in every active group stores its q-tile output:
    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.
    """
    r = _run_prefill(kv_per_cube=kv_per_cube, C=C, T_q=8, S_kv=1024)
    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}"
    )


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 rec in r.engine.op_log
        if rec.op_name == "ipcq_copy"
        and rec.params.get("direction") in ("E", "W")
    )
    assert inter_cube_ipcq == 0, (
        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_helpers.short_ctx._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_helpers.short_ctx._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_helpers.short_ctx._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,
    )


def _run_prefill_composite_fused(*, kv_per_cube: int, C: int, T_q: int,
                                 S_kv: int, h_q: int = 8):
    """Variant (3): GEMM-only composite + softmax_merge prologue fusion."""
    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_prefill_short_composite_fused import (
        _validate_config as _validate_prefill_fuse_config,
        gqa_attention_prefill_short_composite_fused_kernel,
    )
    _validate_prefill_fuse_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_fuse_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_fuse_kv{kv_per_cube}")
        v = ctx.zeros((H_KV * S_kv, D_HEAD),
                      dtype=DTYPE, dp=dp_kv, name=f"v_pre_fuse_kv{kv_per_cube}")
        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
                      name=f"o_pre_fuse_kv{kv_per_cube}")
        ctx.launch(
            f"gqa_prefill_fuse_kv{kv_per_cube}",
            gqa_attention_prefill_short_composite_fused_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_fused(*, kv_per_cube: int, C: int, S_kv: int,
                                h_q: int = 8):
    """Variant (3): GEMM-only composite + softmax_merge prologue fusion."""
    from kernbench.benches.gqa_helpers.short_ctx._gqa_attention_decode_short_composite_fused import (
        _validate_config as _validate_decode_fuse_config,
        gqa_attention_decode_short_composite_fused_kernel,
    )
    T_q = 1
    _validate_decode_fuse_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_fuse_kv{kv_per_cube}")
        k = ctx.zeros((k_rows, TILE_S_KV), dtype=DTYPE, dp=dp_kv,
                      name=f"k_dec_fuse_kv{kv_per_cube}")
        v = ctx.zeros((H_KV * S_kv, D_HEAD), dtype=DTYPE, dp=dp_kv,
                      name=f"v_dec_fuse_kv{kv_per_cube}")
        o = ctx.empty((Q_ROWS, Q_COLS), dtype=DTYPE, dp=dp_o,
                      name=f"o_dec_fuse_kv{kv_per_cube}")
        ctx.launch(
            f"gqa_decode_fuse_kv{kv_per_cube}",
            gqa_attention_decode_short_composite_fused_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):
    """Variant (2) GEMM-only composite prefill — 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):
    """Variant (2) GEMM-only composite decode — 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}"
    )


@pytest.mark.parametrize("kv_per_cube,C", MODES)
def test_decode_composite_fused_smoke(kv_per_cube, C):
    """Variant (3) composite + softmax_merge fused decode — all 4 modes."""
    r = _run_decode_composite_fused(kv_per_cube=kv_per_cube, C=C, S_kv=8192)
    assert r.completion.ok, (
        f"decode-composite-fused kv_per_cube={kv_per_cube}: {r.completion}"
    )


@pytest.mark.parametrize("kv_per_cube,C", MODES)
def test_decode_composite_fused_multitile(kv_per_cube, C):
    """Fused decode with n_tiles_per_pe == 2 in EVERY mode, so the tile-1+
    fused composite loop actually runs.

    The smoke test at S_kv=8192 gives A1 (group_size=8) n_tiles_per_pe=1,
    which never enters the fused path. S_kv = group_size·2·TILE_S_KV forces
    2 tiles per PE for all modes (mirrors test_decode_multitile_per_pe)."""
    group_size = P // kv_per_cube
    S_kv = group_size * 2 * TILE_S_KV
    r = _run_decode_composite_fused(kv_per_cube=kv_per_cube, C=C, S_kv=S_kv)
    assert r.completion.ok, (
        f"decode-composite-fused-multitile kv={kv_per_cube} S_kv={S_kv}: "
        f"{r.completion}"
    )


@pytest.mark.parametrize("kv_per_cube,C", MODES)
def test_prefill_composite_fused_smoke(kv_per_cube, C):
    """Variant (3) composite + softmax_merge fused prefill — all 4 modes.

    S_kv=2048 (n_tiles=2) so the tile-1+ fused composite loop actually
    runs; S_kv=1024 gives n_tiles=1 and never enters the fused path.
    Multi-cube (A1/A2/A4) works now that recv'd K/V slots are pinned."""
    r = _run_prefill_composite_fused(kv_per_cube=kv_per_cube, C=C,
                                     T_q=8, S_kv=2048)
    assert r.completion.ok, (
        f"prefill-composite-fused kv_per_cube={kv_per_cube}: {r.completion}"
    )
