fix(linear): repair B12X kernel registry after the b12x_block rebase

tests/model_executor/kernels/test_b12x_linear.py

@@ -16,51 +16,37 @@
     _POSSIBLE_FP8_KERNELS,
     _POSSIBLE_MXFP4_KERNELS,
     _POSSIBLE_MXFP8_KERNELS,
+    _POSSIBLE_NVFP4_KERNELS,
+    B12xFp8BlockScaledMMKernel,
+    B12xMxFp4LinearKernel,
+    B12xNvFp4LinearKernel,
+    B12xTensorFP8ScaledMMLinearKernel,
     init_fp8_linear_kernel,
+    init_mxfp4_linear_kernel,
     init_mxfp8_linear_kernel,
+    init_nvfp4_linear_kernel,
 )
 from vllm.model_executor.kernels.linear.mxfp8.b12x import (
     B12xMxfp8LinearKernel,
     _b12x_mxfp8_expected_m,
-
     warmup_b12x_mxfp8_linear,
 )
 from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
     Mxfp8LinearLayerConfig,
 )
-from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import (
-    B12xFp8BlockScaledMMKernel,
-    B12xMxFp4LinearKernel,
-    B12xMxfp8LinearKernel,
-    B12xNvFp4LinearKernel,
-    B12xTensorFP8ScaledMMLinearKernel,
-    FP8ScaledMMLinearLayerConfig,
-    Mxfp8LinearLayerConfig,
-    init_fp8_linear_kernel,
-    init_mxfp4_linear_kernel,
-    init_mxfp8_linear_kernel,
-    init_nvfp4_linear_kernel,
-)
 from vllm.model_executor.kernels.linear.nvfp4.marlin import (
     MarlinNvFp4LinearKernel,
 )
+from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
+    FP8ScaledMMLinearLayerConfig,
+)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kFp8Dynamic128Sym,
     kFp8Static128BlockSym,
     kFp8StaticTensorSym,
     kMxfp4Dynamic,
 )
 from vllm.platforms import PlatformEnum
-from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import (
-    B12xTensorFP8ScaledMMLinearKernel,
-    _b12x_tensor_fp8_linear,
-    warmup_b12x_tensor_fp8_linear,
-)
-from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
-    FP8ScaledMMLinearLayerConfig,
-)
-from vllm.platforms import PlatformEnum, current_platform
-from vllm.utils.b12x import b12x_warmup_token_counts
 
 
 def test_b12x_backend_maps_mxfp8_kernel() -> None:
@@ -162,7 +148,6 @@ def test_b12x_backend_registration_priority_and_selection(
     assert isinstance(initializer(**kwargs), kernel_cls)
 
 
-
 def test_b12x_tensor_fp8_can_implement_supported_config() -> None:
     config = FP8ScaledMMLinearLayerConfig(
         activation_quant_key=kFp8StaticTensorSym,
@@ -371,6 +356,8 @@ def mm(
     assert called_bias is bias
     assert out_dtype == torch.bfloat16
     assert expected_m == 6
+
+
 def test_b12x_mxfp8_explicit_backend_selects_kernel(monkeypatch) -> None:
     import vllm.model_executor.kernels.linear as linear_mod
 
@@ -412,7 +399,6 @@ def test_b12x_mxfp8_expected_m_uses_live_m() -> None:
     assert _b12x_mxfp8_expected_m(2048) == 2048
 
 
-
 def test_warmup_b12x_mxfp8_linear_dedupes_weight_signatures(
     monkeypatch,
 ) -> None:

tests/model_executor/kernels/test_b12x_mxfp4_linear.py

@@ -14,7 +14,6 @@
 )
 from vllm.model_executor.kernels.linear.mxfp4.b12x import (
     B12xMxFp4LinearKernel,
-    warmup_b12x_mxfp4_linear,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     kMxfp4Dynamic,
@@ -37,6 +36,7 @@ def test_b12x_mxfp4_fallback_priority() -> None:
         < names.index("EmulationMxfp4LinearKernel")
     )
 
+
 def test_b12x_mxfp4_explicit_backend_selects_native_kernel(monkeypatch) -> None:
     import vllm.model_executor.kernels.linear as linear_mod
 
@@ -94,7 +94,9 @@ def test_b12x_mxfp4_processes_scale_and_preserves_loader(monkeypatch) -> None:
 
     assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr()
     assert layer.weight_scale.weight_loader is weight_loader
-    assert layer.b12x_mxfp4_linear
+    # #52368 replaced the b12x_mxfp4_linear marker with the warmup provider
+    # that b12x_warmup discovers.
+    assert layer.b12x_warmup_provider is kernel
 
 
 def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None:
@@ -143,7 +145,13 @@ def mm_mxfp4(*args, **kwargs):
     assert kwargs == {"out_dtype": torch.bfloat16}
 
 
-def test_warmup_b12x_mxfp4_dedupes_weight_signatures(monkeypatch) -> None:
+def test_b12x_mxfp4_warmup_units_dedupe_weight_signatures(monkeypatch) -> None:
+    """Upstream #52368 replaced the standalone warmup with a unit provider.
+
+    ``b12x_warmup`` keys units by ``unit.key`` and compiles each once, so the
+    dedupe contract now lives in the key: two layers with the same weight
+    signature must produce equal keys, and a different one must not.
+    """
     import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod
 
     calls = []
@@ -152,16 +160,6 @@ def apply(source, weight, weight_scale, bias):
         calls.append((source.shape, weight, weight_scale, bias))
         return source.new_empty((source.shape[0], weight.shape[0]))
 
-    platform = types.SimpleNamespace(
-        is_cuda=lambda: True,
-        is_device_capability_family=lambda family: family == 120,
-    )
-    monkeypatch.setattr(b12x_mod, "current_platform", platform)
-    monkeypatch.setattr(
-        b12x_mod,
-        "_import_b12x_blockscaled",
-        lambda: types.SimpleNamespace(),
-    )
     monkeypatch.setattr(b12x_mod, "_apply_b12x_mxfp4_linear", apply)
 
     def layer(n: int):
@@ -174,17 +172,22 @@ def layer(n: int):
     layer_a = layer(48)
     layer_b = layer(48)
     layer_c = layer(96)
-    model = types.SimpleNamespace(
-        modules=lambda: iter([layer_a, layer_b, layer_c, types.SimpleNamespace()])
-    )
 
-    warmed = warmup_b12x_mxfp4_linear(
-        model,
-        max_tokens=8,
-        cudagraph_capture_sizes=[1, 2],
-    )
+    kernel = object.__new__(B12xMxFp4LinearKernel)
+    token_counts = (1, 2, 8)
+    units = {
+        name: kernel.get_b12x_warmup_unit(obj, token_counts, torch.bfloat16)
+        for name, obj in (("a", layer_a), ("b", layer_b), ("c", layer_c))
+    }
+
+    assert units["a"].key == units["b"].key
+    assert units["a"].key != units["c"].key
+    assert {unit.name for unit in units.values()} == {"MXFP4"}
+
+    # b12x_warmup compiles one unit per distinct key: a (== b) and c.
+    units["a"].compile()
+    units["c"].compile()
 
-    assert warmed == 6
     assert [call[0] for call in calls] == [
         torch.Size([1, 128]),
         torch.Size([2, 128]),

tests/model_executor/kernels/test_b12x_nvfp4_linear.py

@@ -15,15 +15,13 @@
 )
 from vllm.model_executor.kernels.linear.nvfp4.b12x import (
     B12xNvFp4LinearKernel,
-    warmup_b12x_nvfp4_linear,
 )
 from vllm.model_executor.kernels.linear.nvfp4.marlin import (
     MarlinNvFp4LinearKernel,
 )
 from vllm.platforms import PlatformEnum
 
 
-
 def test_b12x_backend_maps_nvfp4_kernel() -> None:
     assert B12xNvFp4LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"]
     assert B12xNvFp4LinearKernel in _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA]
@@ -39,6 +37,7 @@ def test_b12x_nvfp4_fallback_priority() -> None:
         < names.index("EmulationNvFp4LinearKernel")
     )
 
+
 def test_b12x_nvfp4_explicit_backend_selects_native_kernel(monkeypatch) -> None:
     import vllm.model_executor.kernels.linear as linear_mod
 
@@ -110,7 +109,9 @@ def test_b12x_nvfp4_processes_scale_and_preserves_loader(monkeypatch) -> None:
 
     assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr()
     assert layer.weight_scale.weight_loader is weight_loader
-    assert layer.b12x_nvfp4_linear
+    # #52368 replaced the b12x_nvfp4_linear marker with the warmup provider
+    # that b12x_warmup discovers.
+    assert layer.b12x_warmup_provider is kernel
 
 
 def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None:
@@ -173,7 +174,13 @@ def mm_nvfp4(*args, **kwargs):
     assert kwargs == {"out_dtype": torch.bfloat16}
 
 
-def test_warmup_b12x_nvfp4_dedupes_weight_signatures(monkeypatch) -> None:
+def test_b12x_nvfp4_warmup_units_dedupe_weight_signatures(monkeypatch) -> None:
+    """Upstream #52368 replaced the standalone warmup with a unit provider.
+
+    ``b12x_warmup`` keys units by ``unit.key`` and compiles each once, so the
+    dedupe contract now lives in the key: two layers with the same weight
+    signature must produce equal keys, and a different one must not.
+    """
     import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod
 
     calls = []
@@ -198,16 +205,6 @@ def apply(
         )
         return source.new_empty((source.shape[0], weight.shape[0]))
 
-    platform = types.SimpleNamespace(
-        is_cuda=lambda: True,
-        is_device_capability_family=lambda family: family == 120,
-    )
-    monkeypatch.setattr(b12x_mod, "current_platform", platform)
-    monkeypatch.setattr(
-        b12x_mod,
-        "_import_b12x_blockscaled",
-        lambda: types.SimpleNamespace(),
-    )
     monkeypatch.setattr(b12x_mod, "_apply_b12x_nvfp4_linear", apply)
 
     def layer(n: int):
@@ -222,17 +219,21 @@ def layer(n: int):
     layer_a = layer(48)
     layer_b = layer(48)
     layer_c = layer(96)
-    model = types.SimpleNamespace(
-        modules=lambda: iter([layer_a, layer_b, layer_c, types.SimpleNamespace()])
-    )
 
-    warmed = warmup_b12x_nvfp4_linear(
-        model,
-        max_tokens=8,
-        cudagraph_capture_sizes=[1, 2],
-    )
+    kernel = object.__new__(B12xNvFp4LinearKernel)
+    token_counts = (1, 2, 8)
+    units = {
+        name: kernel.get_b12x_warmup_unit(obj, token_counts, torch.bfloat16)
+        for name, obj in (("a", layer_a), ("b", layer_b), ("c", layer_c))
+    }
+
+    assert units["a"].key == units["b"].key
+    assert units["a"].key != units["c"].key
+
+    # b12x_warmup compiles one unit per distinct key: a (== b) and c.
+    units["a"].compile()
+    units["c"].compile()
 
-    assert warmed == 6
     assert [call[0] for call in calls] == [
         torch.Size([1, 128]),
         torch.Size([2, 128]),

tests/model_executor/test_deepseek_v4_mhc_warmup.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Token-size selection for the DeepSeek V4 mHC TileLang warmup.
+
+The tilelang mHC kernels derive their split-k factor from two different K
+dimensions: ``mhc_pre_tilelang`` and ``mhc_fused_post_pre_tilelang`` use
+``hc_mult * hidden_size``, while ``mhc_pre_broadcast_tilelang`` (the b12x mHC
+path) uses ``hidden_size`` alone. Warmup sweeps only the larger K, which is
+sufficient because ``split_k = min(num_sms // grid, K // 64 // 4)`` is
+non-increasing in ``grid``: the smaller K lowers the cap, which merges
+plateaus rather than creating new ones, so the broadcast transitions are a
+subset of the swept ones. These tests pin that relationship, so a future
+change to either the sweep or a kernel's K dimension cannot silently leave a
+kernel to JIT-compile during serving.
+"""
+
+from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
+    _compute_mhc_pre_num_split,
+    _select_mhc_warmup_token_sizes,
+)
+
+HIDDEN_SIZE = 2048
+HC_MULT = 4
+NUM_SMS = 48
+MAX_TOKENS = 16_384
+
+
+def _transitions(*, hidden_size: int, hc_mult: int, num_sms: int) -> list[int]:
+    """Token counts where split-k changes, computed independently of warmup."""
+    sizes: list[int] = []
+    last = None
+    grid = 1
+    while (size := (grid - 1) * 64 + 1) <= MAX_TOKENS:
+        split = _compute_mhc_pre_num_split(
+            num_tokens=size,
+            hidden_size=hidden_size,
+            hc_mult=hc_mult,
+            num_sms=num_sms,
+        )
+        if split != last:
+            sizes.append(size)
+            last = split
+        grid += 1
+    return sizes
+
+
+def test_broadcast_transitions_are_subset_of_swept_transitions():
+    """Broadcast K (hc_mult=1) must not need transitions the sweep misses."""
+    swept = set(_transitions(hidden_size=HIDDEN_SIZE, hc_mult=HC_MULT, num_sms=NUM_SMS))
+    broadcast = set(_transitions(hidden_size=HIDDEN_SIZE, hc_mult=1, num_sms=NUM_SMS))
+    assert broadcast <= swept
+
+
+def test_selected_sizes_cover_both_kernel_families():
+    selected = set(
+        _select_mhc_warmup_token_sizes(
+            max_tokens=MAX_TOKENS,
+            cudagraph_capture_sizes=[],
+            hidden_size=HIDDEN_SIZE,
+            hc_mult=HC_MULT,
+            num_sms=NUM_SMS,
+        )
+    )
+    for hc_mult in (HC_MULT, 1):
+        missing = [
+            size
+            for size in _transitions(
+                hidden_size=HIDDEN_SIZE, hc_mult=hc_mult, num_sms=NUM_SMS
+            )
+            if size not in selected
+        ]
+        assert not missing, f"hc_mult={hc_mult} transitions uncovered: {missing}"
+
+
+def test_small_fma_decode_sizes_are_warmed():
+    """mhc_fused_post_pre_tilelang takes a distinct grid for num_tokens < 8."""
+    selected = set(
+        _select_mhc_warmup_token_sizes(
+            max_tokens=MAX_TOKENS,
+            cudagraph_capture_sizes=[],
+            hidden_size=HIDDEN_SIZE,
+            hc_mult=HC_MULT,
+            num_sms=NUM_SMS,
+        )
+    )
+    assert {1, 2, 4}.issubset(selected)
+    assert 8 in selected
+
+
+def test_selection_respects_max_tokens_and_includes_capture_sizes():
+    selected = _select_mhc_warmup_token_sizes(
+        max_tokens=512,
+        cudagraph_capture_sizes=[7, 96, 1024],
+        hidden_size=HIDDEN_SIZE,
+        hc_mult=HC_MULT,
+        num_sms=NUM_SMS,
+    )
+    assert selected == sorted(set(selected))
+    assert max(selected) <= 512
+    assert {7, 96}.issubset(selected)
+    assert 1024 not in selected
+
+
+def test_non_positive_max_tokens_selects_nothing():
+    assert (
+        _select_mhc_warmup_token_sizes(
+            max_tokens=0,
+            cudagraph_capture_sizes=[64],
+            hidden_size=HIDDEN_SIZE,
+            hc_mult=HC_MULT,
+            num_sms=NUM_SMS,
+        )
+        == []
+    )