# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import enum
import itertools
import json
import math
import os
import pickle
import re
import statistics
import sys
import tempfile
from typing import Any, List

import cloudpickle
import pytest
import torch
from mpi4py import MPI

import tensorrt_llm
import tensorrt_llm._torch.autotuner as autotuner
from tensorrt_llm._torch.autotuner import (AutoTuner, DistributedTuningStrategy,
                                           DynamicDim, DynamicTensorSpec,
                                           FakeTensor, NvMMHConfig,
                                           OptimizationProfile, StaticDim,
                                           TunableRunner, TuningConfig,
                                           autotune)
from tensorrt_llm._torch.distributed import Distributed
from tensorrt_llm._torch.utils import (deep_gemm_jit_warmup_buckets,
                                       get_power_of_2_num_tokens_buckets,
                                       next_positive_power_of_2)
from tensorrt_llm.bindings.internal.runtime import delay_kernel
from tensorrt_llm.logger import logger
from tensorrt_llm.mapping import Mapping

cloudpickle.register_pickle_by_value(sys.modules[__name__])
MPI.pickle.__init__(
    cloudpickle.dumps,
    cloudpickle.loads,
    pickle.HIGHEST_PROTOCOL,
)

# needed since we reuse the mpi executor pool, first test running will leak a thread
pytestmark = pytest.mark.threadleak(enabled=False)


def test_multi_dynamic_dims():
    tuner = autotuner.AutoTuner()
    x = torch.rand([5, 1024])
    w = torch.rand([7, 9])
    dynamic_tensor_specs = (
        DynamicTensorSpec(0, 0, [1, 3, 5]),
        DynamicTensorSpec(0, 1, [16, 24, 1024]),
        # map_to_tuning_buckets is only applied at runtime, not during tuning
        DynamicTensorSpec(1,
                          1, [3, 7, 9],
                          map_to_tuning_buckets=lambda x: x // 2),
    )

    profiles = tuner._optimization_profiles(
        tuning_config=TuningConfig(dynamic_tensor_specs=dynamic_tensor_specs),
        inputs=[x, w])
    # choice(0, 0) * choice(0, 1) * choice(1, 1)
    # 3 * 3 * 3 = 27, input value 9 is already inside the bucket
    assert len(profiles) == 27
    sample_0 = OptimizationProfile(shapes=[[
        DynamicDim(min=1, opt=1, max=3),
        DynamicDim(min=16, opt=16, max=24)
    ], [StaticDim(val=7), DynamicDim(min=3, opt=3, max=7)]])
    sample_26 = OptimizationProfile(shapes=[[
        DynamicDim(min=5, opt=5, max=float('inf')),
        DynamicDim(min=1024, opt=1024, max=float('inf'))
    ], [StaticDim(
        val=7), DynamicDim(min=9, opt=9, max=float('inf'))]])

    assert sample_0 == profiles[0]
    assert sample_26 == profiles[-1]


# For cache testing
"""
tactic 0 is better when x.shape[0] <= M // 2
tactic 1 is better when x.shape[0] > M // 2
"""
M = 32


# add sleep to simulate bad perf
def gemm_0(x, w):
    if x.shape[0] > M // 2:
        delay_kernel(100, torch.cuda.current_stream())
    return x @ w


def gemm_1(x, w):
    if x.shape[0] <= M // 2:
        delay_kernel(100, torch.cuda.current_stream())
    return x @ w


def gemm_fallback(x, w) -> torch.Tensor:
    # always the slowest
    delay_kernel(500, torch.cuda.current_stream())
    return x @ w


def check_gemm_tactic_valid(tactic: int, m: int) -> bool:
    # TODO: CI is not stable for this test. delay_kernel can not guarantee the profiling result.
    # We need to find a more determinist way to test this.
    if m <= M // 2:
        if tactic != 0:
            logger.warning(
                f"Expect tactic 0 but got {tactic} when m ({m}) is small.")
    elif m <= M:
        if tactic != 1:
            logger.warning(
                f"Expect tactic 1 but got {tactic} when m ({m}) is large.")
    else:
        if tactic != -1:
            logger.warning(
                f"Expect fallback tactic (-1) but got {tactic} when m ({m}) > {M}."
            )


class GemmRunner(TunableRunner):

    def get_valid_tactics(self, inputs: List[FakeTensor],
                          profile: OptimizationProfile, **kwargs) -> List[int]:
        # The simulated delay is not deterministic, so we need to return specific tactics here
        return [-1, 0, 1]

    def forward(self,
                /,
                inputs: List[torch.Tensor],
                *,
                tactic: int = -1,
                **kwargs) -> torch.Tensor:
        assert tactic in [-1, 0, 1]
        return [gemm_0, gemm_1, gemm_fallback][tactic](*inputs)


class JitGemmRunner(GemmRunner):
    """GemmRunner that compiles per tactic, like the CuTe DSL runners."""
    kernel_cache: dict[int, bool] = {}
    calls: list[tuple[tuple[int, ...], int]] = []

    def forward(self,
                /,
                inputs: list[torch.Tensor],
                *,
                tactic: int = -1,
                **kwargs) -> torch.Tensor:
        type(self).kernel_cache.setdefault(tactic, True)
        type(self).calls.append((tuple(inputs[0].shape), tactic))
        return super().forward(inputs, tactic=tactic, **kwargs)


@pytest.mark.parametrize("strategy", [
    DistributedTuningStrategy.PARALLEL,
    DistributedTuningStrategy.MERGE,
])
def test_prime_cached_tactics_on_cache_hit(monkeypatch, strategy) -> None:
    """Prime local cached winners without starting collective profiling."""
    monkeypatch.setenv("TLLM_AUTOTUNER_PRIME_CACHED_TACTICS", "1")
    w = torch.randn(64, 128)
    x = torch.randn(3, 64)
    tuning_config = TuningConfig(
        dynamic_tensor_specs=(
            DynamicTensorSpec(input_idx=0,
                              dim_idx=0,
                              gen_tuning_buckets=(3, 4, 5),
                              map_to_tuning_buckets=lambda value: value),
            DynamicTensorSpec(input_idx=1,
                              dim_idx=1,
                              gen_tuning_buckets=(64, 128, 256, 512),
                              map_to_tuning_buckets=lambda value: value),
        ),
        distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL)
    op = "test_prime_cached_tactics"
    tuner = AutoTuner.get()
    tuner.clear_cache()
    with tempfile.TemporaryDirectory() as temp_dir:
        cache_path = os.path.join(temp_dir, "prime.json")
        with autotune(cache_path=cache_path):
            tuner.choose_one(op, [JitGemmRunner()], tuning_config, [x, w])
        winners = {
            value[1]
            for value in tuner.profiling_cache.get_specific_custom_op(
                op).values()
        }
        assert winners

        # Simulate a fresh process that only has the persistent tactic cache.
        tuning_config.distributed_tuning_strategy = strategy
        tuner.profiling_cache.clear()
        tuner._primed_cached_tactics.clear()
        JitGemmRunner.kernel_cache.clear()
        JitGemmRunner.calls.clear()
        if strategy == DistributedTuningStrategy.MERGE:

            def unexpected_prepare(*args, **kwargs) -> None:
                pytest.fail("A collective cache hit must not prepare inputs")

            monkeypatch.setattr(tuner, "_prepare_input_tensors",
                                unexpected_prepare)
        with autotune(cache_path=cache_path):
            tuner.choose_one(op, [JitGemmRunner()], tuning_config, [x, w])
        expected = winners if strategy != DistributedTuningStrategy.MERGE else set(
        )
        assert set(JitGemmRunner.kernel_cache) == expected
        assert len(JitGemmRunner.calls) == len(expected)

        with autotune(cache_path=cache_path):
            tuner.choose_one(op, [JitGemmRunner()], tuning_config, [x, w])
        assert len(JitGemmRunner.calls) == len(expected)

        monkeypatch.setenv("TLLM_AUTOTUNER_PRIME_CACHED_TACTICS", "0")
        tuner._primed_cached_tactics.clear()
        JitGemmRunner.calls.clear()
        with autotune(cache_path=cache_path):
            tuner.choose_one(op, [JitGemmRunner()], tuning_config, [x, w])
        assert JitGemmRunner.calls == []


def test_prime_cached_tactics_after_parallel_cache_merge(monkeypatch) -> None:
    """Prime winners selected by another rank after a cold profile."""
    monkeypatch.setenv("TLLM_AUTOTUNER_PRIME_CACHED_TACTICS", "1")
    x = torch.randn(20, 64)
    w = torch.randn(64, 128)
    tuning_config = TuningConfig(
        dynamic_tensor_specs=(DynamicTensorSpec(
            input_idx=0,
            dim_idx=0,
            gen_tuning_buckets=(3, 20),
            map_to_tuning_buckets=lambda value: value,
        ), ),
        distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL,
    )
    op = "test_prime_cached_tactics_after_parallel_cache_merge"
    tuner = AutoTuner.get()
    runner = JitGemmRunner()
    tuner.clear_cache()
    tuner._primed_cached_tactics.clear()
    JitGemmRunner.kernel_cache.clear()
    JitGemmRunner.calls.clear()

    def merge_remote_winners(strategy, custom_op) -> None:
        assert strategy == DistributedTuningStrategy.PARALLEL
        for m, tactic in ((3, 0), (20, 1)):
            cache_key = tuner.profiling_cache.get_cache_key(
                custom_op,
                runner,
                ((m, 64), (64, 128)),
                tuning_config,
                apply_map_to_tuning_buckets=False,
            )
            tuner.profiling_cache[cache_key] = (0, tactic, 0.5)

    monkeypatch.setattr(tuner, "is_tuning_mode", True)
    # Distributed synchronization is covered separately.
    monkeypatch.setattr(tuner, "_should_current_rank_tune", lambda _: False)
    monkeypatch.setattr(tuner, "cache_pp_recv", lambda: None)
    monkeypatch.setattr(tuner, "_maybe_sync_cache_data", merge_remote_winners)

    selected_runner, tactic = tuner.choose_one(op, [runner], tuning_config,
                                               [x, w])

    assert selected_runner is runner
    assert tactic == 1
    assert set(JitGemmRunner.kernel_cache) == {0, 1}
    assert JitGemmRunner.calls == [((3, 64), 0), ((20, 64), 1)]


def test_failed_cached_tactic_prime_is_retried(monkeypatch) -> None:
    """Do not memoize a cached winner until its prime launch succeeds."""
    monkeypatch.setenv("TLLM_AUTOTUNER_PRIME_CACHED_TACTICS", "1")
    x = torch.randn(3, 64)
    w = torch.randn(64, 128)
    tuning_config = TuningConfig(
        distributed_tuning_strategy=DistributedTuningStrategy.PARALLEL)
    op = "test_failed_cached_tactic_prime_is_retried"
    tuner = AutoTuner.get()
    runner = JitGemmRunner()
    tuner.clear_cache()
    tuner._primed_cached_tactics.clear()
    JitGemmRunner.kernel_cache.clear()
    JitGemmRunner.calls.clear()
    profile = tuner._optimization_profiles(tuning_config, [x, w])[0]
    cache_key = tuner.profiling_cache.get_cache_key(
        op,
        runner,
        profile.get_opt_shapes(),
        tuning_config,
        apply_map_to_tuning_buckets=False,
    )
    tuner.profiling_cache[cache_key] = (0, 0, 1.0)

    original_forward = runner.forward
    attempts = 0

    def fail_once(self, inputs, *, tactic=-1, **kwargs) -> torch.Tensor:
        nonlocal attempts
        attempts += 1
        if attempts == 1:
            raise RuntimeError("injected prime failure")
        return original_forward(inputs, tactic=tactic, **kwargs)

    monkeypatch.setattr(JitGemmRunner, "forward", fail_once)
    tuner._prime_cached_tactics(op, [runner], tuning_config, [x, w])
    tuner._prime_cached_tactics(op, [runner], tuning_config, [x, w])

    assert attempts == 2
    assert set(JitGemmRunner.kernel_cache) == {0}
    assert JitGemmRunner.calls == [((3, 64), 0)]


@torch.library.custom_op("autotuner_test::get_best_gemm_tactic",
                         mutates_args=())
def get_best_gemm_tactic(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
    runners = [GemmRunner()]
    tuner = AutoTuner.get()
    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
        input_idx=0,
        dim_idx=0,
        gen_tuning_buckets=get_power_of_2_num_tokens_buckets,
        map_to_tuning_buckets=next_positive_power_of_2), ), )
    runner, tactic = tuner.choose_one(
        "autotuner_test::get_best_gemm_tactic",
        runners,
        tuning_config,
        [x, w],
    )
    return torch.tensor(tactic)


@get_best_gemm_tactic.register_fake
def _(x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
    return torch.empty(1)


def test_autotuner_cache_basic():
    w = torch.randn(64, 128)

    # tuning with largest M
    AutoTuner.get().clear_cache()
    with autotune():
        torch.ops.autotuner_test.get_best_gemm_tactic(torch.randn(M, 64), w)

    # This tests the logic of print_profiling_cache and print_statistics
    AutoTuner.get().print_profiling_cache()
    AutoTuner.get().print_statistics()

    m = M * 2
    while m >= 1:
        best_tactic = torch.ops.autotuner_test.get_best_gemm_tactic(
            torch.randn(m, 64), w)
        check_gemm_tactic_valid(best_tactic, m)
        m //= 2


def _capture_autotuner_debug_lines(monkeypatch) -> List[str]:
    """Collect what the tuner writes through ``_debug_logger`` during a tuning run.

    The lines are the only record of what the tuner considered: the profiling
    cache keeps just the winner, so a tactic choice can otherwise be explained
    only by re-running the tuner with tactics removed. Patching the instance
    attribute rather than the logger keeps other loggers out of the capture --
    ``_debug_logger`` is bound once in ``__init__`` to ``logger.info`` or
    ``logger.debug`` depending on TLLM_AUTOTUNER_LOG_LEVEL_DEBUG_TO_INFO, and
    this test must not depend on which.
    """
    lines: List[str] = []
    monkeypatch.setattr(AutoTuner.get(), "_debug_logger", lines.append)
    return lines


def test_every_candidate_is_logged_and_the_winner_is_the_best_of_them(
        monkeypatch):
    """Each profiled (runner, tactic) is logged, and Selected agrees with them.

    Two separate things, both needed to read a tactic choice off a run's log:

    1. Candidates the tuner measured appear, not just the winner.
    2. The Selected line is in the same family as the candidates -- it used to
       be the one [Autotuner] line on logger.debug while the rest went through
       _debug_logger, so promoting the family to INFO gave a log with every
       candidate and no winner, which is precisely the log a tuning
       investigation asks for. Capturing through _debug_logger would not see it
       at all if that regressed.

    The agreement check is what makes this more than a formatting test: each
    Selected time must equal the minimum candidate time of some profile, so a
    candidate line that reported a different measurement from the one the
    comparison used would fail here.
    """
    w = torch.randn(64, 128)
    AutoTuner.get().clear_cache()

    lines = _capture_autotuner_debug_lines(monkeypatch)
    with autotune():
        torch.ops.autotuner_test.get_best_gemm_tactic(torch.randn(M, 64), w)

    op = "autotuner_test::get_best_gemm_tactic"
    # Parsed with a regex, not a split on ", ": the shapes field is a tuple of
    # tuples and carries commas of its own.
    candidate_re = re.compile(
        rf"^\[Autotuner\] Candidate: custom_op={re.escape(op)}, runner=.*?, "
        r"tactic=(?P<tactic>-?\d+), shapes=(?P<shapes>.*), "
        r"time=(?P<time>\S+)ms$")
    selected_re = re.compile(
        rf"^\[Autotuner\] Selected: custom_op={re.escape(op)}, runner=.*?, "
        r"tactic=(?P<tactic>-?\d+), time=(?P<time>\S+)ms, fine_grained=")

    candidates = [m for m in map(candidate_re.match, lines) if m]
    selected = [m for m in map(selected_re.match, lines) if m]

    assert candidates, f"no candidate lines were logged; got {lines}"
    assert selected, f"no Selected line was logged; got {lines}"

    best_per_profile = {}
    for m in candidates:
        shapes, time = m["shapes"], m["time"]
        best = best_per_profile.get(shapes)
        if best is None or float(time) < float(best):
            best_per_profile[shapes] = time

    # String equality, deliberately: both lines format through the same
    # "{:.3f}", so the winner's time is the best candidate's time character for
    # character unless the two were computed from different measurements.
    winners = set(best_per_profile.values())
    for m in selected:
        assert m["time"] in winners, (
            f"Selected reports {m['time']}ms but no profile's best candidate "
            f"matches; candidate minima were {sorted(winners)}")


def test_bucket_mapping():
    """Test that map_to_tuning_buckets correctly maps runtime sizes to tuning buckets.

    This test demonstrates the single mapper approach:
    - During tuning: NO mapper is applied, raw bucket values are used as cache keys
    - During runtime: map_to_tuning_buckets is applied to map buffer size to actual work size

    With sparsity=0.25, the buffer contains 25% actual work:
    - Tuning stores buckets: 1, 2, 4, 8, 16, 32 as raw cache keys
    - Runtime buffer 4 -> maps to bucket int(4 * 0.25) = 1
    - Runtime buffer 16 -> maps to bucket int(16 * 0.25) = 4

    In MoE EP, the input buffer is allocated for worst-case but sparsely filled.
    Using map_to_tuning_buckets allows us to map buffer size to actual work size at runtime.
    """
    w = torch.randn(64, 128)
    tuner = AutoTuner.get()
    tuner.clear_cache()

    # Sparsity indicates the fraction of buffer containing valid work
    sparsity = 0.25

    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
        input_idx=0,
        dim_idx=0,
        gen_tuning_buckets=get_power_of_2_num_tokens_buckets(M),
        map_to_tuning_buckets=lambda x: int(x * sparsity)), ), )

    with autotune():
        tuner.choose_one("test_bucket_mapping", [GemmRunner()], tuning_config,
                         [torch.randn(1, 64), w])

    # Verify cache entries use raw tuning bucket values
    cache_entries = tuner.profiling_cache.get_specific_custom_op(
        "test_bucket_mapping")

    # Extract the first dimension of the first input shape from each cache key
    assert len(cache_entries) == len(tuning_config.dynamic_tensor_specs[0].gen_tuning_buckets), \
        f"Expected {len(tuning_config.dynamic_tensor_specs[0].gen_tuning_buckets)} cache entries, got {len(cache_entries)}"

    # Test runtime mapping: buffer size is mapped via map_to_runtime_buckets
    # to find the correct tuning bucket based on actual work size
    test_cases = [
        # size 4 -> valid work size (4*0.25)=1, tactic 0 since 1 <= M//2
        (4, 1, 0),
        # size 8 -> valid work size (8*0.25)=2, tactic 0 since 2 <= M//2
        (8, 2, 0),
        # size 16 -> valid work size (16*0.25)=4, tactic 0 since 4 <= M//2
        (16, 4, 0),
        # size 32 -> valid work size (32*0.25)=8, tactic 0 since 8 <= M//2
        (32, 8, 0),
        # size 64 -> valid work size (64*0.25)=16, tactic 0 since 16 <= M//2
        (64, 16, 0),
        # size 128 -> valid work size (128*0.25)=32, tactic 1 since 32 > M//2
        (128, 32, 1),
        # size 256 -> valid work size (256*0.25)=64, tactic -1 since 64 > M
        (256, 64, -1),
    ]

    for buffer_size, valid_size, expected_tactic in test_cases:
        # Verify cache lookup succeeds with the mapped bucket
        x = torch.randn(buffer_size, 64)
        runner, tactic = tuner.choose_one("test_bucket_mapping", [GemmRunner()],
                                          tuning_config, [x, w])
        assert (
            tactic == expected_tactic
        ), f"buffer size={buffer_size} -> valid work size={valid_size}, expected tactic {expected_tactic} but got {tactic}"


def test_autotuner_try_block():

    class PartialCrashedRunner(TunableRunner):

        def get_valid_tactics(self, inputs: List[FakeTensor],
                              profile: OptimizationProfile,
                              **kwargs) -> List[int]:
            return [-1, 0, 1]

        def forward(self,
                    /,
                    inputs: List[torch.Tensor],
                    *,
                    tactic: int = -1) -> torch.Tensor:
            assert tactic in [-1, 0, 1]
            if tactic == 1:
                raise Exception(
                    "For profiling try block test: Tactic 1 is not suitable. Crash happens."
                )
            return [gemm_0, gemm_1, gemm_fallback][tactic](*inputs)

    x, w = torch.randn(M, 64), torch.randn(64, 128)
    runners = [PartialCrashedRunner()]
    tuner = AutoTuner.get()
    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
        input_idx=0,
        dim_idx=0,
        gen_tuning_buckets=get_power_of_2_num_tokens_buckets,
        map_to_tuning_buckets=next_positive_power_of_2), ), )
    with autotune():
        runner, tactic = tuner.choose_one("test_autotuner_try_block", runners,
                                          tuning_config, [x, w])

    m = M // 2
    while m >= 1:
        _, tactic = tuner.choose_one("test_autotuner_try_block", runners,
                                     tuning_config, [torch.randn(m, 64), w])
        assert tactic in [
            -1, 0
        ], f"Expect only tactic -1, 0 being chosen, but got tactic {tactic}."
        m //= 2


@torch.library.custom_op("autotuner_test::recursive_get_best_gemm_tactic",
                         mutates_args=())
def recursive_get_best_gemm_tactic(x: torch.Tensor, w1: torch.Tensor,
                                   w2: torch.Tensor) -> torch.Tensor:
    # Only the first custom_op is tuned, the second one uses the tuned result in cache
    tactic_1 = get_best_gemm_tactic(x, w1)
    tactic_2 = get_best_gemm_tactic(x, w2)
    return torch.stack([tactic_1, tactic_2])


@recursive_get_best_gemm_tactic.register_fake
def _(x: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor) -> torch.Tensor:
    return torch.empty(2)


def test_recursive_autotuner():
    x, w1, w2 = torch.randn(M, 64), torch.randn(64, 128), torch.randn(64, 128)
    AutoTuner.get().clear_cache()
    with autotune():
        torch.ops.autotuner_test.recursive_get_best_gemm_tactic(
            torch.randn(M, 64), w1, w2)

    m = M * 2
    while m >= 1:
        t1, t2 = torch.ops.autotuner_test.recursive_get_best_gemm_tactic(
            torch.randn(m, 64), w1, w2)
        check_gemm_tactic_valid(t1, m)
        check_gemm_tactic_valid(t2, m)
        m //= 2


class GemmRunnerWithAttributes(TunableRunner):

    def __init__(self, block_size: int, num_warps: int):
        self.block_size = block_size
        self.num_warps = num_warps

    def get_valid_tactics(self, inputs: List[FakeTensor],
                          profile: OptimizationProfile, **kwargs) -> List[int]:
        return [-1, 0, 1]

    def forward(self,
                /,
                inputs: List[torch.Tensor],
                *,
                tactic: int = -1) -> torch.Tensor:
        assert tactic in [-1, 0, 1]
        return [gemm_0, gemm_1, gemm_fallback][tactic](*inputs)


def test_multiple_runners_different_attributes():
    """Test that runners with different attributes get different cache entries"""
    x, w = torch.randn(16, 64), torch.randn(64, 128)

    # Create runners with different attributes
    runner_0 = GemmRunnerWithAttributes(block_size=128, num_warps=4)
    runner_1 = GemmRunnerWithAttributes(block_size=256, num_warps=8)
    runners = [runner_0, runner_1]

    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
        input_idx=0,
        dim_idx=0,
        gen_tuning_buckets=get_power_of_2_num_tokens_buckets,
        map_to_tuning_buckets=next_positive_power_of_2), ), )

    # Do tuning
    with autotune():
        tuner = AutoTuner.get()
        runner_a, tactic_a = tuner.choose_one("test_multiple_runners", runners,
                                              tuning_config, [x, w])

        # Verify different cache keys are generated
        shapes = (x.shape, w.shape)
        cache_key_0 = tuner.profiling_cache.get_cache_key(
            custom_op="test_multiple_runners",
            input_shapes=shapes,
            runner=runner_0,
            tuning_config=tuning_config,
        )
        cache_key_1 = tuner.profiling_cache.get_cache_key(
            custom_op="test_multiple_runners",
            input_shapes=shapes,
            runner=runner_1,
            tuning_config=tuning_config,
        )

        assert cache_key_0 != cache_key_1, "Runners with different attributes should have different cache keys"


def test_multiple_dynamic_shapes_cache():
    """Test that different dynamic shape combinations are properly cached"""
    w = torch.randn(64, 128)
    runners = [GemmRunner()]

    # Define dynamic ranges for both dimensions
    tuning_config = TuningConfig(dynamic_tensor_specs=(
        DynamicTensorSpec(input_idx=0,
                          dim_idx=0,
                          gen_tuning_buckets=(3, 4, 5),
                          map_to_tuning_buckets=lambda x: x),
        DynamicTensorSpec(input_idx=1,
                          dim_idx=1,
                          gen_tuning_buckets=(64, 128, 256, 512),
                          map_to_tuning_buckets=lambda x: x),
    ), )

    # Do tuning with a sample input
    x = torch.randn(3, 64)
    temp_dir = tempfile.TemporaryDirectory()
    cache_path = os.path.join(temp_dir.name,
                              "test_multiple_dynamic_shapes.json")
    with autotune(cache_path=cache_path):
        tuner = AutoTuner.get()
        runner, tactic = tuner.choose_one("test_multiple_dynamic_shapes",
                                          runners, tuning_config, [x, w])

    cache_entries = tuner.profiling_cache.get_specific_custom_op(
        "test_multiple_dynamic_shapes")
    assert len(cache_entries) == 12, \
        f"Expected 12 cache entries for 3x4 shape combinations, got {len(cache_entries)}"
    # Verify cache size - should have 12 entries (3x4 combinations)
    # We also test the cache serialization and deserialization here.
    AutoTuner.get().profiling_cache.clear()
    AutoTuner.get().profiling_cache.load_cache(cache_path, rank=0)
    cache_entries = tuner.profiling_cache.get_specific_custom_op(
        "test_multiple_dynamic_shapes")

    assert len(cache_entries) == 12, \
        f"Expected 12 cache entries for 3x4 shape combinations, got {len(cache_entries)}"


class GemmRunnerComplexTuningConfigs(TunableRunner):

    # test serialization of different types of tactics
    valid_tactic_ids = [-1, 0, 1]
    valid_tile_sizes = [(128, 128), (256, 256)]
    valid_cluster_sizes = [[1, 1, 1], [2, 2, 1]]

    tune_max_num_tokens = 32

    def get_valid_tactics(
        self,
        inputs: List[FakeTensor],
        profile: OptimizationProfile,
        **kwargs,
    ) -> List[Any]:
        # During the tuning process, we verify if the tuning config behaves as expected
        assert inputs[0].shape[0] <= self.tune_max_num_tokens, \
            f"Input shape {inputs[0].shape[0]} is larger than the max num tokens {self.tune_max_num_tokens}"

        assert inputs[0][-1, 0] == inputs[0].shape[0], \
            f"Input shape {inputs[0].shape[0]} is not set through the pre_hook correctly"

        return [{
            "int_tactic_id": tactic_id,
            "tuple_tile_size": tile_size,
            "list_cluster_size": cluster_size,
        } for tactic_id, tile_size, cluster_size in itertools.product(
            self.valid_tactic_ids,
            self.valid_tile_sizes,
            self.valid_cluster_sizes,
        )]

    def forward(
        self,
        /,
        inputs: List[torch.Tensor],
        *,
        tactic: Any = -1,
    ) -> torch.Tensor:
        # Notice that in fallback case tactic is -1
        if tactic == -1:
            # assign default configs for fallback case
            tactic_id, tile_size, cluster_size = -1, (128, 256), [1, 1, 1]
        else:
            tactic_id, tile_size, cluster_size = tactic[
                "int_tactic_id"], tactic["tuple_tile_size"], tactic[
                    "list_cluster_size"]

        assert isinstance(tactic_id, int) and tactic_id in self.valid_tactic_ids
        assert isinstance(tile_size, tuple) and len(tile_size) == 2 \
            and tile_size in self.valid_tile_sizes
        assert isinstance(cluster_size, list) and len(cluster_size) == 3 \
            and cluster_size in self.valid_cluster_sizes
        return [gemm_0, gemm_1, gemm_fallback][tactic_id](*inputs)

    @staticmethod
    def inputs_pre_hook(inputs: List[torch.Tensor]):
        # always set the first element to be the number of tokens in x
        x, w = inputs
        x_hooked = torch.zeros_like(x)
        x_hooked[-1, 0] = x.shape[0]
        return [x_hooked, w]


def test_autotuner_tuning_configs():
    runner_0 = GemmRunnerComplexTuningConfigs()
    runners = [runner_0]
    x, w = torch.randn(64, 64), torch.randn(64, 128)
    tuning_config = TuningConfig(
        dynamic_tensor_specs=(DynamicTensorSpec(
            input_idx=0,
            dim_idx=0,
            gen_tuning_buckets=get_power_of_2_num_tokens_buckets,
            map_to_tuning_buckets=next_positive_power_of_2,
        ), ),
        # Test if the number of tuning tokens is clipped to 32
        tune_max_num_tokens=GemmRunnerComplexTuningConfigs.tune_max_num_tokens,
        inputs_pre_hook=GemmRunnerComplexTuningConfigs.inputs_pre_hook,
        use_cold_l2_cache=True,
        use_cuda_graph=False,
    )
    temp_dir = tempfile.TemporaryDirectory()
    cache_path = os.path.join(temp_dir.name,
                              "test_autotuner_tactic_configs.json")
    with autotune(cache_path=cache_path):
        tuner = AutoTuner.get()
        runner, best_tactic = tuner.choose_one("test_autotuner_tactic_configs",
                                               runners, tuning_config, [x, w])

    runner_0([x, w], tactic=best_tactic)

    # Test if the tactic can be loaded from cache correctly
    AutoTuner.get().profiling_cache.clear()
    AutoTuner.get().profiling_cache.load_cache(cache_path, rank=0)

    # No further tuning should be performed.
    runner, deserialized_tactic = tuner.choose_one(
        "test_autotuner_tactic_configs", runners, tuning_config, [x, w])
    assert best_tactic == deserialized_tactic, "Tactic should be the same after deserialization"

    runner_0([x, w], tactic=deserialized_tactic)


def test_load_cache_skips_non_literal_tactic():
    """Regression: a non-literal tactic repr must be skipped on load, not crash it.

    ``_deserialize_cache_data`` reconstructs tactics with ``ast.literal_eval``,
    which raises ``SyntaxError`` on non-literal reprs (e.g. enum tactic reprs,
    until #16782 serializes enums by value). It must skip such entries -- once
    ``SyntaxError`` was uncaught and had no ``continue``, crashing the load.
    """
    import ast

    class _NonLiteralTactic:

        def __repr__(self):
            return "<_NonLiteralTactic object nvfp4>"

    poisoned_repr = repr(_NonLiteralTactic())  # non-literal object repr
    # Precondition: confirm this repr really does raise SyntaxError.
    with pytest.raises(SyntaxError):
        ast.literal_eval(poisoned_repr)

    cache = AutoTuner.get().profiling_cache
    cache.clear()
    good_key = "('op_good', 'R', '0', ((1, 128),))"
    bad_key = "('op_bad', 'R', '0', ((2, 128),))"
    doc = {
        "metadata": cache._serialize_metadata(),
        "shared": {},
        "rank_0": {
            good_key: {
                "runner_id": 0,
                "tactic": "7",
                "min_time": 0.001
            },
            bad_key: {
                "runner_id": 1,
                "tactic": poisoned_repr,
                "min_time": 0.002
            },
        },
    }
    temp_dir = tempfile.TemporaryDirectory()
    cache_path = os.path.join(temp_dir.name, "poisoned_cache.json")
    with open(cache_path, "w") as f:
        json.dump(doc, f)

    # Must not raise (previously raised SyntaxError out of load_cache).
    cache.load_cache(cache_path, rank=0)

    # The literal-safe entry survived with its exact tactic ...
    good = ("op_good", "R", "0", ((1, 128), ))
    assert good in cache.cache
    assert cache.cache[good][1] == 7
    # ... and the non-literal entry was skipped, not silently mis-decoded.
    bad = ("op_bad", "R", "0", ((2, 128), ))
    assert bad not in cache.cache


def test_kernel_testing_single_context():
    """Test kernel testing with a single choose_one context"""
    x, w = torch.randn(16, 64), torch.randn(64, 128)
    runners = [GemmRunner()]
    tuning_config = TuningConfig(dynamic_tensor_specs=(DynamicTensorSpec(
        input_idx=0,
        dim_idx=0,
        gen_tuning_buckets=get_power_of_2_num_tokens_buckets,
        map_to_tuning_buckets=next_positive_power_of_2), ), )

    tuner = AutoTuner.get()
    tuner.clear_cache()

    # First, do tuning to populate cache
    with autotune():
        runner, tactic = tuner.choose_one("test_kernel_testing_single", runners,
                                          tuning_config, [x, w])

    # Capture execution context
    with tuner.capture() as all_tactics:
        runner, tactic = tuner.choose_one("test_kernel_testing_single", runners,
                                          tuning_config, [x, w])
        reference_output = runner([x, w], tactic=tactic)

    # Test all tactics
    tested_tactics = []
    for (runner, tactic), in all_tactics:
        tested_tactics.append((runner, tactic))
        with tuner.replay(((runner, tactic), )):
            runner_ret, tactic_ret = tuner.choose_one(
                "test_kernel_testing_single", runners, tuning_config, [x, w])
            output = runner_ret([x, w], tactic=tactic_ret)
            # Verify output matches reference
            torch.testing.assert_close(output, reference_output)
            assert runner == runner_ret and tactic == tactic_ret, \
                f"Runner and tactic mismatch: expected ({runner, tactic}), got ({runner_ret, tactic_ret})"

    # Should have tested 3 tactics ([-1, 0, 1])
    assert len(tested_tactics) == len(GemmRunner().get_valid_tactics([x, w], OptimizationProfile([[]]))), \
        f"Expected 3 tactics to be tested, got {len(tested_tactics)}"


class MultiContextRunner(TunableRunner):

    def get_valid_tactics(self, inputs: List[FakeTensor],
                          profile: OptimizationProfile, **kwargs) -> List[int]:
        gemm_idx = kwargs.get("gemm_idx", 0)
        # Different gemm_idx have different number of tactics
        if gemm_idx == 0:
            return [0, 1]
        else:
            return [0, 1, 2]

    def forward(self,
                /,
                inputs: List[torch.Tensor],
                *,
                tactic: int = -1,
                **kwargs) -> torch.Tensor:
        gemm_idx = kwargs.get("gemm_idx", 0)
        # Analogous to CUTLASS MoE trtllm::fused_moe FC1
        if gemm_idx == 0:
            return [gemm_0, gemm_1][tactic](inputs[0], inputs[1])
        # Analogous to CUTLASS MoE trtllm::fused_moe FC2
        else:
            return [gemm_0, gemm_1, gemm_fallback][tactic](inputs[1].T,
                                                           inputs[0].T)


def test_kernel_testing_multiple_contexts():
    """
    Test kernel testing with multiple choose_one contexts
    (e.g., CUTLASS MoE trtllm::fused_moe)
    """

    x, w = torch.randn(16, 64), torch.randn(64, 128)
    runners = [MultiContextRunner()]
    tuning_config = TuningConfig()

    tuner = AutoTuner.get()
    tuner.clear_cache()

    # First, do tuning to populate cache
    with autotune():
        runner, _ = tuner.choose_one("test_multi_context",
                                     runners,
                                     tuning_config, [x, w],
                                     gemm_idx=0)
        runner, _ = tuner.choose_one("test_multi_context",
                                     runners,
                                     tuning_config, [x, w],
                                     gemm_idx=1)

    # Capture execution context (captures both choose_one calls)
    with tuner.capture() as all_tactics:
        runner_0, tactic_0 = tuner.choose_one("test_multi_context",
                                              runners,
                                              tuning_config, [x, w],
                                              gemm_idx=0)
        runner_1, tactic_1 = tuner.choose_one("test_multi_context",
                                              runners,
                                              tuning_config, [x, w],
                                              gemm_idx=1)
        ref_output_0 = runner_0([x, w], tactic=tactic_0, gemm_idx=0)
        ref_output_1 = runner_1([x, w], tactic=tactic_1, gemm_idx=1)

    # Test all tactic combinations (cartesian product)
    tested_tactics = []
    for tactic in all_tactics:
        tested_tactics.append(tactic)
        # Each tactic is ((runner_0, tactic_0), (runner_1, tactic_1))
        assert len(tactic) == 2, f"Expected 2 contexts, got {len(tactic)}"

        with tuner.replay(tactic):
            # Make the same calls in the same order
            runner_0, tactic_0 = tuner.choose_one("test_multi_context",
                                                  runners,
                                                  tuning_config, [x, w],
                                                  gemm_idx=0)
            runner_1, tactic_1 = tuner.choose_one("test_multi_context",
                                                  runners,
                                                  tuning_config, [x, w],
                                                  gemm_idx=1)

            output_0 = runner_0([x, w], tactic=tactic_0, gemm_idx=0)
            output_1 = runner_1([x, w], tactic=tactic_1, gemm_idx=1)

            # Verify each context independently
            # Since we're testing different tactics, outputs will differ
            # Just verify they don't crash and have correct shapes
            assert output_0.shape == ref_output_0.shape
            assert output_1.shape == ref_output_1.shape

    # Should have tested 2*3 = 6 combinations
    num_tactics_for_gemm_idx = lambda gemm_idx: len(runners[
        0].get_valid_tactics([x, w], OptimizationProfile(), gemm_idx=gemm_idx))
    assert len(tested_tactics) == num_tactics_for_gemm_idx(0) * num_tactics_for_gemm_idx(1), \
        f"Expected 6 tactic combinations (2*3), got {len(tested_tactics)}"


def test_kernel_testing_mismatched_ops():
    """
    Correctly raise and capture the exception when captured context != operation performed
    """
    x, w = torch.randn(16, 64), torch.randn(64, 128)
    runners = [GemmRunner()]
    tuning_config = TuningConfig()

    tuner = AutoTuner.get()
    tuner.clear_cache()

    # Capture execution context for operation A
    with tuner.capture() as all_tactics:
        _ = tuner.choose_one("test_op_A", runners, tuning_config, [x, w])

    # Try to test with operation B (should raise RuntimeError)
    try:
        for (runner, tactic), in all_tactics:
            with tuner.replay(((runner, tactic), )):
                # This should raise RuntimeError because custom_op doesn't match
                _ = tuner.choose_one("test_op_B", runners, tuning_config,
                                     [x, w])
        assert False, "Expected RuntimeError for mismatched custom_op, but none was raised"
    except RuntimeError as e:
        # Verify the error message contains useful information
        error_msg = str(e)
        assert "Custom op mismatch" in error_msg, f"Expected 'Custom op mismatch' in error message, got: {error_msg}"
        assert "test_op_A" in error_msg, f"Expected 'test_op_A' in error message, got: {error_msg}"
        assert "test_op_B" in error_msg, f"Expected 'test_op_B' in error message, got: {error_msg}"


class DistributedGemmRunner(TunableRunner):

    def __init__(self, prefer_tactics: List[int] = [0, 1]):
        self.prefer_tactics = prefer_tactics

    def get_valid_tactics(self, inputs, profile, **kwargs):
        # Return all tactics so merge strategy can choose between them
        return self.prefer_tactics

    def forward(self, inputs, *, tactic=-1, **kwargs):
        # tactic 0 is slower
        if tactic % 2 == 0:
            for _ in range(5):
                inputs[0] @ inputs[1]
        return inputs[0] @ inputs[1]

    def unique_id(self):
        return ()


def _distributed_worker_function(world_size, strategy):
    """Worker function to run on each MPI rank."""
    rank = tensorrt_llm.mpi_rank()
    mapping = Mapping(world_size=world_size,
                      rank=rank,
                      tp_size=world_size,
                      pp_size=1)
    dist = Distributed.get(mapping)

    tuner = AutoTuner.get()
    tuner.clear_cache()
    tuner.setup_distributed_state(mapping)

    x = torch.randn(16, 32, device='cuda')
    w = torch.randn(32, 64, device='cuda')
    inputs = [x, w]

    if strategy == DistributedTuningStrategy.PARALLEL:
        # All ranks get the same set of tactics
        prefer_tactics = [0, 1, 2, 3]
    else:
        # Each rank prefers different tactics
        prefer_tactics = [rank]
    runner = DistributedGemmRunner(prefer_tactics=prefer_tactics)
    runner_independent = DistributedGemmRunner()
    config = TuningConfig(distributed_tuning_strategy=strategy)
    config_independent = TuningConfig(
        distributed_tuning_strategy=DistributedTuningStrategy.INDEPENDENT)

    # Keep temp_dir in function scope to prevent premature garbage collection
    temp_dir = None
    if rank == 0:
        temp_dir = tempfile.TemporaryDirectory()
        # rank 0 should broadcast the cache path to all ranks
        cache_path = os.path.join(temp_dir.name, "test_distributed_tuning.json")
        dist.broadcast(cache_path, root=0)
    else:
        cache_path = dist.broadcast(None, root=0)

    with autotune(cache_path=cache_path):
        tuner.choose_one(custom_op=f"test_distributed_{strategy.value}",
                         runners=[runner],
                         tuning_config=config,
                         inputs=inputs)
        # run another normal gemm with INDEPENDENT strategy
        tuner.choose_one(custom_op=f"test_distributed_normal_gemm",
                         runners=[runner_independent],
                         tuning_config=config_independent,
                         inputs=inputs)

    # Check only one cache file is created in the cache path.
    # The sibling ".lock" file is an implementation artifact of
    # _exclusive_cache_lock (see tensorrt_llm/_torch/autotuner.py) and is not
    # a per-rank cache file.
    cache_dir = os.path.dirname(cache_path)
    cache_files = [f for f in os.listdir(cache_dir) if not f.endswith(".lock")]
    assert len(cache_files) == 1, "Only one rank file should be created"

    dist.barrier()

    # Check cache for distributed tuning
    AutoTuner.get().profiling_cache.clear()
    AutoTuner.get().profiling_cache.load_cache(cache_path, rank)

    selected_runner, best_tactic = tuner.choose_one(
        custom_op=f"test_distributed_{strategy.value}",
        runners=[runner],
        tuning_config=config,
        inputs=inputs)

    # Verify cache file structure based on distributed strategy
    with open(cache_path, 'r') as f:
        cache_data = json.load(f)

    # Helper to check if an op name appears in any cache key string
    def has_op_in_section(section_data: dict, op_name: str) -> bool:
        return any(op_name in key_str for key_str in section_data.keys())

    assert 'metadata' in cache_data, "Metadata should be present"
    assert f'rank_{rank}' in cache_data, f"rank {rank} should be present"

    # The INDEPENDENT op "test_distributed_normal_gemm" should always be in rank-specific sections
    assert has_op_in_section(cache_data[f'rank_{rank}'], 'test_distributed_normal_gemm'), \
        f"rank {rank} should have test_distributed_normal_gemm"

    if strategy == DistributedTuningStrategy.INDEPENDENT:
        # Both ops use INDEPENDENT strategy, so no shared section
        assert 'shared' not in cache_data or len(cache_data.get('shared', {})) == 0, \
            "shared should not be present or be empty for INDEPENDENT strategy"
        # Each rank should have 2 entries (the parameterized op + normal_gemm)
        assert len(cache_data[f'rank_{rank}']) == 2, \
            f"rank {rank} should have 2 entries, got {len(cache_data[f'rank_{rank}'])}"
        assert has_op_in_section(cache_data[f'rank_{rank}'], f'test_distributed_{strategy.value}'), \
            f"rank {rank} should have test_distributed_{strategy.value}"

        assert len(
            AutoTuner.get().profiling_cache.independent_op
        ) == 0, f"Non-INDEPENDENT ops should not be present in the cache"
    else:
        # Non-INDEPENDENT ops go to shared section
        assert 'shared' in cache_data, "shared section should be present"
        # Each rank should have only 1 entry (the normal_gemm with INDEPENDENT strategy)
        assert len(cache_data[f'rank_{rank}']) == 1, \
            f"rank {rank} should have 1 entry, got {len(cache_data[f'rank_{rank}'])}"
        # The parameterized op should NOT be in rank-specific section
        assert not has_op_in_section(cache_data[f'rank_{rank}'], f'test_distributed_{strategy.value}'), \
            f"rank {rank} should not have test_distributed_{strategy.value}"
        # The parameterized op should be in shared section
        assert has_op_in_section(cache_data['shared'], f'test_distributed_{strategy.value}'), \
            f"shared should have test_distributed_{strategy.value}"

        assert "test_distributed_normal_gemm" not in AutoTuner.get().profiling_cache.independent_op and \
            f"test_distributed_{strategy.value}" in AutoTuner.get().profiling_cache.independent_op, \
            f"Distributed tuning strategy is not recovered correctly from cache"

    if strategy == DistributedTuningStrategy.BROADCAST:
        # All ranks should select tactic 0
        assert best_tactic == 0, f"Rank {rank} with {strategy} should select tactic 0, got {best_tactic}"
    elif strategy == DistributedTuningStrategy.INDEPENDENT:
        # Each rank should select the tactic it prefers
        assert best_tactic == rank, f"Rank {rank} with {strategy} should select tactic {rank}, got {best_tactic}"
    elif strategy == DistributedTuningStrategy.MERGE:
        # Because tactic 0 is slower, two ranks should always select tactic 1
        assert best_tactic == 1, f"Rank {rank} with {strategy} should select tactic 1, got {best_tactic}"
    elif strategy == DistributedTuningStrategy.PARALLEL:
        # Tactic 1 or 3 should be selected since they are faster.
        # TODO: This might not cover the case that rank1 tunes nothing
        assert best_tactic % 2 == 1, f"Rank {rank} with {strategy} should select tactic 1, got {best_tactic}"
    else:
        assert False, f"Rank {rank} got unknown strategy: {strategy}"

    dist.barrier()
    return True


@pytest.mark.skipif(torch.cuda.device_count() < 2,
                    reason="Requires at least 2 GPUs for this test")
@pytest.mark.parametrize(
    "strategy",
    [
        DistributedTuningStrategy.BROADCAST,
        DistributedTuningStrategy.INDEPENDENT,
        DistributedTuningStrategy.MERGE,
        DistributedTuningStrategy.PARALLEL,
    ],
)
@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True)
def test_autotuner_distributed_strategy(strategy, mpi_pool_executor):
    world_size = 2
    # Use MPIPoolExecutor to run distributed test
    results = mpi_pool_executor.map(
        _distributed_worker_function,
        *zip(*[(
            world_size,
            strategy,
        )] * world_size),
    )
    for r in results:
        assert r is True


@pytest.mark.parametrize("use_cuda_graph", [False, True])
def test_global_timer_vs_cuda_event(use_cuda_graph, monkeypatch):
    """Verify globaltimer and cuda-event backends are statistically indistinguishable."""

    class PureGemmRunner(TunableRunner):

        def get_valid_tactics(self, inputs: List[FakeTensor],
                              profile: OptimizationProfile,
                              **kwargs) -> List[int]:
            return [0]

        def forward(self,
                    /,
                    inputs: List[torch.Tensor],
                    *,
                    tactic: int = 0,
                    **kwargs) -> torch.Tensor:
            assert tactic == 0
            return inputs[0] @ inputs[1]

    # Keep full profiling repeats enabled to reduce measurement noise.
    monkeypatch.setenv("TLLM_AUTOTUNER_DISABLE_SHORT_PROFILE", "1")

    gemm_shapes = [
        (256, 4096, 11008),
        (512, 8192, 8192),
    ]
    num_trials = 6
    rel_tol = 0.05
    stat_zscore = 3.0
    abs_tol_ms = 0.01

    runner = PureGemmRunner()
    tuning_config = TuningConfig(use_cuda_graph=use_cuda_graph)
    tuner = AutoTuner()
    trial_rows = []

    for m, k, n in gemm_shapes:
        x = torch.randn(m, k, device='cuda', dtype=torch.float16)
        w = torch.randn(k, n, device='cuda', dtype=torch.float16)

        event_times = []
        gt_times = []

        # Interleave both backends to avoid drift effects from neighboring load.
        for _ in range(num_trials):
            tuner._use_global_timer = False
            event_times.append(
                tuner._profile_single_kernel(
                    runner=runner,
                    inputs=[x, w],
                    tactic=0,
                    tuning_config=tuning_config,
                    use_cuda_graph=use_cuda_graph,
                ))

            tuner._use_global_timer = True
            gt_times.append(
                tuner._profile_single_kernel(
                    runner=runner,
                    inputs=[x, w],
                    tactic=0,
                    tuning_config=tuning_config,
                    use_cuda_graph=use_cuda_graph,
                ))

            event_ms = event_times[-1]
            gt_ms = gt_times[-1]
            abs_diff = abs(gt_ms - event_ms)
            rel_diff = abs_diff / event_ms if event_ms > 0 else float('inf')
            trial_rows.append((m, k, n, len(event_times), event_ms, gt_ms,
                               abs_diff, rel_diff))

        event_mean = statistics.fmean(event_times)
        gt_mean = statistics.fmean(gt_times)
        event_var = statistics.variance(event_times)
        gt_var = statistics.variance(gt_times)
        mean_diff = abs(gt_mean - event_mean)
        rel_diff = mean_diff / event_mean

        # Two-sample mean delta should be small vs a fixed tolerance and
        # indistinguishable within sampling noise.
        combined_sem = math.sqrt(event_var / num_trials + gt_var / num_trials)
        allowed_diff = max(abs_tol_ms, rel_tol * event_mean,
                           stat_zscore * combined_sem)

        assert event_mean > 0, (
            f"({m},{k},{n}): cuda event mean should be positive, got {event_mean}"
        )
        assert gt_mean > 0, (
            f"({m},{k},{n}): globaltimer mean should be positive, got {gt_mean}"
        )
        assert mean_diff <= allowed_diff, (
            f"({m},{k},{n}): timing backends are distinguishable "
            f"(cuda_event_mean={event_mean:.4f}ms, "
            f"globaltimer_mean={gt_mean:.4f}ms, "
            f"relative_diff={rel_diff * 100:.2f}%, "
            f"allowed_diff={allowed_diff:.4f}ms, "
            f"combined_sem={combined_sem:.4f}ms, "
            f"event_samples={event_times}, gt_samples={gt_times})")

    # Visible with `pytest -s`; otherwise captured by pytest.
    print("\nGlobaltimer vs cuda-event trial comparison")
    print(f"cuda_graph={use_cuda_graph}, trials_per_shape={num_trials}")
    print("-" * 102)
    print(
        f"{'shape (M,K,N)':>21} {'trial':>5} {'cuda_event (ms)':>16} "
        f"{'globaltimer (ms)':>17} {'abs diff (ms)':>14} {'rel diff (%)':>13}")
    print("-" * 102)
    for m, k, n, trial, event_ms, gt_ms, abs_diff, rel_diff in trial_rows:
        print(f"{f'({m},{k},{n})':>21} {trial:>5d} "
              f"{event_ms:>16.4f} {gt_ms:>17.4f} "
              f"{abs_diff:>14.4f} {rel_diff * 100:>13.2f}")
    print("-" * 102)


def _make_shapes(*sizes):
    """Convert size-tuples into Tuple[torch.Size, ...] for _find_nearest_profile."""
    return tuple(torch.Size(s) for s in sizes)


class TestSpecBoundsChecking:
    """Bounds-checking in AutoTuner._find_nearest_profile and AutoTuner._optimization_profiles."""

    def setup_method(self):
        AutoTuner._find_nearest_profile.cache_clear()

    @pytest.mark.parametrize("entry", ["find_nearest", "optimization_profiles"])
    @pytest.mark.parametrize(
        "spec_class,input_idx,dim_idx",
        [
            pytest.param("dynamic", 5, 0, id="dynamic_input_idx_out_of_range"),
            pytest.param("dynamic", 0, 10, id="dynamic_dim_idx_out_of_range"),
            pytest.param("dynamic", -1, 0, id="dynamic_negative_input_idx"),
            pytest.param("dynamic", 0, -1, id="dynamic_negative_dim_idx"),
            pytest.param(
                "constraint", 3, 0, id="constraint_input_idx_out_of_range"),
            pytest.param(
                "constraint", 0, 7, id="constraint_dim_idx_out_of_range"),
            pytest.param(
                "constraint", -1, 0, id="constraint_negative_input_idx"),
            pytest.param("constraint", 0, -1, id="constraint_negative_dim_idx"),
        ],
    )
    def test_oob_spec_skipped(self, entry, spec_class, input_idx, dim_idx):
        from tensorrt_llm._torch.autotuner import ConstraintSpec
        if spec_class == "dynamic":
            spec = DynamicTensorSpec(input_idx=input_idx,
                                     dim_idx=dim_idx,
                                     gen_tuning_buckets=(1, 2))
            dyn_specs = (spec, )
            con_specs = ()
        else:
            spec = ConstraintSpec(input_idx=input_idx,
                                  dim_idx=dim_idx,
                                  infer_shape=lambda shapes: 1)
            dyn_specs = ()
            con_specs = (spec, )

        if entry == "find_nearest":
            shapes = _make_shapes([4, 8])
            result = AutoTuner._find_nearest_profile(
                shapes,
                dynamic_tensor_specs=dyn_specs,
                constraint_specs=con_specs)
            assert result == ((4, 8), )
        else:
            tuner = AutoTuner()
            x = torch.rand([4, 8])
            # Constraint-only configs need a dynamic spec to drive the cartesian product.
            if not dyn_specs:
                dyn_specs = (DynamicTensorSpec(input_idx=0,
                                               dim_idx=0,
                                               gen_tuning_buckets=(1, )), )
            config = TuningConfig(dynamic_tensor_specs=dyn_specs,
                                  constraint_specs=con_specs)
            profiles = tuner._optimization_profiles(config, [x])
            # OOB spec skipped — profile generation still produces at least one profile.
            assert len(profiles) >= 1


def test_single_pair_shortcut(monkeypatch):
    """Single (runner, tactic) candidate must bypass the timed profile loop.

    When ``_profile_runners`` sees exactly one (runner, tactic) pair, it
    must (1) skip ``_profile_single_kernel`` entirely, (2) still fire the
    ``do_preparation`` hook for runners that opt in, (3) fire exactly one
    ``forward()`` to drive any JIT side effect, and (4) record the pair
    in the profiling cache. Multi-tactic ops in the same fixture must
    still use the timed path.
    """

    profile_calls: List[Any] = []

    def _track(self, runner, inputs, tactic, tuning_config, **kwargs):
        profile_calls.append(tactic)
        return 1.0 + len(profile_calls) * 0.01

    monkeypatch.setattr(AutoTuner, "_profile_single_kernel", _track)

    forward_calls: List[tuple] = []

    class PrepRunner(TunableRunner):

        def unique_id(self):
            return ()

        def get_valid_tactics(self, inputs: List[FakeTensor],
                              profile: OptimizationProfile,
                              **kwargs) -> List[int]:
            return [0]

        def forward(self,
                    /,
                    inputs: List[torch.Tensor],
                    *,
                    tactic: int = -1,
                    do_preparation: bool = False,
                    **kwargs) -> torch.Tensor:
            forward_calls.append((tactic, do_preparation))
            if do_preparation:
                return None
            x, w = inputs
            return x @ w

    tuner = AutoTuner.get()
    tuner.clear_cache()
    lines = _capture_autotuner_debug_lines(monkeypatch)
    x = torch.randn(M, 64, device="cuda")
    w = torch.randn(64, 128, device="cuda")

    # Single (runner, tactic): shortcut must fire.
    op_single = "autotuner_test::single_pair_shortcut"
    with autotune():
        _, tactic = tuner.choose_one(op_single, [PrepRunner()], TuningConfig(),
                                     [x, w])
    assert tactic == 0
    assert profile_calls == [], (
        f"_profile_single_kernel must not be called for single-pair op; "
        f"got {profile_calls}")
    assert forward_calls == [
        (-1, True), (0, False)
    ], (f"Expected do_preparation then exactly one forward(tactic=0); "
        f"got {forward_calls}")
    assert len(tuner.profiling_cache.get_specific_custom_op(op_single)) == 1, (
        "single-pair shortcut must still record the (runner, tactic) entry")
    # The shortcut bypasses the timed loop but must not bypass the candidate
    # log: 0.000 is the recorded-without-profiling marker, and the format is
    # the timed path's, so one regex reads a run's log whichever path ran.
    assert any(
        line.startswith(f"[Autotuner] Candidate: custom_op={op_single}, ")
        and line.endswith("time=0.000ms") for line in lines), (
            f"single-pair shortcut must log its candidate; got {lines}")

    # Multi-tactic on the same fixture: timed profile path must still run.
    forward_calls.clear()
    op_multi = "autotuner_test::single_pair_shortcut_multi"
    with autotune():
        tuner.choose_one(op_multi, [GemmRunner()], TuningConfig(), [x, w])
    # GemmRunner exposes 3 tactics -> 3 profile calls.
    assert len(profile_calls) == 3, (
        f"Multi-tactic op must hit _profile_single_kernel per tactic; "
        f"got {len(profile_calls)} ({profile_calls})")


def test_single_pair_shortcut_failure_is_logged_as_inf(monkeypatch):
    """A failing single-pair run still gets a Candidate record, as inf.

    The timed path lists failed candidates as inf next to the ones that ran;
    the shortcut must match, or a run whose only candidate crashed would show
    neither a Candidate nor a Selected line and read as if the op was never
    tuned. choose_one itself must not raise: the fallback (runner 0, tactic
    -1) covers the op, same as an all-failed timed loop."""

    class FailingRunner(TunableRunner):

        def unique_id(self):
            return ()

        def get_valid_tactics(self, inputs: List[FakeTensor],
                              profile: OptimizationProfile,
                              **kwargs) -> List[int]:
            return [0]

        # No do_preparation in the signature: the shortcut fires that hook
        # OUTSIDE its try (same as the timed path), so a runner that raises
        # there crashes choose_one before the candidate is ever attempted --
        # this test is about the attempted-and-failed candidate.
        def forward(self,
                    /,
                    inputs: List[torch.Tensor],
                    *,
                    tactic: int = -1,
                    **kwargs) -> torch.Tensor:
            raise RuntimeError("single-pair candidate crash")

    tuner = AutoTuner.get()
    tuner.clear_cache()
    lines = _capture_autotuner_debug_lines(monkeypatch)
    x = torch.randn(M, 64, device="cuda")
    w = torch.randn(64, 128, device="cuda")

    op = "autotuner_test::single_pair_shortcut_failure"
    with autotune():
        runner, tactic = tuner.choose_one(op, [FailingRunner()], TuningConfig(),
                                          [x, w])
    assert tactic == -1, "an all-failed op must fall back to tactic -1"
    assert any(
        line.startswith(f"[Autotuner] Candidate: custom_op={op}, ")
        and line.endswith("time=infms") for line in lines
    ), (f"a failed single-pair candidate must be logged as inf; got {lines}")


def test_cutedsl_nvfp4_heuristic_matches_full_sweep(monkeypatch):
    """End-to-end guard for the nvMatmulHeuristics tactic generation.

    For one representative problem size, the tactic the AutoTuner selects when
    ``AutoTuner.nvmmh_config`` builds the candidates from nvMatmulHeuristics'
    top configurations must be no slower (up to a small tolerance) than the
    tactic it selects from the full CuteDSL NVFP4 tactic sweep. This validates
    that the nvMMH guidance does not cost performance.

    It additionally compares the heuristic-chosen CuteDSL kernel against the
    cuBLASLt NVFP4 GEMM on the same fp4 inputs: the CuteDSL kernel-only device
    time must be within a small tolerance of cuBLAS. Both kernel symbol names
    are logged for manual inspection -- cuBLAS exposes no API for its selected
    kernel's CTA tile / cluster shape, so that is not asserted.

    Note: the sweep and heuristic paths do NOT profile identical candidate sets
    -- nvMMH may also pick rasterization orders or swizzles the sweep does not
    enumerate -- so the exact winning tactic can differ. Only the achieved
    runtime is a meaningful invariant, hence the tolerance comparisons.

    Blackwell only (SM100/SM103) and requires the nvMatmulHeuristics library;
    skipped otherwise.
    """
    if not torch.cuda.is_available():
        pytest.skip("requires a CUDA device")
    try:
        from tensorrt_llm._utils import get_sm_version
        sm_version = get_sm_version()
    except Exception:
        sm_version = None
    if sm_version not in (100, 103):
        pytest.skip("CuteDSL NVFP4 requires SM100 (B200) / SM103 (B300)")

    from tensorrt_llm._torch.custom_ops import \
        cutedsl_matmul_heuristics as nvmmh
    if not nvmmh.IS_NVMMH_AVAILABLE:
        pytest.skip("nvMatmulHeuristics library not installed")
    from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import \
        CuteDSLNVFP4BlackwellRunner

    # One representative square problem. fp4 packing / scale-factor layout
    # follows shmoo_nvfp4_cutedsl_heuristics.py::_quantize_inputs.
    m = n = k = 4096
    dtype = torch.bfloat16
    sf_vec_size = 16
    torch.manual_seed(0)
    x = torch.randn((m, k), dtype=dtype).cuda()
    w = torch.randn((n, k), dtype=dtype).cuda()
    x_sf_global = (448 * 6) / x.abs().max().float()
    w_sf_global = (448 * 6) / w.abs().max().float()
    x_fp4, x_sf = torch.ops.trtllm.fp4_quantize(x, x_sf_global, sf_vec_size,
                                                False)
    w_fp4, w_sf = torch.ops.trtllm.fp4_quantize(w, w_sf_global, sf_vec_size,
                                                False)
    alpha = torch.tensor([1.0], device="cuda")
    inputs = [x_fp4, w_fp4, x_sf, w_sf, alpha]

    runner = CuteDSLNVFP4BlackwellRunner(output_dtype=dtype)
    tuning_config = runner.__class__.tuning_config

    def _best_tactic():
        tuner = AutoTuner.get()
        tuner.clear_cache()
        with autotune():
            _, tactic = tuner.choose_one(
                "test::cutedsl_nvfp4_heuristic_match",
                [runner],
                tuning_config,
                inputs,
            )
        return tactic

    def _dominant_kernel(fn, iters=20):
        """(name, per-call device-us) of the longest-running CUDA kernel that
        fn launches, isolating kernel time from host/op overhead."""
        from torch.profiler import ProfilerActivity, profile
        fn()
        torch.cuda.synchronize()
        with profile(activities=[ProfilerActivity.CUDA]) as prof:
            for _ in range(iters):
                fn()
            torch.cuda.synchronize()
        best_name, best_total, best_count = "<none>", -1.0, 1
        for e in prof.key_averages():
            t = (getattr(e, "self_device_time_total", 0)
                 or getattr(e, "self_cuda_time_total", 0))
            if t and t > best_total:
                best_name, best_total, best_count = e.key, t, max(1, e.count)
        return best_name, best_total / best_count

    # Full sweep: no nvMMH guidance.
    tuner = AutoTuner.get()
    monkeypatch.setattr(tuner, "nvmmh_config", None)
    sweep_tactic = _best_tactic()

    # nvMatmulHeuristics builds the candidates from its top 5 configurations.
    monkeypatch.setattr(tuner, "nvmmh_config", NvMMHConfig(max_tactics=5))
    heuristic_tactic = _best_tactic()

    # cuBLASLt runs its own heuristic auto-tuning; warm it under autotune().
    def _cublas_call():
        return torch.ops.trtllm.nvfp4_gemm_cublaslt(x_fp4, w_fp4, x_sf, w_sf,
                                                    alpha, dtype)

    with autotune():
        _cublas_call()
    torch.cuda.synchronize()

    # All comparisons use kernel-only device time (isolates the GEMM kernel from
    # host/op dispatch overhead), measured via the CUDA profiler.
    _, sweep_us = _dominant_kernel(lambda: runner(inputs, tactic=sweep_tactic))
    _, heuristic_us = _dominant_kernel(
        lambda: runner(inputs, tactic=heuristic_tactic))
    _, cublas_us = _dominant_kernel(_cublas_call)

    # The guidance must not degrade the achieved kernel runtime beyond this
    # tolerance. With the default max_tactics=5 the nvMMH set includes the
    # empirical best tile (its cluster ranking can still be slightly off, ~2-3%
    # on square 4096), so a tight bound catches gross regressions.
    tolerance = 1.05
    assert heuristic_us <= sweep_us * tolerance, (
        f"heuristic-pruned tactic {heuristic_tactic} ({heuristic_us:.2f} us) is "
        f">{tolerance:.2f}x slower than full-sweep tactic {sweep_tactic} "
        f"({sweep_us:.2f} us) for M={m}, N={n}, K={k}")

    # The heuristic CuteDSL kernel should beat cuBLAS or be within tolerance.
    # This is a cross-library comparison (CuteDSL vs cuBLASLt) at ~28us per call
    # profiled over only 20 iterations; run-to-run jitter easily reaches a few
    # percent from DVFS, L2 residency, and interleaving with cuBLAS autotune
    # warmup, so keep this bound looser than the intra-CuteDSL one above.
    cublas_tolerance = 1.10
    assert heuristic_us <= cublas_us * cublas_tolerance, (
        f"CuteDSL heuristic kernel ({heuristic_us:.2f} us) is "
        f">{cublas_tolerance:.2f}x slower than cuBLAS NVFP4 "
        f"({cublas_us:.2f} us) for M={m}, N={n}, K={k}")


def _nvmmh_problem(case, sm_version):
    """(runner, inputs, run, reference) of one nvMMH-aware CuTe DSL runner of
    this GPU, or None if it has none for the case; run(tactic) returns that
    tactic's output."""
    from _torch.helpers import per_block_cast_to_fp8, per_block_cast_to_fp8_e8m0

    import tensorrt_llm.quantization.utils.fp8_utils as fp8_utils
    from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops as ops

    torch.manual_seed(0)
    bf16 = torch.bfloat16
    rubin = sm_version == 107

    def randn(*shape):
        return torch.randn(*shape, device="cuda", dtype=bf16)

    def into(runner, inputs):

        def run(tactic):
            inputs[-1].fill_(float("nan"))
            runner(inputs, tactic=tactic)
            return inputs[-1]

        return run

    def returned(runner, inputs):
        return lambda tactic: runner(inputs, tactic=tactic)

    if case in ("bf16_gemm", "bf16_gemm_fp32_out"):
        # nvMMH splits K of the FP32-output router GEMM and tiles the
        # BF16-output decode GEMM narrowly.
        fp32_out = case == "bf16_gemm_fp32_out"
        m, n, k = (64, 256, 7168) if fp32_out else (64, 2112, 7168)
        c_dtype = torch.float32 if fp32_out else bf16
        a, b = randn(m, k), randn(n, k)
        runner_class = (ops.CuteDSLBf16RubinGemmRunner
                        if rubin else ops.CuteDSLBf16BlackwellGemmRunner)
        runner = runner_class(output_dtype=c_dtype)
        inputs = [a, b, torch.empty(m, n, device="cuda", dtype=c_dtype)]
        return runner, inputs, into(runner, inputs), a.float() @ b.float().t()
    if case in ("bf16_bmm", "bf16_bmm_fp8_out"):
        fp8_out = case == "bf16_bmm_fp8_out"
        if fp8_out and rubin:
            return None
        batch, m, n, k = 8, 300, 256, 512
        a, b = randn(batch, m, k), randn(batch, n, k)
        if rubin:
            runner = ops.CuteDSLBf16RubinBmmRunner()
        elif fp8_out:
            runner = ops.CuteDSLBf16BlackwellBmmFp8OutRunner()
        else:
            runner = ops.CuteDSLBf16BlackwellBmmRunner()
        c_dtype = torch.float8_e4m3fn if fp8_out else bf16
        inputs = [a, b, torch.empty(batch, m, n, device="cuda", dtype=c_dtype)]
        reference = torch.bmm(a.float(), b.float().transpose(1, 2))
        if fp8_out:
            reference = reference.to(c_dtype).float()
        return runner, inputs, into(runner, inputs), reference
    if case == "fp8_gemm":
        m, n, k = 64, 2048, 7168
        a, b = randn(m, k) / k, randn(n, k) / k
        a_fp8, a_sf = torch.ops.trtllm.fp8_quantize_1x128(a)
        b_fp8, b_sf = per_block_cast_to_fp8(b)
        runner = (ops.CuteDSLFp8RubinGemmRunner()
                  if rubin else ops.CuteDSLFp8BlackwellRunner())
        inputs = [a_fp8, b_fp8, a_sf, b_sf]
        reference = a.float() @ b.float().t()
        return runner, inputs, returned(runner, inputs), reference
    if case == "fp8_bmm":
        # On SM100 nvMMH splits the 64 K tiles in 3 (unequal slices).
        batch, m, n, k = 4, 64, 512, 8192
        a, b = randn(m, batch, k) / k, randn(batch, n, k) / k
        a_fp8, a_sf = torch.ops.trtllm.fp8_batched_quantize_1x128_permute102(a)
        b_fp8 = torch.empty_like(b, dtype=torch.float8_e4m3fn)
        b_sf = torch.empty(batch,
                           n // 128,
                           k // 128,
                           device="cuda",
                           dtype=torch.float32)
        for group in range(batch):
            b_fp8[group], b_sf[group] = per_block_cast_to_fp8(b[group])
        runner = (ops.CuteDSLFp8RubinBmmRunner()
                  if rubin else ops.CuteDSLFp8BlackwellBmmRunner())
        inputs = [
            a_fp8, b_fp8, a_sf, b_sf,
            torch.empty(batch, m, n, device="cuda", dtype=bf16)
        ]
        reference = torch.einsum("mgk,gnk->gmn", a.float(), b.float())
        return runner, inputs, into(runner, inputs), reference
    if case == "mxfp8_gemm":
        if not rubin:
            return None
        # A DeepSeek-V4 decode projection; nvMMH splits its 56 K tiles.
        m, n, k = 32, 2048, 7168
        a, b = randn(m, k) / k, randn(n, k) / k
        a_fp8, a_sf = torch.ops.trtllm.fp8_quantize_1x128_packed_ue8m0(a)
        b_fp8, b_sf_k128 = per_block_cast_to_fp8_e8m0(b)
        b_sf = fp8_utils.transform_k128_scales_to_cutedsl_mxfp8_layout(
            b_sf_k128, mn=n, k=k)
        alpha = ops._get_mxfp8_gemm_alpha(a.device)
        runner = ops.CuteDSLMXFP8RubinLinear(output_dtype=bf16)
        inputs = [a_fp8, b_fp8, a_sf, b_sf, alpha]
        reference = a.float() @ b.float().t()
        return runner, inputs, returned(runner, inputs), reference
    assert case == "nvfp4_gemm"
    # nvMMH splits the 64 K blocks in 2 and 3 on SM100.
    m, n, k = 64, 2048, 16384
    x, w = randn(m, k), randn(n, k)
    x_scale = 2688 / x.abs().max().float()
    w_scale = 2688 / w.abs().max().float()
    x_fp4, x_sf = torch.ops.trtllm.fp4_quantize(x, x_scale, 16, False)
    w_fp4, w_sf = torch.ops.trtllm.fp4_quantize(w, w_scale, 16, False)
    alpha = (1 / (x_scale * w_scale)).reshape(1)
    inputs = [x_fp4, w_fp4, x_sf, w_sf, alpha]
    runner = (ops.CuteDSLNVFP4RubinLinear(bf16)
              if rubin else ops.CuteDSLNVFP4BlackwellRunner(bf16))
    reference = torch.ops.trtllm.nvfp4_gemm_cutlass(*inputs, bf16).float()
    return runner, inputs, returned(runner, inputs), reference


@pytest.mark.parametrize(
    "case",
    [
        "bf16_gemm",
        "bf16_gemm_fp32_out",
        "bf16_bmm",
        "bf16_bmm_fp8_out",
        "fp8_gemm",
        "fp8_bmm",
        "mxfp8_gemm",
        "nvfp4_gemm",
    ],
)
def test_nvmmh_tactics_compute_the_product(case, monkeypatch):
    """Every nvMMH-guided tactic of the SM100/SM103 or SM107 CuTe DSL
    GEMM/BMM runners is correct, including tiles, rasters and split-K
    factors their full sweeps never enumerate."""
    from tensorrt_llm._torch.custom_ops import \
        cutedsl_matmul_heuristics as nvmmh
    from tensorrt_llm._torch.cute_dsl_utils import (
        IS_CUTLASS_DSL_AVAILABLE, IS_CUTLASS_DSL_RUBIN_AVAILABLE)
    from tensorrt_llm._utils import get_sm_version

    sm_version = get_sm_version()
    if not IS_CUTLASS_DSL_AVAILABLE or sm_version not in (100, 103, 107):
        pytest.skip("The test requires SM100/SM103/SM107 and CuTe DSL.")
    if sm_version == 107 and not IS_CUTLASS_DSL_RUBIN_AVAILABLE:
        pytest.skip("The test requires the SM107 CuTe DSL helpers.")
    if not nvmmh.IS_NVMMH_AVAILABLE:
        pytest.skip("nvMatmulHeuristics is not installed")
    from _torch.helpers import calc_diff

    problem = _nvmmh_problem(case, sm_version)
    if problem is None:
        pytest.skip(f"No nvMMH-aware {case} runner on SM{sm_version}.")
    runner, inputs, run, reference = problem
    selected = []
    select_tactics = nvmmh.select_tactics

    def recording_select_tactics(ranked, max_tactics):
        selected.append(select_tactics(ranked, max_tactics))
        return selected[-1]

    monkeypatch.setattr(nvmmh, "select_tactics", recording_select_tactics)
    monkeypatch.setattr(AutoTuner.get(), "nvmmh_config",
                        NvMMHConfig(max_tactics=3))
    tactics = runner.get_valid_tactics(inputs, None)
    # The tactics come from nvMMH, not from the full-sweep fallback.
    assert tactics and tactics == selected[-1]
    # FP8 inputs are quantized from the BF16 reference operands.
    max_diff = 1e-3 if case.startswith(("fp8", "mxfp8")) else 1e-4
    for tactic in tactics:
        output = run(tactic).float()
        torch.cuda.synchronize()
        assert torch.isfinite(output).all(), tactic
        assert calc_diff(output, reference) < max_diff, tactic


@pytest.mark.parametrize("distribution", ["random", "balanced"])
def test_trtllm_gen_moe_dummy_topk_local_experts_less_than_topk(
        distribution, monkeypatch):
    """NVBugs 6457853: autotuner warmup must not fail on EP shards where
    local_num_experts < top_k (e.g. gpt-oss-120b: 128 experts, top_k=4,
    EP64 -> 2 local experts per rank, attention-DP => use_dp=True).
    Dummy rows keep the production shape: top_k distinct ids per row, all
    local experts present, remaining slots padded with out-of-shard ids."""
    from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import \
        prepare_dummy_topk_and_hook

    monkeypatch.setenv("TRTLLM_GEN_MOE_AUTOTUNE_DUMMY_DISTRIBUTION",
                       distribution)
    num_tokens, top_k = 8, 4
    num_experts, local_num_experts, local_expert_offset = 128, 2, 6
    hidden_states = torch.randn(num_tokens,
                                64,
                                dtype=torch.bfloat16,
                                device="cuda")
    topk_ids = torch.randint(0,
                             num_experts, (num_tokens, top_k),
                             dtype=torch.int32,
                             device="cuda")
    topk_weights = torch.ones(num_tokens,
                              top_k,
                              dtype=torch.bfloat16,
                              device="cuda")

    with autotune():
        _, dummy_weights, dummy_ids, _ = prepare_dummy_topk_and_hook(
            topk_weights,
            topk_ids,
            hidden_states,
            None,
            1,
            TuningConfig(),
            top_k,
            num_experts,
            local_num_experts,
            None,
            None,
            None,
            local_expert_offset=local_expert_offset,
            use_dp=True)

    assert dummy_ids.shape == (num_tokens, top_k)
    assert dummy_ids.dtype == torch.int32
    assert dummy_weights.shape == (num_tokens, top_k)
    shard = range(local_expert_offset, local_expert_offset + local_num_experts)
    for row in dummy_ids.tolist():
        assert len(set(row)) == top_k, f"duplicate ids in row {row}"
        assert sum(x in shard for x in row) == local_num_experts, (
            f"expected all {local_num_experts} local experts in row {row}")
        assert all(0 <= x < num_experts for x in row), row


def test_post_tune_merge_tactics_min_time_and_subset_kept():
    tuner = autotuner.AutoTuner()
    tuner.mapping = Mapping(world_size=2, rank=0, tp_size=2)

    r0 = {
        ("gemm", "X"): (0, ("cutlass", 10), 1.0),
        ("q", "A"): (0, ("trtllm", ), 1.0)
    }
    r1 = {
        ("gemm", "X"): (0, ("cublaslt", 0), 0.5),
        ("vae", "B"): (0, ("cutlass", 2), 1.0)
    }

    class _FakeDist:

        def tp_cp_allgather(self, obj):
            return [r0, r1]

    tuner._dist = _FakeDist()
    tuner.profiling_cache.cache = dict(r0)
    tuner.post_tune_merge_tactics()

    cache = tuner.profiling_cache.cache
    assert cache[("gemm", "X")] == (0, ("cublaslt", 0), 0.5)
    assert cache[("q", "A")] == r0[("q", "A")]
    assert cache[("vae", "B")] == r1[("vae", "B")]


def test_post_tune_merge_tactics_single_rank_noop():
    tuner = autotuner.AutoTuner()
    tuner.mapping = Mapping(world_size=1, rank=0, tp_size=1)
    tuner._dist = None
    original = {("gemm", "X"): (0, ("cutlass", 10), 1.0)}
    tuner.profiling_cache.cache = dict(original)
    tuner.post_tune_merge_tactics()
    assert tuner.profiling_cache.cache == original


@pytest.mark.parametrize("tp_size,cp_size", [(1, 4), (2, 2)])
def test_parallel_tactics_split_across_helix_cp_ranks(tp_size, cp_size):
    """Helix CP ranks share the PARALLEL tactic sweep with the TP ranks."""
    from tensorrt_llm.mapping import CpType

    world_size = tp_size * cp_size
    tactics = list(range(11))
    shares = []
    for rank in range(world_size):
        tuner = autotuner.AutoTuner()
        mapping = Mapping(world_size=world_size,
                          rank=rank,
                          tp_size=tp_size,
                          cp_size=cp_size,
                          cp_config={"cp_type": CpType.HELIX})
        tuner.setup_distributed_state(mapping, dist=object())
        assert tuner._is_distributed()
        assert tuner._tuning_group() == (world_size, rank)
        shares.append(
            tuner._maybe_parallelize_tactics(
                tactics, DistributedTuningStrategy.PARALLEL))
        assert tuner._maybe_parallelize_tactics(
            tactics, DistributedTuningStrategy.MERGE) == tactics

    # Every tactic is tuned on exactly one rank.
    assert sorted(t for share in shares for t in share) == tactics


def test_parallel_tactics_split_skips_non_helix_cp_ranks():
    """Non-Helix CP keeps the TP-only tuning group."""
    from tensorrt_llm.mapping import CpType

    tuner = autotuner.AutoTuner()
    mapping = Mapping(world_size=4,
                      rank=3,
                      tp_size=1,
                      cp_size=4,
                      cp_config={"cp_type": CpType.ULYSSES})
    tuner.setup_distributed_state(mapping, dist=object())
    assert tuner._tuning_group() == (1, 0)
    assert not tuner._is_distributed()
    tactics = list(range(5))
    assert tuner._maybe_parallelize_tactics(
        tactics, DistributedTuningStrategy.PARALLEL) == tactics


def test_profiling_cache_enum_tactic_roundtrip(tmp_path):
    # An enum tactic must survive save -> load: repr("<Enum.X: v>") isn't
    # ast.literal_eval-parsable, so the cache serializes tactic.value instead.
    class _Tac(enum.IntEnum):
        TRTLLM = -1

    key = ("op::x", "Runner", "(1,)")
    src = autotuner.AutoTuner().profiling_cache
    src.cache[key] = (0, _Tac.TRTLLM, 1.0)
    path = str(tmp_path / "cache.json")
    src.save_cache(path, rank=0)  # must not raise

    dst = autotuner.AutoTuner().profiling_cache
    dst.load_cache(path, rank=0)
    assert dst.cache[key][1] == _Tac.TRTLLM  # IntEnum compares equal to -1


def test_autotune_post_tune_merge_before_save(tmp_path):
    # autotune(post_tune_merge_dist=...) must merge across ranks and persist the
    # merged winner, then restore the singleton's prior distributed state.
    tuner = AutoTuner.get()
    tuner.clear_cache()
    r0 = {("gemm", "X"): (0, ("cutlass", 10), 1.0)}
    r1 = {("gemm", "X"): (0, ("cublaslt", 0), 0.5)}
    tuner.profiling_cache.cache = dict(r0)

    class _FakeDist:
        mapping = Mapping(world_size=2, rank=0, tp_size=2)

        def tp_cp_allgather(self, obj):
            return [r0, r1]

    prev_mapping, prev_dist = tuner.mapping, tuner._dist
    path = str(tmp_path / "cache.json")
    with autotune(cache_path=path, post_tune_merge_dist=_FakeDist()):
        pass

    # Merged in memory (fastest tactic won) and persisted before the context
    # exits (the saved file carries the merged winner).
    assert tuner.profiling_cache.cache[("gemm", "X")] == (0, ("cublaslt", 0),
                                                          0.5)
    persisted = autotuner.AutoTuner().profiling_cache
    persisted.load_cache(path, rank=0)
    assert persisted.cache[("gemm", "X")][1] == ("cublaslt", 0)
    # The temporary full-world distributed state was restored.
    assert tuner.mapping is prev_mapping
    assert tuner._dist is prev_dist


def _post_tune_merge_worker(world_size):
    """Run on each MPI rank: seed a distinct cache, then merge for real."""
    rank = tensorrt_llm.mpi_rank()
    mapping = Mapping(world_size=world_size,
                      rank=rank,
                      tp_size=world_size,
                      pp_size=1)
    tuner = AutoTuner.get()
    tuner.clear_cache()
    tuner.setup_distributed_state(mapping)

    # Shared key: rank 1's tactic is faster (should win). Plus a rank-unique key.
    shared = ("gemm", "X")
    tuner.profiling_cache.cache = {
        shared: (0, ("cutlass", 10), 1.0) if rank == 0 else
        (0, ("cublaslt", 0), 0.5),
        (f"only_rank{rank}", "K"): (0, ("trtllm", ), 1.0),
    }
    tuner.post_tune_merge_tactics()

    cache = tuner.profiling_cache.cache
    return {
        "shared": cache[shared],
        "has_r0": ("only_rank0", "K") in cache,
        "has_r1": ("only_rank1", "K") in cache,
    }


@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True)
def test_post_tune_merge_tactics_multi_rank(mpi_pool_executor):
    # Real 2-rank merge over an actual process group (not a FakeDist): every
    # rank must converge to the same cache — the min-time winner for the shared
    # key and both ranks' unique keys kept.
    world_size = 2
    results = list(
        mpi_pool_executor.map(_post_tune_merge_worker,
                              *zip(*[(world_size, )] * world_size)))
    assert len(results) == world_size
    for r in results:
        assert r["shared"] == (0, ("cublaslt", 0), 0.5)
        assert r["has_r0"] and r["has_r1"]


def _autotune_dist_allgather_worker(world_size):
    import torch.distributed as dist

    from tensorrt_llm._torch.visual_gen.mapping import _VisualGenAutotuneDist
    rank = tensorrt_llm.mpi_rank()
    os.environ["RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
    os.environ.setdefault("MASTER_PORT", "29557")
    if not dist.is_initialized():
        dist.init_process_group(backend="gloo",
                                world_size=world_size,
                                rank=rank)
    # VG's real mesh has tp_size=1; the communicator must gather over the whole
    # world regardless, so the mapping's tp axis is irrelevant to the gather.
    d = _VisualGenAutotuneDist(
        Mapping(world_size=world_size, rank=rank, tp_size=world_size))
    return d.tp_cp_allgather(f"rank{rank}")


@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True)
def test_visual_gen_autotune_dist_world_allgather(mpi_pool_executor):
    # _VisualGenAutotuneDist.tp_cp_allgather must gather every rank's object over
    # the default world group (not the single-rank tp subgroup), and construct
    # without running TorchDist.__init__.
    world_size = 2
    results = list(
        mpi_pool_executor.map(_autotune_dist_allgather_worker,
                              *zip(*[(world_size, )] * world_size)))
    for got in results:
        assert got == ["rank0", "rank1"]


def test_deep_gemm_jit_warmup_buckets():
    # A bucket left out leaves a DeepGemm layout uncompiled, so it compiles
    # under traffic while holding the GIL. The layout is constant on each
    # 16-aligned window of M, so every window must hold a bucket.
    for max_m in (4096, 6144, 8192):
        b = deep_gemm_jit_warmup_buckets(max_m)
        assert list(b) == sorted(set(b)), "must be sorted and duplicate-free"
        assert b[-1] == max_m, "the bound itself must be a bucket"

        assert all(hi - lo <= 16 for lo, hi in zip(b, b[1:])), \
            "no gap may exceed 16: it would skip a 16-aligned window of M"
        windows = {(m - 1) // 16 for m in b}
        missing = [j for j in range((max_m - 1) // 16 + 1) if j not in windows]
        assert not missing, \
            f"windows with no bucket (M from {[16 * j + 1 for j in missing]})"

    # Clamping keeps a small bound from skipping warmup for larger M,
    # and a huge one from exploding the bucket count.
    assert deep_gemm_jit_warmup_buckets(256) == deep_gemm_jit_warmup_buckets(
        4096)
    # But a decode-only worker (M = batch x MTP) stays in the low band.
    assert max(deep_gemm_jit_warmup_buckets(32)) < 128
    assert deep_gemm_jit_warmup_buckets(16384) == deep_gemm_jit_warmup_buckets(
        8192)