[WIP] test(kernel): verify kernel selection cross-platform and drop expected_kernel_name

python/tokenspeed/runtime/layers/dense/mxfp4.py

@@ -113,7 +113,6 @@ def apply(self, layer, x, bias=None):
             bias=bias,
             out_dtype=x.dtype,
             quant="mxfp4",
-            expected_kernel_name="triton_mm_mxfp4",
         )
         return output.reshape(*output_shape)
 

python/tokenspeed/runtime/layers/dense/nvfp4.py

@@ -209,7 +209,6 @@ def apply(self, layer, x, bias=None):
             quant="nvfp4",
             enable_pdl=_pdl_enabled(),
             override=kernel_override,
-            expected_kernel_name=kernel_override or "cublaslt_mm_nvfp4",
         )
         return out.view(x_fp4.size(0), w_n)
 

tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/flash_attn/__init__.py

@@ -72,6 +72,7 @@
         solution="fa4",
         capability=CapabilityRequirement(
             min_arch_version=ArchVersion(10, 0),
+            max_arch_version=ArchVersion(10, 0),
             vendors=frozenset({"nvidia"}),
         ),
         signatures=format_signatures(
@@ -121,6 +122,7 @@ def fa4_mha_prefill(
         solution="fa4",
         capability=CapabilityRequirement(
             min_arch_version=ArchVersion(10, 0),
+            max_arch_version=ArchVersion(10, 0),
             vendors=frozenset({"nvidia"}),
         ),
         signatures=format_signatures(
@@ -176,6 +178,7 @@ def fa4_mha_extend_with_kvcache(
         solution="fa4",
         capability=CapabilityRequirement(
             min_arch_version=ArchVersion(10, 0),
+            max_arch_version=ArchVersion(10, 0),
             vendors=frozenset({"nvidia"}),
         ),
         signatures=format_signatures(
@@ -233,6 +236,7 @@ def fa4_mha_decode_with_kvcache(
         solution="fa3",
         capability=CapabilityRequirement(
             min_arch_version=ArchVersion(9, 0),
+            max_arch_version=ArchVersion(9, 0),
             vendors=frozenset({"nvidia"}),
         ),
         signatures=format_signatures(
@@ -280,6 +284,7 @@ def fa3_mha_prefill(
         solution="fa3",
         capability=CapabilityRequirement(
             min_arch_version=ArchVersion(9, 0),
+            max_arch_version=ArchVersion(9, 0),
             vendors=frozenset({"nvidia"}),
         ),
         signatures=format_signatures(
@@ -333,6 +338,7 @@ def fa3_mha_extend_with_kvcache(
         solution="fa3",
         capability=CapabilityRequirement(
             min_arch_version=ArchVersion(9, 0),
+            max_arch_version=ArchVersion(9, 0),
             vendors=frozenset({"nvidia"}),
         ),
         signatures=format_signatures(

tokenspeed-kernel/python/tokenspeed_kernel/ops/gemm/__init__.py

@@ -258,7 +258,6 @@ def mm(
     quant: str | None = None,
     enable_pdl: bool = False,
     override: str | None = None,
-    expected_kernel_name: str | None = None,
 ) -> torch.Tensor:
     """Dense matrix multiply with automatic kernel selection.
 
@@ -285,7 +284,6 @@ def mm(
             If ``None``, inferred from input dtypes and scales.
         override: Force selection of a specific kernel by name (e.g.
             ``"cublaslt_mm_nvfp4"``). Bypasses heuristic scoring.
-        expected_kernel_name: Debug hint for expected kernel selection.
     """
     out_dtype = out_dtype or A.dtype
 
@@ -317,7 +315,6 @@ def mm(
         signature,
         traits=traits,
         override=override,
-        expected_kernel_name=expected_kernel_name,
     )
 
     # Online activation quantization

tokenspeed-kernel/python/tokenspeed_kernel/selection.py

@@ -577,7 +577,6 @@ def select_kernel(
     traits: dict[str, Any] | None = None,
     solution: str | None = None,
     override: str | None = None,
-    expected_kernel_name: str | None = None,
 ) -> SelectedKernel:
     """Select the best kernel for an operation.
 
@@ -597,9 +596,6 @@ def select_kernel(
         solution: Restrict selection to a registered solution while preserving
             normal platform, format signature, and trait filtering.
         override: Force a specific kernel name or solution string
-        expected_kernel_name: Debug-only hint. When set, a warning is
-            logged if the selected kernel differs from this name. The
-            selected kernel is still used regardless of the mismatch.
 
     Returns:
         A :class:`SelectedKernel` that is directly callable and also
@@ -706,18 +702,6 @@ def select_kernel(
 
     _log_selection(family, mode, format_signature, winner, scored, platform, objective)
 
-    if expected_kernel_name and winner.name != expected_kernel_name:
-        logger.warning(
-            "[tokenspeed_kernel] select_kernel(%s.%s, %s) chose '%s' but "
-            "expected '%s'. Score breakdown — selected: %s",
-            family,
-            mode,
-            format_signature,
-            winner.name,
-            expected_kernel_name,
-            next((s for sp, s in scored if sp.name == winner.name), "N/A"),
-        )
-
     impl = registry.get_impl(winner.name)
     result = SelectedKernel(name=winner.name, impl=impl)
     registry.cache_put(cache_key, result)

tokenspeed-kernel/test/conftest.py

@@ -197,6 +197,34 @@ def b200_platform() -> PlatformInfo:
     )
 
 
+@pytest.fixture
+def b300_platform() -> PlatformInfo:
+    return PlatformInfo(
+        vendor="nvidia",
+        arch_version=ArchVersion(10, 3),
+        device_name="NVIDIA B300",
+        device_count=8,
+        total_memory=288 * (1024**3),
+        memory_bandwidth=8000.0,
+        sm_count=160,
+        max_threads_per_sm=2048,
+        max_shared_memory_per_sm=262144,
+        sm_features=frozenset(
+            {
+                "tensor_core:f16",
+                "tensor_core:int8",
+                "tensor_core:f8",
+                "tensor_core:f4",
+                "memory:async_copy",
+                "memory:tma",
+                "compute:cluster",
+            }
+        ),
+        runtime_features=frozenset({"runtime:cuda_graph"}),
+        interconnect=InterconnectInfo(topology="nvlink_full"),
+    )
+
+
 @pytest.fixture
 def fresh_registry():
     KernelRegistry.reset()

tokenspeed-kernel/test/ops/test_nvfp4_gemm_swiglu_nvfp4_quant.py

@@ -207,7 +207,6 @@ def test_nvfp4_gemm_swiglu_nvfp4_quant_matches_unfused_model_shapes(
         alpha=fc1_alpha,
         quant="nvfp4",
         enable_pdl=True,
-        expected_kernel_name="flashinfer_mm_nvfp4",
     ).view(m, 2 * i)
 
     silu_out = (
@@ -247,7 +246,6 @@ def test_nvfp4_gemm_swiglu_nvfp4_quant_matches_unfused_model_shapes(
         alpha=fc2_alpha,
         quant="nvfp4",
         enable_pdl=True,
-        expected_kernel_name="flashinfer_mm_nvfp4",
     ).view(m, k)
     actual = tokenspeed_kernel.mm(
         fused_fp4,
@@ -258,7 +256,6 @@ def test_nvfp4_gemm_swiglu_nvfp4_quant_matches_unfused_model_shapes(
         alpha=fc2_alpha,
         quant="nvfp4",
         enable_pdl=True,
-        expected_kernel_name="flashinfer_mm_nvfp4",
     ).view(m, k)
     torch.cuda.synchronize()