kernbench2/tests/analytical_visualization/test_auto_explore.py

"""Interface + smoke tests for auto_explore.

Covers:
  - enumerate_configs yields valid TopologyConfigs with expected pruning
  - score_config returns a ConfigScore with all fields populated
  - pareto_frontier is a subset of feasible and is non-empty for a
    reasonable model+workload
  - run_auto_explore returns a coherent AutoExploreResult (all_scores
    sorted by latency, pareto ⊆ feasible ⊆ all_scores)
"""
from __future__ import annotations

import pytest

from tests.analytical_visualization.auto_explore import (
    ConfigScore,
    enumerate_configs,
    pareto_frontier,
    run_auto_explore,
    score_config,
)
from tests.analytical_visualization.autosuggest import auto_suggest
from tests.analytical_visualization.model_config import (
    FullConfig, MachineParams, TopologyConfig,
)
from tests.analytical_visualization.model_presets import PRESETS


# ── Enumeration ─────────────────────────────────────────────────────


def test_enumerate_yields_valid_configs():
    """Enumerator produces TopologyConfigs with all 9 knobs set."""
    model = PRESETS["Llama 3.1 70B"].model
    it = enumerate_configs(model, s_kv=8192, mode="decode")
    first = next(it)
    for f in ("cp", "tp", "pp", "dp", "kv_shard_mode",
              "ffn_shard_scope", "tp_placement", "cp_placement",
              "cp_ring_variant"):
        assert getattr(first, f) is not None, f"knob {f} not set"
    assert first.s_kv == 8192
    assert first.mode == "decode"


def test_enumerate_prunes_pp_beyond_layers():
    """PP > model.layers is skipped by the enumerator."""
    model = PRESETS["Llama 3.1 70B"].model  # layers=80
    for cfg in enumerate_configs(model, s_kv=8192, mode="decode"):
        assert cfg.pp <= model.layers


def test_enumerate_prunes_ffn_dp_when_dp_is_1():
    """ffn_shard_scope containing 'DP' is skipped when dp=1 (redundant)."""
    model = PRESETS["Llama 3.1 70B"].model
    for cfg in enumerate_configs(model, s_kv=8192, mode="decode"):
        if cfg.dp == 1:
            assert "DP" not in cfg.ffn_shard_scope


# ── Scoring ─────────────────────────────────────────────────────────


def test_score_returns_populated_config_score():
    """score_config returns a fully-populated ConfigScore."""
    model = PRESETS["Llama 3.1 70B"].model
    topo = next(enumerate_configs(model, s_kv=8192, mode="decode"))
    machine = MachineParams()
    cfg = FullConfig(model=model, topo=topo, machine=machine)
    score = score_config(cfg)
    assert isinstance(score, ConfigScore)
    assert score.total_latency_ns > 0
    assert score.pes_used == topo.total_pes
    assert 0.0 <= score.efficiency_score <= 1.0
    assert score.hbm_utilization >= 0.0


# ── Pareto ──────────────────────────────────────────────────────────


def test_pareto_subset_of_feasible():
    """Pareto scores are always feasible (memory + placement)."""
    model = PRESETS["Llama 3.1 70B"].model
    machine = MachineParams()
    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
    for p in res.pareto_scores:
        assert p.fits_memory
        assert p.placement_valid


def test_pareto_non_empty_when_feasible_configs_exist():
    """When at least one config fits, Pareto must have ≥ 1 entry."""
    model = PRESETS["Llama 3.1 70B"].model
    machine = MachineParams()
    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
    assert res.total_feasible > 0
    assert len(res.pareto_scores) > 0


def test_pareto_frontier_non_dominated():
    """No Pareto entry is dominated by another."""
    from tests.analytical_visualization.auto_explore import _dominates
    model = PRESETS["Llama 3.1 70B"].model
    machine = MachineParams()
    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
    for i, a in enumerate(res.pareto_scores):
        for j, b in enumerate(res.pareto_scores):
            if i == j:
                continue
            assert not _dominates(b, a), (
                f"Pareto entry {i} is dominated by entry {j}"
            )


# ── End-to-end sanity ───────────────────────────────────────────────


def test_run_auto_explore_shape():
    """all_scores sorted asc by latency; pareto ⊆ feasible ⊆ all_scores."""
    model = PRESETS["Llama 3.1 70B"].model
    machine = MachineParams()
    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")

    assert res.total_enumerated == len(res.all_scores)
    assert res.total_feasible == sum(
        1 for s in res.all_scores if s.fits_memory and s.placement_valid
    )
    assert len(res.pareto_scores) <= res.total_feasible

    latencies = [s.total_latency_ns for s in res.all_scores]
    assert latencies == sorted(latencies)


def test_pp_does_not_reduce_single_request_latency():
    """A single request traverses all layers regardless of PP.
    Latency-optimal Pareto configs should NOT prefer PP>1 for decode."""
    model = PRESETS["Llama 3.1 70B"].model
    machine = MachineParams()
    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
    # Sort Pareto by latency; the fastest should not need PP>1 to win.
    fastest = res.pareto_scores[0]
    assert fastest.pp == 1, (
        f"expected PP=1 for the fastest decode config; got PP={fastest.pp}"
    )


# ── Parallelism sensitivity ─────────────────────────────────────────


def test_parallelism_sensitivity_has_all_five_knobs():
    """compute_parallelism_sensitivity returns one row per knob."""
    from tests.analytical_visualization.auto_explore import (
        compute_parallelism_sensitivity,
    )
    model = PRESETS["Llama 3.1 70B"].model
    machine = MachineParams()
    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
    baseline = res.pareto_scores[0]
    rows = compute_parallelism_sensitivity(
        baseline, model, machine, s_kv=8192, mode="decode",
    )
    knobs = {r.knob for r in rows}
    assert knobs == {"cp", "tp", "pp", "dp", "ep"}


def test_attention_only_latency_is_lower_than_full():
    """include_ffn=False must produce lower latency than include_ffn=True
    for the same model+workload (FFN cost is dropped)."""
    model = PRESETS["Llama 3.1 70B"].model
    machine = MachineParams()
    r_full = run_auto_explore(model, machine, s_kv=8192, mode="decode",
                                include_ffn=True)
    r_attn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
                                include_ffn=False)
    best_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns)
    best_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns)
    assert best_attn.total_latency_ns < best_full.total_latency_ns, (
        f"attn-only ({best_attn.latency_us:.2f} us) should be less than "
        f"full ({best_full.latency_us:.2f} us)"
    )


def test_attention_only_pareto_non_empty():
    """The attention-only sweep must still produce a non-empty Pareto set."""
    model = PRESETS["Llama 3.1 70B"].model
    machine = MachineParams()
    res = run_auto_explore(model, machine, s_kv=8192, mode="decode",
                             include_attention=True, include_ffn=False)
    assert res.total_feasible > 0
    assert len(res.pareto_scores) > 0


def test_ffn_only_latency_less_than_full_and_different_from_attn():
    """FFN-only latency must be > 0, < full-model latency, and != attn-only."""
    model = PRESETS["Llama 3.1 70B"].model
    machine = MachineParams()
    r_full = run_auto_explore(model, machine, s_kv=8192, mode="decode",
                                include_attention=True, include_ffn=True)
    r_attn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
                                include_attention=True, include_ffn=False)
    r_ffn = run_auto_explore(model, machine, s_kv=8192, mode="decode",
                               include_attention=False, include_ffn=True)
    b_full = min(r_full.pareto_scores, key=lambda s: s.total_latency_ns)
    b_attn = min(r_attn.pareto_scores, key=lambda s: s.total_latency_ns)
    b_ffn = min(r_ffn.pareto_scores, key=lambda s: s.total_latency_ns)
    assert b_ffn.total_latency_ns > 0
    assert b_ffn.total_latency_ns < b_full.total_latency_ns
    assert abs(b_ffn.total_latency_ns - b_attn.total_latency_ns) > 1.0, (
        "FFN-only and attention-only latencies should not be identical"
    )


def test_auto_suggest_cubes_match_default_topology():
    """A TopologyConfig built from auto_suggest's returned (cp, tp, pp,
    cp_placement) with all other fields left at their dataclass defaults
    must produce the same cubes_used the Suggestion reports. Guards the
    sidebar apply: any session_state key auto_suggest doesn't explore
    (tp_placement, kv_shard_mode, cp_ring_variant, ep, ffn_shard_scope)
    must be reset to the default before rerun, else the applied topology
    won't match what the memory-min search actually scored.
    """
    machine = MachineParams()
    matrix = [
        ("Llama 3 70B", 8192),
        ("Llama 3 70B", 32768),
        ("Llama 3 8B",  131072),
        ("Qwen 3 32B",  32768),
    ]
    for name, skv in matrix:
        m = PRESETS[name].model
        for ia, ff in [(True, False), (False, True), (True, True)]:
            sug = auto_suggest(m, machine, s_kv=skv, mode="decode",
                               include_attention=ia, include_ffn=ff)
            topo = TopologyConfig(
                cp=sug.cp, tp=sug.tp, pp=sug.pp,
                s_kv=skv, mode="decode", b=1,
                cp_placement=sug.cp_placement,
            )
            assert topo.cubes_used == sug.cubes_used, (
                f"{name} skv={skv} scope=(ia={ia},ff={ff}): "
                f"topo.cubes_used={topo.cubes_used} != "
                f"sug.cubes_used={sug.cubes_used}"
            )


def test_parallelism_sensitivity_includes_baseline_value():
    """Each row's values contain the baseline_value."""
    from tests.analytical_visualization.auto_explore import (
        compute_parallelism_sensitivity,
    )
    model = PRESETS["Llama 3.1 70B"].model
    machine = MachineParams()
    res = run_auto_explore(model, machine, s_kv=8192, mode="decode")
    baseline = res.pareto_scores[0]
    rows = compute_parallelism_sensitivity(
        baseline, model, machine, s_kv=8192, mode="decode",
    )
    for r in rows:
        # Baseline value must be sweepable OR mapped to something in values.
        assert r.baseline_value in r.values or r.baseline_value == 1, (
            f"knob {r.knob}: baseline_value={r.baseline_value} not in {r.values}"
        )