kernbench2/scripts/verify_case4_dhead_tp.py

"""Phase 1a smoke test for the Case 4 d_head-TP decode kernel.

Runs `gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel` at small
S_kv (2K) to verify:
  1. it imports cleanly,
  2. it runs without error under the same harness as Case 6,
  3. captures op_log_summary (gemm/ipcq/dma counts),
  4. compares vs analytical predictions and vs Case 6 (same memory tier).

Usage:
  python scripts/verify_case4_dhead_tp.py
"""
from __future__ import annotations

import os
import sys
from pathlib import Path

from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_sp import (
    gqa_attention_decode_long_ctx_cube_sp_pe_sp_kernel as _case6_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead import (
    gqa_attention_decode_long_ctx_cube_sp_pe_tp_dhead_kernel as _case4_kernel,
)
from kernbench.benches.gqa_helpers.long_ctx._gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp import (
    gqa_attention_decode_long_ctx_cube_tp_dhead_pe_sp_kernel as _case5_kernel,
)
from kernbench.benches.gqa_helpers.shared._gqa_panel_helpers import (
    _ccl_cfg, _summarize_op_log,
)
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

# Small smoke-test params (S_kv=2K is enough to exercise tile-loop + AR).
_PARAMS = dict(C=8, P=8, T_q=1, S_kv=2_048,
               d_head=128, h_q=8, h_kv=1)


def _bench_fn_case4(ctx):
    p = _PARAMS
    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
    dp_full = DPPolicy(cube="replicate", pe="column_wise",
                       num_cubes=p["C"], num_pes=p["P"])
    dp_kv = DPPolicy(cube="row_wise", pe="column_wise",
                     num_cubes=p["C"], num_pes=p["P"])
    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
                  dtype="f16", dp=dp_full, name="q_c4")
    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
                  dtype="f16", dp=dp_kv, name="k_c4")
    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
                  dtype="f16", dp=dp_kv, name="v_c4")
    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
                  dtype="f16", dp=dp_full, name="o_c4")
    ctx.launch("case4_dhead_tp", _case4_kernel,
               q, k, v, o,
               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
               p["d_head"], p["C"], p["P"],
               _auto_dim_remap=False)


def _bench_fn_case5(ctx):
    p = _PARAMS
    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
    dp_q = DPPolicy(cube="column_wise", pe="replicate",
                    num_cubes=p["C"], num_pes=p["P"])
    dp_kv = DPPolicy(cube="column_wise", pe="row_wise",
                     num_cubes=p["C"], num_pes=p["P"])
    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
                  dtype="f16", dp=dp_q, name="q_c5")
    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
                  dtype="f16", dp=dp_kv, name="k_c5")
    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
                  dtype="f16", dp=dp_kv, name="v_c5")
    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
                  dtype="f16", dp=dp_q, name="o_c5")
    ctx.launch("case5_dhead_tp_inter", _case5_kernel,
               q, k, v, o,
               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
               p["d_head"], p["C"], p["P"],
               _auto_dim_remap=False)


def _bench_fn_case6(ctx):
    p = _PARAMS
    configure_sfr_intercube_multisip(ctx.engine, ctx.spec, _ccl_cfg())
    dp_full = DPPolicy(cube="replicate", pe="replicate",
                       num_cubes=p["C"], num_pes=p["P"])
    dp_kv = DPPolicy(cube="row_wise", pe="row_wise",
                     num_cubes=p["C"], num_pes=p["P"])
    q = ctx.zeros((p["T_q"], p["h_q"] * p["d_head"]),
                  dtype="f16", dp=dp_full, name="q_c6")
    k = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
                  dtype="f16", dp=dp_kv, name="k_c6")
    v = ctx.zeros((p["S_kv"], p["h_kv"] * p["d_head"]),
                  dtype="f16", dp=dp_kv, name="v_c6")
    o = ctx.empty((p["T_q"], p["h_q"] * p["d_head"]),
                  dtype="f16", dp=dp_full, name="o_c6")
    ctx.launch("case6_sp_sp", _case6_kernel,
               q, k, v, o,
               p["T_q"], p["S_kv"], p["h_q"], p["h_kv"],
               p["d_head"], p["C"], p["P"],
               _auto_dim_remap=False)


def main() -> int:
    topology = os.environ.get("GQA_1H_TOPOLOGY", "topology.yaml")
    topo = resolve_topology(topology)

    print(f"Smoke params: {_PARAMS}")
    print()

    for label, bench_fn in (
        ("Case 4 (Cube-SP × PE-TP d_head)",   _bench_fn_case4),
        ("Case 5 (Cube-TP d_head × PE-SP)",   _bench_fn_case5),
        ("Case 6 (Cube-SP × PE-SP S_kv)",     _bench_fn_case6),
    ):
        try:
            res = run_bench(
                topology=topo, bench_fn=bench_fn,
                device=resolve_device(None),
                engine_factory=lambda t, d: GraphEngine(
                    getattr(t, "topology_obj", t), enable_data=True,
                ),
            )
        except Exception as e:
            print(f"  {label:<42} FAIL: {type(e).__name__}: {e}")
            continue
        if not res.completion.ok:
            print(f"  {label:<42} ENGINE FAIL: {res.completion}")
            continue
        s = _summarize_op_log(res.engine.op_log)
        lat = (res.engine.op_log[-1].t_end if res.engine.op_log else 0.0)
        print(f"  {label:<42} "
              f"gemm={s['gemm_count']:>4}  "
              f"ipcq={s['ipcq_copy_count']:>4}  "
              f"dma_r={s['dma_read_count']:>4}  "
              f"dma_w={s['dma_write_count']:>3}  "
              f"latency={lat:.1f} ns  "
              f"(n_ops={len(res.engine.op_log)})")
    return 0


if __name__ == "__main__":
    sys.exit(main())