code/models/common/tests/llm_runtime/test_warmup.py

# SPDX-FileCopyrightText: © 2026 Tenstorrent AI ULC
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import inspect
from dataclasses import replace
from types import SimpleNamespace
from typing import Any, Sequence

import pytest
import torch

from models.common.llm_runtime.config import PageTableLayout, TraceConfig, WarmupConfig
from models.common.llm_runtime.decode import DecodeRuntimeConfig
from models.common.llm_runtime.output_reader import OutputReader
from models.common.llm_runtime.prefill.config import PrefillRuntimeConfig
from models.common.llm_runtime.program_compiler import CompiledProgram, OutputSpec, ProgramKey
from models.common.llm_runtime.warmup import (
    CoverageAlias,
    WarmupCoordinator,
    WarmupCoordinatorConfig,
    _resolve_coverage_manifest,
)


class RecordingExecution:
    def __init__(self, events=None):
        self.prefill_calls = []
        self.decode_calls = []
        self.events = events if events is not None else []
        self.fail_decode_call = None
        self.prefill_replays = []

    def compile_prefill(
        self,
        *,
        tokens: torch.Tensor,  # ↓ Core request
        page_table: torch.Tensor,
        prompt_lens: torch.Tensor | None = None,  # ↓ Sequence metadata
        start_pos: torch.Tensor | None = None,
        empty_slots: Sequence[int] | None = None,  # ↓ Lane routing
        sampling_params: Any = None,  # ↓ Sampling
    ) -> None:
        self.events.append("compile_prefill")
        self.prefill_calls.append(
            {
                "tokens": tokens,
                "page_table": page_table,
                "prompt_lens": prompt_lens,
                "start_pos": start_pos,
                "empty_slots": empty_slots,
                "sampling_params": sampling_params,
            }
        )

    def compile_decode(
        self,
        *,
        tokens: torch.Tensor,  # ↓ Core request
        start_pos: torch.Tensor,
        page_table: torch.Tensor,
        sampling_params: Any = None,  # ↓ Sampling
        reset_batch: bool = False,  # ↓ State transition
    ) -> None:
        call = len(self.decode_calls) + 1
        self.events.append("compile_decode")
        if call == self.fail_decode_call:
            self.fail_decode_call = None
            raise RuntimeError("decode compile failed")
        self.decode_calls.append(
            {
                "tokens": tokens,
                "start_pos": start_pos,
                "page_table": page_table,
                "sampling_params": sampling_params,
                "reset_batch": reset_batch,
            }
        )

    def prefill_forward(
        self,
        *,
        tokens: torch.Tensor,  # ↓ Core request
        page_table: torch.Tensor,
        prompt_lens: torch.Tensor | None = None,  # ↓ Sequence metadata
        start_pos: torch.Tensor | None = None,
        empty_slots: Sequence[int] | None = None,  # ↓ Lane routing
        sampling_params: Any = None,  # ↓ Sampling
    ) -> None:
        self.events.append("prefill_replay")
        self.prefill_replays.append(
            {
                "tokens": tokens,
                "page_table": page_table,
                "prompt_lens": prompt_lens,
                "start_pos": start_pos,
                "empty_slots": empty_slots,
                "sampling_params": sampling_params,
            }
        )


class RecordingTraceCompiler:
    def __init__(self, events=None):
        self.calls = 0
        self.events = events if events is not None else []

    def capture_all(self):
        self.events.append("capture")
        self.calls += 1


class Mesh:
    shape = (1, 1)


def make_runtime_configs(
    *,
    sampling=True,
    lane_capacity=4,
    allow_force_argmax=True,
    page_table_layout=None,
    sampling_config=None,
    model=None,
):
    mesh = Mesh()
    sampling_config = sampling_config or SimpleNamespace(
        allow_force_argmax=allow_force_argmax,
        max_top_k=32,
    )
    sampling_config.max_batch_size = lane_capacity
    model = model or SimpleNamespace(
        config=SimpleNamespace(max_batch_size=lane_capacity, mesh_device=mesh, num_devices=1),
        sampling=SimpleNamespace(
            config=sampling_config,
            decode_forward=lambda logits, *, k=None, p=None, temp=None, seeds=None, tt_out_tok=None, enable_log_probs=False: None,
        ),
        vocab_size=128,
    )
    mesh = model.config.mesh_device
    layout = page_table_layout or PageTableLayout(
        block_size=32,
        raw_capacity_width=128,
        prefill_width=192,
        decode_width=128,
    )
    output_reader = OutputReader(mesh)
    return (
        PrefillRuntimeConfig.resolve(
            model=model,
            output_reader=output_reader,
            page_table_layout=layout,
            max_batch_size=lane_capacity,
            max_prefill_chunk_size=128,
            device_sampling_enabled=sampling,
            can_enable_trace=lambda _sequence_length, _batch_size: True,
        ),
        DecodeRuntimeConfig.resolve(
            model=model,
            output_reader=output_reader,
            lane_capacity=lane_capacity,
            page_table_layout=layout,
            device_sampling_enabled=sampling,
        ),
    )


def make_coordinator(
    *,
    trace_mode="all",
    sampling=True,
    warmup_config=None,
    sequence_lengths=(128, 1024),
    lane_capacity=4,
    execution=None,
    trace_compiler=None,
    events=None,
    allow_force_argmax=True,
    page_table_layout=None,
    sampling_config=None,
):
    events = events if events is not None else []
    execution = execution or RecordingExecution(events)
    if trace_compiler is None and trace_mode != "none":
        trace_compiler = RecordingTraceCompiler(events)
    sampling_calls = []
    bound_calls = []

    def ensure_sampling():
        events.append("sampling")
        sampling_calls.append(True)

    def validate_bound(value):
        bound_calls.append(value)

    layout = page_table_layout or PageTableLayout(
        block_size=32,
        raw_capacity_width=128,
        prefill_width=192,
        decode_width=128,
    )
    prefill_config, decode_config = make_runtime_configs(
        sampling=sampling,
        lane_capacity=lane_capacity,
        allow_force_argmax=allow_force_argmax,
        page_table_layout=layout,
        sampling_config=sampling_config,
    )
    execution.prefill = SimpleNamespace(config=prefill_config)
    execution.decode = SimpleNamespace(config=decode_config)
    execution.eager_executor = execution
    execution.trace_compiler = trace_compiler

    coordinator = WarmupCoordinator(
        config=WarmupCoordinatorConfig.resolve(
            warmup=warmup_config or WarmupConfig(),
            trace=TraceConfig(trace_mode),
            prefill=prefill_config,
            decode=decode_config,
            prefill_sequence_lengths=sequence_lengths,
        ),
        execution=execution,
        ensure_sampling_buffers=ensure_sampling,
        validate_bound_cache=validate_bound,
    )
    return coordinator, execution, trace_compiler, sampling_calls, bound_calls, events


@pytest.mark.parametrize(
    ("method_name", "parameter_names"),
    [
        ("warmup_prefill", ("self", "kv_cache", "can_sample_on_device", "enable_trace")),
        (
            "warmup_decode",
            ("self", "kv_cache", "max_batch_size", "num_blocks", "can_sample_on_device", "enable_trace"),
        ),
    ],
)
def test_warmup_signatures_match_registered_plugin_contract(method_name, parameter_names):
    parameters = inspect.signature(getattr(WarmupCoordinator, method_name)).parameters

    assert tuple(parameters) == parameter_names
    assert parameters["self"].kind is inspect.Parameter.POSITIONAL_OR_KEYWORD
    for name in parameter_names[1:]:
        assert parameters[name].kind is inspect.Parameter.KEYWORD_ONLY
        assert parameters[name].default is inspect.Parameter.empty


def test_registered_plugin_warmup_calls_validate_cache_without_forwarding_it():
    cache = object()
    coordinator, execution, _, _, bound_calls, _ = make_coordinator(
        trace_mode="none",
        sampling=False,
        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
        sequence_lengths=(128,),
        lane_capacity=1,
    )
    prefill_kwargs = {
        "kv_cache": cache,
        "can_sample_on_device": False,
    }
    decode_kwargs = {
        "kv_cache": cache,
        "max_batch_size": 1,
        "num_blocks": 8,
        "can_sample_on_device": False,
    }

    coordinator.warmup_prefill(enable_trace=False, **prefill_kwargs)
    coordinator.warmup_decode(enable_trace=False, **decode_kwargs)

    assert coordinator.coverage_manifest is None
    assert bound_calls == [cache, cache]
    assert all(
        tuple(call) == ("tokens", "page_table", "prompt_lens", "start_pos", "empty_slots", "sampling_params")
        for call in execution.prefill_calls
    )
    assert all(
        tuple(call) == ("tokens", "start_pos", "page_table", "sampling_params", "reset_batch")
        for call in execution.decode_calls
    )


@pytest.mark.parametrize(
    ("method_name", "plugin_kwargs", "unexpected_name"),
    [
        (
            "warmup_prefill",
            {"kv_cache": "cache", "can_sample_on_device": False, "enable_trace": False},
            "greedy_only",
        ),
        (
            "warmup_decode",
            {
                "kv_cache": "cache",
                "max_batch_size": 1,
                "num_blocks": 8,
                "can_sample_on_device": False,
                "enable_trace": False,
            },
            "read_from_device",
        ),
    ],
)
def test_warmup_contract_rejects_unregistered_plugin_keywords(
    method_name,
    plugin_kwargs,
    unexpected_name,
    expect_error,
):
    coordinator, *_ = make_coordinator(trace_mode="none", sampling=False, lane_capacity=1)
    plugin_kwargs[unexpected_name] = False

    with expect_error(TypeError, unexpected_name):
        getattr(coordinator, method_name)(**plugin_kwargs)


def test_configured_prefill_lengths_override_model_supported_defaults():
    coordinator, execution, *_ = make_coordinator(
        warmup_config=WarmupConfig(prefill_seq_lens=(1024,), prefill_batch_sizes=(1,)),
        sequence_lengths=(128,),
        sampling=False,
    )

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)

    regular_lengths = [int(call["tokens"].shape[-1]) for call in execution.prefill_calls if call["start_pos"] is None]
    assert regular_lengths == [1024]


@pytest.mark.parametrize(
    ("sequence_lengths", "message"),
    [
        ((), "non-empty tuple"),
        ((True,), "positive integers"),
        ((128, 128), "unique"),
    ],
)
def test_model_supported_prefill_lengths_are_validated_once(sequence_lengths, message, expect_error):
    with expect_error(ValueError, message):
        make_coordinator(sequence_lengths=sequence_lengths)


def test_sampler_argmax_capability_is_resolved_once():
    class SamplingConfig:
        reads = 0
        max_top_k = 32

        @property
        def allow_force_argmax(self):
            self.reads += 1
            return True

    sampling_config = SamplingConfig()
    coordinator, *_ = make_coordinator(sampling_config=sampling_config)
    resolved_reads = sampling_config.reads

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=True)
    coordinator.warmup_decode(
        kv_cache="cache",
        enable_trace=False,
        max_batch_size=4,
        num_blocks=8,
        can_sample_on_device=True,
    )

    assert sampling_config.reads == resolved_reads


def test_page_table_layout_can_be_reconfigured_only_before_use(expect_error):
    coordinator, execution, *_ = make_coordinator(
        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
        sampling=False,
    )
    final_layout = PageTableLayout(
        block_size=32,
        raw_capacity_width=4,
        prefill_width=64,
        decode_width=8,
    )

    coordinator.configure_page_table_layout(final_layout)
    assert coordinator.config.page_table_layout is final_layout
    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)

    assert all(call["start_pos"] is None for call in execution.prefill_calls)
    with expect_error(RuntimeError, "configuration is sealed"):
        coordinator.configure_page_table_layout(final_layout)


def test_explicit_configuration_seal_precedes_physical_kv_allocation(expect_error):
    coordinator, *_ = make_coordinator()
    coordinator.seal_configuration()

    with expect_error(RuntimeError, "configuration is sealed"):
        coordinator.configure_page_table_layout(
            PageTableLayout(
                block_size=32,
                raw_capacity_width=64,
                prefill_width=128,
                decode_width=64,
            )
        )


def test_page_table_layout_reconfiguration_requires_immutable_layout(expect_error):
    coordinator, *_ = make_coordinator()

    with expect_error(TypeError, "PageTableLayout"):
        coordinator.configure_page_table_layout(SimpleNamespace(block_size=32))


def test_resolved_config_is_frozen_and_owns_both_coverage_plans(expect_error):
    coordinator, *_ = make_coordinator(
        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
        lane_capacity=2,
    )

    assert coordinator.config.eager_plan.decode == (coordinator.config.eager_plan.decode[0],)
    assert [case.sampling_path for case in coordinator.config.sampled_plan.decode] == ["logits", "argmax"]
    with expect_error(AttributeError, "cannot assign"):
        coordinator.config.lane_batch_size = 4


def test_direct_config_construction_rejects_inconsistent_derived_plan(expect_error):
    coordinator, *_ = make_coordinator()

    with expect_error(ValueError, "plans must match"):
        replace(coordinator.config, eager_plan=coordinator.config.sampled_plan)


@pytest.mark.parametrize(
    ("mismatch", "message"),
    [
        ("model", "share one model"),
        ("layout", "share one page-table layout"),
        ("lane", "share one lane capacity"),
        ("sampling", "share device-sampling policy"),
        ("argmax", "share force-argmax capability"),
        ("raw_ceiling", "share one page-table layout ceiling"),
        ("decode_ceiling", "share one page-table layout ceiling"),
    ],
)
def test_resolution_rejects_inconsistent_runtime_configs(mismatch, message, expect_error):
    prefill, decode = make_runtime_configs()
    if mismatch == "model":
        _, decode = make_runtime_configs()
    elif mismatch == "layout":
        decode = decode.with_page_table_layout(PageTableLayout(32, 64, 128, 64))
    elif mismatch == "lane":
        decode = DecodeRuntimeConfig.resolve(
            model=prefill.model,
            output_reader=prefill.output_reader,
            lane_capacity=2,
            page_table_layout=prefill.page_table_layout,
            device_sampling_enabled=True,
        )
    elif mismatch == "sampling":
        decode = DecodeRuntimeConfig.resolve(
            model=prefill.model,
            output_reader=prefill.output_reader,
            lane_capacity=prefill.max_batch_size,
            page_table_layout=prefill.page_table_layout,
            device_sampling_enabled=False,
        )
    elif mismatch == "argmax":
        prefill.model.sampling.config.allow_force_argmax = False
        decode = DecodeRuntimeConfig.resolve(
            model=prefill.model,
            output_reader=prefill.output_reader,
            lane_capacity=prefill.max_batch_size,
            page_table_layout=prefill.page_table_layout,
            device_sampling_enabled=True,
        )
    elif mismatch == "raw_ceiling":
        ceiling = prefill.page_table_layout_ceiling
        prefill = replace(
            prefill,
            page_table_layout_ceiling=replace(
                ceiling,
                raw_capacity_width=ceiling.raw_capacity_width + 1,
                decode_width=ceiling.decode_width + 8,
            ),
        )
    else:
        ceiling = prefill.page_table_layout_ceiling
        prefill = replace(
            prefill,
            page_table_layout_ceiling=replace(ceiling, decode_width=ceiling.decode_width + 8),
        )

    with expect_error(ValueError, message):
        WarmupCoordinatorConfig.resolve(
            warmup=WarmupConfig(),
            trace=TraceConfig("all"),
            prefill=prefill,
            decode=decode,
            prefill_sequence_lengths=(128,),
        )


def test_constructor_rejects_execution_disagreement_with_resolved_config(expect_error):
    coordinator, execution, *_ = make_coordinator()
    config = coordinator.config
    execution.prefill.config = replace(
        execution.prefill.config,
        page_table_layout=PageTableLayout(32, 64, 128, 64),
    )

    with expect_error(ValueError, "warmup config page-table layout"):
        WarmupCoordinator(
            config=config,
            execution=execution,
            ensure_sampling_buffers=lambda: None,
            validate_bound_cache=lambda _: None,
        )


def test_runtime_does_not_copy_static_config_fields():
    coordinator, *_ = make_coordinator()

    assert {
        "page_table_layout",
        "prefill_sequence_lengths",
        "lane_batch_size",
        "device_sampling_enabled",
        "allow_force_argmax",
        "prime_q128_tile_ends",
        "prefill_trace_enabled",
        "decode_trace_enabled",
        "eager_plan",
        "sampled_plan",
    }.isdisjoint(vars(coordinator))


def test_layout_replacement_is_immutable_bounded_and_rebuilds_coverage(expect_error):
    coordinator, *_ = make_coordinator(
        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
        page_table_layout=PageTableLayout(32, 128, 192, 128),
    )
    original = coordinator.config
    replacement = PageTableLayout(32, 4, 64, 8)

    coordinator.configure_page_table_layout(replacement)

    assert coordinator.config is not original
    assert coordinator.config.page_table_layout_ceiling is original.page_table_layout
    assert original.page_table_layout.raw_capacity_width == 128
    assert not any(case.cached_tokens for case in coordinator.config.eager_plan.prefill)
    with expect_error(ValueError, "cannot change block_size"):
        original.with_page_table_layout(PageTableLayout(16, 4, 64, 8))
    with expect_error(ValueError, "capacity ceiling"):
        original.with_page_table_layout(PageTableLayout(32, 129, 192, 136))
    with expect_error(ValueError, "canonical geometry"):
        original.with_page_table_layout(PageTableLayout(32, 128, 200, 128))


def test_q128_batches_are_capped_by_lane_and_non128_is_batch_one():
    config = WarmupConfig(prefill_batch_sizes=(1, 2, 4, 8, 16, 32))
    coordinator, execution, *_ = make_coordinator(warmup_config=config, lane_capacity=8)

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)

    regular_q128 = [
        int(call["tokens"].shape[0])
        for call in execution.prefill_calls
        if int(call["tokens"].shape[-1]) == 128 and call["start_pos"] is None
    ]
    regular_q1024 = [
        int(call["tokens"].shape[0])
        for call in execution.prefill_calls
        if int(call["tokens"].shape[-1]) == 1024 and call["start_pos"] is None
    ]
    assert regular_q128 == [1, 2, 4, 8]
    assert regular_q1024 == [1]


def test_sampling_paths_include_forced_prefill_topk_and_opt_in_true_topk_decode():
    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,), include_decode_top_k=True)
    coordinator, execution, *_ = make_coordinator(warmup_config=config, sequence_lengths=(128,), lane_capacity=2)

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=True)
    coordinator.warmup_decode(
        kv_cache="cache",
        enable_trace=False,
        max_batch_size=2,
        num_blocks=8,
        can_sample_on_device=True,
    )

    assert execution.prefill_calls[0]["sampling_params"] is None
    assert execution.prefill_calls[1]["sampling_params"].top_k.tolist() == [32]
    assert execution.decode_calls[0]["sampling_params"] is None
    assert execution.decode_calls[1]["sampling_params"].top_k.tolist() == [1, 1]
    assert execution.decode_calls[2]["sampling_params"].top_k.tolist() == [32, 32]
    # Preserve the established true-top-k recipe, not merely a top-k label.
    assert execution.decode_calls[2]["sampling_params"].top_p.tolist() == pytest.approx([0.08, 0.08])


def test_q128_single_topk_primes_all_tile_ends():
    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
    coordinator, execution, *_ = make_coordinator(
        warmup_config=config,
        sequence_lengths=(128,),
        lane_capacity=32,
    )
    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=True)

    topk_calls = [
        call
        for call in execution.prefill_calls
        if call["sampling_params"] is not None
        and float(call["sampling_params"].temperature[0]) == 1.0
        and call["start_pos"] is None
    ]
    assert [int(call["prompt_lens"][0]) for call in topk_calls] == [32, 64, 96, 128]
    assert [int(call["tokens"].shape[-1]) for call in topk_calls] == [32, 64, 96, 128]
    argmax_calls = [
        call
        for call in execution.prefill_calls
        if call["sampling_params"] is not None
        and float(call["sampling_params"].temperature[0]) == 0.0
        and call["start_pos"] is None
    ]
    assert [int(call["prompt_lens"][0]) for call in argmax_calls] == [32, 64, 96, 128]


def test_decode_warmup_uses_topk_as_the_platform_greedy_path_when_argmax_is_disabled():
    coordinator, execution, *_ = make_coordinator(
        sequence_lengths=(128,),
        lane_capacity=2,
        allow_force_argmax=False,
    )

    coordinator.warmup_decode(
        kv_cache="cache",
        enable_trace=False,
        max_batch_size=2,
        num_blocks=8,
        can_sample_on_device=True,
    )

    assert execution.decode_calls[0]["sampling_params"] is None
    assert execution.decode_calls[1]["sampling_params"].top_k.tolist() == [32, 32]


def test_eager_and_trace_coverage_are_separately_idempotent():
    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
    coordinator, execution, trace_compiler, *_ = make_coordinator(
        warmup_config=config, sequence_lengths=(128,), lane_capacity=1
    )

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)
    eager_calls = len(execution.prefill_calls)
    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)
    assert len(execution.prefill_calls) == eager_calls

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=True)
    trace_calls = len(execution.prefill_calls)
    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=True)
    assert len(execution.prefill_calls) == trace_calls
    assert trace_compiler.calls == 0


def test_trace_warmup_routes_cached_prefill_through_traced_execution_target():
    coordinator, eager, *_ = make_coordinator(
        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
        sequence_lengths=(128,),
        lane_capacity=1,
        sampling=False,
    )
    traced = RecordingExecution()
    coordinator.execution = traced

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)

    assert not eager.prefill_calls
    assert any(call["start_pos"] is None for call in traced.prefill_calls)
    assert any(call["start_pos"] is not None for call in traced.prefill_calls)


def test_coverage_manifest_uses_compiler_registries_and_deduplicates_trace_identities():
    eager_program = CompiledProgram(ProgramKey("0" * 64), "eager", OutputSpec((1,), torch.float32))
    first_traced = CompiledProgram(ProgramKey("1" * 64), "traced-a", OutputSpec((1,), torch.float32))
    second_traced = CompiledProgram(ProgramKey("2" * 64), "traced-b", OutputSpec((1,), torch.float32))
    shared_trace_key = ProgramKey("a" * 64)
    program_compiler = SimpleNamespace(compiled_programs=(eager_program, first_traced, second_traced))
    eager = SimpleNamespace(program_compiler=program_compiler)
    trace_compiler = SimpleNamespace(
        trace_key_for_program=lambda key: None if key == eager_program.key else shared_trace_key,
        get=lambda key: SimpleNamespace(signature="shared-trace") if key == shared_trace_key else None,
    )

    manifest = _resolve_coverage_manifest(eager, trace_compiler)

    assert manifest.eager_program_signatures == ("eager",)
    assert manifest.traced_source_program_signatures == ("traced-a", "traced-b")
    assert manifest.trace_signatures == ("shared-trace",)
    assert manifest.aliases == (
        CoverageAlias("traced-a", "shared-trace"),
        CoverageAlias("traced-b", "shared-trace"),
    )


def test_coverage_manifest_rejects_any_missing_required_trace_alias(expect_error):
    first_traced = CompiledProgram(ProgramKey("1" * 64), "traced-a", OutputSpec((1,), torch.float32))
    second_traced = CompiledProgram(ProgramKey("2" * 64), "traced-b", OutputSpec((1,), torch.float32))
    trace_key = ProgramKey("a" * 64)
    eager = SimpleNamespace(program_compiler=SimpleNamespace(compiled_programs=(first_traced, second_traced)))
    trace_compiler = SimpleNamespace(
        trace_key_for_program=lambda key: trace_key if key == first_traced.key else None,
        get=lambda key: SimpleNamespace(signature="trace-a") if key == trace_key else None,
    )

    with expect_error(RuntimeError, "required trace alias"):
        _resolve_coverage_manifest(
            eager,
            trace_compiler,
            required_program_keys={first_traced.key, second_traced.key},
            required_trace_program_keys={first_traced.key, second_traced.key},
        )


def test_activation_validates_every_program_returned_by_trace_warmup(expect_error):
    coordinator, execution, trace_compiler, *_ = make_coordinator(
        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
        sequence_lengths=(128,),
        lane_capacity=1,
        sampling=False,
    )
    programs = tuple(
        CompiledProgram(ProgramKey(str(index) * 64), f"program-{index}", OutputSpec((1,), torch.float32))
        for index in range(1, 4)
    )
    execution.program_compiler = SimpleNamespace(compiled_programs=programs)
    prefill_programs = iter(programs[:2])
    execution.compile_prefill = lambda **_kwargs: (next(prefill_programs),)
    execution.compile_decode = lambda **_kwargs: programs[2]
    trace_keys = {
        programs[0].key: ProgramKey("a" * 64),
        programs[2].key: ProgramKey("c" * 64),
    }
    trace_compiler.trace_key_for_program = trace_keys.get
    trace_compiler.get = lambda key: SimpleNamespace(signature=f"trace-{key.digest[0]}")

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
    with expect_error(RuntimeError, programs[1].key.digest):
        coordinator.warmup_decode(
            kv_cache="cache",
            enable_trace=True,
            max_batch_size=1,
            num_blocks=8,
            can_sample_on_device=False,
        )

    assert trace_compiler.calls == 0


@pytest.mark.parametrize("order", [("prefill", "decode"), ("decode", "prefill")])
def test_prefill_decode_order_is_independent_and_capture_waits_for_both(order):
    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
    coordinator, execution, trace_compiler, *_ = make_coordinator(
        warmup_config=config, sequence_lengths=(128,), lane_capacity=1
    )

    def run(operation):
        if operation == "prefill":
            coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=True)
        else:
            coordinator.warmup_decode(
                kv_cache="cache",
                enable_trace=True,
                max_batch_size=1,
                num_blocks=8,
                can_sample_on_device=True,
            )

    run(order[0])
    assert trace_compiler.calls == 0
    run(order[1])
    assert trace_compiler.calls == 1
    run(order[0])
    run(order[1])
    assert trace_compiler.calls == 1


@pytest.mark.parametrize("order", [("prefill", "decode"), ("decode", "prefill")])
def test_capture_uses_phase_specific_sampling_decisions(order):
    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
    coordinator, execution, trace_compiler, *_ = make_coordinator(
        trace_mode="all",
        warmup_config=config,
        sequence_lengths=(128,),
        lane_capacity=1,
    )

    def run(operation):
        if operation == "prefill":
            coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
        else:
            coordinator.warmup_decode(
                kv_cache="cache",
                enable_trace=True,
                max_batch_size=1,
                num_blocks=8,
                can_sample_on_device=True,
            )

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)
    coordinator.warmup_decode(
        kv_cache="cache",
        enable_trace=False,
        max_batch_size=1,
        num_blocks=8,
        can_sample_on_device=True,
    )
    with coordinator.defer_capture():
        run(order[0])
        assert trace_compiler.calls == 0
        run(order[1])
        assert coordinator.capture_pending
        coordinator.activate_pending_capture()

    assert trace_compiler.calls == 1
    assert coordinator.already_warmed_up_prefill
    assert all(call["sampling_params"] is None for call in execution.prefill_calls)
    assert any(call["sampling_params"] is not None for call in execution.decode_calls)
    assert not execution.prefill_replays


def test_capture_deferral_stages_complete_registration_until_explicit_activation():
    coordinator, _, trace_compiler, *_ = make_coordinator(
        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
        sequence_lengths=(128,),
        lane_capacity=1,
        sampling=False,
    )

    with coordinator.defer_capture():
        coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
        coordinator.warmup_decode(
            kv_cache="cache",
            enable_trace=True,
            max_batch_size=1,
            num_blocks=8,
            can_sample_on_device=False,
        )
        assert coordinator.capture_pending
        assert not coordinator.trace_activated
        assert trace_compiler.calls == 0
        coordinator.activate_pending_capture()
        assert coordinator.trace_activated
        assert trace_compiler.calls == 1

    assert not coordinator.capture_pending


def test_capture_deferral_exception_discards_pending_activation(expect_error):
    coordinator, _, trace_compiler, *_ = make_coordinator(
        warmup_config=WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,)),
        sequence_lengths=(128,),
        lane_capacity=1,
        sampling=False,
    )

    with expect_error(RuntimeError, "staging failed"):
        with coordinator.defer_capture():
            coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
            coordinator.warmup_decode(
                kv_cache="cache",
                enable_trace=True,
                max_batch_size=1,
                num_blocks=8,
                can_sample_on_device=False,
            )
            assert coordinator.capture_pending
            raise RuntimeError("staging failed")

    assert not coordinator.capture_pending
    assert not coordinator.trace_activated
    assert trace_compiler.calls == 0


def test_static_all_can_capture_decode_only_runtime_trace():
    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
    coordinator, execution, trace_compiler, *_ = make_coordinator(
        trace_mode="all",
        warmup_config=config,
        sequence_lengths=(128,),
        lane_capacity=1,
    )

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=True)
    assert coordinator.already_warmed_up_prefill
    assert trace_compiler.calls == 0

    coordinator.warmup_decode(
        kv_cache="cache",
        enable_trace=True,
        max_batch_size=1,
        num_blocks=8,
        can_sample_on_device=True,
    )

    assert trace_compiler.calls == 1
    assert not execution.prefill_replays


def test_two_phase_static_all_waits_for_phase_two_decode_before_capture():
    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
    coordinator, _, trace_compiler, *_ = make_coordinator(
        trace_mode="all",
        warmup_config=config,
        sequence_lengths=(128,),
        lane_capacity=1,
    )

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=True)
    coordinator.warmup_decode(
        kv_cache="cache",
        enable_trace=False,
        max_batch_size=1,
        num_blocks=8,
        can_sample_on_device=True,
    )
    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=True)
    assert trace_compiler.calls == 0

    coordinator.warmup_decode(
        kv_cache="cache",
        enable_trace=True,
        max_batch_size=1,
        num_blocks=8,
        can_sample_on_device=True,
    )

    assert trace_compiler.calls == 1


def test_sampling_buffers_are_materialized_before_first_compile_and_capture():
    events = []
    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
    coordinator, _, _, _, _, events = make_coordinator(
        warmup_config=config,
        sequence_lengths=(128,),
        lane_capacity=1,
        events=events,
    )

    coordinator.warmup_decode(
        kv_cache="cache",
        enable_trace=True,
        max_batch_size=1,
        num_blocks=8,
        can_sample_on_device=True,
    )
    coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=True)

    assert events.index("sampling") < events.index("compile_decode")
    assert events.index("sampling") < events.index("compile_prefill")
    assert max(index for index, event in enumerate(events) if event.startswith("compile_")) < events.index("capture")


def test_failed_case_is_not_marked_complete_and_retry_skips_completed_case(expect_error):
    execution = RecordingExecution()
    execution.fail_decode_call = 2
    config = WarmupConfig(prefill_seq_lens=(128,), prefill_batch_sizes=(1,))
    coordinator, execution, *_ = make_coordinator(
        trace_mode="none",
        sampling=True,
        warmup_config=config,
        sequence_lengths=(128,),
        lane_capacity=1,
        execution=execution,
    )

    with expect_error(RuntimeError, "decode compile failed"):
        coordinator.warmup_decode(
            kv_cache="cache",
            enable_trace=False,
            max_batch_size=1,
            num_blocks=8,
            can_sample_on_device=True,
        )
    assert len(execution.decode_calls) == 1

    coordinator.warmup_decode(
        kv_cache="cache",
        enable_trace=False,
        max_batch_size=1,
        num_blocks=8,
        can_sample_on_device=True,
    )
    assert len(execution.decode_calls) == 2


def test_dynamic_hints_cannot_expand_static_trace_or_sampling_ceilings(expect_error):
    coordinator, *_ = make_coordinator(trace_mode="decode_only", sampling=False)

    with expect_error(ValueError, "prefill trace warmup exceeds"):
        coordinator.warmup_prefill(kv_cache="cache", enable_trace=True, can_sample_on_device=False)
    with expect_error(ValueError, "statically disabled"):
        coordinator.warmup_decode(
            kv_cache="cache",
            enable_trace=False,
            max_batch_size=4,
            num_blocks=8,
            can_sample_on_device=True,
        )

    coordinator.warmup_prefill(kv_cache="cache", enable_trace=False, can_sample_on_device=False)