[Core] Shared abstractions for specialized decode kernels

.buildkite/test_areas/models_basic.yaml

@@ -76,10 +76,11 @@ steps:
   - tests/models/test_deepseek_v41_pipeline.py
   - tests/models/transformers/test_layer_registry_e2e.py
   - tests/models/test_deepseek_v41_replay_start.py
+  - tests/models/common/specialized_decode/
   - tests/models/test_deepseek_v4_vl_stream_weights.py
   - tests/models/test_muse_spark_rocm_loading.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v41_pipeline.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_rocm_loading.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v41_pipeline.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_rocm_loading.py models/common/specialized_decode
   mirror:
     amd:
       label: ":amd: (MI355) Basic Models (Other)"
@@ -101,17 +102,19 @@ steps:
       - tests/models/test_deepseek_v4_dspark_rocm.py
       - tests/models/test_deepseek_v4_rocm_wo_a.py
       - tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py
-      - tests/models/test_dsv41_mono_decode.py
+      - tests/models/test_dsv41_specialized_decode.py
       - tests/models/test_muse_spark_rocm_loading.py
       - tests/models/test_deepseek_v4_vl_stream_weights.py
       - tests/models/hy_v4/test_triton_ihc.py
+      - vllm/models/common/specialized_decode/
+      - tests/models/common/specialized_decode/
       commands:
-      - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py models/test_muse_spark_rocm_loading.py models/test_deepseek_v4_vl_stream_weights.py
+      - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_specialized_decode.py models/test_muse_spark_rocm_loading.py models/test_deepseek_v4_vl_stream_weights.py models/common/specialized_decode
       - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed'
       - pytest -v -s models/hy_v4/test_triton_ihc.py
 
-- label: ":amd: (MI355) MiniMax M3 MonoKernel"
-  key: minimax-m3-monokernel-mi355
+- label: ":amd: (MI355) MiniMax M3 Specialized Decode"
+  key: minimax-m3-specialized-decode-mi355
   dind: false
   timeout_in_minutes: 20
   device: mi355_4

tests/engine/test_arg_utils.py

@@ -805,6 +805,55 @@ def test_ir_op_priority():
         ).create_engine_config()
 
 
+@pytest.mark.parametrize(
+    ("argv", "expected"),
+    [
+        ([], None),
+        (["--specialized-decode"], "auto"),
+        (["--specialized-decode", "auto"], "auto"),
+        (["--specialized-decode", "monokernel"], "monokernel"),
+        (["--monokernel"], "monokernel"),
+        (["--no-monokernel"], None),
+        # The named backend wins over the shorthand, whatever the order.
+        (["--monokernel", "--specialized-decode", "auto"], "auto"),
+        (["--specialized-decode", "auto", "--monokernel"], "auto"),
+    ],
+)
+def test_specialized_decode_cli(argv, expected):
+    """``--monokernel`` is shorthand for ``--specialized-decode monokernel``,
+    and a bare ``--specialized-decode`` means ``auto``."""
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    args = EngineArgs.from_cli_args(parser.parse_args(argv))
+    config = args.create_engine_config()
+    assert config.kernel_config.specialized_decode == expected
+
+
+def test_specialized_decode_rejects_an_unknown_backend():
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    with pytest.raises(SystemExit):
+        parser.parse_args(["--specialized-decode", "flydsl"])
+
+
+@pytest.mark.parametrize(
+    ("flags", "expected"),
+    [
+        ({}, "monokernel"),
+        ({"monokernel": False}, None),
+        ({"specialized_decode": "auto"}, "auto"),
+    ],
+)
+def test_specialized_decode_overrides_the_kernel_config(flags, expected):
+    """``--no-monokernel`` turns off a specialized decode that a config file or
+    the Python API asked for, and absent flags leave it alone."""
+    from vllm.config.kernel import KernelConfig
+
+    config = EngineArgs(
+        kernel_config=KernelConfig(specialized_decode="monokernel"),
+        **flags,
+    ).create_engine_config()
+    assert config.kernel_config.specialized_decode == expected
+
+
 @pytest.mark.parametrize(
     ("input_json", "expected_json"),
     [

tests/kernels/models/amd/test_dsv41_gfx942_draft_lm_head.py

@@ -36,7 +36,7 @@ def _set_gate(monkeypatch, on: bool):
 
     # The model records --monokernel when it is created. These tests have no
     # model, so they set the recorded flag themselves.
-    monkeypatch.setattr(ops, "_monokernel", on)
+    monkeypatch.setattr(ops, "_specialized_decode", on)
     ops.enabled.cache_clear()
 
 

tests/kernels/models/amd/test_dsv41_gfx942_fused_markov.py

@@ -37,7 +37,7 @@ def _set_gate(monkeypatch, on: bool):
 
     # The model records --monokernel when it is created. These tests have no
     # model, so they set the recorded flag themselves.
-    monkeypatch.setattr(ops, "_monokernel", on)
+    monkeypatch.setattr(ops, "_specialized_decode", on)
     ops.enabled.cache_clear()
 
 

tests/kernels/models/amd/test_dsv41_gfx942_sparse_indexer.py

@@ -64,7 +64,7 @@ def gfx942_step(monkeypatch):
     if paged_mqa_logits_module() is None:
         pytest.skip("AITER's paged MQA logits module is not installed")
     monkeypatch.setattr(rocm_aiter_ops, "_AITER_ENABLED", True)
-    monkeypatch.setattr(ops, "_monokernel", True)
+    monkeypatch.setattr(ops, "_specialized_decode", True)
     ops.enabled.cache_clear()
     if not is_workspace_manager_initialized():
         init_workspace_manager(torch.device("cuda"))

tests/kernels/models/amd/test_dsv41_gfx942_topk.py

@@ -45,7 +45,7 @@ def gfx942_decode(monkeypatch):
     the recorded flag itself."""
     from vllm.models.deepseek_v41.amd import ops
 
-    monkeypatch.setattr(ops, "_monokernel", True)
+    monkeypatch.setattr(ops, "_specialized_decode", True)
     ops.enabled.cache_clear()
     yield
     ops.enabled.cache_clear()

tests/models/common/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/models/common/specialized_decode/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/models/common/specialized_decode/kit.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""What a model's ``DecodeSpec`` must satisfy."""
+
+import subprocess
+import sys
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any
+
+import pytest
+
+from vllm.config.compilation import CUDAGraphMode
+from vllm.forward_context import ForwardContext, override_forward_context
+from vllm.models.common.specialized_decode.dispatch import SpecializedDecode
+from vllm.models.common.specialized_decode.guards import select_decode_spec
+
+KERNEL_BACKENDS = ("aiter", "flydsl")
+"""Backends a spec module must not pull in of its own: it is imported at config
+time, before anything has decided the spec can run. The package holding it is
+the model's own hardware entry point, which already pulls its platform's
+backend, so only what the spec module adds on top of that package counts."""
+
+
+@contextmanager
+def forward(config, mode: CUDAGraphMode, attn_metadata: Any) -> Iterator[None]:
+    """One forward pass's context, as a runner would set it."""
+    context = ForwardContext(
+        no_compile_layers=config.compilation_config.static_forward_context,
+        attn_metadata=attn_metadata,
+        slot_mapping={},
+        cudagraph_runtime_mode=mode,
+    )
+    with override_forward_context(context):
+        yield
+
+
+class DecodeConformance:
+    """What every model's ``DecodeSpec`` must satisfy. Subclass it in the
+    model's test module and fill in the hooks below."""
+
+    specs: tuple[Any, ...]
+    """The model's variants, as it passes them to ``select_decode_spec``."""
+
+    specs_module: str
+    """Where ``specs`` is defined, for the import test to load on its own."""
+
+    def make_config(self, **overrides) -> Any:
+        """A config this spec serves, with ``overrides`` applied on top."""
+        raise NotImplementedError
+
+    def make_model(self, config) -> Any:
+        """A loaded model this spec serves, with its KV caches bound into
+        ``config.compilation_config.static_forward_context``."""
+        raise NotImplementedError
+
+    def bad_caches(self, config, model) -> dict[str, Any]:
+        """Caches that break at least two of ``refuse_caches``'s checks."""
+        raise NotImplementedError
+
+    @pytest.fixture
+    def config(self):
+        return self.make_config()
+
+    @pytest.fixture
+    def model(self, config):
+        return self.make_model(config)
+
+    @pytest.fixture
+    def decode(self, config, model):
+        return self.make_decode(config, model)
+
+    def make_decode(self, config, model=None, bind: bool = True):
+        """This rank's ``SpecializedDecode``, built the way a model builds it:
+        the variant picked once by ``select_decode_spec``, then the dispatch
+        over the loaded model, then the caches the runner binds."""
+        model = self.make_model(config) if model is None else model
+        spec = select_decode_spec(config, self.specs)
+        assert spec is not None
+        decode = SpecializedDecode(model, config, spec)
+        if bind:
+            decode.bind_caches()
+        return decode
+
+    def test_spec_module_imports_no_kernel_backend(self):
+        module = self.specs_module
+        package = module.rpartition(".")[0]
+        source = (
+            "import sys; sys.path.insert(0, '.'); import importlib;"
+            f"importlib.import_module({package!r});"
+            "loaded = set(sys.modules);"
+            f"importlib.import_module({module!r});"
+            "print('\\n'.join(set(sys.modules) - loaded))"
+        )
+        loaded = subprocess.run(
+            [sys.executable, "-c", source],
+            capture_output=True,
+            text=True,
+            check=True,
+        ).stdout.split()
+        pulled = {n.split(".")[0] for n in loaded} & set(KERNEL_BACKENDS)
+        assert not pulled, f"{module} imports {sorted(pulled)}"
+
+    def test_refuse_caches_reports_every_reason(self, config, model, decode):
+        reasons = decode.spec.refuse_caches(
+            decode.vllm_config, decode.kernel, self.bad_caches(config, model)
+        )
+        assert len(reasons) >= 2, reasons
+
+    def test_build_does_not_move_weights(self, config):
+        model = self.make_model(config)
+        before = {n: p.data_ptr() for n, p in model.named_parameters()}
+        self.make_decode(config, model)
+        after = {n: p.data_ptr() for n, p in model.named_parameters()}
+        assert before == after