lightseekorg--tokenspeed/tokenspeed-kernel/test/test_kernel_api_selection.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Golden selection tests for top-level tokenspeed-kernel public APIs."""

from __future__ import annotations

import importlib
from dataclasses import dataclass
from typing import Callable

import pytest
import tokenspeed_kernel
import tokenspeed_kernel.numerics.reference.gemm as _gemm_reference
import tokenspeed_kernel.ops.attention as _attention_pkg
import tokenspeed_kernel.ops.attention.cuda as _attention_cuda
import tokenspeed_kernel.ops.attention.flash_attn as _attention_flash_attn
import tokenspeed_kernel.ops.attention.flash_mla as _attention_flash_mla
import tokenspeed_kernel.ops.attention.flashinfer as _attention_flashinfer
import tokenspeed_kernel.ops.attention.flashinfer.gated_delta_rule as _attention_flashinfer_gdn
import tokenspeed_kernel.ops.attention.gluon as _attention_gluon
import tokenspeed_kernel.ops.attention.triton as _attention_triton
import tokenspeed_kernel.ops.gemm as _gemm_pkg
import tokenspeed_kernel.ops.gemm.deep_gemm as _gemm_deep_gemm
import tokenspeed_kernel.ops.gemm.flashinfer as _gemm_flashinfer
import tokenspeed_kernel.ops.gemm.gluon as _gemm_gluon
import tokenspeed_kernel.ops.gemm.triton as _gemm_triton
import tokenspeed_kernel.ops.gemm.trtllm as _gemm_trtllm
import tokenspeed_kernel.ops.moe as _moe_pkg
import tokenspeed_kernel.ops.moe.flashinfer as _moe_flashinfer
import tokenspeed_kernel.ops.moe.gluon as _moe_gluon
import tokenspeed_kernel.ops.moe.triton as _moe_triton
import tokenspeed_kernel.ops.quantization as _quantization_pkg
import tokenspeed_kernel.ops.quantization.flashinfer as _quantization_flashinfer
import tokenspeed_kernel.ops.quantization.triton as _quantization_triton
import tokenspeed_kernel.ops.quantization.trtllm as _quantization_trtllm
import tokenspeed_kernel.ops.sampling as _sampling_pkg
import tokenspeed_kernel.ops.sampling.cute_dsl as _sampling_cute_dsl
import tokenspeed_kernel.ops.sampling.gluon as _sampling_gluon
import torch
from tokenspeed_kernel.ops.attention.gdn_utils import GdnChunkPrefillResult
from tokenspeed_kernel.ops.attention.triton import dsa as _attention_triton_dsa
from tokenspeed_kernel.ops.attention.triton import (
    dsa_topk as _attention_triton_dsa_topk,
)
from tokenspeed_kernel.ops.attention.triton import (
    gated_delta_rule as _attention_triton_gdn,
)
from tokenspeed_kernel.ops.attention.triton import (
    merge_state as _attention_triton_merge_state,
)
from tokenspeed_kernel.ops.attention.triton import (
    mha_decode as _attention_triton_mha_decode,
)
from tokenspeed_kernel.ops.attention.triton import (
    mha_prefill as _attention_triton_mha_prefill,
)
from tokenspeed_kernel.ops.attention.triton import (
    mla_decode as _attention_triton_mla_decode,
)
from tokenspeed_kernel.ops.attention.triton import (
    mla_prefill as _attention_triton_mla_prefill,
)
from tokenspeed_kernel.ops.moe.flashinfer import (
    cutedsl_deepep_nvfp4 as _moe_cutedsl_deepep_nvfp4,
)
from tokenspeed_kernel.ops.moe.flashinfer import cutlass_fp8 as _moe_cutlass_fp8
from tokenspeed_kernel.ops.moe.flashinfer import cutlass_nvfp4 as _moe_cutlass_nvfp4
from tokenspeed_kernel.ops.moe.flashinfer import cutlass_unquant as _moe_cutlass_unquant
from tokenspeed_kernel.ops.moe.flashinfer import trtllm_fp8 as _moe_trtllm_fp8
from tokenspeed_kernel.ops.moe.flashinfer import trtllm_mxfp4 as _moe_trtllm_mxfp4
from tokenspeed_kernel.ops.moe.flashinfer import trtllm_mxint4 as _moe_trtllm_mxint4
from tokenspeed_kernel.ops.moe.flashinfer import trtllm_nvfp4 as _moe_trtllm_nvfp4
from tokenspeed_kernel.ops.moe.flashinfer import trtllm_unquant as _moe_trtllm_unquant
from tokenspeed_kernel.ops.moe.gluon import mxfp4 as _moe_gluon_mxfp4
from tokenspeed_kernel.ops.moe.triton import fp8 as _moe_triton_fp8
from tokenspeed_kernel.ops.moe.triton import mxfp4 as _moe_triton_mxfp4
from tokenspeed_kernel.platform import ArchVersion, Platform, PlatformInfo
from tokenspeed_kernel.registry import KernelRegistry
from tokenspeed_kernel.selection import SelectedKernel

_RELOAD_MODULES = [
    # Attention registration modules.
    _attention_cuda,
    _attention_flash_attn,
    _attention_flash_mla,
    _attention_flashinfer_gdn,
    _attention_flashinfer,
    _attention_gluon,
    _attention_triton_mha_prefill,
    _attention_triton_mha_decode,
    _attention_triton_mla_prefill,
    _attention_triton_mla_decode,
    _attention_triton_merge_state,
    _attention_triton_dsa,
    _attention_triton_dsa_topk,
    _attention_triton_gdn,
    _attention_triton,
    _attention_pkg,
    # GEMM registration modules.
    _gemm_reference,
    _gemm_deep_gemm,
    _gemm_flashinfer,
    _gemm_gluon,
    _gemm_triton,
    _gemm_trtllm,
    _gemm_pkg,
    # MoE registration modules.
    _moe_cutedsl_deepep_nvfp4,
    _moe_cutlass_fp8,
    _moe_cutlass_nvfp4,
    _moe_cutlass_unquant,
    _moe_trtllm_fp8,
    _moe_trtllm_mxfp4,
    _moe_trtllm_mxint4,
    _moe_trtllm_nvfp4,
    _moe_trtllm_unquant,
    _moe_flashinfer,
    _moe_gluon_mxfp4,
    _moe_gluon,
    _moe_triton_fp8,
    _moe_triton_mxfp4,
    _moe_triton,
    _moe_pkg,
    # Quantization registration modules.
    _quantization_flashinfer,
    _quantization_triton,
    _quantization_trtllm,
    _quantization_pkg,
    # Sampling registration modules.
    _sampling_cute_dsl,
    _sampling_gluon,
    _sampling_pkg,
    # Top-level public API re-exports.
    tokenspeed_kernel,
]


@pytest.fixture(autouse=True)
def _kernel_registry(fresh_registry):
    """Reload real registrations into the fresh registry for each case."""
    for mod in _RELOAD_MODULES:
        importlib.reload(mod)


def test_builtin_moe_preprocessor_links_are_callables():
    kernel_registry = KernelRegistry.get()
    errors = []
    for kernel_spec in kernel_registry.list_kernels("moe", "apply"):
        preprocessor = kernel_spec.weight_preprocessor
        if preprocessor is not None and not callable(preprocessor):
            errors.append(f"{kernel_spec.name}: non-callable preprocessor")

    process_weight_kernels = kernel_registry.list_kernels("moe", "process_weights")
    assert process_weight_kernels == []

    assert errors == []


def test_moe_process_weights_returns_for_no_preprocessing_plan():
    module = torch.nn.Module()

    result = tokenspeed_kernel.moe_process_weights(
        {"weight_preprocessor": None},
        module,
    )

    assert result is None


def test_moe_process_weights_dispatches_plan_preprocessor_callable():
    calls = []

    def preprocess(plan, w):
        calls.append((plan, w))

    module = torch.nn.Module()
    plan = {"weight_preprocessor": preprocess}

    result = tokenspeed_kernel.moe_process_weights(plan, module)

    assert result is None
    assert calls == [(plan, module)]


@dataclass(frozen=True)
class KernelApiSelectionCase:
    id: str
    family: str
    mode: str
    arch: str
    expected: str
    matches: Callable[[PlatformInfo], bool]
    invoke: Callable[[], object]


def _is_hopper(platform: PlatformInfo) -> bool:
    return platform.is_hopper


def _is_blackwell_sm100(platform: PlatformInfo) -> bool:
    return platform.is_blackwell and platform.arch_version == ArchVersion(10, 0)


def _is_blackwell_non_sm100(platform: PlatformInfo) -> bool:
    return platform.is_blackwell and platform.arch_version != ArchVersion(10, 0)


def _is_blackwell_plus(platform: PlatformInfo) -> bool:
    return platform.is_blackwell_plus


def _is_hopper_plus(platform: PlatformInfo) -> bool:
    return platform.is_nvidia and platform.arch_version >= ArchVersion(9, 0)


def _is_nvidia(platform: PlatformInfo) -> bool:
    return platform.is_nvidia


def _is_nvidia_with_cute_dsl(platform: PlatformInfo) -> bool:
    return platform.is_nvidia and _sampling_cute_dsl.is_available()


def _is_cdna4(platform: PlatformInfo) -> bool:
    return platform.is_cdna4


def _is_supported_gpu(platform: PlatformInfo) -> bool:
    return platform.is_nvidia or platform.is_amd


def _fp8_dtype() -> torch.dtype:
    return Platform.get().fp8e4m3fn.dtype


def _mm_dense() -> torch.Tensor:
    a = torch.empty((4, 16), dtype=torch.bfloat16)
    b = torch.empty((32, 16), dtype=torch.bfloat16)
    return tokenspeed_kernel.mm(a, b)


def _mm_dense_gluon_gfx950() -> torch.Tensor:
    a = torch.empty((16, 64), dtype=torch.bfloat16)
    b = torch.empty((128, 64), dtype=torch.bfloat16)
    return tokenspeed_kernel.mm(a, b)


def _mm_mxfp8() -> torch.Tensor:
    a = torch.empty((4, 128), dtype=_fp8_dtype())
    b = torch.empty((128, 128), dtype=_fp8_dtype())
    a_scales = torch.empty((4, 1), dtype=torch.float32)
    b_scales = torch.empty((1, 1), dtype=torch.float32)
    return tokenspeed_kernel.mm(
        a,
        b,
        A_scales=a_scales,
        B_scales=b_scales,
        out_dtype=torch.bfloat16,
        block_size=[128, 128],
        quant="mxfp8",
    )


def test_gemm_mxfp8_online_activation_signature_uses_quantized_storage() -> None:
    a = torch.empty((4, 128), dtype=torch.bfloat16)
    b = torch.empty((128, 128), dtype=_fp8_dtype())
    b_scales = torch.empty((1, 1), dtype=torch.float32)

    signature = _gemm_pkg._gemm_format_signature(
        a,
        b,
        None,
        b_scales,
        torch.bfloat16,
        "mxfp8",
        [128, 128],
    )

    a_format = signature.format_for("a")
    b_format = signature.format_for("b")
    assert a_format is not None
    assert b_format is not None
    assert a_format.storage_dtype == _fp8_dtype()
    assert b_format.storage_dtype == _fp8_dtype()
    assert a_format.scale is not None
    assert b_format.scale is not None
    assert a_format.scale.block_shape == (128, 128)
    assert b_format.scale.block_shape == (128, 128)


def test_gemm_mxfp8_online_activation_preserves_repeated_rows() -> None:
    if not torch.cuda.is_available():
        pytest.skip("CUDA is required for online mxfp8 GEMM verification")
    if not (Platform.get().is_nvidia or Platform.get().is_cdna4):
        pytest.skip("online mxfp8 GEMM verification requires NVIDIA or AMD CDNA4")

    torch.manual_seed(0)
    num_tokens = 16
    hidden_size = 2048
    output_size = 128
    block_size = [128, 128]
    a = torch.randn((1, hidden_size), device="cuda", dtype=torch.bfloat16).repeat(
        num_tokens, 1
    )
    b = (
        torch.randn((output_size, hidden_size), device="cuda", dtype=torch.float32)
        * 0.1
    ).to(_fp8_dtype())
    b_scales = (
        torch.rand(
            (
                (output_size + block_size[0] - 1) // block_size[0],
                (hidden_size + block_size[1] - 1) // block_size[1],
            ),
            device="cuda",
            dtype=torch.float32,
        )
        + 0.01
    )

    out = tokenspeed_kernel.mm(
        a,
        b,
        B_scales=b_scales,
        out_dtype=torch.bfloat16,
        quant="mxfp8",
        block_size=block_size,
    )
    torch.cuda.synchronize()

    torch.testing.assert_close(out[1:], out[:1].expand_as(out[1:]), rtol=0, atol=0)


def test_gemm_fp8_scaled_signature_uses_fp8_format_with_scale() -> None:
    a = torch.empty((4, 128), dtype=_fp8_dtype())
    b = torch.empty((128, 128), dtype=_fp8_dtype())
    a_scales = torch.empty((1,), dtype=torch.float32)
    b_scales = torch.empty((1,), dtype=torch.float32)

    signature = _gemm_pkg._gemm_format_signature(
        a,
        b,
        a_scales,
        b_scales,
        torch.bfloat16,
        "fp8",
        None,
    )

    for role in ("a", "b"):
        tensor_format = signature.format_for(role)
        assert tensor_format is not None
        assert tensor_format.format == "scaled-fp8"
        assert tensor_format.storage_dtype == _fp8_dtype()
        assert tensor_format.scale is not None
        assert tensor_format.scale.granularity == "tensor"
        assert tensor_format.scale.storage_dtype == torch.float32


def test_gemm_fp8_scaled_signature_uses_channel_granularity() -> None:
    a = torch.empty((4, 128), dtype=_fp8_dtype())
    b = torch.empty((128, 128), dtype=_fp8_dtype())
    a_scales = torch.empty((4,), dtype=torch.float32)
    b_scales = torch.empty((128,), dtype=torch.float32)

    signature = _gemm_pkg._gemm_format_signature(
        a,
        b,
        a_scales,
        b_scales,
        torch.bfloat16,
        "fp8",
        None,
    )

    for role in ("a", "b"):
        tensor_format = signature.format_for(role)
        assert tensor_format is not None
        assert tensor_format.scale is not None
        assert tensor_format.scale.granularity == "channel"


def _mm_nvfp4() -> torch.Tensor:
    a = torch.empty((4, 64), dtype=torch.uint8)
    b = torch.empty((128, 64), dtype=torch.uint8)
    a_scales = torch.empty((4, 1), dtype=torch.float32)
    b_scales = torch.empty((128, 1), dtype=torch.float32)
    alpha = torch.empty((), dtype=torch.float32)
    return tokenspeed_kernel.mm(
        a,
        b,
        A_scales=a_scales,
        B_scales=b_scales,
        out_dtype=torch.bfloat16,
        alpha=alpha,
        quant="nvfp4",
    )


def test_gemm_nvfp4_signature_uses_fixed_block_shape() -> None:
    a = torch.empty((4, 64), dtype=torch.uint8)
    b = torch.empty((128, 64), dtype=torch.uint8)
    a_scales = torch.empty((4, 1), dtype=torch.float32)
    b_scales = torch.empty((128, 1), dtype=torch.float32)

    signature = _gemm_pkg._gemm_format_signature(
        a,
        b,
        a_scales,
        b_scales,
        torch.bfloat16,
        "nvfp4",
        None,
    )

    for role in ("a", "b"):
        tensor_format = signature.format_for(role)
        assert tensor_format is not None
        assert tensor_format.scale is not None
        assert tensor_format.scale.block_shape == (16,)


def _attention_prefill() -> object:
    q = torch.empty((4, 16, 64), dtype=torch.bfloat16)
    k = torch.empty((4, 8, 64), dtype=torch.bfloat16)
    v = torch.empty((4, 8, 64), dtype=torch.bfloat16)
    cu_seqlens = torch.tensor([0, 4], dtype=torch.int32)
    return tokenspeed_kernel.mha_prefill(
        q,
        k,
        v,
        cu_seqlens,
        cu_seqlens_cpu=[0, 4],
        max_seqlen=4,
    )


def _attention_extend() -> object:
    q = torch.empty((4, 16, 64), dtype=torch.bfloat16)
    cu_seqlens_q = torch.tensor([0, 2, 4], dtype=torch.int32)
    cu_seqlens_kv = torch.tensor([0, 64, 192], dtype=torch.int32)
    k_cache = torch.empty((8, 64, 8, 64), dtype=torch.bfloat16)
    v_cache = torch.empty((8, 64, 8, 64), dtype=torch.bfloat16)
    page_table = torch.empty((2, 4), dtype=torch.int32)
    cache_seqlens = torch.tensor([64, 128], dtype=torch.int32)
    return tokenspeed_kernel.mha_extend_with_kvcache(
        q,
        cu_seqlens_q,
        cu_seqlens_kv,
        k_cache,
        v_cache,
        page_table,
        cache_seqlens,
        max_seqlen_q=2,
        max_seqlen_k=128,
    )


def _attention_decode() -> object:
    q = torch.empty((2, 16, 64), dtype=torch.bfloat16)
    k_cache = torch.empty((8, 64, 8, 64), dtype=torch.bfloat16)
    v_cache = torch.empty((8, 64, 8, 64), dtype=torch.bfloat16)
    page_table = torch.empty((2, 4), dtype=torch.int32)
    cache_seqlens = torch.tensor([64, 128], dtype=torch.int32)
    return tokenspeed_kernel.mha_decode_with_kvcache(
        q,
        k_cache,
        v_cache,
        page_table,
        cache_seqlens,
        max_seqlen_k=128,
        max_seqlen_q=1,
    )


def _attention_dsa_decode() -> object:
    q = torch.empty((2, 8, 576), dtype=torch.bfloat16)
    sparse_kv_cache = torch.empty((64, 656), dtype=torch.uint8)
    topk_slots = torch.empty((2, 512), dtype=torch.int32)
    topk_lens = torch.empty((2,), dtype=torch.int32)
    return tokenspeed_kernel.dsa_decode(
        q=q,
        kv_cache=None,
        sparse_kv_cache=sparse_kv_cache,
        topk_slots=topk_slots,
        topk_lens=topk_lens,
        max_seqlen_k=64,
        qk_nope_head_dim=192,
        kv_lora_rank=512,
        qk_rope_head_dim=64,
        softmax_scale=1.0,
        page_size=64,
        solution="triton",
    )


def _attention_dsa_prefill() -> object:
    q = torch.empty((2, 8, 576), dtype=torch.bfloat16)
    sparse_kv_cache = torch.empty((64, 656), dtype=torch.uint8)
    topk_slots = torch.empty((2, 512), dtype=torch.int32)
    topk_lens = torch.empty((2,), dtype=torch.int32)
    return tokenspeed_kernel.dsa_prefill(
        q=q,
        kv_cache=None,
        sparse_kv_cache=sparse_kv_cache,
        topk_slots=topk_slots,
        topk_lens=topk_lens,
        max_seqlen_k=64,
        qk_nope_head_dim=192,
        kv_lora_rank=512,
        qk_rope_head_dim=64,
        softmax_scale=1.0,
        page_size=64,
        solution="triton",
    )


def _attention_dsa_decode_topk() -> object:
    q = torch.empty((2, 2, 128), dtype=torch.bfloat16)
    weights = torch.empty((2, 2), dtype=torch.float32)
    index_k = torch.zeros((128, 132), dtype=torch.uint8)
    seq_lens = torch.tensor([64, 64], dtype=torch.int32)
    block_table = torch.zeros((2, 1), dtype=torch.int32)
    return tokenspeed_kernel.dsa_decode_topk(
        q,
        weights,
        seq_lens,
        block_table,
        page_size=64,
        topk=512,
        softmax_scale=1.0,
        index_k_cache=index_k,
    )


def _attention_dsa_prefill_topk() -> object:
    q = torch.empty((2, 2, 128), dtype=torch.bfloat16)
    weights = torch.empty((2, 2), dtype=torch.float32)
    index_k = torch.zeros((128, 132), dtype=torch.uint8)
    kv_workspace_slots = torch.arange(64, dtype=torch.int64)
    row_starts = torch.tensor([0, 8], dtype=torch.int32)
    row_ends = torch.tensor([8, 16], dtype=torch.int32)
    return tokenspeed_kernel.dsa_prefill_topk(
        q,
        weights,
        kv_workspace_slots,
        row_starts,
        row_ends,
        topk=512,
        softmax_scale=1.0,
        index_k_cache=index_k,
        page_size=64,
    )


def _attention_dsa_plan() -> object:
    seq_lens_2d = torch.tensor([[64], [64]], dtype=torch.int32)
    return tokenspeed_kernel.dsa_plan(seq_lens_2d=seq_lens_2d, page_size=64)


def _attention_merge_state() -> object:
    out_a = torch.empty((4, 16, 64), dtype=torch.bfloat16)
    out_b = torch.empty((4, 16, 64), dtype=torch.bfloat16)
    lse_a = torch.empty((4, 16), dtype=torch.float32)
    lse_b = torch.empty((4, 16), dtype=torch.float32)
    return tokenspeed_kernel.attn_merge_state(out_a, lse_a, out_b, lse_b)


def _attention_gdn_chunk_prefill() -> object:
    q = torch.empty((1, 4, 16, 64), dtype=torch.bfloat16)
    k = torch.empty((1, 4, 16, 64), dtype=torch.bfloat16)
    v = torch.empty((1, 4, 16, 64), dtype=torch.bfloat16)
    g = torch.empty((1, 4, 16), dtype=torch.bfloat16)
    beta = torch.empty((1, 4, 16), dtype=torch.bfloat16)
    initial_state = torch.empty((1, 16, 64, 64), dtype=torch.bfloat16)
    cu_seqlens = torch.tensor([0, 4], dtype=torch.int32)
    return tokenspeed_kernel.gdn_chunk_prefill(
        q,
        k,
        v,
        g,
        beta,
        scale=64**-0.5,
        initial_state=initial_state,
        cu_seqlens=cu_seqlens,
        qk_l2norm=True,
        solution="triton",
    )


def _sampling_argmax() -> object:
    logits = torch.empty((4, 4096), dtype=torch.float32, device="cuda")
    return tokenspeed_kernel.argmax(logits)


def _assert_moe_plan(plan: dict, *, apply: str, preprocessor: str | None) -> None:
    assert plan["apply_kernel_name"] == apply
    actual_preprocessor = plan["weight_preprocessor"]
    actual_name = (
        None
        if actual_preprocessor is None
        else getattr(actual_preprocessor, "__name__", repr(actual_preprocessor))
    )
    assert actual_name == preprocessor


def test_gluon_mxfp4_swiglu_args_default_missing_values_to_standard_swiglu() -> None:
    if not hasattr(_moe_gluon_mxfp4, "_swiglu_args"):
        pytest.skip("Gluon MXFP4 SwiGLU args are AMD-only")

    w = torch.nn.Module()
    w.swiglu_arg = type("SwigluArg", (), {"alpha": None, "limit": None})()

    assert _moe_gluon_mxfp4._swiglu_args(w) == (1.0, 0.0, 0.0)

    w.swiglu_arg = type("SwigluArg", (), {"alpha": 1.702, "limit": 7.0})()
    w.swiglu_beta = 1.0

    assert _moe_gluon_mxfp4._swiglu_args(w) == (1.702, 7.0, 1.0)


def _moe_apply_unquant_trtllm() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "unquant",
        input_dtype=torch.bfloat16,
        activation="swiglu",
        requires_deferred_finalize=True,
        ep_size=2,
        ispp=128,
        internal_activation_dtype="input",
    )
    _assert_moe_plan(
        plan,
        apply="flashinfer_trtllm_unquant_moe_apply",
        preprocessor="flashinfer_trtllm_unquant_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(
        plan,
        x,
        torch.nn.Module(),
        router_logits,
        do_finalize=False,
    )


def _moe_apply_unquant_cutlass() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "unquant",
        input_dtype=torch.bfloat16,
        activation="swiglu",
        ep_size=2,
        ispp=128,
        internal_activation_dtype="input",
    )
    _assert_moe_plan(
        plan,
        apply="flashinfer_cutlass_unquant_moe_apply",
        preprocessor="flashinfer_cutlass_unquant_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(plan, x, torch.nn.Module(), router_logits)


def _moe_apply_fp8_cutlass() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "fp8",
        input_dtype=torch.bfloat16,
        activation="silu",
        ep_size=2,
        ispp=128,
        fp8_scale_block_shape=(128, 128),
        internal_activation_dtype="input",
    )
    _assert_moe_plan(
        plan,
        apply="flashinfer_cutlass_fp8_moe_apply",
        preprocessor="flashinfer_cutlass_fp8_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(plan, x, torch.nn.Module(), router_logits)


def _moe_apply_fp8_trtllm() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "fp8",
        input_dtype=torch.bfloat16,
        activation="silu",
        ep_size=2,
        ispp=128,
        fp8_scale_block_shape=(128, 128),
        internal_activation_dtype="input",
    )
    _assert_moe_plan(
        plan,
        apply="flashinfer_trtllm_fp8_moe_apply",
        preprocessor="flashinfer_trtllm_fp8_moe_process_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(plan, x, torch.nn.Module(), router_logits)


def _moe_apply_nvfp4_trtllm() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "nvfp4",
        input_dtype=torch.bfloat16,
        activation="swiglu",
        requires_deferred_finalize=True,
        ep_size=2,
        ispp=128,
        internal_activation_dtype="input",
    )
    _assert_moe_plan(
        plan,
        apply="flashinfer_trtllm_nvfp4_moe_apply",
        preprocessor="flashinfer_trtllm_nvfp4_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(
        plan,
        x,
        torch.nn.Module(),
        router_logits,
        do_finalize=False,
    )


def _moe_apply_nvfp4_cutlass() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "nvfp4",
        input_dtype=torch.bfloat16,
        activation="swiglu",
        ep_size=2,
        ispp=128,
        internal_activation_dtype="input",
        solution="flashinfer_cutlass",
    )
    _assert_moe_plan(
        plan,
        apply="flashinfer_cutlass_nvfp4_moe_apply",
        preprocessor="flashinfer_cutlass_nvfp4_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(plan, x, torch.nn.Module(), router_logits)


def _moe_apply_nvfp4_deepep_cutedsl() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "nvfp4",
        input_dtype=torch.bfloat16,
        activation="swiglu",
        a2a_backend="deepep",
        ep_size=2,
        ispp=128,
        internal_activation_dtype="input",
        deepep_group=object(),
        solution="flashinfer_cutedsl",
    )
    _assert_moe_plan(
        plan,
        apply="flashinfer_cutedsl_deepep_nvfp4_moe_apply",
        preprocessor="flashinfer_cutedsl_deepep_nvfp4_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(plan, x, torch.nn.Module(), router_logits)


def _moe_apply_mxfp4_trtllm() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "mxfp4",
        input_dtype=torch.bfloat16,
        activation="swiglu",
        ep_size=2,
        ispp=128,
        internal_activation_dtype="input",
        with_bias=True,
    )
    _assert_moe_plan(
        plan,
        apply="flashinfer_trtllm_mxfp4_moe_apply",
        preprocessor="flashinfer_trtllm_mxfp4_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(plan, x, torch.nn.Module(), router_logits)


def _moe_apply_mxfp4_triton() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "mxfp4",
        input_dtype=torch.bfloat16,
        activation="swiglu",
        ispp=128,
        internal_activation_dtype="fp8",
        with_bias=True,
    )
    _assert_moe_plan(
        plan,
        apply="triton_mxfp4_moe_apply",
        preprocessor="triton_mxfp4_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(plan, x, torch.nn.Module(), router_logits)


def _moe_apply_mxfp4_gluon() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "mxfp4",
        input_dtype=torch.bfloat16,
        activation="swiglu",
        ispp=128,
        internal_activation_dtype="fp8",
        with_bias=True,
    )
    _assert_moe_plan(
        plan,
        apply="gluon_mxfp4_moe_apply",
        preprocessor="gluon_mxfp4_gfx950_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(plan, x, torch.nn.Module(), router_logits)


def _moe_apply_mxint4_trtllm() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "mxint4",
        input_dtype=torch.bfloat16,
        activation="swiglu",
        ep_size=2,
        ispp=256,
        internal_activation_dtype="input",
    )
    _assert_moe_plan(
        plan,
        apply="flashinfer_trtllm_mxint4_moe_apply",
        preprocessor="flashinfer_trtllm_mxint4_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(plan, x, torch.nn.Module(), router_logits)


def _moe_apply_mxfp4_dynamic_tp() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "mxfp4",
        input_dtype=torch.bfloat16,
        activation="silu",
        ep_size=1,
        ispp=2048,
        internal_activation_dtype="input",
    )
    _assert_moe_plan(
        plan,
        apply="gluon_mxfp4_dynamic_moe_apply",
        preprocessor="gluon_mxfp4_gfx950_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    return tokenspeed_kernel.moe_apply(
        plan,
        x,
        torch.nn.Module(),
        router_logits,
    )


def _moe_apply_fp8_precomputed_tp() -> object:
    plan = _moe_pkg.moe_plan(
        "fp8",
        input_dtype=torch.bfloat16,
        activation="silu",
        ep_size=1,
        fp8_scale_block_shape=(128, 128),
        solution="triton",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    topk_weights = torch.empty((4, 2), dtype=torch.float32)
    topk_ids = torch.empty((4, 2), dtype=torch.int64)
    return tokenspeed_kernel.moe_apply(
        plan,
        x,
        torch.nn.Module(),
        router_logits,
        topk_weights=topk_weights,
        topk_ids=topk_ids,
    )


def _moe_apply_fp8_precomputed_ep() -> object:
    plan = _moe_pkg.moe_plan(
        "fp8",
        input_dtype=torch.bfloat16,
        activation="silu",
        ep_size=2,
        fp8_scale_block_shape=(128, 128),
        solution="triton",
    )
    _assert_moe_plan(
        plan,
        apply="triton_fp8_ep_precomputed_moe_apply",
        preprocessor="triton_fp8_moe_weights",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    topk_weights = torch.empty((4, 2), dtype=torch.float32)
    topk_ids = torch.empty((4, 2), dtype=torch.int64)
    return tokenspeed_kernel.moe_apply(
        plan,
        x,
        torch.nn.Module(),
        router_logits,
        topk_weights=topk_weights,
        topk_ids=topk_ids,
    )


def _moe_apply_mxfp4_precomputed_ep() -> object:
    plan = tokenspeed_kernel.moe_plan(
        "mxfp4",
        input_dtype=torch.bfloat16,
        activation="silu",
        ep_size=4,
        internal_activation_dtype="input",
    )
    x = torch.empty((4, 16), dtype=torch.bfloat16)
    router_logits = torch.empty((4, 8), dtype=torch.float32)
    topk_weights = torch.empty((4, 2), dtype=torch.float32)
    topk_ids = torch.empty((4, 2), dtype=torch.int64)
    return tokenspeed_kernel.moe_apply(
        plan,
        x,
        torch.nn.Module(),
        router_logits,
        topk_weights=topk_weights,
        topk_ids=topk_ids,
    )


def test_mxfp4_ep_topk_localization_masks_remote_experts() -> None:
    w = torch.nn.Module()
    w.num_experts = 8
    w.num_local_experts = 2
    w.ep_rank = 2
    w.ep_size = 4
    topk_weights = torch.tensor(
        [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]],
        dtype=torch.float32,
    )
    topk_ids = torch.tensor([[0, 4, 5], [6, 7, 3]], dtype=torch.int64)

    local_weights, local_ids, num_experts = _moe_triton_mxfp4._local_topk_for_ep(
        topk_weights,
        topk_ids,
        w,
    )

    assert num_experts == 2
    torch.testing.assert_close(
        local_weights,
        torch.tensor([[0.0, 0.2, 0.3], [0.0, 0.0, 0.0]], dtype=torch.float32),
    )
    assert torch.equal(
        local_ids,
        torch.tensor([[-1, 0, 1], [-1, -1, -1]], dtype=torch.int64),
    )


def _case(
    matches: Callable[[PlatformInfo], bool],
    arch: str,
    family: str,
    mode: str,
    expected: str,
    invoke: Callable[[], object],
) -> KernelApiSelectionCase:
    return KernelApiSelectionCase(
        id=f"{arch}/{family}.{mode}/{expected}",
        arch=arch,
        family=family,
        mode=mode,
        expected=expected,
        matches=matches,
        invoke=invoke,
    )


_CASES = [
    # Attention API x architecture golden cases.
    _case(
        _is_hopper,
        "hopper",
        "attention",
        "mha_prefill",
        "fa3_mha_prefill",
        _attention_prefill,
    ),
    _case(
        _is_hopper,
        "hopper",
        "attention",
        "mha_extend_with_kvcache",
        "fa3_mha_extend_with_kvcache_cached",
        _attention_extend,
    ),
    _case(
        _is_hopper,
        "hopper",
        "attention",
        "mha_decode_with_kvcache",
        "fa3_mha_decode_with_kvcache_cached",
        _attention_decode,
    ),
    _case(
        _is_hopper,
        "hopper",
        "attention",
        "attn_merge_state",
        "cuda_attn_merge_state",
        _attention_merge_state,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "attention",
        "mha_prefill",
        "fa4_mha_prefill",
        _attention_prefill,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "attention",
        "mha_extend_with_kvcache",
        "fa4_mha_extend_with_kvcache_cached",
        _attention_extend,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "attention",
        "mha_decode_with_kvcache",
        "fa4_mha_decode_with_kvcache",
        _attention_decode,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "attention",
        "attn_merge_state",
        "cuda_attn_merge_state",
        _attention_merge_state,
    ),
    _case(
        _is_blackwell_non_sm100,
        "blackwell-non-sm100",
        "attention",
        "mha_extend_with_kvcache",
        "flashinfer_trtllm_mha_extend_with_kvcache",
        _attention_extend,
    ),
    _case(
        _is_blackwell_non_sm100,
        "blackwell-non-sm100",
        "attention",
        "mha_decode_with_kvcache",
        "flashinfer_trtllm_mha_decode_with_kvcache",
        _attention_decode,
    ),
    _case(
        _is_blackwell_non_sm100,
        "blackwell-non-sm100",
        "attention",
        "attn_merge_state",
        "cuda_attn_merge_state",
        _attention_merge_state,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "attention",
        "mha_prefill",
        "gluon_mha_prefill_gfx950",
        _attention_prefill,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "attention",
        "mha_extend_with_kvcache",
        "gluon_mha_extend_gfx950",
        _attention_extend,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "attention",
        "mha_decode_with_kvcache",
        "gluon_mha_decode_gfx950",
        _attention_decode,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "attention",
        "attn_merge_state",
        "triton_attn_merge_state",
        _attention_merge_state,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "attention",
        "dsa_decode",
        "triton_dsa_decode",
        _attention_dsa_decode,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "attention",
        "dsa_prefill",
        "triton_dsa_prefill",
        _attention_dsa_prefill,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "attention",
        "dsa_decode_topk",
        "triton_dsa_decode_topk_fp8",
        _attention_dsa_decode_topk,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "attention",
        "dsa_prefill_topk",
        "triton_dsa_prefill_topk_fp8",
        _attention_dsa_prefill_topk,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "attention",
        "dsa_plan",
        "triton_dsa_plan",
        _attention_dsa_plan,
    ),
    _case(
        _is_supported_gpu,
        "supported-gpu",
        "attention",
        "gdn_chunk_prefill",
        "triton_gdn_chunk_prefill",
        _attention_gdn_chunk_prefill,
    ),
    # GEMM API x architecture golden cases.
    _case(_is_supported_gpu, "supported-gpu", "gemm", "mm", "torch_mm", _mm_dense),
    _case(
        _is_cdna4,
        "cdna4",
        "gemm",
        "mm",
        "gluon_mm_a16w16_gfx950",
        _mm_dense_gluon_gfx950,
    ),
    _case(
        _is_hopper,
        "hopper",
        "gemm",
        "mm",
        "deep_gemm_mm_fp8_blockscale",
        _mm_mxfp8,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "gemm",
        "mm",
        "flashinfer_mm_fp8_blockscale",
        _mm_mxfp8,
    ),
    _case(
        _is_blackwell_plus,
        "blackwell-plus",
        "gemm",
        "mm",
        "cublaslt_mm_nvfp4",
        _mm_nvfp4,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "gemm",
        "mm",
        "triton_mm_fp8_blockscale",
        _mm_mxfp8,
    ),
    # Sampling API x architecture golden cases.
    _case(
        _is_nvidia_with_cute_dsl,
        "nvidia-cutedsl",
        "sampling",
        "argmax",
        "cute_dsl_argmax",
        _sampling_argmax,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "sampling",
        "argmax",
        "gluon_argmax_gfx950",
        _sampling_argmax,
    ),
    # MoE API x architecture golden cases.
    _case(
        _is_hopper,
        "hopper",
        "moe",
        "apply",
        "flashinfer_cutlass_unquant_moe_apply",
        _moe_apply_unquant_cutlass,
    ),
    _case(
        _is_hopper,
        "hopper",
        "moe",
        "apply",
        "flashinfer_cutlass_fp8_moe_apply",
        _moe_apply_fp8_cutlass,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "moe",
        "apply",
        "flashinfer_trtllm_fp8_moe_apply",
        _moe_apply_fp8_trtllm,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "moe",
        "apply",
        "flashinfer_trtllm_unquant_moe_apply",
        _moe_apply_unquant_trtllm,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "moe",
        "apply",
        "flashinfer_trtllm_nvfp4_moe_apply",
        _moe_apply_nvfp4_trtllm,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "moe",
        "apply",
        "flashinfer_cutlass_nvfp4_moe_apply",
        _moe_apply_nvfp4_cutlass,
    ),
    _case(
        _is_blackwell_plus,
        "blackwell-plus",
        "moe",
        "apply",
        "flashinfer_cutedsl_deepep_nvfp4_moe_apply",
        _moe_apply_nvfp4_deepep_cutedsl,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "moe",
        "apply",
        "flashinfer_trtllm_mxfp4_moe_apply",
        _moe_apply_mxfp4_trtllm,
    ),
    _case(
        _is_blackwell_sm100,
        "blackwell-sm100",
        "moe",
        "apply",
        "flashinfer_trtllm_mxint4_moe_apply",
        _moe_apply_mxint4_trtllm,
    ),
    _case(
        _is_hopper,
        "hopper",
        "moe",
        "apply",
        "triton_mxfp4_moe_apply",
        _moe_apply_mxfp4_triton,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "moe",
        "apply",
        "gluon_mxfp4_moe_apply",
        _moe_apply_mxfp4_gluon,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "moe",
        "apply",
        "gluon_mxfp4_dynamic_moe_apply",
        _moe_apply_mxfp4_dynamic_tp,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "moe",
        "apply",
        "triton_mxfp4_ep_precomputed_moe_apply",
        _moe_apply_mxfp4_precomputed_ep,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "moe",
        "apply",
        "triton_fp8_ep_precomputed_moe_apply",
        _moe_apply_fp8_precomputed_tp,
    ),
    _case(
        _is_cdna4,
        "cdna4",
        "moe",
        "apply",
        "triton_fp8_ep_precomputed_moe_apply",
        _moe_apply_fp8_precomputed_ep,
    ),
]


@pytest.fixture
def selected_kernel_spy(monkeypatch):
    active_case: dict[str, KernelApiSelectionCase | None] = {"case": None}
    calls: list[str] = []

    def fake_call(self: SelectedKernel, *args, **kwargs):
        case = active_case["case"]
        assert case is not None, "selected_kernel_spy used without an active case"
        calls.append(self.name)

        if case.family == "gemm":
            a, b, _a_scales, _b_scales, out_dtype = args[:5]
            n = b.shape[-1] if b.shape[0] == a.shape[-1] else b.shape[0]
            return torch.empty((a.shape[0], n), dtype=out_dtype, device=a.device)

        if case.family == "attention":
            if case.mode == "attn_merge_state":
                return torch.empty_like(kwargs["out_a"]), torch.empty_like(
                    kwargs["lse_a"]
                )
            if case.mode == "dsa_plan":
                return torch.empty((1, 4), dtype=torch.int32)
            q = kwargs["q"]
            if case.mode == "gdn_chunk_prefill":
                return GdnChunkPrefillResult(
                    out=torch.empty_like(q),
                    final_state=kwargs.get("initial_state"),
                )
            if kwargs.get("return_lse", False):
                lse = torch.empty(q.shape[:-1], dtype=torch.float32, device=q.device)
                return torch.empty_like(q), lse
            return torch.empty_like(q)

        if case.family == "sampling":
            (logits,) = args[:1]
            out = kwargs.get("out")
            if out is not None:
                return out
            return torch.empty(
                (logits.shape[0],), dtype=torch.int64, device=logits.device
            )

        if case.family == "moe":
            return torch.empty_like(kwargs["x"])

        return None

    monkeypatch.setattr(SelectedKernel, "__call__", fake_call)
    return active_case, calls


def _find_case(*, arch: str, family: str, mode: str) -> KernelApiSelectionCase:
    for case in _CASES:
        if case.arch == arch and case.family == family and case.mode == mode:
            return case
    raise AssertionError(f"missing golden case for {arch}/{family}.{mode}")


def test_attn_merge_state_routes_to_triton_on_cdna4(
    mi350_platform: PlatformInfo,
    selected_kernel_spy,
) -> None:
    case = _find_case(arch="cdna4", family="attention", mode="attn_merge_state")
    registry = KernelRegistry.get()
    expected_spec = registry.get_by_name(case.expected)
    assert expected_spec is not None
    assert expected_spec.capability.satisfied_by(mi350_platform)

    real_platform = Platform.get()
    active_case, calls = selected_kernel_spy
    active_case["case"] = case
    try:
        Platform.override(mi350_platform)
        registry.clear_cache()

        case.invoke()

        assert calls == ["triton_attn_merge_state"]
    finally:
        Platform.override(real_platform)
        registry.clear_cache()


@pytest.mark.parametrize("case", _CASES, ids=lambda case: case.id)
def test_kernel_api_selection(case: KernelApiSelectionCase, selected_kernel_spy):
    platform = Platform.get()
    if not case.matches(platform):
        pytest.skip(
            f"{case.id} only applies to its {case.arch} architecture case; "
            f"current platform is {platform.device_name} ({platform.arch_version})"
        )

    registry = KernelRegistry.get()
    expected_spec = registry.get_by_name(case.expected)
    assert expected_spec is not None, (
        f"{case.expected!r} is not registered on "
        f"{platform.device_name} ({platform.arch_version})"
    )
    assert expected_spec.capability.satisfied_by(platform), (
        f"{case.expected!r} is registered but not compatible with "
        f"{platform.device_name} ({platform.arch_version})"
    )

    active_case, calls = selected_kernel_spy
    active_case["case"] = case
    registry.clear_cache()

    case.invoke()

    assert calls == [case.expected]